突破硬件壁垒：基于Triton的跨平台Attention内核实现5.9倍推理加速，性能达SOTA 105.9%

我们所研究的优化方法累计实现了高达 589%的性能提升 ，并已将相关内核与框架作为开源项目贡献（ ibm.biz/vllm-ibm-triton-lib ）。最终，我们开发的高度优化内核已成为 vLLM 中 AMD 部署的默认注意力后端。

关键词：Triton、Attention Kernel 、Portability 、Large Language Models、Autotuning

• The Anatomy of a Triton Attention Kernel
• 副标题：How to achieve cross-platform state-of-the-art LLM attention using only Triton
• https://arxiv.org/pdf/2511.11581
• https://github.com/foundation-model-stack/vllm-triton-backend/tree/main/ibm-triton-lib
• 1.5 万字，阅读 60 分钟，播客 20 分钟

本文聚焦于解决大语言模型（LLM）推理中硬件依赖与性能优化的核心矛盾，提出一种基于 Triton 领域特定语言的跨平台分页注意力内核方案。该方案旨在突破“硬件 lottery”限制，实现无需底层手动调优、却能在 NVIDIA 和 AMD GPU 上达到顶尖性能的 LLM 推理部署。

图5. vllm-triton-backend中的调优工作流程

现有 LLM 推理框架如 vLLM，依赖大量平台专属闭源库，导致部署复杂且硬件适配成本高。本文通过 Triton 的即时编译特性，结合算法优化、参数自动调优与系统级整合，将通用 Triton 注意力内核的性能从顶尖水平的 19.7%提升至 105.9%。

核心贡献包括：

• 设计功能完备 且性能顶尖 的 【跨平台】分页注意力内核；
• 总结算法优化、编程模型选择 及系统级权衡 等 【可复用经验】；
• 开源内核与微基准测试套件 ，并将其集成到 vLLM 框架，成为 AMD GPU 的默认注意力内核。

图2. vLLM中triton_attn后端的概述

关键技术优化涵盖 Q 块结构设计、并行分块 Softmax、可调节块大小及静态启动网格 等，同时通过 微基准测试驱动的启发式自动调优 ，平衡了调优开销与跨平台适配性。

实验表明，该内核在 NVIDIA H100 和 AMD MI300 GPU 上，对 Llama-3.1-8B 模型的推理性能 媲美甚至超越 FlashAttention 等专用库 ，最长序列场景下提速达 5.9 倍，为 LLM 推理的 硬件无关化部署 提供了高效解决方案。

本文目录

• 关键问题
  • 问题 1：非 GPU 架构适配性与重新调优成本问题
  • 问题 2：静态启动网格在高频小批量短序列场景的资源浪费问题
• 一、引言
• 二、背景与相关工作
  • 2.1 vLLM：一种推理服务器
  • 2.2 Triton：一种瓦片式领域特定语言（tiling DSL）
  • 2.3 可移植性与内核参数自动调优
  • 2.4 注意力内核与库
• 三、解决方案概述
• 四、注意力机制的设计与实现
  • 4.1 核心概念
  • 4.2 术语定义
  • 4.3 基准 Triton 注意力核函数
  • 4.4 预填充与 GQA 优化
  • 4.5 并行分块 Softmax
  • 4.6 可调整分块大小
  • 4.7 静态启动网格
• 五、自动调优与可移植性
  • 5.1 当前自动调优的不足
  • 5.2 vllm-triton-backend 中的自动调优应用
• 六、推理服务器与系统集成
  • 6.1 元数据计算
  • 6.2 Triton 启动开销与 CUDA/HIP 图
• 七、性能评估
  • 7.1 实验方法
  • 7.2 优化步骤评估
  • 7.3 自动调优方法评估
  • 7.4 端到端评估
• 八、见解与未来工作
  • 8.1 Triton 内核需“场景专用”
  • 8.2 tl.dot 的使用
  • 8.3 CUDA/HIP 图并非万能
  • 8.4 未来工作：内核调优中纳入 GPU 特定参数
  • 8.5 未来工作：自动调优与 CUDA/HIP 图的结合
• 九、结论
• 参考文献

关键问题

问题 1：非 GPU 架构适配性与重新调优成本问题

论文【仅】验证了 NVIDIA H100 和 AMD MI300 系列 GPU 的性能，而 LLM 部署还涉及 TPU、自研加速器等多元架构，其提出的自动调优启发式算法和内核设计，在非 GPU 架构上是否仍能维持 SOTA 性能？ 且现有自动调优虽通过微基准测试降低开销，但 面对全新架构时，重新调优的时间成本是否会抵消跨平台优势？

论文提出的方案暂【未验证非 GPU 架构】的 SOTA 性能，但启发式调优框架可降低重新调优成本，仍存在一定适配开销。

1. 论文仅在 NVIDIA H100 和 AMD MI300 等 GPU 上完成了验证。虽然提及 vLLM 支持 TPU、AWS Inferentia 等非 GPU 加速器，但并未将 Triton 注意力内核部署于这些架构，因此无法确认其在这些平台上是否仍能维持 SOTA 性能。
2. 在重新调优成本方面，论文采用了“微基准测试+启发式决策树”的方案。相比传统逐场景调优（如 FlashAttention v2 需 24 小时/GPU），仅需针对新架构的平均与极端场景进行调优，即可通过决策树覆盖多数场景，大幅降低了时间成本。
3. 然而，全新架构仍需重新构建微基准测试场景、适配硬件特性（例如 TPU 的张量核心调度），这仍存在初始适配开销。论文未提供针对非 GPU 架构的具体调优耗时数据。

问题 2：静态启动网格在高频小批量短序列场景的资源浪费问题

论文为兼容 CUDA/HIP graphs 而采用了静态启动网格，虽降低了启动开销，但牺牲了内核对动态批次大小、序列长度的适配灵活性。在真实 LLM 推理的高频小批量、短序列场景中，这种静态配置导致的 GPU 资源浪费（如多余实例启动），是否会让其实际性能反超 FlashAttention 的优势消失？

静态启动网格确实会带来少量资源浪费，但通过启发式内核选择与网格规模优化，在真实推理场景中并未抵消其性能优势，在短序列场景中仍具备竞争力。

静态启动网格的设计初衷是兼容 CUDA/HIP graphs，以解决 Triton 内核 100-300µs 的启动开销问题。其网格规模被限定为“接近但小于 GPU 核心数”，以减少多余实例启动带来的 GPU 调度损耗。

针对高频小批量短序列场景，论文通过启发式策略动态选择内核：在短序列时优先使用基础优化内核（而非并行 tiled softmax 内核），避免额外的归约内核启动开销。同时，决策树会根据序列长度调整 tile 大小，进一步降低资源浪费。

实验验证显示，在 H100 和 MI300 上，短序列（如输出长度 < 1000 tokens）场景中，静态启动网格内核的延迟仅比 FlashAttention V3 高约 1%-5%，且优于未优化的 Triton 内核，说明资源浪费的影响得到了有效控制。

一、引言

