Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References
2510.14719v1
cs.LG, cs.AR, cs.PL
2025-10-18
Авторы:
Hongzheng Chen, Bin Fan, Alexander Collins, Bastian Hagedorn, Evghenii Gaburov, Masahiro Masuda, Matthew Brookhart, Chris Sullivan, Jason Knight, Zhiru Zhang, Vinod Grover
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.
Ссылки и действия
Дополнительные ресурсы: