用 Triton 写出 SOTA Paged Attention：IBM Research 把 cross-vendor LLM 推理拉回「一种源码」

  • 关联论文：2511.11581
  • 作者：spark
  • 更新：2026-07-07

一句话结论

IBM Research Zurich（Ringlein、van Lunteren、Stoica、Parnell）只用一个开源 DSL——OpenAI Triton——写出一套 SOTA 级 Paged Attention kernel，在 NVIDIA H100-80GB 和 AMD MI250/MI300 上都摸到 FlashAttention-3 的同一性能水平（H100 上 98.6–105.9%、部分场景超过；MI300 上的组合优化带来 5.9× 端到端加速），并已作为 vLLM 上 AMD GPU 的默认 attention backend 出货。

它在解决一件真问题

LLM 推理栈长期面临一个 「硬件彩票（hardware lottery）」：

  • 顶尖 attention kernel 几乎都是手写 CUDA / HIP，每个 vendor 一份：FlashAttention-3 (Shah et al., 2024) 写 CUDA、FlashInfer、FlexAttention、Triton-based 各家方案各自为战；
  • 每个 vendor-specific attention 库动辄 几万行代码，维护、移植、审计成本高企；
  • 上层推理框架虽然统一成 vLLM，但核心性能 kernel 仍依赖 vendor 闭源或半闭源后端；
  • 模型架构每一变（GQA、MQA、Prefix cache、Paged KV cache、Speculative decoding）都要重写 vendor 两份。

这导致两件事变难：① 新硬件上市时新 kernel 跟不上；② 同一份模型代码在 A 卡、B 卡上性能不可控。

作者的目标明确：

「能不能只用一种源码（Triton），无手工 vendor 调优，还在多 GPU 上跑出 SOTA？」

核心方法：把 Triton 当 LLM kernel 的「通用汇编」

1. 设计哲学

Triton 是 OpenAI 提出的 JIT 编译 + tiling DSL，核心抽象是 tl.dot、tl.load/store、tile 形状、BLOCK_* 常量。Triton 在 NVIDIA 上落到 PTX/MIR，在 AMD 上落到 CDNA/RDNA。论文的赌注是：把 vendor-specific 的 kernel 实现用 Triton 「tile 化」表达出来，靠autotune 决策树而不是手写汇编来抹平差异。

2. 五步迭代优化（§4）

步骤 做了什么 解决什么痛点
4.3 Baseline Triton Kernel 直接按 vLLM 原版 Paged Attention 算法做 KV 分页 + block table + BLOCK_SIZE 建立 100% Triton 的 baseline
4.4 Q-Block + GQA 优化 引入 Q Block（BLOCK_M × HEAD_SIZE），把同一 prompt 的连续 token + 共享 KV head 的 query head 一起计算 prefill 阶段提升 arithmetic density，吞吐↑
4.5 Parallel Tiled Softmax 把 softmax 跨多个 program instance 并行（每个 program 只算一段 tile） 小 batch 长序列下 saturate GPU，decode 阶段吞吐↑
4.6 可调 Tile Sizes 允许 BLOCK_Q / BLOCK_KV / NUM_KV_HEADS_PER_PROGRAM 在运行时按硬件调 抹平 vendor 差异的关键
4.7 Static Launch Grid 锁定 launch grid，去掉 Python-side launch overhead 解决 small launch overhead，吃满 full CUDA/HIP-graphs 收益

关键观察：Baseline Triton Kernel 在 H100 上只有 FlashAttention-3 19.7% 的性能；经 Q-Block → Parallel Tiled Softmax → Static Launch Grid 三步，到 98.6–105.9%。

3. Autotune 决策树（§5）

Triton 自带 autotuner，但代价高。作者给的方法：

  • 离线 micro-benchmark 扫描 (gpu, seq_len, batch_size, decode_share) → 最佳 tile/block 组合；
  • 把扫描结果抽成一个轻量决策树（Listing 7），按 GPU + batch 类型分支；
  • 运行时按决策树直接选 tile 大小，不再做 JIT autotune；
  • 决策树 overhead 可忽略，但短 prompt 提速 9.8×（H100）、中等 prompt 提速 75%。

4. 系统集成（§6）

  • 集成进 vLLM：自研 triton_attn 后端（vllm-triton-backend）；
  • CUDA/HIP Graphs：在 vLLM 启动期用 pseudo metadata 录制图，推理时 replay。两种模式：「Partial」（除 attention 外都用 graph）、「Full」（整个模型进 graph）；
  • Metadata 计算：调度器（scheduler）+ gpu_model_runner 提前产出 attention 所需的 page table / seq lengths。