近年来，大型语言模型取得了显著发展。除了模型架构和训练流程的改进外，针对现代硬件优化 LLM 应用的创新也层出不穷[7, 36, 21, 44, 8]。

然而，功能和性能上的竞争导致了新人工智能（AI）或机器学习（ML）范式面临“硬件彩票”困境[14]（指 AI/ML 技术的性能高度依赖特定硬件，若技术适配的硬件并非主流平台，其推广和应用会受到严重限制），并使得技术发展向最主流的硬件平台倾斜。AI 算法与 AI 硬件之间的紧密关联，限制了 AI 的部署场景和应用范围——因为大多数功能仅支持少数特定硬件或特定输入问题规模[14]。因此，要实现具有最先进（state-of-the-art, SOTA）性能的 LLM 部署，所需库的数量和规模都大幅增长。

尽管我们欢迎这种专用库“寒武纪大爆发”式的发展，但这一趋势为 AI 民主化[4]带来了额外风险：它构建了一个难以驾驭的依赖“丛林”，而这些依赖往往涉及闭源专有工具（例如[44]）。对第三方库的依赖严重阻碍了新硬件和新 AI 应用的落地。此外，为了移植一行内核代码[41]而需要编写数万行代码的情况，不仅拖慢了研究进度，还不必要地复杂化了新 ML 方法的部署流程。

工业界和学术界长期以来的一个目标是开发完全平台无关的软件栈[4]。

在推理服务器层面，开源框架 vLLM[21, 42]正获得巨大发展动力，其愿景是成为 LLM 推理的事实标准引擎。vLLM 掩盖了部署现代 LLM 的诸多复杂细节，同时拥有活跃且充满活力的开源社区。然而，vLLM 的 LLM 内核编译和运行仍依赖许多专有或闭源依赖项，而这些依赖项是平台和厂商特定的。

图1. vLLM v1的架构以及跨厂商性能可移植性问题：vLLM 支持多种平台，但对于性能关键的内核，它仍依赖以“后端”形式封装在 vLLM 中的外部库。对于许多 LLM 的“核心”——即所谓的（自）注意力层[41]而言，情况尤其如此，该层通常是 LLM 中性能最关键的内核[7, 36, 21, 44, 8]。这些平台特定的注意力库通常包含数万行代码[33, 36]

如图 1 下半部分所示，这一困境引出了一个重要问题：是否有可能设计一种完全平台无关的 LLM 实现，使其在多个平台上都能达到最先进性能？ 如果可能，开发性能关键型、平台无关内核需要哪些关键步骤？本研究的目标就是回答这些问题。

近期的研究工作[33, 11]表明，有可能回归到单源开源库（指通过一个代码库适配多种平台，而非为不同平台编写多个代码库），同时整合所有 LLM 内核的最先进创新。在本研究中，我们延续这一研究方向，阐述并展示如何开发一款可投入生产、开源、跨平台且性能最先进的 LLM 注意力内核。我们的方法是基于 OpenAI Triton 领域特定语言构建，并展示实现跨平台性能关键型内核所需的关键步骤，以达到最先进性能。

我们的贡献总结如下：

1. 展示了一个功能完整、跨平台的分页注意力内核（paged attention，一种通过“分页”方式管理 KV 缓存，减少 GPU 内存消耗的注意力实现方式），且性能达到最先进水平。
2. 全面总结了关键经验教训，包括算法优化、所采用的编程与内存模型，以及具有广泛适用性的系统级权衡策略。
3. 开源了我们的内核和微基准测试套件（ibm.biz/vllm-ibm-triton-lib），并将这些内核集成到最广泛使用的推理框架之一 vLLM（vllm.ai）中，该内核已被采纳为 AMD GPU 的默认注意力内核。

二、背景与相关工作

2.1 vLLM：一种推理服务器

通常，LLM 通过服务或推理框架（如 vLLM[21, 42]）部署，这些框架掩盖了模型部署和请求调度的诸多细节。

图1. vLLM v1的架构以及跨厂商性能可移植性问题：vLLM 支持多种平台，但对于性能关键的内核，它仍依赖以“后端”形式封装在 vLLM 中的外部库。对于许多 LLM 的“核心”——即所谓的（自）注意力层[41]而言，情况尤其如此，该层通常是 LLM 中性能最关键的内核[7, 36, 21, 44, 8]。这些平台特定的注意力库通常包含数万行代码[33, 36]

图 1 展示了 vLLM 版本 1（“V1”）的架构示意图。推理服务器的核心作用是支持多用户并行使用同一 LLM，从而降低成本，因此是 AI 技术栈中的关键组成部分。此外，推理服务器还能降低整体应用的延迟并提高吞吐量——因为模型权重加载等初始启动成本只需承担一次。

如图 1 顶部所示，用户通常只需提供待部署的 LLM、可选的量化配置（若需），以及模型权重（若权重未在 Hugging Face[15]等公共仓库中公开）。如今，vLLM 已成为 LLM 服务的行业事实标准，正越来越多地被用于生产环境，且可运行在 NVIDIA GPU、AMD GPU 以及 AWS Inferentia[1]、谷歌 TPU[39]、IBM Spyre[17, 43]等定制加速器上。

图1. vLLM v1的架构以及跨厂商性能可移植性问题：vLLM 支持多种平台，但对于性能关键的内核，它仍依赖以“后端”形式封装在 vLLM 中的外部库。对于许多 LLM 的“核心”——即所谓的（自）注意力层[41]而言，情况尤其如此，该层通常是 LLM 中性能最关键的内核[7, 36, 21, 44, 8]。这些平台特定的注意力库通常包含数万行代码[33, 36]

为实现这种灵活性，vLLM 采用了复杂的内部结构，将调度、模型预处理/后处理、运行时执行等功能分离。

为达到最先进性能，vLLM 的运行时执行分为两大部分，如图 1 下半部分所示：

• 部署的 LLM 中较简单的层（如归一化层或投影层）由平台无关的 PyTorch 函数编写，并通过 torch.compile 自动编译优化；
• 而复杂且性能最关键的注意力层，则被封装在一个名为“注意力后端”（attention backend）的抽象模块中。

在 vLLM 中，注意力后端有多种实现，通常是对手动优化库（如 flash_attn 或 flashinfer）的封装，如图 1 中部所示。因此，许多后端依赖外部库，若没有这些库，vLLM 无法正常使用。

2.2 Triton：一种瓦片式领域特定语言（tiling DSL）

领域特定语言（DSL）Triton 近年来逐渐成为编写定制 CUDA 内核的热门开源替代方案。Triton 支持用高层级 Python 代码编写和调试内核，这些代码可编译并运行在多种 GPU 架构上。已有研究表明，Triton 内核既能实现高性能，又能跨不同 GPU 平台移植。正因如此，Triton 的普及率不断提高，已被用于许多 LLM 技术栈，并且是 pytorch.compile 的核心组成部分。

Triton 利用即时（JIT）编译器，并基于“分层瓦片”的思想，自动实现内存合并、共享内存分配以及线程间同步。

清单1. 一个简单的 Triton 向量加法程序。其中 BLOCK_SIZE 是全局瓦片大小，可手动设置或通过自动调优确定

