MLA 完全拆解：DeepSeek 低秩 KV 压缩注意力的原理与全部实现细节

背景：KV Cache 为什么是大模型推理的第一瓶颈#

大语言模型做推理时有一个绕不开的机制：KV Cache。自回归生成第 $t$ 个 token 时，注意力层需要计算当前 query 与前面所有历史 token 的 key、value 的相关性。如果每个 step 都重新计算一遍历史 token 的 key 和 value，计算量会随上下文长度平方增长——这是完全不可接受的。所以推理系统会把每一层每个历史 token 的 key 和 value 缓存在显存里，每次生成新 token 时只计算当前 token 的 key/value 并追加到缓存中。

这个缓存就是 KV Cache。它的代价是显存容量和显存带宽：缓存随 token 数线性增长，并且 decode 阶段每个 step 都要把整段 KV 从显存读一遍。

先算一笔账，看标准 Multi-Head Attention（MHA）的 KV Cache 有多大。设模型有 $l$ 层，每层 $n_h$ 个注意力头，每头维度 $d_h$ 。每个 token 每层要缓存 K 和 V 两份，即 $2 n_h d_h$ 个元素。以 DeepSeek-V2 的配置为例（ $n_h = 128$ ， $d_h = 128$ ， $l = 60$ ）：

2 n_h d_h l = 2 \times 128 \times 128 \times 60 = 1{,}966{,}080 \text{ 个元素/token}

BF16 下每个元素占 2 字节，一个 token 的 KV 就要约 4 MB 显存。于是：

128K 上下文，单序列： $32{,}768 \times 131{,}072 \times 60 \times 2 \text{ B} \approx 515 \text{ GB}$ 。一张 H100/H800 的 80 GB 显存连一条长序列的 KV 都装不下；
32K 上下文，8 个并发请求：约 1030 GB，需要几十张卡；
即使只算单序列 32K：MHA 也要约 128 GB。

这就是长上下文推理一直做不起来的根本原因。而 KV Cache 还不是只占空间——decode 阶段每个 step 都要把这整段 KV 从 HBM 读进计算单元，显存带宽同样被它占满。后面会看到，MLA 的另一个巨大收益恰恰来自带宽。

DeepSeek 团队在 2024 年 5 月发布的 DeepSeek-V2（论文 arXiv:2405.04434）中给出了他们的答案：MLA（Multi-head Latent Attention，多头潜变量注意力）。V2 总参 236B、每 token 激活 21B，上下文 128K。官方报告的数字是：相比 DeepSeek 67B，KV Cache 减少 93.3%，最大生成吞吐提升到 5.76 倍，训练成本节省 42.5%，而模型能力反而更强。之后 DeepSeek-V3（671B/37B 激活）、V3.1、V3.2（DeepSeek-V3 技术报告）全部沿用 MLA，直到今天它仍是 DeepSeek 全系模型推理的基石。

DeepSeek-V2 整体架构：注意力用 MLA 压缩 KV Cache，FFN 用 DeepSeekMoE 稀疏化

DeepSeek-V2 架构图：注意力模块使用 MLA 显著压缩 KV Cache 保证推理效率，FFN 使用 DeepSeekMoE 以稀疏结构降低训练成本。（来源：DeepSeek-V2 论文 Figure 2）

本文把 MLA 拆成四个层次讲透：低秩联合压缩（它是什么）、RoPE 解耦（为什么必须这么设计）、矩阵吸收（推理时怎么省算力）、Kernel 与框架实现（FlashMLA、官方推理代码、vLLM 等落地方式）。

核心思想：把 K 和 V 一起压成一个潜变量#

在 MHA 之前，业界已经有两个减 KV Cache 的经典方案：

MQA（Multi-Query Attention），由 Shazeer 在 2019 年提出（论文 arXiv:1911.02150）：所有 query 头共享同一组 K/V，KV 从 $2 n_h d_h$ 降到 $2 d_h$ ，压缩 $n_h$ 倍，但表达能力明显下降——所有头看到同样的键值，头间的差异化注意力模式被牺牲掉；
GQA（Grouped-Query Attention），Google 在 2023 年提出（论文 arXiv:2305.13245）：把 $n_h$ 个 query 头分成 $n_g$ 组，每组共享一份 K/V，KV 为 $2 n_g d_h$ 。LLaMA-2 70B 用 $n_g = 8$ （即每 8 个 query 头共享一组 KV），工程上验证性能损失不大，但压缩比也只有 8 倍。

这两个方案的本质都是减少头的数量：让多个 query 头共享同一个 key、value。它们的共同问题是压缩比受限于”共享程度”，而共享过头就损失性能。

MLA 换了一个完全不同的思路。它的核心洞察是：K 和 V 的信息冗余主要不在”头之间”，而在”维度之间”。 $n_h d_h$ 维的 K/V 向量中，大部分信息可以压缩到一个更低维的表示里，需要时再通过一个线性投影恢复出来。换句话说，MQA/GQA 砍的是”头”，MLA 压的是”表示的秩”。

具体做法是低秩联合压缩：对第 $t$ 个 token 的输入 $h_t \in \mathbb{R}^{d}$ （ $d$ 是隐藏维度），先做一个下投影，把 K 和 V 一起压进一个低维潜变量 $c_t^{KV} \in \mathbb{R}^{d_c}$ ：

c_t^{KV} = W^{DKV} h_t

k_t^{C} = W^{UK} c_t^{KV}, \qquad v_t^{C} = W^{UV} c_t^{KV}

