vLLM Triton Attention 后端深度解析

本文改编自红帽托管的 vLLM 办公时间 (Office Hours) 会话，由 IBM Research 的 Burkhard Ringlein 主讲，对 vLLM Triton Attention 后端进行了深入的技术剖析。探索过往主题并点击此处参加未来的办公时间会话。

在过去的一年中，来自 IBM Research、红帽和 AMD 的团队共同开发并将一个基于 Triton 的 Attention 后端合并进了 vLLM 主线，旨在实现业界领先的性能，同时在不同 GPU 厂商之间保持强大的可移植性。这项工作是由加速器硬件日益多样化以及维护大量高度专业化内核的成本上升所驱动的。

本文提供了该工作的深入技术解析。我们解释了为什么 Triton 非常适合 vLLM，描述了 Triton Attention 后端及其应用场景，随后深入探讨了高性能 Paged Attention 内核的实现。在此过程中，我们将涵盖内核级优化、并行化策略、CUDA 图交互以及基准测试结果，最后简要介绍 Helion。

为什么 Triton 对 vLLM 有所助益

vLLM 旨在跨平台、模型和执行策略提供尽可能最佳的推理性能。在实践中，这意味着要支持多种加速器及其代际、广泛的模型架构，以及多样的负载特性，例如变化的批处理大小、序列长度和 Attention 模式。

一种方法是编写许多高度专业化的内核，每个内核都针对特定的模型和 GPU 架构进行调优。虽然有效，但这种方法无法扩展。在多个 GPU 平台（例如 NVIDIA Hopper 和 Blackwell、AMD MI300、Intel 或任何未来平台）上维护数百个内核很快就会变得不切实际。

相反，我们倾向于使用性能可移植的内核，能够自动适应它们运行的硬件。Triton 后端正是遵循了这一理念。

Triton 是一种特定领域语言，允许开发人员用 Python 编写 GPU 内核（例如矩阵乘法或 Attention）。这些内核会被编译为适用于多个平台的底层高效 GPU 代码。Triton 的分块 (tiled) 编程模型达到了一种平衡：它既足够底层，足以表达与硬件相关的优化，又足够高层，在很大程度上保持了硬件无关性。

如图 1 所示，开发人员以逻辑块（tiles）来描述计算。Triton 编译器和自动调优器（autotuner）确定这些块如何映射到底层硬件上。块的形状和执行布局在不同 GPU 之间可能存在显著差异，但这些决策是自动完成的，通常由自动调优引导（更多详细信息可在我们的论文 GPU 性能可移植性需要自动调优 (arxiv.org) 中找到）。

vLLM 中的 Triton Attention 后端

Attention 通常是大语言模型中性能最关键的操作。为了管理复杂性，vLLM 引入了一个名为 Attention 后端的抽象层，它将 Attention 实现隔离在一个通用 API 之后，并将其与线性层或层归一化等更简单的组件分离开来。

在此抽象层内，vLLM 支持多个 Attention 后端，包括 CUDA 平台上的 FlashAttention 和 FlashInfer、基于 ROCm 的 Attention 后端，以及用于 MLA 风格 Attention 的专用后端（完整列表见此处）。Triton Attention 后端完全由 Triton 实现，并且是 vLLM 的原生部分。

引入该后端是为了解决性能可移植性和依赖性问题。它在 NVIDIA、AMD 和 Intel GPU 上运行相同的源代码，仅依赖 PyTorch 和 Triton，并且始终作为 vLLM 的一部分提供。尽管最初由 IBM Research 和红帽 AI 开发，但现在它由更广泛的社区进行维护和扩展。

何时使用 Triton Attention 后端

Triton Attention 后端是在 ROCm 上运行的 AMD GPU 上的默认后端；在 Intel XPU 上，当运行 float32 格式时，vLLM 也会回退到 Triton Attention，因为 Flash Attention 在该平台不支持 fp32。它还支持需要特定功能（如 StepFun 音频模型使用的 ALiBi sqrt，或 sink tokens 和 GPT-OSS 行为）的模型，特别是在非 Hopper 架构的 NVIDIA GPU（如 A100）上。

此外，它还支持具有小头大小（head sizes）的模型、编码器和解码器 Attention 以及多模态前缀 Attention。由于 Triton Attention 后端始终存在，它也作为回退后端使用，以防 FlashAttention、FlashInfer 或其他依赖项不可用或无法导入。它还支持批处理不变性等功能。

在 Triton 中编写高性能、可移植的 Paged Attention 内核

在 Triton Attention 后端开发初期，内核首先在 vLLM 之外实现，并使用大量的微基准测试进行了评估。内核 API 被设计为符合 vLLM 的要求，但在端到端集成之前，性能调优是在隔离环境下完成的。

微基准测试对于理解预填充（prefill）密集型、解码（decode）密集型和混合工作负载，以及跨不同批处理大小和上下文长度的性能行为至关重要。

图 2 显示了具有代表性的微基准测试结果。x 轴代表总 token 数，y 轴代表延迟。不同的子图区分了仅预填充、混合和仅解码工作负载。这些结果表明，不同的内核变体在不同的方案下表现优异，没有单一配置能在所有场景下占据主导地位。

微基准测试通过揭示可能被系统级效应掩盖的内核级行为，是对端到端基准测试的补充。

回顾：Paged Attention 内核的作用