列表 1 展示了一个用 Triton 实现的一维并行向量加法程序。Triton 内核可通过超参数针对不同工作负载规模或目标架构进行微调。例如，在列表 1 中，BLOCK_SIZE 是一个配置参数，会影响 GPU 核心间的调度。然而，在实际应用中，Triton 内核仍需针对特定工作负载进行手动优化，且在不同 GPU 平台上的性能表现并不一致。为解决这一问题，Triton 内核及其编译器参数可通过自动调优进行优化。

2.3 可移植性与内核参数自动调优

自动调优是编译技术的补充手段，可进一步提高性能可移植性。它通过在编译期间或编译前利用微基准测试进行“试错”，帮助编译器找到最优或接近最优的内核配置参数集。

自动调优的主要优势是避免了像 2.4 小节所述库那样，手动编写数万行高度优化的代码。在新硬件上移植或部署应用时，自动调优尤为有用。经过自动调优的内核通过“经验性能调优”增强编译过程：生成并基准测试大量内核变体，为目标硬件和目标场景选择性能最佳的配置。自动调优减少了编译器为特定内核编译时需考虑的参数空间，能够探索比纯编译器方法多一个数量级的优化空间，从而实现更高性能和更优的代码专用性。

因此，与纯编译器方法相比，自动调优在“可移植性”与“编译时间”的权衡上更具优势——纯编译器方法可能导致极长的编译时间，甚至因问题复杂度而无法实现。

在先前的研究中，已证明自动调优可实现 LLM 的 Triton 内核在 GPU 上的可移植性。该工作将 Flash Attention v2 的 Triton 实现与 flash_attn 库、ROCm flash attention 实现进行了对比，结果表明：使用相同的内核代码，Triton 实现在 NVIDIA A100 和 AMD MI250 上均能达到与厂商特定最先进库相当的性能。

在学术文献之外，LLM 应用中也在使用自动调优技术，例如 PyTorch Inductor。PyTorch Inductor 是 torch.compile 的调优前端，它本身也是一种 JIT 编译器，可将 Triton 用作后端。PyTorch Inductor 通过“顺序尝试所有选项”的方式选择不同算法，且最近已添加了对 Triton 内核的支持。

然而，在实际应用中，通过自动调优 Triton 内核来提高性能可移植性的做法仍很少见或未达最优，主要原因是自动调优的开销过大。 在第 5 节中，我们将介绍解决这一限制的方案。

2.4 注意力内核与库

Flash Attention 是注意力算法领域流行的开源库，包含超过 7 万行代码，其中大部分是 CUDA 代码。Flash Attention 的核心创新是“flash 技巧”——通过优化注意力算法的内存局部性来提升性能，这一创新使其广为人知。Flash Attention 主要针对最新的 NVIDIA GPU 优化。此外，还存在 RocmAttention——它是 Flash Attention 的衍生版本，包含 AMD 特定优化，并通过 hipify 工具对 CUDA 代码进行交叉编译。

然而，早期版本的 Flash Attention 存在严重的内存碎片化问题：对于每个请求，必须预留足够内存以存储可能生成的最大令牌数。因此，“分页注意力”开发了分页版本的 Flash Attention，利用“分页”思想减少 GPU 内存消耗。 其核心思路是：
* 只为新请求预留少量内存（例如 16 个令牌），存储在一个名为“页”的数据结构中；
* 若请求生成的令牌数超过 16，则分配新的页。

分页注意力提高了推理服务器的效率，是 vLLM 等推理平台的核心特性。 此后，flash_attn、flashinfer 等许多其他注意力库也纷纷跟进，为其算法添加了分页版本。

Flashinfer 是一个部分开源的库，包含 5.1 万行代码，且依赖许多专有二进制组件。它用 CUDA 编写，仅能部署在最新的 NVIDIA 硬件上。

Flex Attention 试图减少注意力库的数量，覆盖更广泛的注意力类型。它提供了一种优化的注意力实现，支持自定义评分方法，从而能够支持不同的滑动窗口、分页/非分页模式、软截断修改或不同的注意力掩码。

Aiter 是近期推出的 AMD 专用推理库，目标是为针对最新 AMD GPU 优化的各类内核提供完整封装。在 Aiter 中，内核用多种语言编写，包括 Triton、HIP、CK 甚至汇编语言。在 AMD GPU 上部署 vLLM 时，可选择使用 Aiter。

三、解决方案概述

本文提出的注意力机制解决方案 及其在 vLLM 中的集成架构如图 2 所示。

图2. vLLM中triton_attn后端的概述。该图涵盖了实现最佳性能所需的所有关键步骤与组件：作为 vLLM 核心组件的调度器（scheduler）① 和gpu_model_runner②、包含三个核函数（kernel）③a 的triton_attn 后端③，以及后端中用于提升性能可移植性的核函数配置启发式规则（configuration heuristics，根据经验或数据制定的决策规则，用于快速选择适配不同硬件的核函数配置）③b。此外，triton_attn 后端在调用这些启发式规则时，还会考虑其部署所依赖的GPU 类型⑤（如 NVIDIA 或 AMD GPU）。与手动优化的注意力后端不同，vLLM 中的大多数其他层均采用原生 PyTorch 代码编写，并通过torch.compile④（PyTorch 的即时编译工具，可将 PyTorch 代码自动编译为高性能 GPU 可执行代码）编译为高性能核函数。

典型的 vLLM 部署会在启动阶段（startup time）记录CUDA 图（CUDA-graphs）或 HIP 图（HIP-graphs） ①a（NVIDIA 和 AMD 分别推出的 GPU 任务优化技术，可将一系列 GPU 操作记录为“图”，后续重复执行时直接“回放”图以减少软件开销），之后推理服务器才会进入就绪状态。

为实现这一过程，需先创建伪元数据（dummy metadata） ①b（用于模拟真实输入的虚假元数据，仅用于启动阶段记录图）并输入模型（包括注意力后端）。在推理阶段，这些图会被直接回放，而无需重新执行原始代码。vLLM 支持两种图模式：

• 部分 CUDA 图（Partial CUDA-graphs） ：除注意力层和 mamba 层（一种基于状态空间模型的高效序列处理层）外，所有层的操作均以 CUDA 图形式执行；
• 完整 CUDA 图（Full CUDA-graphs） ：所有层的操作均被记录到一个 CUDA 图中，执行时直接回放整个图。

四、注意力机制的设计与实现

4.1 核心概念

注意力机制由 Vaswani 等人在其开创性论文[41]中提出，其计算方式如下：

其中， 分别代表查询矩阵（Query） 、键矩阵（Key） 和值矩阵（Value） ，均由输入嵌入通过线性投影得到； 表示序列的 token 数量（序列长度）； 表示键向量的维度。中间矩阵乘积 存储了所有 token 对之间的注意力得分（attention score） ，而 是用于稳定 softmax 操作（一种激活函数，可将数值转换为概率分布，使所有结果之和为 1）的缩放因子。公式（1）的朴素实现会产生 的计算复杂度和内存复杂度。为提升 GPU 上注意力计算的性能，研究者开发了多种优化技术，下文将简要介绍这些技术，并说明其在本文注意力核函数中的集成方式。

