SampleAttention: Near-Lossless Acceleration of Long Context LLM Inference with Adaptive Structured Sparse Attention (2406.15486v2)

Abstract: LLMs now support extremely long context windows, but the quadratic complexity of vanilla attention results in significantly long Time-to-First-Token (TTFT) latency. Existing approaches to address this complexity require additional pretraining or finetuning, and often sacrifice model accuracy. In this paper, we first provide both theoretical and empirical foundations for near-lossless sparse attention. We find dynamically capturing head-specific sparse patterns at runtime with low overhead is crucial. To address this, we propose SampleAttention, an adaptive structured and near-lossless sparse attention. Leveraging observed significant sparse patterns, SampleAttention attends to a fixed percentage of adjacent tokens to capture local window patterns, and employs a two-stage query-guided key-value filtering approach, which adaptively select a minimum set of key-values with low overhead, to capture column stripe patterns. Comprehensive evaluations show that SampleAttention can seamlessly replace vanilla attention in off-the-shelf LLMs with nearly no accuracy loss, and reduces TTFT by up to $2.42\times$ compared with FlashAttention.

Near-Lossless Acceleration of Long Context LLM Inference with Adaptive Structured Sparse Attention

The paper "Near-Lossless Acceleration of Long Context LLM Inference with Adaptive Structured Sparse Attention" addresses a critical bottleneck in the inference of LLMs: the quadratic complexity of the attention mechanism that leads to long Time-to-First-Token (TTFT) latency. This research presents SampleAttention, an approach characterized by adaptive structured sparse attention. This technique aims to replace vanilla attention dynamically, ensuring minimal accuracy loss while substantially speeding up inference.

Introduction and Motivation

Contemporary LLMs such as ChatGLM-6B and InternLM2-7B have achieved remarkable advances in supporting extensive context windows, extending to millions of tokens. These models are pivotal for applications like document analysis, code copilot functionalities, and prolonged conversations. However, their quadratic attention mechanism imposes a significant latency burden, predominantly influencing TTFT. This paper identifies the necessity of mitigating this computational overhead without compromising the accuracy of model predictions. Existing methods—ranging from static sparse attention to low-rank approximations and memory augmentation—require pretraining or finetuning and often introduce accuracy trade-offs.

Theoretical Foundation and Empirical Observations

The authors establish the theoretical foundation of near-lossless sparse attention by proving that it's feasible to find an attention mask $M$ that yields a sparse attention matrix $\tilde{P}$ very close to the full attention matrix $P$. They introduce the concept of sparsity degree (SD) and cumulative residual attention (CRA) to quantify the effectiveness of sparsity in ensuring model performance. More importantly, they delineate conditions under which sparse attention guarantees near-lossless results regarding model output.

Empirical observations show that attention scores in LLMs for long contexts are inherently sparse, exhibit significant head-specific and content-aware variability, and possess evident local window and column stripe patterns. This implies that an effective sparse attention mechanism must adapt dynamically to content variations and different attention head behaviors.

SampleAttention: Method and Implementation

SampleAttention leverages these insights by introducing an adaptive structured sparse attention mechanism. It consists of two primary components:

1. Tuned Window Size: A fixed percentage of the sequence length is used to determine the window size, which can dynamically adjust to various contexts.
2. Query-Guided KV Filtering: This innovative two-stage process first samples attention scores from a subset of queries and then derives key-value indices from accumulated column scores. By focusing on these indices, SampleAttention can capture critical tokens ensuring a high CRA while minimizing computation.

To achieve practical speedup, SampleAttention also incorporates hardware-efficient kernel implementations, based on optimizing FlashAttention. It reduces I/O operations and computation complexities, fostering greater inference efficiency even for extremely long sequences (up to 1 million tokens).

Experimental Validation

The evaluation across diverse benchmarks—LongBench, BABILong, and the "Needle in a Haystack" task showcases SampleAttention's capacity to sustain model accuracy (within 99% of the benchmark) while reducing inference latency substantially. The approach delivers up to $2.42\times$ reduction in TTFT compared to FlashAttention. The findings highlight that SampleAttention is robust across different payload sizes and remains effective for extensive sequence lengths, thus making it a versatile solution for modern LLM deployments.

Implications and Future Directions

The shift towards adaptive structured sparse attention marks a significant advancement in scalable LLM operations, providing a pathway to deploy LLMs in real-time, low-latency applications without sacrificing performance. The implications for AI applications are vast, enabling improved user experiences in interactive and resource-intensive tasks such as live document analysis and extensive dialogue management systems.

Future research can explore optimizing the hyperparameters used in SampleAttention to further refine the balance between performance and speed. Additionally, incorporating more nuanced sparsity patterns and sampling strategies could unlock even greater efficiency. Auto-tuning mechanisms during runtime can be another avenue to adapt quickly to varying input sizes and patterns, ensuring optimal performance dynamically.

In conclusion, the paper presents a comprehensive approach to mitigating one of the biggest challenges in scaling LLMs, setting the stage for more responsive and efficient AI systems capable of handling ultra-long context windows effectively.