Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References

文章标题: Tawa：利用异步引用为现代GPU提供自动Warp专业化
作者/机构: Hongzheng Chen∗†, Bin Fan, Alexander Collins, Bastian Hagedorn, Evghenii Gaburov, Masahiro Masuda, Matthew Brookhart, Chris Sullivan, Jason Knight, Zhiru Zhang†, Vinod Grover
所属机构: NVIDIA, †Cornell University

A1 主要贡献

核心问题: 现代GPU架构（如NVIDIA Hopper和Blackwell）已经从同构处理器演变为集成了多个专用硬件单元（如Tensor Cores和Tensor Memory Accelerator, TMA）的复杂异构系统，支持细粒度的异步数据流执行模型。然而，主流的GPU编程范式，即单指令多线程（SIMT），其为同步、数据并行执行而设计，与硬件的异步、任务并行能力存在根本性的错位。这导致了一个显著的“可编程性鸿沟”：虽然硬件层面的Warp专业化（即不同的warp承担不同的角色，如数据生产和计算消费）是释放硬件潜力的关键，但它迫使开发者手动编排复杂的底层通信和软件流水线，这一过程耗时、易错且难以维护。

研究目标: 本文旨在解决上述可编程性鸿沟，目标是创建一个能够从高级、基于Tile的程序中自动、系统地生成高性能、Warp专业化的代码的编译器。该编译器应能将开发者从繁琐的内核重写工作中解放出来，自动处理角色划分、数据流管理和流水线编排。

创新点与主要贡献:
为了应对上述挑战，本文提出了Tawa，一个自动化编译器，它通过分析从基于Tile的DSL生成的计算图，将程序划分为Warp专业化的组件，并管理通信和流水线。其主要贡献如下：
1. 提出异步引用（aref）: 这是一种用于Warp级通信的新型IR抽象。aref 表达了异步数据传输的意图，而无需暴露底层的硬件原语（如mbarriers），在IR层面提供了一种清晰、可移植的表示方法。
2. 实现Tawa自动编译流程: Tawa是一个用于NVIDIA GPU上Warp专业化的自动编译流程。它以Triton程序作为输入，进行任务感知的Warp组划分，应用多粒度流水线，并生成高性能的PTX代码。
3. 高性能评估结果: 在NVIDIA H100 GPU上对代表性的LLM内核进行评估，结果表明Tawa实现了高硬件利用率。
* 对于GEMM，Tawa的性能比高度优化的cuBLAS GEMM内核最高提升1.1倍。
* 对于Attention负载，Tawa比Triton快1.2倍，并且与手工优化的CUTLASS C++ FlashAttention-3内核性能相当，但编程工作量远小于后者。

A3 背景知识

A. 现代NVIDIA GPU架构

流式多处理器（SM）作为基本计算单元。如图1所示，流式多处理器（SM）是现代GPU的基本计算单元，包含了计算和内存资源。每个SM集成了一个大型寄存器文件、一个每个时钟周期可向一个warp（32个线程）分发指令的warp调度器，以及多个处理通用算术的CUDA核心。此外，SM还包括用于加速矩阵计算的专用Tensor Cores和一个作为软件管理的共享内存（SMEM）的内存块。一个线程块（CTA）被调度到一个SM上，并共享其寄存器和SMEM。在像Hopper这样的最新GPU代中，对Tensor Cores的依赖性增加，其贡献了总计算吞吐量的90%。因此，在现代GPU上实现高性能的关键在于用充足且及时的数据来充分饱和这些Tensor Cores。


图1：简化的NVIDIA H100 SXM5 GPU架构。

Hopper架构的两项关键硬件创新。为了提高Tensor Core的利用率，Hopper架构引入了两项关键硬件创新：Tensor Memory Accelerator（TMA）和Warp Group Matrix Multiply-Accumulate（WGMMA）。TMA能够实现从全局内存（GMEM）到共享内存的硬件管理的多维数据传输，将地址生成和传输协调任务从软件卸载到一个专用单元。为了进一步实现这种重叠，Hopper通过异步事务屏障（asynchronous transaction barriers）扩展了编程模型，该屏障允许线程不仅基于到达情况进行同步，还基于交换的数据量进行同步。每个事务携带一个字节计数，只有当所有生产者都到达并且达到预期的总事务计数时，屏障才会释放线程。这种机制在协调跨线程块的异步内存拷贝和集体数据交换方面尤其强大。

WGMMA与Warp专业化的兴起。与此同时，WGMMA引入了Warp组级别的矩阵乘法操作，允许多个warp（4个warp作为一个warp组）以异步高效的方式协同执行大块（large-tile）GEMM内核。这些机制共同为warp专业化奠定了基础，即一个SM内的不同warp承担数据通信和计算的不同角色，从而提高LLM工作负载的吞吐量【【10，Thunderkittens: Simple, fast, and adorable kernels，2025，The Thirteenth International Conference on Learning Representations】–【12，Flashinfer: Efficient and customizable attention engine for LLM inference serving，2025，arXiv preprint arXiv:2501.01005】】。这些趋势揭示了更广泛的架构转向异步数据流执行的趋势，这一方向在后续的Blackwell等架构中得以延续【【7，Nvidia blackwell architecture technical brief，2025，https:// http://resources.nvidia.com/en-us-blackwell-architecture】】。

B. 现有的GPU编程模型

