MLA 逐步推导，从低秩 KV 到矩阵吸收

从 MHA 和 GQA 接着往下推。GQA 让多个 Q head 共用一组已经投影好的 K/V。MLA 让多个 head 的 K/V 都从一个共享的低维向量生成，缓存保存这个低维向量。

只看这一句，还解释不了 MLA 为什么值得做。每一步都把所有历史重新展开，可能会引入大量计算。我们先按最直接的方式写出计算，再一步步移动矩阵乘法，看看哪些历史计算可以消掉。

全文使用行向量。输入 token 写成 \(x\in\mathbb R^{1\times D}\)，投影在右边相乘。符号 \(C\) 单独出现时表示历史 latent 矩阵，上标 C 表示 content，即内容分支。推导中的矩阵维度会在使用前定义。

系列路线 · 前置内容 MHA、MQA、GQA

1. 先算清楚原来的缓存

先沿用前一讲的演示配置：

MHA 每个 token、每层保存 32 个 key 和 32 个 value，所以元素数是

一个 BF16 元素占 2 B，因此是 16 KiB。历史长度为 8192 时，一层占 128 MiB。8 个 KV head 的 GQA 则是一层 32 MiB。

MLA 先问一个表示层面的问题。每个 head 的 K/V 是否可以由同一个较短的向量，通过各自的线性变换生成？假设这个向量只有 512 维，那么每 token 先保存 512 个元素。后面处理位置编码时还会再加一部分缓存。

2. 从输入得到 latent，再得到 K/V

对一段历史输入

定义降投影矩阵

于是共享 latent 为

每个 head 再有自己的两个上投影：

本文先取 \(d_v=d_h=128\)。该 head 的内容 key、value 为

相同的 C 经过不同的 \(U_{K,h}\)，会生成不同 head 的内容 key；经过不同的 \(U_{V,h}\)，又会生成不同的 value。因此缓存共享并不要求所有 head 最终使用相同的 K/V。

把所有 head 的上投影横向拼起来，可写成

当 \(d_c=512,H=32,d_h=d_v=128\) 时，两个矩阵都是 \([512,4096]\)。乘法 \([8192,512]\times[512,4096]\) 可以一次生成全部 head 的内容 K 或 V。

这个结构需要模型学习相应的投影。对一个任意的已训练 MHA，把 K/V 强行压到较小维度通常不能保持原模型输出。接下来证明的等价性，是同一个 MLA 模型内部两种计算顺序之间的等价性。MLA 的结构与解耦位置编码来自 DeepSeek-V2。DeepSeek-V2，§2.1 与附录 C

3. 最直接的 Decode 为什么有问题

假设新 token 已经产生当前 query \(q_h\in\mathbb R^{1\times d_h}\)，历史 C 为 \([T,d_c]\)。直接做法是先重建该 head 的历史 K/V：

然后计算

暂时忽略位置编码，单看重建 K 的计算量就是 \(2Td_cd_h\) FLOPs。取 \(T=8192,d_c=512,d_h=128\)，得到每 head 1,073,741,824 FLOPs，约 1.074 GFLOPs。V 的重建还要再做一遍。

这时缓存虽然小了，每步却要对整个历史做上投影。问题出在我们把 \(CU_K\) 放到了最前面。下面把表达式代进去，检查乘法顺序能否调整。

4. K 的吸收，一步一步移动转置

从单 head 的未缩放分数开始：

代入 \(K_h^C=CU_{K,h}\)：

使用 \((AB)^{\mathsf T}=B^{\mathsf T}A^{\mathsf T}\)：

再使用矩阵乘法的结合律：

先检查维度，避免只看符号：

现在只为当前 Q 做一次 \([1,128]\times[128,512]\)。历史 latent 直接参加打分，省去了本步对所有历史位置的 K 上投影。一次 query 变换约 131,072 FLOPs，比刚才那次历史 K 重建少了 \(T\) 倍。

