Attention, Watch Your Progress: Balancing Warp Specialized GPU Pipelines

Abstract

Contemporary GPU kernels for deep learning workloads, particularly attention mechanisms in transformers, employ sophisticated parallel programming techniques to maximize hardware utilization. These techniques create software pipelines that coordinate multiple specialized hardware units (Tensor Cores, Tensor Memory Accelerators, and SIMT cores) through warp-level synchronization. However, current warp scheduling mechanisms are not designed for these new programming paradigms, leading to suboptimal performance due to imbalance and inefficient instruction issue patterns. We propose Progress-Aware Warp Scheduling (PAWS), a novel scheduling mechanism that adapts to modern attention kernels by leveraging pre-existing machine code information to transmit priority information to the warp scheduler. Our approach enables the warp scheduler to identify and prioritize warps executing slower program phases, thereby improving overall pipeline throughput. We demonstrate that compilers can automatically implement this mechanism, providing significant performance improvements for attention kernels on modern GPU architectures. We perform a holistic parameter sweep of 1800 attention implementations representing contemporary and future attention kernels. PAWS consistently performs better than state-of-the-art hardware warp schedulers, delivering 28% speedup when phases are highly skewed and 15% improvement on attention configurations found in QWen3, Llama4, and Grok 1.0. RTL synthesis results indicate low hardware overhead to implement PAWS in silicon.