MLA 如何压缩 KV Cache 并加速推理

MLA 把每个 token 的多头 Key、Value 压进一个共享的低维向量，再通过矩阵乘法的结合律，直接在压缩表示上计算注意力。这样可以显著缩小 KV Cache，也能减少生成阶段反复读取历史信息的开销。

但它有一个值得追问的细节：矩阵吸收后，每个头的内容点积可能从 128 维变成 512 维，Value 聚合也跟着变宽。计算量增加，为什么推理仍可能更快？沿着低秩压缩、矩阵吸收和解耦 RoPE 走一遍，再把 prefill 与 decode 分开计算，这个问题就能解释清楚。

一 先看注意力需要保存什么

大模型每生成一个新 token，都要结合前面的上下文。注意力中的 Query（Q）表示当前要找什么，Key（K）用于匹配，Value（V）提供加权汇总的内容。一个 token 在不同层、不同头上都有各自的表示。

生成下一个 token 时，历史 K、V 可以直接复用。把它们保留在 GPU 显存中，就得到 KV Cache。它省掉了重复计算，却带来另一项成本：上下文越长，每一步要读取的历史数据越多。

先把几个容易混在一起的机制放到同一张表里：

机制	全称	主要改变什么
MHA	Multi-Head Attention	每个 Q 头对应自己的 K、V
MQA	Multi-Query Attention	所有 Q 头共享一组 K、V
GQA	Grouped-Query Attention	每组 Q 头共享一组 K、V
MLA	Multi-head Latent Attention	所有头的 K、V 由共同的低维表示生成
DSA	DeepSeek Sparse Attention	先选择相关历史位置，再做主注意力
CSA	Compressed Sparse Attention	先沿序列压缩历史，再选择相关压缩条目

MQA、GQA 调整头之间的共享关系；MLA 压缩每个 token 的联合 KV 表示；DSA、CSA 进一步改变历史访问方式。它们涉及不同的优化维度。例如 DeepSeek-V3.2 将 MLA 与 DSA 结合使用，DeepSeek-V4 的 CSA 则把序列压缩和稀疏选择结合起来。MQA 论文、GQA 论文、DSA 的 vLLM 实现说明、DeepSeek-V4 官方模型说明

下面只讨论稠密 MLA：每个 Query 仍然访问全部可见历史位置，不计入稀疏选择的收益。

同一个历史 token 怎样存

MHA 按头保存

头 1
K₁ · V₁

头 2
K₂ · V₂

头 3
K₃ · V₃

头 4
K₄ · V₄

图示 4 个头；示例配置实际有 32 个头。

MLA 保存共享表示

一份 cᴷⱽ512 维潜在向量

一份位置 Key64 维，所有头共享

不同头使用不同投影，仍有各自的注意力权重。

图 1：缓存的组织方式示意。方框数量表示共享关系，面积不表示字节数。MLA 的位置分支在第四节展开。

用同一组维度贯穿全文

以下是教学配置，不对应某个具体模型的全部参数。公式采用列向量约定，t 是当前 Query 的位置，j 是被访问的位置，i 是头编号。

符号	含义	数值
D	隐藏状态维度	4096
H	注意力头数	32
d_C	每头 Q、K 的内容维度	128
d_R	每头 Q、K 的 RoPE 维度	64
d_V	每头 Value 维度	128
r	KV 潜在向量维度	512
r_Q	Query 潜在向量维度	1024
T	decode 时可访问的历史长度，含当前 token	8192

内容维度不显式应用 RoPE；位置维度负责这部分旋转计算。所有缓存大小均按 BF16、每元素 2 字节计算。GFLOPs 按十进制，MiB 按 2²⁰ 字节；一次乘加计为 2 FLOPs。

二 低秩联合压缩怎样保留多头表示

先压到一个共同的瓶颈

暂时不考虑位置分支。一个 token 的隐藏状态有 4096 维，先下投影到 512 维：

