[Tech] 为什么 MLA + MTP 容易吃亏？从 Arithmetic Intensity 看 Attention

前几天在看苏剑林老师的这篇文章：https://kexue.fm/archives/11848，里面有一句话让我停了一下：

但现在除 KV Cache 外，Decoding 还有一个新的变数——MTP，或者说推测解码，其思想是计算换速度。然而 MLA 在 Decoding 时表现为 head_dims=512+ 的 MQA，已经提前消耗了大部分算力，所以“MLA+MTP”很容易吃亏。

我第一反应是：这是什么意思？为什么 MLA 在 decode 时会“提前消耗算力”？为什么它在 decode 时的 FLOPs，相当于一个 head dim 512+ 的 MHA，反而和 MTP 有冲突？

顺着这个问题往下看，最后发现背后其实是一个非常经典、但放在 Attention 上又很有意思的概念：Arithmetic Intensity（AI，计算密度）。

而且最后得到的结论意外地简单（下面都按 BF16 的 KV cache 算）：

把 MHA 一路约下来，AI 居然刚好就是 1；
GQA、MQA 也一样干净，只和 head 的数量有关 —— context length、head dim 全部约掉了，一个都不剩；
MLA 还是同一个形状，同样和 latent dim 无关，只是多了一个不到 2 的常数。

把四种结构排在一起，single-token decode 下 attention core 的 AI 是这样的：

Attention	历史 cache 里存什么	这时的 AI 大约是
MHA	每个 Query head 各有 K、V	1
GQA	一组 Query head 共用 K、V	Query head 数 / KV head 数
MQA	所有 Query head 共用 K、V	Query head 数
MLA	只存一份 latent，K、V 都从它展开	约 2 × Query head 数

第 2 节和第 4 节会把这四行分别推一遍。前三行其实是同一个公式的三个 special case，只是 KV head 数取值不同；最后一行多出来的那个常数，来源完全不同。

而这个不到 2 的常数，恰好足以把很多现代 GPU 上的 Attention decode 从明显 memory-bound 推到 Roofline knee 附近。这样再叠 MTP，就很容易把 workload 推进 compute-bound —— 所以苏神才说 MLA 对 MTP 不友好。

这篇文章把这个推导完整写一遍。

这篇会把推导写得比较细 —— 从「矩阵乘法怎么数 FLOPs」讲起，每个矩阵的形状都摆出来。如果你对 attention 的结构和 decode 时的计算已经很熟，第 2 节可以只看 2.5 的结论，直接跳到第 3 节。

1. 什么是 Arithmetic Intensity

Arithmetic Intensity 定义很简单：

\boxed{ AI = \frac{\text{FLOPs}}{\text{HBM Bytes}} }

也就是：每从 HBM 搬 1 Byte 数据进来，做了多少浮点运算。

一个 workload 的 AI 很低，意味着 GPU 大部分时间在搬数据；AI 很高，意味着拿到一份数据后会反复做很多计算。

硬件本身也有一个对应的临界值：

\boxed{ AI_{\text{hardware}} = \frac{\text{Peak Compute}}{\text{HBM Bandwidth}} }

在最理想的 Roofline 模型下：

AI_{\text{workload}} < AI_{\text{hardware}} \quad\Rightarrow\quad \text{memory-bound}

而

AI_{\text{workload}} > AI_{\text{hardware}} \quad\Rightarrow\quad \text{compute-bound}

下面用 dense BF16 Tensor Core throughput 来粗略算几张常见卡的 theoretical balance point。

GPU	Dense BF16 Peak	HBM Bandwidth	Theoretical Balance Point
H100 SXM	~989.5 TFLOP/s	3.35 TB/s	~295 FLOP/B
H200 SXM	~989.5 TFLOP/s	4.8 TB/s	~206 FLOP/B
B200 (HGX)	~2.25 PFLOP/s	~8 TB/s	~281 FLOP/B

来源：NVIDIA H100、NVIDIA H200、NVIDIA HGX B200、NVIDIA DGX B200。

当然，真实 kernel 不可能同时完美打满理论 FLOPs 和 HBM 带宽，因此这些只是一个用于建立直觉的 Roofline 上界，而不是 profiler 中的真实分界线。记住「几百 FLOP/B」这个量级就够，后面要拿它对照。

这里要先提醒一句：AI 是一个比值，它只回答「你卡在算力还是带宽这一侧」，不回答「哪种做法更快」。 分子分母一起涨，AI 可以纹丝不动，而两边都变慢了。第 5 节会给一个具体例子：两种算法算同一件事，AI 高的那个反而多花了 120 倍 FLOPs。

本文接下来只算 attention 里和 KV 有关的那一段：

decode 一个 token 时，从它的 hidden state 算出 Q、K、V，读进历史 KV cache，一路算到这一层的 attention output。

softmax 的算术量相对两个大矩阵乘小得多，先不算进去。这里也不打算估整个 Transformer layer 的 latency，只想把一个问题单独拎出来看：换一种 Attention 结构，上面这段的计算密度会变成什么样。

第 2 到第 4 节先只算 decode（一次生成一个 token，历史从 cache 读）。prefill（一次把整段序列算完）留到第 5 节 —— 到那时会看到，同一个模型在两个阶段该用的算法正好相反。

PS：为什么不把最后那个 W_O 也算进去？因为它和「用哪种 attention 结构」无关。W_O 拿到的永远是各个 head 拼好的那个向量，宽度只由 H_q 和 d_v 决定，KV 那边怎么组织它都看不见 —— MHA、GQA、MQA、MLA 喂给它的东西一样宽。而且它和 Q/K/V projection 一样是权重乘向量：不碰 KV cache、不随 L 增长，batch = 1 时 AI 恒等于 2/b。把它算进来，只会给每种结构加上同一个常数项，反而冲淡要比较的东西。MLP 和通信同理。

2. MHA / GQA / MQA 的统一公式

2.1 FLOPs 怎么数

(m × k) 的矩阵乘 (k × n) 的矩阵，要做 m·n·k 次乘加（MAC）。一次乘加是一次乘法加一次加法，算 2 FLOPs，所以：

\text{FLOPs} = 2mnk

2.2 MHA 简介