传统库与专用库。CUDA是GPU编程的主流接口【【13，Scalable parallel programming with cuda: Is cuda the parallel programming model that application developers have been waiting for?，2008，Queue】】，但随着架构的演进，编程复杂性显著增加。NVIDIA的CUTLASS C++库【【9，Cutlass，2025，https://github.com/NVIDIA/cutlass】】为GEMM及相关计算提供了模板化的内核，但利用它来编写神经网络算子需要复杂的配置和手动调优。为了减少领域工作负载的工程量，出现了一些专用库。ThunderKittens【【10 ，Thunderkittens: Simple, fast, and adorable kernels，2025，The Thirteenth International Conference on Learning Representations】】引入了基于16×16 tile抽象的紧凑C++接口，提供了性能有竞争力的GEMM和attention内核。FlashInfer【【12，Flashinfer: Efficient and customizable attention engine for LLM inference serving，2025，arXiv preprint arXiv:2501.01005】】专注于高效的attention和KV-cache管理，并通过JIT专业化支持大规模LLM推理。这些系统简化了对高性能内核的访问，但仍以库为中心：它们提供优化的算子模板，而不是针对不同应用的通用策略，这可能导致为不同尺寸矩阵手动调优性能时出现大量的维护问题。

领域专用编译器基础设施。为了进一步提高抽象层次，领域专用的编译器基础设施被提了出来。TVM【【14，Tvm: An automated end-to-end optimizing compiler for deep learning，2018，Proceedings of the 13th USENIX Conference on Operating Systems Design and Implementation】】开创了一个带有自动调度【【15，Ansor: Generating high-performance tensor programs for deep learning，2020，Proceedings of the 14th USENIX Conference on Operating Systems Design and Implementation】、【16，Tensor program optimization with probabilistic programs，2022，Advances in Neural Information Processing Systems】】的端到端编译框架，而Relax【【17，Relax: composable abstractions for endto-end dynamic machine learning，2025，Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2】】通过图级IR将其扩展到动态工作负载。Fireiron【【18，Fireiron: A data-movement-aware scheduling language for GPUs，2020，Proceedings of the ACM International Conference on Parallel Architectures and Compilation Techniques】】也利用调度语言来指定分块和数据移动，Graphene【【19，Graphene: An IR for optimized tensor computations on GPUs，2023，Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】】通过提供一个将具有各种布局的张量映射到Tensor Core指令的IR来进一步增强这一思想。虽然这些系统功能强大，但它们是在Hopper之前开发的，并未直接利用TMA和硬件warp专业化，因此留下了显著的性能未被发掘。同样，许多关于GPU warp专业化的工作【【20，Harmonic CUDA: Asynchronous programming on GPUs，2023，Proceedings of the 14th International Workshop on Programming Models and Applications for Multicores and Manycores】–【23，Wasp: Exploiting gpu pipeline parallelism with hardware-accelerated automatic warp specialization，2024，2024 IEEE International Symposium on High-Performance Computer Architecture (HPCA)】】也是在Hopper架构之前提出的，没有有效地利用最新的TMA特性。

新兴的基于Tile的内核语言。为了实现更精细的控制，一类新的基于Tile的内核语言应运而生。Triton【【8，Triton: An intermediate language and compiler for tiled neural network computations，2019，Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages】】建立在MLIR【【24，MLIR: Scaling compiler infrastructure for domain specific computation，2021，Proceedings of the 2021 IEEE/ACM International Symposium on Code Generation and Optimization】】之上，允许程序员用Python编写高效内核，同时抽象化分块和内存管理。尽管在普及内核开发方面取得了成功，Triton仍然缺乏对warp专业化的内在抽象，导致在Hopper架构上存在性能差距。最近的工作如TileLang【【25，Tilelang: A composable tiled programming model for ai systems，2025，arXiv preprint arXiv:2504.17577】】和Mosaic【【26，Pallas:Mosaic GPU，2025，https://docs.jax.dev/en/latest/pallas/gpu/ index.html】】通过引入用于数据流调度和流水线组合的显式API扩展了这一方向。Cypress【【27，Task-based tensor computations on modern GPUs，2025，Proceedings of the ACM on Programming Languages】】提出了一个基于任务的编程模型，通过逻辑计算描述和映射规范来抽象复杂的数据移动和同步。然而，程序员仍然需要手动管理线程、warp和SM的多级层次结构，以及相应的数据分解。Gluon【【28，Gluon，2025，https://github.com/triton-lang/triton/tree/ main/python/examples/gluon】】为Triton提出了一个添加了布局抽象的IR，但仍需要程序员显式地集成通信操作。这些方法为专家用户提供了富有表现力的控制，但它们继续将同步和资源管理的重大责任强加给程序员，这限制了程序员的生产力。

Tawa的定位与贡献。本文主张用户不应该处理底层的硬件细节，而是应自动生成warp专业化的程序。Tawa通过引入aref作为warp通信的IR级抽象，并自动生成支持TMA的warp专业化流水线，弥补了这些差距，从而在不增加底层编程负担的情况下实现了高性能。

A2 方法细节

A. 概述