$$\underbrace{c_j^{KV}}_{512\times1} =\underbrace{W^{DKV}}_{512\times4096} \underbrace{h_j}_{4096\times1}$$

再用两组上投影生成所有头的内容 Key 和 Value：

$$\begin{aligned} k_j^C&=W^{UK}c_j^{KV}, & W^{UK}&\in\mathbb R^{4096\times512},\\ v_j&=W^{UV}c_j^{KV}, & W^{UV}&\in\mathbb R^{4096\times512}. \end{aligned}$$

4096 维拆成 32 个头，每头各 128 维。因此，第 i 个头的两块上投影矩阵都是 128×512：

$$k_{j,i}^C=W_i^{UK}c_j^{KV},\qquad v_{j,i}=W_i^{UV}c_j^{KV}$$

各头共用同一个潜在向量，但使用不同的投影矩阵。共享的是生成 K、V 的基础表示，最终的各头 K、V 可以不同。DeepSeek-V2 的低秩联合压缩

4096 维输入 → 512 维共享表示 → 多头输出

隐藏状态 h4096 维

→

潜在向量 cᴷⱽ512 维历史缓存保留这里

→

32 个内容 K
每头 128 维

32 个 V
每头 128 维

图 2：右侧两个分支分别生成 4096 个元素。图中暂不包含 RoPE；decode 的吸收路径会避免实际生成整段历史的右侧输出。

低秩限制在哪里

把所有头的内容 K、V 叠起来：

$$\begin{bmatrix}k_j^C\\v_j\end{bmatrix} =\underbrace{\begin{bmatrix}W^{UK}\\W^{UV}\end{bmatrix}}_{8192\times512} \underbrace{W^{DKV}}_{512\times4096}h_j$$

这两个矩阵相乘，秩最多为 512。虽然联合输出有 8192 维，它只能落在一个最多 512 维的线性子空间中。这里的“低秩”相对于所有头的联合表示而言，512 大于单头的 128 并不矛盾。

这是通过训练学习的结构约束，不要求用 SVD，也不能保证把任意已有 MHA 无损改造成 MLA。需要区分两件事：低秩结构限制了表示能力；在已经训练好的 MLA 内改变乘法顺序，则可以保持数学等价。

实际实现还可以在下投影之后加入 RMSNorm。此时从隐藏状态到 KV 的全过程不再是上述纯线性映射，但上投影输出仍由 512 维瓶颈生成。后面的吸收公式对归一化后的潜在向量同样成立。

三 矩阵吸收把历史展开移到当前计算

只把缓存变小还不够。如果每生成一个 token，都把全部历史潜在向量展开成 K、V，展开计算和临时内存开销仍然很大。MLA 利用线性运算的结合律，绕过这个步骤。

Key 侧把上投影移到当前 Query

暂不考虑 RoPE，一个头的内容打分为：

$$\begin{aligned} (q_{t,i}^C)^{\mathsf T}k_{j,i}^C &=(q_{t,i}^C)^{\mathsf T}W_i^{UK}c_j^{KV}\\ &=\left((W_i^{UK})^{\mathsf T}q_{t,i}^C\right)^{\mathsf T}c_j^{KV}. \end{aligned}$$

定义变换后的 Query：

$$\underbrace{\widetilde q_{t,i}}_{512\times1} =\underbrace{(W_i^{UK})^{\mathsf T}}_{512\times128} \underbrace{q_{t,i}^C}_{128\times1}$$

此后直接计算 512 维 Query · 512 维历史潜在向量。一次 Query 变换就能用于所有历史位置，不需要逐位置生成 128 维的内容 Key。

用一个缩小的数值例子核对：

$$q=\begin{bmatrix}1\\2\end{bmatrix},\quad W=\begin{bmatrix}1&0&2\\0&1&1\end{bmatrix},\quad c=\begin{bmatrix}3\\4\\5\end{bmatrix}$$