其中 $W^{DKV} \in \mathbb{R}^{d_c \times d}$ 是下投影矩阵， $W^{UK}, W^{UV} \in \mathbb{R}^{n_h d_h \times d_c}$ 是 K、V 各自的上投影矩阵， $d_c \ll n_h d_h$ 是压缩维度。注意这里有个关键细节：K 和 V 共享同一个潜变量 $c_t^{KV}$ ，但各自用自己的上投影矩阵恢复。所以每个头仍然有自己独立的 K/V（通过 $W^{UK}$ 、 $W^{UV}$ 的对应行段），不像 MQA 那样所有头强制共享同一个向量——这是 MLA 表达能力强于 MQA/GQA 的根源。

用线性代数的语言说，这个”下投影 + 上投影”的结构就是矩阵分解：把完整的 K/V 投影矩阵分解成 $W^{DKV}$ 与 $W^{UK}$ 、 $W^{UV}$ 的低秩乘积。潜变量维度 $d_c$ 就是秩。下投影丢弃”不重要的方向”（相当于 SVD 截断），上投影把保留的方向展开回完整维度。这听起来像有损压缩，但论文的消融实验（后面有数据）显示，MLA 在把 KV 压到 1/57 之后，效果不仅没有变差，甚至比完整 MHA 还好一点——这是因为低秩压缩本身带有正则化效果，且 $d_c = 512$ 对 $n_h d_h = 16384$ 的压缩并没有伤到模型真正依赖的低维子空间。

MHA、GQA、MQA 与 MLA 的对比：MLA 将 K 和 V 联合压缩进一个潜向量 c，推理时只缓存该潜向量（虚线框），彻底绕开了”缓存完整 K/V”。（来源：DeepSeek-V2 论文 Figure 3）

上面这张对比图值得逐格看：MHA 每个 query 头有自己独立的 K/V（蓝色方块每个头一份）；MQA 所有头共享一对 K/V；GQA 每几个头共享一组；而 MLA 的 K 和 V 先由同一个下投影压成一个小的潜向量（图最右侧的 $c$ ），再由不同的上投影矩阵分别恢复成每个头的 K、V——但推理时根本不缓存恢复出来的完整 K/V，只缓存潜向量本身。

Query 侧也可以做同样的低秩压缩：

c_t^{Q} = W^{DQ} h_t, \qquad q_t^{C} = W^{UQ} c_t^{Q}

$W^{DQ} \in \mathbb{R}^{d_c' \times d}$ 、 $W^{UQ} \in \mathbb{R}^{n_h d_h \times d_c'}$ 分别是 query 的下投影和上投影， $d_c'$ 是 query 压缩维度。这里有个值得记住的差异：query 压缩并不减少 KV Cache（推理时缓存的是历史 token 的 K/V，不是 Q），它的作用在训练侧——减少激活（activation）内存占用，后面”训练细节”一节会展开。

KV 缓存账：MLA 到底省了多少#

先明确 MLA 推理时每层每 token 到底缓存什么。答案目前只有一部分： $c_t^{KV}$ （ $d_c$ 维）。但还差一块——用于位置编码的解耦 key $k_t^{R}$ （ $d_h^R$ 维），它的来龙去脉下一节详细讲。先接受这个结论：MLA 每层每 token 缓存

(d_c + d_h^{R}) \text{ 个元素}

DeepSeek-V2 取 $d_c = 4 d_h = 512$ ， $d_h^{R} = d_h / 2 = 64$ ，即每层每 token 缓存 $512 + 64 = 576$ 个元素。对照 MHA 的 $2 n_h d_h = 2 \times 128 \times 128 = 32{,}768$ 个元素：

\frac{32{,}768}{576} \approx 56.9 \times

论文用了一个更简洁的表达： $d_c + d_h^{R} = 4 d_h + d_h / 2 = \tfrac{9}{2} d_h$ ，即 MLA 的 KV Cache 相当于 GQA 只有 2.25 组时的规模（因为 GQA 是 $2 n_g d_h$ ，令 $2 n_g = 9/2$ 得 $n_g = 2.25$ ）。论文 Table 1 把这四种机制放在一起对比：

注意力机制	每 token 每层 KV 缓存（元素）	表达能力
MHA	$2 n_h d_h$	强
GQA	$2 n_g d_h$	中等
MQA	$2 d_h$	弱
MLA（V2 配置）	$(d_c + d_h^{R}) \approx \tfrac{9}{2} d_h$	更强

换成年费化的数字更直观（bf16，60 层）：

单 token：MHA 约 64 KB vs MLA 约 1.1 KB；
单序列 32K 上下文：MHA 约 128 GB vs MLA 约 2.2 GB；
单序列 128K 上下文：MHA 约 515 GB vs MLA 约 9.1 GB；
decode 阶段每步读取的 KV 量：同样差 56.9 倍——这是 MLA 提升生成吞吐最直接的来源。

关于官方”减少 93.3%“这个数字需要说明口径：它是 DeepSeek 相对 DeepSeek 67B 的系统级对比数字（67B 是 95 层、64 头、稠密 MHA，KV 总量为 $2 \times 64 \times 128 \times 95 = 1{,}556{,}480$ 元素/token，是 V2 的 576 × 60 = 34,560 元素/token 的 45 倍左右），并不是同一模型上直接替换注意力得到的”每层压缩比”。社区里常引用的”28.4 倍”则是按每层、64 头口径算的（ $16{,}384 / 576 \approx 28.4$ ）。无论哪个口径，MLA 把 KV 压到 MHA 的几十分之一这点是一致的。

论文附录 C.2 的消融实验才是同口径的证据：分别训练只差注意力的两组 MoE 模型（约 16B 小模型和约 250B 大模型），对比 MHA 与 MLA：

