随着智能体与长上下文工作负载成为常态，上下文长度不断增加，注意力在推理时间中所占比例也越来越大。在一组实测中，DeepSeek-R1 的 prefill 时间构成中，注意力占比从 4K 上下文的 18% 上升到 128K 上下文的 85%。当注意力已成为主导成本时，决定模型推理性能的不仅是如何实现注意力，更是如何设计注意力——围绕 GPU 的执行方式塑造模型架构，这正是 AI 模型协同设计的前提。
NVIDIA 的分析聚焦于稠密注意力，即每个查询都会关注序列长度上的全部键与值，并考察分组大小（每个 KV 头对应的查询头数量）、头维度与序列长度如何影响性能。所有分析都建立在两类依据之上：来自 GEMM 形状算术的解析公式，以及使用 FP8 进行注意力计算与 KV 缓存时的 prefill、decode 内核实测数据。
prefill 与 decode 是两个不同的问题。prefill 并行处理完整提示词，产生大 GEMM-M（等于输入序列长度乘以分组大小）的矩阵乘法，属于算力受限；而在没有投机解码的情况下，decode 一次生成一个 token，产生小的 GEMM-M（等于分组大小），并因从高带宽显存读取 KV 缓存而变为显存带宽受限。投机解码会提高 GEMM-M，并可能把 decode 推向算力受限一侧。由于查询长度、KV 访问模式与瓶颈在 prefill 和 decode 中均不相同，每个参数都需要分阶段单独分析。
需要注意的是，在智能体与多轮应用中常见的前缀缓存场景下，新一轮对话可能输入序列长度很短，却要关注很长的前缀缓存；当输入序列短而前缀缓存长时，prefill 的行为会更接近 decode。
算力受限还是显存受限，由算术强度决定。roofline 模型用算力上限与带宽上限界定 GPU 性能，而算术强度（总浮点运算数除以总访问字节数）决定了哪一条上限起约束作用。脊点（ridge point）标记了从显存受限到算力受限的转折：prefill 位于脊点之上属于算力受限，decode 位于脊点之下属于显存受限，而投机解码提高了 decode 的算术强度，可使其向脊点移动。
在 GPU 上实际执行注意力的 FlashAttention 内核，通过不物化完整注意力矩阵来降低开销：它把 Q、K、V 的分块从高带宽显存流式读入片上 SRAM，并把三步融合为一次遍历——先用第一个批量矩阵乘法（BMM1）计算查询与键的得分，再用在线 softmax 借助运行中的最大值与和完成归一化，最后用第二个批量矩阵乘法（BMM2）对值加权。两个矩阵乘法运行在 Tensor Core 上，softmax 的指数运算运行在特殊函数单元上。这些矩阵乘法的形状决定了后续的算术强度分析。
在 decode 阶段，GEMM-M 等于分组大小，通常为 8 到 16，远低于 GPU 分块所需的 64 或 128，从而限制了每个分块可并行的工作量。更大的分组大小意味着每个 token 需要加载的 KV 更少，并且每次加载可以摊薄到更多查询头上，从而提升利用率。分组大小即共享同一个 KV 头的查询头数量：MHA 的分组大小为 1，GQA 通常为 4、8、16 等，MQA 则等于查询头总数。
在 prefill 阶段，随着分组大小增长，公式中的 1/G 项逐渐消失，算术强度趋近于输入序列长度的两倍。当输入序列长度为 32K 时，把分组大小从 8 提高到 16 对算术强度的改善不足 6%。换言之，prefill 由输入序列长度而非分组大小主导：实测中把分组大小从 1（MHA）变化到 64（MQA），prefill 运行时间的变化不足 1%。NVIDIA 把这些分析连同注意力在多 GPU 上的并行方式，提炼为一份包含四条准则的协同设计清单，帮助模型开发者提升在 NVIDIA GPU 上的推理吞吐与交互性，并表示后续还将发布一篇覆盖稀疏注意力的文章。
WeChat
Profile