快速解码的硬件高效注意力机制 Hardware-Efficient Attention for Fast Decoding

Hardware-Efficient Attention for Fast Decoding

快速解码的硬件高效注意力机制

https://arxiv.org/pdf/2505.21487v1

摘要
大型批次和长上下文下的LLM解码受到从高带宽内存加载键值（KV）缓存的限制，这增加了每个token的延迟，而解码的串行特性也限制了并行性。我们分析了计算强度、并行化与模型质量之间的相互关系，并质疑当前架构是否充分挖掘了现代硬件的潜力。本工作重新设计注意力机制，以在每次从内存加载数据时执行更多计算，从而最大化硬件效率，同时不牺牲并行可扩展性。我们首先提出分组绑定注意力（Grouped-Tied Attention, GTA），这是一种简单的变体，通过合并并重用键和值状态，在不损害模型质量的前提下减少内存传输。接着我们引入分组潜在注意力（Grouped Latent Attention, GLA），这是一种对并行友好的潜在注意力机制，并结合底层优化以实现快速解码，同时保持高模型质量。实验表明，GTA在模型质量上与分组查询注意力（GQA）相当，但使用的KV缓存大约仅为后者的一半；而GLA在质量上与多头潜在注意力（MLA）相当，且更易于分片。我们的优化版GLA内核在推测解码场景中，当查询长度超过1时，速度比FlashMLA快达2倍。此外，由于每设备获取的KV缓存更小，GLA在在线服务基准测试中将端到端延迟降低，并将吞吐量提高最多2倍。

1 引言
鉴于测试时计算（OpenAI, 2024）的重要性，推理效率如今正推动人工智能领域的进步，要求更加注重面向推理的架构设计。逐token解码的串行特性限制了并行化的机会。在解码过程中，多头注意力（MHA）（Vaswani等, 2017）会缓存所有先前token的键值（KV）状态。这些缓存状态随批大小和序列长度线性增长，并迅速耗尽高带宽内存（HBM）。此外，从片外内存中获取如此庞大的KV缓存主导了执行时间，显著超过了每个解码步骤中矩阵-向量运算所涉及的少量计算。随着KV缓存的增长，内存访问增加了延迟，导致GPU利用率长期处于较低水平（He 和 Zhai, 2024）。这一瓶颈阻碍了多种应用场景的发展：（i）对延迟敏感的交互式应用；（ii）用于多步推理的大批量LLM智能体（Yu等, 2024b）；（iii）测试时计算扩展（OpenAI, 2024；Snell等, 2024）；（iv）高吞吐量批量推理；以及（v）长上下文视频建模（Wu等, 2024）。总体而言，这些问题凸显了对注意力机制进行硬件高效重构的迫切需求。理想的注意力机制应满足三个条件：（1）实现高模型质量；（2）能够跨多个设备高效扩展；（3）在推理阶段有效利用现代硬件。

面向推理的注意力变体可加速解码过程。多查询注意力（MQA）（Shazeer, 2019）仅缓存一个KV头，显著减少了内存使用量。分组查询注意力（GQA）（Ainslie等, 2023）提供了一种折中方案，通过在较小的查询头组之间共享KV头来提升质量。最近由DeepSeek（DeepSeek-AI, 2024, 2025）提出的多头潜在注意力（MLA），通过低秩因子分解投影将隐藏状态压缩为一个联合潜在向量，缓存单个高维潜在头，然后在计算注意力前将其上投影恢复。如上所述，这些变体通过减少KV缓存，缓解了MHA在解码期间因数据移动过多而导致的内存受限性能问题（Ivanov等, 2021；Ootomo 和 Yokota, 2023；Gholami等, 2024）。遗憾的是，现代加速器的计算能力增长速度已超过内存带宽的增长速度（Ghose等, 2018）。计算强度（Williams等, 2009），即算术操作数与内存访问字节数之比（每字节浮点运算数，FLOPs/byte），常用于分析工作负载是内存受限还是计算受限。MQA通过在所有查询头之间复用单个KV头，减少了KV缓存的内存占用，从而缩短了从高带宽内存重新加载的时间，同时保持FLOPs不变，因此提高了计算强度，但牺牲了模型质量和并行性（Pope等, 2022）。GQA适度提升了计算强度（与组大小成正比），并在推理过程中具备良好的可扩展性。然而，在中等张量并行度下，每块GPU仍需存储较大的KV缓存，且当组尺寸过大时，模型质量会下降。另一种方式是，DeepSeek在其低秩MLA投影矩阵解码过程中进行吸收，直接使用单个潜在头计算注意力，使计算强度相对于MQA有效翻倍。然而，MLA本质上会在所有设备上复制该潜在头，限制了并行推理能力。

在本研究中，我们从计算强度的角度重新设计高效的硬件注意力变体，重点关注解码阶段的可扩展性，同时保持模型质量。我们将组大小gq定义为每个独立KV头对应的查询头数量；该组大小在很大程度上决定了计算强度。增大 gq可提升计算强度，并成比例地缩小KV缓存。然而，一旦超过某个阈值，进一步增加 gq就会在更高的每字节内存访问操作数（即更高的GPU利用率）与并行化之间产生权衡，迫使投影权重和KV缓存在设备间重复，从而削弱并行可扩展性。基于这些设计原则，我们提出了两种注意力变体，它们兼具高计算强度和跨设备高效扩展能力，并辅以底层优化。为了将这些设计选择付诸实践，我们的内核通过异步软件流水和warp特化实现计算与内存的重叠，采用协作式偏移计算器处理分页KV，从而始终保持张量核满载运行，推动内核从内存受限向计算受限转变。

• 我们首先探索分组绑定注意力（GTA），该方法将键和值表示绑定为一个共享状态，供少量查询头组使用。GTA可在保持相同组大小的情况下，将KV缓存大小减少最多一半，并使计算强度相对于对应的GQA提升最多2倍，同时保持模型质量和并行性。