指标	小模型 w/ MHA	小模型 w/ MLA	大模型 w/ MHA	大模型 w/ MLA
激活参数	2.5B	2.4B	25.0B	21.5B
总参数	15.8B	15.7B	250.8B	247.4B
KV Cache/token（元素）	110.6K	15.6K	860.2K	34.6K
BBH（EM，3-shot）	37.9	39.0	46.6	50.7
MMLU（5-shot）	48.7	50.0	57.5	59.0
C-Eval（5-shot）	51.6	50.9	57.9	59.2
CMMLU（5-shot）	52.3	53.4	60.7	62.5

（数据来源：DeepSeek-V2 论文 Table 7）

大模型上 MLA 的 KV Cache 只有 MHA 的 4%，但 MMLU 从 57.5 涨到 59.0、BBH 从 46.6 涨到 50.7。这正是论文敢说”MLA 性能比 MHA 更强”的依据。注意小模型的 C-Eval 一项 MLA 略低（50.9 vs 51.6），说明低秩压缩并非零代价，只是整体上收益远大于损失。

RoPE 解耦：MLA 设计里最微妙的一环#

上面所有的压缩都是对”内容”（content）做的。但 Transformer 还需要位置信息——token 在序列中的位置必须编码进注意力分数里，否则”我打你”和”你打我”在注意力层看来没有区别。DeepSeek 沿用了 RoPE（Rotary Position Embedding，旋转位置编码，RoFormer 论文 arXiv:2104.09864）。

问题来了：RoPE 和低秩 KV 压缩是不兼容的。这是 MLA 设计中最容易被忽略、也最致命的一个坑，值得完整推导一遍。

RoPE 的做法是把 key 和 query 按位置旋转：对位置 $t$ 的向量 $x$ ，施加一个只由位置 $t$ 决定的分块对角旋转矩阵 $\mathcal{R}_t$ 。旋转后， $t$ 位置的 query 与 $s$ 位置的 key 的内积 $(\mathcal{R}_t q)^T (\mathcal{R}_s k) = q^T \mathcal{R}_{s-t} k$ 只依赖相对位置 $s - t$ ，这正是位置编码想要的性质。

假设我们对恢复出来的 key 直接施加 RoPE，即把式中的 $k_t^{C}$ 换成 $\mathcal{R}_t k_t^{C} = \mathcal{R}_t W^{UK} c_t^{KV}$ 。那么注意力分数就变成：

q_{t}^{T} (\mathcal{R}_s W^{UK} c_s^{KV}) = (W^{UK,T} \mathcal{R}_s^{T} q_t)^{T} c_s^{KV}

问题出在中间这个 $\mathcal{R}_s$ ：旋转矩阵随位置 $s$ 变化，它”插”在 $W^{UK}$ 和 query 之间。原本 $W^{UK}$ 可以被吸收进 query 侧（下节讲），因为矩阵乘法可以重结合—— $q^T (W c) = (W^T q)^T c$ 。但一旦中间夹了一个位置相关的 $\mathcal{R}_s$ ，就变成 $(W^{UK,T} \mathcal{R}_s^{T} q)^T c$ ，query 侧无法预先算出一个与位置无关的变换，因为 $\mathcal{R}_s$ 依赖的是历史 token 的位置 $s$ ，而 $s$ 对每个缓存里的 token 都不同。矩阵乘法不满足交换律， $\mathcal{R}_s$ 既不能挪到 $W^{UK}$ 右边，也不能被吸收。

后果是灾难性的：要么推理时对每个历史 token 实时做一次 $W^{UK}$ 上投影加旋转（等于放弃”只缓存潜变量”的收益，每步都要把所有历史 key 重算一遍，计算量随上下文线性增长）；要么把完整 key 缓存下来（回到 MHA 的内存问题）。两个都不能接受。

DeepSeek 的解法是解耦 RoPE（Decoupled RoPE）：把”内容”和”位置”彻底分成两条通路。内容部分走低秩压缩（可吸收），位置部分单独用一组很小的向量承载：

q_t^{R} = \operatorname{RoPE}(W^{QR} c_t^{Q}), \qquad k_t^{R} = \operatorname{RoPE}(W^{KR} h_t)