先算 Wc 得到 [13,9]，与 q 点积为 31；先算 Wᵀq 得到 [1,2,4]，与 c 点积也为 31。变化的是计算位置，数值没有改变。

若内容 Query 本来由 q=W_Q h 生成，数学上还可以把 W_UKᵀ W_Q 合成新权重。工程上不一定真的合成：新矩阵可能更大，保留原权重、只调整执行顺序，也能避免展开历史 K。

Value 侧先聚合再上投影

注意力权重记为 α。普通路径先得到历史 Value，再加权求和：

$$\begin{aligned} o_{t,i} &=\sum_{j\le t}\alpha_{t,i,j}v_{j,i}\\ &=\sum_{j\le t}\alpha_{t,i,j}W_i^{UV}c_j^{KV}\\ &=W_i^{UV}\underbrace{\left(\sum_{j\le t}\alpha_{t,i,j}c_j^{KV}\right)}_{z_{t,i}\in\mathbb R^{512}}. \end{aligned}$$

因为同一个头的上投影矩阵不随历史位置变化，可以先把历史潜在向量聚成一个 512 维的 z，再对 z 做一次 512→128 投影。各个头的 α 不同，所以它们虽然共享缓存，聚合出来的 z 也不同。

同一个上投影矩阵 换一个执行位置

Key 侧

展开路径每个历史 c → Kᵢ → 与 Qᵢ 点积

吸收路径当前 Qᵢ → Q̃ᵢ → 与所有历史 c 点积

把作用于历史的投影移到当前 Query。

Value 侧

展开路径每个历史 c → Vᵢ → 加权求和

吸收路径历史 c → 加权求和 → 上投影

先得到一个聚合结果，再做 Value 投影。

图 3：两处都利用线性运算的结合关系。Softmax 仍按原定义计算，没有被移过矩阵或省略。

再把输出矩阵按头切成 W_O,i，最终输出可以写成：

$$y_t=\sum_i W_i^O W_i^{UV}z_{t,i}$$

因此 Value 上投影还可以在数学上与输出投影合并。实际代码也可以保留两个矩阵，只在聚合之后依次执行它们。DeepSeek-V3 官方参考代码的 naive 与 absorb 两条路径

四 解耦 RoPE 怎样保留吸收能力

普通 RoPE 在中间插入了位置相关矩阵

RoPE（Rotary Position Embedding，旋转位置编码）给不同位置使用不同的旋转矩阵。把位置 p 的旋转写作 R_p；高维旋转可以理解为多个二维旋转块。

若直接对内容 Q、K 应用 RoPE：

$$q'_{t,i}=R_tq_{t,i}^C,\qquad k'_{j,i}=R_jW_i^{UK}c_j^{KV}$$

点积变成：

