DeepSeekV2之MLA

鱿鱼圈 Lv4
  2026-04-20 23:50 2026-04-20 23:50 创建   2026-06-30 22:44:58 2026-06-30 22:44:58 更新      

论文链接：arxiv.org/pdf/2405.04434

MLA(Multi-head Latent Attention)完整详解

结合 DeepSeek-V2 / V3 论文数学公式、知乎讲解(zhuanlan.zhihu.com/p/16730036197)与 sglang 代码实现(python/sglang/srt/models/deepseek_v2.py、deepseek_common/attention_forward_methods/forward_mla.py)。

一句话:MLA 通过把 KV 压缩成一个低秩潜向量(latent)来减少 KV-cache 显存,再借助"矩阵吸收"在解码时退化成 MQA 的访存/计算模式,从而做到又省、又快、又(几乎)不掉质量。

公式以 LaTeX 书写( 为行内, 为块级),在支持 KaTeX/MathJax 的 Markdown 渲染器中可正常显示。

img

第一部分:为什么需要 MLA —— KV-cache 是推理瓶颈

1.1 LLM 推理的两个阶段

  • Prefill(预填充):对整个 prompt 的所有 token 一次性并行计算,产出第一个输出 token。计算密集。
  • Decode(解码):逐个 token 生成(token-by-token),直到 EOS。访存密集。

1.2 KV-cache 的由来:用空间换时间

LLM 是 Causal Attention(因果注意力):位置 的 token 只与它前面的 token 交互。用一个下三角 mask 实现。

观察两个事实:

  1. 前面 token 的 不受后面 token 影响;
  2. 每生成一个新 token,都要用到所有前序 token 的 ,而这些值始终不变。

所以为避免重复计算,把前序的 缓存起来 —— 这就是 KV-cache。

代价:它把"计算密集型"任务变成了"访存密集型"任务。 不再现算,而是从显存(HBM)读出来,GPU 核心和显存之间频繁读写 → 访存速率直接决定推理速度。

1.3 显存里存了什么

推理阶段显存主要三部分:

  1. 模型参数:固定大小。
  2. KV-cache:随 (Batch × Seq_len)动态增长 —— 长上下文 / 高并发下会爆炸。
  3. 运行时中间数据:即用即释放,很小。

1.4 KV-cache 到底多大(以 Qwen-72B 朴素 MHA 估算)

配置:层数 ,每层 个头,每头维度 ,bf16(2 字节)。

单 token 的 KV 元素数:

单 token 存储:

总量随 Batch、序列长度变化:。

场景 Batch Seq KV-cache 总量
短文本 1 2048
并发长文本 32 4096

对比模型参数 ():

  • 短文本场景:参数占主导(2 张 A100-80G 够用)。
  • 并发长文本场景:KV-cache 占主导(335 GB,至少 7 张 A100)。

结论:减少 KV-cache → 单次读写更少(更快)+ 数据能塞进单卡/单机用更高带宽(更快)。这正是 MLA 的目标。

访存分级(2 机 × 8×A100):卡内带宽 单机卡间的 3 倍 跨机的 20 倍。所以数据要尽量留在卡内 → 越省显存,越能享受高带宽。

第二部分:已有的 KV-cache 优化路线

类别 思路 代表
共享 KV 多个头共享一组 KV MQA、GQA、(MLA)
窗口 KV 只保留窗口内 KV,丢弃超窗的 Longformer
量化压缩 更低 bit 存 KV INT8
计算优化 减少换入换出,多在 SRAM 算 FlashAttention

MLA 属于共享 KV 分支的变体。先看它要对比的两个基线:

2.1 MHA / MQA / GQA(单 token KV 缓存量, 层、每层 头)

方法 KV 头数 单 token 缓存量 含义
MHA 个向量 每个 query 头有专属 K/V → 最占显存,质量最好
GQA 组 个向量 每 个头共享一组 K/V → 折中
MQA 个向量 所有头共享一组 K/V → 最省显存,质量下降
  • 时 GQA 退化成 MQA; 时退化成 MHA。

核心权衡:KV 头越少越省显存越快,但越掉质量。 MLA 的野心:显存像 MQA,质量像 MHA。

第三部分:MLA 的数学原理(论文公式)

img

img

img

img

3.0 符号

  • :第 个 token 的隐层输入(DeepSeek-V3:)。
  • :KV 低秩压缩维度(,V3 取 )。
  • :Q 低秩压缩维度(,V3 取 )。
  • :单头维度()。
  • :每层头数。
  • :RoPE 解耦维度(,V3 取 )。
  • 下标 表示第 个 attention head 的切片。