分块 Softmax

Softmax 函数针对 矩阵的每一行（对应一个查询 token）独立计算，通过对所有键的注意力得分进行归一化，得到概率分布。为保证数值稳定性，Softmax 的计算方式如下：

其中， 表示得分向量 （即 矩阵的一行）的第 个元素， 表示该得分向量中的最大值。减去最大值可避免指数运算时因数值过大导致的上溢（overflow），或因数值过小导致的下溢（underflow 至 0），从而提升计算稳定性。

Softmax 可通过分块方式（tiled approach，也称 online softmax） 高效计算： 将每行得分向量划分为多个小“块（tile）”，逐块增量处理，而非一次性处理整行 。分块 Softmax 将“除以指数和”的操作延迟到所有块处理完毕后执行，过程中会持续维护该行的最大值（）和指数和（），并在处理每个块后更新这两个值（若最大值发生变化，还需对中间结果进行重新缩放）。

这种分块策略能让核函数最大限度（或完全）利用 GPU 的高速共享内存（shared memory） 和寄存器（register） ，避免频繁访问速度较慢的全局内存（global memory） ，从而显著提升性能。在实际实现中，分块 Softmax 会与公式（1）中的分块矩阵乘法融合，下文将这种融合实现统称为“分块 Softmax”。分块 Softmax 是 FlashAttention[8]中的核心优化技术之一。

键值缓存（KV Cache）

LLM 注意力层的计算可通过键值缓存（KV Cache） 加速：将已处理输入 token 的 K 和 V 矩阵缓存起来，后续注意力计算时直接复用这些缓存的矩阵，无需重新计算。

• 键值缓存的初始化发生在预填充阶段（prefill phase） ——此时需为提示（prompt）中的所有 token 计算 K 和 V 矩阵；
• 而在解码阶段（decode phase） ，仅需为新生成的每个 token 计算 K 和 V 矩阵。

分组查询注意力与多查询注意力（GQA 与 MQA）

当前主流的 Transformer 架构基于多头注意力（multi-head attention） ——每个层中多个注意力头（attention head）并行工作，以捕捉更丰富的 token 间关系。但这种方式需在多个头间重复执行注意力计算，且每个头都有独立的 Q、K、V 投影矩阵，内存和计算开销较大。

• 分组查询注意力（Grouped Query Attention, GQA） 通过减少键头（key head）和值头（value head）的数量来优化：让多个查询头（query head）共享同一组 K 和 V 投影矩阵，从而减少 K 和 V 矩阵的计算量，降低键值缓存的内存占用。
• 多查询注意力（Multi Query Attention, MQA） 则是 GQA 的极端形式——所有查询头共享一个键值头。

批处理（Batching）

最后（且不仅限于注意力计算），通过批处理（将多个序列合并为一个批次处理） 可提升并行性，通常能更高效地利用 GPU 资源（如计算核心、内存带宽）。

4.2 术语定义

本文使用以下术语（均针对单个序列，与 vLLM 中的术语一致）：

• 上下文长度（Context Length） ：序列中已处理 token 的数量，其 K 和 V 矩阵存储在键值缓存中；
• 查询长度（Query Length） ：当前正在处理的新 token 数量，需针对缓存的上下文计算注意力；
• 序列长度（Sequence Length） ：序列的总 token 数量，等于上下文长度与查询长度之和。

对于预填充注意力（prefill attention） （处理输入提示的注意力计算），上下文长度为 0，查询长度等于提示的 token 数量；对于解码注意力（decode attention） （生成新 token 时的注意力计算），查询长度为 1。

此外，前缀长度（Prefix Length） 指序列中某个 token 之前的 token 数量，包括已处理的 token 和当前 token 之前的新 token（例如提示中的前序 token）。

4.3 基准 Triton 注意力核函数

本文在 Triton 中实现的首个分页注意力核函数，遵循了原始分页注意力算法[21]的设计思路。该实现假设：Q、K、V 矩阵已在核函数启动前计算完成，并以“块”的形式存储在分页内存（模仿操作系统分页机制，将内存划分为固定大小的块管理，以提升内存利用率）结构的键值缓存中。vLLM 中的键值缓存通过块表（类似操作系统中的页表，记录缓存块的位置信息，用于快速查找）访问，参数 BLOCK_SIZE 定义了单个键值缓存块可存储的最大 token 数量。

尽管预填充和解码阶段使用完全相同的核函数，但两个阶段的启动网格（GPU 中用于分配线程块的结构，决定并行执行的线程块数量和组织方式）不同：
* 预填充阶段：启动网格的维度为“批处理中的 token 数 × 查询头数”，即二维启动网格。
* 解码阶段：启动网格的维度为“批处理中的序列数 × 查询头数”，此设置更适配解码阶段的计算特征。

该实现的详细说明（包括代码清单 3）如下。

清单3. 基准Triton注意力内核的高级简化表示。它假定𝑄、𝐾和𝑉在启动内核之前已经计算完成，并存储在KV缓存中。vLLM中的KV缓存通过块表（类似于页表）进行访问，不过这里不直接讨论这一点，仅提及参数BLOCK_SIZE，该参数定义了单个KV缓存块中存储的最大令牌数量。从第36至40行可以看出，针对处于预填充阶段和解码阶段的序列，内核是分别启动的。在启动网格中，num_seqs表示批处理序列的数量，而tot_query_length指的是需要计算注意力的这些序列中的令牌总数（即所有批处理序列的查询长度之和）。对于预填充注意力，后者的值是通过汇总所有序列的提示长度得到的，而对于解码注意力，它简单地等于批处理中的序列数量。

4.4 预填充与 GQA 优化

4.3 节描述的基准注意力核函数，每个程序实例仅处理“一个查询 token + 一个查询头”的组合。本节介绍的优化通过以下方式，增加单个程序实例可处理的组合数量：
* 针对预填充注意力：处理同一提示中的多个连续 token。
* 针对GQA：处理共享同一个键值头的多个查询头。

图3. Q 块和 KV 瓦片

这些组合共同构成了本文所称的查询块。一个 Q Block 包含 BLOCK_M 个组合，涵盖 BLOCK_Q 个查询 token 和 BLOCK_Q 个查询头。通过设置 BLOCK_Q = num_query_heads_per_kv_head，每个 Q Block 可覆盖映射到同一个键值头的所有查询头。

从结构上看，Q Block 是一个三维块，维度分别对应“查询 token 数、查询头数、头维度”；但为提升实现效率，Q Block 被表示为形状为“BLOCK_M × HEAD_SIZE”的二维张量——这种扁平化处理可简化内存访问模式，更适配 Triton 的编程模型。

对于一个序列，处理所有查询 token 需 ceil(seq_len / BLOCK_Q) 个 Q Block。在解码注意力中（每次仅处理 1 个查询 token，即 BLOCK_Q=1），每个序列对应一个 Q Block。

Q Block 结构使核函数能并行处理多个注意力计算，从而提升效率：
* 在预填充注意力中，多个查询 token 会关注相同的前序 token，可高效复用 K 和 V 矩阵。
* 在 GQA 中，多个查询头对应同一个键值头，每个 Q Block 只需加载一次该键值头的 K 和 V 矩阵。

