RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

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.

Recent studies [9,10,11] reveal that in many long-context applications, such as document question answering (QA) and in-context learning, a significant portion of the context remains constant and is reused across multiple requests. For instance, when conducting QA on a corpus, these documents are introduced to the LLM as a prefix and then re-utilized to answer different questions (i.e., user queries). As another example, some fixed contexts containing domain-specific knowledge are used to † Work performed during the internship at Microsoft Research.