3.1 KV 的低秩压缩(论文公式 41 / 42 / 45)

压缩(down-projection) —— 把 维输入压成 维潜向量:

解压(up-projection) —— 用时再展开回每个头的 K/V:

类比 LoRA:先压缩()再扩展()。但 LoRA 关心"减少参数量",MLA 关心减少 KV 激活值(cache):推理时只缓存 ( 维),而不是缓存完整的 K/V。

疑点(此时还看不出省在哪): 解压后维度和 MHA 一个量级,还多了一步计算。省显存的秘密在 3.4 的矩阵吸收。

3.2 Q 的低秩压缩(论文公式 37 / 38)

Q 也压缩(主要为省参数,不影响 cache):

3.3 解耦 RoPE(论文公式 39 / 43 / 40 / 44)—— MLA 最关键的设计

RoPE 单独算两个带位置的小向量,而不是乘在上面的 上:

所有头共享一份式

拼接成最终的 Q、K:

注意公式 44 中 不带头下标 (全头共享)。所以每个头 Q/K 的总维度 (V3:)。

3.4 矩阵吸收(Weight Absorption)—— 省 cache 的真正原理

先理解吸收:矩阵乘满足结合律。要算 ,可以先把变换矩阵合并:

提前算好

也就是把 "吸收"进 ,提前算好 ,运行时省掉对 的变换。

a) 不加 RoPE 时,K 的解压矩阵可被吸收。 attention 的内容部分:

提前合并吸收

把 吸收进 Q 侧后:只需缓存 ( 维),不必缓存解压出的 ( 维)。 这就是 MLA 压缩 KV-cache 的核心。

类似地,V 的解压矩阵 可吸收进输出投影 ,于是 V 也不必缓存,只用同一个 即可。

b) 加 RoPE 后为什么不能吸收(必须解耦的原因)。 设位置旋转矩阵为 (满足 ),若对 直接套 RoPE:

随相对位置变化

中间多出一个 ,它随相对位置变化,不是固定矩阵,无法提前算好;且

旋转矩阵与解压矩阵不可交换

于是 被 RoPE 卡在 和 Q 之间,吸收链断裂 → 推理时必须对所有前缀 token 重算 K,效率大降。

论文原文:“RoPE 与低秩 KV 不兼容…与当前生成 token 相关的 RoPE 矩阵会位于 和 之间,而矩阵乘法不满足交换律,因此推理时必须重算所有前缀 token 的 K。”

c) 解法:用一个很小的维度 、以 MQA 方式单独承载 RoPE。 把 Q/K 各劈成两半:

  • nope(内容)半:,不带 RoPE → 可压缩、可吸收。
  • rope(位置)半:,带 RoPE,维度小(), 全头共享一份(不解压,无 可吸收,故 RoPE 套在它上面无冲突)。

最终 attention 分数 = 两半点积之和(论文公式 8 的形式):

走吸收共享式共享

完整的注意力输出:

3.5 MLA 实际缓存的东西

每层只缓存两个向量:

缓存项 维度 含义
KV 低秩压缩潜向量(K、V 共用)
共享的、带 RoPE 的位置 K

单 token、单层缓存量 (等价约 个 MQA 的缓存)。

方法 单 token/层 缓存元素 相对 MLA
MHA (如 )
MQA
MLA

MLA 比 MQA 多约 缓存,但保留了"恢复全 K/V"的能力,表达力远强于 MQA/GQA → 又省又快又强。

注:论文称 MLA 效果甚至超过 MHA,这点知乎作者存疑(缺消融对比),原理上也较难解释。客观结论是:MLA MHA 质量、 MQA 显存。

第四部分:sglang 代码实现对照

文件:python/sglang/srt/models/deepseek_v2.py(DeepseekV2AttentionMLA)、 python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py。

4.1 维度定义(deepseek_v2.py:1505-1510)

1
2
3
4
5
6
self.qk_nope_head_dim = qk_nope_head_dim   # d_h 内容半,128
self.qk_rope_head_dim = qk_rope_head_dim   # d_R 位置半,64
self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim   # 192 = 公式40/44 的拼接维度
self.v_head_dim = v_head_dim               # 128
self.kv_lora_rank = kv_lora_rank           # d_c, 512
self.q_lora_rank  = q_lora_rank            # d'_c, 1536

4.2 投影层(对应公式 41 / 37 / 43 / 42 / 45)

KV 压缩 + 共享 一把投影(公式 41 + 43;deepseek_v2.py:1566):

