DynaX: Sparse Attention Acceleration with Dynamic X: M Fine-Grained Structured Pruning

Owning to the mechanism of self-attention, Transformers have exhibited incredible performance in a wide range of artificial intelligence tasks. With the growth of sequence length, attention computation with quadratic complexity becomes the bottleneck, and dynamic sparsity is an effective technique to alleviate this problem. However, dynamic attention sparsity for long-sequence tasks suffers from two challenges, i.e., irregular sparse patterns and heavy prediction overhead. To this end, this paper proposes DynaX, an algorithm-hardware co-design framework that accelerates attention computation via dynamic X:M fine-grained structured pruning. Different from traditional N:M pruning, DynaX dynamically selects variable X (rather than a fixed N) important scores from a group via a 2-step pruning method, which results in high sparsity and less prediction memory overhead while maintaining pattern regularity to a certain extent. After that, DynaX performs block scheduling to reorganize score blocks into hardware blocks that can perfectly match the size of the processing element array (PEA), resulting in a higher utilization rate. Experimental results show that DynaX can achieve average sparsity of 89.54% and 91.77% for short-sequence tasks and long-sequence tasks, respectively, with less than 1% accuracy loss. Compared to Sanger and SALO2, DynaX achieves a speedup of 1.99X and 1.50X on the BERT-base model, and an energy efficiency improvement of 5.16X and 4.20X, respectively.