其中 $W^{QR} \in \mathbb{R}^{d_h^{R} n_h \times d_c'}$ 、 $W^{KR} \in \mathbb{R}^{d_h^{R} \times d}$ 是额外的投影矩阵， $d_h^{R}$ 是解耦向量每头的维度（V2 中取 64）。注意两个不对称的细节：

$q^R$ 每头一份（ $d_h^{R} \times n_h$ 维，随 query 头变化），而 $k^R$ 所有头共享一份（ $d_h^{R}$ 维）——这是论文式 (14)、(15) 的不对称之处；
$k^R$ 直接从 $h_t$ 投影，没有经过 $c_t^{KV}$ 的压缩，因此它无法（也不需要）被吸收，推理时作为唯一的”位置分量”原样缓存。

最终每个头参与注意力计算的 query 和 key 都是内容与位置的拼接：

q_{t,i} = [q_{t,i}^{C}; q_{t,i}^{R}], \qquad k_{t,i} = [k_{t,i}^{C}; k_t^{R}]

$[\cdot;\cdot]$ 表示沿最后一维拼接，下标 $i$ 表示第 $i$ 个头。于是注意力分数为：

o_{t,i} = \sum_{j=1}^{t} \operatorname{Softmax}_j \left( \frac{q_{t,i}^{T} k_{j,i}}{\sqrt{d_h + d_h^{R}}} \right) v_{j,i}^{C}

u_t = W^{O} [o_{t,1}; o_{t,2}; \dots; o_{t,n_h}]

softmax 的缩放因子用 $\sqrt{d_h + d_h^{R}} = \sqrt{128 + 64} = \sqrt{192}$ ，因为参与内积的向量拼接后维度是 $d_h + d_h^{R}$ ——这是一个很容易写错的小细节：不是 $\sqrt{d_h}$ ，也不是 $\sqrt{d_c}$ 。

解耦 RoPE 的代价也要说清楚：位置信息只由 64 维承载，其余 128 维（内容）对位置不敏感。这被称为 partial RoPE。在长上下文中，位置信号相对内容信号更稀疏，MLA 的位置敏感性天然弱于全维度 RoPE 的 MHA——DeepSeek 用 128K 的 NIAH（大海捞针）测试证明了这个折中在 128K 内是够用的（见”质量验证”一节），但”解耦的位置分量是不是在更长上下文里的瓶颈”至今仍是开放问题。

矩阵吸收：推理时的灵魂优化#

现在到了 MLA 最关键、也最容易被误解的工程优化：矩阵吸收（Matrix Absorption）。

先想一个问题：如果推理时每生成一个 token，都要先把潜变量 $c_t^{KV}$ 通过 $W^{UK}$ 上投影成完整的 128 头 K（ $512 \rightarrow 16{,}384$ 个元素），再和每个头算注意力，那么解码虽然省了显存，但每次都要多算一次巨大的上投影 GEMM，而且对历史 token 的 K 也要重新投影——省下的带宽又吐回去了。

吸收技巧的本质是：利用线性代数把上投影挪走，让注意力计算直接在 512 维潜空间里完成。

先看 K 侧。第 $i$ 个头的分数（暂时忽略 RoPE 部分）是：

q_{t,i}^{C,T} k_{j,i}^{C} = q_{t,i}^{C,T} (W^{UK}_{i} c_j^{KV}) = (W^{UK,T}_{i} q_{t,i}^{C})^{T} c_j^{KV}

其中 $W^{UK}_{i} \in \mathbb{R}^{d_h \times d_c}$ 是 $W^{UK}$ 中产生第 $i$ 头 key 的行段。定义吸收后的 query：

\tilde{q}_{t,i} = W^{UK,T}_{i} q_{t,i}^{C} \in \mathbb{R}^{d_c}

那么 $q^T k = \tilde{q}^T c^{KV}$ ——K 的上投影被吸收进了 query 侧。query 每步只需要算一次固定开销的 $W^{UK,T}$ 变换（ $16{,}384 \times 512$ 的矩阵乘），之后与历史 token 的交互全部发生在 512 维潜空间。对缓存里的每个历史 token，我们再也不需要它的完整 K，只需要 $c^{KV}$ （512 维）加 $k^R$ （64 维）。

再看 V 侧。注意力输出是：

o_{t,i} = \sum_j w_{j} v_{j,i}^{C} = \sum_j w_j (W^{UV}_{i} c_j^{KV}) = W^{UV}_{i} \left( \sum_j w_j c_j^{KV} \right) = W^{UV}_{i} o_{t,i}^{latent}

注意权重 $w_j$ 不依赖 $i$ 的 V 分支（V 没有 RoPE），所以可以先把分数加权求和作用在 512 维的 $c^{KV}$ 上，得到潜空间输出 $o^{latent}$ ，最后再一次性应用 $W^{UV}$ ：

u_t = W^{O} [o_{t,1}; \dots; o_{t,n_h}] = W^{O} W^{UV} o_t^{latent}

$W^{O}$ 与 $W^{UV}$ 都是线性变换，可以预先合并成一个大矩阵 $W^{O} W^{UV} \in \mathbb{R}^{d \times d_c}$ （一个 $5120 \times 512$ 的 GEMM）。论文原话是：“推理时 $W^{UK}$ 可以被吸收进 $W^{Q}$ ， $W^{UV}$ 可以被吸收进 $W^{O}$ ，我们甚至不需要把 K 和 V 完整算出来”——这里" $W^{Q}$ "指的就是 query 侧的变换，吸收后我们不再显式构造 $W^{UK}$ 和 $W^{UV}$ 的完整输出。

吸收之后的 decode 流程（单 token）可以概括为：

用 $W^{DQ}$ 、 $W^{UQ}$ 算 $q^{C}$ （128 维/头），用 $W^{QR}$ 算 $q^{R}$ （64 维/头），对 $q^{R}$ 做 RoPE；
用 $W^{UK,T}$ 把 $q^{C}$ 变换到潜空间： $\tilde{q} = W^{UK,T} q^{C}$ （每个头 $128 \rightarrow 512$ ）；
分数 $= \tilde{q}^{T} c^{KV} + q^{R,T} k^{R}$ ，softmax；
输出 $= W^{O} W^{UV} \left( \sum_j w_j c_j^{KV} \right)$ 。

全程只读取每 token 576 个元素（ $c^{KV}$ + $k^{R}$ ），注意力矩阵乘在 512 维潜空间进行。对比不吸收的朴素路径（官方代码里的 naive 模式）：它要把每个头的完整 K/V 都缓存下来，每 token 每层缓存 $n_h \times (192 + 128) = 40{,}960$ 个元素——只比 MHA 的 32{,}768 略少，压缩的意义荡然无存。所以 naive 模式只适合训练对照，推理必须走吸收模式。

这里有一个反直觉的事实值得讲清楚：吸收模式每步的注意力 FLOPs 其实比 naive 更多，省下的是 HBM 带宽，不是算力。naive 的注意力在每头 $192 + 128 = 320$ 维上进行，吸收模式则在 $576 + 512$ 维上进行——每 token 的潜变量 $c^{KV}$ 在潜空间里同时扮演 K 和 V 两个角色——每头 FLOPs 大约是 naive 的 $1088 / 320 \approx 3.4$ 倍。但两者从 HBM 读取的数据差了 71 倍（576 元素 vs 40{,}960 元素/token）。按 FlashMLA 博客的计算/访存比公式（下一节展开）：naive 模式比值约为 1 FLOP/byte，是重度 memory-bound；吸收模式约 242 FLOPs/byte，逼近 H800 的算力带宽比（约 258），恰好进入 compute-bound 区间。吸收的本质是用更多的浮点运算换掉两个数量级的显存搬运——对 decode 阶段这种带宽即生命的场景，这笔交易稳赚。这也是 FlashMLA 分析里”K 和 V 是同一个东西”的由来。

这套吸收逻辑在 DeepSeek-V3 官方推理代码（inference/model.py）里写得很直白。下面是吸收了关键注释的精简版：

1
# 注意：以下为官方代码的精简摘录，仅保留 MLA 核心路径（absorb 模式）
2
q = self.wq_b(self.q_norm(self.wq_a(x)))      # [b, s, n_h * 192]：wq_a 下投影 + wq_b 上投影
3
q_nope, q_pe = torch.split(q, [128, 64], dim=-1)   # 内容 128 维 + 位置 64 维
4
q_pe = apply_rotary_emb(q_pe, freqs_cis)      # 只旋转 64 维位置分量
5

6
kv = self.wkv_a(x)                            # 融合投影：[b, s, 512 + 64]，一次算 c^KV 和 k^R 的原始量
7
kv, k_pe = torch.split(kv, [512, 64], dim=-1)
8
k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)
9

