面向快速交互式长上下文推理的AI模型注意力协同设计

NVIDIA发布了一份新的分析，研究注意力机制设计选择如何影响其GPU上的快速、交互式长上下文推理。该指南是公司AI模型协同设计系列的一部分，聚焦于密集注意力，并为模型开发者提出了一份检查清单。它研究了组大小、头维度和序列长度如何影响预填充和解码性能。

随着上下文长度扩展和代理工作负载增长，注意力现在主导推理成本。研究人员认为，注意力的设计方式（而不仅仅是实现方式）决定了推理性能。他们主张围绕GPU执行协同设计模型架构。

该分析区分了两个阶段：预填充并行处理提示，受计算限制；而解码逐个生成标记，受KV缓存读取的内存限制。算术强度决定了适用的瓶颈。屋顶线模型显示预填充在峰值点之上，解码在峰值点之下。推测解码可以将解码转向计算受限。

像FlashAttention这样的注意力内核将Q、K和V的片从HBM流式传输到SRAM，并融合三个步骤：得分、softmax和值加权。批量矩阵乘法在Tensor Cores上运行，其形状决定了算术强度。对于解码，GEMM-M通常为8-16，远低于GPU片大小，限制了并行工作。

组大小（即共享一个KV头的查询头数量）至关重要。多头注意力的组大小为1，分组查询注意力使用4、8或16，多查询注意力共享一个KV头。较大的组大小每个标记加载更少的K/V，并在查询头之间分摊加载，从而提高利用率。

该论文提供了来自预填充和解码内核的分析公式和测量数据，使用FP8进行注意力和KV缓存。四条实用指南帮助开发者提高吞吐量和交互性。前缀缓存可以在序列长度较短但前缀缓存较长时使预填充表现得像解码。

这些建议