重新理解 MLA：缓存 X 而非 KV

很多文章在介绍 MLA（Multi-head Latent Attention）时，会说它是把 KV 投影到一个低维空间，需要时再还原回去。这个说法不算错，但没有抓住要点。如果低秩投影就是核心思想，那 GQA 早就在做了。KV 头数少于 Q 头数，从输入到 KV 的映射本身就是降维。为什么还需要 MLA¹？

苏神在他的解析文章²中给出了一个 insight：MLA 的本质改进不是低秩投影，而是用可学习投影替代 GQA 的分割和复制，让每个头从同一个低维向量中提取出不同的 KV，在低缓存开销下达到接近 MHA 的多头表达能力。

本文从一个思想实验出发³：如果我们不缓存 KV，而是缓存更小的输入 X 会怎样？这个过程会自然地引出矩阵吸收和降维的思路。

如果你对 Attention 的基本原理还不熟悉，建议先阅读《为什么 KV 缓存没有 Q》。

出发点：KV Cache 太大了

KV Cache 的本质是用空间换时间。在自回归生成的 Decode 阶段，每生成一个新 token，都需要用当前的 $q$ 去和历史所有位置的 $k, v$ 做注意力计算。不缓存就要把每一层的历史 token 全部重算一遍；缓存则要把每一层、每一个头、每一个历史 token 的 $k, v$ 向量全部存下来。注意，缓存节省的不仅仅是 Attention 的计算量，还包括 FFN 在内的整个 layer。

问题在于，这个交换做得太激进了。KV Cache 将 Decode 阶段的计算瓶颈从 compute-bound 变成了 memory-bound¹：

显存容量瓶颈：KV Cache 的总量随 batch size 和序列长度线性增长，很快超过模型参数本身的存储占用。
带宽瓶颈：Decode 阶段每步只算一个 token，算力强度极低，GPU 的计算单元大量空转在等待数据搬运。
互相制约：为了提高算力利用率想增大 batch size，又会进一步加剧显存压力。

虽然模型参数在整个 batch 间共享，加大 batch 可以有效摊薄参数读取的带宽成本。但 KV Cache 是每个请求独立的：batch 翻倍，KV Cache 的总读取量也翻倍。这使得 Decode 阶段 attention 的算力强度（arithmetic intensity）始终在 1 附近⁴，远低于任何硬件的 roofline。KV Cache 的大小直接决定了推理吞吐的天花板。

MQA 和 GQA 试图通过让多个 Q 头共享同一组 KV 来缓解这个问题。它们确实有效，但都是在”保留 KV Cache”这一范式内做优化。MLA 则换了一个思路：不缓存 KV。

思想实验：缓存 X 而非 KV

在进入 MLA 的实际设计之前，我们先做一个思想实验³。它不等同于 MLA，但展示了 MLA 背后的核心思路：可以不缓存 KV，而是缓存更小的东西。

一个等价变换

在标准的 MHA 中（省略 scaling factor，只看单个头），Decode 阶段的注意力输出可以写成：

o = \text{softmax}\left(q K^\top\right) V W_O

其中 $q = x W_Q$ ， $K$ 和 $V$ 分别是历史所有位置的 Key 和 Value 矩阵， $W_O$ 是输出投影。展开 $K$ 和 $V$ ：

o = \text{softmax}\left(x W_Q (X W_K)^\top\right) \cdot X W_V W_O

转置 $(XW_K)^\top = W_K^\top X^\top$ 。应用矩阵乘法的结合律，得到：

o = \text{softmax}\left(x \underbrace{W_Q W_K^\top}_{\text{合并}} X^\top\right) \cdot X \underbrace{W_V W_O}_{\text{合并}}

重组之后，KV 消失了。 注意力不再需要显式计算 $K = XW_K$ 和 $V = XW_V$ ，而是直接用输入 $X$ 参与计算。 $W_K$ 被吸收进了 $W_Q$ 侧， $W_V$ 被推迟到输出端与 $W_O$ 合并。

此时，不需要缓存 KV，只需要缓存每一层的输入 $X$ 。

矩阵吸收与代价分析

softmax 内部的运算 $x W_Q W_K^\top X^\top$ 有两种计算顺序（以下计算量省略矩阵乘法的常数 2，主要关心复杂度和缩放比例）：

顺序 A（显式计算 K）：对每个历史位置先投影出 $k_i$ ，再做内积：

\underbrace{1 \times d \times d_h}_{x W_Q} + \underbrace{N \times d \times d_h}_{X W_K} + \underbrace{1 \times d_h \times N}_{q K^\top} \simeq O(N \cdot d \cdot d_h)

