引言

Multi-head Latent Attention（MLA）是 DeepSeek-V2/V3 系列模型中最核心的架构创新之一。它在标准 Multi-Head Attention（MHA）的基础上引入了低秩压缩，大幅降低了 KV cache 的显存占用，同时保持了与 MHA 相当的模型质量。

本文从 QKV head dimension 处理差异 这一视角切入，深入分析 MLA 的设计原理、vLLM 中的具体实现，以及这种设计带来的性能收益。

1. 标准 MHA 回顾

在标准 MHA 中，给定输入序列 X ∈ ℝ^{S×D}，Q、K、V 通过三个独立的线性变换得到：

Q = X @ W_Q   → [S, H, d_head]
K = X @ W_K   → [S, H, d_head]
V = X @ W_V   → [S, H, d_head]

其中 d_head = D / H，三者完全相等。注意力计算为：

Attention(Q, K, V) = softmax(Q @ K^T / √d_head) @ V   → [S, H, d_head]

输出再经 W_O 映射回 D 维。这种 Q/K/V 共享同一 head dim 的设计深入人心，以至于很多人默认这是注意力机制的"必须要求"。

2. MLA 的核心思路：低秩压缩

MLA 的关键洞察在于：K 和 V 的信息存在大量冗余，可以压缩到一个低秩的 latent 空间中。

具体来说，MLA 引入了一个压缩矩阵 W_DKV ∈ ℝ^{D×d_c}（以及后续的升维矩阵），将 K 和 V 的联合信息压缩到维度为 d_c（例如 512，称为 kv_lora_rank）的 latent 向量 c_kv 中：

c_kv = W_DKV @ [k; v]   → [d_c]

然后在注意力计算时将其解压回完整维度：

  • k_nope = W_UK @ c_kv → [P] 其中 P = qk_nope_head_dim
  • k_pe = RoPE(W_KR @ c_kv) → [R] 其中 R = qk_rope_head_dim
  • v = W_UV @ c_kv → [V] 其中 V = v_head_dim

3. QKV Head Dimension 的差异

3.1 Prefill 路径（forward_mha）

在 prefill 的 forward_mha 中，各张量的 shape 如下：

张量Shape说明
q[Sq, H, P+R]Q 拼接 nope + rope
k_nope[Sq, H, P]从 kv_b_proj 解压（全 H 头）
k_pe[Sq, 1, R]单头 RoPE，广播到 H
k[Sq, H, P+R]concat(k_nope, k_pe)
v[Sq, H, V]从 kv_b_proj 解压，V 不一定等于 P+R

关键代码（vllm/vllm/models/deepseek_v4/nvidia/model.py）：

kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
    -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
)
k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

Q 和 K 的 head dim 是 P+R（例如 128+64=192），而 V 的 head dim 是独立的 V（例如 128）。 这与标准 MHA 形成鲜明对比——标准 MHA 中三者均为 d_head。

3.2 Decode 路径（forward_mqa）

在 decode（自回归生成）路径中，MLA 走 MQA 路径（单 KV 头），head dim 的差异更为显著：

张量Shape说明
q[B, H, kv_lora_rank + R]q_nope 经 W_UK 投影 + q_pe
k_buffer[B, page, 1, kv_lora_rank + R]cache 中 [c_kv | k_pe]
v_buffer[B, page, 1, kv_lora_rank]cache 中仅 c_kv 部分

这里 Q 和 K 的 head dim 是 kv_lora_rank + R（例如 512+64=576），而 V 的 head dim 是 kv_lora_rank（例如 512）。Q/K 和 V 不仅数值不同，维度性质也不同——Q/K 包含了 RoPE 部分而 V 没有。

4. 为什么不同的 Head Dim 可以计算注意力？

注意力公式本身对此给出了明确的答案：

Attention(Q, K, V) = softmax(Q @ K^T / √d) @ V
  1. Q @ K^T 是 dot product，要求 Q 和 K 的 head dim 相同——这确实成立，两者都是 qk_head_dim。
  2. softmax 权重 × V 是 V 行的线性组合，输出维度自然等于 V 的 head dim，数学上从未要求与 Q/K 一致。

标准 MHA 恰好将三者设为相等只是一个设计惯例，而非数学约束。MLA 充分利用了这一自由度。

FlashAttention 等主流 attention kernel 原生支持 Q/K 与 V 的 head dim 不同：

# FlashAttention 调用示例（伪代码）
output = flash_attn_func(q, k, v)  # q/k head_dim ≠ v head_dim 完全合法
# output shape: [Sq, H, v_head_dim]

5. Rope 的不可吸收性——Head Dim 差异的根源

MLA 最精妙的设计在于 nope 部分的吸收（absorption）。

5.1 Nope 吸收

在 decode 时，nope 部分的计算为：

q_nope [H, P] @ k_nope^T = q_nope @ (W_UK @ c_kv)^T
                          = (q_nope @ W_UK) @ c_kv^T

因为 W_UK 是纯线性变换，可以自由交换结合律。将 W_UK 吸收到 Q 侧后，KV cache 中只需存低秩的 c_kv（例如 512 维），无需存完整的 k_nope（例如 128×H 维）。这是 MLA 压缩 KV cache 的核心机制。

5.2 Rope 的不可吸收

Rope 部分则完全不同：

q_pe [H, R] @ k_pe^T

