Tawa自动优化GPU任务并行,提升大模型计算效率。
Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References
- 用异步引用抽象自动划分数据流角色
- 在H100上比cuBLAS快1.1倍,比Triton快1.2倍
- 大幅降低编程负担,适合大模型开发者
现代GPU配备专用硬件单元,支持高性能异步数据流执行。但传统的SIMT编程模型与这种任务并行硬件严重不匹配,造成显著的可编程鸿沟。虽然硬件级线程组专化是释放峰值性能的关键,却迫使开发者手动处理复杂的低层通信和软件流水线——这一过程耗时、易错且难以维护。为此,我们提出Tawa,一个自动编译器,能从高层的基于分块的程序自动生成高性能的线程组专化代码。其核心是新型IR抽象‘异步引用(aref)’,可在不暴露底层硬件细节的前提下表达线程组级通信。借助该抽象,Tawa自动将程序划分为生产者-消费者角色,并管理复杂的数据流流水线,免除开发者侵入式内核重写。在NVIDIA H100 GPU上的代表性LLM核函数评估显示,Tawa实现高硬件利用率,相比高度优化的cuBLAS GEMM核最多提速1.1倍;在注意力运算中,较Triton提速1.2倍,且性能媲美手工优化的CUTLASS C++ FlashAttention-3内核,而编程工作量远低于后者。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。