Triton 编译器旨在为 AI 内核生成跨硬件的性能可移植代码和运行时。为了保持编译器生成代码的最先进（SOTA）水平，Triton 开发者社区一直在推动算子调度、内存分配、布局和管理方面的改进——涵盖 Triton DSL 层级以及底层（例如 Gluon 和 TLX）。
随着内核、优化空间和加速器硬件日趋复杂，内核编写者和维护者很难维持 SOTA 性能。Warp 特化（Warp specialization）已成为一种在 GPU 上实现开箱即用式性能提升的热门技术——其核心思想是为每个 Warp 提供特化的代码路径，而不是执行相同的代码。这减少了控制流分歧带来的性能损失，改进了延迟隐藏，并更有效地利用了 GPU 上的硬件单元。
Warp 特化在编译器中实现为一系列降低（lowering）过程，在 JIT 时刻对操作进行特化，通过搜索计算与内存管理、调度、针对底层硬件单元的特化以及同步的空间来完成。生成最优的 Warp 特化代码以达到硬件 Roofline 性能是一个组合优化问题。
Warp 特化是许多用例的基础设施。它帮助内核作者专注于算法优化，而不必担心“具体如何实现”，这在内核变得复杂时尤为重要。它可以针对硬件拓扑结构和工作负载异构性进行特化。它支持对复杂内核和优化进行特化，包括大型融合内核（“megakernels”）。
在本文中，我们将概述 Triton 中 Warp 特化的当前设计（“autoWS”），并讨论我们的未来规划。我们将与 Triton 开发者社区共同构建这一功能，并欢迎对文中提出的计划提供反馈。
当前的 Warp 特化
实现说明： autoWS 是基于 OSS Triton 构建的。它目前正在 Meta 的 OSS 镜像 facebookexperimental/triton 中积极开发，并已部分 合入上游。autoWS 可通过调优配置（在 ForOp 中设置 warp_specialize=True）来启用，适用于手动编写、TorchInductor 和 Helion 生成的内核。autoWS 中的编译器支持仍有限且处于实验阶段，我们正致力于将其通用化（详见下文路线图）。
@triton.jit
def mykernel(...):
...
for start_n in tl.range(lo, hi, BLOCK_N,
warp_specialize=warp_specialize):
…
编译器通过对 warp_specialize 代码区域内的代码路径（算子和数据）进行分区，为每个 Warp 特化代码路径，从而优化控制流分歧、延迟隐藏和专用硬件单元的利用，同时保持内核代码指定的正确性和数值特性。
当前及未来的 GPU 架构包含日益复杂的流水线，需要先进的编译器支持，才能从极其复杂的内核中获得 Roofline 性能。目前的 autoWS 实现支持 Hopper 和 Blackwell 加速器。我们首先对当前的 autoWS 设计进行高层概述。
Warp 特化使用以下过程（按顺序）：
- 数据分区：增加可调度的 GEMM/算子数量，以实现资源使用的完全重叠。
- 基于启发式方法创建软件流水线（SWP）调度，并通过属性传递决策（循环调度器）。
- 将代码划分为不同的 Warp 分区（分区调度器）——决策通过属性传递。
- 分析并创建分区间的通信缓冲区（缓冲区创建），缓冲区可位于共享内存（SMEM）或张量内存（TMEM，针对更新一代的 NVIDIA 加速器）。
- 对缓冲区拷贝和缓冲区重用做出决策（内存规划器）——决策通过缓冲区分配上的属性传递。
- 创建封装了分区间数据流和同步的生产者-消费者通道。将通道正确地降低为缓冲区和屏障操作，并将代码拆分为多个 Warp 分区（代码分区器）。
决策制定过程包括循环调度器、分区调度器、缓冲区创建中的启发式算法（选择将通道放置在 SMEM 还是 TMEM），以及内存规划器。
分区调度器。分区调度器将代码划分为一个或多个 Warp 分区，并将这些分区作为算子属性传递。它目前使用简单的启发式算法，基于 Triton 中的 NVIDIA Warp 特化。分区策略的示例包括当前支持的计算分区（如用于张量核心 MMA 的 Gen5 操作）、数据分区（TMA 加载）以及尾声（epilogue）和修正分区；未来工作包括混合分区（数据分区内的 CUDA 操作）、TMA 加载的拆分/混合、硬件 SFU 操作以及带有计算的原子加法。Warp 分区形成后，使用缓冲区/通道设置分区间的数据通信。
软件流水线（SWP）调度器。我们修改了 SWP 调度器以利用数据分区产生的数据并行性。改进后的调度器通过在循环中将依赖操作重新排序，使其尽可能远离，并用数据无关的副本替换它们，从而减少等待数据的时间。Flash Attention 前向传递的示例如下。
| SWP 之前 | SWP 之后 |
def dp_fa_fwd(...): ... for (...) qk1 = tl.dot(q1, k) qk2 = tl.dot(q2, k) p1 = softmax(qk1) p2 = softmax(qk2) acc1 = tl.dot(p1, v) acc2 = tl.dot(p2, v) |
def dp_fa_fwd(...): ... qk1 = tl.dot(q1, k0) qk2_prev = tl.dot(q2, k0) p1 = softmax(qk0) acc1 = tl.dot(p1, v0) for (...) qk1 = tl.dot(q1, k) p2 = softmax(qk2_prev) acc2 = tl.dot(p2, v_prev) qk2_prev = tl.dot(q2, k) p1 = softmax(qk0) acc1 = tl.dot(p1, v) p2 = softmax(qk2_prev) acc2 = tl.dot(p2, v_prev) |
在上面的示例中，Attention 前向循环的连续迭代被流水线化，重叠了来自不同循环迭代的两个点积和 Softmax 计算。
流水线语义通过 Warp 特化过程得到保留。目前，这已针对 Attention 的前向传递实现，因为它主要基于 tl.dot 的独立链。我们计划对该实现进行通用化。
内存规划器。内存规划器决定每个通道使用多少缓冲区以及通道应如何重用缓冲区。请注意，早期的过程会在缓冲区创建期间为通道分配 TMEM 或 SMEM。
内存规划器使用通道感知的活跃度和依赖链分析，在生产者-消费者模式之间分配和重用 TMEM 和 SMEM 缓冲区。它根据每个最内层循环对分配进行分类；最内层循环获得多重缓冲分配。TMEM 中的操作按累加器类型（操作数 D，即存储 MMA 结果的累加器）、大小和活跃范围进行优先级排序。
如果活跃范围不重叠，内存规划器会积极地重用已分配的缓冲区。它会在同一循环的依赖链上，或同一分区内的不同循环上重用缓冲区。如果缓冲区与所有已分配缓冲区重叠且现有分配无法重用，则分配新空间。
代码分区器。一旦内核操作完成分区且缓冲区规划完成，下一步就是连接各分区以设置内核计算。代码分区器进行小范围的指令排序和同步。重新排序的一个示例是，如果来自不同数据分区的操作最终落在同一个 Warp 分区中，则对它们进行重新排序，以减少活跃范围和寄存器压力。
代码分区器按如下方式设置同步机制：对于每个通道，我们假设源操作和目标操作处于相同作用域，并将执行相同的次数。代码分区器使用一个累积索引计数来跟踪执行次数，并根据累积索引计数和通道缓冲区数量计算缓冲区索引和相位。代码分区器使用两个屏障实现通道源和目标之间的同步。如果源或目标来自 TMA/gen5，它可以使用与 TMA/gen5 GEMM 相关联的现有屏障。
内存规划器对缓冲区重用做出决策，并使用缓冲区 ID、偏移量、拷贝等标注分配操作。如果两个通道具有相同的缓冲区 ID，则使用相同的内存分配。在为通道生成屏障操作时，如果两个通道 A 和 B 重用相同的空间，代码分区器会为两个通道使用同一组屏障和单个累积索引计数。累积索引计数是两个通道的源或目标操作的合并执行次数。这保证了正确性：通道被链接并使用正确的缓冲区索引，或者在 A 和 B 之间强制执行正确的同步。
autoWS 的实现展示了 Flash Attention 前向内核性能接近手动调优的底层实现 [Triton Conference, 2025]。我们在 B200 上针对各种 Attention 和序列长度配置进行的 Flash Attention 前向传递内核基准测试显示，TFLOPS 数字接近 Gluon 和 cuDNN 实现，是原生 Triton 的 1.5-2 倍（cuDNN 仍领先 10-20%）。autoWS 基准测试基于 Helion 自动调优和 ptxas 高级编译器配置。
近期方向（< 一年）
基于配置文件的分区调度
我们计划将基于配置文件的优化作为未来方向（见下文“未来方向”）。我们打算从一个战术性但循序渐进的步骤开始——在内核内实现更接近 Roofline 的算子分区。我们从算子性能（如 TMEM 和 SMEM 读写）以及区域和端到端执行时间的离线估计开始。在分区过程中，autoWS 将估计分区之间算子到算子通道的通信开销。
基于配置文件的 SWP 调度
当前的 SWP 调度可以通过使用运行时配置文件生成更接近 Roofline 的调度。我们计划扩展 SWP 调度器，使其使用带有数据依赖分析的算子配置文件来改进生成的调度。这可以利用具有更好外部循环支持（例如，通过循环展平）的标准模调度（modulo scheduling）。我们可以通过自动调优代价评估中排名靠前的调度来找到最接近 Roofline 的调度。
以 Flash Attention 后向传递内核为例。我们可以构建依赖图并用延迟信息标注每个算子。这使得编译器流水线优化（如模调度）成为可能。我们还可以对少数最佳调度进行排名以进行自动调优。下图可视化了一个玩具示例，展示了如果调度器使用延迟配置文件，Flash Attention 后向传递的 SWP 调度改进。图的上半部分显示了没有基于配置文件的流水线执行操作（箭头显示数据依赖）；有了延迟信息，编译器可以像下半部分那样“装箱”操作的执行。
内存规划器改进
内存规划器需要找到带有缓冲区重用的最优 TMEM/SMEM 分配通道集，这往往处于组合搜索空间中（随着内核和硬件复杂性增加尤甚）。我们正在考虑改进内存规划的机制：在 DSL 层级提供用户注解以指导规划，并估计内存规划的代价，并在选定的规划上进行自动调优。
乒乓（Ping-Pong）调度
乒乓调度是指在调度需要关键硬件资源（如 SFU、SMEM/TMEM）的长运行、高占用代码区域时强制执行排他性。当这些资源发生争用时，优先考虑产生即需数据的 Warp 可以显著提高硬件利用率和内核性能。我们正在开发一个识别并调度这些关键部分的过程（在代码分区之前执行并设置算子上下文）。
我们目前通过对有限的一组操作进行模式匹配（基于硬件类型上的离线性能调优实验）来识别关键操作部分。该过程使用预定义规则（例如，包含算术操作的内存操作）识别区域边界。关键部分周围的屏障同步可能会降低性能——我们目前开放了乒乓调优配置供自动调优使用。
基于区域的显式子分块（Subtitling）
如果分块（tiling）导致寄存器压力、Bank 冲突和停顿等瓶颈，子分块可以作为分块程序的性能优化转换。目前作者可以在内核代码中手动进行子分块。如果子分块区域跨越多个 Warp 组分区，则需要改进编译器以正确处理它，并使用更细粒度的同步将通道细分为更小的通道。这解锁了生产者-消费者流水线的改进——当生产者的一部分完成后，消费者的相应部分即可开始。为了简化内核编写，我们可以为基于区域的子分块添加语法糖，即将子分块因子的规范作为 DSL 中内核作者的显式区域级原语公开。
可调试性和工具
Warp 特化使模型和内核作者更难调试数值和性能问题，因为它会将转换优化为 Warp 粒度的调度、内存规划和布局。我们正在构建工具和 IR 支持，使作者能够调试内核代码生成和执行。
新工具。我们计划将 Triton TTGIR 转换为可读的 TLX 内核，以便于调试和进一步的性能手动调优。为了让作者易于理解编译器决策，我们正在构建可视化 Warp 特化代码的工具，例如 Warp 分区、带有通道的依赖图、内存分配和规划决策以及 SWP 调度。我们正在探索为作者添加交互式工具的可能性，以便进行编译时决策并继续后续的流程处理。
IR 改进。我们的实现目前在 IR 层级不使用通道的抽象表示。我们正在考虑使用上游 Triton 的 aref 作为额外的抽象。aref 过程流水线在插入 aref 后立即进行 aref 降低，我们计划添加对 aref 的支持以表示缓冲区重用，并推迟降低，以便流程能够利用该抽象。
通用性和稳定性
我们正在致力于通用化启发式算法并稳定转换过程，以支持更多种类的内核：Flash Attention 后向、Flex Attention、Jagged Attention 内核等。
SOTA 硬件特化
添加对 SOTA 和新兴硬件功能的支持对于实现内核的 Roofline 性能是必要的。我们正在添加对 NVIDIA Blackwell 的 Cluster Launch Control、分布式共享内存、多 CTA 和片上 TMA 描述符流水线等功能的支持；以及 AMD 波特化、乒乓和多流调度。这是一个持续的努力，因为我们期望支持新的硬件发布。
未来方向
基于模型的全局优化
生成能够达到内核和硬件类型 Roofline 性能的代码是一个组合搜索问题。我们正在探索使用硬件和操作的代价模型作为机制，来剪枝分区和同步、调度、内存规划和张量布局的联合搜索空间。代价模型可以是静态指定的（例如操作延迟和硬件规范），也可以是基于运行时基准测试和性能分析数据学习的。我们计划使用 Triton-MPP（多过程性能分析器）来测量跨硬件类型的内核基准测试的操作延迟，并使用配置文件来指导优化过程。
准确的代价模型将使编译器能够为任何硬件类型生成高性能的复杂内核代码。我们正在探索构建一个全局规划器，可以通过结合数据依赖分析和代价模型，在调度、同步和内存管理之间进行联合优化。全局规划器可以输出最优的 SWP 调度、通道和缓冲区重用配置、操作分区和内存分配。
内核/算子融合与 Megakernels
内核和算子融合可以通过减少 GPU-CPU 上下文切换开销并提高内存局部性和压力，显著优化模型端到端性能。实现融合往往是在算法、调度、张量布局和内存规划维度上的组合搜索。Megakernels 将内核融合的思想带到模型中，为前向和后向传递计算实现一个或少数几个内核。
内核/算子融合由 PyTorch 使用 Inductor 层的 Triton 模板通过算法优化完成。虽然这在许多情况下有效，但它不支持与用户定义的 Triton 内核的融合。Triton 编译器可以在 IR 层级实现手写内核和自动生成内核之间的瓦片（tile）级融合。此外，对内核内的底层操作进行特化可以通过联合算法和性能优化，使模型工作负载更接近 Roofline 性能。
我们的目标是在一系列依赖或独立的内核之上构建积极的融合支持。在内核内部和跨内核融合及分块循环是该目标的一部分。毋庸置疑，算法转换应该是可证明正确的。
数值确定性
我们正在考虑添加对 确定性 Warp 特化 的支持，这将允许内核/模型作者推理和控制程序特化及编译器决策。作者将能够利用 Triton 和 PyTorch DSL 中的数值支持作为数值确定性和稳定性的杠杆。我们在上文的近期方向部分涵盖了确定性的其他方面。
语言支持
语言支持使作者能够将预期工作负载和操作的领域知识传递给编译器（以生成高性能代码），同时控制和推理编译器决策及代码转换。我们正在探索包括 DSL 抽象在内的语言支持，这些抽象将计算与数据分离（例如 Cypress, PLDI 2025; Halide, PLDI 2013），并独立于数据属性表达“任务图”和调度；以及使编译器能够将程序区域和操作特化得更接近 Roofline 的作者提示。用于指定调度、数据属性和特化的声明式语言是一个潜在方向。语言选择往往很广泛，我们期待社区的反馈和建议。
结束语
在本文中，我们介绍了 Triton 中 autoWS Warp 特化的现状，以及我们改进 Triton 编译器、工具和语言支持的路线图和想法。我们邀请社区对我们的路线图提供反馈和建议，并分享您对如何改进 autoWS 的想法。
致谢： 感谢我们的领导团队 Alexey Loginov, Bill Yoshimi, Ian Barber 和 Parthiv Patel；以及 Triton 团队和 Meta 内部客户对 autoWS 开发的支持。我们感谢 NVIDIA 和 OpenAI 在 autoWS 上的合作。