$$(q'_{t,i})^{\mathsf T}k'_{j,i} =(q_{t,i}^C)^{\mathsf T}R_t^{\mathsf T}R_jW_i^{UK}c_j^{KV}$$

问题在于 R_tᵀR_j 夹在 Query 和上投影矩阵之间，而且随历史位置 j 改变。矩阵乘法通常不能交换顺序；若把这段运算搬到 Query 一侧，得到的变换 Query 就会依赖 j，无法再计算一次后服务全部历史位置。

普通 RoPE 仍可以与压缩缓存一起使用，但会破坏前面这条高效的统一吸收路径。直接缓存展开并旋转后的 Key，又会增大缓存。

把显式位置计算放进单独的分支

MLA 将每个头的 Q、K 拼成两部分：128 维的内容分支不应用 RoPE，64 维的位置分支单独旋转。位置 Query 各头不同，位置 Key 在所有头之间共享：

$$\begin{aligned} q_{t,i}&=[q_{t,i}^C;q_{t,i}^R], & q_{t,i}^R&=R_tW_i^{QR}c_t^Q,\\ k_{j,i}&=[W_i^{UK}c_j^{KV};k_j^R], & k_j^R&=R_jW^{KR}h_j. \end{aligned}$$

Query 分支先计算 c_Q = W_DQ h，其中 W_DQ 为 1024×4096；再通过每头 128×1024 的 W_UQ,i 生成内容 Query。位置 Query 的 W_QR,i 为 64×1024，共享位置 Key 的 W_KR 为 64×4096。历史 Q 无需缓存，KV Cache 的主要收益来自 KV 分支。

拼接向量的点积等于两部分点积之和，最终打分为：

$$s_{t,i,j} =\frac{\widetilde q_{t,i}^{\mathsf T}c_j^{KV}+(q_{t,i}^R)^{\mathsf T}k_j^R} {\sqrt{d_C+d_R}},\qquad \alpha_{t,i,:}=\operatorname{softmax}(s_{t,i,:})$$

内容分支继续吸收，位置分支直接计算；两部分先加分数，再做同一次 Softmax。DeepSeek-V2 的解耦 RoPE

两条分支算分数 相加后统一 Softmax

内容分支

当前 Q̃ᵢ · 历史 cᴷⱽ
512 维点积

不显式应用 RoPE
保留矩阵吸收

位置分支

旋转后的 Qᴿᵢ · Kᴿ
64 维点积

Qᴿ 各头不同
Kᴿ 在各头之间共享

两部分分数相加 → 除以 √192 → Softmax → 聚合历史 cᴷⱽ

图 4：每个历史 token 缓存 512 维 cᴷⱽ 和 64 维 Kᴿ，共 576 个元素。“内容”指该分支不显式应用 RoPE，输入隐藏状态仍可能携带之前层传递的位置信息。

缩放因子仍是 √192。原始 Q、K 的维度为 128+64，吸收只改变计算顺序，不能因为潜在空间点积变宽，就把缩放改成 √576。

五 Decode 为何能接受更多计算

注意力核心 FLOPs 确实增加

一次 decode 中，每个请求只有一个新 Query，要访问 T 个缓存位置。展开后的 QK 点积为 192 维，Value 聚合为 128 维；吸收后分别为 576 和 512 维：

$$\begin{aligned} F_{\mathrm{expanded}}&\approx2HT(d_C+d_R+d_V),\\ F_{\mathrm{absorbed}}&\approx2HT(2r+d_R). \end{aligned}$$

代入 H=32、T=8192，分别为 0.168 和 0.570 GFLOPs，后者是前者的 3.4 倍。QK 增加 3 倍，Value 聚合增加 4 倍，合计并非恰好 4 倍。

这里暂时只算 QK 和加权聚合，不包含投影、Softmax、缓存写入和 kernel 启动。展开路径可以是普通 MHA，也可以是已经展开 K、V 的同一个 MLA；两者在这个核心的形状相同。

历史数据体积降得更多

若 MHA 也采用每头 192 维 K 和 128 维 V，单层、单请求的缓存大小为：

$$\begin{aligned} M_{\mathrm{MHA}}&=2TH(d_C+d_R+d_V)=160\ \mathrm{MiB},\\ M_{\mathrm{MLA}}&=2T(r+d_R)=9\ \mathrm{MiB}. \end{aligned}$$

MLA 的核心计算仍然乘以头数 H，共享缓存的数据体积却不再乘以 H。多个头可以使用读入的同一份潜在向量，这是它提高数据复用率的基础。

单层 单请求 历史长度 8192

注意力核心 FLOPs

展开

0.168 G

吸收

0.570 G

历史 KV 缓存体积

MHA

160 MiB

MLA

9 MiB

图 5：按本文配置推算。两组条形分别按各自最大值缩放，不能跨单位比较长度；缓存体积也不等于真实 HBM 读取量。

9 MiB 是缓存体积，也是理想复用情况下读取一次的量。若 QK 与聚合分别读一次潜在向量，仅缓存读取就约为 17 MiB；头之间复用不足、分块方式不同还会产生重复流量。融合 kernel、共享内存和硬件缓存决定了实际能省下多少 HBM（高带宽显存）访问。

用时间判断瓶颈

Roofline 模型把计算与访存放到同一套单位中。令 F 为 FLOPs，M 为实际访问字节数，P 为可达到的计算吞吐，BW 为可达到的带宽，则忽略其他成本时，耗时至少受下面较大的一项限制：

$$t\gtrsim\max\left(\frac FP,\frac M{BW}\right)$$

假设系统可以达到 100 TFLOPs/s 算力和 3 TB/s 带宽，并用上面的缓存单次读取量做理想估算：

时间项	MHA	吸收后的 MLA
计算量 ÷ 算力	1.68 μs	5.70 μs
理想读取量 ÷ 带宽	55.92 μs	3.15 μs
两项的最大值	55.92 μs	5.70 μs

图 6：图中每组较长的柱决定简化时间下界。数据由本文假设推算，采用理想 KV 单次读取，不含投影、Softmax、权重访存、通信等成本，不是硬件实测或端到端加速预测。

这个算例里，MHA 的少量计算很快完成，读取却更慢；MLA 多做了一些计算，仍可能因为大幅减少数据读取而缩短时间。瓶颈甚至可能从带宽转到计算。

吸收还避免了每步重新展开历史

前面的 3.4 倍，是相对于“已经拥有展开 K、V”的注意力核心。若只保存压缩缓存，却在每次 decode 时重新展开全部历史，代价还包括：

$$F_{\mathrm{expand}}=2THr(d_C+d_V)\approx68.7\ \mathrm{GFLOPs}$$

吸收后，对当前 Query 做变换、对当前聚合结果做 Value 上投影，两项合计约为：

$$F_{\mathrm{current}}=2Hr(d_C+d_V)\approx8.39\ \mathrm{MFLOPs}$$

它们都不再随历史长度 T 增长。吸收让压缩缓存可以直接参加注意力计算，避免每步付出整段历史的展开成本。若把展开结果长期缓存，又会失去主要的缓存压缩收益。

六 Prefill 为何经常选择展开路径

Prefill 一次处理许多输入 token。同一份 K、V 可以被许多 Query 使用，分块算法还能把读入片上存储的数据反复用于计算；较大的矩阵乘也更容易利用 Tensor Core。它通常比逐 token decode 更偏向计算受限。vLLM 对 prefill 与 decode 的说明

设没有已有前缀，prompt 长度为 L。展开路径的核心矩阵形状是：

$$\begin{aligned} QK^{\mathsf T}:&\quad[L,192]\times[192,L],\\ AV:&\quad[L,L]\times[L,128]. \end{aligned}$$

吸收路径则是 [L,576]×[576,L] 和 [L,L]×[L,512]。对大量 Query 继续支付更宽的点积、聚合成本，计算增量就明显了。FlashAttention 可以避免把完整的 L×L 分数矩阵写到显存，但稠密注意力的核心 FLOPs 仍随 L² 增长。

对 causal attention，有效 Query-Key 对数为 N=L(L+1)/2：

$$F_{\mathrm{expanded}}\approx2HN(d_C+d_R+d_V),\qquad F_{\mathrm{absorbed}}\approx2HN(2r+d_R)$$

取 L=8192：

Prefill 路径	注意力核心 FLOPs
展开的 K、V	687 GFLOPs
吸收后的潜在空间计算	2337 GFLOPs

展开一次全部 K、V 需要约 68.7 GFLOPs，随 L 线性增长；它随后可以服务大量 Query。对于足够长的普通 prefill，付出一次展开成本、降低主要的二次方计算量，通常更划算。两条路径还各有投影、临时存储和 kernel 实现成本，表格只比较注意力核心。

因此，同一个 MLA 模型可以使用两条数学等价的执行路径：

参数相同 计算顺序随阶段改变

Prefill 大量 Query

临时展开 K、V较窄维度的标准注意力

展开一次 多个 Query 复用优先减少核心 FLOPs

长期缓存仍可保存压缩表示。

Decode 每请求一个新 Query

直接使用潜在缓存更宽维度的吸收路径

共享缓存 多头复用优先减少历史数据搬运

避免每一步重新展开整段历史。

图 7：常见的阶段选择，不是固定调度规则。长前缀上的短 prefill、推测解码等多 Query 情形，还需结合实际尺寸选择算法。

临时展开 K、V 不要求把它们作为长期缓存保存，也可以分块处理以控制峰值内存。vLLM 的 MLA 说明区分了偏计算效率和偏数据搬运效率的路径，并为 prefill、decode 提供不同后端。MLA 实现说明、MLA 后端选择

七 怎样判断 MLA 在实际服务中的收益

MLA 缩小缓存，还可能让同一张 GPU 容纳更多并发请求。更大的 batch 能让一次权重读取服务更多 token，提高整体吞吐。这项收益需要与单请求的生成延迟分开看。

具体效果取决于工作负载。上下文较短时，权重读取和投影可能占据更大比例；Query 数量增多后，注意力更容易受计算限制。不同请求的历史 KV 通常不同，增加 batch 并不会像共享模型权重那样，自动让它们复用同一份 KV 数据。

比较基线也很重要。本文的缓存算例使用 MHA；GQA、MQA 本身已经共享 KV，MLA 相对它们的带宽优势会小一些。各模型的头数、潜在维度、精度和模型质量也不相同，不能直接从这组教学数字推出某个模型的领先幅度。

观察一个真实实现时，可以沿着三件事检查：长期缓存是否只保留潜在表示和共享位置 Key，decode 是否避免展开历史 KV，以及 kernel 是否让多个头有效复用已读入的数据。随后分别测量 prefill 的处理速度、decode 的每 token 延迟和并发吞吐，再结合 HBM 流量与计算利用率判断瓶颈。

附录 一次 Decode 的核心张量形状

这段代码展示吸收路径的计算顺序，省略投影、归一化、RoPE 计算和缓存写入。qR、KR 已经旋转，缓存只包含当前 Query 可见的位置。

# qC: [B, H, d_C]   = [B, 32, 128]
# qR: [B, H, d_R]   = [B, 32, 64]
# C:  [B, T, r]     = [B, T, 512]
# KR: [B, T, d_R]   = [B, T, 64]
# UK: [H, d_C, r]   = [32, 128, 512]
# UV: [H, d_V, r]   = [32, 128, 512]
# WO: [D, H * d_V]  = [4096, 4096]

q_lat = torch.einsum("bhd,hdr->bhr", qC, UK)   # [B, 32, 512]
scores = (
    torch.einsum("bhr,btr->bht", q_lat, C)
    + torch.einsum("bhe,bte->bht", qR, KR)
)                                           # [B, 32, T]
alpha = torch.softmax(scores / (128 + 64) ** 0.5, dim=-1)
z = torch.einsum("bht,btr->bhr", alpha, C)     # [B, 32, 512]
o = torch.einsum("bhr,hvr->bhv", z, UV)        # [B, 32, 128]
y = o.flatten(1) @ WO.T                      # [B, 4096]

代码用于核对维度和代数关系，不代表优化 kernel 的访存行为。实际融合实现会分块计算分数、在线 Softmax 和输出累加，避免显式物化整段 scores、alpha，并尽量复用同一块潜在缓存。

参考资料

• DeepSeek-V2 论文：MLA 的低秩压缩、矩阵吸收和解耦 RoPE。
• DeepSeek-V3 官方推理代码：对照 naive 和 absorb 两条执行路径。
• vLLM 的 MLA 实现说明：计算效率、数据搬运与阶段选择。
• FlashAttention 论文：通过分块与片上复用减少注意力的 HBM 访问。
• 从 CUDA 代码到 GPU 执行：理解线程分组、SM 与硬件执行。
• 视频稀疏注意力 VSA：进一步了解如何通过选择历史区块减少注意力访问。