• 我们提出分组潜在注意力（GLA），这是潜在注意力的一种对并行友好的扩展形式，并受益于底层优化。GLA在模型质量上与MLA相似，在推测解码等场景中（当查询序列长度为2或更大时）速度可快达2倍。在在线服务基准测试中，GLA将端到端延迟降低，并将token吞吐量提高最多2倍。

• 我们在FineWeb-Edu数据集上进行了中等规模的语言建模实验，验证了这些变体的有效性。在一个XL模型（1.47B参数）中，GTA实现了10.12的困惑度（GQA为10.20），GLA达到10.21困惑度和60.0%的平均下游准确率（MLA为10.25困惑度和59.1%准确率）。在一个大模型（876M参数）中，GTA实现11.2困惑度和57.6%平均下游准确率（优于GQA的11.3和56.9%）。对于一个中等模型（433M参数），GLA取得55.4%的平均下游准确率，略高于MLA的54.9%。

• 我们将上述算法设计与一系列底层优化相结合，开发出的注意力解码内核比FlashMLA 1（Li, 2025）快1.22倍。我们以宽松的开源许可证发布这些优化内核，以造福研究人员和从业者。代码地址为：https://github.com/Dao-AILab/grouped-latent-attention

2 预备知识 2.1 面向推理的注意力机制

MQA（Shazeer, 2019）通过减少MHA中的KV缓存大小来加速解码。它缓存一个由所有查询头共享的单一KV头。然而，这种激进的共享方式牺牲了模型质量；在分布式推理中，若按头维度划分工作负载，则每个设备必须复制该单一KV头以确保其可访问性，从而抵消了内存节省的效果（Pope等, 2022）。GQA（Ainslie等, 2023）通过将查询头分组并共享一个独立的KV头来解决这些问题。这种方法在分布式推理中消除了大部分GPU间的重复，且不增加额外内存开销。因此，它相对于标准MHA减少了KV缓存大小，同时在模型质量上优于MQA。为了保持模型质量，GQA需要使用适中的组数量；因此，在仅使用少量设备（例如张量并行度为2）时，对于长序列或大批次，每GPU的KV缓存可能迅速超出HBM容量。只有在使用大量设备（例如八路分割）时，每设备的KV缓存才会显著缩小。

2.2 分布式推理

串行解码将注意力组件的并行化机会限制在头轴上。张量并行（TP）将注意力层在设备之间进行划分，并频繁同步激活值（例如all-gather操作）。这一开销可通过高速GPU互连（如NVLink）（NVIDIA Corporation, 2024）缓解。TP避免了像数据并行（DP）（Li等, 2020）那样在每块GPU上复制整个模型，后者对于大模型而言变得不切实际。它还规避了流水线并行（PP）（Narayanan等, 2021）中存在的问题——即在逐token解码过程中部分GPU处于空闲状态，因为PP将模型划分为多个顺序阶段。在解码阶段，张量并行更受青睐，因为它将权重分布在多个GPU上，从而减轻了在逐个处理token时加载KV缓存的瓶颈（Su等, 2025）。

2.3 硬件特性

GPU架构概述：现代GPU由大量计算单元（例如浮点运算单元）组成，这些单元被组织成流式多处理器（SM），并配有内存层次结构。除了常规浮点单元外，NVIDIA GPU还包含专用于低精度矩阵乘法的张量核心（Tensor Cores）。内存层次结构包括高带宽内存（HBM）和片上SRAM（通常称为共享内存）。例如，NVIDIA H100 GPU拥有80 GB的HBM3内存，带宽为3.35 TB/s，芯片上共132个SM，每个SM配备256 KB SRAM，使得芯片内SRAM总带宽达到约33 TB/s（NVIDIA, 2022）。

执行模型：GPU运行内核（kernels），每个内核由大量线程组成，这些线程被分组为线程块（thread blocks），并调度到SM上执行。在线程块内部，线程被组织为每组32个线程的warp，它们可以通过快速的shuffle指令或通过片上共享内存共享数据。每个内核从HBM加载数据到寄存器和片上SRAM（即共享内存），执行必要的计算，并将结果写回HBM。

性能特性：根据计算与内存访问之间的平衡关系，操作可分为计算受限（compute-bound）或内存受限（memory-bound），其衡量指标为计算强度（arithmetic intensity）（Williams等, 2009），定义为每访问一byte内存所执行的算术操作数。当一个操作的执行时间主要由算术运算主导，而内存访问时间相对较短时，该操作为计算受限。相反，如果内存访问时间主导了其执行时间，则该工作负载为内存受限。

3 方法论

我们通过聚焦于计算强度，阐述了设计硬件高效注意力变体的设计视角。随后，我们展示了如何通过绑定键和值状态，并对缓存的潜在表示进行分片，从而在最大化计算强度的同时，实现跨设备的高效并行化。

3.1 解码中的计算强度：一种面向硬件效率的视角

在串行解码过程中，训练或预填充阶段使用的大型GEMM（通用矩阵乘法）工作负载转变为较小批量的GEMV（通用矩阵-向量乘法）。缓存的键矩阵中每个加载的BF16（2字节）元素，会与已位于寄存器中的单token查询元素执行一次MAC运算（2 FLOPs），从而产生1:1的FLOP-to-byte比率。这一计算强度远低于Nvidia Hopper H100 SXM GPU（NVIDIA, 2022）的密集BF16理论峰值性能，即约每字节295 FLOPs（989 TFLOPs / 3.35 TB/s），导致张量核心严重利用率不足。本质上，在分布式推理中，整体延迟受限于以下三者中最慢的一个：在GPU峰值计算能力下完成所有FLOP所需的时间、在GPU峰值内存带宽下传输必要数据所需的时间，或在可用互连带宽下的设备间通信延迟（Pope等, 2022；Austin等, 2025）。实际上，在适度的张量并行度（TP）下（例如八路分片），反复加载庞大的KV缓存主导了解码延迟（Recasens等, 2025）。

