[ICLR 2025] MLRA：突破 MLA 并行瓶颈，长文本解码提速 280%
Multi-Head Low-Rank Attention
本文提出了 Multi-Head Low-Rank Attention (MLRA)，这是一种新型注意力机制，旨在解决 Multi-Head Latent Attention (MLA) 在分布式推理中无法进行张量并行（Tensor Parallelism）分片的痛点。MLRA 通过将潜在状态分解为多个独立的分支并进行并行计算，实现了原生 4 路张量并行支持，并在 2.9B 模型规模下达到了 SOTA 的 Perplexity 和推理速度。
TL;DR
在大语言模型（LLM）推理中，KV Cache 的加载速度是制约长文本生成的最大屏障。DeepSeek 提出的 MLA 虽好，但无法在分布式环境下进行张量并行（TP）分片。本文提出的 Multi-Head Low-Rank Attention (MLRA) 通过将潜在状态分解为 4 个独立分支，在保留 MLA 压缩优势的同时，解锁了极致的并行效率。在 2.9B 模型测试中，MLRA 不仅精度超越 MLA，解码速度更是提升了 2.8 倍。
1. 痛点：为什么 MLA 难以并行？
在长文本推理阶段，算力不再是瓶颈，内存带宽才是（Memory-bound）。虽然 MLA 通过 Latent Compression 将 KV Cache 压缩到了极小，但它存在一个“系统级”缺陷：
- 分片限制：MLA 的 Latent Head 是单体结构。在使用 Tensor Parallelism (TP) 时，传统的权重分片方法无法作用于这个单一的压缩头，导致每个 GPU 都必须载入一份完整的 KV Cache 复本。
- 冗余加载：当跨 GPU 推理时，这种冗余加载白白浪费了宝贵的 HBM 带宽，限制了推理吞吐量的进一步提升。
2. 核心机制：MLRA 的分支化直觉
MLRA 的核心思想极其优雅：既然一个大头不好分片，那就把它拆成 4 个独立的小头。
2.1 架构拆解
MLRA 将 KV 投影矩阵拆分为四个子块（Block），每个子块对应一个独立的分支：
- 独立投影：每个分支持有自己的低秩 Latent Head。
- 注意力求和：每个分支独立计算其 Attention Output。
- 结果归集：最后的注意力输出是各分支结果的加权和。
2.2 方差校准（Variance Calibration）
多分支求和会带来一个数学挑战：方差爆炸。作者通过数学推导发现，由于 MLA 引入了 NoPE（非位置编码）和 RoPE（旋转位置编码）的混合，两者的方差极易失衡。 针对 MLRA-4，作者应用了如下缩放校准： 这种设计确保了模型在训练过程中的数值稳定性（Initial Training Stability）。
3. 实验表现：SOTA 的精度与效率
3.1 精度碾压
在 2.9B 规模的 FineWeb-Edu 等多项数据集评估中，MLRA-4 展示了极强的语言建模能力。
| Method | Avg Perplexity | Avg Zero-shot Acc |
|---|---|---|
| GQA | 14.139 | 57.89% |
| MLA | 13.727 | 58.75% |
| MLRA-4 | 13.672 | 58.84% |
3.2 解码效率的飞跃
得益于对 4 路张量并行的原生支持，MLRA-4 显著降低了单位 Token 的 KV 加载量。对比 FlashMLA 优化后的 MLA，MLRA-4 在不同上下文长度下均保持了 2.8 倍左右的加速比。
4. 深度洞察：为什么 MLRA 是未来的方向？
- 硬件感知设计：MLRA 并不是单纯为了刷榜，而是从分布式系统设计的角度倒推出来的算法。它解决的是“算法架构与大规模集群互联”之间的摩擦。
- 更灵活的扩展性：相比于 GQA 需要 8 路 TP 才能将 KV 加载降低到 1.5，MLRA 仅用 4 路 TP 就达到了同样的资源效率。
- 门控机制（Gated Attention）的潜力：作者进一步证明，结合 Gated 机制后，MLRA-4 的平均 Perplexity 能进一步下探至 13.621，这预示着该机制与现代非线性激活层有极佳的兼容性。
5. 总结与局限
总结：MLRA 完美解决了 MLA 的并行化痼疾，它是目前处理长文本推理最平衡、最具有工程价值的 Attention 变体之一。
局限性：多分支设计虽然提升了并行度，但在单卡（Batch=1）推理时，由于算子切分的细碎化，可能需要专属优化的 Triton Kernel（作者已开源）才能完全释放其性能。
本文基于 arXiv 论文 《MULTI-HEAD LOW-RANK ATTENTION》 创作。