Tawa的整体设计。图2a展示了Tawa的整体设计。由于Triton【【8，Triton: An intermediate language and compiler for tiled neural network computations，2019，Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages】】已成为PyTorch【【29，Pytorch: An imperative style, high-performance deep learning library，2019，Proceedings of the 33rd International Conference on Neural Information Processing Systems】、【30，Pytorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation，2024，Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS’24)】】事实上的后端，并被广泛用于编写高效的GPU内核，我们选择Triton作为我们编译器的前端接口。程序员在Triton-Python中编写带有分块计算和TMA通信的内核（图2b），这些内核首先被翻译成标准的Triton-MLIR表示。在此基础上，Tawa引入了一系列编译器遍（passes），直接从高级Python代码中自动生成warp专业化的程序。


图2：Tawa从Triton前端到内部MLIR表示的编译流程示例。tl是Triton语言的缩写。tt是Triton MLIR方言。代码片段为演示目的进行了简化。

Tawa的三大关键转换。为了解决第一节中概述的挑战，Tawa应用了三个关键的转换：（1）任务感知的划分（§III-C）。我们引入了一种新颖的算法，将程序划分为生产者和消费者warp组，确保正确的角色分配和通信边界，直接解决了挑战1。（2）异步引用抽象（§III-B）。我们设计了aref，一个显式捕获warp间数据流的中间表示。通过在IR级别表达通信意图，aref允许编译器自动插入和协调必要的同步和数据移动操作，解决了挑战2。我们还提出了一个新的Tawa MLIR方言来编码这些aref操作。（3）多粒度软件流水线（§III-D）。在aref程序之上，我们应用了一个软件流水线遍，它在多个粒度上编排通信和计算之间的重叠，解决了挑战3。

编译流程的最终阶段。最终生成的程序如图2c所示。一旦这些遍被应用，aref表示就会被降级到标准的Triton-MLIR方言，之后现有的Triton编译流水线将继续进行：代码被逐步降级到LLVM IR，最终到PTX，准备在NVIDIA GPU上执行。


(c) 高度简化的Tawa MLIR代码 (D = 2)

B. 异步引用（Asynchronous Reference）

aref抽象的概念。为了更好地为warp之间的通信建模，我们引入了异步引用（asynchronous reference）【【31，It’s about time - temporal abstractions for asynchronous gpu tensor computation，2025，In submission】】（简称aref），这是一种IR抽象，它将GPU上生产者和消费者之间的一个单槽通道（one-slot channel）模型化，如图3所示。它将一个数据缓冲区与两个由硬件mbarriers实现的同步原语打包在一起，通常命名为empty和full。在任何时刻，这两个屏障中恰好有一个编码了该槽的状态：当empty mbarrier持有一个信用（credit）时，生产者可以写入该槽；当full mbarrier持有一个信用时，该槽包含一个已发布的值，可供消费者读取。


图3：aref抽象及其相关操作。

aref接口的操作与语义。aref接口暴露了三个操作：put、get和consumed，其形式化语义在图4中定义。这些语义为将操作降级为GPU特定指令提供了严谨的基础（§III-E）。put是生产者的发布步骤：它要求槽为空（empty mbarrier E持有信用），将有效负载写入缓冲区，并通过在full mbarrier F上以释放顺序（release ordering）到达来将状态翻转为full。这一转换使数据对消费者可见，并防止后续写入过早地重用该槽。get是消费者的获取步骤：它要求槽为full（观察到full mbarrier F），执行将缓冲区数据读入程序变量的操作，并将aref转换为一个借用状态（borrowed state），此时E和F都不持有信用，反映出该值正在使用中，但该槽尚不可重用。consumed完成了握手：一旦消费者不再需要该值，它就在empty mbarrier E上到达，恢复empty信用并启用下一次put；这引发了从生产者的写操作到消费者的读操作，再到生产者后续重用的“发生于之前”（happens-before）链。因此，aref作为一个一级IR值，代表了一个具体的通信路径，强制了跨warp组的数据移动和计算的正确顺序。

图4：aref操作的运算语义。σ是存储映射，将aref标识符a映射到实际结构σ(a) = ⟨buf, F, E⟩，其中F是full mbarrier标志，E是empty mbarrier标志。初始状态E = 1, F = 0。
$$ \begin{array}{l} \text { PUT } \\ \frac{\sigma(a) . E=1}{\langle\sigma . p u t(a, v)\rangle \rightarrow\langle\sigma[a \mapsto(b u f=v, F=1, E=0)], e\rangle} \\ \text { GET } \\ \frac{\sigma(a) . F=1}{\langle\sigma . g e t(a)\rangle \rightarrow\langle\sigma[a \mapsto(b u f=\sigma(a) . b u f, F=0, E=0)], \sigma(a) . b u f\rangle} \\ \text { CONSUMED } \\ \langle\sigma . c o n s u m e d(a)\rangle \rightarrow\langle\sigma[a \mapsto(b u f=\sigma(a) . b u f, F=0, E=1)], e\rangle \end{array} $$