10
wkv_b = self.wkv_b.weight.view(n_h, -1, 512)  # wkv_b: [n_h, 256, 512]（前 128 行是 W^UK，后 128 行是 W^UV）
11
q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :128])   # 吸收 W^UK：q 进入 512 维潜空间
12

13
self.kv_cache[..., :] = self.kv_norm(kv)      # 只缓存 512 维潜变量
14
self.pe_cache[..., :]  = k_pe.squeeze(2)      # 加 64 维 RoPE key
15

16
scores = (torch.einsum("bshc,btc->bsht", q_nope, self.kv_cache) +
17
          torch.einsum("bshr,btr->bsht", q_pe, self.pe_cache)) * self.softmax_scale
18
scores = scores.softmax(dim=-1, dtype=torch.float32)
19

20
x = torch.einsum("bsht,btc->bshc", scores, self.kv_cache)   # 潜空间加权和（V 侧吸收前）
21
x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -128:])      # 应用 W^UV，回到 128 维/头
22
x = self.wo(x.flatten(2))                                   # 最后做 W^O

这段代码值得逐行对照上面的推导：

wkv_a 是 Linear(dim, kv_lora_rank + qk_rope_head_dim)——一次投影同时产出潜变量和未旋转的 RoPE key，这是工程上把两个线性层融合成一个 GEMM 的常规优化；
kv_norm（RMSNorm）加在潜变量上，对应论文里”在压缩潜向量后加额外的 RMSNorm 保证训练稳定”的说明；q_norm 同理；
einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :128]) 就是 $\tilde{q}_{t,i} = W^{UK,T}_i q^{C}_{t,i}$ ，hdc 里 h 是头、d 是 128 维内容、c 是 512 维潜空间，正好是 $W^{UK}$ 的转置作用；
分数计算拆成两项：bshc,btc->bsht（潜空间内容分数）加 bshr,btr->bsht（位置分数），r 是 64 维 RoPE 维度；
最后 wkv_b[:, -128:] 是 $W^{UV}$ ，先把它作用在潜空间输出上，再乘 $W^{O}$ ——官方实现没有把 $W^{UV}$ 和 $W^{O}$ 预合并成一个大矩阵（这样省一次 GEMM 的权重搬运，代价是多一次小 GEMM），而 vLLM 等框架选择了预合并，两种做法数学上等价。

官方代码还保留了 attn_impl = "naive" 的对照路径：naive 模式缓存完整的 k_cache（每 token 每头 192 维）和 v_cache（128 维），即回到”缓存完整 K/V”的老路。attn_impl 默认是 "absorb"——官方部署用的就是吸收模式。这个开关的存在也说明：MLA 的推理路径不是唯一的，实现者要在”缓存小 vs 计算少”之间做选择，而吸收模式两边通吃，成为所有主流框架的默认实现。

另外注意代码里的 softmax_scale：DeepSeek-V3 在扩展上下文（YaRN 长度外推）时还会乘一个与 rope 缩放因子相关的修正系数 mscale，即 softmax_scale * mscale^2——这是 RoPE 长文本外推时一个容易踩的坑，官方代码里专门处理了。

训练侧细节：query 压缩、RMSNorm 与激活显存#

MLA 在训练时和推理时的关注点不同。推理的瓶颈是 KV Cache（常驻显存）；训练的瓶颈是激活（activation）显存——训练时每步要把整个 batch 的所有中间结果（q、k、v、注意力输出等）存下来做反向传播，batch 大、序列长时激活显存同样爆炸。

query 侧的低秩压缩（ $d_c' = 1536$ ）就是为训练准备的。 $W^{DQ}$ 把 $h_t$ 压到 1536 维，注意力计算时再通过 $W^{UQ}$ 展开成 $128 \times 192 = 24{,}576$ 维的完整 query。反向传播只需要对 1536 维的 $c_t^{Q}$ 求梯度，而不是对 24,576 维的完整 q 求梯度，激活显存大幅下降。论文没有单独给出”query 压缩省了多少”的数字，但 V2 相对 DeepSeek 67B 的训练成本节省 42.5%，query 压缩的贡献就在其中。

