vLLM Triton Attention Backend Deep Dive

发表时间: 2026-03 · Blog post by vLLM Team at IBM Research (blog.vllm.ai)

原文: https://blog.vllm.ai/blog/2026-03-04-vllm-triton-backend-deep-dive

作者/机构：vLLM Team at IBM Research (Burkhard Ringlein, Jan van Lunteren, Chih-Chieh Yang, Sara Kokkila Schumacher, Thomas Parnell, Mudhakar Srivatsa, Raghu Ganti)

速读

一句话结论 本文为 vLLM 推理框架开发了一个完全基于 Triton 编写的 Paged Attention 后端，通过引入 Q Blocks 块划分、并行分块 Softmax 以及持久化内核设计，在仅用约 800 行代码的情况下实现了跨 GPU 厂商的性能可移植性，并在 NVIDIA 和 AMD 硬件上达到了与高度定制化专用内核相当甚至更优的端到端推理延迟。

要解决什么问题 原有的大语言模型推理框架在处理注意力机制时，通常依赖为特定模型和特定 GPU 架构高度定制的底层内核。这种做法面临两个极其严重的卡点。首先是工程维护层面的扩展性灾难：随着加速器硬件生态的日益多样化（如 NVIDIA Hopper 与 Blackwell、AMD MI300、Intel 芯片等），以及工作负载特征的碎片化（包括纯预填充、纯解码、混合负载，以及不断变化的批处理大小和上下文长度），维护成百上千个专用内核的成本极高且极不现实。微基准测试表明，没有任何一种单一的内核配置能够统治所有场景。其次是系统级调度与硬件底层执行机制的冲突：在推理引擎中，为了降低内核启动的开销，通常会使用 CUDA Graphs 来记录并回放固定的执行图；然而，注意力内核的启动网格（Launch Grids）大小往往依赖于动态变化的批处理大小和序列长度。当 GPU 执行这些内核时，如果启动的线程数超过了物理流式多处理器（SM）的数量，计算就会分波次（Waves）执行。如果直接将动态变化的网格捕获到 CUDA Graphs 中，即使后续实际负载变小，系统仍会盲目按照记录时的最大网格进行回放，导致出现大量多余的执行波次，进而引发严重的计算资源浪费和延迟激增。

怎么做的 核心思路是利用 Triton 这种领域特定语言，在 Python 层面以逻辑分块（Logical Tiles）的形式表达注意力计算，将硬件相关的执行布局交由编译器和自动调优器处理，并通过重新设计任务调度机制来彻底消除 CUDA Graphs 回放时的资源浪费。该方法主要由三个关键部件构成。第一是用于优化矩阵乘法（tl.dot）的 Q Blocks 设计。在 Paged Attention 遍历分页 KV 缓存计算注意力分数时，由于 KV 侧的分块大小被严格限制为缓存的页大小，作者将优化重心放在了查询（Query）侧。对于分组查询注意力（GQA），该设计将与同一个 KV 头关联的所有查询头合并处理以提升缓存复用率，同时把多个查询 token 组合成一个单一的工作项（即 Q Block），以此提高计算单元的利用率，并依赖自动调优为不同平台选择合适的块大小。第二是应对解码阶段的并行分块 Softmax（被称为 3D 内核）。在仅需处理单个查询 token 的纯解码场景下，Q Blocks 无法提供并行收益。为此，该部件将 KV 缓存的遍历过程拆分给多个内核实例并行执行，每个实例计算出局部结果后，再启动第二个归约内核（因为 Triton 缺乏全局屏障机制）来汇总最终输出，通过牺牲少量的内核启动开销换取了极大的并行度。第三是持久化内核（Persistent Kernels）机制。为了解决动态网格与 CUDA Graphs 的冲突，该机制彻底放弃了随负载大小变化的传统启动网格，改为启动数量固定且与当前 GPU 可用计算资源（SM 数量）完全相等的内核实例。每个实例在启动后，会动态地从 GPU 内存中读取元数据来决定自己当前需要处理的工作量。这种设计使得启动网格保持绝对恒定，从而让 CUDA Graphs 能够被高效、无损地重复调用。

