SADIMM: Accelerating Sparse Attention Using DIMM-Based Near-Memory Processing

Self-attention mechanism is the performance bottleneck of Transformer based language models. In response, researchers have proposed sparse attention to expedite Transformer execution. However, sparse attention involves massive random access, rendering it as a memory-intensive kernel. Memory-based architectures, such as <i>near-memory processing</i> (NMP), demonstrate notable performance enhancements in memory-intensive applications. Nonetheless, existing NMP-based sparse attention accelerators face suboptimal performance due to hardware and software challenges. On the hardware front, current solutions employ homogeneous logic integration, struggling to support the diverse operations in sparse attention. On the software side, token-based dataflow is commonly adopted, leading to load imbalance after the pruning of weakly connected tokens. To address these challenges, this paper introduces SADIMM, a hardware-software co-designed NMP-based sparse attention accelerator. In hardware, we propose a heterogeneous integration approach to efficiently support various operations within the attention mechanism. This involves employing different logic units for different operations, thereby improving hardware efficiency. In software, we implement a dimension-based dataflow, dividing input sequences by model dimensions. This approach achieves load balancing after the pruning of weakly connected tokens. Compared to NVIDIA RTX A6000 GPU, the experimental results on BERT, BART, and GPT-2 models demonstrate that SADIMM achieves 48<inline-formula><tex-math notation="LaTeX">$\boldsymbol{\times}$</tex-math></inline-formula>, 35<inline-formula><tex-math notation="LaTeX">$\boldsymbol{\times}$</tex-math></inline-formula>, 37<inline-formula><tex-math notation="LaTeX">$\boldsymbol{\times}$</tex-math></inline-formula> speedups and 194<inline-formula><tex-math notation="LaTeX">$\boldsymbol{\times}$</tex-math></inline-formula>, 202<inline-formula><tex-math notation="LaTeX">$\boldsymbol{\times}$</tex-math></inline-formula>, 191<inline-formula><tex-math notation="LaTeX">$\boldsymbol{\times}$</tex-math></inline-formula> energy efficiency improvement, respectively.