上述优化减少了内存带宽消耗，并通过提升计算密度（单位内存访问对应的计算操作数量）提高了计算效率。

图3. Q 块和 KV 瓦片

对于每个 Q Block，采用分块 Softmax 处理的 K 和 V“块”结构，与 4.3 节基准核函数保持一致（如图 3 所示）。这些 K/V 块对应 Q Block 中查询头所映射的键值头，并覆盖 Q Block 中所有 token 的前序 token（范围不超过 Q Block 中任意 token 的最大前缀长度）。代码清单 4 提供了基于 Q Block 的优化注意力核函数示例代码。

清单4. 为预填充和GQA优化的注意力内核。这是基于QBlock的优化注意力内核的Triton实现，采用了与基准内核类似的高级且简化的表示方式。一个关键差异在于启动网格，它现在包含了批次中所有序列的Q块总数，以及KV头的数量（第38行）。对于解码，Q块的总数等于批次序列的数量（第41行）。此外，现在查询头索引是从KV头派生出来的（第5-6行）。序列索引的确定方式与之类似，是通过对存储批次中所有序列的Q块累计数量的张量进行二分查找来实现的（第9行）。

4.5 并行分块 Softmax

若使用 4.4 节的核函数处理解码注意力，启动网格的第一维度等于批处理中的序列数。当批处理大小较小时，启动的程序实例数量有限，会导致 GPU 资源未充分利用，性能下降。

这一问题在预填充注意力中不存在——因为提示通常包含大量 token，生成的 Q Block 数量足以饱和 GPU 的计算资源。

为在解码注意力中挖掘足够的并行性，可将分块 Softmax 的迭代处理过程跨多个程序实例并行化。这种优化对“小批量、长序列”场景尤为有效：此类场景中需处理大量分块，但 GPU 易处于未充分利用状态；通过将分块分配给多个程序实例并行处理（而非在单个实例中串行处理），可显著提升速度。

为说明并行分块 Softmax 的实现，先定义以下术语：
* 块：分块 Softmax 中 Q 与 K、V 迭代计算的粒度（即每次处理的 K/V 子矩阵大小）。
* 段：由单个程序实例处理的一组块，多个段可并行处理。

图4. 并行分块Softmax

图 4 通过示例展示了这一过程：
* 图 4（a）：将一次注意力计算划分为 32 个块，按 4.1 节所述的分块 Softmax 逻辑迭代处理，过程中维护并更新中间结果、最大值和指数和。
* 图 4（b）：将 32 个块分配到 4 个段中（每个段含 8 个块），每个段在独立的程序实例中按常规迭代方式处理；处理完成后，需通过归约步骤（将多个段的中间结果合并为最终结果的操作，如合并最大值、重新计算指数和等）合并所有段的输出——程序实例会将段的中间结果存储到内存中，最后通过类似迭代分块 Softmax 的方式，读取并合并这些中间结果，重新缩放后得到最终注意力输出。

下面给出该实现的示例代码。

清单5. 支持并行分块Softmax的内核：为了在解码注意力阶段提取足够的并行度，可以对 Listing 4 中第 22 行 for 循环体内分块 Softmax 的迭代处理进行并行化。Listing 5 展示了相应的 Triton 实现。在此版本中，使用三维启动网格来启动注意力内核，其中每个 Q 块与 KV 头组合所包含的段数构成了第三维度（第 61 行）。分配给每个程序实例的段索引（第 5 行）决定了在该程序实例内将迭代处理哪一组分块（第 21 至 26 行）。一旦某个段的注意力计算完成，中间结果便会被存储到内存中（第 37 至 40 行）。第一个内核执行完毕后，会启动归约内核以从段级结果计算出最终的注意力输出。

图2. vLLM中triton_attn后端的概述

4.6 可调整分块大小

在4.3至4.5节描述的核函数中，分块Softmax所使用的分块大小（tile size）受限于键值缓存块的大小（由BLOCK_SIZE参数定义），必须与BLOCK_SIZE保持一致。

本文对核函数实现进行了扩展，实现了分块大小与BLOCK_SIZE的解耦。现在，分块大小可以独立于块大小进行配置（可小于、等于或大于BLOCK_SIZE），这带来了两大优势：

1. 性能优化：分块大小可以独立调整，并且可以为预填充和解码两个阶段分别设置最优的分块大小，从而进一步提升性能。
2. 适配混合模型：支持混合模型（例如Mamba[12]这类结合了传统Transformer注意力层与状态空间模型（SSM）层的模型）。在此类模型中，注意力层通常需要使用“非2的幂次”的大块大小，以实现注意力层与SSM层之间的页对齐（确保不同层的内存块在地址上对齐，以减少内存访问冲突）。

4.7 静态启动网格

作为最后一项优化，本文将核函数修改为使用静态启动网格——即每次启动的程序实例数量是固定的（通过调整每个实例处理的Q Block数量来实现）。这种设计的主要优势在于提升了与完整CUDA/HIP图的兼容性，具体将在6.2节中讨论。

五、自动调优与可移植性

5.1 当前自动调优的不足

已有研究[33, 11]表明，基于Triton的单源核函数能够在多种硬件平台上实现有竞争力的性能，但这依赖于自动调优（通过试错和微基准测试自动寻找最优核函数配置参数的过程，参见2.3节）。

然而，在实际应用中，Triton核函数的自动调优存在调优时间开销巨大的问题[33, 30, 22]，导致其在许多场景中无法使用（例如vLLM[20]）。例如，为了使FlashAttention v2[7]在每种GPU上达到最佳性能，需要进行近24小时的全面调优[33]。

在某些场景中，可以通过自动调优结果缓存（将调优得到的最优配置存储起来，供后续部署时复用）来减少开销[33, 30, 22, 27, 16]。Triton自动调优器的缓存中存储了调优结果，其核心是“场景到调优结果”的映射。这里的“场景”指的是Triton核函数输入参数的特定组合（如张量指针、形状、步长、标量参数等）。但在vLLM注意力后端中，部分核函数参数（如序列长度）会随请求变化，因此缓存仅在“完全相同的场景再次出现”时有效。例如，针对32个token的场景调优后，再次遇到32个token的请求可以复用缓存；但若遇到33个token的请求，则需要重新调优。

除了调优时间开销，在推理服务器中使用自动调优还存在其他问题：

1. 查找开销：即使提前对所有可能场景完成了调优，在实际推理时查找“当前场景对应的最优配置”仍会增加约数十微秒（µs）的核函数启动时间。对于处理时间短的负载，这一开销可能会抵消调优带来的性能提升。
2. 与CUDA/HIP图不兼容：如第3节所述，CUDA/HIP图一旦被记录就只能“回放”，无法在回放时调用Triton自动调优器来查找最优配置（参见6.2节）。这意味着在使用图技术时，自动调优的灵活性完全丧失。

5.2 vllm-triton-backend 中的自动调优应用

图5. vllm-triton-backend中的调优工作流程

