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

MLRA通过拆分KV缓存为四个并行分支，显著降低显存占用并实现4路张量并行。推理速度比MLA最高快2.8倍，支持百万级上下文，且模型质量更优。无需牺牲性能，即可高效扩展长文本处理能力。

随着大语言模型（LLM）越来越多地应用于长文本任务——如检索增强生成（RAG）、多步思维链（CoT）推理以及超长对话——在每个解码步骤中必须处理的token数量显著增加。

其核心问题在于：自回归生成（Autoregressive Generation）受限于显存带宽（Memory-bound），而非计算能力（Compute-bound）。

在每一个解码步中，模型必须将整个Key-Value (KV) 缓存从片外存储（如HBM）重新加载到片内存储（如SRAM）中，这种数据搬运主导了推理延迟，从而导致解码过程总 GPU 利用率低下。

在标准的多头注意力（MHA）机制下，KV缓存的大小随注意力头数、头维度以及序列长度线性增长。当上下文长度超过10万（100K）个token时，KV缓存会迅速演变成一个严重的性能瓶颈。

为了减少KV缓存的开销，业界已经尝试了多种方法：

• 分组查询注意力（GQA）：通过在多组查询头之间共享Key和Value，成比例地减少了缓存量。

• 多查询注意力（MQA）：将这一思路推向极端——让所有头共享同一组KV对，但这往往会损害模型质量。

• 多头潜在注意力（MLA）：由DeepSeek-V2近期提出。它将整个KV缓存压缩进一个低维的「潜在头」（Latent Head）中，在保持模型质量的同时，实现了惊人的缓存节省效果。

MLA（Multi-Head Latent Attention）通过低秩压缩将隐藏状态映射至单一的潜在向量（Latent Head），并在推理时仅需缓存该潜在向量，极大地降低了KV Cache的显存占用。

在解码阶段，通过将Key的上投影矩阵（Up-projection Matrix）吸收（Absorb）到Query中，可以避免显式生成（Materialize）所有注意力头的完整Keys和Values

然而，MLA的设计存在两个关键缺陷：

张量并行（TP）受阻：由于采用单一潜在向量架构，其KV Cache无法在多个计算设备间进行切分，迫使SGLang等开源推理框架只能采用数据并行（DP）模式处理MLA解码，导致模型权重在各设备间冗余存储，无法实现高效的张量并行解码。

可扩展性与架构适配受限：矩阵吸收后，MLA的解码在算子层面等效于一个超大维度（如576维）的MQA（Multi-Query Attention）。

由于GPU片上资源（如SRAM）难以支撑极大的单头维度，导致其难以进一步扩展Latent Dimension，也是目前FlashMLA等高性能kernel主要局限于NVIDIA Hopper架构（H100/H200）的原因。

宾夕法尼亚州立大学、康涅狄格大学、卡内基梅隆大学、加利福尼亚大学洛杉矶分校的研究人员提出的MLRA (Multi-Head Low-Rank Attention) 提供了一个巧妙的代数视角：将原本的一个大矩阵运算，拆解为四个独立的小块运算。

论文: https://arxiv.org/pdf/2603.02188

博客: https://SongtaoLiu0823.github.io/mlra

代码: https://github.com/SongtaoLiu0823/MLRA

数据与权重: https://huggingface.co/Soughing/MLRA

块分解视角 (The Block Decomposition)

研究人员将原本4倍维度的 KV 潜在向量逻辑上划分为四个相等的子块。与其对应的权重矩阵也相应地垂直切割为四个块。

物理含义：现在的Key和Value不再被视为一个不可分割的整体，而是四个子块投影结果的累加。

求和顺序外移 (Moving the Sum Outside)

这是MLRA最关键的改进。

• MLA的逻辑： 先对四个子块进行投影并求和，生成一个完整的 Key/Value，再做 Softmax 注意力计算。

• MLRA的逻辑： 既然求和是线性的，直接对每一个子块独立进行投影和注意力计算，再将这四个独立分支产生的输出结果进行求和。