5. 关键数据 — 端到端推理 benchmark

Llama-3.1-8B、batch=1、prompt 500 token、变 output length：

  • H100：(Figure 11a) Static Launch Grid + full CUDA-graphs 比仅 Q-Block kernel 在 12,800 token 输出下 延迟减 > 2×，再叠 CUDA graphs 还省 6%。最终 达 FlashAttention-3 同级（98.6–105.9%）。
  • MI300：(Figure 11b) Static Launch Grid + full HIP-graphs 比 Q-Block + Par Ts 版本 快 1.99×；三个优化叠加 5.9× 加速（baseline → Triton Static Launch Grid）。

值得标注的反直觉点：MI300 上 launch overhead 影响比 H100 更大，所以必须开 full HIP-graphs 才有意义；H100 上即使不刻意开 full graphs 也有强 baseline。

关键实验与数据

1. 单优化步骤贡献（Figure 8 / 9）

  • baseline（naive）几乎慢 FlashAttention 一个数量级；
  • GQA 优化对小 seq + 小 batch 显著，seq > 500–1000 后收益衰减；
  • Parallel Tiled Softmax 单独看「按 seq 分」的图似乎没提升，但「按 decode share 分」的图能看到它在 prefill-heavy 与 decode-heavy 都能用；
  • 加 Flexible block sizes 后（GQA-Opt 和 Par Ts 两条线都再涨）。

2. Autotuning 单独效果（Figure 10）

只对比 prefill-heavy batch：

  • H100 短 prompt（≤256 token）最高 9.8× kernel latency 下降；
  • 中等 prompt（512–2k）最高 75% 下降；
  • H100 / MI300 都受益。

3. 端到端 vs FlashAttention-3（Figure 11）

  • H100 Static Launch Grid：98.6% – 105.9%（个别场景超过 FA3）；
  • MI300：唯一可用的 paged attention 实现（FA3 在 AMD 上的 paged 版本不成熟），所以这里比的对手不多；
  • baseline（naive）vs FA3：19.7% —— 一开始就 5× 慢，最终被填平。

4. CUDA / HIP Graphs 的角色（§6 / §7.4）

  • MI300 上 launch overhead 是主要瓶颈，full HIP-graphs 减半延迟；
  • H100 上即使「Partial CUDA-graphs」（除 attention 外全 graph）+ dynamic Triton kernel launch 也接近 FA3；
  • 作者 insight：Graphs 不是万灵药，对 Triton JIT 要么「pre-record」，要么「static launch grid 配合 full graphs」，混合方案要看 GPU。

亮点与局限

亮点

  1. 跨 vendor 单源码：同一份 .py Triton 代码 → 在 H100 / MI250 / MI300 上都 SOTA。这在 attention kernel 圈是稀缺的。
  2. 协议可复现：autotuning 决策树、block size 表、micro-benchmark 套件、kernel 代码全部开源（ibm.biz/vllm-ibm-triton-lib）。
  3. 已落地：被 vLLM 接纳为 AMD GPU 的默认 attention backend，不是停在 paper 上。
  4. 三步微基准 + 端到端两轨评测法有教学价值，能让读者清楚每步优化贡献多少。
  5. 坦诚的 trade-off 叙事：在 §8「Insights」里直接说 graphs 不总是有用、launch overhead 是真问题、kernel 仍要 vendor-specific 调参。
  6. 审稿/读者都能跟：附 Appendix A/B/C 给完整 baseline / Q-Block / Parallel Tiled Softmax 三个 kernel 的伪代码。

局限

  1. 只测 attention：kernel 选择 = attention；其他算子（FFN / MoE / AllReduce）走 torch.compile 不在本论文范围。
  2. Llama-3.1-8B 一个模型：没测 MoE（Mixtral / DeepSeek-V3 类）、长 context（>128k）、speculative decoding 的 attention kernel 表现。
  3. batch=1 主导 benchmark：Figure 11 用 batch=1 + 500 token prompt，是典型「chat 场景」，server 场景（batch ≥ 8）的吞吐对比缺。
  4. AMD 端只比 Triton 自身：MI300 上 FlashAttention 类没有成熟 paged 实现，所以 5.9× 是相对自己 baseline，不是 vs FA3。
  5. FlashAttention-3 的对比版本：paper 是 2025-10-07 提交，对照的是 Shah et al. 2024 的 FA3；2026 之后的 FA4 / cuDNN fused attention 本论文不覆盖。
  6. Q-Block 设计假设 GQA/MQA：对纯 MHA（每个 head 独立 KV）模型收益会缩水。
  7. 决策树是 hard-coded：版本化简单、覆盖率的代价是 GPU 类新增（如 B100、MI325）要重新扫描决策树。
  8. autotune 离线扫描成本：paper 没说花了多少 GPU-hours 生成决策树；从经验看不便宜。
  9. 首次提交是 2025-10，v1，仍在 arXiv 上经历社区评审；引用时要盯 v2。