为解决Triton自动调优“时间开销大”和“灵活性不足”的问题，本文采用了两步法（如图5所示）。

第一步：基于微基准测试的核函数调优

首先，构建一个微基准测试框架，在vLLM运行时组件之外独立执行核函数调优。这种方式避免了增加vLLM的启动延迟，也不会加重持续集成（CI）流水线的负担（频繁在CI中执行调优会大幅增加耗时）。

该微基准测试框架的设计特点包括：
* 调用与vLLM中完全相同的核函数代码。
* 能够模拟真实请求模式（如不同上下文长度、提示长度、批处理大小的组合）和LLM架构。
* 能够揭示端到端测试中难以发现的性能问题。例如，部分现有核函数仅针对“批处理中所有请求的token数量相同”的场景优化，但实际LLM推理中这种情况极少。通过微基准测试可以模拟“token数量可变”的真实场景，并据此调整核函数。

第二步：将调优结果转化为启发式规则

调优完成后，分析自动调优结果，并将其转化为启发式规则（此处指简单的“if-else决策树”，如图5右侧所示）。代码清单2给出了一个启发式规则的示例。

清单2. 作为Triton启发式方法的决策树（已调优）

与“自动调优缓存静态复用”[33, 30, 22, 2, 32]（仅复用完全匹配的场景）相比，决策树的优势在于：能够覆盖未纳入调优的新场景。无需对所有场景逐一调优，仅需针对“平均场景”和“极端场景”进行调优，即可生成适用于大多数场景的决策树。因此，这种方式在运行时和调优阶段均有优势：运行时无需查找缓存，调优阶段可减少工作量。

六、推理服务器与系统集成

要使优化后的注意力核函数达到最佳性能，需要将其集成到vLLM中。这一过程需要解决两个关键问题：
1. 注意力元数据的创建与计算适配。
2. CUDA/HIP图的使用权衡（即是否使用图技术，以及如何平衡其优缺点）。

6.1 元数据计算

在vLLM中，调度器确定下一批次需要处理的（部分）请求后，会执行以下操作：
1. 将数据复制到GPU（若尚未在GPU中）。
2. 将数据整理为分页数据结构（如键值缓存的块表）。
3. 计算对应的元数据（描述数据的数据，此处指注意力计算所需的辅助信息）。

vLLM中的注意力元数据包括“批处理中各请求的长度列表”等信息；对于部分后端，还会将批处理中的请求按“解码请求”或“预填充请求”排序（集中处理同一类型请求，以减少切换开销）。

为了适配triton_attn后端，需要对元数据计算进行以下修改：
1. 统计批处理中的解码请求数量：用于判断是否启用并行分块Softmax核函数（参见4.5节）。
2. 构建“查询块累积数量张量”：该张量存储批处理中每个序列对应的“查询块（Q Block）累积数量”。在每个启动的程序实例中，通过二分查找该张量，可以根据“查询块索引”快速定位对应的“序列索引”（参见4.4节末尾）。同时，该张量也用于计算处理整个批处理所需的查询块总数。

6.2 Triton 启动开销与 CUDA/HIP 图

集成 triton_attn 后端时需解决的核心问题是CUDA/HIP 图的使用[5, 13]。CUDA/HIP 图的核心价值是：若模型前向传播（forward pass，从输入到输出的计算过程）的操作序列固定，可通过图记录这些操作，后续执行时直接回放以消除软件开销。

矛盾点：Triton 的灵活性与图的固定性

Triton 核函数的性能优势之一是“可根据批处理形状（batch shape，如请求数量、序列长度）动态选择最优配置”（参见 5 节）。

但这意味着 triton_attn 后端的行为会随每次前向传播变化， 与“图一旦记录就无法修改”的特性冲突 ——若不使用 CUDA/HIP 图，Triton 核函数的启动开销（launch overhead，如参数检查、上下文切换等软件层面的开销） 会显著增加。

根据本文的性能分析，不使用 CUDA/HIP 图时，Triton 核函数的启动开销约为 100-300 微秒（µs）；对于 token 数量少于约 1000 的序列，这一开销会主导核函数的总执行时间（参见 7.2 节）。需注意的是，这一开销不包括“首次执行时的 JIT 编译时间”（将 Triton 代码编译为 GPU 可执行代码的时间），仅指每次核函数启动时 Triton 产生的软件开销 （如检查是否需要重新 JIT 编译）。即使通过缓存规避部分检查[18]， 启动开销仍约为 80 微秒。

使用 CUDA/HIP 图的权衡

为降低启动开销，需使用 CUDA/HIP 图，但这会引入新的权衡：

1. 参数冻结：图中所有核函数的参数（包括张量指针）会被“冻结”[37, 5, 13]——一旦记录图，后续回放时无法修改这些参数；
2. 内存占用增加：若记录大量图，GPU 内存会被“预留内存（reserved memory，为图执行预先分配的内存，即使实际请求不需要也无法释放）”占据，且图本身也需占用内存。

因此，vLLM 采取了限制图数量的策略：仅为“2 的幂次批大小（power-of-two batch sizes，如 2、4、8、…、128）”各记录一个图[37]。具体来说：

• vLLM 启动时，会为“批大小 ≤128 的所有 2 的幂次”分别创建一个“伪请求”，并记录对应的图；
• 记录图时，需使用“模型支持的最大序列长度”分配核函数所需的 GPU 内存——这意味着，即使实际请求的序列长度远小于最大值，使用图时仍会按“最大序列长度”分配内存并执行核函数。

静态启动网格的必要性

对于本文的 Triton 核函数，使用 CUDA/HIP 图会带来额外性能损耗：Triton 核函数通常根据批处理元数据（如序列长度、查询块数量）动态确定“启动的程序实例数量”；但图记录后，启动网格被固定——即使实际请求所需的实例数量更少，也会按“最大序列长度对应的实例数量”启动核函数。

多余的实例会立即退出（不执行计算），虽不影响结果正确性，但会导致 GPU 调度器（负责分配 GPU 资源的组件）调度过多的“GPU 波前（waves，即 GPU 流多处理器（SM）的一轮核函数执行）”，增加不必要的调度开销。

本文的性能评估显示：多余实例带来的延迟增加，几乎在所有场景中都超过了图技术节省的启动开销。因此，需将核函数修改为使用“静态启动网格”（参见 4.7 节）——启动的实例数量固定，且接近（但不超过）GPU 的核心数量，以避免多余实例的调度开销。

七、性能评估

我们采用双轨方法对内核进行评估：首先通过微基准测试套件测量各优化步骤带来的增量性能变化，随后进行全面的端到端测试。本次评估旨在回答以下研究问题：

1. 各个优化步骤（参见第 4 节）对性能的 影响程度如何？
2. 以简单启发式规则形式实现的自动调优，能否进一步提升性能（参见第 5 节）？
3. 整个推理服务器的端到端性能表现如何？
4. CUDA/HIP 图与动态即时编译（JIT，Just-In-Time）相比，会产生哪些影响？

7.1 实验方法