最基础的 attention 就是 MHA（Multi-Head Attention）：把 hidden state 切成若干个 head，每个 head 各自算自己的 Query、Key、Value，各自做一遍 attention，最后把结果拼回去。

先把形状摆出来。这一层拿到的是当前 token 的 hidden state：

h \in \mathbb{R}^{1 \times d_{\text{model}}}

三个 projection 矩阵把它变成 Query、Key、Value。MHA 里三者的 head 数是一样的，但后面 GQA、MQA 会让 Key/Value 的 head 数变少，所以这里就给它单独一个符号：Query head 数记作 H_q，KV head 数记作 H_kv，MHA 就是 H_kv = H_q 的情形。

这里还要顺一件事：Key head 数和 Value head 数，数学上并不要求相等 —— 让 K 分 8 组、V 分 4 组，式子照样成立。但真实模型里没人这么设计：K 和 V 在 cache 里总是成对存的，存了某个位置的 k，就一定同时存了它的 v。本文顺着这个惯例，认为两者相等，统一记作 H_kv。

至于 head dim：Query 和 Key 的必须相同，否则 q 和 k 没法做点积，记作 d_k；Value 的可以不同，记作 d_v。

W^Q \in \mathbb{R}^{d_{\text{model}} \times H_q d_k}, \quad W^K \in \mathbb{R}^{d_{\text{model}} \times H_{kv} d_k}, \quad W^V \in \mathbb{R}^{d_{\text{model}} \times H_{kv} d_v}

单个 head 上：

q_h \in \mathbb{R}^{1\times d_k}, \qquad k_h \in \mathbb{R}^{1\times d_k}, \qquad v_h \in \mathbb{R}^{1\times d_v}

新算出来的 k、v 追加进 KV cache。设历史长度为 L，那么这一层的 cache 是：

K \in \mathbb{R}^{L \times H_{kv} \times d_k}, \qquad V \in \mathbb{R}^{L \times H_{kv} \times d_v}

然后每个 Query head 做三步。记 g(h) 为第 h 个 Query head 用的那个 KV head（MHA 里就是 g(h) = h）：

一、打分。 当前 query 和全部 L 个历史 key 做点积，形状是 (1 × d_k) 乘 (d_k × L)：

s_h = q_h K_{g(h)}^\top \in \mathbb{R}^{1\times L}

二、归一化。 softmax 把分数变成权重，形状不变：

p_h = \operatorname{softmax}(s_h) \in \mathbb{R}^{1\times L}

三、加权求和。 拿权重去乘 Value，形状是 (1 × L) 乘 (L × d_v)：

o_h = p_h V_{g(h)} \in \mathbb{R}^{1\times d_v}

H_q 个 head 各自算出一个 o_h，拼起来就是这一层的 attention output。

2.3 FLOPs

三部分，每一部分都是往 2mnk 里代形状。

一、算出当前 token 的 Q、K、V。 三次 (1 × d_model) 乘 (d_model × ·)，代进去 m = 1：

F_{\text{proj}} = 2d_{\text{model}}\left(H_q d_k + H_{kv} d_k + H_{kv} d_v\right)

二、打分。 每个 Query head 一次 (1 × d_k) 乘 (d_k × L)，共 H_q 个：

F_{QK} = 2H_q L d_k

三、加权求和。 每个 Query head 一次 (1 × L) 乘 (L × d_v)：

F_{PV} = 2H_q L d_v

三部分加起来，就是这一层 attention 的全部计算量：

F = F_{\text{proj}} + 2H_q L (d_k + d_v)

这两项的性质完全不同：F_proj 和历史长度 L 无关，它只处理当前这一个 token；后一项正比于 L，它要把整段历史扫一遍。两者相比，量级上是：

\frac{F_{\text{proj}}}{2H_q L (d_k+d_v)} \;\sim\; \frac{d_{\text{model}}}{L}

（代进 MHA 的 H_kv = H_q、d_k = d_v，正好是 1.5 d_model/L。）

这个比值没有想象中小。d_model = 7168 时它是 10752/L：L = 8K 时 F_proj 还比 F_attn 大（1.3 倍），32K 时是 33%，要到十几万 token 才降到 10% 以下。

所以略去 F_proj，靠的并不是「它小」这个理由。真正起作用的是：它和我们要比的东西无关 —— F_proj 是权重乘向量，读的是权重不是 KV cache，AI 恒等于 2/b，换哪种 Attention 结构它都在那儿。而 F_attn 是唯一随 L 增长、也是唯一被 Attention 结构改变的那一项。下面算 AI 只留它：

\boxed{ F_{\text{attn}} = 2H_q L (d_k + d_v) }

话虽如此，F_proj 这条路径本身别忘掉。第 3 节会看到，MLA 做的事情恰恰是把本来要乘上 L 的计算，搬回到这条与 L 无关的路径上。

2.4 HBM Bytes

设每个 element 占 b Bytes，BF16 时 b = 2。

历史 cache 里有 H_kv 份 Key 和 H_kv 份 Value，每个位置的 Key 宽 d_k、Value 宽 d_v，一共 L 个位置：

\boxed{ B_{KV} = bLH_{kv}(d_k + d_v) }

这里假设 fused attention 不把 L 长度的 score / probability 矩阵写回 HBM。

三个 projection 权重当然也要从 HBM 读。但它们读的是权重、不是 KV cache，而且 batch = 1 时矩阵乘向量的 AI 恒等于 2/b，和选哪种 attention 结构无关，所以不影响下面几种结构之间的横向对比。

2.5 两式相除

只看随 L 增长的那部分：

AI = \frac{2H_qL(d_k+d_v)}{bLH_{kv}(d_k+d_v)}

L 约掉，(d_k + d_v) 整块约掉：

\boxed{ AI_{\text{MHA/GQA/MQA}} = \frac{2}{b}\frac{H_q}{H_{kv}} }

BF16 下 b = 2：

\boxed{ AI_{\text{BF16}} = \frac{H_q}{H_{kv}} }

漂亮的地方在这儿：context length L、head dim d_k 和 d_v 全部约掉了。

剩下的只有一个量：

一个 KV head 被多少个 Query head 复用

2.6 三个 special case

MHA、GQA、MQA 的差别，只在 H_kv 取什么值：