效果如何 实验在 NVIDIA H100 和 AMD MI300 硬件平台上进行，测试模型为 Llama 3.1 8B，设置批处理大小为 1，输入长度为 500 tokens。在 NVIDIA H100 上，对比基线为高度优化的、代码量高达 70000 行的 FlashAttention 3，本方法（仅约 800 行代码）在长序列解码请求下达到了基线 100.7% 的性能表现。在 AMD MI300 平台上，对比基线为早期的注意力实现版本，本方法实现了约 5.8 倍的端到端延迟加速。目前，该 Triton 注意力后端已成为 vLLM 在 AMD GPU（基于 ROCm）上的默认选项，并在 Intel XPU 运行 float32 时作为后备方案。除了常规场景，该方法还成功支持了需要特定功能的模型，例如 StepFun 音频模型使用的 ALiBi sqrt、在 pre-Hopper 架构（如 A100）上的 Sink tokens 和 GPT-OSS 行为，以及小头尺寸、编码器-解码器注意力和多模态前缀注意力。该方法的局限性在于，并行分块 Softmax 引入了额外的归约内核启动开销，因此在实际应用中需要依赖启发式规则来判断何时启用该 3D 内核才能获得正向收益。此外，作者还提到正在实验基于 PyTorch 团队新推出的领域特定语言 Helion 来实现简化版的 Paged Attention，并已取得初步成果。

一、 主要贡献

随着加速器硬件的多样性不断增加（例如 NVIDIA Hopper 和 Blackwell、AMD MI300、Intel 及其未来平台），为特定模型和 GPU 架构编写和维护数百个高度专业化的底层算子（Kernels）的成本急剧上升，这种传统方法已无法规模化。为了解决这一核心问题，IBM Research、Red Hat 和 AMD 的团队合作开发并向开源社区贡献了基于 Triton 的 vLLM 注意力后端（Attention Backend）。

该研究的核心目标是实现具有强大跨 GPU 供应商可移植性的最先进（State-of-the-art）推理性能。其主要创新点和系统设计如下：

  1. 引入 Triton 作为统一后端：Triton 是一种领域特定语言（DSL），允许开发者用 Python 编写矩阵乘法或注意力等 GPU 算子，并将其编译为适用于多个平台的高效 GPU 代码。Triton 采用分块编程模型（Tiled programming model），这种模型在表达硬件相关优化的底层细节与保持硬件无关的高层抽象之间取得了平衡。开发者以逻辑块（Logical tiles）的形式表达计算，而 Triton 编译器和自动调优器（Autotuner）会自动决定如何将这些块映射到特定的底层硬件布局上。
  2. 构建统一的注意力抽象层：vLLM 引入了注意力后端抽象层，将注意力机制的实现封装在通用 API 之后，与线性层或层归一化等简单组件隔离开来。Triton 注意力后端完全使用 Triton 原生实现，依赖项仅为 PyTorch 和 Triton，能够在 NVIDIA、AMD 和 Intel GPU 上运行相同的源代码。
  3. 广泛的场景支持与后备机制：该后端目前是运行在 ROCm 上的 AMD GPU 的默认选项；在 Intel XPU 上运行 float32 时，由于 Flash Attention 不支持 fp32，系统也会回退到 Triton Attention。此外，它还支持各种特殊模型需求，例如 StepFun 音频模型使用的 ALiBi sqrt、sink tokens、GPT-OSS 行为（特别是在 A100 等 pre-Hopper 架构的 NVIDIA GPU 上）、小注意力头尺寸（Small head sizes）、编码器和解码器注意力、以及多模态前缀注意力。当 FlashAttention、FlashInfer 等依赖项不可用或导入失败时，它作为常驻组件提供可靠的后备支持，同时支持批处理不变性（Batch invariance）等特性。

Figure 1Figure 1 Triton的分块编程模型，逻辑块由编译器和自动调优器映射到特定于硬件的执行布局。

二、 背景知识与关键设计原则

Paged Attention 算子的工作原理
Paged Attention 通过对 KV 缓存进行分页（Paging）来实现内存高效的注意力机制。其核心处理流程如下：首先，针对批处理（Batch）中的每个查询（Query），算子会处理每一个查询标记（Query token）。接着，对于每个标记，算子会遍历查询头（Query heads）及其对应的 KV 头。最后，算子遍历分页的 KV 缓存，计算注意力分数（Attention scores）并应用值向量（Value vectors）。在逻辑布局上，查询标记沿 x 轴排列，查询头沿 y 轴排列，而对分页 KV 缓存的遍历构成了最内层的循环。