我们在两款 GPU 上进行评估，分别是 NVIDIA H100-80GB 和 AMD MI250-128GB。选择这两款 GPU 的原因在于：它们采用相同的工艺节点（由台积电制造的 5 纳米工艺），代表了两大主流硬件厂商，同时也因其普及度和可获取性。

在微基准测试中，我们的内核参数基于 Llama3-8B 大语言模型架构[23]（128 维头维度、32 个查询头、8 个键值头），并根据真实场景样本调整序列长度和批处理大小。在真实的在线推理场景中，一个批处理内的序列通常具有可变长度，我们的测试也模拟了这一特征。每个微基准测试的测量前，都会对内核进行 20 次迭代热身，随后取 100 次迭代的平均值作为最终结果。

我们使用 vLLM 自带的基准测试套件[42]，在 NVIDIA H100 和 AMD MI300 GPU 上进行端到端实验。为了专注评估内核改进效果（而非提示序列冗余带来的影响），我们 禁用了 vLLM 的前缀缓存功能。测试使用随机数据，并忽略模型的序列结束（end-of-sequence）token。热身迭代次数和测量迭代次数保持默认值，分别为 10 次和 30 次。

7.2 优化步骤评估

为精确测量单一组件（内核）的性能，我们通过微基准测试评估各个优化步骤的效果。

图6. 不同内核优化与基准的性能比较

图 6a 展示了在 H100 上评估的四种不同的分页注意力（paged attention）内核实现：flash_attn代表最先进的库 Flash Attention 3[36]，另外三种是我们基于 Triton 实现的版本——Triton (naive)是第 4.3 节所述的朴素实现，Triton (GQA opt.)是第 4.4 节的优化版本，Triton (parallel tiled)是第 4.5 节介绍的并行分块优化版本。

图 6b 展示了在 AMD MI300 GPU 上的相同实现，但由于 AMD GPU 目前缺乏可与之竞争的分页注意力实现，因此仅呈现了三种 Triton 实现。在图 6a 和图 6b 中，每个子图代表“顶部标注的最大序列长度”对应的批处理，x 轴为批处理大小，y 轴为归一化延迟（值越低越好）。所有微基准测试的延迟值均以朴素实现的最左侧数据点为基准进行归一化。

观察代表朴素实现的橙色曲线可知：

• 在所有序列长度下，其速度均比 FlashAttention 慢近一个数量级。
• 而针对数据局部性（尤其针对分组查询注意力（GQA）模型）的优化版本（紫色曲线），在短序列长度和小批处理大小场景下，性能优于朴素实现，有时甚至比 FlashAttention 更快；
• 但对于长度超过 500-1000 个 token 的序列，该优化带来的性能提升并不显著。并行分块 softmax 优化对性能的提升则更为有限。

图6. 不同内核优化与基准的性能比较

然而，若我们不按序列长度、而是按批处理的组成结构（如图 6c 和图 6d 所示）重新绘制相同数据，会观察到完全不同的性能特征。在这两张图中，x 轴表示批处理内“批大小 × 序列长度”的总和，子图按批处理中“仅解码请求”的占比（0%、50%、100%）划分。由于 vLLM 通常会优先处理解码请求，因此包含高比例解码请求的批处理在实际场景中十分常见。

如图 6c 和 6d 所示，GQA 优化版内核的优势主要体现在“预填充密集型”批处理中——这符合预期，因为预填充阶段是计算密集型（compute-bound）任务。而内存密集型（memory-bound）的解码阶段，则需要通过并行分块 softmax 实现更高的并行度。

如图 6c 最右侧子图所示：

• 蓝色曲线（并行分块优化版）的性能已与 FlashAttention 几乎持平；
• 在第 7.4 节中还将看到，对于极长的解码序列，并行分块 softmax 优化版的性能甚至超过了 GQA 优化版内核。

图6. 不同内核优化与基准的性能比较

图7. 可调瓦片大小优化的性能比较（参见4.6小节）

图 7 对比了第 4.6 节所述的“灵活块大小”方法，同时也呈现了图 6 中性能最佳的两种实现作为参照。由于按“解码请求占比”划分更能体现性能差异，因此图 7 仅展示了该划分方式下的结果。正如预期，两种“灵活块大小”版本的性能均优于其对应的原实现版本。

7.3 自动调优方法评估

图8. 比较经过优化的QGA内核在有预填充大小调优和无预填充大小调优情况下的性能

初步测试表明，自动调优主要影响预填充阶段，因此图 8 仅展示了预填充密集型批处理的评估结果，并按序列长度划分。图 8 中的“调优版”（tuned）遵循第 5 节所述的流程，最终采用了基于自动调优结果的简单决策树（如代码清单 2 所示）。

清单2. 作为Triton启发式方法的决策树（已调优）

可以看出，这种有限的自动调优方法在两种平台上，均降低了短提示序列（H100 上最高可达 9.8 倍）和中长提示序列（最高可达 75%）的内核延迟 。由于更复杂的启发式规则需要更深入的开销与权衡分析，超出了本文的讨论范围，因此我们未对其进行评估。但仅通过这种简单的启发式规则，已能体现出自动调优的价值，后续的端到端评估也采用了该规则。

7.4 端到端评估

图9. Llama-3.1-8B模型在批量大小为1、输入长度为500个标记时的vLLM延迟基准测试结果

图 9 展示了 metallama/Llama-3.1-8B-Instruct 模型[23]的延迟基准测试结果：批处理大小为 1、提示序列长度为 500 个 token，测试变量为生成的输出 token 数量。该测试配置能详细分析第 4 节所述各优化的影响，且排除了调度决策对测量结果的干扰。

• “基于静态启动网格的 Triton 内核”（Triton Static Launch Grid）整合了其他内核中的 Q 块（Q-Block）和并行分块 softmax 优化，还加入了静态启动网格（参见第 4.7 节），并且在执行时启用了完整的 CUDA/HIP 图（参见第 6.2 节）；
• 其他所有实验均使用部分 HIP/CUDA 图。此外，该内核还融入了自动调优过程中确定的分块大小与段大小启发式规则（参见第 7.3 节）。

图 9a 展示的 H100 测试结果表明，经过逐步优化的内核性能持续提升：

• 对于输出长度为 12800 个 token 的场景，并行分块 softmax 优化相比仅实现 Q 块优化的内核，端到端延迟降低了两倍多；
• 而启用完整 CUDA 图的“基于静态启动网格的内核”，在相同输出长度下，延迟又进一步降低了约 6%。

这一结果表明，序列长度越长，延迟降低效果越显著。作为参照，图 9 还包含了 FlashAttention V3 的实验结果——可以看出，在该测试配置下，“基于静态启动网格的内核”性能已与 FlashAttention V3 相当 。此外，图 9 数据显示：

• 朴素实现的性能仅为 FlashAttention3 的 19.7%；
• 我们的第一个 Triton 优化内核（见前文代码清单 4）将性能提升 2.1 倍，达到 FlashAttention3 的 49%；
• 而最终的静态启动网格优化（参见第 4.7 节），性能更是达到了 FlashAttention3 的 98.6%–105.9%。

图9. Llama-3.1-8B模型在批量大小为1、输入长度为500个标记时的vLLM延迟基准测试结果