其中 k_pe = RoPE(W_KR @ c_kv, pos_k)，RoPE 引入的旋转矩阵依赖于 key token 的位置 pos_k。RoPE 具有 RoPE(a, pos1)^T @ RoPE(b, pos2) = a^T @ RoPE(pos2 - pos1)(b) 的性质，所以：

q_pe(pos_q)^T @ k_pe(pos_k) = (W_QR @ h_q)^T @ RoPE(pos_k - pos_q, W_KR @ c_kv)

这仍然依赖于 pos_k - pos_q，无法像 W_UK 那样简单地分解到 Q 侧。k_pe 必须写入 cache，decode 时 q_pe 直接与缓存的 k_pe 做 dot product。

5.3 Cache 布局的物理解释

这就是 KV cache 中每个 entry 存储 [c_kv | k_pe]（例如 512+64=576 维）的根本原因：

cache_entry = [  c_kv (kv_lora_rank)  |  k_pe (qk_rope_head_dim)  ]
                  ↑                         ↑
              可吸收（低秩压缩）          不可吸收（必须缓存）

相比之下，标准 MHA 需要缓存完整的 K 和 V（每头 2 × d_head × H 维），MLA 通过吸收将每 token 的 KV cache 从 2 × d_head × H 压缩到 kv_lora_rank + qk_rope_head_dim。

6. vLLM Triton Kernel 中的 MLA 优化

在 decode 路径中，vLLM 的 Triton decode attention kernel 对 MLA 做了专门的优化（vllm/vllm/v1/attention/ops/triton_decode_attention.py）。

6.1 两阶段计算

def decode_attention_fwd_grouped(...):
    # Stage 1: Q @ K^T attention socres × V
    _decode_grouped_att_m_fwd(q, k_buffer, v_buffer, attn_logits, ...)
    # Stage 2: 在线 softmax 合并多个 split 的结果
    _decode_softmax_reducev_fwd(attn_logits, q, o, lse, v_buffer, ...)

6.2 MLA 特定的 Tile 划分

当 is_mla=True 时，tile 尺寸分开设置：

if is_mla:
    BLOCK_DMODEL = kv_lora_rank   # nope 部分（如 512）
    BLOCK_DPE    = qk_rope_head_dim  # rope 部分（如 64）

Q 和 K 分两段加载：

  • q[..., :BLOCK_DMODEL] → ql_nope（投影后的 latent query）
  • q[..., BLOCK_DMODEL:BLOCK_DMODEL+BLOCK_DPE] → q_pe（RoPE query）

注意力分数分段计算：

qk = tl.dot(q, k.to(q.dtype))              # nope × nope
if BLOCK_DPE > 0:
    qk += tl.dot(qpe, kpe.to(qpe.dtype))    # rope × rope

6.3 省掉 V 的显存加载

这是 MLA decode kernel 最精妙的优化。对 IS_MLA=True，K buffer 的前 kv_lora_rank 维恰好就是 c_kv（也是 V），所以：

if not IS_MLA:
    v = tl.load(V_Buffer, ...)      # 普通 GQA：从 V_Buffer 加载
else:
    v = tl.trans(k)                 # MLA：直接转置已加载的 k → 就是 V

因为 k 以 [BLOCK_DMODEL, BLOCK_N] 形状加载了 c_kv 部分，转置后 v = tl.trans(k) 给出 [BLOCK_N, kv_lora_rank] 形状的 V 数据。省掉一次显存加载。

7. Nope vs Rope 在 MHA/MQA 中的归属

MLA 中 nope 和 rope 在 MHA/MQA 的归属不同：

部分Prefill（forward_mha）Decode（forward_mqa）
nopeMHA：k_nope [Sq, H, P] 解压到全头MQA：ql_nope [H, kv_lora_rank] × c_kv [1, kv_lora_rank]
ropeMQA：k_pe [Sq, 1, R] 广播到 HMQA：q_pe [H, R] × k_pe [1, R]

在 prefill 中，kv_b_proj 将 c_kv 解压回完整的 H 头 k_nope（MHA），而 k_pe 始终保持单头广播（MQA）。在 decode 中，两者都是 MQA。是否能吸收取决于 RoPE 能否分解，与 MHA/MQA 无关。

8. 性能收益总结

MLA 的 head dim 差异化设计带来的性能收益：

方面标准 MHAMLA节省
KV cache（每 token）2 × d_head × Hkv_lora_rank + R约 4-8 倍
计算量（decode）完整 K/V 投影c_kv 低秩计算 + 吸收显著降低
显存加载（decode kernel）分别加载 K 和 Vv = tl.trans(k) 省一次加载1 次显存读取

以 DeepSeek-V3 的具体配置为例（d_head=128, H=128, kv_lora_rank=512, R=64）：

  • 标准 MHA KV cache：2 × 128 × 128 = 32,768 维/token
  • MLA KV cache：512 + 64 = 576 维/token
  • 压缩比：约 57 倍

9. 结语

MLA 通过将 Q/K 与 V 的 head dim 解耦，在注意力机制的数学框架内找到了一个优雅的优化空间。nope 部分的吸收压缩了 KV cache，rope 部分的不可吸收则保留了位置编码的必要信息。这种"能省则省，该留则留"的设计哲学，使得 MLA 成为当前大模型架构中最具实用价值的注意力变体之一。

理解这些底层细节，不仅是正确阅读和修改 vLLM 代码的前提，也能帮助我们在设计新架构时，更清醒地判断哪些"惯例"可以突破、哪些"约束"必须遵守。