结构	H_kv 取值	含义	BF16 AI
MHA	H_kv = H_q	每个 Query head 各有一份	1
GQA	1 < H_kv < H_q	一组 Query head 共用	H_q / H_kv
MQA	H_kv = 1	全部 Query head 共用	H_q

所以 MHA → GQA → MQA 就是同一个公式里滑动 H_kv：FLOPs 没有因为 KV head 变少而下降，但要从 HBM 读的历史 KV 变少了，因为一份 KV 被更多 Query head 重用。

这就是第一层 data reuse。

2.7 代入真实模型

例如：

Qwen2.5-7B：H_q = 28, H_kv = 4；
Mixtral-8x7B：H_q = 32, H_kv = 8；
Falcon-7B：71 个 Query heads，且 multi_query=true，即 true MQA。

那么 BF16 single-token decode 下，这部分的 AI 分别大约是：

Model	Attention	H_q	H_kv	BF16 AI
Qwen2.5-7B	GQA	28	4	7 FLOP/B
Mixtral-8x7B	GQA	32	8	4 FLOP/B
Falcon-7B	MQA	71	1	71 FLOP/B

到这里，MQA 已经把跨 head 的 KV reuse 做到头了：所有 Query head 共用一份 K 和一份 V，AI 等于 Query head 数。

问题也就卡在这儿：MQA 的 AI 上限就是 Query head 数，而这个数长不大。 常见模型也就 32、64、128 个 Query head，Falcon-7B 的 71 已经算多的了，而且 Query head 数是模型结构定死的，不能为了 AI 随便加。对照第 1 节那几张卡两三百 FLOP/B 的 balance point，光靠堆 Query head 数根本够不到。

而且注意最后这几个字：

一份 K，和一份 V。

K 和 V 仍然是两份不同的数据。

这正是 MLA 下一步要动的地方。

2.8 那 prefill 呢？

上面整节算的都是 decode。同样的式子代进 prefill 会怎样？

prefill 一次要把整段输入算完。设输入序列长度是 L（和上面同一个符号 —— decode 时它是历史长度，prefill 时就是这段输入本身），那么每个 token 都要跟它前面所有 token 做 attention，于是每一份 K、V 会被大约 L 个 query token 用到 —— 复用是白送的。同样一除，得到：

AI_{\text{prefill}} \approx \frac{H_q}{H_{kv}}\cdot\frac{L}{b}

也就是decode 的 AI 再乘上 L/b。代进 H100 那条 295 FLOP/B 的线：

	decode AI	prefill 里越过 295 需要的 L
MHA（H_q = H_kv）	1	590
GQA（32 / 8）	4	148
MQA（32 / 1）	32	18

prefill 里连最朴素的 MHA 都是 compute-bound —— 输入超过六百来个 token 就过线了，GQA/MQA 几十个 token 就过。所以 prefill 阶段根本没有「memory-bound」这个问题可谈。

还有一点值得注意：GQA 和 MQA 完全不改变 prefill 的 FLOPs。 因为 FLOPs 只和 H_q 有关，跟 H_kv 无关。它们只是把一个本来就已经过线的 AI 抬得更高，外加把 cache 变小。

所以：

MHA → GQA → MQA → MLA 这整条线，从头到尾是一个 decode 的故事。 prefill 里复用本来就免费，没有东西需要修；这些结构会被发明出来，正是因为 decode 只有一个 query token、复用为零。

3. MLA 的 latent cache

编者注：为了先把最核心的结构说清楚，下面几节都忽略 RoPE。它不会改变主结论，补回来的那一步放在文末的附录 A。

MLA 的关键不是改变 Attention 最终的数学形式：

\operatorname{softmax}(QK^\top)V

而是改变历史 K/V 的 parameterization。

对于历史 token j，先把 hidden state 压缩到一个所有 head 共用的 KV latent。这里的 D 是 down-projection（降维）的意思，W^DKV 就是「把 hidden state 降维成 KV latent」的那个矩阵，它把 d_model 维压到窄得多的 d_c 维：

\boxed{ c_j=W^{DKV}h_j }

其中：

c_j\in\mathbb{R}^{d_c}

然后第 h 个 head 的 K/V 都从这一份 latent 展开：

k_{j,h}=W_h^Kc_j

v_{j,h}=W_h^Vc_j

于是过去原本需要 cache 的

K_j, V_j

现在可以只 cache：

\boxed{c_j}

拿 Kimi K3 代一下：96 个 head，K 和 V 的 head dim 都是 128，所以一份 KV 是 128 + 128 = 256 个数。每个 token、每层要存多少，就看存几份：

结构	每 token 每层存的数	32K context、93 层、BF16
MHA（96 份 KV）	256 × 96	139.5 GB
GQA（8 份 KV）	256 × 8	11.6 GB
MQA（1 份 KV）	256 × 1	1.45 GB
MLA（一份 latent）	512	2.91 GB

比 MHA 小 48 倍，比 8 组的 GQA 小 4 倍 —— 但比 MQA 还大一倍。

所以 MLA 的卖点并不是「cache 最小」，纯比 cache 大小它输给 MQA。它换来的是别的东西：这一份 latent 同时承担了 K 和 V 两边的角色，而 MQA 那 256 里，128 是 K、128 是 V，各干各的。这一点值多少，第 4 节来算。

但这马上带来一个问题：如果每生成一个新 token，都对历史 L 个 latent 重新做一遍 W^K 和 W^V 展开，那只是把 HBM 问题换成了巨大的计算问题。

MLA decode 真正精妙的地方是 matrix absorption / reassociation。

3.1 K-side absorption

原来的 content score：

q_h^Tk_{j,h}

代入：

q_h^TW_h^Kc_j

根据矩阵乘法结合律，可以改写成：

\boxed{ q_h^TW_h^Kc_j = \left((W_h^K)^Tq_h\right)^Tc_j }

也就是说，不需要对 L 个历史 latent 逐个生成 K。

只需要先对当前这一个 Query token做一次变换，把它拉进 latent 空间 —— 得到的 q̃_h 可以理解成「latent 空间里的 query」：

\tilde q_h=(W_h^K)^Tq_h

然后直接在 latent space 里和全部历史做点积：

\tilde q_h^Tc_j

方向正好反过来了：原来是把每个历史 latent 展开成 K，再拿 query 去和它比；现在是把 query 一次性拉进 latent 空间，直接和 cache 里存着的 latent 比。算出来的分数完全一样，但后者根本不用碰历史。

