随着智能体与长上下文工作负载日趋普遍，上下文长度不断增加，注意力机制在推理时间中所占比重也随之上升（图1）。由于注意力机制如今已成为主要性能瓶颈，模型的架构设计方式——而非仅仅是实现方式——越来越大程度地决定着推理性能。围绕GPU执行方式对模型架构进行优化，正是AI模型协同设计的核心理念。关于模型设计选择如何在不牺牲精度的前提下同时影响吞吐量与交互性，可参阅此前发布的文章《AI模型协同设计：面向硬件友好的大语言模型设计》。
本文重点分析组大小（每个KV头对应的查询头数量）、头维度以及序列长度如何影响稠密注意力（Dense Attention）的性能——在稠密注意力中，每个查询都会关注序列全长范围内的所有键与值。基于上述分析，结合注意力机制在GPU间的并行化方式，本文提炼出四条实用准则，形成一份协同设计清单，旨在帮助模型开发者在英伟达GPU上提升推理吞吐量与交互性。关于稀疏注意力的相关内容，敬请期待后续文章。
所有分析均基于两类来源：通过GEMM形状运算推导的解析公式，以及使用FP8精度（涵盖注意力计算与KV缓存）的预填充与解码内核实测数据。
预填充阶段与解码阶段的区别
预填充阶段并行处理完整的提示词，产生大型GEMM-M（= ISL × G）矩阵乘法，属于计算密集型操作。在不使用投机解码的情况下，解码阶段每次仅生成一个Token，产生较小的GEMM-M（= G）矩阵乘法，受限于从高带宽内存（HBM）读取KV缓存，属于内存密集型操作。
投机解码可增大GEMM-M，并使解码阶段向计算密集型转变。由于预填充与解码在查询长度、KV访问方式和性能瓶颈上存在差异（表2），本文对每个参数分别按阶段进行分析。
注意：在智能体应用和多轮对话应用中，前缀缓存（Prefix Caching）十分常见。新一轮对话的ISL可能较短，但需要关注较大的前缀缓存。当ISL较短但前缀缓存较长时，预填充的行为模式与解码相近。
屋顶线模型与算术强度
屋顶线模型（Roofline Model）以算力上限和带宽上限约束GPU性能，算术强度（Arithmetic Intensity）决定哪个因素构成瓶颈（公式1）：
算术强度 = 总浮点运算量 / 总访存字节数
脊点（Ridge Point）标志着从内存密集型到计算密集型的转变边界。预填充阶段远高于该点，属于计算密集型；解码阶段低于该点，属于内存密集型（图2）。投机解码可提升解码阶段的算术强度，使其向脊点靠近。
FlashAttention的工作原理
FlashAttention无需实体化完整的注意力矩阵即可完成注意力计算。它将Q、K、V的分块从HBM流式传输至片上SRAM，并将三个步骤融合为单次计算过程：
第一步，批量矩阵乘法（BMM1）计算查询与键的相关性得分；第二步，在线Softmax利用滑动最大值与求和对得分进行归一化；第三步，第二次批量矩阵乘法（BMM2）对值进行加权聚合。
其中，批量矩阵乘法在Tensor Core上执行，Softmax指数运算则在特殊功能单元上运行。后续算术强度分析以这两个矩阵乘法的形状为基础展开。
组大小（G）的影响
组大小（G）是共享同一KV头的查询头数量。多头注意力（MHA）的G = 1，分组查询注意力（GQA）的G = 4、8、16……，多查询注意力（MQA）的G = QH。
算术强度关于G的函数分析如下（以下公式中"字节"指HBM移动的字节数，为简化起见，假设每个元素占1字节，即FP8 KV缓存）：
随着G增大，1/G项趋近于零，算术强度趋近于2 × ISL。当ISL = 32K时，将G从8提升至16，算术强度提升幅度不足6%。换言之，预填充性能主要由ISL决定，而非G。图4也证实了这一点：将G从1（MHA）变化至64（MQA），预填充运行时间变化幅度不足1%。公式2、3、4如下：
浮点运算量 = 4 × PB × QH × ISL? × Hsz（与G无关）
访存字节数 = 2 × PB × KH × Hsz × ISL × (G + 1)
算术强度 = 2 × G × ISL / (G + 1) = 2 × ISL / (1 + 1/G) → 当G趋近于无穷大时趋近于2 × ISL
G每翻倍，解码算术强度也随之翻倍。将G从1提升至8可获得8倍增益，原因在于减少了内存访问量并提升了GPU计算利用率。该效应与KVSL无关：算术强度始终接近2 × G，因此除非G非常大，否则解码阶段仍维持内存密集型。英伟达Nemotron 3等模型采用具有两个KV头的GQA设计，从而提升了解码效率。公式5、6、7如下：
浮点运算量 = 4 × DB × QH × KVSL × Hsz（与G无关）
访存字节数 = 2 × DB × KH × Hsz × (G + KVSL)
算术强度 = 2 × G × KVSL / (G + KVSL) ≈ 2 × G（当KVSL远大于G时）
图4显示，G每翻倍，解码运行时间约下降2倍，原因是KV头数量减半意味着每个Token加载的数据量减半。当G超过16时，KVSL = 32K的曲线趋于平缓，此时两类开销占据主导：固定的启动和后处理开销，以及通过将KV分布到多个流处理器（SM）以实现并行处理的Flash解码归约开销。KVSL = 128K的内核能更好地分摊这些开销，并持续保持2倍的趋势。
注意：投机解码将有效GEMM-M提升至 (1 + D) × G，其中D为草稿Token数量。当M足够大以填满计算分块时，解码阶段将向计算密集型转变。
准则一：为解码效率选择G值，并尽量将其提高。预填充运行时间对G不敏感，而解码阶段的算术强度约为2 × G，因此更高的G可提升解码速度和GPU利用率。投机解码是在给定G下进一步提升性能的另一有效手段。
头维度（Hsz）的影响
与组大小不同，头维度（Hsz）不影响算术强度。Hsz翻倍会同时使浮点运算量（公式2和5）和访存字节数（公式3和6）翻倍，二者之比保持不变。然而，图5显示运行时间随Hsz增大而上升，这是因为注意力内核包含三类工作，它们随Hsz的缩放规律各不相同。
矩阵乘法：随Hsz增长，但以对齐步长为单位递增。此前的文章建议模型维度选取128的倍数，以与GPU分块尺寸和缓存行宽度对齐。部分填充的分块与完整分块的代价相同，因此Hsz = 64实际上需要支付128的代价。Hsz ≥ 512时将接近张量内存（TMEM）容量上限，因此128和256是最优选择。
内存（KV状态）：同样随Hsz增长。由于GPU以128字节为单位移动数据，内存访问与矩阵乘法类似，当Hsz为128的倍数时效率最高。
Softmax则有所不同：其开销与Hsz无关，因为Softmax作用于注意力得分矩阵（查询Token × 键），该矩阵不含头维度。公式8和9如下：
Softmax运算量（预填充）≈ PB × QH × ISL?
Softmax运算量（解码）≈ DB × QH × KVSL
上述三类开销的相对权重决定了不同阶段中Hsz的实际代价（图5）。
预填充阶段属于计算密集型（矩阵乘法加Softmax）。随着Hsz增大，矩阵乘法的浮点运算量增加，而Softmax保持固定。若预填充为纯矩阵乘法，Hsz翻倍则运行时间翻倍；但由于固定的Softmax开销不随Hsz缩放，实际运行时间的增幅小于Hsz的增幅。图5证实了这一点：预填充时间随Hsz增大而上升，但增速低于Hsz的增长率。较大的Hsz能够分摊Softmax开销，使内核更多时间用于矩阵乘法，从而降低预填充阶段对Softmax的依赖程度。
解码阶段属于内存密集型（流式读取KV缓存）。更大的Hsz会增加每个Token的KV字节数（公式6），因此运行时间应与Hsz呈线性关系。图5对此予以确认，但存在轻微的次线性现象，原因是启动、后处理和Flash解码归约开销不随Hsz缩放，在KVSL = 32K这类较短内核中权重更为显著。
准则二：使用Hsz = 128或256。Hsz不改变算术强度，但必须与硬件对齐。Hsz = 64通常仍需支付128宽分块的代价，而Hsz ≥ 512则接近TMEM容量上限。因此，128和256是最优选择区间。
序列长度的影响
序列长度（ISL / KVSL）对预填充和解码的影响方式不同，因为它通过不同变量进入各阶段：预填充并行处理所有ISL输入Token，而每个解码步骤读取长度为KVSL的KV缓存。因此，两者的缩放规律存在差异（图6）。
预填充的ISL?缩放特性
预填充的计算量随ISL?增长（每个Token关注所有其他Token），而KV访问量仅与ISL成正比（公式2、3）。因此，算术强度随ISL线性增长，使预填充始终处于脊点以上，维持计算密集型特征。ISL翻倍时，运行时间大约增长至原来的四倍（图6）。在ISL较短时，缩放比例低于4倍，原因是固定的启动和后处理开销占主导；当ISL足够大以分摊这些开销时，平方缩放规律才得以体现。
解码的KVSL线性缩放特性
每个解码步骤需要读取完整的KV缓存以生成一个Token，因此访存字节数随KVSL增长，而每步计算量保持较小（公式5和6）。算术强度维持在约2 × G，远低于脊点，因此解码在所有长度下均维持内存密集型特征。KVSL翻倍时，运行时间应翻倍，图6对此予以确认。在KVSL较短时，缩放比例低于2倍，原因是启动、后处理和Flash解码归约开销不随KVSL缩放，其占比随KVSL增大而逐渐收缩。
准则三：尽量压缩有效KV状态。使用场景决定序列长度，但代价是不对称的：预填充随ISL?增长，而解码随KVSL线性增长。可通过KV缓存压缩、稀疏注意力或滑动窗口注意力、或混合模型架构（如Nemotron 3）等方式压缩有效KV状态——在混合架构中，仅部分层承载持续增长的全局KV状态。
张量并行与KV头数量
张量并行（TP）将注意力头分布到多个GPU上，每个GPU分配QH/TP个查询头和KH/TP个KV头，分片对象是头而非Token。在表3中，仅包含KH/TP的批次维度缩小，单GPU的GEMM形状和算术强度保持不变。
TP存在实际限制：KV头数量必须能被GPU数量整除。一旦TP > KH，某个组的查询头将跨越多个设备，每个设备都需要一份共享KV头的副本，这会造成KV状态的重复，增加内存和带宽开销而无任何收益（图7）。因此，应保持TP ≤ KH，确保每个GPU至少拥有一个完整的组：即一个KV头及其对应的G个查询头。
对于KV头较少的模型（例如Nemotron 3仅有两个KV头），TP的扩展空间很快耗尽，因为KV缓存无法在不引入重复的情况下进一步分片至每个GPU少于一个KV头。此时，注意力机制需要采用不同的扩展方式：注意力数据并行（ADP）对请求进行分片，KV并行（KVP）将长序列KV缓存分布到多个GPU；前馈网络（FFN）则单独使用专家并行（EP）进行扩展。
TensorRT-LLM将上述方式组合为宽EP（Wide EP，注意力使用ADP加FFN使用EP）和螺旋并行（Helix Parallelism，注意力使用KVP加FFN使用EP）。在这两种情况下，KH均决定高效扩展的策略选择。
准则四：根据KV头数量（KH）制定并行策略。保持TP ≤ KH，确保每个GPU拥有完整的KV头。KV头较少的模型（MQA为1个，或GQA为2个）很快耗尽TP空间，更适合采用ADP或KVP处理注意力，配合EP处理MoE前馈网络（在TensorRT-LLM中分别对应宽EP和螺旋并行）。
模型协同设计检查清单
将以下四条准则作为模型设计检查清单，开始进行AI模型注意力机制的协同设计。这些设计选择可在相同硬件条件下提升GPU利用率、推理速度、吞吐量与交互性。
准则一：为解码效率选择组大小（G）并尽量提高。预填充对组大小不敏感。
准则二：使用头维度（Hsz）= 128或256，以与GPU分块和128字节传输对齐，同时保持在TMEM预算范围内。更大的头维度也有助于在预填充阶段隐藏Softmax开销。
准则三：通过KV缓存压缩、稀疏注意力或滑动窗口注意力、或混合模型架构来压缩有效KV状态。
准则四：将并行策略与KV头数量（KH）匹配。保持TP ≤ KH，对KV头较少的模型采用宽EP和螺旋并行进行扩展。
致谢
本文是英伟达跨团队协作的成果。感谢Timmy Liu、Jatin Mitra、Tiyasa Mitra、Bhargava Gopireddy、Brian Pharris、Julien Demouth和Eduardo Alvarez的大力支持。
Q&A
Q1：GQA（分组查询注意力）相比MHA（多头注意力）在推理效率上有哪些优势？
A：GQA通过让多个查询头共享同一个KV头（即增大组大小G），在解码阶段可显著减少KV缓存的内存读取量。G从1增大到8可带来约8倍的算术强度提升，解码运行时间随G翻倍而约减少一半。英伟达Nemotron 3等模型正是采用了GQA设计来提升解码效率。而在预填充阶段，G的变化对运行时间的影响不足1%，因此提高G对解码有显著收益，对预填充几乎无负面影响。
Q2：为什么注意力机制的头维度（Hsz）推荐选择128或256，而不是其他数值？
A：头维度（Hsz）必须与GPU硬件特性对齐。GPU以128字节为单位进行数据传输，矩阵乘法也以固定分块尺寸执行，若Hsz未对齐（如Hsz = 64），实际上仍需支付128宽分块的计算代价。Hsz ≥ 512则会接近张量内存（TMEM）的容量上限，影响性能稳定性。此外，较大的Hsz能分摊固定的Softmax开销，使预填充阶段更加高效。综合以上因素，128和256是兼顾硬件对齐与性能的最优选择。
Q3：TensorRT-LLM中的宽EP（Wide EP）和螺旋并行（Helix Parallelism）分别是什么，适用于哪类模型？
A：宽EP（Wide EP）将注意力数据并行（ADP）与前馈网络的专家并行（EP）结合；螺旋并行（Helix Parallelism）则将KV并行（KVP）与专家并行（EP）结合。两者都适用于KV头数量较少的模型，例如仅有1个KV头的MQA模型或只有2个KV头的GQA模型。这类模型的张量并行（TP）扩展空间有限，一旦TP超过KV头数量便会产生KV状态重复的问题。宽EP和螺旋并行可在不触发该问题的前提下，实现对注意力和MoE前馈网络的高效扩展。
