NVIDIA 发布 AI 模型注意力协同设计指南：四条准则提升长上下文推理效率

随着智能体和长上下文工作负载日益普遍，上下文长度不断增加，注意力在推理时间中占据的比例也越来越大（图 1）。既然注意力已成为主要开销，那么它如何设计——而不仅仅是如何实现——就越来越多地决定着模型的推理性能。围绕 GPU 的执行方式来塑造模型架构，正是 AI 模型协同设计（co-design）的前提。关于模型设计选择如何在不牺牲精度的前提下同时影响吞吐量和交互性，请参阅上一篇文章《AI 模型协同设计：硬件友好的 LLM 设计》。

本文探讨分组大小（每个 KV 头对应的查询头数）、头维度和序列长度如何影响稠密注意力（即每个查询都要关注序列长度范围内的所有键和值）的性能。我们把这些分析，连同注意力在多块 GPU 间的并行方式，提炼为四条实用准则：这是一份协同设计检查清单，可帮助模型开发者在 NVIDIA GPU 上提高推理吞吐量和交互性。后续文章将介绍稀疏注意力，敬请期待。

所有分析都基于两类来源：由 GEMM 形状运算得出的解析公式，以及在注意力计算和 KV 缓存均采用 FP8 的情况下，对预填充和解码内核的实测数据。

表 1. 本文公式中使用的符号

预填充和解码为何是两个不同的问题？

预填充并行处理整个提示词，产生 GEMM-M（= ISL × $G$）很大的矩阵乘法，属于计算受限。在不使用推测解码时，解码每次只生成一个词元，产生 GEMM-M（= $G$）很小的矩阵乘法，并因需从高带宽内存（HBM）读取 KV 缓存而成为内存受限。

推测解码会增大 GEMM-M，可使解码向计算受限转变。由于预填充和解码的查询长度、KV 访问方式和瓶颈各不相同（表 2），我们针对每个阶段分别分析各个参数。

表 2. 预填充与解码在查询长度、KV 上下文、注意力 GEMM-M 和主要瓶颈上的差异

注：在智能体和多轮对话应用中常见的前缀缓存（prefix caching）下，新一轮对话的 ISL 可能很短，却要关注很大的前缀缓存。当 ISL 很短而前缀缓存很长时，预填充的表现类似解码。

算术强度如何决定计算受限还是内存受限

屋顶线（roofline）模型用计算上限和带宽上限来界定 GPU 性能，如此前所述。算术强度决定哪一个起约束作用（公式 1）：

算术强度 = 总 FLOPs / 总访存字节数

拐点（ridge point）标志着从内存受限到计算受限的转变。预填充远高于拐点，属于计算受限；解码则低于拐点，属于内存受限（图 2）。推测解码会提高解码的算术强度，并可能使其向拐点靠近。

FlashAttention 内核如何在 GPU 上计算注意力？

FlashAttention 无需实例化完整的注意力矩阵即可计算注意力。它把 $Q$、$K$、$V$ 的分块从 HBM 流式载入片上 SRAM，并将三个步骤融合为一趟完成：

• 第一步，批量矩阵乘法（BMM1）用查询对键打分
• 第二步，在线 softmax 利用滚动的最大值和求和对分数做归一化
• 第三步，第二次批量矩阵乘法（BMM2）对值加权

BMM 在 Tensor Core 上运行，而 softmax 的指数运算在特殊函数单元上运行。BMM 的形状决定了下面的算术强度分析。

GEMM 形状

注意力性能取决于其两次矩阵乘法的形状。表 3 列出了 BMM1 和 BMM2 在各阶段的 (Batch, M, N, K) 维度。

对于解码，GEMM-M = $G$，通常为 8-16，远小于 GPU 的分块 tile-M（64 或 128），限制了每个分块内的并行工作量。更大的 $G$ 意味着每个词元加载的 KV 更少，并且每次加载可在更多查询头之间摊销，从而提高利用率。下一节将对这一效应进行量化。

分组大小

分组大小（$G$）是共享同一个 KV 头的查询头数。MHA 的 $G$ = 1，GQA 的 $G$ = 4、8、16……，MQA 的 $G$ = QH。

算术强度随 $G$ 的变化。在下列公式中，“Bytes”指在 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：

