从 MHA 到 MLA

MQA/GQA 靠「少存几组 KV」换显存，MLA 换的是表示形式——把 KV 压成低秩 latent、再把上投影矩阵吸收进 Q 侧，于是缓存里只剩一份压缩向量；代价是 RoPE 必须从中拆出来单走一路。

为什么盯着 KV Cache

自回归解码每步只算一个 token 的 Q，但要与全部历史 K/V 做内积，所以历史必须缓存。单层单 token 的开销是

$2n_h{d}_h\text{（K 与 V 各一份）}$

乘上层数与 batch 后，长上下文场景下 KV Cache 而非权重成为显存主项，且每步解码都要把它整体读一遍——解码是彻底的内存带宽受限（背景见 KV Cache）。降本有三条路：算子侧减少 HBM 往返（FlashAttention）、限制每个 query 能看到的范围（稀疏注意力）、以及本文这条——改造头结构，直接把 $2n_h{d}_h$ 这个系数打下来。

MQA 与 GQA：共享 KV

MQA（Shazeer, 2019）最激进：$n_h$ 个 Q 头共用一组 KV，缓存缩小 $n_h$ 倍。但 K/V 的表达能力被压到单头，训练不稳、质量掉点明显。

GQA（2023）做插值：Q 头分成 $g$ 组，每组共享一组 KV，缓存为 $2g{d}_h$。$g=n_h$ 退化成 MHA，$g=1$ 退化成 MQA。实践中 $g=8$ 就能基本保住质量，LLaMA-2 70B、Llama-3 全系、Qwen2 均采用，是当前事实标准。

两者的共同问题：压缩手段是丢弃头间差异，KV 的秩被硬性砍掉，质量与显存严格互斥。

MLA：低秩 latent 替代共享

DeepSeek-V2（arXiv:2405.04434）换了思路——不共享，而是压缩。隐状态 $h_t$ 先下投影到 $d_c$ 维 latent：

$c_t^{KV}=W^{DKV}{h}_t\in {\mathbb{R}}^{d_c},d_c\ll n_h{d}_h$

每个头的 K、V 由两个上投影矩阵从同一份 latent 恢复：

$[k_{t,1}^C;\dots ;k_{t,n_h}^C]=W^{UK}{c}_t^{KV},[v_{t,1}^C;\dots ;v_{t,n_h}^C]=W^{UV}{c}_t^{KV}$

各头仍有独立参数（$W^{UK}$ 按头分块不同），只是被约束在一个 $d_c$ 维的公共子空间里——这是低秩近似，不是共享。Q 侧同样压缩到 $d_c^{\prime }$ 维（DeepSeek-V2 取 1536），目的不同：Q 不进缓存，压缩是为了减少训练期激活值。

RoPE 解耦：本设计的关键约束

MLA 的推理红利来自矩阵吸收：内容项分数可以改写成

$q_{t,i}^{C\top }{k}_{j,i}^C=h_t^{\top }\left(W_i^{UQ\top }{W}_i^{UK}\right)c_j^{KV}$

括号里是常量矩阵，可在加载时预乘进 Q 的投影权重。于是 $k^C$ 永远不必显式解压，缓存里只放 latent 就够了；同理 $W^{UV}$ 可被吸收进输出投影 $W^O$。

import torch
 
d_model, n_h, d_h, d_c = 256, 4, 32, 64
torch.manual_seed(0)
W_dkv = torch.randn(d_c, d_model) / d_model**0.5          # h -> latent
W_uk = torch.randn(n_h * d_h, d_c) / d_c**0.5             # latent -> K（按头分块）
W_uq = torch.randn(n_h * d_h, d_model) / d_model**0.5
 
h = torch.randn(6, d_model)
c_kv = h @ W_dkv.T                                        # (6, d_c)：唯一进 cache 的内容量
 
# 朴素路径：先解压出 K，再逐头打分
k_c = (c_kv @ W_uk.T).view(6, n_h, d_h)
q_c = (h @ W_uq.T).view(6, n_h, d_h)
naive = torch.einsum("thd,jhd->htj", q_c, k_c)
 
# 吸收路径：W_uk 并入 W_uq，直接对 latent 打分，K 从不显式出现
W_absorb = W_uq.view(n_h, d_h, d_model).transpose(1, 2) @ W_uk.view(n_h, d_h, d_c)
absorbed = torch.einsum("td,hdc,jc->htj", h, W_absorb, c_kv)
 
print(torch.allclose(naive, absorbed, atol=1e-4))         # True

但 旋转位置编码（RoPE） 会破坏这个结构。若对 $k^C$ 施加旋转，分数变成

$h_t^{\top }{W}_i^{UQ\top }{\mathbf{R}}_{j-t}{W}_i^{UK}{c}_j^{KV}$

中间那个 ${\mathbf{R}}_{j-t}$ 依赖 $t$ 与 $j$ 两个下标，把两侧投影矩阵隔开——乘积无法预计算成一个常量矩阵，每对 $(t,j)$ 都得重算，吸收失效，只能退回显式解压 K，MLA 的全部收益归零。

DeepSeek 的解法是把位置信息拆出独立通路：额外一个 $d_h^R=64$ 维的分支，从 $h_t$ 直接投影并旋转，且全头共享

$k_t^R=\mathrm{RoPE}(W^{KR}{h}_t),k_{t,i}=[k_{t,i}^C;k_t^R]$

Q 侧对称地拆出逐头的 $q_{t,i}^R$。拼接后的内积自然分解为两项：

$q_{t,i}^{\top }{k}_{j,i}=\underset{\text{可吸收，走 latent}}{\underbrace{q_{t,i}^{C\top }{k}_{j,i}^C}}+\underset{\text{带位置，走独立小通路}}{\underbrace{q_{t,i}^{R\top }{k}_j^R}}$

内容与位置各行其道：前者享受低秩压缩，后者保留精确的相对位置语义。KV Cache 因此只需存 $c_t^{KV}$ 与 $k_t^R$ 两项，共 $d_c+d_h^R$ 个元素。这是位置编码与注意力结构耦合的教科书案例——位置编码的选择直接决定了哪些代数变换合法。

数字

DeepSeek-V2 取 $n_h=128$、$d_h=128$、$d_c=4d_h=512$、$d_h^R=64$，每层每 token 缓存 $512+64=576$ 个元素，即 $4.5d_h$，论文指出这等价于只有 2.25 组的 GQA——却在评测上与 MHA 持平甚至略优。整体口径：相比 DeepSeek 67B，KV Cache 降 93.3%，最大生成吞吐提升至 5.76 倍，训练成本省 42.5%，236B 总参/21B 激活，支持 128K 上下文。DeepSeek-V3 沿用同一套 MLA 配置。

对比

方案	单层单 token 缓存	质量	代表模型
MHA	$2n_h{d}_h$	基线	GPT-2/3、Llama-1
MQA	$2d_h$	明显下降	PaLM、Falcon
GQA	$2g{d}_h$	接近 MHA	Llama-2 70B / Llama-3、Qwen2
MLA	$d_c+d_h^R$	持平或略优 MHA	DeepSeek-V2/V3

MLA 的实现复杂度显著高于 GQA——两套投影、两条通路、推理期还要做权重吸收与算子改写，这是它至今未被广泛复刻的主因。但对 MoE 长上下文模型，激活参数已经很省、瓶颈全在 KV Cache 上，MLA 恰好补上了这块：它是 DeepSeek 把 128K 上下文做到可负担推理成本的必要配套，而非可选优化。