3.2 V-side absorption

同样，原来的 Attention output：

o_h = \sum_jp_{h,j}v_{j,h}

代入 v_j,h = W_h^Vc_j：

o_h = \sum_jp_{h,j}W_h^Vc_j

把固定的矩阵提到外面：

\boxed{ o_h = W_h^V \left( \sum_jp_{h,j}c_j \right) }

所以 PV 也可以直接在 latent space 中完成：

P_hC

最后再对得到的一个 latent output做一次 W_h^V projection。

3.3 为什么 projection 没有变贵

这是我一开始最容易误解的地方。

如果不做 absorption，那么每生成一个 token，都得把历史里的每一个 c_j 重新展开一遍 —— 对每个 head 算出 k_j,h = W_h^Kc_j 和 v_j,h = W_h^Vc_j。历史有多长就要算多少遍，成本乘 L：

O(Ld_cd_kH)

但 absorption 之后，额外的 K-side / V-side projection 都只作用在当前 query 或最终 output 上：

O(Hd_cd_k)+O(Hd_cd_v)

它们不再乘 context length L。

真正随历史长度增长的部分变成：

O(HLd_c)

因此长 context 下，当前 token 上这几次 projection，相对读取整段历史 KV 的开销会越来越小。

换句话说：

MLA 并没有让 projection 消失。它改的是矩阵乘法的结合顺序，把这笔 expensive projection 从历史序列的 O(L) 维度上挪走，只留在每个 decode step 的 O(1) current-token path 上。

结果完全一样，计算图的代价却完全不一样。

4. MLA 的 AI 为什么翻倍

按第 3 节说的，历史每个 token 只存一份 latent。

MLA 里所有 Query head 共用同一份 latent，没有「KV head 数」这个维度了，所以下面只需要一个 head 数，记作 H，它就是前面的 H_q。

历史每个 token cache：

c_j\in\mathbb{R}^{d_c}

4.1 FLOPs

一、和历史长度无关的那几笔。 这一个新 token 要：把 hidden state 压成 latent（2d_modeld_c）、算出自己的 query（2d_modelHd_k）、把 W^K 折进 query（3.1 那步 absorption，2Hd_kd_c）、最后把 latent 输出投影回每个 head（2Hd_cd_v）。这几笔都只作用在当前这一个 token 上，记作 F_fixed。

二、要扫历史的那两笔。 absorbed 之后，打分和加权求和都在 latent 上做，宽度都是 d_c：

F_{QK} = 2HLd_c, \qquad F_{PV} = 2HLd_c

合起来：

F = F_{\text{fixed}} + 4HLd_c

代进 DeepSeek-V3（d_model = 7168、H = 128、d_c = 512、d_k = d_v = 128），F_fixed ≈ 276 MFLOP，而后一项是 0.26 MFLOP × L，比值约 1052/L —— L = 8K 时是 13%，32K 时是 3%。

略去 F_fixed 的理由和前面一样：它不随 L 变，也不参与我们要比的那件事。只留随 L 增长的那一项：

\boxed{ F_{\text{MLA}}=4HLd_c }

4.2 HBM Bytes

关键来了。

历史不再有一份 K cache 加一份 V cache，而只有一份：

C\in\mathbb{R}^{L\times d_c}

因此 fused kernel 下：

\boxed{ B_{\text{MLA}}=bLd_c }

注意分母里没有 K + V 的那个 2。

同一份 latent 从 HBM 读进来后：

参与一次 QK；
再参与一次 PV。

所以：

AI_{\text{MLA}} = \frac{4HLd_c}{bLd_c}

得到：

\boxed{ AI_{\text{MLA}} = \frac{4H}{b} }

BF16 下 b = 2：

\boxed{ AI_{\text{MLA,BF16}}=2H }

和 2.5 节一样，这里 latent dim d_c 也被约掉了：AI 和 latent 开多宽完全无关，只和 head 数量有关。所以把 latent rank 从 512 调到 1024，KV cache 会翻倍，但这个 AI 不动。

这就是整个 MLA AI 故事最简洁的一句话：

MQA 已经让所有 Query heads 共用 KV，所以 AI 达到 H；MLA 再进一步，把原来的 K 和 V 两份历史表示压成一份 latent。这份 latent 从 HBM 搬进来以后，K 那边的计算和 V 那边的计算都用它，同一份数据被消费两次，AI 又获得一个 2× 的常数提升。

这里需要特别澄清：这并不是说“MLA 的 cache 一定比 MQA 小一半”。

例如一个 d_k = d_v = 128 的普通 MQA，每个历史 token 的 KV width 是：

128+128=256

而 DeepSeek/Kimi 常见的 latent rank 是 512，实际 cache width 甚至更大。

这个 2× 的 AI 提升来自的是：

一份 latent 同时当 K 和 V 用

而不是“字节数机械地减半”。

还有一个重要的实现条件：如果 kernel 先读一遍 latent 做 QK，扔掉，然后 PV 又从 HBM 重新读一遍，那么这层 reuse 就没了，AI 会重新掉回接近 H。所以这个结果隐含的是 FlashAttention/FlashMLA 风格的 fused/streaming execution，让同一 tile 的 latent 能在片上服务两部分计算。

拿 AI = 2H 直接代两个真实模型：DeepSeek-V3 / R1 有 128 个 head，得到 256 FLOP/B；Kimi K3 的 MLA layer 有 96 个 head，得到 192 FLOP/B。

对照第 1 节那几张卡两三百 FLOP/B 的 balance point —— 一个 batch = 1 的 decode，居然已经顶到 roofline 的拐点附近了。

5. 同一个 MLA，两种算法

第 3 节讲 absorption 的时候，是这么说的：如果每生成一个新 token 都把历史 L 个 latent 重新展开成 K 和 V，那只是把 HBM 的问题换成了一个巨大的计算问题 —— 而 absorb 正好把这笔展开省掉了。

那句话里其实藏着一个前提：「每生成一个新 token」。decode 一次只算一个 query token，展开出来的 K、V 用一次就扔，这笔钱没有任何人可以分摊。省掉它当然是纯赚。

