MLA 如何压缩 KV Cache：低秩潜变量、矩阵吸收与解耦 RoPE

在自回归推理中，每生成一个新 token，都要读取历史 token 的 Key 和 Value。上下文越长、batch 越大，KV Cache 带来的显存占用和内存带宽压力就越明显。

MQA 和 GQA 通过减少完整 KV head 的数量来缓解这个问题；MLA 则进一步改变缓存对象，只保存能够生成多头 K、V 的低维 latent。本文先从 Single-Head Attention、MHA、MQA 和 GQA 的关系出发，再解释 MLA 的低秩压缩、矩阵吸收与解耦 RoPE。

从 Single-Head Attention 到 GQA

Single-Head Attention

从最基础的 Single-Head Attention 计算公式开始，假设输入序列为 $X\in \mathbb{R}^{n\times d}$，特征矩阵为 $W_Q \in \mathbb{R}^{d\times d_q}$、$W_K \in \mathbb{R}^{d\times d_k}$ 和 $W_V \in \mathbb{R}^{d\times d_v}$。其中：

• $n$ 表示序列长度
• $d$ 表示每个 token 的特征维度
• $d_q$、$d_k$、$d_v$ 分别表示 Q、K、V 的特征维度

为简化后续记号，本文将 Q、K、V 的列维度统一用 $d_h$ 表示。

则 Attention 的计算公式可表示为：

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$ $$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_h}}\right)V$$

然而，对于 Single-Head Attention 来说，Q、K、V 只能在一个子空间中进行注意力计算，这可能限制模型的表达能力。MHA 将这个总宽度切成多个子空间，让不同注意力头学习不同的关系。

MHA：更丰富的表达能力

在 MHA 中，我们把表达空间划分成多个子空间：

$$Q_i=XW_{Q_i}, \quad K_i=XW_{K_i}, \quad V_i=XW_{V_i}, \quad i=1,2,\ldots,n_h$$

为了与总宽度为 $d$ 的 Single-Head Attention 对齐，常见实现令：

$$n_h d_h=d,\qquad Q=\operatorname{Concat}(Q_1,\ldots,Q_{n_h})\in\mathbb{R}^{n\times d}$$

K 和 V 同理。也就是说，MHA 不是复制 $n_h$ 份宽度为 $d$ 的 Q、K、V，而是把总宽度 $d$ 切成 $n_h$ 份，每份宽度为 $d_h$。

每个头都有独立的注意力矩阵：

$$\text{head}_i=\text{softmax}\left(\frac{Q_iK_i^T}{\sqrt{d_h}}\right)V_i$$

最终再拼接所有头的输出，并通过一个线性变换得到最终的输出：

$$\text{Output} = \text{Linear}\left(\left[\text{head}_1; \text{head}_2; \ldots; \text{head}_{n_h}\right]\right)$$

MHA 通过多个独立的注意力头丰富了不同子空间中的表达能力，但每个头都需要保存自己的 Key 和 Value。由于总宽度仍是 $n_h d_h=d$，它并没有减少 KV Cache。

在长上下文场景下，这会带来较大的显存占用与内存带宽压力。为了解决这一问题，作为 KV Cache 优化方案的 Multi-Query Attention（MQA）便被正式提出。

MQA：所有 Query head 共享一组 K、V

在 MQA 中，所有注意力头共享同一组 Key 和 Value：

$$Q_i=XW_{Q_i}, \quad K=XW_K, \quad V=XW_V, \quad i=1,2,\ldots,n_h$$

后续的计算与 MHA 相同，不再赘述。

MQA 让所有 Query head 共享同一组 Key 和 Value，虽然显著降低了 KV Cache，但也可能带来一定的表达能力损失。为在推理效率与模型质量之间取得更好的平衡，Grouped-Query Attention（GQA）被提出。

GQA：MHA 和 MQA 的折中

GQA 在 MQA 的基础上引入了介于 1 和 Query head 数量之间的多个 KV head：每组 Query head 共享一组 K、V。