顺序 B（矩阵吸收）：先将 $W_Q W_K^\top$ 合并为一个矩阵，再与 $X$ 做内积：

\underbrace{1 \times d \times d_h}_{x W_Q} + \underbrace{1 \times d_h \times d}_{q W_K^\top} + \underbrace{1 \times d \times N}_{q' X^\top} \simeq O(Nd)

顺序 B 比 A 快 $d_h$ 倍。矩阵吸收的本质是将 $W_K$ 吸收进 $W_Q$ ，把 N 次 $d \to d_h$ 投影变为 1 次投影。代价是 query 从 $d_h$ 维”升维”回 $d$ 维（ $q W_K^\top$ 这一步），从而能直接与 $d$ 维的 $X$ 做内积，有效头大小因此从 $d_h$ 变成了 $d$ 。

经过矩阵吸收，缓存 $X$ 方案的完整单头 Decode 计算量为：

\underbrace{1 \times d \times d_h}_{计算\, q} + \underbrace{1 \times d_h \times d}_{q W_K^\top} + \underbrace{1 \times d \times N}_{q'X^\top} + \underbrace{1 \times N \times d}_{s \cdot X} + \underbrace{1 \times d \times d_h}_{\cdot\, W_V} + \underbrace{1 \times d_h \times d}_{W_O} = 4dd_h + 2Nd

对比带 KV Cache 单头 Decode 计算量（只需计算当前 token 的 $q, k, v$ ，历史的 $k, v$ 已缓存）：

\underbrace{3 \times 1 \times d \times d_h}_{计算\, q,k,v} + \underbrace{1 \times d_h \times N}_{qK^\top} + \underbrace{1 \times N \times d_h}_{s \cdot V} + \underbrace{1 \times d_h \times d}_{计算\, o} = 4dd_h + 2Nd_h

投影部分不变（ $4dd_h$ ），增量 $2N(d - d_h)$ 来自 attention 计算部分（ $QK^\top$ 打分 + 加权求和），有效头大小从 $d_h$ 增大到 $d$ 。

代价分析

以 DeepSeek-V3（ $d = 7168$ ， $n_h = 128$ ， $d_h = 128$ ）为例：

缓存：KV Cache 每层每 token 缓存 $2 n_h d_h$ 维，缓存 $X$ 只需 $d$ 维，压缩 $\frac{2 n_h d_h}{d} = \frac{2 \times 128 \times 128}{7168} \approx 4.57$ 倍。
计算：attention 计算（ $QK^\top$ + 加权求和）的有效头大小从 $d_h$ 变为 $d$ ，变为原来的 $\frac{d}{d_h} = \frac{7168}{128} = 56$ 倍。

缓存压缩 4.57 倍，attention 计算变为 56 倍。不够划算，但这只是第一步。

对 X 降维

经过矩阵吸收，上述方案的计算量已经可控。但 cache 的压缩比只有 4-5 倍，不够。

思路很直接：用一个矩阵把 $X$ 降维。

c_i = x_i W_c, \quad W_c \in \mathbb{R}^{d \times d_c}, \quad d_c \ll d

K 和 V 的计算改为从 $c$ 出发： $k = c W_K$ ， $v = c W_V$ （投影矩阵的维度相应调整）。矩阵吸收同样适用： $W_K$ 被吸收进 $W_Q$ 侧， $W_V$ 被推迟到输出端。最终只需缓存 $d_c$ 维的 $c$ ，attention 的有效头大小从 $d$ 降至 $d_c$ 。

以 DeepSeek-V3 的设置（ $d_c = 512$ ）：

缓存量： $\frac{2 n_h d_h}{d_c} = \frac{2 \times 128 \times 128}{512} = 64$ ，相比传统 KV Cache 压缩 64 倍（此前是 $\frac{2 n_h d_h}{d} \approx 4.57$ 倍）。
计算量：有效头大小从 $d$ 降到 $d_c$ ，attention 计算量增加为 $d_c/d_h = 4$ 倍（此前是 $d/d_h = 56$ 倍）。

缓存压缩 64 倍，计算量只增加 4 倍。 现在十分划算了。

但这个思想实验有一个根本局限：所有注意力头共享同一份 K 和 V（ $W_K, W_V$ 没有头的上标），等效于 MQA。这限制了模型的表达能力。接下来我们看实际的 MLA 如何解决这个问题。

从思想实验到实际 MLA

可学习投影替代分割和复制

上面的思想实验揭示了一条路径：缓存低维的 $c$ 而非 KV，通过矩阵吸收控制计算开销。但它止步于 MQA：所有头被迫使用相同的 KV。

GQA 面对类似的问题，采用的是分割和复制：将低秩向量分成 $g$ 组，每组内的头共享一份 KV⁵。这种方式的信息提取能力有限，同组头被迫看到完全相同的 KV²。

MLA 的做法是替换为每个头独立的可学习投影矩阵（上标 $(s)$ 表示第 $s$ 个注意力头）：

k_i^{(s)} = c_i W_K^{(s)}, \quad v_i^{(s)} = c_i W_V^{(s)}

每个头从同一个 $c_i$ 中提取出不同的 KV，信息利用率远高于 GQA 的分割复制。这为将 $d_c$ 压到 512 提供了条件。

MHA 的表达能力，接近 MQA 量级的缓存开销

引入 per-head 投影后，KV 变成了每个头各不相同的，缓存量岂不是回到了 MHA 的水平？

这里正是矩阵吸收的用武之地：

q^{(s)} k_i^{(s)\top} = (x W_Q^{(s)})(c_i W_K^{(s)})^\top = x (W_Q^{(s)} W_K^{(s)\top}) c_i^\top

$W_K^{(s)}$ 被吸收进 $W_Q^{(s)}$ ， $k_i^{(s)}$ 无需显式计算。同理， $W_V^{(s)}$ 可以吸收进输出投影。于是只需缓存所有头共享的 $c_i$ 。KV Cache 大小与头数无关。

Per-head 投影赋予了 MLA 接近 MHA 的多头表达能力（每个头看到不同的 KV），矩阵吸收又让缓存只需存共享的 $c$ （接近 MQA 量级）。两者并不矛盾。对同一个 MLA 计算来说，先显式恢复 per-head KV 再计算 attention，和把恢复矩阵吸收到 Query / 输出投影中、直接基于共享的 $c$ 计算，是同一组矩阵乘法的两种执行顺序，结果等价。

这也解释了为什么 DeepSeek-V2 能将头数设到 128（远超常规设置）：增加头数只增加计算量和模型能力，不增加缓存压力。

这一特性也改变了多卡部署的缓存分布方式。传统 GQA 的 KV Cache 按头分割，每个 TP rank 只持有自己负责的那几个头的缓存，天然可分。MLA 的缓存是所有头共享的 $c_i$ ，per-head 的 KV 在推理时通过各 rank 自己的 $W_K^{(s)}$ 临时重建，因此每个 rank 都需要完整的 $c_i$ 副本。缓存从分割变成复制后，会吃掉一部分全局缓存收益：以 8 卡 TP 为例，常见 $n_g=8$ 的 GQA 全局缓存是 $2 \times 8 \times 128 = 2048$ 维，MLA 全局缓存是 $8 \times (d_c + d_r) = 8 \times 576 = 4608$ 维。它不比 8-group GQA 更省，但仍远低于 MHA 的 $2 \times 128 \times 128 = 32768$ 维，并换来了更接近 MHA 的表达能力。

MLA 与 MHA、GQA、MQA 的结构对比（来自 DeepSeek-V2 技术报告）

MLA 与其他方法的 KV Cache 大小及模型效果对比（来自 DeepSeek-V2 技术报告）

上图中 MLA 每层每 token 需缓存 $d_c + d_r = 576$ 个元素（BF16 下为 1152 Bytes），令 GQA 的对应量 $2 n_g d_h = d_c + d_r$ 解出 $n_g = 2.25$ ，即 MLA 缓存大小仅相当于 2.25 个组的 GQA。

另外，ablation 实验（论文 Appendix D, Table 9）显示其 benchmark 表现略优于 MHA。

总结：思想实验与实际 MLA 的对照

前面的思想实验帮助理解了”缓存 X 而非 KV”的核心思路。以下是思想实验与实际 MLA 之间的对应关系：

思想实验	实际 MLA
缓存输入 $X$ （ $d$ 维）	缓存 $c = xW_c$ （ $d_c$ 维）
所有头共享同一份 KV	共享 $c$ ，每头用独立的 $W_K^{(s)}, W_V^{(s)}$ 提取不同 KV
$W_Q W_K^\top$ 合并为一个矩阵	每个头独立合并 $W_Q^{(s)} W_K^{(s)\top}$
不含位置编码	部分维度加 RoPE（所有头共享），部分走矩阵吸收
等效 MQA（共享 KV，表达能力受限）	兼具 MHA 的表达能力与接近 MQA 量级的缓存开销

RoPE 兼容性：一个精巧的折中

到此为止，MLA 的核心思想已经完整了。但还有一个绕不开的问题：位置编码（RoPE）⁶。

在实际的 MLA 中，Q 侧也有自己的降维投影： $c_t^Q = x_t W_c^Q$ （ $d_c^Q = 1536$ ），与 KV 侧的 $c_i = x_i W_c$ （ $d_c = 512$ ）独立。Q 的降维主要是为了减少训练阶段的参数量和激活值显存，与推理缓存无关。

矩阵吸收的前提是： $W_Q^{(s)} W_K^{(s)\top}$ 可以合并为一个与位置无关的固定矩阵。但 RoPE 在 Q 和 K 上各施加了一个位置相关的旋转矩阵 $R$ （下标 $t$ 为当前 token 位置， $i$ 为历史 token 位置），使得：

q_t^{(s)} k_i^{(s)\top} = (c_t^Q W_Q^{(s)} R_t)(c_i W_K^{(s)} R_i)^\top = c_t^Q W_Q^{(s)} \underbrace{R_{t-i}}_{\text{相对位置}} W_K^{(s)\top} c_i^\top

中间插入了一个与相对位置 $t - i$ 相关的 $R_{t-i}$ ， $W_Q^{(s)} W_K^{(s)\top}$ 无法再被合并为固定矩阵。如果强行缓存 $c_i$ 而不缓存 $k_i$ ，推理时就必须对所有历史位置重新计算带 RoPE 的 Key，完全抵消了缓存的意义。

DeepSeek 的解法是一种混合方案：每个头的 Q、K 被拆成两部分拼接，一部分不加 RoPE（走矩阵吸收路线），另一部分加 RoPE：

q_t^{(s)} = \left[c_t^Q W_{qc}^{(s)},\; c_t^Q W_{qr}^{(s)} R_t\right], \quad k_i^{(s)} = \left[c_i W_{kc}^{(s)},\; x_i W_{kr} R_i\right]

其中，不带 RoPE 的部分可以照常做矩阵吸收，推理时只缓存 $c_i$ 。带 RoPE 的部分所有头共享同一个 $W_{kr}$ （没有上标 $(s)$ ），相当于 MQA，只需额外缓存一个 $d_r$ 维的向量（ $d_r = d_h / 2 = 64$ ）。

最终每层需要缓存的是 $c_i$ （ $d_c = 512$ 维）加上带 RoPE 的共享 Key（ $d_r = 64$ 维），总共 $576$ 维，相比纯 $c_i$ 增幅不大，但成功引入了位置信息。

本文符号与论文符号的对照

下表列出本文符号与 DeepSeek-V2 论文¹符号的映射（D = down-projection，U = up-projection）：

本文	论文	说明
$x_i$	$h_t$	当前层输入
$c_i = x_i W_c$	$c_t^{KV} = W^{DKV} h_t$	KV 压缩向量（ $d_c$ 维）
$c_t^Q = x_t W_c^Q$	$c_t^Q = W^{DQ} h_t$	Q 压缩向量（ $d_c^Q$ 维）
$q^{(s)} = c_t^Q W_Q^{(s)}$ （ $W_{qc}^{(s)}$ ）	$q_t^C = W^{UQ} c_t^Q$	per-head Query（non-RoPE）
$k_i^{(s)} = c_i W_K^{(s)}$ （ $W_{kc}^{(s)}$ ）	$k_t^C = W^{UK} c_t^{KV}$	per-head Key（non-RoPE）
$v_i^{(s)} = c_i W_V^{(s)}$	$v_t^C = W^{UV} c_t^{KV}$	per-head Value
$c_t^Q W_{qr}^{(s)} R_t$	$R_t W^{QR} c_t^Q$	per-head RoPE Query
$x_i W_{kr} R_i$	$R_i W^{KR} h_t$	共享 RoPE Key（输入为 $h_t$ ）
$q^{(s)} k_i^{(s)\top}$	${q_t^{(s)}}^\top k_i^{(s)}$	attention 打分（转置在不同侧）

现在可以看懂 DeepSeek-V2 论文里的这张图了。

MLA 架构（来自 DeepSeek-V2 技术报告）

Prefill 与 Decode 的计算策略

一个自然的追问：既然矩阵吸收这么好，为什么不在 Prefill 和训练阶段也用？

答案藏在计算复杂度的阶次里³。Attention 的计算可以拆为两部分：

投影部分（计算 Q、K、V 和输出投影）：复杂度与 $N$ 线性。
attention 计算部分（ $QK^\top$ 打分 + 加权求和）：当 $N_q > 1$ 时，复杂度与 $N^2$ 成正比。

矩阵吸收减少投影部分的计算量（消掉 $d_h$ 因子），但增加 attention 计算部分的有效头维度（从 $d_h$ 变成 $d_c$ ）。

在 Decode 阶段（ $N_q = 1$ ）， $QK^\top$ 只是向量与矩阵的乘法， $O(d_c N)$ ， $N^2$ 项不存在。主导项来自投影部分，矩阵吸收消掉的 $d_h$ 因子（128 倍）远大于 attention 增加的计算量（3-4 倍），划算。

在 Prefill/训练 阶段（ $N_q = N$ ）， $QK^\top$ 变成 $N \times N$ 矩阵乘法， $O(d_c N^2)$ 。 $N^2$ 项占主导，矩阵吸收把头维度从 $d_h$ 增大到 $d_c$ ，直接让这一项膨胀 3-4 倍，而投影部分的节省只是 $O(N)$ 级别，不划算。

所以：Decode 用矩阵吸收形式（只需读取共享的 $c$ ），Prefill 和训练用非吸收形式（显式计算 per-head 的 KV）。 两种形式数学上完全等价，各自在不同的计算场景下更高效。

结语

回到开头的问题：MLA 到底做了什么？

MLA 重新定义了”该缓存什么”。 MQA 和 GQA 在 KV Cache 内部做取舍，减少 KV 头数，用表达能力换存储。MLA 跳出了这个框架：放弃缓存 KV，转而缓存一个所有头共享的低维向量 $c$ ，通过矩阵吸收消除额外的计算开销，再通过 per-head 的可学习投影从 $c$ 中恢复出多头的表达能力。最终的结果是缓存压缩比约 56.9x（算上 RoPE 的额外维度），而模型能力不降反升。

整个设计用到的数学工具都是线性代数的基本操作：结合律、低秩投影、恒等变换。真正的巧妙在于把这些简单组件组合成了一个系统，在计算能力和显存带宽间的 trade-off 里找到了一个此前不存在的平衡点。

如果希望进一步了解工程实现、源码级细节，推荐阅读⁷。

附录：为什么说 GQA 是”分割和复制”

正文中提到”GQA 采用的是分割和复制”，这是苏剑林²提出的一个视角，把 GQA 和 MLA 放进同一个”压缩→恢复”的统一框架中比较。

常规的理解是：GQA 减少 KV 头数， $n_h$ 个 Q 头分成 $g$ 组，每组复用同一份 KV。但苏剑林指出，如果把 GQA 所有组的 KV 拼成一个向量：

c_i = \left[k_i^{(1)}, \cdots, k_i^{(g)}, v_i^{(1)}, \cdots, v_i^{(g)}\right] = x_i \left[W_K^{(1)}, \cdots, W_K^{(g)}, W_V^{(1)}, \cdots, W_V^{(g)}\right]

这个 $c_i$ 的维度是 $g(d_k + d_v) < d$ ，所以 $x_i \to c_i$ 本身就是一个低秩投影，只不过投影矩阵恰好是各组 $W_K, W_V$ 的拼接。

接下来，GQA 从 $c_i$ 恢复每个头的 KV 的方式是：分割（按组切片取出 $k^{(j)}, v^{(j)}$ ）+ 复制（同组内的 $n_h/g$ 个 Q 头共用这份 KV）。这些都是最简单的线性操作。

而 MLA 把这一步换成了每个头独立的可学习投影 $W_K^{(s)}, W_V^{(s)}$ ，自由度高得多。从这个视角看，GQA 和 MLA 的区别不在于是否做低秩投影，而在于从低维向量中恢复 per-head KV 的机制：分割复制 vs 可学习投影。

参考资料

DeepSeek-AI, DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, 2024. arXiv:2405.04434 ↩ ↩² ↩³
苏剑林, 缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA, 2024. kexue.fm/archives/10091 ↩ ↩² ↩³
我是那谁, 带你从头发明MLA, 2025. 知乎专栏 ↩ ↩² ↩³
Sholto Douglas, Anselm Levskaya et al., Scaling Book: Inference, 2025. jax-ml.github.io/scaling-book/inference ↩
这里的”分割和复制”是苏剑林提出的回溯视角，把 GQA 纳入和 MLA 相同的”低秩投影→恢复 per-head KV”框架中比较。如果这个说法不好理解，可以参考文末附录的展开说明。 ↩
苏剑林, Transformer升级之路：2、博采众长的旋转式位置编码, 2021. kexue.fm/archives/8265 ↩
罗西的思考, 探秘Transformer系列之（28）--- DeepSeek MLA, 2025. 博客园——覆盖了矩阵吸收、RoPE 解耦、源码实现和 MHA→MLA 转换等细节，适合作为补充阅读。 ↩