在工程上实现了彻底的解耦：

1. 完美适配4路张量并行 (4-way TP)： 每一张显卡（设备）现在只需要处理其中一个子块。由于计算是在分支末端才聚合，四张显卡可以完全并行工作，互不干扰。
2. 极低的显存占用： 每张卡只需要缓存对应子块的维度（约为原来的1/4），显著降低了单卡的KV Cache压力。
3. 计算流水线优化： 通过将「求和」操作移出注意力核心，消除了计算过程中的等待时间，极大提升了大规模模型推理的吞吐量。

实验结果

研究人员将MLRA与一系列全面的基准模型进行了对比，包括MHA、MQA、GQA、MLA、MFA10、TPA11、GLA-2、GLA-4以及GTA12。所有模型均在Llama-3架构下，使用来自FineWeb-Edu的983亿（98.3B）token从零开始训练，参数规模均为29亿（2.9B）。

为确保公平比较，通过调整FFN的中间层维度，使所有模型的参数量保持一致。

验证集困惑度 (Validation Perplexity)

研究人员评估了在Wikipedia、C4、Pile、RefinedWeb、Cosmopedia、FineWeb和FineWeb-Edu这七个数据集上的困惑度。

主要发现如下：

• 性能最优： MLRA-4在所有七个数据集上取得了最佳的平均困惑度（13.672），超越了包括MLA（13.727）在内的所有基准模型。

• 广泛领先： 在七个数据集中的六个（除Pile外），MLRA-4均排名第一。

• 分支优势： MLRA-4在所有数据集上的表现始终优于MLRA-2，这证明了增加分支数量对模型性能有益。

• 关键结论： 这是一个显著的结果——在4路张量并行（TP）下，MLRA-4将单设备的KV缓存读取量降低至1.5d_h，仅为MLA的三分之一，同时还实现了更优的模型质量。

常识推理能力

研究人员在七个常识推理基准测试（ARC-E/C、OpenBookQA、BoolQ、HellaSwag、Winogrande和PIQA）上进行了零样本（Zero-shot）性能评估。结果与困惑度的测试发现高度一致：

MLRA-4在所有对比的注意力变体中，取得了最高的平均零样本准确率。

解码速度 (Decoding Speed)

研究人员在单块NVIDIA H100 80GB GPU上，针对128K到2M不等的上下文长度测试了单序列解码延迟。MLRA-4基于FlashAttention-3实现，而MLA则使用官方的FlashMLA内核。

• 持续领先： 在所有上下文长度下，MLRA-4的表现均优于所有基准模型。

• 加速效果： 相比GQA，MLRA-4实现了1.05倍至1.26倍的加速，且加速比随序列长度增长而增加。

• 对比MLA： 相比MLA，MLRA-4稳定保持了2.8倍的加速，证实了在4路张量并行下，减少单设备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在所有序列长度下均实现了最高的解码吞吐量。

• 短文本优势： 在注意力计算前的预计算占主导的短序列中，MLRA-4得益于更少的Q/K/V参数量。

• 长文本优势： 在长序列中，无需重复加载KV缓存的4路张量并行提供了巨大的吞吐优势。尽管GQA使用了TP=8，但在长文本下仍落后于MLRA-4，因为MLRA-4更小的单设备缓存实现了更高效的显存利用。

至关重要的一点是，MLRA具备极佳的可扩展性！

在MLA中，增加KV潜在头维度（latent-head dimension）往往会导致高性能解码kernel 难以部署；而在固定的激活值/参数预算下，单纯增加头数可能会降低模型质量。相比之下，MLRA的多分支低秩架构支持显著更多的头数，同时保持了对张量并行（TP）的友好性和内核执行效率。

研究人员发布了完整的训练代码，数据、预训练权重以及基于FlashAttention-3的高性能解码内核，以方便开发者复现和部署MLRA。

参考资料：

https://arxiv.org/pdf/2603.02188

文章来自于“新智元”，作者 “LRST”。