$$Q_i=XW_{Q_i},\quad K_j=XW_{K_j},\quad V_j=XW_{V_j}, \quad i=1,2,\ldots,h_q,\quad j=1,2,\ldots,h_k$$

其中，$h_q$ 表示 Query head 的数量，$h_k$ 表示 Key/Value head 的数量。GQA 中通常满足：

$$h_q = g h_k$$

其中，$g=\frac{h_q}{h_k}$ 表示每个 Key/Value head 对应的 Query head 数量。也就是说：

$$Q_{(j-1)g+1},\ldots,Q_{jg} \quad\text{共享}\quad K_j,V_j$$

在这个统一表示下，MQA 与 MHA 都是 GQA 的边界情况。

KV head 数量如何决定缓存量

具体而言，Single-Head Attention、MHA、MQA 与 GQA 的 Q、K、V 关系如下图所示：

从 MHA 到 MQA、GQA，优化 KV Cache 的方式都是减少需要保存的 KV head 数量：MHA 为每个 Query head 保留一组独立的 K、V，MQA 让所有 Query head 共享一组 K、V，而 GQA 则介于两者之间。

对于长度为 $n$ 的序列，统一记 KV head 数量为 $h_k$，则每层需要缓存的元素数量为：

$$N_{\mathrm{KV}}=2n h_k d_h$$

因此，MHA 对应 $h_k=h_q$，MQA 对应 $h_k=1$，GQA 则取 $1<h_k<h_q$。

需要注意的是，KV Cache 保存的不是权重矩阵 $W_K$ 和 $W_V$，而是每一层为历史 token 计算出的 Key 和 Value。接下来讨论的缓存量，都以这些中间结果所包含的元素数量为准。

MLA：进一步压缩 KV Cache

MQA 和 GQA 通过共享完整的 K、V 来减少缓存，但共享程度越高，越可能限制不同注意力头的表达能力。DeepSeek 团队在《DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model》中提出了另一种思路：不再继续减少 KV head 的数量，而是把所有头的 Key 和 Value 联合压缩到一个低维 latent 中。这就是多头潜在注意力机制（Multi-Head Latent Attention，MLA）。

MLA 的基本思路

先看 MHA 需要缓存什么。设第 $t$ 个 token 在当前注意力层的输入为 $h_t\in\mathbb{R}^{1\times d}$，共有 $n_h$ 个注意力头，每个头的维度为 $d_h$。所有头拼接后的总宽度为 $n_h d_h$；在常见的 MHA 实现中它通常等于 $d$，但这并不是数学上的必要条件。该 token 的 Key 和 Value 分别为：

$$k_t=\operatorname{Concat}(k_{t,1},\ldots,k_{t,n_h}),\qquad v_t=\operatorname{Concat}(v_{t,1},\ldots,v_{t,n_h})$$

其中 $k_t,v_t\in\mathbb{R}^{1\times n_h d_h}$。因此，对于长度为 $n$ 的序列，MHA 每层需要缓存的元素数量为：

$$N_{\mathrm{MHA}}=2n\,n_h d_h$$

MLA 不再直接保存这些完整的 Key 和 Value，而是先把 $h_t$ 下投影为一个共享的 KV latent：

$$c_t^{KV}=h_tW^{DKV},\qquad W^{DKV}\in\mathbb{R}^{d\times d_c},\qquad c_t^{KV}\in\mathbb{R}^{1\times d_c}$$

然后，每个注意力头再通过各自的上投影，从同一个 $c_t^{KV}$ 中得到所需的 Key 和 Value：

$$k_{t,i}^{C}=c_t^{KV}W_i^{UK},\qquad v_{t,i}^{C}=c_t^{KV}W_i^{UV}, \qquad i=1,2,\ldots,n_h$$

其中 $W_i^{UK},W_i^{UV}\in\mathbb{R}^{d_c\times d_h}$。虽然所有头共享同一个 latent，但每个头仍然拥有独立的上投影，因此 MLA 并不等同于让所有头直接共享同一组完整的 K、V。

从整个序列看，令 $C^{KV}\in\mathbb{R}^{n\times d_c}$ 表示所有 token 的 KV latent，并将各头的上投影矩阵横向拼接为 $W^{UK}$ 和 $W^{UV}$，则：

