A Slice and Dice Approach to Accelerate Compound Sparse Attention on GPU

Transformer-based models are vital to various research domains, including NLP, computer vision, and recommendation systems. However, because the attention mechanism with quadratic complexity limits computation and memory footprint in long sequences, numerous sparse attention-based transformer models are proposed to alleviate these problems. To efficiently infer these models on GPUs, prior solutions such as Triton and Sputnik have accelerated sparse attention in the sparse transformers, including optimizations for sampled dense-dense matrix multiplication (SDDMM), sparse softmax (SpSoftmax), and sparse matrix-matrix multiplication (SpMM). Although the existing optimizations have achieved impressive improvement, they have shortcomings, including low data reuse, unnecessary computation, and unnecessary memory accesses in the latest sparse transformer models based on compound sparse attention.We propose Multigrain to accelerate sparse attention by applying compound sparse GPU kernels with multi-stream for compound sparse patterns. First, we take sparse patterns with high spatial locality and execute them on our customized coarse-grained kernels, which exploit data reuse and high-performance tensor core units. Second, we execute the sparse patterns with low spatial locality on the fine-grained kernels to reduce unnecessary computation and memory accesses. Last, we process the coarse-grained and fine-grained kernels in parallel using multi-stream for the SDDMM and SpMM. In the latest sparse transformer models such as Longformer and QDS-transformer, Multigrain achieved 2.07× and 1.55× end-to-end speedup, respectively, over Triton, which only uses the coarse-grained method. We also achieved 2.08× and 1.08× end-to-end speedup, respectively, compared to Sputnik, which only uses the fine-grained kernel.