Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References
By: Hongzheng Chen , Bin Fan , Alexander Collins and more
Potential Business Impact:
Makes computer graphics run much faster.
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.
Similar Papers
Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References
Machine Learning (CS)
Makes computer graphics run much faster.
Optimal Software Pipelining and Warp Specialization for Tensor Core GPUs
Programming Languages
Makes computer graphics run much faster.
From Loop Nests to Silicon: Mapping AI Workloads onto AMD NPUs with MLIR-AIR
Computation and Language
Makes computers use special chips much faster.