Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References

TL;DR

This work tackles the programmability gap between SIMT GPU programming and warp-specialized, asynchronous hardware. It introduces Tawa, an automated compiler that transforms high-level tile-based programs into warp-specialized code using an aref (asynchronous reference) IR to express inter-warp data movement and synchronization. Through task-aware partitioning and multi-granularity pipelining, Tawa automatically coordinates producer-consumer roles and overlaid compute pipelines, lowering to efficient PTX without user code changes. Empirical results on NVIDIA H100 show Tawa delivering strong hardware utilization and speedups up to 1.2x over Triton and matching hand-optimized FlashAttention/CuBLAS kernels, while reducing programming effort significantly.

Abstract

Modern GPUs feature specialized hardware units that enable high-performance, asynchronous dataflow execution. However, the conventional SIMT programming model is fundamentally misaligned with this task-parallel hardware, creating a significant programmability gap. While hardware-level warp specialization is the key to unlocking peak performance, it forces developers to manually orchestrate complex, low-level communication and software pipelines--a process that is labor-intensive, error-prone, and unsustainable. To address this challenge, we present Tawa, an automated compiler that systematically generates high-performance, warp-specialized code from a high-level, tile-based program. Central to our approach is a novel IR abstraction, asynchronous references (aref), which expresses warp-level communication without exposing low-level hardware details. Using this abstraction, Tawa automatically partitions programs into producer-consumer roles and manages the intricate dataflow pipeline, relieving developers of invasive kernel rewriting. Evaluation on NVIDIA H100 GPUs across representative LLM kernels shows that Tawa delivers high hardware utilization, achieving up to 1.1$\times$ speedup over highly optimized cuBLAS GEMM kernels. For attention workloads, Tawa attains 1.2$\times$ speedup over Triton and matches the performance of the hand-optimized CUTLASS C++ FlashAttention-3 kernel with far less programming effort.

Figures (12)

• Figure 1: A simplified NVIDIA H100 SXM5 GPU architecture.
• Figure 2: Example Tawa compilation flow from the Triton frontend to the internal MLIR representation. tl is the shorthand for the Triton language. tt is the Triton MLIR dialect. Code snippets are simplified for demonstration purposes.
• Figure 3: aref abstraction and associated operations.
• Figure 4: Operational semantics of the aref operations. $\sigma$ is the store map from an aref identifier $a$ to the actual structure $\sigma(a) = \langle \textit{buf}, F, E\rangle$, where $F$ is the full mbarrier flag and $E$ is the empty mbarrier flag. Initially $E=1$, $F=0$.
• Figure 5: Task-aware partition for the GEMM kernel in Fig. \ref{['subfig:triton']}. Notice warp specialization is just one way of achieving the overlapping in (c), and the actual latency is not shown to scale.