aref与其他模型的比较及其特性。与OpenCL pipe【【32，OpenCL 2.0 reference guide，2015，https://www.khronos.org/ assets/uploads/developers/presentations/opencl20-quick-reference-card. pdf】】或其他加速器编程模型中的类似机制【【33，oneAPI toolkits: Pipe，2023，https://www.intel.com/content/ www/us/en/docs/oneapi-fpga-add-on/optimization-guide/2023-1/ the-pipe-class-and-its-use.html】–【37，Heteroflow: An accelerator programming model with decoupled data placement for software-defined fpgas，2022，Proceedings of the 2022 ACM/SIGDA International Symposium on Field-Programmable Gate Arrays】】相比，aref抽象是专门为GPU warp间通信设计的。其语义植根于硬件支持的mbarrier同步，提供了一个清晰的排序模型。除了同步语义，aref还是类型通用的，支持类型为<T>的缓冲区，这使其能够携带结构化的有效负载（如张量或元组），同时重用相同的mbarriers。为了便于深度流水线，多个aref实例可以被分组到一个深度为D的循环缓冲区中，有效地形成一个支持跨流水线阶段存储和传输中间数据的环形结构。

C. 任务感知的划分（Task-Aware Partitioning）

划分流程概述。我们首先遍历MLIR计算图（如图5a），其中每个节点是一个MLIR操作，每条边捕获了使用-定义（use-def）依赖关系。我们分两步进行划分过程：划分注解和循环分布。

1) 划分注解（Partition Annotation）

基于语义标签的节点划分。我们从内核的具有副作用的汇点（sink）（例如，Store）开始，沿着使用-定义链进行向后遍历。在遍历过程中，我们通过检查每个节点的效果为其附加一个语义标签。那些对地址计算有贡献的节点，例如数据传输的指针算术，被标记为迭代语句（橙色边）。那些为实际计算转换或消费一个tile的节点，例如WGMMA，被标记为tile语句（蓝色边）。重要的是，迭代语句在IR中不一定是连续的，如图2b中高亮部分所示。

Hopper架构上的具体划分策略。在Hopper上，一个典型的划分是创建一个加载WG作为生产者和一个计算WG作为消费者，如图4所示。Tawa将迭代语句及其支配的TMALoads注解在生产者分区（WG0）中，将剩余的tile语句注解在消费者分区（WG1）中。注意，结尾部分（epilogue）也附加到WG1，以确保它被某个WG实际执行。对于Blackwell及后续支持更多WG的GPU代，Tawa也可以通过在图遍历期间注解语义角色，基于GPU上可用的warp角色来划分IR。

2) 循环分布（Loop Distribution）

通过aref建立跨分区通信。在操作被注解后，我们需要构建实际的IR来为每个WG划分工作负载。对于每个跨分区边（例如，图5a中的TMALoad→LocalAlloc），我们创建一个大小为D的aref张量，表示一个D槽循环缓冲区。在生产者侧创建一个put操作，将每个TMALoad结果放入一个aref槽中。槽索引计算为迭代k模缓冲区大小D，从而允许缓冲区槽被重用（图2c中的L15, L23）。类似地，在消费者侧创建一个get操作，从同一个aref槽中获取数据。在数据用于WGMMA中的计算后，在末尾插入一个额外的consumed操作，以指示该槽已被释放。如§III-B所述，aref支持在缓冲区中对多个值进行分组以共享mbarrier。因此，在创建aref时一个直接的优化是检查加载的张量是否在同一个WGMMA中使用。如果是，我们可以将张量的元组放入aref缓冲区，并共享WGs之间的通信通道。由此产生的计算图如图5b所示。由于aref已经携带了缓冲区，LocalAlloc操作也可以被消除，而是直接从aref缓冲区读取。

克隆循环以实现独立进展。为了使两个分区都能独立地在原始循环嵌套中前进，我们还需要围绕切分点执行循环分布。在IR中创建了两个WG区域，如图2c所示（L10-32）。scf.for循环被克隆，以便WG0和WG1各自携带一个关于k的同构循环。迭代语句保留在WG0中，tile语句保留在WG1中。warp间的通信通过代码中高亮的aref完成（L16, L24）。最后，计算结果从for循环中产生，并在结尾部分存回全局内存。最终的执行过程是图5c中warp专业化的时间线，该图显示了TMA加载如何与Tensor Core计算重叠。


图5：图2b中GEMM内核的任务感知划分。注意，warp专业化只是实现(c)中重叠的一种方式，实际延迟未按比例显示。

D. 多粒度流水线（Multi-Granularity Pipelining）

计算Warp组的进一步优化。在获得warp专业化程序后，我们对计算warp组进行进一步优化，引入两种流水线机制以确保硬件组件并行工作。具体来说，我们考虑一个重叠MMA地址计算和计算的细粒度流水线，以及一个重叠CUDA和Tensor Core计算的粗粒度流水线。

1) 细粒度流水线（Fine-Grained Pipeline）

问题：CUDA Core和Tensor Core的顺序执行。现代GPU同时拥有处理标量和地址计算的CUDA Cores，以及加速MMA操作的Tensor Cores；然而，简单的顺序执行迫使这些单元按序操作，导致地址生成会阻塞张量执行，反之亦然，从而导致利用率低下。

解决方案：自动流水线重叠地址计算与MMA。为缓解这种低效，我们引入了一种自动流水线算法，通过一个深度为P的有界流水线，将CUDA Cores上的地址计算与Tensor Cores上的MMA执行重叠。在每次迭代中，编译器异步发出下一个MMA指令（WGMMA.issue(k)），并且只有在未完成操作的数量达到最大流水线深度时才阻塞（WGMMA.wait() {pendings=P}）。由于第k次等待操作不需要在每次MMA发出后立即完成，这使得更多的MMA指令可以在等待前处于飞行状态。一旦流水线被填满，来自迭代k-P的结果引用被释放（aref.consumed(k-P)），这通过强制数据流依赖性并释放资源以供重用来维持正确性。在最后一次迭代之后，编译器在结尾部分通过消费剩余的待处理操作来排空流水线。这种设计有效地解耦了地址计算和矩阵乘法，维持了两个功能单元之间持续的工作流，并与顺序基线相比有效地隐藏了延迟并减少了空闲时间。


