过去一年中，IBM研究院、红帽和超威半导体（AMD）的团队合作开发并向vLLM主仓库贡献了一款基于Triton的注意力后端。其目标非常明确：在实现顶尖性能的同时，确保代码能在不同GPU厂商的硬件上无缝移植，无需为每款芯片重写大量专用内核。
这项工作的背后，是当下加速器硬件生态日益复杂的现实。维护海量高度专用化核函数的成本正不断攀升，开发团队迫切需要一个更优雅的解决方案。
本文将对这项工作进行全面的技术解析，阐释Triton为何适配vLLM的技术需求，介绍Triton注意力后端的相关情况及其使用场景，随后深入讲解高性能分页注意力核函数的实现过程。文中还将依次介绍核函数级别的优化手段、并行化策略、与CUDA图的交互方式以及基准测试结果，最后对Helion做简要介绍。
关键问题
问题一：性能可移植的边界——Triton能否全面替代手写专用内核？
Triton的抽象是否真的能在所有场景下无损替代专家优化的CUDA内核？虽然作者给出了H100上100.7%的性能对比，但这只是特定模型、特定批次大小下的结果。现实中，像FlashAttention这样的库之所以有数万行代码，正是因为它们在为不同架构（如Hopper、Blackwell）做极致的手工调优。
那么，Triton能否在更复杂的模型、更长的序列、更细粒度的硬件特性（如Tensor Core、TMA、WGMMA）上真正匹敌这些手工调优？如果不能，那“可移植”是否意味着“平庸的性能”？
事实上，Triton用约800行代码实现了与7万行FlashAttention 3相当的性能（在H100上达到100.7%），并在AMD MI300上取得了5.8倍的加速，这确实证明了其“性能可移植”的潜力。
但作者也承认，微基准测试显示不同内核变体在不同负载（预填充/解码/混合）下各有优劣，没有单一配置能统治所有场景。Triton依赖autotuner自动选择分块（tile）大小，但autotuner的搜索空间有限，且新硬件特性可能尚未被Triton编译器充分映射。
因此，彻底放弃手写内核为时尚早——在极端长序列、稀疏注意力或未来新架构上，专家优化的CUDA内核仍可能保持领先。但Triton大幅降低了维护成本，使vLLM能在绝大多数常见场景下接近最优性能，这是走向硬件无关化的重要一步。
问题二：动态调度与CUDA图的平衡——持久化内核会引发负载不均吗？
“持久化内核”的设计是为了解决CUDA Graph与可变启动网格（launch grid）的冲突，但这种设计是否引入了新的调度开销或负载不均问题？
在极端动态的工作负载下（如混合了长短不一的请求），这种固定线程块数量、动态分配任务的策略是否会导致GPU资源利用率下降？我们是否在用一种“通用”的调度机制来掩盖底层硬件调度器的不足？持久化内核通过固定数量的线程块来适配CUDA Graph，但这也意味着某些线程块可能处理更多任务，而另一些可能空闲。这种“静态启动+动态负载”的模式在高度动态的在线推理场景中是否真的稳健？
持久化内核通过固定线程块数量、让每个块动态从全局内存中获取任务，确实解决了CUDA Graph与可变启动网格的冲突。作者未直接讨论负载不均问题，但理论上，如果某个线程块需要处理超长的键值（KV）缓存遍历，而其他块早早空闲，确实会导致部分流多处理器（SM）闲置，降低整体利用率。
然而，这种设计隐含了任务窃取或动态分配的机制——每个线程块循环获取下一个任务，直到队列为空，这天然具备一定的负载均衡能力。
- 在典型的推理负载中，请求长度分布通常较为集中，这种机制的效果应可接受。
- 但在极端动态的场景（如长短请求剧烈混合）下，可能仍需引入启发式规则或更精细的分组策略。
作者提到使用“启发式方法（heuristics）用于确定并行化方法”，表明团队已经考虑了这种权衡，实际效果有待后续更详尽的性能数据验证。
一、为何Triton能为vLLM提供助力
vLLM的目标是在各类平台、模型和执行策略下，实现最优的推理性能。这意味着它需要支持多种加速器及不同代际的硬件、各类模型架构，还要适配多样化的工作负载特征，比如不同的批次大小、序列长度和注意力模式。
一种直观的实现思路是编写大量高度专用化的核函数，为每种模型和GPU架构单独调优。这种方式虽能实现极致性能，但可扩展性极差。要在多款GPU平台（如英伟达Hopper、Blackwell，超威半导体MI300，英特尔GPU，以及未来平台）上维护数百个核函数，很快就会变得不可操作。
因此，团队更倾向于采用性能可移植的核函数，而Triton后端正是遵循这一思路设计的。
Triton是一门领域特定语言（DSL），开发者可以用Python编写GPU核函数（如矩阵乘法、注意力计算相关核函数），这些代码会被编译为可在多款平台上高效运行的GPU代码。Triton的分块编程模型实现了一种精妙的平衡：其底层特性足以让开发者实现与硬件紧密相关的优化，而高层特性又能让代码在很大程度上与硬件解耦。
如图1所示，开发者以逻辑分块为单位编写计算逻辑，Triton编译器和自动调优工具（autotuner）会决定这些逻辑分块如何映射到底层硬件。
图1：Triton的分块编程模型。编译器和自动调优工具会将逻辑分块映射为适配特定硬件的执行布局。代码通过tl.arange生成BLOCK_SIZE_N与BLOCK_SIZE_M维度的索引，拼接为2D tile坐标。不同GPU上的分块形状和执行布局可能差异显著，这些决策由工具自动完成。
不同GPU上的分块形状和执行布局可能差异显著，但这些决策会由工具自动完成，通常由自动调优工具指导。这项关于性能可移植性的研究强调了自动调优在跨平台GPU编程中的必要性。
二、vLLM中的Triton注意力后端
在大语言模型中，注意力计算通常是对性能影响最大的核心操作。为了降低系统复杂度，vLLM引入了注意力后端这一抽象层，通过统一的API封装各类注意力计算的实现，使其与线性层、层归一化等相对简单的组件解耦。
在这一架构下，vLLM支持多款注意力后端，包括CUDA平台上的FlashAttention和FlashInfer、基于ROCm的注意力后端，以及适用于MLA（多线性注意力）风格的专用后端。Triton注意力后端完全基于Triton实现，是vLLM的原生后端。
研发该后端的初衷，正是为了解决前文提到的性能可移植性和依赖项管理难题。它的同一份源码可在英伟达、超威半导体、英特尔的GPU上运行，仅依赖PyTorch和Triton两个库，且作为vLLM的内置组件随时可用。 该后端最初由IBM研究院和红帽人工智能团队开发，目前已由更广泛的社区负责维护和扩展。
三、Triton注意力后端的使用场景
- 在基于ROCm运行的超威半导体GPU上，Triton注意力后端是默认选择。
- 在英特尔XPU上也会启用该后端。
- 当运行32位浮点精度（FP32）计算时，vLLM会回退到Triton注意力后端，因为FlashAttention在该精度下无法在上述部分硬件上运行。
- 支持需要特定特性的模型，比如StepFun音频模型所使用的ALiBi平方根注意力机制、sink tokens（用于稳定注意力计算的特殊令牌）和GPT-OSS相关行为，尤其在英伟达Hopper架构之前的GPU（如A100）上表现良好。
- 它还支持小注意力头尺寸的模型、编码器-解码器注意力计算以及多模态前缀注意力计算。
由于Triton注意力后端是vLLM的内置组件，当FlashAttention、FlashInfer或其他依赖项无法使用或导入失败时，它会作为兜底后端生效，也支持批次不变性等特性。
四、基于Triton编写高性能、可移植的分页注意力核函数
Triton注意力后端的开发启动时，团队首先在vLLM框架外实现了核心核函数，并通过大量微基准测试对其进行独立评估。核函数的API按照vLLM的需求设计，而性能调优工作则在端到端集成前完成。
微基准测试至关重要，它能帮助团队理解核函数在预填充密集型（模型推理初期填充上下文的计算密集型场景）、解码密集型（模型生成新令牌的逐步计算场景）、混合负载场景下的性能表现，也能分析不同批次大小和上下文长度对性能的影响。
下面图2展示了具有代表性的微基准测试结果，横轴为令牌总数，纵轴为归一化延迟，不同子图分别对应纯预填充、混合、纯解码负载场景。
图2：多款Triton分页注意力核函数变体在预填充、解码、混合负载下的微基准测试对比。结果表明，不同的核函数变体在不同场景下各有优势，没有单一的配置能在所有场景中表现最优。
微基准测试是对端到端基准测试的补充，它能暴露核函数级别的性能细节，而这些细节往往会被系统级的干扰所掩盖。
五、分页注意力核函数的作用
分页注意力通过对键值（KV）缓存进行分页管理，实现了内存高效的注意力计算。对于批次中的每个查询，核函数会处理每个查询令牌；针对每个令牌，核函数会遍历查询头及对应的键值头，随后遍历分页的键值缓存，计算注意力分数并应用值向量。
图3展示了这一计算结构：查询令牌沿横轴排列，查询头沿纵轴排列，对分页键值缓存的遍历构成最内层循环。
图3：分页注意力的概念示意图，展示了查询令牌、查询头以及对分页键值缓存的遍历过程。
六、基于查询块优化tl.dot的分块尺寸
注意力计算的核心是矩阵乘法，在Triton中由tl.dot实现。但要实现高性能，需要设置足够大的分块尺寸以充分利用硬件资源，而单纯加载分页的键值缓存并不能达到理想的性能效果。
键值侧的分块尺寸受键值缓存的页大小限制，因此优化工作主要集中在查询侧。对于分组查询注意力（GQA），可将单个键值头对应的所有查询头批量处理，提升缓存的复用率。为了进一步提升并行度，团队将多个查询令牌组合为一个工作单元，称为查询块（Q block）。
图4阐释了这一方法：启动网格覆盖批次大小和键值头维度，而查询块则决定每个核函数实例处理的查询令牌数和查询头数，自动调优工具会为各平台选择合适的块尺寸。
图4：查询块将多个查询头和查询令牌组合为单个工作单元，提升tl.dot的硬件利用率和缓存复用率。
七、借助并行分块Softmax实现并行化
批量处理多个查询令牌的方式在预填充负载场景下表现优异，但在解码负载场景中毫无优势——因为解码时每次仅处理一个查询令牌。为解决这一问题，团队通过并行分块Softmax引入了额外的并行化机制，该核函数也被称为“3D核函数”。
这一方法将对键值缓存的遍历操作拆分到多个核函数实例中执行，每个实例计算部分结果，后续通过归约操作得到最终输出。由于Triton不提供全局屏障，该归约操作需要启动第二个核函数，这就需要在提升并行度和增加核函数启动开销之间做权衡，团队通过启发式算法判断该方法的适用场景。
八、CUDA图、启动网格与GPU执行波次
CUDA图通过记录并复现固定的执行图，能显著减少核函数的启动开销。但注意力核函数存在一个难点：其启动网格通常由批次大小和序列长度决定，具有可变性。
GPU通过固定数量的流多处理器（SM）执行核函数，当启动的线程块数量超过SM数量时，计算会以波次（wave）为单位依次执行。图5展示了这一行为，后续的波次（尤其是未完全占满SM的尾波次）会导致硬件资源利用不充分。
图5：启动线程块数超过可用流多处理器数量时的GPU执行波次示意图。尾波次（wave 1）可能导致SM利用不充分。
当该执行过程被捕获到CUDA图中后，即便实际的工作负载规模减小，这种资源利用的低效问题也会被复现。图6展示了固定的启动网格如何导致额外的无效计算，进而增加延迟。
图6：通过CUDA图复现固定启动网格时产生的额外无效计算。
九、从可变启动网格到持久化核函数
早期版本的分页注意力核函数采用可变启动网格，网格规模随工作负载大小动态调整，如图7所示。这种方式虽具备灵活性，但与CUDA图的兼容性较差。
图7：早期分页注意力核函数所使用的可变启动网格。
为解决CUDA图兼容性问题，团队设计了持久化核函数。该方案会启动固定数量的核函数实例（线程块），实例数与硬件的可用计算资源（SM数量）相等，每个实例通过读取GPU全局内存中的元数据，动态确定自身需要处理的计算量。这一设计让启动网格保持固定，从而实现了CUDA图的高效捕获与复用。
图8：持久化核函数方案，采用固定启动网格和动态任务分配机制。
十、基准测试结果
2025年末的基准测试结果验证了该方案的有效性。图9展示了在英伟达H100和超威半导体MI300上，运行Llama 3.1 8B模型、批次大小为1、输入长度为500令牌时的端到端延迟结果，横轴为输出令牌长度。
图9(a)：H100 GPU上不同注意力实现方案的归一化延迟对比。优化后的Triton方案与FlashAttention V3性能相当。
图9(b)：MI300 GPU上不同Triton注意力实现的归一化延迟对比。优化方案相比Baseline有显著提升。
- 在H100上，对于长解码请求，Triton注意力后端的性能达到了FlashAttention 3的100.7%。
- 在MI300上，其性能较早期实现提升了约5.8倍。
值得注意的是，两款平台使用的是同一份Triton核函数源码。基于Triton的分页注意力实现仅有约800行代码，而FlashAttention 3的代码量约为7万行。这充分展示了通过精心设计和高层抽象实现高性能核函数的潜力。
十一、前瞻：Helion中的分页注意力
Helion是PyTorch团队推出的一款全新的领域特定语言，可以将其视为更高层级的Triton，或者支持分块编程的PyTorch。它提升了抽象层级，旨在让开发者更容易编写出高效且正确的内核，同时优化自动调优流程。
Helion会通过约10分钟的自动调优来评估海量可能的实现方案，以提升内核在不同硬件上的性能移植性。它可以自动完成张量索引、掩码处理、网格尺寸计算、搜索空间定义等繁琐工作。
团队已在Helion中实现了一个简化版的分页注意力核函数作为实验，初期结果表现喜人。图10和图11展示了在H100和MI300X上，使用Helion实现的注意力后端（分为静态形状和动态形状版本）与原生Triton后端的性能对比。
图10：H100上Helion实现与Triton注意力后端的端到端性能对比。
图11：MI300X上Helion实现与Triton注意力后端的端到端性能对比。
相关工作已在PyTorch博客发布，代码也以草稿形式提交了拉取请求至vLLM代码仓库，展示了未来开源实战中进一步简化高性能内核开发的可能性。
结论
随着模型、推理优化技术和硬件平台的不断发展，性能可移植性的重要性日益凸显。vLLM中的Triton注意力后端证明，通过单一、可移植的核函数实现，也能达到顶尖的注意力计算性能。
通过精心的核函数设计、大量的微基准测试，以及持久化核函数、CUDA图等系统级优化手段，Triton后端在保持跨GPU厂商（英伟达、超威半导体、英特尔）可移植性的同时，其性能持平甚至超越了部分高度专用化的实现方案。
本文概述了Triton注意力后端中最重要的优化手段。随着像Helion这样更高层级的DSL出现，编写高性能、可移植的人工智能计算内核的门槛有望进一步降低，这对于整个开发者社区和云栈生态来说都是一个积极的信号。
发表于 3 小时前
|
查看: 3|
回复: 0
