Reproduce “Accelerating Sparse Transformer Inference on GPU” on OpenResearch

About this paper

Wenhao Dai, Haodong Deng, Mengfei Rong, Xinyu Yang, Hongyu Liu, Fangxin Liu, Hailong Yang, Qianwen Cao and 1 more

The STOF framework accelerates sparse Transformer inference on GPUs by integrating unified Multi-Head Attention (MHA) kernels and an adaptive operator fusion module. It achieves average speedups of 1.8x for MHA over state-of-the-art FlexAttention and 1.3-1.4x for end-to-end inference over PyTorch Compile on RTX 4090 and A100 GPUs, while reducing tuning time by 6.7x over other auto-tuning frameworks.

Key insights

• A hierarchical sparse storage format combined with dynamic kernel selection enables efficient processing of diverse sparse MHA patterns by enhancing data locality and computation skipping.
• Adaptive operator fusion, guided by a two-stage hierarchical search and reward-based parameter sampling, effectively navigates the complex optimization space, delivering consistent end-to-end performance gains.
• The synergistic optimization of both MHA kernels and downstream operator fusion, rather than isolated approaches, unlocks greater performance potential for sparse Transformer inference on GPUs.

Summary by alphaXiv.