如表1所示，标准MHA的计算强度约为1，在解码期间其GPU利用率可能低至7%（Recasens等, 2025）。理论上，在不增加延迟的前提下，它可以承受约两个半数量级更多的FLOPs。然而，由于内核开销、重叠有限以及其他低效因素的存在，实际加速效果自然会更小。这种低利用率反映了近年来GPU代际之间计算吞吐量与内存带宽之间的差距日益扩大。现代GPU在计算上投入的硅片面积远多于内存带宽，这一趋势自2017年起愈发明显（见附录B.7，图15（右））。硬件FLOPs每两年增长约3倍，而HBM内存带宽在同一时期仅增长约1.6倍（Gholami等, 2024）。为了缓解带宽限制并在解码期间提高计算强度，实践者可以通过融合内核、数据分块以最大化片上数据复用、选择性地重新计算小型中间结果而非存储它们，或采用特殊的注意力重排序技术来减少片外数据传输（Dao等, 2022），而无需改变注意力机制本身的设计。

3.2 最大化计算强度的设计策略

GQA 是一种通用的注意力形式，它在每一组查询头之间复用一个已加载的 KV 头。由于共享 KV 头并不会减少运算数量（每个查询头仍需计算其注意力分数），但确实减少了内存读取，因此计算强度与组大小 gq（即每个独立 KV 头对应的查询头数量）成正比地增加（见表 1）。例如，在 MQA 的情况下，所有查询头共享一个 KV 头，其计算强度近似等于查询头的数量 hq。

DeepSeek 提出了 MLA，其设计面向推理，通过吸收低秩矩阵避免显式生成 KV 状态，从而通过三个因素实现高计算强度：缓存单个潜在头、对键和值复用该潜在表示，以及使用大量查询头。最初，MLA 加载一个单一的潜在头，并在所有查询头之间复用，类似于 MQA。现在，加载到片上内存中的确切潜在表示在计算注意力时同时服务于键和值状态，与最激进的 MQA 设计相比，有效使计算强度翻倍（见图 1 左侧，展示 MLA 内存加载示意图）。尽管 DeepSeek 的论文（DeepSeek-AI, 2024）未明确讨论这一点，但使用低秩投影（减少参数量）来缩小潜在缓存，也为增加上投影矩阵的列维度以恢复丢失的参数提供了可能。

3.3 硬件高效且可并行的注意力机制 3.3.1 分组绑定注意力（Grouped-Tied Attention, GTA）

奇异值图显示，键缓存呈现出急剧衰减的特性，几乎所有的方差都集中在少数几个主方向上，因此键位于一个低秩子空间中，具有高度冗余性（Saxena等, 2024）。在应用RoPE之前，这一效应更为显著，此时键坍缩到更小的子空间中（Yu等, 2024a；Sun等, 2024）。与此同时，多项研究表明，仅对头维度的一部分进行RoPE旋转即可保持精度，因此对整个头宽度进行旋转带来的质量提升微乎其微（Black等, 2022；Barbero等, 2025）。如果键本质上是低秩的，并且每个头只需部分通道进行旋转以实现位置编码，那么对每个通道都进行旋转并缓存完整秩的键张量就会浪费内存。相反，我们可以仅对头维度中需要位置信息的部分进行旋转；其余未旋转的通道通常处于低秩子空间且冗余，可以与值状态共享或绑定。

GQA 已通过让多个查询头共享一个独立的 KV 头，减少了 KV 缓存和内存传输，并能在多个设备上高效扩展。基于 GQA 的分组设计，并结合上述关于低秩性和部分 RoPE 的洞察，我们提出分组绑定注意力（GTA），该方法将分组、键值绑定为单一状态以及部分应用 RoPE 三者统一起来，旨在在保持模型质量的同时减少 KV 缓存大小。

在 GTA 中，与 GQA 类似，每组查询头共享一个独立的 KV 头。GTA 更进一步：将键和值的投影参数绑定，生成一个单一状态，称为“绑定的 KV”（tied KV），其形状与单个键或值向量相同（见图2，展示 GTA 与 GQA 架构差异的示意图）。值路径使用该绑定 KV 状态的全部维度，而键路径仅使用其前半部分作为未旋转的部分。键的剩余 RoPE 成分来自隐藏状态的另一个单头投影，该投影在所有组之间广播，并与未旋转的前半部分拼接，形成完整的键向量。实证消融实验表明，即使在后续将该共享部分用于值路径前将其旋转逆变换回来，对共享部分应用 RoPE 仍会降低模型质量，因此绑定部分从不进行旋转。经过这些步骤后，查询、键和值状态定义如下：

通过绑定KV状态，我们将单一状态加载到片上内存中，并将其同时用于键和值，且在少量查询头之间共享。这种复用减少了内存传输，使计算强度大致翻倍，并相对于具有相同组数的GQA将其KV缓存占用减少约一半。GQA-4表示四个独立的键头和值头，而GTA-4表示四个绑定的KV头。实验表明，其困惑度（见5.1.1节）以及在下游任务中的性能（见5.1.2节）与对应的GQA相当。

3.3.2 分组潜在注意力（Grouped Latent Attention, GLA）

4 系统优化：异步与分布式偏移计算

我们描述了系统优化，以实现在现代硬件（如 H100 GPU）上对 MLA、GTA 和 GLA 的峰值性能。由于非常快速的专用矩阵乘法单元，如张量核心（第 2.3 节），我们需要仔细的软件流水线化内存加载和张量核心指令，以始终保持张量核心忙碌。这是通过利用现代硬件上的异步性的 warp 专业化技术实现的。

4.1 通过软件流水线化和 warp 专业化实现异步

我们使用两种技术来重叠计算和内存加载：

1. 软件流水线化：我们在当前 KV 块用于计算时加载下一个 KV 块。这种经典技术（Lam，1988）避免了张量核心等待内存加载。

2. Warp 专业化：我们有单独的 warp 执行内存加载，使用 TMA（张量内存加速器）或异步复制（cp.async 指令），以及单独的 warp 执行矩阵乘累加（MMA）。前者作为生产者 warp，后者作为消费者 warp（Bauer 等人，2014）。这在矩阵乘法（Thakkar 等人，2023）和注意力（Shah 等人，2024a）中常用。这种解耦简化了软件流水线化，允许 warp 调度器重叠内存加载和计算。