prefill 就完全是另一回事了。一次要算成千上万个 query token，而展开出来的 K 和 V 对它们是同一份、可以全体共用 —— 这笔本来「用一次就扔」的开销，一下子被几千个 query token 分掉了。这时候还要不要省它，就不一定了。

结论有点反直觉：同一个 MLA 模型、同一份权重，decode 和 prefill 该用的算法正好相反。 这一节把两条路摊开算一遍，看看分界线到底在哪。

5.1 同一个乘法，两种括号

回到 3.1 那个式子。对一个 query head、一个历史 token，content score 就是三样东西连乘：

q_h^{\top} W_h^{K} c_j

括号加在哪一边，结果完全一样（矩阵乘法结合律），但算法完全不同：

\underbrace{q_h^{\top} \left( W_h^{K} c_j \right)}_{\text{expand}} \qquad\text{vs.}\qquad \underbrace{\left( (W_h^{K})^{\top} q_h \right)^{\top} c_j}_{\text{absorb}}

左边先把历史 latent 展开成 K，再和 query 点积；右边先把 W^K 折进 query，再直接和 latent 点积。V 那边同理，只是方向相反（3.2 节）。

第一个区别：那笔 projection 落在哪一侧。

展开：W^Kc_j 是对历史 token 做的 —— 每个历史 token 算一次，这一批里所有 query token 共用；
absorb：(W^K)^⊤q_h 是对 query token 做的 —— 每个 query token 算一次，对全部历史共用。

先把宽度记一下：

d_c：latent 的宽度，也就是第 3 节那个 c_j —— cache 里每个历史 token 就是这么宽；
d_k：展开之后每个 head 的 K；
d_v：展开之后每个 head 的 V。

第二个区别：剩下那个点积有多宽。 一对「query token × 历史 token」上，两条路都要做两次乘加 —— 一次打分，一次加权求和。把形状摆出来就清楚了。

展开之后：

打分：q_h（1 × d_k）和 k_j,h（1 × d_k）做点积，得到一个标量 —— d_k 次乘加；
加权求和：标量 p_h,j 乘上 v_j,h（1 × d_v），累加进 o_h（1 × d_v）—— d_v 次乘加。

合计 d_k + d_v。

absorb 之后：

打分：q̃_h（1 × d_c）和 c_j（1 × d_c）做点积，得到一个标量 —— d_c 次乘加；
加权求和：标量 p_h,j 乘上 c_j（1 × d_c），累加进一个 d_c 宽的 latent 输出 —— d_c 次乘加。

合计 2d_c。

打分那步的差别是显然的（d_k 对 d_c）。加权求和那步容易被忽略：展开之后累加的是 d_v 宽的 v，而 absorb 累加的是 d_c 宽的 latent。

这里要强调一句：两条路最后吐出来的 o_h 是同一个东西，都是 d_v 宽。这正是 3.2 节那一步 —— 先乘 W^V 再求和，和先求和再乘 W^V，结果完全相等：

\sum_j p_j \left( W_h^{V} c_j \right) \;=\; W_h^{V} \left( \sum_j p_j c_j \right)

区别只在 W^V 用在求和的哪一边：展开是先用，所以对每个历史 token 各用一次，之后累加的自然是 d_v 宽；absorb 是后用，所以累加过程停在 d_c 宽的 latent 上，最后补一次 W^V 变回 d_v —— 而这一次是每个 query token 只做一次，不摊在每一对上。

和 5.1 开头 K 那边的道理是同一个：括号加在哪边，答案不变，代价不同。

代进 Kimi K3（d_c = 512、d_k = d_v = 128）：展开之后是 128 + 128 = 256，absorb 是 2 × 512 = 1024。

同样一对 query/历史 token，absorb 要多算 4 倍。 latent 宽，在它上面点积就是更贵 —— 这一层直觉是对的。

两条路的区别就这么两句话：展开把 projection 记在历史头上，换来每一对更窄的点积；absorb 把 projection 记在 query 头上，代价是每一对都得在宽的 latent 上算。

5.2 先算一个 query token 对一个历史 token

5.1 只算了点积。但两条路各自还有那一次 projection，把它也算上，才看得出一对一时到底谁便宜。

展开这条： 把这个历史 token 的 latent 变成 K 和 V，要 d_c(d_k + d_v)。

absorb 这条： 把 W^K 折进 query、再把输出投影回这个 head，也是 d_c(d_k + d_v) —— 一模一样。

代进 K3（d_c = 512、d_k = d_v = 128；DeepSeek-V3 这几个宽度完全相同，只有 head 数是 128 而不是 96），两边的 projection 都是 512 × 256 = 131072：

	projection	点积	一对一合计
展开	131072	256	131328
absorb	131072	1024	132096

所以一对一的时候，胜负完全由点积那一栏决定 —— 也就是 5.1 那个 4 倍。projection 这一栏对谁都一样，抵消掉了。

到这里，展开看起来是纯赚的。

5.3 放大到 S 个 query、L 个历史

一对一算不出胜负，因为真实情况下这两栏放大的倍数不一样。

点积那一栏好说：每一对 query token × 历史 token 都要算一次，所以乘 S·L。

projection 那一栏就是 5.1 说的关键了 —— 它落在哪一侧，就按哪一侧的数量收费：

展开：W^Kc_j 是对历史 token 做的，每个历史 token 算一次就够，这一批里所有 query token 共用 → 乘 L；
absorb：(W^K)^⊤q_h 是对 query token 做的，每个 query token 算一次，对全部历史共用 → 乘 S。

这里要说明一句：W^K 是每个 head 各一份的，不是所有 head 共用 —— 在实现里，那个上投影矩阵的输出维度就是 H × (d_k + d_v)（sglang 里 kv_b_proj 的形状 kv_lora_rank → num_heads * (qk_nope_head_dim + v_head_dim)）。所以展开一个历史 token 得给每个 head 各展一次，没法在 head 之间共享。

不过 absorb 那边同理：(W^K)^⊤q_h 也是每个 head 各算一次。两条路的 projection 都按 head 计费，H 对两边一视同仁 —— 所以拿一个 head 来讨论不会偏袒谁，这也正是后面交叉点里 H 会约掉的原因。

把 H 个 head 乘上去：

F_{\text{MHA}} = 2H\left[\, L\,d_c(d_k+d_v) \;+\; S\,L\,(d_k+d_v) \,\right]

