DIAS: Distance-based Attention Sparsity for Ultra-Long-Sequence Transformer with Tree-like Processing-in-Memory Architecture

Long-context inference has become a central focus in recent self-regressive Transformer research. However, challenges still remain in performing decode stage due to the memory bandwidth bottleneck of attention mechanisms and the substantial memory overhead associated with KV cache. Although attention sparsity has been proposed as a potential solution, conventional sparsity methods that rely on heuristic algorithms often suffer from accuracy degradation when applied to ultra-long sequences. To break through the dilemma between accuracy-performance and bandwidth-capacity, this work proposes DIAS, a distancebased irregular attention sparsity approach with processing-inmemory (PIM) architecture. DIAS employs approximate topK attention (AKAttention) scores through graph-based search to improve inference efficiency while maintaining accuracy. Furthermore, a scalable tree-like PIM (TreePIM) architecture is introduced to achieve both memory capacity and bandwidth improvement by isolating enormous memory access for KV cache into the PIM units. Evaluations on various configurations of DIAS for Longbench with Llama3-405B models with 1 M sequence length show up to 75 times speedup compared with the state-of-the-art LLM accelerator, with accuracy drop of less than $1 \%$. Index Terms-AI and Machine Learning, Architecture & System Design