网站设计应遵循的原则,企业3合1网站建设公司,做包装看什么网站,哪个地方网站建设的公司多首先#xff0c;本文回顾了MHA的计算方式以及KV Cache的原理#xff0c;然后深入到了DeepSeek V2的MLA的原理介绍#xff0c;同时对MLA节省的KV Cache比例做了详细的计算解读。接着#xff0c;带着对原理的理解理清了HuggingFace MLA的全部实现#xff0c;每行代码都去对应… 首先本文回顾了MHA的计算方式以及KV Cache的原理然后深入到了DeepSeek V2的MLA的原理介绍同时对MLA节省的KV Cache比例做了详细的计算解读。接着带着对原理的理解理清了HuggingFace MLA的全部实现每行代码都去对应了完整公式中的具体行并且对每个操作前后的Tensor Shape变化也进行了解析。我们可以看到目前的官方实现在存储KV Cache的时候并不是存储隐向量而是把隐向量都解压缩变成了标准的MHA的KV Cache实际上是完全不能节省显存的。接着就继续学习了一下清华大学的ZHANG Mingxing组实现的MLA矩阵吸收的工程实现在这一节也详细分析了原理包括 W U K W^{UK} WUK和 W U V W_{UV} WUV分别如何吸收到 W U Q W_{UQ} WUQ和 W o W_o Wo中分析了实现了矩阵吸收的每行代码的原理以及操作发生前后相关Tensor的维度变化。接着对矩阵吸收代码实现里的矩阵乘法的性质进行分析可以看到MLA在大多数阶段都是计算密集型而非访存密集型的。最后引用了作者团队的Benchmark结果以及说明为何不是直接保存吸收后的大投影矩阵而是在forward里面重新计算两个矩阵的吸收。 这里提一下我维护的几个记录个人学习笔记以及社区中其它大佬们的优秀博客链接的仓库都获得了不少star感谢读者们的认可我也会继续在开源社区多做贡献。github主页https://github.com/BBuf 欢迎来踩 0x0. 前言
这篇文章主要是对Deepseek2提出的优化KV Cache的MLA方法做个人理解特别是关于MLA的矩阵吸收部分这部分Paper以及官方开源实现没有给出。然后开源社区大佬确实出手很快在知乎的《如何看待 DeepSeek 发布的 MoE 大模型 DeepSeek-V2》问题下清华大学的ZHANG Mingxing组就提出了MLA的矩阵吸收实现并且给出了一个兼容现有Transformers实现的PRhttps://huggingface.co/deepseek-ai/DeepSeek-V2-Chat/discussions/12。 翻译
这个修补过的DeepseekV2Model包含了对DeepseekV2Attention的以下修改以减少VRAM消耗并提高效率
不再缓存解压缩的Key/Value状态而仅缓存低秩Key-Value联合压缩以及Key的解耦RoPE部分。 为了重用transformers库的缓存实用程序我们将k_pe视为key_states将compressed_kv视为value_states。采用DeepseekV2论文中描述的吸收技术通过改变计算Query和Output向量时的乘法顺序。这不仅节省了中间张量的内存消耗还减少了浮点运算的次数。分别计算RoPE部分和非RoPE部分的注意力分数然后将它们相加。原始实现将Query/Key向量的两部分连接起来但在缓存压缩Key/Value状态时由于不必要的数据广播和内存往返而被证明效率较低。
通过应用上述更改MLA模块在解码阶段可以在NVIDIA A100-PCIE-40GB GPU上对单个请求实现高达20.4倍的加速对32个批处理请求实现3.63倍的加速在NVIDIA GeForce RTX 4080上对单个和批处理请求分别实现26.2倍和3.52倍的加速。
这篇博客就是以我的视角来理解清楚这个PR中矩阵吸收的工程实现以及为何它可以加速现有的Deepseek2 MLA实现。本文先回顾一下MHA的Decode阶段KV Cache由来然后根据paper的公式去理解Transformers中的DeepseekV2Attention类的实现。再接下来就学习一下ZHANG Mingxing大佬组所提出的MLA矩阵吸收工程实现。
Paper Linkhttps://arxiv.org/pdf/2405.04434
0x1. MHA 解码阶段KV Cache由来
首先回顾一下 MHA 机制在Decode阶段的原理和实现。 下面的公式来源也是DeepSeek2 paper我做了更精细的解释。 假设batch_size为1另外由于是解码阶段输入只有一个token所以序列的长度也是1所以输入可以表示为 h t ∈ R d h_t \in \mathbb{R}^d ht∈Rd。接着假设embedding词表维度为 d d d并且有 n h n_h nh表示注意力头的数量 d h d_h dh表示每个注意力头的维度。 t表示解码阶段当前是第几个token。 然后通过 W Q , W K , W V ∈ R d h n h × d W^Q, W^K, W^V \in \mathbb{R}^{d_h n_h \times d} WQ,WK,WV∈Rdhnh×d三个参数矩阵得到 q t , k t , v t ∈ R d h n h q_t, k_t, v_t \in \mathbb{R}^{d_h n_h} qt,kt,vt∈Rdhnh具体方法就是三个矩阵乘 q t W Q h t , k t W K h t , v t W V h t , q_t W^Q h_t, \newline k_t W^K h_t, \newline v_t W^V h_t, qtWQht,ktWKht,vtWVht,
在 MHA 的计算中这里的 q t , k t , v t q_t, k_t, v_t qt,kt,vt 又会分割成 n h n_h nh 个注意力头即 [ q t , 1 ; q t , 2 ; ⋯ ; q t , n h ] q t [ k t , 1 ; k t , 2 ; ⋯ ; k t , n h ] k t [ v t , 1 ; v t , 2 ; ⋯ ; v t , n h ] v t \begin{bmatrix} q_{t,1}; q_{t,2}; \cdots ; q_{t,n_h} \end{bmatrix} q_t \newline \begin{bmatrix} k_{t,1}; k_{t,2}; \cdots ; k_{t,n_h} \end{bmatrix} k_t \newline \begin{bmatrix} v_{t,1}; v_{t,2}; \cdots ; v_{t,n_h} \end{bmatrix} v_t [qt,1;qt,2;⋯;qt,nh]qt[kt,1;kt,2;⋯;kt,nh]kt[vt,1;vt,2;⋯;vt,nh]vt
这里 q t , i , k t , i , v t , i ∈ R d h q_{t,i}, k_{t,i}, v_{t,i} \in \mathbb{R}^{d_h} qt,i,kt,i,vt,i∈Rdh 分别表示query、key和value的第 i i i个头的计算结果。
接下来就是计算注意力分数和输出了公式如下 o t , i ∑ j 1 t Softmax j ( q t , i k j , i d h ) v j , i , u t W O [ o t , 1 ; o t , 2 ; ⋯ ; o t , n h ] o_{t,i} \sum_{j1}^{t} \text{Softmax}_j \left( \frac{q_{t,i} k_{j,i}}{\sqrt{d_h}} \right) v_{j,i}, \newline u_t W^O [o_{t,1}; o_{t,2}; \cdots ; o_{t,n_h}] ot,ij1∑tSoftmaxj(dh qt,ikj,i)vj,i,utWO[ot,1;ot,2;⋯;ot,nh]
这里 W O ∈ R d × d h n h W^O \in \mathbb{R}^{d \times d_h n_h} WO∈Rd×dhnh 表示输出映射矩阵。从上面的公式可以看出来对于当前的第 t t t 个 token的query会和 t t t之前所有token的key, value做注意力计算并且由于token by token的生成 t t t之前所的有token对应的 k k k v v v我们都可以Cache下来避免重复计算这就是KV Cache的由来。
对于一个 l l l层的标准MHA的网络来说每个token需要的KV Cache大小为 2 n h d h l 2n_hd_hl 2nhdhl其中2表示bf16的字节。
为了改进KV Cache演化了一系列AI Infra的工作比如Paged Attention, GQA, MLA包括最新的GQAMLA之外的另一种KV Cache压缩方式动态内存压缩DMCvAttention用于在没有Paged Attention的情况下Serving LLM 。
0x2. DeepSeek2 MLA 原理介绍 上面的图是Deepseek2 paper中对集中常见KV Cache压缩方法的对比可以看到MLA的核心是对keys和values进行低秩联合压缩来减少KV Cache。对应paper的公式9-11。 其中 c t K V ∈ R d c \mathbf{c}_{t}^{KV} \in \mathbb{R}^{d_c} ctKV∈Rdc表示对 key 和 value 压缩后的隐向量 latent vector 这里 d c ( ≪ d h n h ) d_c (\ll d_h n_h) dc(≪dhnh) 表示 KV Cache压缩的维度。 W D K V ∈ R d c × d \mathbf{W}^{DKV} \in \mathbb{R}^{d_c \times d} WDKV∈Rdc×d表示向下映射 down-projection 矩阵 W U K , W U V ∈ R d h n h × d c \mathbf{W}^{UK}, \mathbf{W}^{UV} \in \mathbb{R}^{d_h n_h \times d_c} WUK,WUV∈Rdhnh×dc表示向上映射 up-projection 矩阵
这样在推理时只需要缓存隐向量 c t K V \mathbf{c}_{t}^{KV} ctKV 即可因此 MLA 对应的每一个 token 的 KV Cache 参数只有 2 d c l 2d_c l 2dcl 个其中 l l l是网络层数 2 2 2是bfloat16的字节。
此外为了降低训练过程中的激活内存DeepSeek2还对query进行低秩压缩即便这并不能降低KV Cache 其中 c t Q ∈ R d c ′ \mathbf{c}_{t}^{Q} \in \mathbb{R}^{d_c} ctQ∈Rdc′表示将 queries 压缩后的隐向量 d c ′ ( ≪ d h n h ) d_c (\ll d_h n_h) dc′(≪dhnh) 表示 query 压缩后的维度 W D Q ∈ R d c ′ × d , W U Q ∈ R d h n h × d c ′ \mathbf{W}^{DQ} \in \mathbb{R}^{d_c \times d}, \mathbf{W}^{UQ} \in \mathbb{R}^{d_h n_h \times d_c} WDQ∈Rdc′×d,WUQ∈Rdhnh×dc′ 分别表示 down-projection 和 up-projection 矩阵
接下来MLA讨论的一个问题是在上面的压缩过程中我们并没有考虑到RoPE。原始的RoPE需要在query和key中融入相对位置信息。在MLA中在query中融入相对位置信息是比较容易的但是由于KV Cache缓存的是压缩后的低秩key-value信息这里面是没办法融入相对位置信息的。 关于RoPE为何不兼容MLA苏神的博客里有更深刻的解释建议阅读 https://kexue.fm/archives/10091 下面是对 Paper 的 Decoupled Rotary Position Embedding 章节进行解释。 识别图中的所有文本和公式如下
由于对 query 和 key 来说RoPE 都是位置敏感的。如果对 k t C \mathbf{k}_{t}^{C} ktC 采用 RoPE那么当前生成 token 相关的 RoPE 矩阵会在 W Q \mathbf{W}^{Q} WQ 和 W U K \mathbf{W}^{UK} WUK 之间并且矩阵乘法不遵循交换律因此在推理时 W U K \mathbf{W}^{UK} WUK 就无法整合到 W Q \mathbf{W}^{Q} WQ 中。这就意味着推理时我们必须重新计算所有之前 tokens 的 keys这将大大降低推理效率。 这里的 W U K \mathbf{W}^{UK} WUK 就整合到 W Q \mathbf{W}^{Q} WQ 请看下面截图的解释来自苏神的博客。我会在下一大节再仔细讨论这个原理。 因此DeepSeek2提出了解耦 RoPE 策略具体来说
使用额外的多头 queries q t , i R ∈ R d h R \mathbf{q}_{t, i}^{R} \in \mathbb{R}^{d_h^R} qt,iR∈RdhR 以及共享的 key k t R ∈ R d h R \mathbf{k}_{t}^{R} \in \mathbb{R}^{d_h^R} ktR∈RdhR 来携带 RoPE 信息其中 d h R d_h^R dhR 表示解耦的 queries 和 key 的一个 head 的维度。
基于这种解耦的 RoPE 策略MLA 遵循的计算逻辑为 其中 W Q R ∈ R d h R n h × d c ′ \mathbf{W}^{QR} \in \mathbb{R}^{d_h^R n_h \times d_c} WQR∈RdhRnh×dc′ 和 W K R ∈ R d h R × d \mathbf{W}^{KR} \in \mathbb{R}^{d_h^R \times d} WKR∈RdhR×d 分别表示计算解耦后的 queries 和 key 的矩阵RoPE( ⋅ \cdot ⋅) 表示应用 RoPE 的操作 [ ⋅ ; ⋅ ] [ \cdot ; \cdot ] [⋅;⋅] 表示拼接操作
推理时只需要缓存解耦后的 key 即可当然还有上面的隐向量 c t K V \mathbf{c}_{t}^{KV} ctKV因此对于单个 token 的 KV Cache 只包含 ( d c d h R ) l (d_c d_h^R)l (dcdhR)l 个元素这里没考虑层数和bf16的字节数。具体可以看一下Table 1的数据对比 翻译一下 表1 | 各种注意力机制中每个token的KV Cache对比。 n h n_h nh 表示注意力头的数量 d h d_h dh 表示每个注意力头的维度 l l l 表示层数 n g n_g ng 表示GQA中的组数 d c d_c dc 和 d h R d_h^R dhR 分别表示KV压缩维度和MLA中解耦后queries和key的每头维度。KV Cache的数量以元素的数量来衡量而不考虑存储精度。对于DeepSeek-V2 d c d_c dc 被设置为 4 d h 4d_h 4dh 而 d h R d_h^R dhR 被设置为 d h 2 \frac{d_h}{2} 2dh。因此其KV Cache等于只有2.25组的GQA但其性能强于MHA。 原理的话应该就是这些了接下来就带着原理阅读DeepseekV2Attention的实现。
这里再特别说明一下Paper中相比于Dense的Deepseek 67B或者LLaMa3 70B节省93.3% KV Cache的计算方法 首先是层数DeepSeek2是60层而Deepseek 67B为95层层数的节省比例为 60 / 95
然后是单层的KV Cache比例是(4.5 x 128) / (2 x 8 x 128)其中2表示K和V8表示num_attention_heads128表示head_size4.5则是上面MLA中的9/2的压缩隐向量。
此外DeepSeek2针对KV Cache使用了6Bit量化节省比例为 6 / 16
把这三个比例乘起来再用1减掉就是93.3%的由来。 这里的6bit感觉是考虑了量化参数 scale 和 zero-point, 如果采用4bit量化, 而scale / zero-point 为fp32, 则当group_size32时, 根据group-wise量化规则每32个元素对应一组float32的scale和zero_point那么每个元素平摊的位宽就多了2Bit量化位宽等同于6Bit。 0x3. MLA HuggingFace官方实现代码解读
为了便于描述代码这里直接把完整的公式贴出来Paper的附录C 对照上面的原理介绍下面对 DeepseekV2Attention 模块进行解读代码链接https://huggingface.co/deepseek-ai/DeepSeek-V2-Chat/blob/main/modeling_deepseek.py#L680
首先来看初始化部分为了快速理解代码直接忽略掉RoPE计算相关的代码部分。
# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama-DeepseekV2
class DeepseekV2Attention(nn.Module):Multi-headed attention from Attention Is All You Need paperdef __init__(self, config: DeepseekV2Config, layer_idx: Optional[int] None):super().__init__()self.attention_dropout config.attention_dropoutself.hidden_size config.hidden_sizeself.num_heads config.num_attention_headsself.max_position_embeddings config.max_position_embeddingsself.rope_theta config.rope_theta# 对应 query 压缩后的隐向量的维度 d_cself.q_lora_rank config.q_lora_rank# 对应$d_h^R$ 表示应用了rope的 queries 和 key 的一个 head 的维度。self.qk_rope_head_dim config.qk_rope_head_dim# 对应 key-value 压缩后的隐向量维度 d_cself.kv_lora_rank config.kv_lora_rank# value 的一个注意力头的隐藏层为度self.v_head_dim config.v_head_dim# 表示query和key的隐藏向量中应用rope部分的维度self.qk_nope_head_dim config.qk_nope_head_dim# 每一个注意力头的维度应该是两部分只和self.q_head_dim config.qk_nope_head_dim config.qk_rope_head_dimself.is_causal True# MLA 中对 Q 投影矩阵也做了一个低秩分解对应生成 q_a_proj 和 q_b_proj 两个矩阵。# 其中 q_a_proj 大小为 [hidden_size, q_lora_rank] [5120, 1536]# 对应上面公式中的W^DQself.q_a_proj nn.Linear(self.hidden_size, config.q_lora_rank, biasconfig.attention_bias)self.q_a_layernorm DeepseekV2RMSNorm(config.q_lora_rank)# q_b_proj 大小为 [q_lora_rank, num_heads * q_head_dim] # [q_lora_rank, num_attention_heads * (qk_nope_head_dim qk_rope_head_dim)] [1536, 128*(12864)] [1536, 24576] # 对应上述公式中的W^UQ和W^QR合并后的大矩阵self.q_b_proj nn.Linear(config.q_lora_rank, self.num_heads * self.q_head_dim, biasFalse)# 与Q向量类似KV向量的生成也是先投影到一个低维的 compressed_kv 向量对应c_t^{KV}# 再升维展开。具体的代码涉及 kv_a_proj_with_mqa 和 kv_b_proj 两个参数矩阵。# 其中 kv_a_proj_with_mqa 大小为 [hidden_size kv_lora_rank qk_rope_head_dim]# [5120, 512 64] [5120, 576]对应上述公式中的W^{DKV}和W^{KR}。self.kv_a_proj_with_mqa nn.Linear(self.hidden_size,config.kv_lora_rank config.qk_rope_head_dim,biasconfig.attention_bias,)self.kv_a_layernorm DeepseekV2RMSNorm(config.kv_lora_rank)# kv_b_proj 大小为 [kv_lora_rank num_heads * (q_head_dim - qk_rope_head_dim v_head_dim)] # [512, 128*((12864)-64128)] [512, 32768]对应上述公式中的W^{UK}和W^{UV}。# 由于 W^{UK} 只涉及 non rope 的部分所以维度中把 qk_rope_head_dim 去掉了。self.kv_b_proj nn.Linear(config.kv_lora_rank,self.num_heads* (self.q_head_dim - self.qk_rope_head_dim self.v_head_dim),biasFalse,)# 对应完整公式的第 47 行self.o_proj nn.Linear(self.num_heads * self.v_head_dim,self.hidden_size,biasconfig.attention_bias,)根据0x2节的原理介绍现在已经可以把 DeepseekV2Attention 模块里面所有的权重矩阵都和初始化的代码对应起来了如果你想继续看下去一定要理解到初始化的每行代码。
为了方便理解forward代码时回看公式这里再重复贴一下完整公式 接下来再看一下forward的代码这对应了完整公式里面的详细计算过程
def forward(self,hidden_states: torch.Tensor,attention_mask: Optional[torch.Tensor] None,position_ids: Optional[torch.LongTensor] None,past_key_value: Optional[Cache] None,output_attentions: bool False,use_cache: bool False,**kwargs,) - Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:if padding_mask in kwargs:warnings.warn(Passing padding_mask is deprecated and will be removed in v4.37. Please make sure use attention_mask instead.)# hidden_states对应公式中的h_t的shape是(batch_size, seq_length, # hidden_size)其中 hidden_size 具体为 5120假设batch_size和seq_length都为1bsz, q_len, _ hidden_states.size()# 计算Q对应完整公式中的 37-39 行先降维再升维好处是相比直接使用大小为 [5120, 24576] 的矩阵# [5120, 1536] * [1536, 24576] 这样的低秩分解在存储空间和计算量上都大幅度降低q self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))q q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)# 切分 rope 和非 rope 部分完整公式中 40 行反过来q_nope, q_pe torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim-1)# 对应公式中的 41 和 43 行只是还没有加 rope# 一个优化的 MLA KVCache 实现只需要缓存这个 compressed_kv 就行# kv_a_proj_with_mqa shape 为[hidden_size kv_lora_rank qk_rope_head_dim]# [5120, 512 64] [5120, 576]# 所以compressed_kv的shape就是[1, 1, 576]compressed_kv self.kv_a_proj_with_mqa(hidden_states)# 对应完整公式的 44 行反过来compressed_kv, k_pe torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim-1)# 这里的 k_pe 和 上面的 q_pe 要扔给 RoPE模块所以需要重整下shapek_pe k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)# 对应公式中的 42 和 45 行将 MLA 展开成标准 MHA 的形式kv (self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(bsz, q_len, self.num_heads, self.qk_nope_head_dim self.v_head_dim).transpose(1, 2))# 因为 kv_b_proj 打包了 W^{UK} 和 W^{UV} 把他们分离出来k_nope, value_states torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim-1)# 获取key/value的序列长度kv_seq_len value_states.shape[-2]if past_key_value is not None:if self.layer_idx is None:raise ValueError(fThe cache structure has changed since version v4.36. If you are using {self.__class__.__name__} for auto-regressive decoding with k/v caching, please make sure to initialize the attention class with a layer index.)kv_seq_len past_key_value.get_usable_length(kv_seq_len, self.layer_idx)# 给需要 rope 的部分加 ropecos, sin self.rotary_emb(value_states, seq_lenkv_seq_len)q_pe, k_pe apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)# 更新和拼接历史 KVCache可以看到这里存储的是展开后的 MHA KVCache# 其中 q_head_dim 等于 qk_nope_head_dim qk_rope_head_dimquery_states k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)query_states[:, :, :, : self.qk_nope_head_dim] q_nopequery_states[:, :, :, self.qk_nope_head_dim :] q_pekey_states k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)key_states[:, :, :, : self.qk_nope_head_dim] k_nopekey_states[:, :, :, self.qk_nope_head_dim :] k_pe# Transformers库中标准的 KV Cache 更新代码if past_key_value is not None:cache_kwargs {sin: sin, cos: cos} # Specific to RoPE modelskey_states, value_states past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)# 后续就是标准的多头自注意力计算了为了篇幅忽略这部分代码...读完这一节基本就能将MLA完整公式的每一行都搞清楚了并且我们可以看到目前的官方实现在存储KV Cache的时候并不是存储隐向量而是把隐向量都解压缩变成了标准的MHA的KV Cache实际上是完全不能节省显存的。
0x4. 矩阵吸收
这一节就是学习和理解一下清华大学的ZHANG Mingxing组实现的MLA矩阵吸收 https://zhuanlan.zhihu.com/p/700214123 。它的代码是直接应用在HF实现上的所以可以很方便进行应用。为了理解再重复贴一下完整公式 以及Paper中提到的矩阵吸收 0x4.1 W^{UK}的吸收
对于 W U K \mathbf{W}^{UK} WUK 矩阵我们有 atten_weights q t ⊤ k t ( W U Q c t Q ) ⊤ W U K c t K V c t Q ⊤ W U Q ⊤ W U K c t K V \text{atten\_weights} \mathbf{q}_{t}^\top \mathbf{k}_{t} (\mathbf{W}^{UQ} \mathbf{c}_{t}^{Q})^\top \mathbf{W}^{UK} \mathbf{c}_{t}^{KV} \mathbf{c}_{t}^{Q^\top} \mathbf{W}^{UQ^\top} \mathbf{W}^{UK} \mathbf{c}_{t}^{KV} atten_weightsqt⊤kt(WUQctQ)⊤WUKctKVctQ⊤WUQ⊤WUKctKV
也就是说我们实际上不需要将低维的 c t K V \mathbf{c}_{t}^{KV} ctKV 展开再计算而是直接将 W U K \mathbf{W}^{UK} WUK 通过结合律先和左边做乘法。
# 以下和上一节的MLA forward部分实现相同
# hidden_states对应公式中的h_t的shape是(batch_size, seq_length,
# hidden_size)其中 hidden_size 具体为 5120假设batch_size为1seq_length为q_len
bsz, q_len, _ hidden_states.size()# 计算Q对应完整公式中的 37-39 行先降维再升维好处是相比直接使用大小为 [5120, 24576] 的矩阵
# [5120, 1536] * [1536, 24576] 这样的低秩分解在存储空间和计算量上都大幅度降低
q self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
q q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)
# 切分 rope 和非 rope 部分完整公式中 40 行反过来
q_nope, q_pe torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim-1
)# 对应公式中的 41 和 43 行只是还没有加 rope
# 一个优化的 MLA KVCache 实现只需要缓存这个 compressed_kv 就行不过后面实际上展开
# hidden_states 的 shape 为 (1, past_len, hidden_size)
# kv_a_proj_with_mqa shape 为[hidden_size kv_lora_rank qk_rope_head_dim]
# [5120, 512 64] [5120, 576]
# 所以compressed_kv的shape就是[1, past_len, 576]
compressed_kv self.kv_a_proj_with_mqa(hidden_states)
# 对应完整公式的 44 行反过来
compressed_kv, k_pe torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim-1
)
# 这里的 k_pe 和 上面的 q_pe 要扔给 RoPE模块所以需要重整下shape
k_pe k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)以下部分是 W U K W^{UK} WUK的吸收需要做的改动省略掉了把compressed_kv和k_pe加入到Transformers KV Cache的改动
# 从 kv_b_proj 中分离的 W^{UK} 和 W^{UV} 两部分他们要分别在不同的地方吸收
kv_b_proj self.kv_b_proj.weight.view(self.num_heads, -1, self.kv_lora_rank)
q_absorb kv_b_proj[:, :self.qk_nope_head_dim,:]
out_absorb kv_b_proj[:, self.qk_nope_head_dim:, :]cos, sin self.rotary_emb(q_pe)
q_pe apply_rotary_pos_emb(q_pe, cos, sin, q_position_ids)
# W^{UK} 即 q_absorb 被 q_nope 吸收
q_nope torch.einsum(hdc,bhqd-bhqc, q_absorb, q_nope)
# 吸收后 attn_weights 直接基于 compressed_kv 计算不用展开。
attn_weights torch.matmul(q_pe, k_pe.transpose(2, 3)) torch.einsum(bhqc,blc-bhql, q_nope, compressed_kv)
attn_weights * self.softmax_scale主要是吸收这里的代码需要仔细理解
从0x4节的讲解已经知道kv_b_proj就是 W U K W^{UK} WUK 和 W U V W^{UV} WUV两部分这里是把 W U K W^{UK} WUK吸收到 W U Q W^{UQ} WUQ所以需要先把两者分离出来。注意到 self.kv_b_proj weight shape为 [kv_lora_rank num_heads * (q_head_dim - qk_rope_head_dim v_head_dim)] [512, 128*((12864)-64128)] [512, 32768]所以kv_b_proj的shape为[num_heads,q_head_dim - qk_rope_head_dim v_head_dim , kv_lora_rank], q_absorb的shape为[num_heads, qk_nope_head_dim , kv_lora_rank][128, 128, 512]同样out_absorb的shape为[num_heads, v_head_dim , kv_lora_rank][128, 128, 512]。q_nope torch.einsum(hdc,bhqd-bhqc, q_absorb, q_nope) 这行代码中q_nope的shape是[batch_size, num_heads, q_len, q_head_dim]。所以这行代码就是一个矩阵乘法把 W U K W^{UK} WUK吸收到 W U Q W^{UQ} WUQ。吸收后 attn_weights 直接基于 compressed_kv 计算不用展开。对应torch.einsum(bhqc,blc-bhql, q_nope, compressed_kv)这行代码。其中q_nope 的维度是 [batch_size, num_heads, q_len, kv_lora_rank]compressed_kv 是 [batch_size, past_len, kv_lora_rank]输出的维度是 [batch_size, num_heads, q_len, past_len]。此外我们还可以观察到torch.matmul(q_pe, k_pe.transpose(2, 3))这行代码是分开计算了RoPE部分的q和k的注意力计算再求和没有和原本的实现一样将加上了 rope 的 q_pe/k_pe 和没加 rope 的 q_nope/k_nope 拼接起来一起也就是下面的代码。作者团队把这个拆分叫做 Move Elision 的优化后续有性能对比。
# 更新和拼接历史 KVCache可以看到这里存储的是展开后的 MHA KVCache
# 其中 q_head_dim 等于 qk_nope_head_dim qk_rope_head_dim
query_states k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
query_states[:, :, :, : self.qk_nope_head_dim] q_nope
query_states[:, :, :, self.qk_nope_head_dim :] q_pekey_states k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
key_states[:, :, :, : self.qk_nope_head_dim] k_nope
key_states[:, :, :, self.qk_nope_head_dim :] k_pe除了压缩KV Cache之外我们还可以观察到上面涉及到的2个矩阵乘法实际上都来到了计算密集的领域例如对于 torch.einsum(bhqc,blc-bhql, q_nope, compressed_kv) 。由于不同 head 的 q_nope 部分 share 了共同的 compressed_kv 部分实际计算的是 batch_size 个 [num_heads * q_len, kv_lora_rank] 和 [past_len, kv_lora_rank] 的矩阵乘法。计算等价于一个 MQA 操作计算强度正比于 num_heads 的也就是 128。 0x4.2 W^{UV}的吸收
对于 W U V W^{UV} WUV我们有 v W U V c t K V v W^{UV}c_t^{KV} vWUVctKV (对应公式的45行) u u u attn_weights 和 v v v 的矩阵乘法 o u W o o uW_o ouWo
这里把 attn_weights 记作 M M M那么有 o M W U V c t K V W o o MW^{UV}c_t^{KV}W_o oMWUVctKVWo
和 W U K W^{UK} WUK的吸收过程类似利用结合律改变计算顺序那么有 o M c t K V W U V W o o Mc_t^{KV}W^{UV}W_o oMctKVWUVWo
我们可以把 W U V W^{UV} WUV吸收到 W o W_o Wo中对应的代码实现
# attn_weight的shape是[batch_size, num_heads, q_len, past_len]
# compressed_kv的shape是[batch_size, past_len, kv_lora_rank]
# attn_output的shape是[batch_size, num_heads, q_len, kv_lora_rank]
attn_output torch.einsum(bhql,blc-bhqc, attn_weights, compressed_kv)
# out_absorb的shape是[num_heads, v_head_dim , kv_lora_rank]
# out_absorb.mT的shape是[num_heads, kv_lora_rank, v_head_dim]
# 最终attn_output的shape是[batch_size, num_heads, q_len, v_head_dim]
attn_output torch.matmul(attn_output, out_absorb.mT)注意.mT 方法用于获取张量的转置transpose。对于二维张量矩阵转置操作会交换其行和列。而对于高维张量.mT 会交换最后两个维度。 同样这里除了压缩KV Cache之外我们还可以观察到上面涉及到的2个矩阵乘法实际上也来到了计算密集的领域例如对于 attn_output torch.einsum(bhql,blc-bhqc, attn_weights, compressed_kv) 。由于不同 head 的 attn_weights 部分 share 了共同的 compressed_kv 部分实际计算的是 batch_size 个 [num_heads * q_len, kv_lora_rank] 和 [past_len, kv_lora_rank] 的矩阵乘法。计算等价于一个 MQA 操作计算强度正比于 num_heads 的也就是 128。因此相比 MHA吸收后的 MLA 计算强度要大得多因此也可以更加充分的利用 GPU 算力。 0x4.3 MLA MatMul的性质
上面几乎分析了每个矩阵乘法的计算shape可以发现除了在对q做计算时涉及到gemv之外也就是q self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))其它地方的矩阵乘运算q_len维度都是和num_heads在一起做计算而num_heads在Deepseek2的配置里面已经是128了导致其它的Matmul几乎都落在了计算密集的范畴。
综上对于MLA模块来说有很大比例的MatMul都达到了计算密集的范畴这改变了之前MHA的访存密集的性质。然而在整个网络中间由于有MoE模块的存在如果BatchSize不够大无法激活所有的expert导致计算和访存比计算密度低还是无法整体达到计算密集的范围但MLA节省的KV Cache本就可以让DeepSeek2的Batch大幅度提升所以均摊到每个token的带宽需求相比于Dense的LLaMa3 70B也会大幅度下降。
0x4.4 Benchmark
最后作者团队在他们的Blog中给出了一些Benchmark结果可以看到这个矩阵吸收的有效性。 图中的标签分别表示 原始的解压缩版本CacheDecompressed (CD)KV缓存压缩后的CacheCompressedCC吸收后直接使用 compressed_kv 计算的 Absorbed_CacheCompressed (A_CC) 版本和增加了 move elision 优化的最终版本 Absorbed_CacheCompressed_MoveElision (A_CC_ME)。
0x4.5 是否需要重计算
以 W U K W^{UK} WUK的吸收为例子这里实际上是说在矩阵吸收的时候是否要提前把 W U K W^{UK} WUK和 W U Q W^{UQ} WUQ的矩阵乘结果保存下来而不是在forward的时候重计算。作者在评论区回复过这个问题意思就是直接在forward的时候重计算速度会更优。博客里面的解释如下 0x5. 总结
这就是本篇博客的全部内容了这里总结一下。首先本文回顾了MHA的计算方式以及KV Cache的原理然后深入到了DeepSeek V2的MLA的原理介绍同时对MLA节省的KV Cache比例做了详细的计算解读。接着带着对原理的理解理清了HuggingFace MLA的全部实现每行代码都去对应了完整公式中的具体行并且对每个操作前后的Tensor Shape变化也进行了解析。我们可以看到目前的官方实现在存储KV Cache的时候并不是存储隐向量而是把隐向量都解压缩变成了标准的MHA的KV Cache实际上是完全不能节省显存的。接着就继续学习了一下清华大学的ZHANG Mingxing组实现的MLA矩阵吸收的工程实现在这一节也详细分析了原理包括 W U K W^{UK} WUK和 W U V W_{UV} WUV分别如何吸收到 W U Q W_{UQ} WUQ和 W o W_o Wo中分析了实现了矩阵吸收的每行代码的原理以及操作发生前后相关Tensor的维度变化。接着对矩阵吸收代码实现里的矩阵乘法的性质进行分析可以看到MLA在大多数阶段都是计算密集型而非访存密集型的。最后引用了作者团队的Benchmark结果以及说明为何不是直接保存吸收后的大投影矩阵而是在forward里面重新计算两个矩阵的吸收。
0x6. 参考资料
https://www.zhihu.com/question/655172528https://arxiv.org/pdf/2405.04434