4.2 分布式偏移计算用于分页 KVA

新的注意力变体，如 MLA、GTA 和 GLA，对计算和内存子系统都造成了压力，必须尽可能快地执行内存加载。分页 KV（Kwon 等人，2023）已成为存储 KV 缓存的标准方式。然而，分页 KV 使得使用 TMA 变得困难，TMA 是一种执行地址计算和边界检查以加载连续内存块的专用硬件单元。相反，我们使用异步复制指令（cp.async），其中每个线程分别发出单独的加载指令。挑战在于地址计算出奇地昂贵，因为它需要 64 位整数索引，这转化为每个整数乘法需要多个指令。我们反而让同一个 warp 中的多个线程合作计算地址。例如，对于头维度 128，要从全局内存加载一个 128 x 128 大小的块到共享内存，我们使用 128 个线程，每行 16 个线程加载：

在解码过程中，如表1所示，GLA的计算强度约为 2⋅gq（是GQA的两倍）；GLA-2的计算强度达到约 hqFLOPs/byte（每字节内存访问对应的浮点运算数），与MQA在其最激进共享设计下的计算强度相近，但GLA-2具有更好的模型质量。参见图1（右），展示了GLA-2在计算注意力时的内存加载示意图。GLA保持了高计算强度（见图3，屋顶线分析），具备高效的并行能力（见图4（右）），并提供高模型质量（见第5.1节）。我们在第5.2节对GLA的token吞吐量进行了基准测试；详细的延迟和吞吐量基准数据见附录B.6。

5 实验与结果

我们通过实验验证了我们的两种简单方法——GTA 和 GLA——（1）在模型质量上可与 GQA 和 MLA 相媲美，（2）易于并行化，（3）在现代硬件（如 H100 GPU）上运行高效。例如，GLA 在上游和下游任务中的质量与 MLA 相当，但更易于分片，且我们的 GLA 内核在推测解码设置下速度比 DeepSeek 的 FlashMLA 快达 2 倍。GLA 表示查询潜在头被分片的配置，该配置消除了其下投影矩阵在设备间的重复，减少了每设备的参数数量；我们主要将其作为消融实验纳入研究，因为查询潜在头并不被缓存。

实验设置：我们在 FineWeb-Edu-100B 数据集（Lozhkov 等, 2024）上训练了四种规模的模型：小型（1.83 亿参数）、中型（4.33 亿参数）、大型（8.76 亿参数）和 XL 型（14.71 亿参数），采用 GPT-3 模型（Brown 等, 2020）的配置和 Llama 3 架构（Grattafiori 等, 2024）。小型模型在 250 亿个 token 上进行训练，而中型、大型和 XL 型模型各自在 500 亿个 token 上进行训练。有关架构超参数和训练设置的更多细节见附录 B.1。

5.1 模型质量 5.1.1 验证困惑度

除了在 FineWeb-Edu 验证集的 1 亿个 token 上报告验证困惑度外，我们还在另外四个数据集上评估了困惑度：RedPajama v1（Weber 等, 2024）中的 Wikipedia 和 C4（Colossal Clean Crawled Corpus）部分、Cosmopedia（Ben Allal 等, 2024）以及 Pile（Gao 等, 2020），每个数据集均使用 1 亿个 token 进行评估。我们在表 2 和表 5 中报告了这五个数据集上的平均困惑度。各数据集的详细验证困惑度结果见附录 B.2.1。对于大型模型中的 MLA 和 GLA，默认情况下我们使用旋转位置编码维度 dR=32；dR=48时的验证困惑度结果见附录 B.2.1。

GLA-2 在中等和大规模模型上通常优于 MLA，在 FineWeb-Edu 数据集和五个数据集的平均困惑度上均实现了更低的验证困惑度，而在小模型上表现与 MLA 相当。值得注意的是，GTA 解决了 GQA 的关键局限性，并在中等和大规模模型上进一步降低了困惑度。在大型模型中，GLA-2 与 GLA-2 表现相当，平均困惑度均为所有变体中最低（24.49–24.51），优于次优的 MLA（24.93）。GTA 解决了 GQA 的一个关键限制，即在较低张量并行度（TP）设置下无法有效减少每设备的 KV 缓存。对于大型模型，GLA-2 和 GLA-2 平均表现相似，相较于其他变体展现出强大的性能。这一趋势在 XL 规模（14.7 亿参数）模型上依然成立：GTA-4 在困惑度上略优于 GQA-4（FineWeb-Edu 上为 10.12 对比 10.20），而 GLA 在困惑度上持续优于 MLA（例如，FineWeb-Edu 上为 10.21 对比 10.25；见表 5）。

5.1.2 下游任务评估

我们在标准基准上评估了零样本性能：SciQ（Welbl 等, 2017）、OpenBookQA（Mihaylov 等, 2018）、ARC-Easy 子集（Yadav 等, 2019）、HellaSwag（Zellers 等, 2019）、PIQA（Bisk 等, 2020）、WinoGrande（Sakaguchi 等, 2020）和 MMLU（Hendrycks 等, 2021）。

在各项下游基准测试中，我们提出的注意力变体保持或超过了基线准确率。对于4.33亿参数的中等模型（见表3），GLA-2取得了最高的平均准确率55.4%，略高于MLA的54.9%。在8.76亿参数的大模型上（见表4），GLA-2的平均准确率为57.5%，仅比GTA-4低0.1个百分点，且与GQA-4基本持平，表明分组或绑定操作并未降低模型质量。这一趋势在14.71亿参数的XL模型上依然存在：GLA-2的平均准确率达到60.0%，而MLA为59.1%；GTA-4和GQA-4均达到60.2%（见表5）。这些结果证实，我们的硬件高效变体在从中等规模到XL规模的模型上，均能保持甚至提升下游任务的性能。有关下游评估的更多细节，见附录B.2.2。有关小规模和中等规模模型的消融实验，见附录B.3。

5.2 并行化