FLOPs = 4 × PB × QH × ISL² × Hsz（与 $G$ 无关）
Bytes = 2 × PB × KH × Hsz × ISL × ($G$ + 1)
算术强度 = 2 × $G$ × ISL / ($G$ + 1) = 2 × ISL / (1 + 1/$G$) → 2 × ISL（当 $G$ → ∞）

解码（GEMM-M = $G$）

$G$ 翻倍，解码算术强度也翻倍。把 $G$ 从 1 提高到 8，可通过减少内存流量、提高 GPU 计算利用率获得 8 倍收益。这与 KVSL 无关：算术强度始终接近 2 × $G$，因此除非 $G$ 非常大，解码仍然是内存受限。NVIDIA Nemotron 3 等模型采用了只有两个 KV 头的 GQA，使解码更高效。公式 5、6 和 7：

FLOPs = 4 × DB × QH × KVSL × Hsz（与 $G$ 无关）
Bytes = 2 × DB × KH × Hsz × ($G$ + KVSL)
算术强度 = 2 × $G$ × KVSL / ($G$ + KVSL) ≈ 2 × $G$（当 KVSL ≫ $G$）

图 4 显示，$G$ 每翻一倍，解码运行时间约下降一半，因为 KV 头数减半，每个词元加载的数据也随之减半。当 $G$ 超过 16 后，KVSL = 32K 的曲线趋于平缓。其单步内核规模很小，以下两项开销占据主导：固定的启动与后处理开销，以及为了在 KV 头很少时保持并行而将 KV 拆分到多个 SM 上所需的 flash-decoding 归约。更长的 KVSL = 128K 内核能更好地摊销这些开销，因此继续遵循 2 倍的趋势。

注：推测解码会把有效 GEMM-M 提高到 (1+$D$) × $G$，其中 $D$ 是草稿词元数。一旦大到足以填满计算分块，解码就会向计算受限转变。

准则 1：按解码效率选择 $G$，并尽量取高。预填充运行时间不随 $G$ 变化，而解码算术强度 ≈ 2 × $G$，因此更高的 $G$ 能提高解码速度和 GPU 利用率。在给定 $G$ 下，推测解码是另一个提升性能的手段。

头维度

与分组大小不同，头维度（Hsz）不影响算术强度。Hsz 翻倍会使 FLOPs（公式 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 无关，因为它作用于注意力分数矩阵（查询词元 × 键），而该矩阵没有头维度。公式 8 和 9：

Softmax 运算量（预填充）≈ PB × QH × ISL²
Softmax 运算量（解码）≈ DB × QH × KVSL

这几部分之间的平衡决定了 Hsz 在各阶段的代价（图 5）。

预填充是计算受限的（矩阵乘法加 softmax）。随着 Hsz 增大，矩阵乘法的 FLOPs 增加，而 softmax 保持不变。如果预填充是纯矩阵乘法，Hsz 翻倍会使运行时间翻倍；但固定的 softmax 不随之扩展，因此运行时间的增幅小于 Hsz 的增幅。图 5 证实了这一点：预填充随 Hsz 上升，但慢于 Hsz 的增长。更大的 Hsz 能摊销 softmax 的开销，把更多内核时间转移到矩阵乘法上，使预填充不那么受 softmax 制约。

解码是内存受限的（流式读取 KV 缓存）。更大的 Hsz 会增加每个词元的 KV 字节数（公式 6），因此运行时间应随 Hsz 扩展。图 5 证实了这一点，只是略呈次线性，因为启动、后处理和 flash-decoding 归约的开销不随 Hsz 扩展，在较短的 KVSL = 32K 内核中占比更大。

准则 2：使用 128 或 256 的 Hsz。Hsz 不改变算术强度，但必须与硬件对齐。通常，Hsz = 64 仍要付出 128 宽分块的代价，而 Hsz ≥ 512 则逼近 TMEM 的容量上限。这使 128 和 256 成为最佳选择。

序列长度

序列长度（ISL / KVSL）对预填充和解码的影响不同，因为它在两个阶段中通过不同的变量起作用：预填充同时处理全部 ISL 个输入词元，而每个解码步骤读取长度为 KVSL 的 KV 缓存。因此二者的扩展速率不同（图 6）。

预填充呈二次方扩展