图6：细粒度流水线重叠MMA地址计算和计算。实际延迟未按比例显示。

2) 粗粒度流水线（Coarse-Grained Pipeline）

引入三阶段流水线。细粒度流水线在循环主要由矩阵乘法操作组成时最有效。然而，当CUDA Cores上执行额外计算时，这些操作也可以与MMA重叠。因此，我们进一步提出了算法1中的一个粗粒度的三阶段生产者-转换者-（可选）消费者流水线。每个循环迭代j有一个产生中间tile的Tensor Core阶段$T_j$（例如，一个矩阵乘法片段），一个转换该中间体的CUDA Core阶段$C_j$（例如，归一化、激活、规约），以及一个可选的消耗转换后结果的下游Tensor Core阶段$U_j$（例如，第二次矩阵乘法/结尾部分）。

算法 1: 粗粒度CUDA和Tensor Core流水线

输入: Tile数量 N; 布尔值 USE_U
1  Prologue;
2  MAYBEAREFGET(T0);
3  ISSUEANDCOMMIT(T0);
4  DOTWAIT(T0) ; // 物化T0输出
5  MAYBEAREFCONSUMED(T0);
6  COMPUTE(C0);
7  Steady state;
8  for j ← 1 to N − 1 do
9      MAYBEAREFGET(Tj ); ISSUEANDCOMMIT(Tj );
10     if USE_U then
11         MAYBEAREFGET(Uj−1);
12         ISSUEANDCOMMIT(Uj−1);
13     DOTWAIT(Tj−1) ; // 确保Tj−1结果可见
14     MAYBEAREFCONSUMED(Tj−1); COMPUTE(Cj−1);
15     if USE_U then
16         DOTWAIT(Uj−1);
17         MAYBEAREFCONSUMED(Uj−1)
18 Epilogue;
19 DOTWAIT(TN−1);
20 MAYBEAREFCONSUMED(TN−1);
21 COMPUTE(CN−1);
22 if USE_U then
23     MAYBEAREFGET(UN−1);
24     ISSUEANDCOMMIT(UN−1);
25     DOTWAIT(UN−1);
26     MAYBEAREFCONSUMED(UN−1)
27 FINALEPILOGUEANDSTORE();

流水线调度逻辑。序言（prologue）通过运行$T_0$至完成并立即计算$C_0$来启动流水线。稳态（steady-state）则重叠$T_j$与$C_{j−1}$（以及，如果存在，$U_{j−1}$）：在每次迭代中，我们首先将当前tile的Tensor Core工作入队并提交，可选地将前一个tile的下游U入队并提交，然后在产生$T_{j−1}$结果的点积组上执行精确的完成等待，并运行CUDA Core转换$C_{j−1}$。结尾部分（epilogue）排空最后一个CUDA Core转换，如果使用了U，则发出并完成$U_{N−1}$。正确性依赖于两种同步机制：DOTWAIT(Tj−1)（和DOTWAIT(Uj−1)）确保在Tensor Core边界处计算组的完成，而MAYBEAREFGET(·)是包装器，仅当一个阶段实际消耗跨WG的aref对象时才执行get操作。当不需要读取外部数据时，包装器编译为空操作，保持模板的最小化并避免不必要的停顿。

从MLIR图自动构建调度。调度是通过几个有原则的分析和随后的机械合成步骤从MLIR计算图中构建的。阶段识别使用方言和类型级别的线索将每次迭代的子图划分为$T_j$、$C_j$和（可选地）$U_j$：Tensor Core微块及其粘合逻辑构成T（如果存在第二个tensor-core阶段，则可能构成U），而读取T输出的算术、归一化、激活、规约或布局转换构成C。这产生了一个微小的迭代内DAG，边为$T_j \rightarrow C_j$，如果存在，则为$C_j \rightarrow U_j$。一个attention【【38，Attention is all you need，2017，Proceedings of the 31st International Conference on Neural Information Processing Systems】】的例子是，$QK^T$构成T，softmax操作构成C，第二个GEMM PV构成U。最后，一个aref-use检查会标记任何读取由不同warp组产生的数据的阶段；只有对于那些被标记的阶段，才会发出包装器MAYBEAREFGET(·)和MAYBEAREFCONSUMED(·)。

E. aref降级（aref Lowering）

aref到底层指令的转换。aref的降级过程将高级的异步引用转换为在GPU上直接执行的显式同步和内存传输指令。在IR层面，create_aref声明一个缓冲区并分配必要的mbarriers，作为通信子图的根。编译器随后将put、get和consumed操作重写为具体指令：put扩展为一个由empty/full屏障保护的异步TMA加载，get变为对相应full屏障的阻塞等待，而consumed通过到达empty屏障来信号缓冲区的可用性。这种模式匹配确保一旦整个子图被重写，抽象的aref节点可以被安全地移除，只留下可执行的底层同步和数据移动指令。