1
2
3
4
self.kv_a_proj_with_mqa = ReplicatedLinear(
    self.hidden_size,
    self.kv_lora_rank + self.qk_rope_head_dim,   # 512 + 64 = 576 ← 就是缓存的全部
)

Q 压缩(公式 37/38);sglang 还把 Q-down、KV-down、 融成一个 GEMM(deepseek_v2.py:1538-1555):

1
2
3
4
5
6
self.fused_qkv_a_proj_with_mqa = ReplicatedLinear(
    self.hidden_size,
    self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim,  # 1536+512+64
)
self.q_a_layernorm = RMSNorm(self.q_lora_rank, ...)
self.q_b_proj = ColumnParallelLinear(q_lora_rank, self.num_heads * self.qk_head_dim, ...)

KV 解压矩阵 (公式 42/45;deepseek_v2.py:1638):

1
2
3
4
self.kv_b_proj = ColumnParallelLinear(
    self.kv_lora_rank,                                            # 512
    self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),   # 解压成 K_nope + V
)

4.3 两个注意力对象:MHA 形态 vs MQA 形态

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 解码用:吸收后的"压缩空间 MQA",num_kv_heads=1(deepseek_v2.py:1678)
self.attn_mqa = RadixAttention(
    self.num_local_heads,
    self.kv_lora_rank + self.qk_rope_head_dim,   # K 维 = 512 + 64
    num_kv_heads=1,                              # ★ 只有 1 个 KV 头
    v_head_dim=self.kv_lora_rank,                # V 就是 c^{KV} 本身(512)
)
# 预填充用:解压成完整 K/V 的标准 MHA(deepseek_v2.py:1689)
self.attn_mha = RadixAttention(
    self.num_local_heads,
    self.qk_nope_head_dim + self.qk_rope_head_dim,  # 192
    num_kv_heads=self.num_local_heads,
    v_head_dim=self.v_head_dim,
)

4.4 吸收用的权重 w_kc / w_vc 从哪来(deepseek_weight_loader.py:567-590)

、 对应的 w_kc、w_vc 不是新权重,而是把 kv_b_proj 的权重 reshape/split 出来,分别用于"吸收进 Q"和"吸收进输出":

1
2
3
4
5
w_kc, w_vc = w.unflatten(
    0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
self_attn.w_kc = ...   # W^{UK} 切片 → 吸收进 Q
self_attn.w_vc = ...   # W^{UV} 切片 → 吸收进输出

所以"吸收"几乎不增加权重存储:同一份 kv_b_proj 换个形状用而已。

4.5 解码路径 forward_absorb 关键步骤(forward_mla.py)

  1. 取压缩潜向量并 split(:157-164、:302-306):
1
2
3
q, latent_cache = (...).split([self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], -1)
k_nope = latent_cache[..., : self.kv_lora_rank]    # c^{KV}(512)
k_pe   = latent_cache[..., self.kv_lora_rank :]    # 共享 k^R(64)
  1. RMSNorm,然后把 Q 拆成内容/位置两半(:303、:305):
1
q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
  1. 吸收:q_nope 先乘 w_kc 进入压缩空间,对应 中的 合并(:397):
1
q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc)   # → 512 维,可直接点积 c^{KV}
  1. 只对位置半套 RoPE(:418):
1
q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)   # nope 半完全不碰
  1. 在压缩空间跑 MQA: 对 点积(attn_mqa,K 维 )。
  2. 输出用 w_vc()解压回 V 空间(:656-659),再过 o_proj:
1
attn_bmm_output = torch.bmm(attn_output.transpose(0, 1), self.w_vc * self.w_scale)

4.6 两条路径:prefill 用 MHA,decode 用 absorb(dispatch:deepseek_v2.py:1831-1852)

forward_normal(MHA) forward_absorb(MLA/MQA)
用在 prefill / 长 extend decode
做法 解压出完整 K/V,标准 MHA(可用 FlashAttention) 吸收权重,压缩空间跑 MQA
为什么 prefill 计算密集,解压一次划算 decode 访存密集,省 KV 读取最关键

第五部分:常见疑问澄清

Q1:吸收后退化成 的 MQA,为什么不掉质量? 缓存里只有 1 份共享的 (像 MQA → 省显存、访存少 → 快)。但每个 query 头点积之前先乘上自己专属的 头切片(),等于戴着各自的"眼镜"去解读同一份共享数据 → 个头看出 种不同内容 → 保住 MHA 的多头表达力。用 MQA 的代价,拿到接近 MHA 的效果。