F_{\text{MLA}} = 2H\left[\, S\,d_c(d_k+d_v) \;+\; S\,L\,\cdot 2d_c \,\right]

差别就在第一项：展开乘的是 L，absorb 乘的是 S。 谁划算，取决于这一次 forward 里 S 和 L 谁大 —— 而 decode 和 prefill 恰好站在这个问题的两端。

5.4 decode：MLA 算法完胜

decode 时 S = 1，只有一个新 token —— 那笔 projection 没有任何东西可以摊。

下面把 Kimi K3 的参数代进上面两个式子（96 个 head、d_c = 512、d_k = d_v = 128），算的是一层的量：

L	MLA 算法 FLOPs / bytes / AI	MHA 算法 FLOPs / bytes / AI
2,048	0.43 G / 2.10 M / 204	51.6 G / 103 M / 503
8,192	1.64 G / 8.39 M / 195	207 G / 411 M / 503
32,768	6.47 G / 33.6 M / 193	826 G / 1.64 G / 503

MLA 算法在两个轴上同时赢：FLOPs 少 120 倍，bytes 少 71 倍。

而且这里没有「短序列用另一条」的余地。把两个式子在 S = 1 时令相等，交叉点落在 L ≈ 1.006 —— 只有历史里就一个 token 时 MHA 算法才微弱领先，从 L ≥ 2 开始 MLA 算法就一直赢。道理也直白：L 变长只是把那个 120 倍的差距按比例放大，不可能反转。

5.5 prefill：反过来

prefill 时一次要算一大片 token，S 就等于序列长度。这时展开那笔 projection 摊到了成千上万个 query token 头上。

同样代 K3 的参数，同样是一层：

序列长度	MLA 算法	MHA 算法	MHA 算法便宜
2,048	464 G	155 G	3.00×
8,192	6.80 T	1.86 T	3.67×
32,768	106 T	27.2 T	3.91×

prefill 里没有交叉点，MHA 算法在所有长度上都赢。 原因挺漂亮：S = L 时，两条路的「每 token 一次性开销」在代数上完全相等，都是 2HL·d_c(d_k+d_v) —— MLA 算法花在「把 W^K 折进每个 query、再把每个输出投影回去」，MHA 算法花在「把每个历史 token 展开成 K、V」。两边直接抵消，只剩二次项之比：

\frac{2d_c}{d_k+d_v} = \frac{1024}{256} = 4

序列越长越贴近这个 4 倍，短序列只是让优势变小（2048 时是 3 倍），但不会翻转。

5.6 那 S 多大才该换算法？

前面两节各算了一个极端：S = 1 的 decode，absorb 赢 128 倍；S 上千的 prefill，展开赢 4 倍。那中间呢 —— S 到底大到多少，就该从 absorb 换成展开？

这个问题有确定的答案。5.3 那两个式子里，MLA 那条随 S 涨，MHA 那条几乎不涨，那么必然存在一个 S 让它们打平；这个 S 就是该换算法的临界点，令 F_MLA = F_MHA 解出来就是。

先看看它长什么样。拿 K3 的参数、历史长度固定取 L = 32768，把两条算法的 FLOPs 相除画出来 —— 横轴是 S，纵轴是比值，穿过 1 的地方就是打平的地方：

两条算法的 FLOPs 之比随 S 变化。比值在 S = 1 时约 0.008，随 S 上升，在 S 约 170 处穿过 1。decode 在 S = 1、MTP 在 2–8，都远在交点左侧；prefill 的 S 上千，在交点右侧。

左端 decode 的比值只有 0.008，absorb 便宜 100 多倍；S 越大比值越往上爬 —— 因为 MLA 那条的两项都带 S，每多一个 query token 都要重算一遍，而 MHA 那条的大头是展开，只带 L、不带 S，多算几个 query token 基本不额外花钱。曲线在 S ≈ 170 穿过 1。

真解一遍会发现 H 直接约掉了（两条路的 projection 都按 head 计费，5.3 说过），交叉点只由几个 head dim 决定：

S^{*} = \frac{a L}{a + (c-e)L}, \qquad a = d_c(d_k+d_v),\; c = 2d_c,\; e = d_k+d_v

代进去，L 从 512 涨到 128K，S* 只从 128 挪到 171。也就是说：

决定用哪条算法的是「这一次算几个 query token」，不是「历史有多长」

而且因为 H 约掉了，DeepSeek-V3（128 head）和 Kimi K3（96 head）的交叉点一模一样。

图上那三条竖线就是这个意思：decode 在 S = 1、MTP 在 2–8，都死死钉在交点左边；prefill 动辄上千，稳稳在右边。两个 regime 各自待在自己那一侧，中间那段几乎没人会踩到 —— 所以现实里这从来不是一个「算一下再决定」的问题，而是按 decode / prefill 直接分派。

5.7 这一节正好说明：AI 会骗人

第 1 节留了个尾巴，说 AI 只能判断落在 roofline 哪一侧。这里就是现成的反例。

回头看 5.4 那张表的 decode 行：MHA 算法的 AI 是 503，MLA 算法只有 193。 光看 AI，你会选 MHA 算法 —— 而它实际上多花 120 倍 FLOPs、多碰 71 倍 bytes，两样都差一两个数量级。它的 AI 高，恰恰是因为分子炸了，不是因为它更划算。

prefill 那边则是另一种失灵：两条路的 AI 都在 10⁵~10⁶ 这个量级，AI 只会告诉你「都深度 compute-bound」，根本区分不出该选哪条。

一个给错答案，一个给不出答案。所以：

AI 是一个比值，它只回答「你卡在算力还是带宽这一侧」。要判断哪条路更快，得把 FLOPs 和 bytes 的绝对值摆出来看。

5.8 sglang 就是这么分的

这不是纸上推演。sglang 里 DeepSeek 的 attention 分派逻辑（python/sglang/srt/models/deepseek_common/attention_backend_handler.py）就是按 forward mode 切的：

prefill 走 MHA_ONE_SHOT / MHA_CHUNKED_KV，也就是展开那条；
decode、以及投机解码的 verify / draft 走 MLA，也就是 absorb 那条。

