超越MLA！新架构MLRA百万token，解码最高2.8倍速 | ICLR'26

新智元报道

随着大语言模型（LLM）在检索增强生成（RAG）、多步思维链（CoT）及超长对话等任务中广泛应用，解码阶段需处理的token数量大幅增加。其核心瓶颈在于：自回归生成受显存带宽限制（Memory-bound），而非算力限制（Compute-bound）。

每个解码步均需将完整Key-Value（KV）缓存从片外存储（如HBM）加载至片内缓存（如SRAM），数据搬运主导推理延迟，导致GPU利用率低下。

在标准多头注意力（MHA）机制下，KV缓存规模随头数、头维度与序列长度线性增长；当上下文超过10万token时，即成为显著性能瓶颈。

为缓解该问题，业界已提出多种方案：

• 分组查询注意力（GQA）：通过多组查询头共享KV，成比例减少缓存量；
• 多查询注意力（MQA）：所有查询头共享单组KV，但常以牺牲模型质量为代价；
• 多头潜在注意力（MLA）：由DeepSeek-V2提出，将KV压缩至单一低维“潜在头”，在保障质量的同时大幅节省缓存。

MLA通过低秩映射将隐藏状态投射为一个潜在向量，并仅缓存该向量；解码时将Key的上投影矩阵吸收进Query中，避免显式生成全部KV，从而降低显存压力。

MLA的两大局限

尽管高效，MLA存在两个关键工程瓶颈：

• 张量并行受限：因采用单一潜在向量结构，KV缓存无法切分，导致SGLang等框架只能使用数据并行（DP），造成模型权重冗余，无法实现高效张量并行（TP）解码；
• 可扩展性受限：矩阵吸收后，其解码逻辑等效于超高维MQA（如576维）。受限于GPU片上资源（如SRAM），难以进一步扩大潜在维度，致使FlashMLA等高性能Kernel目前仅适配NVIDIA Hopper架构（H100/H200）。

MLRA：多头低秩注意力

宾夕法尼亚州立大学、康涅狄格大学、卡内基梅隆大学、加州大学洛杉矶分校联合提出MLRA（Multi-Head Low-Rank Attention），从代数视角重构计算流程：将原本的大矩阵运算拆解为四个独立子块并行运算。

块分解视角（The Block Decomposition）

将原4倍维度的KV潜在向量逻辑划分为四个相等子块，对应权重矩阵亦垂直切分为四块；物理含义是：当前Key/Value为四个子块投影结果的累加和。

求和顺序外移（Moving the Sum Outside）

这是MLRA的核心改进：

• MLA逻辑：先对四个子块分别投影并求和，再执行Softmax注意力；
• MLRA逻辑：利用线性可加性，对每个子块独立完成投影与注意力计算，最后汇总四路输出。

该设计带来三大工程优势：

1. 完美支持4路张量并行：每张GPU仅处理一个子块，四卡完全解耦、并行无等待；
2. 显存占用降至约1/4：单卡只需缓存对应子块维度，大幅缓解KV Cache压力；
3. 计算流水线优化：将“求和”移出注意力核心，消除同步等待，显著提升大规模推理吞吐量。

实验结果

研究人员在Llama-3架构下，基于FineWeb-Edu 983亿token，从零训练了29亿参数（2.9B）模型，对比MHA、MQA、GQA、MLA、MFA、TPA、GLA-2/4、GTA等主流注意力变体。所有模型通过调整FFN中间维度确保参数量一致。

验证集困惑度（Validation Perplexity）

在Wikipedia、C4、Pile、RefinedWeb、Cosmopedia、FineWeb、FineWeb-Edu共7个数据集上评估：

• MLRA-4取得最低平均困惑度13.672，优于MLA（13.727）及其他所有基线；
• 在6/7个数据集上排名第一（Pile除外）；
• MLRA-4持续优于MLRA-2，表明分支数量增加有利于模型性能；
• 关键结论：4路张量并行下，MLRA-4将单设备KV缓存读取量降至1.5d_h，仅为MLA的1/3，同时实现更优模型质量。

常识推理能力

在ARC-E/C、OpenBookQA、BoolQ、HellaSwag、Winogrande、PIQA共7项零样本常识推理任务中，MLRA-4同样取得最高平均准确率，与困惑度结果高度一致。

解码速度（Decoding Speed）

在单块NVIDIA H100 80GB GPU上，测试128K–2M上下文长度的单序列解码延迟（MLRA-4基于FlashAttention-3，MLA采用官方FlashMLA）：

• MLRA-4在所有长度下均领先全部基线；
• 相比GQA加速1.05–1.26倍，且加速比随序列增长而提升；
• 相比MLA稳定加速2.8倍，验证了减少单设备KV读取可直接转化为更高解码速度。

解码吞吐量（Decoding Throughput）

在8块H100 GPU上评估批量解码吞吐量（隐藏层维度7168，参考DeepSeek-V3设置）：MLA（DP=8）、GLA-2（TP=2/DP=4）、MLRA-4（TP=4/DP=2）、GQA（TP=8）：

• MLRA-4在所有序列长度下均达成最高吞吐；
• 短序列中受益于更少的Q/K/V参数量；
• 长序列中因无需重复加载KV缓存，4路TP显著提升显存效率——即使GQA启用TP=8，仍落后于MLRA-4。

可扩展性与开源支持

MLRA具备优异可扩展性：MLA增大潜在头维度易导致Kernel部署困难；而MLRA的多分支低秩结构可在固定参数/激活预算下支持更多头数，同时保持对张量并行的友好性与内核执行效率。

研究人员已开源完整训练代码、预训练权重、数据集及基于FlashAttention-3的高性能解码Kernel，便于复现与工业部署。