GLA 通过划分潜在头在多个 GPU 上实现扩展，减少内存通信，从而加快解码速度。而采用张量并行（TP）的 MLA 则在每块 GPU 上复制其单个高维潜在头；为了控制这一开销，早期系统退而采用张量并行与数据并行（DP）结合的混合方案，将不同的批次序列分配给不同的 GPU，这种方案在大批量或高并发请求场景下具有一定优势。我们使用 DeepSeek Coder V2 Base 模型（2360亿参数，210亿激活参数）进行验证，该模型被量化为 FP8，并通过 SGLang 框架（Zheng 等, 2024b）结合我们的 FlashAttention 3 内核进行部署以进行基准测试。实验比较了在八块 H100 GPU 上纯张量并行（TP）与混合 TP 加两路或四路数据并行（DP）的性能。在混合设置中，仅注意力子模块在 DP 组间被复制；其输出在进入 MoE 前馈层前通过 all-gather 汇聚，并重新分发，以缓解 MLA 中 KV 缓存的重复问题。在相同并行配置下，采用零冗余分片策略的 GLA 配置（潜在头在 TP 秩之间均匀分布）由于每设备加载的潜在 KV 缓存更小，性能优于 MLA。

在混合序列负载下，若有四个并发序列，其中一个非常长，另外三个较短，MLA 的混合设置可能会因设备等待长序列完成而停滞。如图5（右）所示，在预填充长度均匀采样至最多131K个token的情况下，GLA-8的吞吐量约为MLA（TP=2, DP=4）的2.5倍。GLA-4（TP=4, DP=2）也优于MLA在（TP=2, DP=4）下的表现，并且由于其更低的DP秩，对工作负载不平衡具有更强的容忍能力。图5（左）显示，在16个并发请求、每个请求的预填充长度固定为32K和64K token、解码长度固定为4K token的条件下，采用纯TP=8的GLA-8在吞吐量上高于采用混合并行的MLA。更普遍地，无论在纯TP还是在更小DP秩的混合并行配置下（例如，GLA-4（TP=4, DP=2）对比 MLA（TP=2, DP=4）），GLA均能达到或超过MLA的性能，证实了GLA对工作负载不平衡具有潜在的鲁棒性。有关更详细的端到端延迟、吞吐量结果及基准测试设置，请参见附录B.6。

5.3 速度

我们对 MLA（1 个维度为 512 的潜在头，RoPE 维度为 64）的解码内核进行了基准测试，结果如图 4（左）所示；对 GLA（2 个潜在头，每个维度为 256，RoPE 维度为 64）的测试结果如附录 B.7 中图 15（左）所示，测试中使用了分页 KV（page size 为 64）。我们将我们的 GLA 实现与目前能找到的最优化的 MLA 实现——DeepSeek 的 FlashMLA（Li, 2025）——进行对比。在标准解码设置下（查询长度为 1），我们的 GLA 内核比 FlashMLA 快约 20%；在推测解码设置下（查询长度为 2），速度超过 FlashMLA 2 倍以上。我们的 GLA 内核在 H100 GPU 上达到了最高内存带宽的 93%，以及最高 TFLOPS 的 70%，充分最大化了内存和计算子系统的利用率。

6 讨论

我们讨论相关工作、局限性以及一些未来方向。关于相关工作，附录A回顾了在预训练期间减少KV缓存的注意力机制，综述了更广泛的硬件高效设计，并扩展至涵盖用于加速LLM解码的后处理策略、部分RoPE应用的研究以及低秩投影方法，这些见解共同塑造了我们的方法。

GTA通过为每个头缓存一个绑定状态，将与相同数量KV头的GQA相比的未分片KV缓存大小减少了约一半。当TP=8且有8个KV头时，GTA-8每token存储1.5ℎ，其中ℎ/2来自独立的RoPE，而GQA-8每token存储2ℎ，仅节省25%。然而，当有8个KV头且TP=4时，差距扩大：GTA-8每token分配2.5ℎ，而GQA-8使用4ℎ，因此中等TP度数可带来更大的内存缓解。我们的GLA实验使用两个潜在头以匹配MLA的未分片KV缓存大小，当TP≥2时，GLA将每设备的KV缓存减半，而MLA在设备间复制KV缓存。在高达14.71亿参数的模型规模上，模型质量保持相当。下一步值得评估更大规模的模型，使用更多潜在头，同时将头维度固定为2ℎ，以进一步验证GLA的性能。例如，Llama 4系列模型（高达4000亿参数）采用GQA-8，其中ℎ=40，ℎ=8，在TP=8时每设备每token的KV缓存大小为2ℎ，超过此度数则会复制缓存。匹配的GLA-8若使用8个潜在头，则每设备每token缓存略多，为2.5ℎ（其中ℎ/2来自解耦的RoPE）；在大致相同的缓存预算下，GLA-8是否能在质量上优于GQA-8仍是一个开放问题。对于一半的KV缓存，基于我们在14.71亿参数规模上的实验，GLA-2在FineWeb-Edu上的验证困惑度为10.218，而GQA-4为10.202。此外，解耦RoPE带来的每token额外ℎ/2 KV缓存占用可通过仅在部分层应用RoPE来缓解，如Cohere（Yang等, 2025）和Llama 4（Meta AI, 2025）所示。这些扩展研究留待未来工作。

第3.2节表明，计算强度随查询头数量增加而提升。附录B.3提出了一个消融实验，用低秩版本替换完整的查询和输出投影以减少参数。通过增加列维度以在GQA和GTA中每组添加更多查询头来补偿减少的容量。该消融旨在提高质量（或在KV缓存进一步减少时保持质量），同时提升计算强度，有效提高GPU利用率。在少数秩和查询头上的消融实验中，验证困惑度比基线略差0.1至0.2。对此权衡的详细研究留待未来工作。最后，探索并行化和计算强度在其他架构（如Mamba（Gu和Dao, 2024；Dao和Gu, 2024）和线性注意力（Yang等, 2024））中的相互作用，可能进一步挖掘现代硬件的潜力。

7 结论

