随着智能体和长上下文工作负载变得普遍，上下文长度增加，注意力在推理时间中所占的比例也越来越大（图1）。由于注意力现在主导了这部分成本，它的设计方式——而不仅仅是实现方式——越来越决定模型的推理性能。围绕GPU如何执行来塑造模型架构，是AI模型协同设计的前提。关于模型设计选择如何在不牺牲准确性的情况下影响吞吐量和交互性的讨论，请参阅之前的文章AI模型协同设计：硬件友好的LLM设计。
本文探讨分组大小（每个KV头对应的查询头数）、头维度和序列长度如何影响稠密注意力的性能，其中每个查询都会关注序列长度上的所有键和值。我们将这一分析，连同注意力如何在GPU上并行化，提炼为四条实用指南：一份协同设计检查清单，帮助模型开发者在NVIDIA GPU上提高推理吞吐量和交互性。敬请期待一篇涵盖稀疏注意力的文章。
图1. DeepSeek-R1在4K、32K和128K上下文长度下的预填充时间分解，其中注意力占比从18%上升到85%
每项分析都基于两个来源：来自GEMM形状算术的解析公式，以及使用FP8进行注意力计算和KV缓存的预填充和解码内核的实测数据。
缩写 | 定义 |
PB | 预填充批大小 |
DB | 解码批大小 |
QH | 查询头数 |
KH | KV头数（MHA时KH = QH，GQA时KH = QH/G，MQA时KH = 1） |
(G) | 分组大小 = QH/KH（共享一个KV头的查询头数） |
Hsz | 头维度（通常为64、128或256） |
ISL | 输入序列长度（预填充中的查询token数） |
KVSL | 解码迭代中的平均KV缓存序列长度 |
表1. 本文公式中使用的符号
预填充和解码为何是两个不同的问题？
预填充并行处理整个提示，产生大型GEMM-M（= ISL × (G)）矩阵乘法，这些是计算受限的。在没有推测解码的情况下，解码一次生成一个token，产生小型GEMM-M（= (G)）矩阵乘法，并因从高带宽内存（HBM）读取KV缓存而成为内存受限的。
推测解码会增加GEMM-M，并可能使解码转向计算受限。由于预填充和解码的查询长度、KV访问和瓶颈不同（表2），每个参数都针对每个阶段单独分析。
预填充 | 解码 | |
查询长度 | 完整输入（ISL个token） | 1个token |
KV上下文 | 提示（ISL个token） | 完整KV缓存（KVSL个token） |
注意力GEMM-M | ISL × (G)（大） | (G)（小） |
主要瓶颈 | 计算（矩阵乘法 + softmax） | HBM带宽（内存） |
表2. 预填充和解码在查询长度、KV上下文、注意力GEMM-M和主要瓶颈方面不同
注意：在智能体和多轮应用中常见的前缀缓存下，新一轮可能具有较短的ISL，同时关注一个大的前缀缓存。在短ISL但长前缀缓存的情况下，预填充的行为类似于解码。
算术强度如何决定计算受限与内存受限行为
如此前所述，roofline 模型通过计算和带宽上限来界定 GPU 性能。算术强度决定了哪一项成为瓶颈（公式 1）：
算术强度 = 总 FLOPs / 总访问字节数
脊点（ridge point）标志着从内存受限到计算受限的转变。Prefill 远高于脊点，属于计算受限；而 decode 低于脊点，属于内存受限（图 2）。投机解码提高了 decode 的算术强度，可将其推向脊点。
图 2. Roofline 模型显示 decode 位于内存受限斜坡上，prefill 位于计算受限平台上
FlashAttention 内核如何在 GPU 上计算注意力？
FlashAttention 在不物化完整注意力矩阵的情况下计算注意力。它将 (Q)、(K) 和 (V) 的分块从 HBM 流式传输到片上 SRAM，并将三个步骤融合为一次遍历：
- 首先，批量矩阵乘法（BMM1）计算查询与键的得分
- 其次，在线 softmax 使用运行最大值和运行和来归一化得分
- 第三，第二次批量矩阵乘法（BMM2）对值进行加权
BMM 在 Tensor Core 上运行，而 softmax 指数运算在特殊函数单元上运行。BMM 的形状驱动了后续的算术强度分析。
图 3. FlashAttention 内核，BMM1 和 BMM2 在 Tensor Core 上运行，融合的在线 softmax 在特殊函数单元上运行。图片改编自
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness## GEMM 形状
注意力性能取决于其两个矩阵乘法的形状。表 3 列出了 BMM1 和 BMM2 各阶段（Batch、M、N、K）的维度。
BMM | 阶段 | Batch | M | N | K | 含义 |
BMM1 | Prefill | PB × KH | ISL × (G) | ISL | Hsz | Q · Kᵀ：计算查询与键的得分 |
Decode | DB × KH | 1 × (G) | KVSL | Hsz | ||
BMM2 | Prefill | PB × KH | ISL × (G) | Hsz | ISL | 权重 · V：聚合值 |
Decode | DB × KH | 1 × (G) | Hsz | KVSL |
表 3. BMM1 和 BMM2 的 GEMM（Batch、M、N、K）形状对于 decode，GEMM-M = (G)，通常为 8-16，远低于 GPU 的 tile-M 64 或 128，限制了每个 tile 的并行工作量。更大的 (G) 意味着每个 token 加载更少的 KV，并将每次加载分摊到更多查询头上，从而提高利用率。下一节将量化这一效果。
组大小
组大小（(G)）是共享一个 KV 头的查询头数量。MHA 的 (G) = 1，GQA 的 (G) = 4、8、16、…，MQA 的 (G) = QH。
算术强度作为 (G) 的函数。** **在以下公式中，“字节”指移动的 HBM 字节数。为简化起见，假设每个元素 1 字节（即 FP8 KV 缓存）。
Prefill
随着 (G) 增大，1/(G) 项消失，算术强度趋近于 2 × ISL。在 ISL = 32K 时，将 (G) 从 8 提高到 16 仅使算术强度提升不到 6%。换言之，prefill 主要由 ISL 决定，而非 (G)。图 4 证实了这一点：将 (G) 从 1（MHA）变化到 64（MQA），prefill 运行时间变化不到 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)) → 当 (G) → ∞ 时趋近于 2 × ISL
解码（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) 每翻倍一次，解码运行时间约下降 2 倍，因为将 KV 头减半会使每个 token 加载的数据减半。超过 (G) = 16 后，KVSL = 32K 曲线趋于平缓。其每步内核足够小，以至于两项成本占主导：固定的设置和后处理开销，以及为在少量 KV 头下保持并行而将 KV 拆分到多个 SM 上所带来的 flash-decoding 归约开销。较长的 KVSL = 128K 内核能更好地摊销这些成本，并继续遵循 2 倍趋势。
注意：推测解码将有效 GEMM-M 提升到 (1+(D)) × (G)，其中 (D) 是草稿 token 的数量。一旦大到足以填满计算瓦片，解码就会转向计算受限。
图 4. 归一化运行时间与(G)
（QH=64，Hsz=128，PB=1，DB=8；预填充 ISL 和解码 KVSL 分别为 32K 和 128K）。预填充在 (G) 上保持平坦（计算受限）；解码在 (G) 每翻倍时约下降 2 倍（内存受限），从 MHA（(G)=1）到 MQA（(G)=64）
准则 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 无关，因为它作用于注意力分数矩阵（查询 token × 键），该矩阵没有头维度。公式 8 和 9：
Softmax 操作数（预填充）≈ PB × QH × ISL²
Softmax 操作数（解码）≈ DB × QH × KVSL
这些因素之间的平衡决定了每个阶段的 Hsz 成本（图 5）。
图 5. 归一化运行时与 Hsz 的关系（QH=64，(G)=32，PB=1，DB=8；预填充 ISL 与解码 KVSL 分别为 32K 和 128K）。运行时随 Hsz 上升
预填充是计算受限的** **（矩阵乘法加 softmax）。随着 Hsz 增大，矩阵乘法的 FLOPs 增长，而 softmax 保持不变。如果预填充是纯矩阵乘法，Hsz 翻倍会使运行时翻倍；但固定的 softmax 不随之扩展，因此运行时的上升幅度小于 Hsz 的增幅。图 5 证实了这一点：预填充随 Hsz 上升，但慢于 Hsz 的增长。更宽的 Hsz 摊薄了 softmax，将更多内核时间转移到矩阵乘法上，使预填充不再那么受 softmax 限制。
解码是内存受限的（流式读取 KV 缓存）。更大的 Hsz 会增加每个 token 的 KV 字节数（公式 6），因此运行时应当随 Hsz 扩展。图 5 证实了这一点，不过略呈次线性，因为设置、后处理以及 flash-decoding 归约开销不随 Hsz 扩展，且在较短的 KVSL = 32K 内核上占比更大。
准则 2：使用 128 或 256 的 Hsz。** **Hsz 不改变算术强度，但必须与硬件对齐。通常，Hsz = 64 仍需为 128 宽的 tile 付出代价，而 Hsz ≥ 512 则逼近 TMEM 容量上限。这使得 128 和 256 成为最佳选择。
序列长度
序列长度（ISL / KVSL）对预填充和解码的影响不同，因为它在每个阶段通过不同的变量进入：预填充一次性处理所有 ISL 输入 token，而每个解码步骤读取长度为 KVSL 的 KV 缓存。因此两者的扩展速率不同（图 6）。
预填充呈二次方扩展
预填充执行 ISL² 的工作量（每个 token 关注每个 token），而 KV 流量仅与 ISL 成正比增长（公式 2、3）。因此算术强度随 ISL 线性上升，使预填充远高于脊点，处于计算受限状态。ISL 翻倍应大致使运行时变为四倍（图 6）。在较短的 ISL 下，扩展低于 4 倍，因为固定的设置和后处理开销占主导；一旦 ISL 大到足以摊薄这些开销，二次方扩展便显现出来。
解码呈线性扩展
每个步骤读取完整的 KV 缓存以生成一个 token，因此字节数随 KVSL 增长，而每步的工作量仍然很小（公式 5 和 6）。算术强度保持在约 2 × (G)，远低于脊点，因此解码在所有长度下都保持内存受限。KVSL 翻倍应使运行时翻倍，图 6 证实了这一点。在较短的 KVSL 下，扩展低于 2 倍，因为设置、后处理以及 flash-decoding 归约开销不随 KVSL 扩展。随着 KVSL 增长，它们的占比会缩小。
图 6. 归一化运行时与序列长度的关系（QH=64，(G)=32，Hsz=256，PB=1，DB=8）。预填充随 ISL 按 O(n²) 扩展；解码随 KVSL 按 O(n) 扩展；两者在较短长度下均低于理想扩展速率
准则 3：尽可能减少有效 KV 状态。用例决定了序列长度，但代价是不对称的：预填充随 ISL² 增长，而解码随 KVSL 线性增长。通过 KV 缓存压缩、稀疏或滑动窗口注意力，或混合模型架构（如 Nemotron 3，其中只有部分层承载不断增长的全局 KV 状态）来减少有效 KV 状态。
张量并行将注意力头拆分到多个 GPU 上
张量并行（TP）将注意力头拆分到多个 GPU 上，每个 GPU 获得 QH/TP 个查询头和 KH/TP 个 KV 头。它分片的是头，而不是 token。在表 3 中，只有包含 KH/TP 的批次维度会缩小；每个 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 决定了高效扩展的方式。
图 7. 注意力分片策略：当 TP > KH 时，KV 缓存被复制，增加内存和带宽开销
准则 4：让 KH 决定并行策略。保持 TP ≤ KH，使每个 GPU 都有一个完整的 KV 头。KV 头很少的模型（MQA 为 1，或 GQA 为 2）会很快耗尽 TP，更适合用 ADP 或 KVP 处理注意力，加上 EP 处理 MoE FFN（在 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，并使用Wide EP和Helix Parallelism来扩展少量 KH 模型。
致谢
本文是 NVIDIA 跨团队合作的成果。我们感谢 Timmy Liu、Jatin Mitra、Tiyasa Mitra、Bhargava Gopireddy、Brian Pharris、Julien Demouth 和 Eduardo Alvarez 的帮助。