预填充执行 ISL² 量级的工作（每个词元都要关注所有词元），而 KV 流量只与 ISL 成正比增长（公式 2、3）。因此算术强度随 ISL 线性上升，使预填充远高于拐点，保持计算受限。ISL 翻倍，运行时间应大致变为四倍（图 6）。在 ISL 较短时，扩展幅度低于 4 倍，因为固定的启动与后处理开销占主导；只有当 ISL 大到足以摊销这些开销时，才会表现出二次方扩展。

解码呈线性扩展

每一步都要读取完整的 KV 缓存来生成一个词元，因此字节数随 KVSL 增长，而单步计算量仍然很小（公式 5 和 6）。算术强度保持在 2 × $G$ 附近，远低于拐点，所以解码在所有长度下都是内存受限。KVSL 翻倍，运行时间应翻倍，图 6 证实了这一点。在 KVSL 较短时，扩展幅度低于 2 倍，因为启动、后处理和 flash-decoding 归约的开销不随 KVSL 扩展。随着 KVSL 增大，它们的占比会逐渐缩小。

准则 3：尽可能减小有效 KV 状态。序列长度由使用场景决定，但代价并不对称：预填充按 ISL² 增长，而解码随 KVSL 线性增长。可以通过 KV 缓存压缩、稀疏或滑动窗口注意力，或者混合模型架构（如 Nemotron 3，只有部分层携带不断增长的全局 KV 状态）来减小有效 KV 状态。

张量并行在多块 GPU 间切分注意力头

张量并行（TP）把注意力头切分到多块 GPU 上，使每块 GPU 拥有 QH/TP 个查询头和 KH/TP 个 KV 头。它切分的是头，而不是词元。在表 3 中，只有包含 KH/TP 的 batch 维度会缩小；单块 GPU 上的 GEMM 形状和算术强度保持不变。

TP 有一个实际限制：KV 头必须能在 GPU 间均匀切分。一旦 TP > KH，同一分组的查询头就会跨越多个 rank，每个 rank 都需要一份共享 KV 头的副本。这会复制 KV 状态，增加内存和带宽开销却没有任何收益（图 7）。因此，应保持 TP ≤ KH，使每块 GPU 至少拥有一个完整的分组：一个 KV 头及其 $G$ 个查询头。

KV 头很少的模型（例如只有两个的 Nemotron 3）很快就会用尽 TP 的扩展空间，因为在不复制的前提下，缓存无法切分到每块 GPU 少于一个 KV 头。此时注意力必须采用其他扩展方式：注意力数据并行（ADP）对请求进行切分，KV 并行（KVP）则把长序列 KV 缓存切分到多块 GPU 上。FFN 则通过专家并行（EP）单独扩展。

TensorRT-LLM 将这些方式组合起来，形成 Wide EP（注意力用 ADP，FFN 用 EP）和 Helix Parallelism（注意力用 KVP，FFN 用 EP）。在这两种方案中，KH 都决定了高效的扩展方式。

准则 4：让 KH 决定并行策略。保持 TP ≤ KH，使每块 GPU 拥有完整的 KV 头。KV 头很少的模型（MQA 为 1，GQA 为 2）会很快用尽 TP，更适合对注意力采用 ADP 或 KVP，并对 MoE FFN 采用 EP（在 TensorRT-LLM 中实现为 Wide EP 和 Helix Parallelism）。

开始协同设计 AI 模型注意力

可将下面总结的四条准则作为模型设计检查清单，开始对 AI 模型注意力进行协同设计。这些选择能够在相同硬件上提高 GPU 利用率、推理速度、吞吐量和交互性。

• 准则 1：按解码需求选择分组大小（$G$），并尽量取高。预填充对分组大小不敏感。
• 准则 2：使用头维度（Hsz）= 128 或 256，以匹配 GPU 分块和 128 字节传输，同时不超出 TMEM 预算。更大的头还能在预填充中掩盖 softmax 的开销。
• 准则 3：通过 KV 缓存压缩、稀疏或滑动窗口注意力，或混合模型来减小有效 KV 状态。
• 准则 4：让并行方式与 KV 头数（KH）相匹配。保持 TP ≤ KH，对 KH 很少的模型采用 Wide EP 和 Helix Parallelism 进行扩展。

致谢

本文是 NVIDIA 跨团队合作的成果。感谢 Timmy Liu、Jatin Mitra、Tiyasa Mitra、Bhargava Gopireddy、Brian Pharris、Julien Demouth 和 Eduardo Alvarez 的帮助。