我们证明，在解码时关注高计算强度和有效并行化可产生硬件高效的注意力机制。我们提出分组绑定注意力（GTA），通过绑定键值（KV）状态，将KV缓存需求减少约一半，并将计算强度相对于相同组数的GQA提高一倍。GTA实现了比GQA基线更低的困惑度和更好的平均下游准确率。我们随后提出分组潜在注意力（GLA），这是一种可并行化的面向推理的注意力变体，具有高计算强度，在解码速度上比基线DeepSeek FlashMLA（Li, 2025）快达2倍。GLA在所有评估的模型规模和基准上持续匹配或超过MLA的准确率。此外，通过在设备间分片潜在头，且每设备的头维度小于MLA，GLA在解码期间获取更小的KV缓存，从而在在线服务基准测试中将端到端延迟降低，并将吞吐量提高最多2倍。因此，GTA可高效替代GQA（仅需一半的KV缓存内存）。由于其在GPU间的可扩展分片和更快的解码速度，GLA是MLA的实用替代方案。

A 相关工作 A.1 注意力变体 A.1.1 算法层面

预训练阶段：对 DeepSeek 的 MLA 的后续工作包括多矩阵分解注意力（Multi-matrix Factorization Attention, MFA）（Hu 等, 2025），其结构类似于 MQA，但使用更大的头维度和因子化的查询投影。然而，它与 MQA 有类似的局限性，特别是由于在设备间复制导致每设备的 KV 缓存利用效率低下，且与张量并行不兼容。张量积注意力（Tensor Product Attention, TPA）（Zhang 等, 2025）使用每状态两个秩为 R的投影矩阵对查询、键和值进行因子化，因此每个 token 的缓存大小为 (RK+RV)(h+dh)个元素。我们在附录 B.3 中报告了 TPA 的验证困惑度消融实验结果。TPA 的低秩结构支持直接的头级别张量并行分片，但其 KV 缓存随 R线性增长，一旦 R≥4就已超过 MLA；而为了保证质量，通常需要四或更高的秩，尤其是在更大模型规模下，因此它很快会失去内存节省的优势。

后处理适配：Lin 等（2025）提出 SIGMA，该方法在微调阶段使用一种新颖的差分缩放 QKV 模块，专门优化以通过压缩键（K）、轻微压缩值（V）并增强查询（Q）来提升推理效率。Slim Attention（Graef 和 Wasielewski, 2025）是一种训练后方法，仅保留键向量并在运行时动态重建值，从而将任意多头注意力的上下文内存减少一半。我们提出的方法比这两种方法进一步减少了 KV 缓存。与上述技术不同，我们的方法试图在预训练阶段重构注意力架构，而不是修改已预训练完成的模型的现有注意力结构。不过，类似的知识蒸馏方法已被应用于 MLA（Meng 等, 2025；Ji 等, 2025），这些方法也可在训练后阶段用于 GLA，以实现低 KV 缓存占用和易于并行化的优势。

A.1.2 系统层面

FlashAttention（Dao 等, 2022）通过一种 I/O 感知的分块策略重新组织注意力计算，将数据保留在高速内存中，避免显式生成整个注意力矩阵。它显著降低了内存开销，并在长序列解码时带来显著加速。FlashAttention-2（Dao, 2023）进一步优化注意力内核，减少非矩阵乘法操作并改进并行任务划分，提升了硬件利用率。其系统级改进使吞吐量相比原始 FlashAttention 显著提高。随后，FlashAttention-3（Shah 等, 2024b）利用下一代 GPU 特性，如异步内存操作和低精度计算，将注意力效率推向硬件极限。NSA（Natively trainable sparse attention，原生可训练稀疏注意力）（Yuan 等, 2025）引入了与硬件对齐的动态分层稀疏注意力模式，该设计在保持接近全注意力精度的同时降低了计算复杂度，支持极长序列的高效解码。我们提出的方法与这些系统级注意力优化正交，可与其结合使用。

A.2 加速解码的其他方法

以下后处理适配设计可与 GLA 和 GTA 的架构设计协同工作，增强其性能表现。

算法层面：已有大量算法层面的努力，例如 token 驱逐（Zhang 等, 2023；Xiao 等, 2024）、在相邻层之间共享 KV 缓存（Brandon 等, 2024）、批处理以提升 GPU 利用率（Mukherjee 等, 2023），以及推测解码（Xia 等, 2023）。

系统层面：在系统方面，已有量化（Hooper 等, 2024）、CPU 卸载（Aminabadi 等, 2022；Sheng 等, 2023；He 和 Zhai, 2024）、以及使用 PagedAttention（Kwon 等, 2023）进行内存管理以缓解内存碎片问题的研究。

硬件层面：此外，硬件方面也有利于推理的进展，例如支持 FP4（NVIDIA, 2024）、NVLink（NVIDIA Corporation, 2024），以及专为快速推理设计的硬件芯片（Groq, 2024）。

A.3 低秩投影

实证研究表明，在应用 RoPE 之前，键的激活值具有急剧衰减的奇异值谱（Yu 等, 2024a；Chen 等, 2024），意味着许多维度贡献极小。此外，Singhania 等（2024）表明，键在不同模型中表现出显著降低的内在维度，暗示其本质上处于低秩空间。例如，Kobayashi 等（2024）发现，权重衰减的正则化会驱动键-查询联合映射降至更低的秩。相比之下，值的激活表现出混合的低秩趋势。一些研究显示，若不对缓存的值进行压缩，否则会导致严重精度损失（Chang 等, 2024；Singhania 等, 2024）。然而，也有研究发现，部分压缩可在可接受的性能下降范围内实现（Saxena 等, 2024；Sun 等, 2024）。这些不一致表明，值的有效秩依赖于具体模型和方法。

GTA 利用了上述研究的洞察，在每个查询组内绑定键和值状态，从而减少 KV 缓存大小。

A.4 旋转位置编码（RoPE）