但后续打分的规约长度从 128 变成 512。原来的 \(qK^{\mathsf T}\) 是约 2.097 MFLOPs，现在 \(q'C^{\mathsf T}\) 是约 8.389 MFLOPs。这两项变化要放在一起看。

5. V 的吸收，为什么能移到求和后面

此时 \(p_h\in\mathbb R^{1\times T}\) 已经算好。输出为

仍然使用结合律：

先得到 latent 空间里的加权汇总

再投影回该 head 的 value 维度

按位置展开，也能看清原因：

同一个线性变换 \(U_{V,h}\) 作用于每个位置，因此可以移到加权求和之后。Softmax 仍在产生 \(p_h\) 的位置执行。

到这里，历史 C 会参与两个主要操作。一条路径用 \(q'_hC^{\mathsf T}\) 产生分数，另一条用 \(p_hC\) 汇总内容。两条路径在数学上共享 C。具体内核若能把一个 C tile 留在片上，才有机会减少两次使用之间的 HBM 重读。

6. 用一个小矩阵把两条路径算到相同结果

取两个历史位置，latent 维度为 2，单 head 的内容与 value 维度都为 1：

展开历史 key：

吸收后先算 query：

两种顺序产生相同分数。这里 head 维度为 1，缩放因子也是 1。令

直接展开 value：

吸收路径先聚合 C：

数值约为 10.9852。这个例子同时验证了 K 侧和 V 侧的重排。浮点实现中，规约顺序和精度会带来小误差。

7. RoPE 为什么会打断刚才的吸收

先看普通 RoPE。用 \(R_j\) 表示第 \(j\) 个位置的旋转矩阵，按行向量记法，旋转后的 key 为

如果 query 也在位置 \(i\) 做旋转，打分里就出现

要像前面那样形成一个 query-side 投影，投影中会包含 \(R_j^{\mathsf T}\)。它随着被读取的历史位置 \(j\) 改变，因此当前 query 无法只变换一次就覆盖所有历史位置。矩阵维度相同也不能随意交换这些矩阵的乘法次序。

MLA 为此把内容匹配和位置匹配分成两条路径。内容 key 保持为 \(c_jU_{K,h}\)。位置 key 单独生成，经过 RoPE 后保存为 \(k_j^R\)。该位置 key 跨 Q head 共享，各 head 有自己的位置 query \(q_{i,h}^R\)。

内容 query 为 \(q_{i,h}^C\)，两项分数相加：

内容分支再次应用刚才的吸收：

先合并两项，再做一次 Softmax。两个分支各自 Softmax 再相加，会得到另一个函数。对于使用长上下文缩放的具体模型，还应按其配置保留额外的缩放规则。这里展示基础形式。

因此每个历史位置实际保存

8. 真实尺寸与演示配置分开算

前面用 32 个 head 解释结构。现在换到 DeepSeek-V2 的这组 MLA 参数，专门分析 Attention 形状：

本文的 \(D\) 继续表示模型 hidden size，数值分析时使用具体 checkpoint 的配置。每 token、每层的缓存为

8K 历史是一层 9 MiB。其中 latent 8 MiB，位置 key 1 MiB。这个 576 元素的缓存由所有 head 共享。相同 128 head、128 维的 MHA 是每 token 64 KiB，一层 8K 历史为 512 MiB；8 KV head 的 GQA 则为 32 MiB。比较时应明确 head 数和每层维度，避免把 32-head 与 128-head 的 MHA 混为同一基线。

9. Q 的低秩投影怎样接进来

MLA 还可以压缩 Q 的中间激活。先得到

随后各 head 从 \(c_i^Q\) 产生两份 query：

对应权重形状为 \([d_q,d_h]\) 和 \([d_q,d_r]\)。实际模型中的 RMSNorm 要放在配置规定的低秩激活处。它影响激活生成，但当 C 和 Q 已经生成后，前面证明的 K/V 乘法重排仍然成立。

内容 query 还可以继续代入吸收：

括号中是两个固定权重的乘积，可以预先合成。然而它的输出宽度从 \(d_h\) 增到了 \(d_c\)。两次窄 GEMM 与一次宽 GEMM 的权重存储、计算和融合代价并不相同。运行时可以选择只变换当前 query，也可以选择合成权重，需要根据实际后端判断。

输出也有类似选择。将输出投影按 head 切成 \(W_{O,h}:[d_v,D]\)，则

预合成后的每 head 权重是 \([d_c,D]\)。这展示了线性结构允许的重排，也说明“可以吸收”之后还要检查合成权重是否更大。

10. Prefill 完整走一遍

取首次 Prefill 长度 \(S=4096\)，此时没有更早的缓存。省略 batch，按上述 128-head 配置：

1. 输入 \([4096,D]\) 经 KV 降投影，得到 \(C:[4096,512]\)。
2. 同一输入产生共享的位置 key，形状 \([4096,64]\)，按位置做 RoPE。
3. Q 降投影得到 \([4096,1536]\)，上投影产生内容 Q \([128,4096,128]\) 和位置 Q \([128,4096,64]\)。
4. 一种执行方式先把 C 展开为每 head 的内容 K 和 V，形状各为 \([128,4096,128]\)。
5. 每 head 将内容、位置打分相加，执行因果 Softmax 与 PV。
6. 合并各 head 输出，执行输出投影。供后续 Decode 使用的持久缓存可以保存 C 与位置 key。

展开路径每 head 的主要形状是

也可将内容和位置 Q/K 拼成 192 维打分。图中完整分数矩阵只是数学表示，融合内核可以按 tile 处理并避免把它整体写回 HBM。

Prefill 有 4096 行 query，展开出的 K/V 能被很多行反复使用。吸收路径则要对每行 query 做内容变换，且后续内容 QK 与 latent PV 都用更宽的 512 维。哪条路径更快取决于历史长度、投影成本、临时存储和内核实现。已有前缀的 chunked Prefill 还要计入旧历史读取，不能只看新 chunk 的长度。

11. Decode 完整走一遍

假设本步追加后历史长度 \(T=8192\)。当前输入是一行：

1. 为新 token 生成 512 维 latent 和 64 维位置 key，写入各自缓存。
2. 为当前 token 生成 128 个内容 Q 和位置 Q。
3. 对每个内容 Q 做吸收，得到 512 维 \(q'_h\)。
4. 计算内容分数与位置分数，合并、缩放、mask，再统一 Softmax。
5. 用概率聚合 C，得到每 head 的 512 维 \(z_h\)。
6. 做 value 上投影得到 128 维 \(o_h\)，再做输出投影。

对应单 head 的矩阵形状如下：

内容 QK、位置 QK 和 latent 聚合按一次乘加 2 FLOPs 计算，合计约

代入得到约 2.282 GFLOPs。若同样 128 个 Q head 的普通 Attention 使用 192 维 Q/K 和 128 维 V，主 QK+PV 则为 \(2HT(192+128)\)，约 0.671 GFLOPs。这里仅对比扫描主体，另有投影、Softmax 和数据转换。它说明缓存变小之后，部分矩阵计算可能增加。

12. 为什么 NPU 上不能只看 K=512 是否整齐

假设矩阵阵列的规约 tile 宽 256。内容 QK 的 512 可以分两段，位置 QK 的 64 在简单 padding 映射下只占四分之一。也可以在逻辑上拼接

把两项分数合成 576 维的一次规约。它可能减少分支合并与中间写回，但会引入拼接布局、对齐和尾块处理。实际硬件也可能支持更细的规约形状。

再看 M 维。单 head Decode 是 \(M=1\)，规约维填满也不能保证整个阵列填满。因为各 head 共享 C，可以把多个 head 的吸收后 query 堆在一起，形成 \([H,512]\times[512,T]\)。这些行各自归一化，概率矩阵再乘同一份 C。这样可以利用 head 维提供更多行，并把同一 latent tile 供给多行。

取 128-token 的 latent tile，BF16 的 C 占 \(128\times512\times2=128\) KiB，位置 key 再占 16 KiB。若一次处理 16 个 Q head，FP32 输出累积 \([16,512]\) 还占 32 KiB，分数、行统计量及双缓冲也要空间。一次处理更多 head 能提高共享复用，也会挤占片上容量。

因此应测量实际的 C 读量、矩阵单元活跃度、Softmax/重缩放用时，以及跨核部分输出归约。沿历史拆分可以增加任务数，但需要合并每段的归一化状态与 512 维部分输出。这些是硬件分析模型，文中的数值没有使用 NPU 真机跑分。

13. 用代码验证两种顺序得到相同输出

下面使用小尺寸 FP64 数组，只验证已生成的 Q、latent、位置分支之间的代数等价性。加入了 causal mask 和两条打分分支。用真实权重验证时，再把投影、RMSNorm、RoPE 的生成部分接到前面。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

import numpy as np

rng = np.random.default_rng(7)
query_length, history_length = 3, 5
latent_dim, content_dim, rope_dim, value_dim = 4, 2, 2, 3

latent = rng.normal(size=(history_length, latent_dim))
key_up = rng.normal(size=(latent_dim, content_dim))
value_up = rng.normal(size=(latent_dim, value_dim))
query_content = rng.normal(size=(query_length, content_dim))
query_rope = rng.normal(size=(query_length, rope_dim))
key_rope = rng.normal(size=(history_length, rope_dim))

# Query rows correspond to the last three positions in the history.
query_positions = np.arange(history_length - query_length, history_length)
visible = np.arange(history_length)[None, :] <= query_positions[:, None]
scale = np.sqrt(content_dim + rope_dim)

def softmax(scores):
    scores = np.where(visible, scores / scale, -np.inf)
    probabilities = np.exp(scores - scores.max(axis=-1, keepdims=True))
    return probabilities / probabilities.sum(axis=-1, keepdims=True)

expanded_key = latent @ key_up
expanded_value = latent @ value_up
scores_expanded = query_content @ expanded_key.T + query_rope @ key_rope.T
output_expanded = softmax(scores_expanded) @ expanded_value

query_absorbed = query_content @ key_up.T
scores_absorbed = query_absorbed @ latent.T + query_rope @ key_rope.T
output_absorbed = (softmax(scores_absorbed) @ latent) @ value_up

np.testing.assert_allclose(scores_expanded, scores_absorbed, rtol=1e-12, atol=1e-12)
np.testing.assert_allclose(output_expanded, output_absorbed, rtol=1e-12, atol=1e-12)
print("Matrix absorption verified")

14. MLA 与 GQA 的选择怎样比较

把比较放回每 token、每层。BF16、8 个 128 维 KV head 的 GQA 为 4096 B，前面 MLA 的缓存为 1152 B。若 GQA 只有两个 KV head，载荷变成 1024 B，已经低于这个 MLA 配置。缓存宽度与模型质量、投影方式必须一起比较。

GQA 的计算路径沿用多头 Attention 的基本形状，MLA 则多了 latent 表示、解耦 RoPE，以及展开或吸收两种执行路径。模型训练是否能学到足够好的低秩表示、推理后端是否支持这些形状、业务中 Prefill/Decode 各占多少，都会影响选择。

已有模型也使用混合方案。例如 Kimi Linear 将 KDA 与全局 MLA 层组合。全局层仍可读取历史，而递推层使用自己的状态。它说明 MLA 可以成为架构中的一种层，而不必覆盖全部层。Kimi Linear

接下来沿路线学习 YOCO。前面共享发生在 head 之间，YOCO 会把问题推进到层之间。要继续检查的是哪些层生成共享 K/V、哪些层读取，以及 Prefill 中哪些位置还需要通过后半段网络。返回系列路线