RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

2024-09-16

arXiv.org

AI Summary

To address the high computational cost of attention computation and excessive KV cache memory consumption in long-context large language model inference, this paper proposes a training-free retrieval-augmented attention mechanism. The method integrates CPU-based approximate nearest neighbor search (ANNS) indexing, attention-distribution-aware adaptive retrieval, offline indexing of KV vectors, and online retrieval. Its core innovation is the first attention-aware vector retrieval algorithm, which mitigates query-key distribution shift (out-of-distribution, OOD) and enables dynamic sparse attention approximation. Evaluated on a 128K-context, 8B-parameter model, the approach achieves 0.188 seconds per token generation latency on a single RTX 4090 GPU (24 GB), reduces GPU memory usage by 97%, and maintains accuracy nearly equivalent to the full-attention baseline.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Transformer-based Large Language Models (LLMs) have become increasingly important. However, due to the quadratic time complexity of attention computation, scaling LLMs to longer contexts incurs extremely slow inference speed and high GPU memory consumption for caching key-value (KV) vectors. This paper proposes RetrievalAttention, a training-free approach to both accelerate attention computation and reduce GPU memory consumption. By leveraging the dynamic sparsity of attention mechanism, RetrievalAttention proposes to build approximate nearest neighbor search (ANNS) indexes for KV vectors in CPU memory and retrieve the most relevant ones through vector search during generation. Unfortunately, we observe that the off-the-shelf ANNS indexes are often ineffective for such retrieval tasks due to the out-of-distribution (OOD) between query vectors and key vectors in the attention mechanism. RetrievalAttention addresses the OOD challenge by designing an attention-aware vector search algorithm that can adapt to the distribution of query vectors. Our evaluation demonstrates that RetrievalAttention achieves near full attention accuracy while only requiring access to 1--3% of the data. This leads to a significant reduction in the inference cost of long-context LLMs, with a much lower GPU memory footprint. In particular, RetrievalAttention only needs a single NVIDIA RTX4090 (24GB) to serve 128K tokens for LLMs with 8B parameters, which is capable of generating one token in 0.188 seconds.

Problem

• Large Language Models
• Long Document Processing
• Computational Efficiency

Innovation

• RetrievalAttention
• Efficient Long Document Processing
• Optimized Memory Usage