按头维度（Barbero 等, 2025）指出，可能无需对每个头维度都应用 RoPE（Su 等, 2023），因为最高频率分量已能提供足够的位置区分能力。其他研究（DeepSeek-AI, 2024；Black 等, 2022；Wang 和 Komatsuzaki, 2021）发现，仅对一半维度应用 RoPE 仍可保持模型质量。受此启发，GTA 仅对头维度的一部分进行旋转，其余部分与值头维度的一半绑定，从而在保持精度的同时减少 KV 缓存大小。

按层应用（Chen 和 Yan, 2024）表明，RoPE 在 Transformer 的早期层中贡献最大，因为这些层的注意力集中在局部句法关系上。相比之下，深层更倾向于语义线索，意味着在模型堆叠的后期，位置旋转的收益递减。（Yang 等, 2025）提出 RNoPE，该设计将 RoPE 层与无位置编码（NoPE）层交错，并将 RoPE 限制在滑动窗口内，在极长上下文长度下显著提升了检索能力，并影响了 Llama 4（Meta AI, 2025）的架构选择。总体而言，对部分层应用 RoPE 对 GLA 有益，因为对于解耦的 RoPE（单头且维度为头维度的一半），可以考虑将其移除。

B 完整实验结果 B.1 实验设置

B.3 消融实验

不同的注意力变体减少了每层的可学习参数和表示能力；因此，这些节省下来的参数需要在模型其他部分重新分配。例如，Llama 2（Touvron 等, 2023）在其消融实验中通过增加前馈网络（FFN）的宽度，使 MQA 和 GQA 与 MHA 进行公平比较。此外，MQA 最初提出时也建议增加宽度以匹配 MHA 的参数量（Shazeer, 2019）。同时，DeepSeek-AI（2024）通过调整模型深度，增加层数以实现公平比较。改变模型深度的做法相对较少，因为这使得直接对比变得困难，尤其是在中等规模模型中，参数匹配的灵活性较低。（Pope 等, 2022）通过缩小 MHA 的头维度来匹配 MQA 和 TPA 的参数量，而（Zhang 等, 2025）则通过增加查询头的数量来对齐参数总数，本质上是将节省的参数重新分配到查询投影中。

例如，在 GQA 和 GTA 中，KV 头的数量必须能整除查询头数量，因此在调整查询头数量以尽可能接近基线进行公平比较时，灵活性较低。而对于 MLA 和 GLA，增加查询头数量是有益的，因为在解码过程中我们不会显式生成 KV 状态。本质上，单个潜在头会在所有或一组查询头之间共享；因此，增加查询头数量可以简单地提升解码时的 GPU 利用率，正如我们之前在表 1 中所展示的那样：MLA 和 GLA 的计算强度最终取决于查询头的数量。

B.3.1 消融实验：小型模型（1.88 亿参数）

B.5 速度

我们在此展示，我们的技术（分布式偏移计算）在使用分页 KV（paged KV）时显著加快了注意力内核的运行速度。通常情况下，当页面大小较小时，注意力内核速度会变慢，因为地址计算的开销更大（Kwon 等, 2023）。然而，更小的页面大小可以减少内存碎片，并解锁新的应用场景，例如前缀缓存（prefix caching）（Zheng 等, 2024b）。我们在图6中展示了使用分页 KV 时 GLA（2 个潜在头，每个维度为 256，RoPE 维度为 64）解码内核的速度基准测试结果。我们比较了页面大小为 1 和 64 的情况，分别在启用和不启用分布式偏移计算下的性能表现。启用分布式偏移计算后，页面大小为 1 时不会出现性能下降，其速度与页面大小为 64 时相当。相反，在不启用分布式偏移计算的情况下，页面大小为 1 的速度比页面大小为 64 慢 1.3 倍。我们观察到，对于页面大小为 64 的情况，分布式偏移计算带来了 1.2 倍的加速；而对于页面大小为 1，则实现了 1.5 倍的加速。

B.6 端到端延迟与吞吐量

服务部署环境设置：我们使用 SGLang —— 一个面向生产的推理服务框架，并在其实时服务模式下运行所有实验，确保 HTTP 解析、动态请求队列和 GPU 内核调用的耗时被统一计入（Zheng 等, 2024a）。在此模式下进行评估，能够暴露离线测试中无法体现的排队开销和网络延迟，从而真实反映 GLA 和 MLA 在实际部署约束下的行为表现。负载生成器发送 1280 个提示（prompts），并设置并发请求数上限，以控制同时处于活跃状态的请求数量。服务器会动态将这些活跃请求组合成小批次，因此该上限影响的是负载压力，而非固定批大小。我们使用 DeepSeek-Coder-V2 Base 模型的预训练权重（总参数 2360 亿，激活参数 210 亿），量化至 FP8，并使用我们的 FlashAttention3 内核进行服务。在基准测试中，我们将页面大小设为 64。为了模拟 GLA，我们将 MLA 的潜在维度重构为 GLA 结构，并随机初始化权重，因为在本阶段我们仅评估性能而非准确率。我们还采用了分块预填充（chunked-prefills）（Agrawal 等, 2023），每块长度为 8192 个 token，并逐块运行预填充内核。解码批次独立形成，因此默认情况下预填充 token 与解码 token 不会混合。

并行化与评估指标：每个 Transformer 块在八块 GPU 上采用张量并行（tensor parallelism, TP）进行分片，而 MoE 前馈层进一步通过专家并行（expert parallelism）进行划分。我们也测试了数据并行（DP）与张量并行结合的混合配置；每当启用数据并行时，仅注意力子模块在数据并行组间复制，其输出在进入 MoE 前馈层前通过 all-gather 汇聚，然后重新分发，以缓解 MLA 的 KV 缓存复制问题。我们测试了广泛的推理工作负载，以评估在相同并行配置下，以及 GLA 仅使用张量并行而 MLA 使用张量与数据并行混合配置时，GLA 与 MLA 的表现。我们报告四项服务级指标：端到端（E2E）延迟、首 token 时间（TTFT）、token 间延迟（ITL）和输出吞吐量。图中所有数值均以中位数（median）汇总，因其对大规模交互系统中的长尾行为更不敏感。此外，表格中还提供了均值（mean）数据。

B.6.1 张量并行：GLA 与 MLA 对比