$$K^C=C^{KV}W^{UK},\qquad V^C=C^{KV}W^{UV},\qquad \operatorname{rank}(K^C),\operatorname{rank}(V^C)\le d_c$$

也就是说，MLA 假设所有头的内容 K、V 都可以由同一个低维 latent 表示生成。这是模型训练时学习到的低秩约束，并不是对任意 MHA K、V 的无损压缩。

这里的压缩关系是：

$$d_c\ll n_h d_h$$

$d_c$ 比较的是所有注意力头拼接后的宽度 $n_h d_h$，而不是单个头的宽度 $d_h$。事实上，DeepSeek-V2 设置的是 $d_c=4d_h$；它依然远小于所有头拼接后的总宽度。

为什么只需要缓存 KV latent？

由上面的上投影关系可知，保存 $c_j^{KV}$ 足以恢复 MLA 自身生成的内容 K、V。MLA 的关键并不是无法恢复这些 K、V，而是推理时不必真的把它们恢复出来。采用解耦 RoPE 时，位置分支的 Key 还需要单独缓存，后文再讨论。

如何通过结合律避免恢复完整 K/V？

继续只看内容分支。对于第 $i$ 个头，有：

$$K_i^{C}=C^{KV}W_i^{UK},\qquad V_i^{C}=C^{KV}W_i^{UV}$$

整个序列的 内容 Query $Q_i^C$ 与 Key 的分数矩阵可以重新结合为：

$$Q_i^{C}(K_i^{C})^T =\left(Q_i^{C}(W_i^{UK})^T\right)(C^{KV})^T$$

因此，$W_i^{UK}$ 可以预先并入 Query 一侧的投影，计算注意力分数时直接读取 $C^{KV}$，不必显式恢复完整的 Key。

记 $A_i\in\mathbb{R}^{n\times n}$ 为第 $i$ 个头经过 Softmax 和因果掩码后的注意力矩阵。标准 Value 路径为：

$$H_i=A_iV_i^{C} \quad\Longrightarrow\quad H_i=A_iC^{KV}W_i^{UV}$$

令

$$W^{UV}_{\mathrm{blk}} =\operatorname{BlockDiag} \left(W_1^{UV},\ldots,W_{n_h}^{UV}\right), \qquad W'_O=W^{UV}_{\mathrm{blk}}W_O$$

则所有注意力头的输出可以写成：

$$O =\operatorname{Concat} \left(A_1C^{KV},\ldots,A_{n_h}C^{KV}\right)W'_O$$

因此，$W_i^{UV}$ 可以预先与 $W_O$ 合并；运行时只需对缓存的 $C^{KV}$ 做注意力加权聚合，无需显式恢复完整的 Value。

解耦 RoPE：单独保留位置信息

RoPE（Rotary Position Embedding）通过位置相关的旋转矩阵，将位置信息注入 Query 和 Key。对于位置 $t$ 的 Query 与位置 $j$ 的 Key，可简写为：

$$\widehat q_t=q_tR_t,\qquad \widehat k_j=k_jR_j$$

其中 $R_t$ 和 $R_j$ 分别表示对应位置的旋转矩阵。内积 $\widehat q_t\widehat k_j^T$ 中包含 $R_tR_j^T$，因此能够体现 token 之间的相对位置关系。

如果直接对 MLA 的内容 Key 使用 RoPE，注意力分数应为：

$$s_{tji}^{\mathrm{full}} =\left(q_{t,i}^{C}R_t\right) \left(c_j^{KV}W_i^{UK}R_j\right)^T =q_{t,i}^{C}R_tR_j^T(W_i^{UK})^T(c_j^{KV})^T$$

这种写法保留了完整内容 Key 的 RoPE，但 $R_j$ 随历史位置 $j$ 变化，使 Query 侧的变换也依赖 $j$。因此，$W_i^{UK}$ 无法预先吸收到一个与历史位置无关的 Query 投影中。DeepSeek 在 MLA 中做出的取舍是：不再对内容分支的 Key 直接使用 RoPE，而是增加一条独立的位置分支：

