MLA 数学公式推导 — 从 MHA 到 Absorbed MQA 的等价变换

1. 起点：标准 MHA 的注意力公式

对于第 h 个 head，标准 MHA 的计算为：

(1) Q h = X \cdot W Q h, K h = X \cdot W K h, V h = X \cdot W V h

(2) O h = softmax( Q h \cdot K h T / \sqrtd ) \cdot V h

KV Cache 需要存储：每个 head 的 K_h 和 V_h。对于 n_h = 128 个 head，每 token 存 128 × (d_k + d_v) 维。

MHA 的问题：128 heads × (128 + 128) dims × 2B = 65,536 B/token。KV Cache 太大！

2. GQA 的解决方案（对照组）

GQA 将 n_h 个 Q head 分为 n_g 组，每组共享一对 KV head：

(3) K g = X \cdot W K g [S, d k], V g = X \cdot W V g [S, d v] (g = 1, ..., n g)

KV Cache = n_g × (d_k + d_v)。当 n_g = 1 时退化为 MQA（所有 head 共享 1 对 KV）。

3. MLA 第一步：联合低秩压缩 KV

1 MLA 的核心思想：用一个共享的低秩投影把所有 head 的 KV 压缩到同一个 latent 向量 c_t：

压缩 (4) c t = x t \cdot W DKV [1, d model] \times [d model, d c] \to [1, d c] 其中 d c = 512 ≪ n h \times d h

2 需要做注意力时，再解压缩回每个 head 的 K 和 V：

解压 (5) K h,nope = c t \cdot W UK h, V h = c t \cdot W UV h [1, 512] \times [512, 128] \to [1, 128] (per head)

关键：KV Cache 只需要存 c_t [512]，而非 128 个 head 的独立 KV。存储量从 128 × 256 = 32,768 降到 512，压缩 64 倍！

但这与 GQA 有什么关系？请注意：所有 head 共享同一个 c_t，但通过不同的 W_UK^h 和 W_UV^h 解压出不同的 K_h 和 V_h。

对比 GQA：GQA 的 n_g=1（即 MQA）时，所有 head 共享同一对 K, V（重复复制）。而 MLA 共享的是同一个 latent，但每个 head 解压出不同的 K, V。MLA 严格比 MQA/GQA 更具表达力。

4. MLA 第二步：分离 RoPE (Decoupled RoPE)

位置编码 RoPE 需要作用在 K 上，但 RoPE 有位置相关的旋转矩阵 R_t，如果直接对 c_t · W_UK 施加 RoPE，会破坏压缩结构（后面矩阵吸收时无法消除 R_t）。

解决方案：把 K 拆成两部分——不带 RoPE 的 K_nope 和带 RoPE 的 K_rope：

Key 拆分 (6) K h,nope = c t \cdot W UK h [128 dims, 不带位置信息] K rope = RoPE( x t \cdot W KR) [64 dims, 带位置信息, 所有 head 共享]

完整 Key (7) K h = [ K h,nope; K rope] [128 + 64 = 192 dims]

同样，Q 侧也做对应的拆分：

Query 拆分 (8) Q h,nope = q c \cdot W UQ h [128 dims] Q h,rope = RoPE( q c \cdot W QR h) [64 dims]

KV Cache 内容：存储 c_t [512] + k_rope [64] = 576 dims/token。k_rope 必须单独存，因为它带了 RoPE 的位置信息，无法被压缩进 c_t。

5. 核心推导：矩阵吸收 → 等价 MQA

这是整个推导最关键的部分。我们来看 Attention Score 的 nope 部分如何通过矩阵吸收变成 MQA 形式。

5.1 原始公式（Non-Absorbed, 即标准 MHA 形式）

原始 (9) Score nope = Q h,nope \cdot K h,nope T

展开 Q 和 K 的定义：

(10) = ( q c \cdot W UQ h) \cdot (c t \cdot W UK h) T

利用转置的性质 (AB)^T = B^TA^T：

(11) = q c \cdot W UQ h \cdot W UK h T \cdot c t T

⬇ 关键一步：矩阵结合律

5.2 矩阵吸收（Absorbed Form）

利用矩阵乘法结合律 (ABC) = A(BC)，把 W_UQ 和 W_UK 合并：

吸收 (12) = q c \cdot (W UQ h \cdot W UK h T) \cdot c t T