流水线下的正确性保证。这一转换的一个关键方面是在流水线下的正确性，其中多个tile在飞行中被加载、计算和消费。为了避免循环等待导致的死锁，降级过程引入了一种奇偶校验机制：每个操作在由迭代奇偶性索引的两组屏障之间交替。当奇偶性切换时，如果数据已经被生产，消费者可以跳过等待，而生产者可以重用缓冲区槽而不会覆盖仍在使用中的值。这种机制不仅保证了活性（liveness），还实现了高效的多缓冲，允许通信和计算在多个阶段中重叠。通过这种方式，aref降级提供了一个从异步数据流的抽象、声明式视图到高性能GPU执行所需的精确TMA和mbarrier指令的系统性桥梁。

IV. 进一步优化

为了在warp专业化场景中进一步提高Tensor Core的利用率，Tawa还集成了额外的优化，这些优化被纳入编译流程，确保程序能够自动受益而无需手动重构。

A. 协作式计算Warp组（Cooperative Compute Warp Groups）

协作Warp组解决寄存器预算限制。warp专业化执行的一个关键限制是每个warp组可用的寄存器预算有限，这限制了最大tile大小并可能限制整体计算强度。协作式warp组优化通过允许多个warp组协同计算同一个tile来缓解这一瓶颈。如图7a所示，两个消费者warp组（WG1和WG2）共同消费由一个TMA warp组（WG0）生产的数据。通过汇集它们的寄存器，协作的warp可以形成更大的tile，从而增加算术强度，改善数据重用，并减少内存流量。从概念上讲，这种优化将一种标准warp并行的形式（即多个warp在同一个tile的不同片段上执行相同的计算）重新引入到warp专业化的设置中。Tawa通过通知后端代码生成器协作映射关系来透明地处理这一点，以便线程索引被一致地分配。重要的是，aref抽象无需任何更改，因为通信语义保持不变。


图7：(a) 协作式warp组；(b) 持久化内核。

B. 持久化内核（Persistent Kernels）

持久化内核减少启动开销。GPU流水线中另一个主要的低效来源是重复的内核启动，每次都会引入显著的CTA启动开销。持久化内核通过只启动与SM数量一样多的CTA并将它们在整个执行期间保持驻留来减少这种开销。如图7b所示，Tawa将执行模型转换为一个软件管理的流水线，其中驻留的CTA负责迭代处理tile。这使得TMA驱动的加载、Tensor Core计算和CUDA Core结尾部分能够在迭代间重叠，消除了重新启动的延迟并减少了空闲时间。Tawa在每个结尾部分之前自动插入同步以保持正确性。这种优化在具有大K维度的深流水线中尤其有效，因为在这些情况下，累积的启动成本否则会占主导地位。

A4 实验环境

  • 软件配置:
    • Tawa基于Triton（commit: 0c7edf）构建，增加了约4000行C++代码和一个封装aref操作的MLIR方言。
    • 对比基线包括：商业闭源库cuBLAS（v12.7），NVIDIA开源库CUTLASS（v4.0.0），同版本的Triton，以及两个学术框架ThunderKittens（6c27e2）和TileLang（bcfc83）。
    • 所有基线都针对一组固定的tile尺寸（{64, 128, 256}的组合）进行了手动调优。Triton基线启用了官方软件流水线。Tawa的aref大小和MMA流水线深度也经过手动选择以最大化性能。
  • 硬件配置:
    • GPU: 一台NVIDIA H100 SXM5 GPU，配备80GB HBM3e内存。
    • 软件栈: CUDA 12.7。
  • 数据集/工作负载:
    • 实验通过25次预热运行和1000次测量运行来计算平均执行时间。
    • GEMM及其变体: 评估FP16和FP8精度的通用矩阵乘法（GEMM）。
    • Attention核: 评估具有不同序列长度的Attention核，以反映真实的LLM工作负载。

A4 实验结果

B. 矩阵乘法（Matrix Multiplication）

  • 实验内容: 在M=N=8192的矩阵尺寸下，对内层维度K从256到16384进行扫描，测试FP16和FP8精度的GEMM吞吐量。
  • 实验结果: 如图8所示，在两种精度下，Tawa都达到了与高度优化的cuBLAS相当的性能，并在大多数形状上优于其他通用框架。Tawa在FP16和FP8下都达到了高达79%的硬件利用率。
    • FP16: Tawa相比cuBLAS、Triton、TileLang和ThunderKittens平均分别取得了1.01×、1.13×、1.23×和1.09×的加速。对Triton的提升显示了重叠TMA驱动的数据移动与WGMMA执行的优势。
    • FP8: Tawa的优势更明显，尤其是在小K值时。相比cuBLAS、Triton、TileLang和ThunderKittens，平均加速比分别为1.06×、1.02×、2.40×和1.24×。TileLang在FP8 WGMMA的布局管理上遇到困难，性能较差。
  • 分析结论: Tawa的aref抽象和自动warp专业化提供了一个强大的编排机制，能够有效地在不同精度下重叠TMA加载和WGMMA计算。


图8：不同框架的GEMM性能结果。M = N = 8192，K值变化。

C. GEMM变体

  • 实验内容: 评估了两种在Triton中实现的GEMM变体：批处理GEMM（Batched GEMM，单个核内执行多个相同形状的小GEMM）和分组GEMM（Grouped GEMM，单个核内处理多个不同形状的GEMM）。
  • 实验结果: 如图9所示，Tawa在两种情况下都持续优于Triton基线，实现了高达7%的加速。
  • 分析结论: 性能提升源于Tawa基于aref的划分和自动warp专业化，它允许一个GEMM的数据移动与另一个GEMM的计算重叠，从而减少了空闲时间并更有效地调度异构形状。


