vLLM Triton 注意力后端深度解析

vLLM 引入了基于 Triton 的注意力后端，以在多种 GPU 硬件上提供高性能、可移植的注意力 kernel。该后端使 vLLM 能够维护一个单一的注意力操作代码库，在 NVIDIA、AMD 和 Intel GPU 上高效运行，降低了硬件特定 kernel 的维护负担，同时匹配专用实现的性能。

通过 Triton 实现性能可移植性

为每种 GPU 架构（如 NVIDIA Hopper、Blackwell 和 AMD MI300）维护数百个专用 kernel 并不现实。vLLM 使用 Triton，这是一种领域特定语言，允许使用 Python 编写 GPU kernel，并编译为适用于多平台的高效代码。

Triton 采用 tiled 编程模型，开发者定义逻辑计算块。Triton 编译器和自动调优器随后将这些块映射到硬件特定的执行布局。此抽象使后端能够保持硬件无关性，同时仍可进行底层优化。

Triton 注意力后端架构

为了将注意力实现与线性层等其他组件隔离，vLLM 使用注意力后端抽象层。Triton 注意力后端是 vLLM 原生的，仅依赖 PyTorch 和 Triton，并完全使用 Triton 实现。

使用场景

Triton 后端在以下场景中部署：

AMD GPU 的默认后端： 它是运行在 ROCm 上的 AMD GPU 的主要后端。
Intel XPU： 用于 float32 操作，因为 FlashAttention 在该平台上不支持 fp32。
特定模型特性： 支持 ALiBi sqrt（StepFun 音频模型使用），sink tokens，以及 GPT-OSS 行为，尤其在 Hopper 之前的 NVIDIA GPU（如 A100）上。
专用需求： 处理头尺寸较小的模型、编码器/解码器注意力以及多模态前缀注意力。
回退机制： 当 FlashAttention、FlashInfer 或其他依赖不可用时，作为通用回退后端。

分页注意力的技术实现

分页注意力通过对 KV 缓存进行分页来优化内存。kernel 处理查询 token，遍历查询和 KV 头，并遍历分页的 KV 缓存以计算注意力得分。

Q 块优化

为了最大化 tl.dot（Triton 的矩阵乘法）的利用率，后端使用 “Q 块”。由于 KV 缓存页大小限制了 KV 端的 tile 大小，kernel 将多个查询 token 和头组合成一个工作项（即 Q 块），以提升并行度和缓存复用，尤其针对 Group Query Attention（GQA）。

并行 tiled Softmax（3D Kernel）

虽然 Q 块对预填充工作负载有益，但解码工作负载仅处理单个查询 token。为此，vLLM 实现了使用并行 tiled softmax 的 “3D kernel”。该方法将 KV 缓存遍历拆分到多个 kernel 实例中。每个实例计算部分结果，然后通过第二次 kernel 启动进行归约，生成最终输出。

持久 Kernel 与 CUDA Graphs

CUDA 图通过记录固定的执行图来降低启动开销。然而，标准注意力 kernel 通常根据批大小和序列长度使用可变的启动网格，这在通过 CUDA 图重放时会导致低效（例如工作浪费或流式多处理器利用率不足）。

为了解决此问题，vLLM 开发了 持久 kernel。不使用可变网格，而是启动固定数量的 kernel 实例——等于可用计算资源的数量。这些实例动态读取 GPU 内存中的元数据以确定工作负载，从而保证固定的启动网格并实现高效的 CUDA 图复用。

基准测试与结果

2025 年底的基准测试表明，Triton 后端在代码复杂度远低于专用方案的情况下实现了高效率。Triton 分页注意力实现约有 800 行代码，而 FlashAttention-3 大约有 70,000 行。

性能指标（Llama 3.1 8B，500 输入 token，批大小 1）：

NVIDIA H100： 在长解码请求中达到 FlashAttention-3 性能的 100.7%。
AMD MI300： 相比之前的实现提升约 5.8 倍。

未来方向：Helion

实验性工作已在 Helion 上启动，Helion 是 PyTorch 团队推出的新领域特定语言，被描述为更高级的 Triton 或 tiled PyTorch。使用 Helion 实现的简化分页注意力 kernel 已展示出有前景的早期结果，当前以草案 Pull Request 形式在 vLLM 仓库中提供。