Post-Training Sparse Attention with Double Sparsity

Shuo Yang,Ying Sheng,Joseph E. Gonzalez,Ion Stoica,Lianmin Zheng

2024-08-19

Abstract:The inference process for large language models is slow and memory-intensive, with one of the most critical bottlenecks being excessive Key-Value (KV) cache accesses. This paper introduces "Double Sparsity," a novel post-training sparse attention technique designed to alleviate this bottleneck by reducing KV cache access. Double Sparsity combines token sparsity, which focuses on utilizing only the important tokens for computing self-attention, with channel sparsity, an approach that uses important feature channels for identifying important tokens. Our key insight is that the pattern of channel sparsity is relatively static, allowing us to use offline calibration to make it efficient at runtime, thereby enabling accurate and efficient identification of important tokens. Moreover, this method can be combined with offloading to achieve significant memory usage reduction. Experimental results demonstrate that Double Sparsity can achieve $\frac{1}{16}$ token and channel sparsity with minimal impact on accuracy across various tasks, including wiki-2 perplexity, key-value retrieval, and long context benchmarks with models including Llama-2-7B, Llama-2-70B, and Mixtral-8x7B. It brings up to a 14.1$\times$ acceleration in attention operations and a 1.9$\times$ improvement in end-to-end inference on GPUs. With offloading, it achieves a decoding speed acceleration of 16.3$\times$ compared to state-of-the-art solutions at a sequence length of 256K. Our code is publicly available at <a class="link-external link-https" href="https://github.com/andy-yang-1/DoubleSparse" rel="external noopener nofollow">this https URL</a>.

Machine Learning,Artificial Intelligence,Computation and Language

What problem does this paper attempt to address?

This paper attempts to solve the problems of slow speed and high memory consumption in the reasoning process of large - scale language models (LLMs), especially targeting the key bottleneck of excessive Key - Value (KV) cache access in the self - attention layer. The paper introduces a new post - training sparse attention technique named "Double Sparsity", aiming to alleviate this bottleneck by reducing KV cache access. Double Sparsity combines two sparsity methods: token sparsity, that is, only using important tokens for self - attention calculation; and channel sparsity, that is, using important feature channels to identify important tokens. The key insight of the paper is that the pattern of channel sparsity is relatively static and can be made efficient at runtime through offline calibration, thus achieving accurate and efficient identification of important tokens. In addition, this method can also be combined with offloading techniques to significantly reduce memory usage. Experimental results show that Double Sparsity can achieve 1/16 token and channel sparsity in various tasks with minimal impact on accuracy, and achieve up to 14.1 - fold acceleration of attention operations and 1.9 - fold overall inference acceleration on the GPU. When the offloading function is enabled, in the case of a sequence length of 256K, the decoding speed is 16.3 times faster than the state - of - the - art solutions.