MLA 推理提速 13 倍：一个被忽略的矩阵恒等式

在 DeepSeek-V2 等模型采用的多头潜在注意力（MLA）架构中，最常见的「naive」实现在长上下文解码阶段性能远逊于标准多头注意力（MHA）。一篇技术博客指出，只要在加载模型时利用一个简单的矩阵乘法恒等式对权重做一次变换，就能把 MLA 的解码吞吐从约 5,000 tok/s 提升到 71,000 tok/s，在作者测试条件下达到约 13 倍加速。

问题：naive 实现为何慢

MLA 的核心思路是在 KV cache 中只保存低秩 latent 表示（形状为 (B, H, T, r)，其中 r 远小于 head dim d），解码时再通过 W_k_up、W_v_up 两个矩阵把它上投影回完整的 K 和 V，再做标准的缩放点积注意力。

直觉上，cache 更小，解码理应更快。但在 32K 上下文、batch=16 的设置下，作者实测 naive MLA 的吞吐只有约 5,000 tok/s，而相同参数规模的 MHA 能跑到 42,000 tok/s。

瓶颈来自上投影：在每个 decode 步，都要对整段缓存执行两次量级为 O(B·H·T·r·d_head) 的 einsum，T 越大代价越高。结果是延迟从 1K 上下文的 0.15 ms 一路涨到 32K 的 3 ms，单纯上投影就比注意力本体还贵。

解法：把上投影吸收进 Q

作者利用了线性代数的结合律，把 K[t] = K_lat[t] · W_k_up 代入注意力打分式后得到：

• score[t] = q · (K_lat[t] · W_k_up)ᵀ = (q · W_k_upᵀ) · K_lat[t]ᵀ

其中 q · W_k_upᵀ 与时间步 t 无关，可以在循环外只算一次。V 侧同理：先在 latent 空间做加权聚合，最后再把结果投影回 head 维度。

重构后的 decode 步只剩三步：

• 把 q 与 W_k_up 相乘，得到形状 (B, H, 1, r) 的 q_abs；
• 用 q_abs 直接与缓存中的 K_lat 做注意力打分；
• 在 latent 空间做加权求和，再用 W_v_up 一次性投影回输出维度。

整个过程不再需要在每个 cached token 上重新构造完整的 K 和 V。

实测数据

作者在 H100、bf16 下做了上下文长度扫描：

• 作者自用模型规模（H=8, d_head=48）：naive 与 absorbed 之间的差距随上下文拉长持续扩大，长上下文下加速比达到约 13–29 倍；
• 生产规模（H=32, d_head=128, r=32, batch=16）：在 128K 上下文附近，absorbed 版本相对 naive 快近 30 倍；
• 峰值吞吐约 71,000 tok/s，是同设置下 naive 5,000 tok/s 的约 14 倍。

与 GQA、MQA 的横向对比

吸收版 MLA 解决了 naive MLA 的解码开销问题，但并没有让它在同缓存量下反超 GQA。在 131K 上下文、batch=16 的生产规模下：

• absorbed MLA 比同等 cache 占用的 GQA(kv=8) 慢约 27%；
• MQA 由于 cache 再小 8 倍，仍是长上下文下速度最快的方案，因为此时瓶颈已转移到带宽而非算力。

不过 absorbed MLA 相比 GQA 仍保留两点潜在优势：

• latent 表示是 per-head 的，保留了 GQA 通过硬共享丢失的 head 特定结构；
• r 是一个连续可调的旋钮，比 GQA 离散的 group_size 粒度更细，便于在 cache 与质量之间做更精细的权衡。

一个值得注意的省略

作者在 absorbed 版本的 benchmark 中没有对 K 应用 RoPE，原文也提示读者关注「the catch」对结果的影响。实际生产实现中，RoPE 与 MLA 的兼容方式会直接影响这一优化能否原样落地，是否能完全复现文中的 13–30 倍加速比仍需结合具体框架与算子实测验证。