图9：FP16批处理GEMM和分组GEMM的结果。

D. 多头注意力（Multi-Head Attention）

  • 实验内容: 在序列长度L从1024到16384，批大小为4，头维度为128的设置下，评估了（非）因果多头注意力（MHA）。
  • 实验结果: 如图10所示，Tawa与包括手工编写的FlashAttention-3（FA3）在内的强基线进行了比较。
    • FP16: Tawa达到了FA3吞吐量的96%，并比Triton稳定高出1.21×。Triton基线类似于FlashAttention-2，未能充分利用Hopper的warp专业化流水线。在长序列下，Tawa比TileLang和ThunderKittens分别快1.13×和1.23×。
    • FP8: 在大L值时，Tawa的优势更明显，达到FA3的89%，比Triton快1.11×。TileLang和ThunderKittens无法执行FP8配置，表明其内核主要为FP16调优。
  • 分析结论: aref提供了一个原则性的抽象，能将通用内核转换为数据流流水线，通过重叠TMA传输、共享内存物化和WGMMA，在保持warp组占用率和稳定MMA吞吐量的同时实现了高性能。这些优势在不同精度和语义下都成立，表明Tawa的自动warp专业化策略具有良好的泛化性。


图10：不同框架的FP16和FP8 MHA性能结果。

E. 超参数选择

  • 实验内容: 研究了aref大小（D）和MMA流水线深度（P）对（非）持久化GEMM内核（K=16384）性能的影响。
  • 实验结果: 如图11所示，在可行区域（D≥P）内，增加D能持续提升性能，因为它允许TMA生产者预取更多数据，平滑延迟。固定D时，将P从1增加到2通常有帮助，但增加到3则会因寄存器压力增大而降低吞吐量。持久化版本始终比非持久化版本快5-10%。
  • 分析结论: 建议使用较大的D和适中的P（1-2）来平衡重叠和资源压力，并使用持久化内核来稳定流水线并利用缓存/重用优势。


图11：不同aref和MMA流水线大小的影响。

F. 消融研究

  • 实验内容: 在最大的GEMM（K=16384）和MHA（L=16384）内核上，逐步添加Tawa的优化技术，研究其效果。
  • 实验结果: 如图12所示：
    • GEMM:
      • 基线Triton（无WS）: 104 TFLOPs/s。
      • +Auto WS (自动Warp专业化): 性能提升3.78倍至393 TFLOPs/s。
      • +Large Tile Size (使用协作Warp组并增大Tile): 性能再提升1.46倍。
      • +Persistent Kernel (持久化内核): 性能再提升10%。
      • +Better Aref Size (调优Aref大小): 达到峰值718 TFLOPs/s，相比基线提升近7倍。
    • MHA:
      • +Auto WS 和 +Cooperative WGs (协作Warp组) 共同将性能提升了2.84倍。
      • 添加粗粒度流水线和调优aref大小后，最终性能达到654 TFLOPs/s。
  • 分析结论: Tawa通过精心结合warp专业化、协作warp组和流水线技术，能够显著提升现代GPU硬件的利用率。


图12：不同优化对FP16内核的影响。

A5 结论与未来工作

结论

本文介绍了Tawa，一个通过引入异步引用（aref）为现代GPU生成高效warp专业化内核的自动化编译器。评估表明，Tawa在一系列基准测试中达到了与高度优化的手写库相当的性能，同时显著减少了编程工作量。这项工作为异步GPU编程的原则性编译器支持迈出了一步，并能启发未来为下一代架构自动生成高性能内核的研究。

未来工作

  1. 扩展通信模式: 当前Tawa主要针对基于双缓冲的生产者-消费者流水线，未来可以支持更高级的模式，如“乒乓内核”（ping-pong kernels）和“多播通信”（multicast communication），以进一步提高效率。乒乓内核允许warp组在迭代间交替角色，动态平衡负载；多播则允许单个生产者向多个消费者分发数据，减少冗余传输。
  2. 支持下一代架构: 计划将Tawa推广到支持如Blackwell等下一代架构。特别是，Blackwell引入了张量内存（tmem）作为新的硬件管理内存层次。有效利用tmem需要扩展aref以模拟多级数据移动，并重新审视跨异构内存资源的调度决策。
  3. 增强可扩展性: 未来的工作负载需要更高的可扩展性，这促使我们支持多个生产者和多个消费者的warp组，以及更复杂的图划分算法，该算法能同时考虑负载均衡、寄存器压力和内存占用。通过这些扩展，Tawa可以演变为一个跨越多代硬件的通用、warp专业化、异步GPU编程编译器框架。

方法细节中的引用汇总