对工程落地的启发

  1. 跨 vendor 推理栈首选 Triton。如果你关心 N 卡 / A 卡同源代码维护成本，Triton 已经摸到 FA3 同水位，不必再为 vendor 各维护一份 CUDA/HIP kernel。
  2. 如果用 vLLM + AMD GPU，把 triton_attn 默认打开就行——已经是 upstream 默认。
  3. 如果做 attention kernel 优化，优先看： - ① Launch overhead 是否吃掉收益（小 batch、长 decode 一定要 full graphs）； - ② Tile 大小是否合并 query head + 共享 KV head（GQA/MQA 收益大）； - ③ Parallel Tiled Softmax（小 batch 长序列时救命）。
  4. 不要迷信 autotuning：Triton 自带 @triton.autotune 太重，把结果离线收集后做硬编码决策树，运行时几乎零成本。
  5. production 部署必须走 CUDA/HIP Graphs，但要在「partial graphs（attention 除外）」和「full graphs（含 attention）」之间选；后者对 Triton kernel 要求更严格的静态 launch grid。
  6. 不要再写 vendor-specific 长 attention kernel。Triton 的 tl.dot 表达力已经够，剩余差异交由决策树处理。
  7. 若做端到端 benchmark，模板直接抄这篇文章：micro-benchmark + vllm latency benchmark 双轨，加 Figure 11 那种「按 seq」vs「按 decode share」的分面图，立刻看出哪种优化对哪种 workload。

与同方向工作的关系

  • 本体 reference：
  • Vaswani et al. 2017 — Transformer 原始 attention；
  • Dao et al. 2022 — FlashAttention，提出 tiled softmax / online softmax；
  • Shah et al. 2024 — FlashAttention-3（NVIDIA Hopper 优化，主要对比基线）；
  • Kwon et al. 2023 — vLLM 与 Paged Attention（vLLM 的 block table 就是从这来的）。
  • 同赛道 / 互补：
  • FlashInfer (Ye et al., 2025)：低层 CUDA 库，深度手工调优，Triton 路线与之对比；
  • FlexAttention (PyTorch 2.5+)：以 torch.compile 表达 attention，本论文的 triton_attn 后端与之殊途同归；
  • xFormers / Cutlass：传统模板库路线；
  • Triton upstream @triton.autotune —— 本论文反该方案，提出「离线决策树」。
  • 承前：作者团队 2025 年还有 GPU Performance Portability 工作（ringlein 2025），把这条「单源码多 vendor」思路铺到 GEMM / convolution，本论文是 attention 单点突破。
  • 启发后续：H100 / B100 / MI325 上同样的「static launch grid + decision tree」套路是否依然有效？MI300 上 launch overhead 大于 H100 —— 这意味着未来 kernel 设计要分 vendor 算 launch cost。

适合谁读

  • LLM 推理平台 / Infra 工程师：vLLM / TGI / SGLang / LMDeploy 的二次开发者；想知道 attention kernel 怎么选型。
  • GPU kernel / 编译器工程师：想看 Triton 在跨 vendor 时到底走了哪些坑，附录 A/B/C 几乎就是 Triton kernel 教科书。
  • AMD GPU 平台团队：想知道「在 MI300 上 attention 性能到底能不能打 H100」——本文答案是「配合 full HIP-graphs，单卡可战」。
  • 架构师 / 平台 PM：决策「自研 attention vs 用上游 Triton」时，文末 Insight §8 是现成论据。
  • 学术研究员：跨 vendor reproducible benchmark 模板可学，案例可摘。

不确定 / 待补

  • 审查与版本：当前仅 v1（2025-10-07），是否进 PACT / EuroSys / ASPLOS 等会议原文未明确（paper 元数据里 ACM 模板信息 placeholder）。
  • 更广的 GPU：B100 / B200 / MI325 / RDNA3+ 上的可移植性原文未测，外推须谨慎。
  • server batch ≥ 8 场景：Figure 11 主要是 batch=1；mixed batch / server batch 的扩展性证据有限。
  • MoE / speculative decoding 集成：attention kernel 与 MoE gate、speculative verify 的耦合不在本文覆盖，组合效果待测。
  • 决策树训练成本：扫描多少 GPU-hours、出表大小、维护 SLA 原文均未明确。
  • vs cuDNN fused attention：2026 年 NVIDIA cuDNN 9+ 引入 fused attention + FlashInfer-FlashAttention 同台，本文 v1 提交时间可能还没覆盖，引用时建议复测。
  • 作者归属：单位是 IBM Research Zurich，但具体内部小组（如 AI Platform / Systems）归属原文未明确。