训练侧另外两个实现细节（论文 “Model Hyper-Parameters” 一节明确提到）：

压缩后加 RMSNorm：在 $c_t^{KV}$ 、 $c_t^{Q}$ 之后分别加 kv_norm、q_norm（即上面代码里的两个 RMSNorm）。原因：低秩压缩改变了层的输出尺度，直接接后续投影会导致数值不稳定，归一化把尺度拉回来；
宽度瓶颈处乘缩放因子：在压缩潜向量和 MoE 路由专家的中间态等”宽度瓶颈”位置，乘上额外的缩放系数，目的同样是稳定训练。这两个都是”让低秩结构训得动”的工程兜底，论文没有给出具体系数。

还有一个对称性的细节：训练时也可以做吸收。论文原话是”为了减少训练时的激活内存，我们也把 $W^{UK}$ 吸收进 $W^{UQ}$ ，把 $W^{UV}$ 吸收进 $W^{O}$ “——即训练前向里同样不显式构造完整的 K/V，直接在潜空间算注意力，从而连 q、k、v 的完整激活都不用落盘。所以”吸收”不是推理专属的技巧，训练和推理共用同一套数学。

质量验证：MLA 的长上下文能力#

压缩 57 倍之后能力会不会掉？除了附录消融（上面 Table 7 的数据），DeepSeek-V2 还专门做了 128K 的”大海捞针”（Needle In A Haystack）测试——把一句话随机藏进 128K 文档的不同深度，要求模型找到它。这个任务对位置编码的敏感度极高，正好检验 partial RoPE 是否够用。

NIAH 测试：横轴是文档长度（最长 128K），纵轴是”针”埋藏深度，颜色越深表示召回越准。DeepSeek-V2 在全部长度与深度组合上保持深色（高召回）。（来源：DeepSeek-V2 论文 Figure 4）

结果是全矩阵深色——128K 范围内 MLA 的召回几乎不受影响。这给”partial RoPE 够不够用”提供了实证答案：至少在 128K 内，64 维位置分量足以支撑精确指代。DeepSeek-V3 继续把 MLA 用于 128K 上下文（训练时扩展到 128K 再加 YaRN 外推），V3 的多语言与代码评测全面优于同规模模型，说明 MLA 的压缩在更大模型上依然成立。

FlashMLA：MLA 的高性能 Kernel 实现#

理解了 MLA 的数学之后，再看它的工程落地。2025 年 3 月 DeepSeek 开源周期间，DeepSeek 开源了 FlashMLA——为 Hopper GPU 优化的 MLA 解码 kernel 库（GitHub 仓库，论文引用为 FlashMLA: Efficient Multi-head Latent Attention Kernels，2025），它是 DeepSeek-V3 与 V3.2 线上推理实际使用的注意力内核。到今天它已经历了三次大的迭代：

版本	时间	内容	H800 SXM5 性能
v0.1	2025-03-01	Dense MLA 解码 kernel（BF16/FP16）	3000 GB/s（memory-bound）/ 580 TFLOPS（compute-bound）
更新版	2025-04-22	新调度 kernel（seesaw）	3000 GB/s / 660 TFLOPS
稀疏版	2025-09-29	V3.2 DSA 稀疏注意力 + FP8 KV	解码 410 TFLOPS / 预填充 640 TFLOPS（B200 预填充 1450 TFLOPS）

（数据来源：FlashMLA README 与新 kernel 深挖博客）

为什么 MLA 解码 kernel 是 compute-bound#

FlashMLA 深挖博客里第一个让人意外的结论是：MLA 解码 kernel 是计算密集（compute-bound）而不是带宽密集（memory-bound）的。解码阶段每个请求每步只处理 1 个 query token，直觉上应该是”读历史 KV 的带宽”主导。推导如下。

设每请求 query 头数 $h_q$ 、每请求 query token 数 $s_q$ （不开 MTP/投机解码时为 1）、历史 KV token 数 $s_k$ （ $s_k \gg h_q s_q$ ），K 和 V 的每头维度分别是 $d_k$ 、 $d_v$ 。一个解码 step 的 FLOPs 约为：

\text{FLOPs} \approx 2 h_q s_q s_k (d_k + d_v)

（一次 $q \times K^T$ 加一次 $P \times V$ 的矩阵乘）。读入的数据（bf16）约为：

\text{bytes} \approx 2 \times (h_q s_q d_k + s_k d_k + h_q s_q d_v) \approx 2 s_k d_k

（历史 KV 占主导）。两者相除，计算/访存比：

\frac{\text{FLOPs}}{\text{bytes}} \approx \frac{2 h_q s_q s_k (d_k + d_v)}{2 s_k d_k} \approx h_q s_q \cdot \frac{d_k + d_v}{d_k} \approx 2 h_q s_q

H800 SXM5 的峰值算力 990 TFLOPS、带宽 3.35 TB/s，理论上当 $2 h_q s_q \geq 990 / 3.35 \approx 296$ 时计算才会成为瓶颈——但实际降频（约 1600 MHz）后有效算力约 865 TFLOPS，所以阈值是 $865 / 3.35 \approx 258$ ，即 $h_q s_q \geq 129$ 就进入 compute-bound。DeepSeek 的解码实例不用张量并行， $h_q = 128$ （所有 128 个 query 头都在一张卡上）， $128 \times 1 \geq 129$ 差一点……不过博客给出的结论是 $h_q s_q \ge \tfrac{1}{2} \cdot \tfrac{865}{3.35} \approx 129$ ，实际 128 恰好在这个临界点上，所以”kernel 处于 compute-bound 配置”。