本节汇总了在“方法细节”章节（III. TAWA COMPILER, IV. FURTHER OPTIMIZATIONS）和“背景知识”章节（II. BACKGROUND AND RELATED WORK）中引用的参考文献，并说明了其引用上下文。

  • 引用 [7]: NVIDIA, “Nvidia blackwell architecture technical brief,” 2025.

    • 引用段落: A3.背景知识 - A. 现代NVIDIA GPU架构
    • 原文描述: "These trends illustrate a broader architectural shift toward asynchronous dataflow execution, a direction continued in subsequent architectures such as Blackwell [7]."
    • 引用说明: 该文献被引用以佐证GPU架构正朝着异步数据流执行的方向发展，并且这种趋势在Blackwell架构中得以延续。

  • 引用 [8]: P. Tillet, H. T. Kung, and D. Cox, “Triton: An intermediate language and compiler for tiled neural network computations,” 2019, Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages.

    • 引用段落: A3.背景知识 - B. 现有的GPU编程模型; A2.方法细节 - A. 概述
    • 原文描述: "Triton [8], built on MLIR [24], allows programmers to write efficient kernels in Python while abstracting tiling and memory management." "Since Triton [8] has become the de facto backend for PyTorch [29], [30] and is widely adopted for writing efficient GPU kernels, we select Triton as the frontend interface of our compiler."
    • 引用说明: 多次被引用，作为Tawa的前端和对比基准。它被描述为一个基于MLIR的、允许用Python编写高效GPU内核的语言和编译器，但缺乏对warp专业化的内在抽象。

  • 引用 [9]: NVIDIA, “Cutlass,” 2025.

    • 引用段落: A3.背景知识 - B. 现有的GPU编程模型
    • 原文描述: "NVIDIA’s CUTLASS C++ [9] library provides templated kernels for GEMM and related computations, but leveraging it to write neural network operators requires intricate configuration and hand-tuning."
    • 引用说明: 引用CUTLASS作为一个提供模板化内核的C++库，但指出其使用复杂，需要手动精细配置和调优。

  • 引用 [10], [11], [12]:

    • [10] B. F. Spector, et al., “Thunderkittens: Simple, fast, and adorable kernels,” 2025, The Thirteenth International Conference on Learning Representations.
    • [11] J. Shah, et al., “Flashattention-3: Fast and accurate attention with asynchrony and lowprecision,” 2024, Advances in Neural Information Processing Systems.
    • [12] Z. Ye, et al., “Flashinfer: Efficient and customizable attention engine for LLM inference serving,” 2025, arXiv preprint arXiv:2501.01005.
    • 引用段落: A3.背景知识 - A. 现代NVIDIA GPU架构; A3.背景知识 - B. 现有的GPU编程模型
    • 原文描述: "Together, these mechanisms provide a foundation for warp specialization, where different warps within an SM assume distinct roles in data communication and computation, thereby improving throughput for LLM workloads [10]–[12]."
    • 引用说明: 这些文献被一同引用，以说明warp专业化能够提高LLM工作负载的吞吐量。同时，ThunderKittens和FlashInfer在“现有模型”部分被单独提及，作为简化高性能内核开发的专用库的例子。

  • 引用 [13]: J. Nickolls, et al., “Scalable parallel programming with cuda: Is cuda the parallel programming model that application developers have been waiting for?” 2008, Queue.

    • 引用段落: A3.背景知识 - B. 现有的GPU编程模型
    • 原文描述: "CUDA is the dominant interface for GPU programming [13], but as architectures evolve, programming complexity has increased substantially."
    • 引用说明: 引用CUDA作为GPU编程的主流接口。

  • 引用 [14] - [28]: 这一系列引用在 A3.背景知识 - B. 现有的GPU编程模型 中被提及，用于描述各种现有的GPU编程模型和编译器框架，如TVM, Relax, Fireiron, Graphene, Triton, TileLang等，并指出它们在利用现代GPU（如Hopper）的warp专业化特性方面存在的局限性，从而凸显Tawa的贡献。

  • 引用 [29], [30]:

    • [29] A. Paszke, et al., “Pytorch: An imperative style, high-performance deep learning library,” 2019, Proceedings of the 33rd International Conference on Neural Information Processing Systems.
    • [30] J. Ansel, et al., “Pytorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation,” 2024, ASPLOS’24.
    • 引用段落: A2.方法细节 - A. 概述
    • 原文描述: "Since Triton [8] has become the de facto backend for PyTorch [29], [30]..."
    • 引用说明: 引用PyTorch以说明Triton作为其后端的重要性，这是选择Triton作为Tawa前端的理由之一。

  • 引用 [31]: B. Hagedorn and V. Grover, “It’s about time - temporal abstractions for asynchronous gpu tensor computation,” 2025, In submission.

    • 引用段落: A2.方法细节 - B. 异步引用（Asynchronous Reference）
    • 原文描述: "To better model communication between warps, we introduce asynchronous reference [31], or aref..."
    • 引用说明: 这是aref概念的直接来源引用。

  • 引用 [32] - [37]: 这一系列引用在 A2.方法细节 - B. 异步引用（Asynchronous Reference） 中被提及，用于将aref与OpenCL pipe等其他加速器编程模型中的类似通信机制进行比较，强调aref是专门为GPU warp间通信设计的。

  • 引用 [38]: A. Vaswani, et al., “Attention is all you need,” 2017, NIPS’17.

    • 引用段落: A2.方法细节 - D. 多粒度流水线（Multi-Granularity Pipelining）
    • 原文描述: "An example of attention [38] is that the QKT forms T, the softmax operation forms C, and the second GEMM P V constructs U."
    • 引用说明: 引用经典的Attention论文，以Attention机制为例来说明粗粒度流水线中的T, C, U三个阶段如何划分。