图 9b 展示的 MI300 测试结果表明，启动开销对 AMD GPU 性能的影响比 H100 更显著 ：结合启发式规则和完整 HIP 图的“基于静态启动网格的 Triton 内核”，比“Q 块+并行分块 softmax”（Q Block+Par Ts）版本快 1.99 倍。

这一点从“短输出长度场景下，并行分块 softmax 内核的延迟高于 Q 块内核”可明显看出——因为前者需要额外启动一个归约（reduction）内核（参见第 4.5 节）。但对于较长的序列，启动开销在总处理时间中的占比变得可忽略，此时并行分块 softmax 优化的优势开始显现。 启用完整 HIP 图（缓解启动开销的关键方案）的“基于静态启动网格的内核”，其性能提升更为显著 ：在所有输出长度下，相比并行分块 softmax 内核，延迟均降低了约一半。

综合来看，图 9b 中三种优化的累计性能提升达到了 5.9 倍 。

八、见解与未来工作

在本文所述的研究工作中，我们阐述了如何开发一个仅基于 Triton 的最先进注意力内核。同时，我们也获得了超出本研究场景的见解（例如，这些经验已被用于改进 vLLM 中 mamba/状态空间模型（SSM）层的内核[19]）。

下文将讨论这些见解，并介绍未来的研究方向。

8.1 Triton 内核需“场景专用”

在研究过程中，我们尝试过合并第 4 节所述的三种不同内核——例如，开发一种能在启动时检测是否为“仅解码”场景，并分支到预填充内核（参见第 4.4 节）或解码优化版内核（类似第 4.3 节的朴素实现）的融合内核。

我们的目标是减少内核启动次数 ，从而降低启动相关的控制开销与软件开销。然而，实验发现这类融合内核的性能至少下降了 2 倍，其性能损失远超过减少启动次数所节省的延迟 （约 150 微秒）。对 Triton 中间表示的后续分析表明， 融合内核的软件流水线无法生成有效的流水线结构。

由此我们得出结论：Triton 内核必须围绕“单一特定问题”设计 ，且该问题需具有强数据依赖性 ；否则，宁可承受多次内核启动的开销，也不应采用融合内核设计。

8.2 tl.dot 的使用

注意力计算中涉及公式 1 中的 乘法，在 Triton 中可通过以下两种方式实现：

1. K 与 Q 的广播版本进行按元素乘法（element-wise multiplication），随后求和：tl.sum(K * Q[:, None], axis=0)；
2. 矩阵乘法指令：tl.dot(Q, K)——该方法也适用于 softmax 输出与 V 的乘法。

第二种方法（使用tl.dot）要求输入矩阵的维度必须是特定矩阵乘加分块大小（如 8、16、32）的整数倍，这通常需要通过填充（padding）来满足该要求。尽管如此，tl.dot仍是更优选择，因为它几乎总能带来更好的性能——编译器会将其直接映射到 MMA 单元。

相比之下， “按元素乘法+求和”的方式通常无法被编译器识别为矩阵乘法，因此性能较差。

8.3 CUDA/HIP 图并非万能

我们观察到社区中存在一种普遍趋势：当遇到 Triton 的即时编译（JIT）开销、启动开销（或其他 JIT 框架如 CUTLASS 的开销）时，人们会倾向于使用 HIP/CUDA 图。

但正如第 6.2 节所讨论的，CUDA/HIP 图存在权衡取舍。在“内核/二进制层面”使用这类记录图，会强制应用程序选择“单一版本的 Triton 内核二进制文件”，这与 JIT 语言的设计理念相悖—— Triton 不仅会根据常量，还会根据内存访问模式（例如步长是否可被 16 整除）对内核进行特化；而自动调优（参见第 5 节）会进一步增加内核的多样性。

因此我们发现：当单一内核的运行时间与 Triton 的平均启动开销（约 200 微秒）处于同一量级时， 除非内核是专门为 CUDA/HIP 图设计的（如第 6.2 节所述），否则使用 CUDA/HIP 图无法降低整体延迟。

8.4 未来工作：内核调优中纳入 GPU 特定参数

除了num_warps（与线程数量相关）、num_stages（与软件流水线深度相关）等平台无关参数外，较新版本的 Triton 还引入了 GPU 特定参数。

• 针对 Hopper 和 Blackwell GPU 的“线程束特化”（Warp specialization）技术，引入了num_consumer_groups和num_buffers_warp_spec参数[25,3]；
• 又如 AMD GPU 的waves_per_eu参数（每个执行单元的波前数，波前是 GPU 中并行执行的线程组）[26]。

虽然这些参数可通过第 5 节所述的方法进行探索，但超出了本文的讨论范围。不过，在早期实验中，将这些参数纳入调优已使 部分内核的性能提升高达 80%； 文献中的相关研究（如在 flash_attn3 中使用线程束特化[36]）甚至报告了超过 100%的性能提升。

8.5 未来工作：自动调优与 CUDA/HIP 图的结合

如前所述，CUDA/HIP 图限制了 运行时对内核变体 的选择灵活性。

因此，若使用 CUDA/HIP 图，无法根据批处理中的实际 token 数量调整内核块大小及对应的启动网格（launch grid，定义 GPU 上执行内核的线程块分布）。

为了在这种情况下仍能利用自动调优的灵活性，我们需要将“决策树（如块大小决策）”迁移到 Triton 内核内部——这样，Triton 二进制文件保持不变，但可通过调整内部循环来适应不同场景（需进行一些微调）。

然而，Triton 循环【不支持】通过break或return语句实现提前退出， 因此我们必须插入空操作指令（NOPs，即不执行任何有效操作的指令）来实现所需的掩码操作（masking，用于过滤掉无效数据的处理）。在实验中，这种调优方式的性能提升不如“转向静态启动网格”显著，因此该优化方向仍留待未来研究。

九、结论

人工智能民主化不仅依赖于大语言模型（LLM）和人工智能模型的广泛应用，还依赖于人工智能部署的灵活性与可移植性。因此，LLM 解决方案的实现【必须】支持对【不同硬件】平台的选择。

本文认为，Triton 和 vLLM 等开源解决方案是实现这一目标的关键推动因素。我们的实验表明：基于 Triton 实现的分页注意力内核，使用相同的源代码，在 NVIDIA 和 AMD GPU 上均实现了最先进的性能。

在研究过程中，我们获得了多方面的经验：

• 包括 如何处理 Triton 的内存层级
• 如何提升数据局部性（指数据在内存中的分布是否便于 CPU/GPU 快速访问，好的数据局部性可减少内存访问延迟）
• 如何通过 自动调优提升性能可移植性
• 以及 如何平衡“即时编译二进制文件”与“仅数据流图（如 CUDA 和 HIP 图）”的集成

我们所研究的优化方法累计实现了高达 589%的性能提升，并已将相关内核与框架作为开源项目贡献（ibm.biz/vllm-ibm-triton-lib）。最终，我们开发的高度优化内核已成为 vLLM 中 AMD 部署的默认注意力后端。

参考文献

关注“鲸栖”小程序，掌握最新AI资讯