RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

Fan Yang (Kuaishou- 快手科技) · Kai Zhang (The Ohio State University) · Meng Chen (Fudan University) · Chen Chen (University of Central Florida) · Qi Chen (Johns Hopkins University) · Huiqiang Jiang (Microsoft Research Asia) · Chengruidong Zhang (Microsoft) · Yuqing Yang (Fudan University) · Lili Qiu (Microsoft) · Di Liu (Shanghai Jiao Tong University) · Baotong Lu (Microsoft Research) · Zhenhua Han (Microsoft) · Qianxi Zhang (Microsoft) · Bailu Ding (Microsoft Research)

approximate nearest neighbor searchattention accuracyattention computationattention-aware vector indexdecoding speedgpu memory inference costs inference speedkey-value vectorskv vector indexesnvidia rtx4090out-of-distributionretrievalattentiontraining-free transformer

Transformer-based Large Language Models (LLMs) have become increasingly important. However, scaling LLMs to longer contexts incurs slow inference speed and high GPU memory consumption for caching key-value (KV) vectors. This paper presents RetrievalAttention, a training-free approach to both accelerate the decoding phase and reduce GPU memory consumption by pre-building KV vector indexes for fixed contexts and maintaining them in CPU memory for efficient retrieval. Unlike conventional KV cache methods, RetrievalAttention integrate approximate nearest neighbor search (ANNS) indexes into attention computation. We observe that off-the-shelf ANNS techniques often fail due to the out-of-distribution (OOD) nature of query and key vectors in attention mechanisms. RetrievalAttention overcomes this with an attention-aware vector index. Our evaluation shows RetrievalAttention achieves near full attention accuracy while accessing only 1-3\% of the data, significantly reducing inference costs. Remarkably, RetrievalAttention enables LLMs with 8B parameters to handle 128K tokens on a single NVIDIA RTX4090 (24GB), achieving a decoding speed of 0.107 seconds per token.