定义吸收后的 Q 投影矩阵：

定义 (13) W Q' h = W UQ h \cdot W UK h T [d nope, d nope] \times [d nope, d c] T \to [d nope, d c] 即 [128, 128] \times [128, 512] T \to [128, 512]

代入后，Score 变成：

吸收后 (14) Score nope = ( q c \cdot W Q' h) \cdot c t T = Q' h \cdot c t T

关键发现！吸收后，K 侧直接使用原始的 c_t（不需要通过 W_UK 解压缩），而且 c_t 在所有 128 个 head 之间是共享的。这就是 MQA 的形式——Q 是 multi-head 的，K 是 single-head 共享的！

5.3 V 侧同理

Output 的计算也可以做同样的吸收：

(15) O h = softmax(Score) \cdot V h = softmax(Score) \cdot c t \cdot W UV h

吸收 W_UV 到 W_O 侧：

V 侧吸收 (16) O = \sum h O h \cdot W O h = \sum h [ softmax(Score) \cdot c t] \cdot (W UV h \cdot W O h)

吸收后，V 也直接使用 c_t（共享），无需解压！

6. 完整公式：Absorbed MLA = 分数合并的 MQA

把 nope 部分和 rope 部分合起来，Absorbed MLA 的完整计算为：

Score (17) Score h = Q' h,nope \cdot c t T + Q h,rope \cdot K rope T [1, 512] \times [512, S] + [1, 64] \times [64, S] \to [1, S]

等价写法 (18) Score h = [Q' h,nope; Q h,rope] \cdot [c t; K rope] T [1, 576] \times [576, S] \to [1, S]

7. 一张表看懂：MHA → GQA → MQA → Absorbed MLA

维度	Standard MHA	GQA (n_g 组)	MQA (n_g=1)	Absorbed MLA
Q heads	n_h = 128	n_h = 128	n_h = 128	n_h = 128
KV heads	n_h = 128	n_g (e.g., 8)	1	1 (= c_t)
head_dim_k	128	128	128	576 (= 512 + 64)
head_dim_v	128	128	128	512 (= d_c)
KV Cache/token	128 × 256 = 32,768	8 × 256 = 2,048	1 × 256 = 256	1 × 576 = 576
表达力	最强	中等	最弱	接近 MHA
Score GEMM	[128] × [128, S]	[128] × [128, S]	[128] × [128, S]	[576] × [576, S]

Absorbed MLA 本质上等价于一个 head_dim=576 的 MQA (n_kv=1)。
相比标准 MQA (head_dim=128)：Score GEMM 的 K 维度更大 (576 vs 128)，计算量约 4.5 倍，但 KV Cache 只多存 576 vs 256（2.25 倍），而表达力远超 MQA 接近 MHA。这就是 MLA 的精妙之处——用略多的计算量换取了远超 MQA 的表达力，同时保持了 MQA 级别的 KV Cache 效率。

8. 为什么 RoPE 必须分离？

假设不分离 RoPE，直接对 K = c_t · W_UK 施加 RoPE：

假设 K h = R t \cdot ( c t \cdot W UK h)

尝试吸收：

Q \cdot K T = q \cdot W UQ \cdot (R t \cdot c t \cdot W UK) T = q \cdot W UQ \cdot W UK T \cdot c t T \cdot R t T

问题：R_t^T 在 c_t^T 的右边！因为 R_t 依赖于 token 位置 t，每个 token 的 R_t 不同，它无法被合并进静态的权重矩阵中。矩阵吸收失败！

这就是为什么 DeepSeek 必须将 RoPE 分离（Decouple）出来——让 RoPE 作用在一个独立的 64 维 k_rope 上，而非作用在 c_t · W_UK 上。这样 nope 部分可以正常吸收，rope 部分单独计算后相加。

9. 一张图总结全部推导

最终结论：
MLA Absorbed Form = MQA，其中 n_kv=1，head_dim_k=576，head_dim_v=512。
MLA Non-Absorbed Form = MHA，其中 n_kv=128，head_dim_k=192，head_dim_v=128。

两者在数学上完全等价（产生相同的 attention output），区别仅在于矩阵乘法的结合顺序不同，导致计算特征（compute vs memory bound）不同。GQA 的任何配置都可以转化为等价的 MLA（TransMLA 论文已证明），但反过来不一定成立——MLA 严格比 GQA 更具表达力。