值得注意的是 decode 路径上完全没有按 L 分支：整个分派函数里唯一看序列长度的地方被挡在 prefill 分支里。这和 5.4 的结论对得上 —— decode 侧本来就没有值得切换的窗口。

6. Sparse attention 把方向反过来

第 5 节的结论建立在一个前提上：attention 要看全部 L 个历史 token。DeepSeek-V3.2 的 DSA、以及 GLM 的同类做法，都在这个前提上动刀 —— 先用一个轻量的 selector 给历史 token 打分，只挑出 top-k 个真正参与 attention。V3.2 的配置里这个 k 是 index_topk = 2048。

这件事对两条算法不是对称的：

MLA 算法吃得到稀疏。 cache 里存的就是一个个 latent，选出哪 2048 个就按 index gather 哪 2048 个，后面的打分和加权求和只在这 2048 个上做。成本从 L 变成 k。
MHA 算法吃不到。 它得先把历史展开成 K、V，再交给一个稠密的大 GEMM —— 而 MHA 这条路上的算子只提供稠密 GEMM，没有「只对选中的这几个 token 做」的 selective GEMM。所以展开那笔钱还是按整个 L 付。

于是 L 越长，两者差得越离谱。这里代 DeepSeek-V3 的参数（128 个 head），decode、一层：

L	MLA 算法 + sparse	MHA 算法（稠密）	差距
2,048	0.57 G	68.9 G	121×
8,192	0.57 G	275 G	483×
32,768	0.57 G	1.10 T	1931×
131,072	0.57 G	4.41 T	7725×

sparse 那条一旦超过 k 就是平的，稠密那条一直是线性的。

所以稀疏之后，选择的方向和第 5 节反过来了：稠密时是「序列越长越该用 MHA 算法」，稀疏时变成「序列越长越该用 MLA 算法」。sglang 里 DSA 的分派（dsa_backend.py）正好印证：只有当 max_kv_len 不超过一个阈值时才走 MHA 算法，超过就走 sparse MLA，而 decode / verify 一律走 MLA。

那个阈值的默认值是 2048 —— 和 index_topk 一模一样。含义很清楚：kv 长度还不到 2048 时，top-k 等于全选，稀疏一点便宜都占不到，那就不如用更快的稠密 kernel；一旦超过，稀疏才开始真的省东西。

7. 四种结构对比

下面全部指 BF16、single-token decode，只算 attention core 对历史 KV 的那一遍读取和计算，并假设 fused kernel，不计 softmax / projection / output projection 等低阶项。

Attention	Historical Cache 结构	主要 reuse	BF16 Decode AI
MHA	每个 Q head 各有 K/V	基本无跨-head reuse	1
GQA	一组 Q heads 共用 K/V	KV 跨 Query heads reuse	H_q/H_kv
MQA	所有 Q heads 共用一份 K + 一份 V	最大 cross-head reuse	H_q
MLA	一份 latent	cross-head reuse + 一份 latent 同时当 K 和 V 用	2H_q
DeepSeek V3/R1	一份 latent，128 个 head	同上	256 FLOP/B
Kimi K3 的 MLA layer	一份 latent，96 个 head	同上	192 FLOP/B

把这条路径压成一行就是：

\boxed{ 1 \rightarrow \frac{H_q}{H_{kv}} \rightarrow H_q \rightarrow \sim2H_q }

我觉得这个视角比“某个 Attention variant 省了多少 KV cache”更统一：

Attention 架构的演化，很大一部分其实是在设计 data reuse。GQA/MQA 复用的是不同 Query heads 之间的 KV；MLA 再进一步，让一份 latent 同时当 K 和 V 用。

8. 为什么 MLA + MTP 容易吃亏

现在可以重新看开头那句话了。

普通 autoregressive decode 每次只有一个 query token。MTP / speculative decoding 的系统直觉则是：既然历史 KV cache 已经从 HBM 搬进来了，能不能一次让它服务多个待验证 token？

设一次 verify 同一份历史 cache 服务 S 个 query positions。在一个非常粗略的模型下，从 HBM 读的数据量没有按 S 同比例增加，但 QK/PV computation 会接近乘 S，因此：

AI(S)\approx S\cdot AI(S=1)

对于普通 MQA：

AI_{\text{MQA}}\approx SH

而 MLA：

AI_{\text{MLA}} \approx SH\frac{2d_c+d_s}{d_c+d_s}

这里可以顺便回答一个问题：MTP 把 query token 数从 1 抬到了 S，那会不会抬过第 5 节那个交叉点、让它改走 MHA 算法？不会。 投机窗口一般是 2 到 8，撑死几十，离 171 还差得远。所以 MTP 期间仍然停在 MLA 算法这一侧 —— 每个 query token 还是要在 576 宽的 latent 上算一遍。这正是它的 AI 会老老实实跟着 S 往上顶的原因。

也就是说，MTP 和 MLA 其实都在花同一种资源：原本 memory-bound decode 中闲置的 GPU compute。

MLA 用额外 compute 换更强的 cache reuse；
MTP 用额外 speculative / verify compute 换更少的串行 decoding steps。

如果 workload 原本是一个 AI 很低的 MQA，例如 AI ≈ 70~100，GPU 离 Roofline knee 还很远，MTP 多塞一些计算往往是在利用原本闲置的 Tensor Core。

但 DeepSeek-style MLA 单 query 就已经可以达到约：

256\ \text{FLOP/B}

Kimi K3 的 Gated MLA layer 也约为：

192\ \text{FLOP/B}

对比前面的 theoretical balance point：H200 约 206 FLOP/B，H100/B200 也就在两三百 FLOP/B 的量级。

所以 MLA 在 S = 1 时，本来就已经把 decode 从“明显 memory-bound”推到了接近 compute/memory 平衡的位置。再让 S 变成 2：

\text{DeepSeek: }256\rightarrow512

\text{Kimi K3: }192\rightarrow384

很容易直接越过 Roofline knee，进入 compute-bound。

这时 MTP 增加的计算就不再是“免费吃掉 spare compute”了，而会真实地开始占 latency。

所以开头那句“MLA 在 Decoding 时表现为 head_dims=512+ 的 MQA，已经提前消耗了大部分算力”，我现在的理解可以浓缩成：