Figure 3Figure 3 Paged attention的概念视图，展示了查询标记、查询头以及对分页KV缓存的遍历。

CUDA Graphs、启动网格与 GPU 执行波（Execution Waves）
CUDA Graphs 通过记录和重放固定的执行图来显著减少算子的启动开销。然而，注意力算子的启动网格（Launch grids）通常依赖于批处理大小（Batch size）和序列长度（Sequence length），这给 CUDA Graphs 的使用带来了挑战。
在硬件层面，GPU 使用固定数量的流式多处理器（Streaming Multiprocessors, SMs）来执行算子。当启动的线程数超过可用的 SM 数量时，执行会以“波”（Waves）的形式进行。如果产生第二波执行，往往会导致硬件利用率不足。当这种行为被捕获到 CUDA Graph 中时，即使后续实际的有效工作负载规模减小，这种低效的固定启动网格依然会被重放，从而导致额外的无效工作和延迟增加。

Figure 5Figure 5 当启动的线程数超过可用的流式多处理器时产生的GPU执行波。在此示例中，GPU有8个SM，而我们需要执行12个线程。
Figure 6Figure 6 通过CUDA graphs重放固定启动网格时导致的额外无效工作。

三、 方法细节

基于微基准测试（Microbenchmarks）的算子开发
在 Triton 注意力后端开发的初期，算子的实现独立于 vLLM 之外，并通过广泛的微基准测试进行评估。API 的设计严格匹配 vLLM 的需求，但在进行端到端集成之前，性能调优是在隔离环境中完成的。微基准测试对于理解算子在不同工作负载下的性能行为至关重要，包括重预填充（Prefill-heavy）、重解码（Decode-heavy）以及混合工作负载，同时也覆盖了不同的批处理大小和上下文长度。测试结果表明，不同的算子变体在不同的运行区间内各有优势，没有任何单一配置能够在所有场景下占据绝对主导地位。这种微基准测试补充了端到端基准测试的不足，暴露了可能被系统级效应掩盖的算子级底层行为。

Figure 2Figure 2 多种Triton paged attention算子变体在预填充、解码和混合工作负载下的微基准测试比较。

使用 Q Blocks 优化 tl.dot 的分块大小
注意力机制的核心计算是矩阵乘法，在 Triton 中通过 tl.dot 实现。为了充分利用硬件性能，必须使用足够大的分块（Tiles），仅仅加载分页 KV 缓存无法获得理想的性能。
由于 KV 侧的分块大小受到 KV 缓存页面大小（Page size）的严格限制，因此优化的重点转向了查询（Query）侧。对于分组查询注意力（Group Query Attention, GQA），通过将与单个 KV 头关联的所有查询头组合在一起处理，可以有效增加缓存的重用率。为了进一步提升并行度，设计中将多个查询标记分组为一个单一的工作项（Work item），这一结构被称为 Q block。在执行映射上，启动网格覆盖了批处理大小和 KV 头，而 Q blocks 则决定了每个算子实例处理多少个查询标记和查询头。最终，通过自动调优（Autotuning）为每个硬件平台选择最合适的 block 大小。

Figure 4Figure 4 Q blocks将多个查询头和查询标记组合成单个工作项，以提高tl.dot的利用率和缓存重用。

通过并行分块 Softmax（Parallel Tiled Softmax）增加并行度
一次性处理多个查询标记的方法在预填充（Prefill）工作负载中表现出色，但对于解码（Decode）工作负载毫无益处，因为解码阶段每次只处理一个查询标记。为了解决这一问题，设计引入了并行分块 Softmax（被称为“3D 算子”）来提供额外的并行度。
该方法将对 KV 缓存的遍历拆分到多个算子实例中执行。每个实例计算出部分结果，随后这些部分结果被归约（Reduced）以生成最终输出。由于 Triton 语言本身不提供全局屏障（Global barrier），这种归约操作必须通过启动第二个算子来实现。这引发了额外并行度与算子启动开销之间的权衡（Trade-off），系统通过内部启发式算法（Heuristics）来动态决定何时采用这种方法最为有利。

