vLLM Triton Attention 后端深度解析

本文改编自 Red Hat 主办的 vLLM 办公时间 (Office Hours) 会议，由来自 IBM Research 的 Burkhard Ringlein 主讲，深入剖析了 vLLM Triton Attention 后端。您可以浏览往期话题，并在此处加入未来的办公时间会议。

在过去的一年中，来自 IBM Research、Red Hat 和 AMD 的团队开发并向主干贡献了一个基于 Triton 的 Attention 后端，旨在实现卓越的性能，同时在各类 GPU 供应商之间保持强大的可移植性。这项工作的驱动力在于加速器硬件的多样性日益增长，以及维护大量高度专业化核函数的成本不断上升。

本文深入探讨了这一工作。我们阐述了为何 Triton 是 vLLM 的绝佳选择，描述了 Triton Attention 后端及其应用场景，随后深入讲解高性能 Paged Attention 核函数的实现。在此过程中，我们涵盖了核函数级别的优化、并行策略、CUDA 图交互和基准测试结果，最后简要介绍了 Helion。

为什么 Triton 对 vLLM 有所助益

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

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

相反，我们更倾向于性能可移植的核函数，它们能自动适应所运行的硬件。Triton 后端正是采用了这种方法。

Triton 是一种领域专用语言，允许开发者用 Python 编写 GPU 核函数（如矩阵乘法或 Attention）。这些核函数被编译成跨平台的高效 GPU 代码。Triton 的分块编程模型在二者之间取得了平衡：它既底层到足以表达与硬件相关的优化，又足够高层，基本保持了硬件无关性。

如图 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 和 Red Hat AI 开发，但现在由更广泛的社区进行维护和扩展。

何时使用 Triton Attention 后端

Triton Attention 后端是运行在 ROCm 上 AMD GPU 的默认后端；在 Intel XPU 上运行 float32 时，vLLM 会退回到 Triton Attention，因为 Flash Attention 在那里不支持 fp32。它还支持需要特定功能的模型，例如 StepFun 音频模型使用的 ALiBi sqrt，或在较旧的 NVIDIA GPU（如 A100）上使用 Sink Token 和 GPT-OSS 行为的模型。

此外，它支持小头大小（head size）、编码器和解码器 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。对于批次中的每个 Query，核函数处理每个 Query Token。对于每个 Token，它会遍历 Query 头和对应的 KV 头，然后遍历分页的 KV 缓存以计算 Attention 分数并应用 Value 向量。

图 3 展示了该结构。Query Token 排列在 X 轴上，Query 头在 Y 轴上，分页 KV 缓存的遍历构成了最内层循环。为清晰起见，省略了因果掩码和滑动窗口等细节。

关于核函数底层优化的详细解释，我们建议参阅核函数作者撰写的 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 块（Q Blocks）优化 tl.dot 的分块大小

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

KV 侧的分块大小受 KV 缓存页大小的限制，因此优化重点在于 Query 侧。对于分组查询 Attention（Group Query Attention），通过同时处理与单个 KV 头相关的所有 Query 头，可以提高缓存重用率。为了进一步提高并行性，多个 Query Token 被组合成一个工作单元，称为 Q 块。

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

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

一次处理多个 Query Token 对于预填充工作负载非常有效，但对于仅处理单个 Query Token 的解码工作负载则没有益处。为了解决这个问题，我们通过并行分块 Softmax（所谓的“3D 核函数”）引入了额外的并行化。

这种方法将 KV 缓存的遍历拆分到多个核函数实例中。每个实例计算部分结果，随后归约产生最终输出。由于 Triton 不提供全局屏障，此归约需要启动第二个核函数，从而在增加并行性和启动开销之间产生权衡。启发式方法被用来决定何时采用这种方法。

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

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

GPU 使用固定数量的流式多处理器（SM）执行核函数。当启动的线程数多于 SM 数量时，执行会以波次（Waves）形式进行。图 5 展示了这种行为，其中第二波次会导致利用率不足。

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

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

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

为了解决这个问题，我们设计了持久化核函数（Persistent Kernels，正在提交至 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 博客上发表，代码可在 vLLM 仓库的草案 Pull Request 中查看。

总结

随着模型、推理优化和硬件平台的不断进步，性能可移植性变得愈发重要。vLLM 中的 Triton Attention 后端证明，使用单一的、可移植的核函数实现，完全可以达到最先进的 Attention 性能。

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

虽然本篇博客概述了 Triton Attention 后端中最重要的优化，但您可以在我们相关的论文《Triton Attention 核函数解析》(The Anatomy of a Triton Attention Kernel, arxiv.org) 中找到所有细节和更多的基准测试结果。

致谢

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