AMS: Pushing the Limits of Attention Acceleration with Dynamic Multi-level Structured Sparsity

Transformer models have achieved remarkable contextual comprehension capabilities in recent years, mainly thanks to the self-attention mechanism. As the input sentence becomes longer, the computation of the attention part quadratically increases and becomes the bottleneck. The dynamic sparsity is a prevalent approach to alleviate this problem. However, existing sparse attention accelerator incur substantial redundant computations, reaching the performance limitation. They typically employ single-level sparse structure, supporting multiple or finegrained N:M sparse patterns for acceleration, which fails to accommodate the diverse sparse pattern with low hardware overhead. By reexamining the distribution of sparsity patterns, we observe that more high-level sparsity exists across multi-level and within each level with negligible hardware overhead. With multi-levels, more hierarchical sparsity patterns can be exploited without increasing the tile size. Within each level, we further prune redundant elements in an NM-aware manner that fits sparsity patterns dynamically with global information. Motivated by the observation, this paper proposes AMS, an efficient Attention accelerator with Dynamic Multi-level Structured Sparsity. AMS exploits structured sparsity patterns using dynamic multi-level structure in a small tile size, resulting in both high sparsity and low hardware overhead. In the following formal compute stage, these structured patterns enable a coordinated optimization of both scheduling algorithms and processing elements array architecture, achieving a high hardware utilization rate. Compared to NVIDIA V100 GPU, AMS achieves a 52.3× speedup and 20.7× higher power efficiency. Compared to the state-of-the-art works DynaX, SOFA, Highlight and DFSS, AMS can achieve a speedup of 1.5×, 2.1×, 6.9× and 32.0× and a power efficiency improvement of 1.2×, 2.2×, 2.7× and 12.7×, respectively.