从可变启动网格（Variable Launch Grids）演进到持久化算子（Persistent Kernels）
早期版本的 Paged Attention 算子采用的是随工作负载规模动态缩放的可变启动网格。虽然这种方式具有灵活性，但正如背景知识所述，它与 CUDA Graphs 的交互效果极差。
为了彻底解决这个问题，设计了持久化算子（Persistent kernels）（相关 PR 正在提交至 vLLM）。具体流程为：首先，启动一个固定数量的算子实例，该数量严格等于当前可用的计算资源。接着，每个算子实例通过读取 GPU 内存中的元数据（Metadata），动态地决定自己需要处理多少工作量。这种机制使得启动网格保持绝对恒定，从而允许 CUDA Graphs 被极其高效地重用，消除了重放过程中的无效工作。

Figure 7Figure 7 早期paged attention算子中使用的可变启动网格。
Figure 8Figure 8 采用固定启动网格和动态工作分配的持久化算子方法。

四、 实验环境

  • 评估模型：Llama 3.1 8B。
  • 工作负载配置：批处理大小（Batch size）设置为 1，输入上下文长度为 500 tokens，输出长度可变。
  • 硬件配置：NVIDIA H100 GPU 和 AMD MI300 GPU。
  • 软件配置：基于 PyTorch 和 Triton 的 vLLM，对比基线为 FlashAttention 3 及早期算子实现。

五、 实验结果

在 2025 年末进行的基准测试验证了上述方法的有效性：

  1. NVIDIA 平台表现：在 H100 上，针对长解码请求，Triton 注意力后端的性能达到了 FlashAttention 3 的 100.7%。
  2. AMD 平台表现：在 MI300 上，与早期的实现相比，Triton 注意力后端实现了约 5.8 倍的加速。
  3. 代码复杂性与可移植性：最关键的是，这两个平台使用的是完全相同的 Triton 算子源代码。Triton 版本的 paged attention 实现仅用约 800 行代码，而作为对比，FlashAttention3 的代码量约为 70,000 行。

Figure 9Figure 9
Figure 9Figure 9 Triton paged attention与FlashAttention 3在NVIDIA H100和AMD MI300上的端到端延迟比较。结果以最左侧的基线进行了归一化处理。

六、 结论

随着模型架构、推理优化技术和硬件平台的持续演进，性能可移植性（Performance portability）变得前所未有地重要。vLLM 中的 Triton 注意力后端证明了，使用单一且可移植的算子实现，完全可以达到业界最先进的注意力计算性能。
通过精心的算子设计、详尽的微基准测试，以及持久化算子和 CUDA Graphs 等系统级优化，Triton 后端在保持跨 GPU 供应商可移植性的同时，性能持平甚至超越了高度专业化的底层实现。目前，它已成为 AMD 平台的默认注意力后端，并能使用同一套源代码在 NVIDIA 和 Intel 平台上高效运行。

未来展望（Helion 预览）
研究团队正在探索 PyTorch 团队推出的新领域特定语言 https://github.com/pytorch/helion（可视为更高抽象级别的 Triton 或分块的 PyTorch）。作为实验，团队已经在 Helion 中实现了一个简化的 paged attention 算子，并取得了非常有希望的早期结果，相关代码已作为草案 PR 提交至 vLLM 仓库。

七、 补充细节

  • 自动调优原理参考文献：关于 Triton 编译器和自动调优器如何决定逻辑块到硬件特定执行布局的映射细节，详见文献【1，GPU Performance Portability needs Autotuning + 2025 + https://arxiv.org/abs/2505.03780】 。
  • 核心算子实现源码：Triton paged attention 算子的核心代码位于 vLLM 仓库的 vllm/v1/attention/ops/triton_unified_attention.py 路径下。
  • 底层优化原理论述：关于算子底层优化的详细解释，作者推荐参考由算子原作者撰写的 PyTorch 官方博客《Enabling vLLM v1 on AMD GPUs with Triton》。
  • 完整优化与基准测试细节：本文涵盖了 Triton 注意力后端中最重要的优化概述，所有技术细节及更全面的基准测试结果均发表于相关论文【2，The Anatomy of a Triton Attention Kernel + 2025 + https://arxiv.org/abs/2511.11581】 。