Paged attention 通过对 KV 缓存进行分页，以内存高效的方式实现 Attention。对于批处理中的每个查询，内核会处理每个查询 token。对于每个 token，它会遍历查询头和相应的 KV 头，然后遍历分页的 KV 缓存以计算 Attention 分数并应用值向量。

这种结构如图 3 所示。查询 token 排列在 x 轴上，查询头排列在 y 轴上，分页 KV 缓存的遍历构成了最内层循环。为清晰起见，省略了因果掩码 (causal masking) 和滑动窗口等细节。

有关内核底层优化的详细说明，我们推荐内核作者撰写的 PyTorch 博客：https://pytorch.ac.cn/blog/enabling-vllm-v1-on-amd-gpus-with-triton/

代码可在此处找到：https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/ops/triton_unified_attention.py

利用 Q 块优化 tl.dot 的分块大小

Attention 中的核心计算是矩阵乘法，在 Triton 中使用 tl.dot 实现。然而，高性能需要足够大的 tile 来充分利用硬件，而简单地加载分页 KV 缓存并不能带来良好的结果。

KV 侧的 tile 大小受限于 KV 缓存的页大小，因此优化重点在于查询侧。对于组查询 Attention（Group Query Attention），通过将与单个 KV 头关联的所有查询头放在一起处理，可以提高缓存重用率。为了进一步提高并行度，多个查询 token 被组合成一个工作项，称为 Q 块 (Q block)。

图 4 展示了这种方法。启动网格跨越批处理大小和 KV 头，而 Q 块决定了每个内核实例处理多少查询 token 和头。自动调优器会为每个平台选择合适的块大小。

通过并行分块 Softmax 增加并行度

同时处理多个查询 token 对于预填充工作负载效果良好，但对于仅处理单个查询 token 的解码工作负载则没有好处。为了解决这个问题，通过“3D 内核”即并行分块 Softmax 引入了额外的并行化。

这种方法将 KV 缓存的遍历拆分到多个内核实例中。每个实例计算部分结果，随后被归约 (reduced) 以产生最终输出。由于 Triton 不提供全局屏障 (global barrier)，这种归约需要启动第二个内核，从而在额外的并行度和启动开销之间引入了权衡。启发式方法被用于确定何时采用这种方法有益。

CUDA 图、启动网格与 GPU 执行波次

CUDA 图通过记录和重放固定的执行图来减少内核启动开销。然而，Attention 内核带来了挑战，因为它们的启动网格通常取决于批处理大小和序列长度。

GPU 使用固定数量的流式多处理器 (SM) 执行内核。当启动的线程数多于 SM 数时，执行会分波次进行。图 5 说明了这种行为，其中第二波次会导致利用率不足。

当在 CUDA 图中捕获时，即使有效工作负载大小减少，这种低效率也会被重放。图 6 显示了固定的启动网格如何导致额外的浪费工作和增加延迟。

从可变启动网格到持久化内核

早期版本的 Paged Attention 内核使用了随工作负载大小缩放的可变启动网格，如图 7 所示。虽然灵活，但这种方法与 CUDA 图的兼容性较差。

为了解决这个问题，我们设计了持久化内核（相关 vLLM PR 待处理）。启动固定数量的内核实例，等同于可用的计算资源。每个实例通过从 GPU 内存读取元数据，动态确定需要处理多少工作。这保持了启动网格恒定，并允许高效重用 CUDA 图。

基准测试结果

2025 年末的基准测试结果证明了这种方法的有效性。图 9 显示了 Llama 3.1 8B 在 NVIDIA H100 和 AMD MI300 上，批处理大小为 1、输入长度为 500 token 时的端到端延迟结果，输出长度在 x 轴上表示。

在 H100 上，Triton Attention 后端对于长解码请求达到了 FlashAttention 3 性能的 100.7%。在 MI300 上，它比早期实现提升了约 5.8 倍。重要的是，两个平台使用了相同的 Triton 内核源代码。请注意，Triton 中的 Paged Attention 实现大约有 800 行代码，而 FlashAttention 3 约有 70,000 行代码。

预览：Helion 中的 Paged Attention

Helion 是 PyTorch 团队开发的一种新的特定领域语言，可以被视为更高级的 Triton 或分块 PyTorch。作为实验，我们在 Helion 中实现了一个简化的 Paged Attention 内核，并取得了令人鼓舞的早期结果。该工作已发表在 PyTorch 博客上，代码以 草案 PR 的形式提供在 vLLM 仓库中。

总结

随着模型、推理优化和硬件平台的不断进步，性能可移植性变得日益重要。vLLM 中的 Triton Attention 后端证明了使用单一的、可移植的内核实现来获得业界领先的 Attention 性能是可能的。

通过精心的内核设计、广泛的微基准测试以及持久化内核和 CUDA 图等系统级优化，Triton 后端不仅匹配甚至超过了高度专业化的实现，同时保持了跨 GPU 厂商的可移植性。如今，它是 AMD 上的默认 Attention 后端，并使用相同的源代码在 NVIDIA 和 Intel 平台上高效运行。

虽然这篇博客概述了 Triton Attention 后端中最重要的优化，您可以在我们的相关论文 Triton Attention 内核的剖析 (arxiv.org) 中找到所有细节和更多基准测试结果。

致谢

这项工作由 IBM Research 的 AI 平台团队完成——感谢所有相关人员：Burkhard Ringlein, Jan van Lunteren, Chih-Chieh Yang, Sara Kokkila Schumacher, Thomas Parnell, Mudhakar Srivatsa, Raghu Ganti。