在此配置中，使用 8 块 H100 GPU，张量并行度为 8。GLA-8 采用 8 个潜在头，每个 token 需缓存维度为 256 的潜在表示；而 MLA 则在每台设备上复制一个 512 维的潜在缓存。两种方法均将 RoPE 维度设为 64 并进行解耦处理。图 7 和表 27 显示，GLA-8 在所有负载水平下均表现出稳定优势。当并发请求数为 16 时，GLA-8 将中位数端到端延迟从 136 秒降低至 117 秒，减少了约 15%，同时 token 吞吐量提升约 17%。当并发数上升至 64 时，GLA-8 的完成时间为 179 秒，而 MLA 为 381 秒，延迟降低 53%；首 token 到达时间从约 3 分钟缩短至 12 秒，吞吐量提升约 70%，达到每秒 1461 个 token。即使在 128 个并发请求下，GLA-8 仍能将延迟降低约 24%，并保持近 60% 的吞吐量优势。这些优势源于 GLA-8 每设备更小的 KV 缓存占用，减少了内存通信量，允许更多活跃请求驻留于 GPU，从而缩短了计算开始前的等待时间。

B.6.2 数据并行 + 张量并行：GLA 与 MLA 对比

图8和表29展示了在引入数据并行注意力后，计算能力与内存通信之间的平衡如何变化。在混合TP=2 + DP=4的配置下，如图8和表29所示，GLA-2（两个潜在头，每个维度为256）在16个并发请求的轻负载下，将中位数端到端延迟从137秒缩短至120秒，吞吐量从每秒477个token提升至每秒544个token。当并发数达到64时，优势进一步扩大：延迟降低约16%（196秒 vs. 166秒），吞吐量提高19%，达到每秒1584个token（相比MLA的1334个token）。

类似地，在混合TP=4 + DP=2的配置下，如图9和表31所示，GLA-4在各种并发请求数下始终优于MLA。

然而，当并发请求数达到128时，如图10和图11所示，采用TP=2与DP=4混合配置的MLA，通过额外的副本将批处理分散到更多设备上，从而充分利用所有计算单元，其性能超过了采用纯TP=8的GLA-8：MLA每秒输出约2122个token，比GLA-8的1363个token高出约56%；完成时间约为247秒，比GLA-8的433秒提前约43%。性能反转的原因在于，当服务器负载较重时，数据并行带来的额外计算通路弥补了其KV缓存复制带来的开销。相比之下，纯TP的GLA-8已达到内存带宽上限，无法进一步扩展。这表明，数据并行仅在高并发场景下才具有优势。

B.6.3 数据并行：工作负载不平衡

“随机比例”（random ratio）参数表示基准测试中的随机请求生成器可为任意单个预填充或解码序列分配的最小长度所占的比例。例如，当随机比例为 0.125 且序列长度为 4096 个 token 时，每个请求的长度将从 512 到 4096 的整数范围内均匀采样，从而在保证每个批次有统一下限的同时，保留真实场景中序列长度的分布差异。该随机比例应用于预填充序列长度（131K）和解码序列长度（4K）。本节实验展示了在批次内序列长度不一的情况下出现的工作负载不平衡问题，这可能导致部分 GPU 处于空闲状态。

在图13和表35中，我们展示了在预填充长度为131K、解码长度相对较长（4K）的场景下，当批次内序列长度均匀采样时，采用张量并行度为8的GLA-8，其吞吐量约为在四路数据并行、张量并行度为2的混合配置下MLA的2.7倍。因为在每个数据并行组中的每一次NCCL通信集合操作都必须等待所有设备参与，只要有一个副本仍在处理一个极长的序列，就会迫使其他所有副本及其张量并行分片一同等待，导致整体吞吐量退化为最慢那个“拖尾”（straggler）设备的速度。而在纯TP-8配置中，不存在额外的数据并行同步屏障，只有持有权重的八个分片之间需要相互等待；某个长序列只会拖慢其所在的分片组，而集群其余部分仍可继续工作，从而保持高得多的GPU利用率。

B.6.4 延迟敏感型工作负载

在延迟敏感型工作负载中，主要目标是最小化端到端响应时间，特别是首 token 时间（time to first token），以满足严格的服务级别目标（SLO），而非最大化总吞吐量。由于较大的批处理会增加排队和预填充延迟，因此延迟敏感的服务通常保持非常小的批大小，以牺牲吞吐量为代价来实现更快的响应。在表38中，采用纯八路张量并行的 GLA-8 相比于使用张量并行与数据并行混合的 MLA，在必须缓解 KV 缓存复制问题的情况下，成功将端到端延迟降低了约2倍，并将首 token 时间缩短了近4倍。

B.6.5 解码密集型工作负载

在解码密集型工作负载中，生成的续写内容非常长，导致串行解码阶段占据主要的挂钟时间（wall-clock time），此时延迟和KV缓存的内存带宽成为主要瓶颈。由于模型大部分时间都在进行串行解码，批处理带来的收益极小。在图14中，预填充较短（256个token），解码过程长达32K个token，在八路并行度下对比GLA-8与MLA，GLA-8的吞吐量最高可达MLA的2.5倍。

B.6.6 小上下文与短对话

“小上下文”描述的是提示（prompt）和生成的续写内容都远小于模型上下文窗口的推理请求。例如，语音助手对一个简短问题做出单次简短回复。在表40中，在八路并行配置下，使用八个潜在头的GLA-8相比采用TP=2与四路数据并行（DP）混合配置的MLA具有更低的延迟。由于这是单一批次场景，四路DP中三个分组的GPU处于空闲状态，而GLA-8每层只需加载一半的KV缓存，因此性能优于MLA。

B.7 内核执行时间

我们在 H100 GPU 上对这两种设置下的注意力内核延迟进行了基准测试（忽略通信开销），结果如表 44 和表 45 所示。在这些设置下，采用 TP = 2 的 GLA 比采用 DP 的 MLA 快 1.3 至 1.5 倍。

原文链接https：//arxiv.org/pdf/2505.21487v1