$$q_{t,i}=\operatorname{Concat}(q_{t,i}^{C},q_{t,i}^{R}),\qquad k_{j,i}=\operatorname{Concat}(k_{j,i}^{C},k_{j}^{R})$$

其中 $q_{t,i}^{R},k_j^{R}\in\mathbb{R}^{1\times d_h^R}$ 使用 RoPE，$k_j^{R}$ 由所有注意力头共享。因此，注意力分数可以拆成内容相似度与位置相似度两部分：

$$q_{t,i}k_{j,i}^{T} =q_{t,i}^{C}(k_{j,i}^{C})^{T} +q_{t,i}^{R}(k_j^{R})^{T}$$

将各位置的 $q_{t,i}$ 和 $k_{j,i}$ 沿序列维度堆叠为 $Q_i$ 和 $K_i$，最终的注意力矩阵为：

$$A_i=\operatorname{softmax}\left( \frac{Q_iK_i^T}{\sqrt{d_h+d_h^R}}+M \right)$$

其中 $M$ 是因果掩码。这就是前文 Value 路径中使用的 $A_i$：解耦 RoPE 改变了注意力权重的生成方式，但不影响 Value 侧的矩阵吸收。

这并不是对 $s_{tji}^{\mathrm{full}}$ 的等价展开，而是 MLA 的结构性取舍：内容分支不使用 RoPE，以保留 latent 缓存和矩阵吸收优化；位置分支单独承担位置建模，并额外缓存 $k_j^{R}$。这就是“解耦 RoPE”的含义。

MLA 的 KV Cache 节省效果

以 DeepSeek-V2 的配置为例：

$$n_h=128,\qquad d_h=128,\qquad d_c=512=4d_h,\qquad d_h^R=64$$

MHA 每个 token、每层需要缓存完整的多头 Key 和 Value：

$$N_{\mathrm{MHA}}=2n_h d_h =2\times128\times128=32768$$

MLA 需要缓存 KV latent 和解耦 RoPE 的位置 Key：

$$N_{\mathrm{MLA}}=d_c+d_h^R =512+64=576$$

因此，MLA 的缓存量约为 MHA 的：

$$\frac{N_{\mathrm{MLA}}}{N_{\mathrm{MHA}}} =\frac{576}{32768}\approx1.76\%$$

也就是说，在这一配置下，理论缓存量约减少 $98.24\%$。论文中的实际模型对比如下：

模型规模	MHA KV Cache	MLA KV Cache	减少比例
小型 MoE	$110.6\mathrm{K}$	$15.6\mathrm{K}$	$85.9\%$
大型 MoE	$860.2\mathrm{K}$	$34.6\mathrm{K}$	$96.0\%$

此外，相比 DeepSeek 67B，DeepSeek-V2 的整体 KV Cache 减少了 $93.3\%$。

小结

MHA、MQA、GQA 与 MLA 解决的是同一个 KV Cache 问题，但路径不同：MQA/GQA 减少完整 KV head 的数量，MLA 则改变缓存对象，把所有头的内容 K、V 联合压缩到共享的 KV latent 中。

在 MLA 内部，$c_j^{KV}$ 足以生成模型自身的内容 K、V；推理时再利用矩阵乘法的结合律，把 $W^{UK}$ 吸收到 Query 侧、把 $W^{UV}$ 吸收到输出投影侧，从而避免显式恢复完整 K、V。解耦 RoPE 则用一条额外的位置分支换取这些吸收关系，最终每个 token、每层只需缓存 $d_c+d_h^R$ 个元素。

参考资料

1. MHA · Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention Is All You Need. arXiv:1706.03762.
2. MQA · Noam Shazeer. 2019. Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150.
3. GQA · Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. 2023. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245.
4. MLA · DeepSeek-AI et al. 2024. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434.
5. 引用博客 · Ray. 2025. Learn and Code MHA, MQA, and GQA: A Guide to Efficient Attention Mechanisms in Transformers. Medium, February 7, 2025.