MLA 和 MTP 都是在做 compute-for-bandwidth / compute-for-latency 的交换。MLA 先用掉了一大块 decode 阶段原本闲置的 compute headroom，因此留给 MTP 的免费算力更少。

数学上，MLA 相对 MQA 的 AI 其实“只”多了一个不到 2 的常数；但系统上，这个常数刚好可能把 workload 推过 Roofline 的分界线。

这也是我觉得这个结论最有意思的地方。

致谢

感谢 Yangmin 在 MLA inference 上的洞见与讨论。

写完之后我才发现，Zyphra 那篇 Compressed Convolutional Attention（arXiv:2510.04476）把同一件事独立地说了一遍，而且用词几乎一样：

MLA has an arithmetic intensity of 2n_heads at inference time, which is very large, and targets the ridge of the roofline plot on a H100 for a single query. … The arithmetic intensity required to breach the ridge of the roofline for an Nvidia H100 under bfloat16 is 295 FLOPs per byte. Deepseek seemingly chose their number of heads for the DeepseekV3 model to accordingly saturate the roofline to approach compute bound inference at batch size 1. … This falls short in cases where speculative decoding is utilized such that the arithmetic intensity passes the roofline.

AI = 2H、295 FLOP/B、以及「DeepSeek 是照着 roofline 挑的 head 数」——三条都对上了。它还提了一个本文没展开的角度：MLA 在 tensor parallel 下也吃亏，因为那份共享的 KV 得按 TP rank 在每张卡上复制一遍，MQA 省下来的复用就没了。

那篇文章里还有一句话，正好可以用来收尾第 5 节那个「AI 会骗人」的观察：

while MLA is capable of better compute utilization on decode, the increase in FLOPs that would otherwise go unused does not automatically result in victory. Model quality and latency, not SM utilization, is the end goal.

附录 A：把 RoPE 加回来

前面的 AI = 2H 是最干净的主结论，但它建立在「忽略 RoPE」之上。这里补回来。

问题出在 3.1 节那步 absorption：它靠的是矩阵乘法的结合律，把 W^K 挪到 query 那一侧。而 RoPE 是一个跟位置有关的旋转，插在 query 和 key 中间，两边转的角度不一样，这个结合就做不成了 —— 一旦带上 RoPE，就没法再直接在 latent 上打分。

DeepSeek 的解法是分工：latent 那部分照常 absorb、不带 RoPE；另外再单独留一小段专门带 RoPE 的 key，所有 head 共用一份，跟着 latent 一起存进 cache。这一小段只参与打分，不参与最后的加权求和（它没有对应的 Value）。

把这一小段的宽度记作：

d_s

那么每个历史 token 实际 cache width 是：

d_c+d_s

QK contraction width 是：

d_c+d_s

PV contraction width 仍然只有：

d_c

因此：

F = 2HL(d_c+d_s)+2HLd_c

即：

\boxed{ F=2HL(2d_c+d_s) }

HBM traffic：

\boxed{ B=bL(d_c+d_s) }

所以：

\boxed{ AI_{\text{MLA}} = \frac{2H}{b} \frac{2d_c+d_s}{d_c+d_s} }

BF16 下：

\boxed{ AI_{\text{MLA,BF16}} = H \frac{2d_c+d_s}{d_c+d_s} }

当 d_s ≪ d_c 时，这个因子自然接近：

2H

A.1 DeepSeek V3 / R1

DeepSeek-V3 config 中：

H=128,\quad d_c=512,\quad d_s=64

因此：

AI = 128\times\frac{2\times512+64}{512+64}

\boxed{ AI_{\text{DeepSeek MLA}}\approx241.8\ \text{FLOP/B} }

A.2 Kimi K3

Kimi K3 config 的 Gated MLA layer 中：

H=96,\quad d_c=512,\quad d_s=64

所以：

AI = 96\times\frac{1088}{576}

\boxed{ AI_{\text{Kimi K3 MLA}}\approx181.3\ \text{FLOP/B} }

K3 有个实现上的特殊点：配置里保留了 qk_rope_head_dim=64 的 layout，但同时设了 mla_use_nope=true，所以那 64 维并不是真的在做 RoPE。对本文的推导来说这不影响 —— 我们只需要知道它同样是所有 head 共用、只参与打分、不参与加权求和的一小段，宽度 64，代进上面的式子就行。

A.3 第 5 节那几个数

第 5 节比较两条算法时也用了同样的简化。把这一小段加回来之后：

	正文（简化）	加上这一小段
每一对，absorb	1024	1088
每一对，展开后	256	320
展开的准备费	131072	131072
每一对的比	4.0×	3.4×

准备费不变（那一小段是所有 head 共用的，不用按 head 展开）。两边都变宽了，所以比值从 4 降到 3.4，prefill 那边 MHA 算法的优势也从 3.9 倍降到 3.3 倍左右。

交叉点则完全不动：它由准备费除以「每一对的差」决定，而这个差是 1024 − 256 = 768，加上之后是 1088 − 320 = 768 —— 一模一样。所以 S* ≈ 171 这个结论不受影响。

1. 什么是 Arithmetic Intensity#

2. MHA / GQA / MQA 的统一公式#

2.1 FLOPs 怎么数#

2.2 MHA 简介#

2.3 FLOPs#

2.4 HBM Bytes#

2.5 两式相除#

2.6 三个 special case#

2.7 代入真实模型#

2.8 那 prefill 呢？#

3. MLA 的 latent cache#

3.1 K-side absorption#

3.2 V-side absorption#

3.3 为什么 projection 没有变贵#

4. MLA 的 AI 为什么翻倍#

4.1 FLOPs#

4.2 HBM Bytes#

5. 同一个 MLA，两种算法#

5.1 同一个乘法，两种括号#

5.2 先算一个 query token 对一个历史 token#

5.3 放大到 S 个 query、L 个历史#

5.4 decode：MLA 算法完胜#

5.5 prefill：反过来#

5.6 那 S 多大才该换算法？#

5.7 这一节正好说明：AI 会骗人#

5.8 sglang 就是这么分的#

6. Sparse attention 把方向反过来#

7. 四种结构对比#

8. 为什么 MLA + MTP 容易吃亏#

致谢#

附录 A：把 RoPE 加回来#

A.1 DeepSeek V3 / R1#

A.2 Kimi K3#

A.3 第 5 节那几个数#