这里 $d_k = 576$ （512 潜空间 + 64 RoPE）、 $d_v = 512$ （潜空间），正是 FlashMLA README 里说的 “MQA mode: head_dim_k = 576, head_dim_v = 512”——FlashMLA 把 MLA 当成一个 K 维度 576、V 维度 512 的 MQA 来算，因为吸收之后所有 query 头共享同一个潜空间 V。这也是为什么 README 的术语表里 MLA 有两种模式：解码用 MQA 模式（576/512），预填充用 MHA 模式（head_dim 192/128）——预填充阶段每 token 要算完整的多头注意力、且没有”历史 KV 带宽”压力，直接用完整维度反而更好。

既然 compute-bound，优化目标就从”压带宽”变成”让 Tensor Core 尽量不停”。

Seesaw 调度：单输出矩阵的 ping-pong#

FlashMLA 新 kernel 的核心调度技巧是博客里称为 seesaw（跷跷板） 的调度，这是对 FlashAttention-3 ping-pong 调度的改编。背景问题：FlashAttention 的在线 softmax 需要在寄存器里维护一个输出累加矩阵，而 MLA 的输出矩阵是 $64 \times 512$ ，要占 $64 \times 512 = 32{,}768$ 个 32 位寄存器；一个 SM 只有 65,536 个寄存器，只够放一个输出矩阵。FlashAttention-3 的 ping-pong（双缓冲两个输出矩阵、一个给 CUDA Core 用、一个给 Tensor Core 用）在这里直接失效——没有第二个矩阵的寄存器空间。

seesaw 的办法是把输出矩阵竖切成 $O_L$ 、 $O_R$ 两半（各 $64 \times 256$ ），两个 warpgroup 各持一半，每次取两个 KV 块（ $K_0$ 、 $K_1$ 与 $V_0$ 、 $V_1$ ）交错计算：warpgroup 0 算 $O_L$ 的一半更新，warpgroup 1 算 $O_R$ 的一半更新，通过精心安排的缩放因子（ $scale_0$ 、 $scale_1$ ）保持与标准在线 softmax 完全等价。同时，每个 KV 块的数据在不再需要后立即发起 TMA（Tensor Memory Accelerator）预取下一块，把访存和计算重叠起来。

FlashMLA 的 seesaw 调度：两个 warpgroup 轮流处理两个 KV 块，输出矩阵竖切为 O_L/O_R 各持一半，CUDA Core 的 softmax 操作与 Tensor Core 的 GEMM 交错执行，数据用完即发起 TMA 预取。（来源：FlashMLA 新 kernel 深挖博客）

配合的其他优化：细粒度 TMA 流水（一个 $64 \times 576$ 的 K 块拆成 9 个 $64 \times 64$ 的小块分别搬运，第一块搬完就开始 GEMM，不必等整块）、缓存提示 EVICT_FIRST（提高 L2 命中率）、programmatic dependent launch（把 split-k 的 splitkv_mla 与合并 kernel 重叠）、以及tile scheduler（把”请求 × KV 块”的任务均匀分给各 SM，解决变长序列的负载不均）。博客报告这些优化把 Tensor Core 利用率推到理论峰值的 80%。

FlashMLA 的 API 设计也值得一看，它把”分页 KV Cache”和”变长序列”直接暴露给用户：

1
from flash_mla import get_mla_metadata, flash_mla_with_kvcache
2

3
# 解码循环前先算一次调度元数据
4
tile_scheduler_metadata, num_splits = get_mla_metadata(
5
    cache_seqlens,          # 每条序列的缓存长度（变长）
6
    s_q * h_q // h_kv,      # 每个 KV 头对应的 query 数
7
    h_kv, h_q, is_fp8, topk,
8
)
9

10
for i in range(num_layers):
11
    o_i, lse_i = flash_mla_with_kvcache(
12
        q_i, kvcache_i, block_table, cache_seqlens, dv,
13
        tile_scheduler_metadata, num_splits,
14
        is_causal, is_fp8_kvcache, indices,
15
    )

kvcache_i 是分页的潜空间 KV 缓存，页大小（block size）为 64——即每页存 64 个 token 的 $c^{KV}$ + $k^{R}$ 。分页让缓存可以按请求动态分配/回收，避免预留整条序列的连续内存（这是 vLLM 的 PagedAttention 论文 arXiv:2309.06180 提出的思路在 MLA 上的延续）；
s_q 是每请求的 query token 数：不开 MTP 时为 1，开 MTP（DeepSeek-V3 的多 token 预测）时可以大于 1——FlashMLA 天然支持这种”一个请求多 query token”的解码形态；
返回 lse（log-sum-exp），方便上层实现投机解码或拆分合并的注意力。

2025 年 9 月的稀疏版还引入了一种量化 KV Cache 格式：DeepSeek-V3.2 的 DSA 稀疏注意力用 FP8 存潜变量。每 token 的 KV 固定 656 字节：前 512 字节是量化后的 512 维潜变量（FP8 e4m3，每 128 个值配一个 FP32 scale，共 4 个 scale 占 16 字节），最后 128 字节是不量化的 64 维 RoPE key（BF16）——位置分量对精度敏感，刻意保留高精度。kernel 读进来后先反量化成 BF16 再做注意力。这个”FP8 潜变量 + 全精度位置分量”的格式是 MLA 量化实践中最重要的经验之一。

主流框架中的 MLA 实现#

MLA 自 2024 年 5 月发布以来，所有主流推理框架都已适配。它们在数学上完全一致——缓存 576 维（或更少）+ 吸收上投影——差异在工程细节：