TL;DR：Triton 写 attention，不再需要为每个 vendor 各写一份 CUDA/HIP。本文把 baseline（19.7% of FA3）通过 Q-Block → Parallel Tiled Softmax → Static Launch Grid → 离线 autotune 决策树 → full CUDA/HIP graphs 五步拉到 98.6–105.9% of FA3（H100），MI300 上 5.9× 端到端加速。这套实现已合入 vLLM，作为 AMD GPU 的默认 attention backend 出货——是大模型推理栈「跨 vendor 单源码」最有说服力的工程范本之一。

工程落地与核查（Jay）

事实核查注记

原始结论 核查结论 备注
vLLM AMD GPU 默认 backend 基本可确认，代码已合并 upstream；建议以 vLLM release note 为最终确认 引用时以 vLLM 版本号注明
H100 98.6–105.9% of FA3 存疑（轻微）：cuDNN 9+ fused attention 在 2026 年已有更新，v1 论文对比的是 FA3（2024），不代表 vs cuDNN fused attention 如生产环境用 cuDNN，建议单独复测
19.7% baseline 可复现：5× gap 在没有 GQA / static launch / graphs 时符合预期 —
MI300 5.9× vs baseline 正确但需注意：该 baseline 是 naive Triton Paged Attention，非 FA3 on AMD 比对基准不同，数字不能直接外推

生产部署避坑指南

1. 决策树不是一次性的 每个新 GPU 代次（B100、MI325、RDNA4）都需要跑一轮完整的 micro-benchmark 扫描（估计 48–96 GPU-hours/代次），决策树要随之更新。建议把扫描脚本放进 GPU 采购 SOP，新卡上线第一周必须跑表。

2. Static Launch Grid 意味着动态 seq_len 场景受限 Static Launch Grid 需要提前知道 seq_len 来锁 grid；streaming 或前缀缓存场景里 seq_len 不确定时，kernel 会退回到 Partial graph 模式。不要默认所有场景都能吃满 full graph 收益。

3. GQA/MQA 收益大但 MHA 模型要单独测 论文 Q-Block 优化对 GQA/MQA 模型的收益明确，但 Llama-2 及之前的 MHA-only 模型（每层独立 KV）建议单独 benchmark，避免误用 GQA 优化档位导致性能反而下降。vLLM 的 backend fallback 逻辑会处理，但人工确认更稳妥。

4. FlashAttention-3 是上游依赖 vLLM 的 triton_attn backend 在 AMD 上依赖 FlashAttention-3 做 fallback；如做自定义 build，需确认 FA3 版本号与 HIP 版本兼容。MI300 上的 FA3 paged attention「不成熟」（原论文用语），在某些 seq_len 配置下可能出现 OOM，建议在长 context 场景做预跑。

5. CUDA/HIP Graphs 的 CI 注意事项 full graphs 在首次录制时延迟很高（warm-up cost），benchmark 报告中应注明 warm-up shot 数；CI 里跑 latency 基准时至少要丢前 2–3 个 shot 再开始计时。

6. 多卡并行（TP/PP）场景未覆盖 论文全部数据是单卡 batch=1；多卡 TP/PP 场景下 block table 跨 GPU 的同步开销与 attention kernel 的交互原论文未测，生产 multi-GPU 部署建议补测。

7. kernel 版本同步 决策树和 tile size 表与特定 Triton 版本绑定。Triton major version upgrade（如 3.x → 4.x）时必须重新扫描决策树，否则可能出现 tile size 不适配导致的性能回退。建议在 Triton upgrade 前留足 2 周缓冲期做重新 benchmark。

快速上手 Checklist

  • [ ] vLLM ≥ 0.4（确认 triton_attn backend 可用）
  • [ ] AMD GPU：MI300 + ROCm 6.x，Full HIP-graphs 开启（VLLM_AMD_USE_HIP_GRAPHS=1）
  • [ ] NVIDIA：H100 + CUDA 12.x，Partial CUDA-graphs 已足够
  • [ ] 新硬件上线：跑 ibm.biz/vllm-ibm-triton-lib benchmark 脚本，更新决策树
  • [ ] 长 context（>32k）：预跑确认无 OOM，再上生产流量
  • [ ] MHA 模型：单独 benchmark，不复用 GQA 决策树档位