Q2: 不也要存吗? 要,但它是权重(每层 1 份、固定、和 token 数/batch 无关),不是 KV-cache(每 token 每请求都涨)。MLA 省的是后者。 每层 固定,而 KV-cache 每层可达数 GB 且随上下文爆炸。且 只是 kv_b_proj 换形状,本就存在 → 几乎零额外开销。省显存只看"随 token 增长的部分"。

Q3:为什么非要把 RoPE 单独拆出来(nope/rope)? RoPE 是依赖相对位置的旋转矩阵 。若把它套在解压得到的 上, 会卡在 和解压矩阵 之间,而 (不可交换)→ 吸收失效,必须重算所有前缀 K。冲突的根源是 。解法:把位置信息放到一条不需要吸收的小通道(,直接投影、共享、不解压),RoPE 套在它上面无冲突。位置共享走 rope 半,语义分头走 nope 半。

第六部分:一图流总结

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
hidden h_t (d=7168)
  │
  ├─►(公式41) W^DKV ─► c^{KV}(512) ──┐  RMSNorm           ← ★ 缓存
  ├─►(公式43) W^KR ─► RoPE ─► k^R(64)─┤  全头共享          ← ★ 缓存
  │
  └─►(公式37) W^DQ ─► c^Q(1536) ─►(38) q ─split─► q_nope(128)   q_pe(64)
                                                  │             │
   decode(absorb):                       吸收 w_kc(=W^UK)     RoPE(公式39)
                                                  ▼             ▼
                          q_nope_out(512) ┐                q_pe' ┐
                          c^{KV}(512) ─────┤── 点积         k^R' ─┤── 点积
                                  score_nope                 score_pe
                                          └──── 相加 ────────┘
                                                  │  softmax · c^{KV}
                                          attn(512) ─吸收 w_vc(=W^UV)─► V空间 ─► o_proj

三句话记住 MLA:

  1. 省显存 —— 每 token/层只缓存 ,约为 MHA 的 、MQA 的 。
  2. 拆 RoPE —— 位置半单独存(,共享、不解压),因为 RoPE 与解压矩阵不可交换,会破坏吸收。
  3. 吸收 + MQA —— decode 时把 吸进 Q、 吸进输出,在压缩空间跑 的 MQA;每头各自的 保住多头表达力 → 又省、又快、又(几乎)不掉质量。
  • 标题: DeepSeekV2之MLA
  • 作者: 鱿鱼圈
  • 创建于 : 2026-04-20 23:50:00
  • 更新于 : 2026-06-30 22:44:58
  • 链接: https://yuyanqi.com/2026/04/20/DeepSeekV2之MLA/
  • 版权声明: 本文章采用 CC BY-NC-SA 4.0 进行许可。
评论
目录
DeepSeekV2之MLA
  1. MLA(Multi-head Latent Attention)完整详解
    1. 第一部分:为什么需要 MLA —— KV-cache 是推理瓶颈
      1. 1.1 LLM 推理的两个阶段
      2. 1.2 KV-cache 的由来:用空间换时间
      3. 1.3 显存里存了什么
      4. 1.4 KV-cache 到底多大(以 Qwen-72B 朴素 MHA 估算)
    2. 第二部分:已有的 KV-cache 优化路线
      1. 2.1 MHA / MQA / GQA(单 token KV 缓存量, 层、每层 头)
    3. 第三部分:MLA 的数学原理(论文公式)
      1. 3.0 符号
      2. 3.1 KV 的低秩压缩(论文公式 41 / 42 / 45)
      3. 3.2 Q 的低秩压缩(论文公式 37 / 38)
      4. 3.3 解耦 RoPE(论文公式 39 / 43 / 40 / 44)—— MLA 最关键的设计
      5. 3.4 矩阵吸收(Weight Absorption)—— 省 cache 的真正原理
      6. 3.5 MLA 实际缓存的东西
    4. 第四部分:sglang 代码实现对照
      1. 4.1 维度定义(deepseek_v2.py:1505-1510)
      2. 4.2 投影层(对应公式 41 / 37 / 43 / 42 / 45)
      3. 4.3 两个注意力对象:MHA 形态 vs MQA 形态
      4. 4.4 吸收用的权重 w_kc / w_vc 从哪来(deepseek_weight_loader.py:567-590)
      5. 4.5 解码路径 forward_absorb 关键步骤(forward_mla.py)
      6. 4.6 两条路径:prefill 用 MHA,decode 用 absorb(dispatch:deepseek_v2.py:1831-1852)
    5. 第五部分:常见疑问澄清
    6. 第六部分:一图流总结