vLLM：DeepseekV2Attention 后端按吸收模式实现，每 token 缓存 kv_lora_rank + qk_rope_head_dim = 576 维到分页内存中；query 侧先算 $q^{C}$ 与 $q^{R}$ ，把 $W^{UK}$ 吸收进 query、把 $W^{UV}$ 与 $W^{O}$ 融合成单个线性层再作用在潜空间输出上。较新的版本还支持基于 FlashAttention-3 的 MLA 实现与级联（cascade）MLA 缓存；
SGLang：通过 FlashInfer 的 MLA wrapper 调用，同样是潜空间注意力 + 576 维缓存，页大小 64；
TensorRT-LLM：把吸收后的 MLA 作为专用 attention 节点实现，权重融合在权重转换阶段完成；
llama.cpp：支持 MLA 解码，并可对其 KV 缓存做分块量化（i-quants）压缩潜变量，配合 CPU/GPU 混合推理；
Hugging Face Transformers：2024 年 6 月即合并了 MLA 支持（DeepSeek-V2 的 DeepseekV2ForCausalLM），用标准注意力 + 缓存 576 维的等效写法。

一个统一的事实：所有实现的性能都取决于同一个东西——每 token 读取的 576 维数据能否喂饱带宽。这也解释了为什么 FlashMLA 的 3000 GB/s 数字（接近 H800 的 3.35 TB/s 峰值）是整个 MLA 生态的性能标杆。

局限与未解决的问题#

MLA 不是没有争议和代价。把它讲透，也要把它没解决的问题讲清楚：

1. 维度选择缺乏消融支撑。 $d_c = 4 d_h = 512$ 、 $d_c' = \tfrac{3}{2} d_h = 1536$ 、 $d_h^{R} = \tfrac{1}{2} d_h = 64$ 这三个关键数字在论文里没有系统消融——没有扫描过 $d_c \in \{2, 4, 8, 16\} \times d_h$ 的 Pareto 曲线，也没有论证为什么 query 压缩用 1536 而 KV 用 512。它们是”工程上 work 的组合”，不是”推出来的最优”。 $d_c$ 的最优值应当依赖模型规模、上下文长度、训练数据，但目前没有公开数据回答” $d_c$ 与 $n_h$ 的 trade-off 拐点在哪”。

2. Partial RoPE 的位置敏感性问题。 位置信息只占 64/192 的维度。NIAH 证明了 128K 内够用，但更长上下文或对位置极敏感的任务（精确的坐标指代、长程共指）上，MLA 相对全维度 RoPE 的差距没有被量化过。V3.2 转向稀疏注意力后，位置子空间的选择又和稀疏掩码的构造耦合，问题更复杂了。

3. 短上下文下收益有限。 KV 压缩主要解放的是长上下文和高并发下的显存与带宽。对短 prompt、低并发的场景，MLA 的每 token 固定开销（两次吸收变换）反而略多于 MHA，且低秩路径引入额外的投影层，模型参数里有 60 层 × 若干投影矩阵的开销。DeepSeek 报告 5.76 倍吞吐是在长上下文高并发生产配置下测的，短场景要打折扣。

4. 低秩是隐式假设。 压缩能成立的前提是”K/V 的信息集中在 $d_c$ 维子空间”。这个假设在 128 头、16384 维的配置下被验证有效，但它不是免费的——附录消融里小模型 C-Eval 的 -0.7 分就是代价的痕迹。如果未来模型需要更细粒度的头级信息（比如某些头专门编码特定位置模式），低秩压缩的墙会先被撞到。

5. 生态里所有优化都在”潜空间”这条路上越走越远。 吸收技巧让注意力完全绕开完整 K/V，但也意味着未来任何想”看完整 K/V”的优化（例如基于完整向量的 KV 驱逐/量化）都无法直接在 MLA 上做——只能在 576 维潜变量上做。QEvict、MAC-Attention 等 KV 管理方法在 MLA 模型上全部要重新适配，这是 MLA 流行之后社区一直在补的功课。

小结#

把 MLA 的完整图景收束成一张表：

设计	解决的问题	手段	代价
低秩 KV 联合压缩	KV Cache 太大	下投影到 512 维潜变量，缓存潜变量而非完整 K/V	隐含低秩假设，压缩有少量精度损失
Query 低秩压缩	训练激活显存大	$c^{Q}$ 1536 维 + 上投影	额外投影层参数
解耦 RoPE	RoPE 与低秩压缩不兼容	内容/位置分离，位置只走 64 维共享 key	partial RoPE 位置敏感度下降
矩阵吸收	避免每步重算上投影	$W^{UK}$ 吸收进 query、 $W^{UV}$ 吸收进 $W^{O}$ ，注意力全程在潜空间	吸收变换的固定开销
FlashMLA	潜空间注意力的算力瓶颈	seesaw 调度、TMA 细粒度流水、分页 64、FP8 潜变量	实现复杂度高

MLA 给我们的启示可以概括为一句：KV Cache 优化不只有”共享头”一条路，压缩表示的秩是一条更通用的路。MQA/GQA 省缓存靠”砍头”，MLA 省缓存靠”压缩”，两者可以组合（MLA 的 $n_h$ 和 $d_c$ 本来就是独立变量）。沿着这条路，DeepSeek-V3.2 的 DSA 稀疏注意力在 MLA 的 576 维缓存之上再做 token 级稀疏，Kimi K3 等后续模型也推出了自己的混合注意力变体——MLA 打开的”低秩潜空间”思路，至今仍是长上下文推理优化最重要的技术路线之一。

音乐

音乐