PaperlayerView exact revision on arXiv

Source-linked AI summary

RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chengruidong Zhang, Bailu Ding, Kai Zhang, Chen Chen, Fan Yang, Yuqing Yang, Lili Qiu

arXiv:2409.10516v3cs.LGcs.CL

TL;DR

Long-context inference is costly in attention latency and KV-cache memory, and query–key distribution shift makes standard ANNS inefficient for attention retrieval. RetrievalAttention uses query-aware vector search with CPU–GPU execution. At 128K on RTX4090, it reports lower decoding latency than exact KNN and traditional ANNS while matching full-attention accuracy.

  • Problem

    Long-context inference faces high attention latency and KV-cache memory needs, while query–key distribution shift makes conventional ANNS require scanning 30–50% of keys for acceptable accuracy.

    Source: §1 · pp. 1–2

  • Method

    RetrievalAttention uses query-distribution-aware ANNS to retrieve critical keys for dynamic sparse attention, offloads most KV vectors to CPU, and merges CPU and GPU attention results.

    Source: §1 · p. 2

  • Results

    4.9× and 1.98× decoding-latency reductions over exact KNN and traditional ANNS, respectively, while maintaining full-attention accuracy at 128K on RTX4090.

    Source: §1 · p. 3

  • Takeaways & Limitations

    RetrievalAttention supports 8B-level LLMs with 128K tokens on a single 24GB RTX4090 with acceptable latency and without compromising model accuracy.

    Source: §6 · p. 11

    The work focuses on token generation and assumes long-context prompt prefill has already been completed.

    Source: §3 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

Long-context inference strains KV-cache memory and grows in latency, while query–key distribution shift makes ordinary ANNS ineffective. RetrievalAttention adapts vector search to query distributions and reports near-full-attention accuracy with faster decoding at 128K.

1 INTRODUCTION

  • KV caching reduces redundant computation but requires substantial GPU memory and leaves latency increasing linearly with context size.The Llama-3-8B example requires approximately 125GB per million tokens.Source: §1 · p. 1
  • Dynamic sparsity means each query interacts strongly with a limited, query-dependent subset of key and value vectors, but prior static or heuristic selection can reduce accuracy (Li et al., 2024; Xiao et al., 2024a).The paper presents dynamic sparsity as an opportunity for more efficient attention computation.Source: §1 · p. 1
  • 4.9× and 1.98× decoding-latency reductions over exact KNN and traditional ANNS, respectively, while matching full-attention accuracy at 128K on RTX4090.Evaluation covers three long-context LLMs on RTX4090 and A100 GPUs across long-context benchmarks.Source: Figure 1 · p. 2, §1 · pp. 2–3
  • Inner-product ANNS over key vectors can retrieve high-scoring vectors in sub-linear time and run on CPU, supporting GPU-memory savings.The paper says this search aligns with attention's similarity measurement and can improve accuracy over static or heuristic methods.Source: §1 · p. 2
  • Query and key vectors are out of distribution because attention uses different projection weights, and conventional ANNS needs to scan 30–50% of keys for acceptable accuracy.The paper reports query vectors deviate more than 10× farther from keys than in-distribution query vectors do.Source: §1 · p. 2
  • RetrievalAttention tailors vector indexing to query distributions, traverses 1–3% of key vectors, and offloads most KV vectors to CPU while merging retrieved CPU attention with partial GPU attention.A small number of KV vectors remain on GPU following static patterns.Source: §1 · p. 2

2 BACKGROUND AND MOTIVATION

Long-context inference is costly because attention scales quadratically and KV caches demand substantial GPU memory. Although attention is sparse, critical tokens vary with each query, and query-key distribution mismatch makes conventional vector indexes inefficient for finding them.

2.1 LLM AND ATTENTION OPERATION

  • Each decoding step compares the current query with preceding keys, normalizes scores with Softmax, then uses them to weight values into the output.KV caching avoids recomputing past key-value states, reducing decoding attention complexity to O(n).Source: §2.1 · p. 3

2.2 EXPENSIVE LONG-CONTEXT SERVING

  • At one million tokens, generating each token takes 1,765 seconds without KV caching, while the KV cache requires 125 GB of GPU memory.Attention accounts for over 96% of latency in the no-cache setting.Source: §2.2 · p. 3

2.3 DYNAMIC AND SPARSE ATTENTION

  • Top-1000 query-specific tokens recover an average 89% of attention scores, but reusing the first token’s selection drops recovery to 71%.Critical keys change with the current query, so token importance must be selected dynamically.Source: §2.3 · pp. 3–4, Figure 2 · p. 4

2.4 CHALLENGES OF OFF-THE-SHELF VECTOR SEARCH

  • IVF needs ~30–50% scanning for recall above 0.95, while HNSW reaches a local optimum; K-to-K search reaches high recall with 1–5% scanning, unlike Q-to-K retrieval.The query-key distribution gap explains why key-only index construction does not align with attention retrieval.Source: Figure 3 · p. 4, §2.4 · p. 5
  • Attention sparsity makes retrieval of only high-inner-product key vectors a potential route to lower-cost attention, but conventional indexes scan too much for effective Q-to-K search.ANNS searches based on inner product align with identifying critical keys, yet query-key distribution mismatch undermines off-the-shelf indexes.Source: §2.4 · pp. 4–5

3 RETRIEVALATTENTION DESIGN

RetrievalAttention approximates attention by dynamically retrieving important key-value vectors from CPU-indexed context while retaining consistently important vectors on the GPU. Its query-aware index scans only a small fraction of keys, and CPU-GPU co-execution combines partial attention outputs.

3 RETRIEVALATTENTION DESIGN

  • The work focuses on accelerating token generation and assumes long-context prompt prefill has already been completed.This setting is supported by existing services that cache context or separate prefill from decoding.Source: §3 · p. 5

3.1 APPROXIMATED ATTENTION

  • RetrievalAttention approximates full attention using high-score KV vectors, with most vectors offloaded to CPU indexes and relevant ones retrieved dynamically.Vectors with attention scores above ε define the selected token subset; persistent and retrieved vectors are used for the approximation.Source: §3 · pp. 5–6, §3.1 · p. 6

3.2 ATTENTION-AWARE VECTOR SEARCH

  • Index construction links prefill query vectors to their exact nearest keys, then projects those links onto keys to streamline later search.A decoding query searches nearby query vectors and uses their mapped key neighbors; projecting links removes the need to store query vectors.Source: §3.2 · p. 6
  • Scanning 1–3% of key vectors achieves high recall and reduces index search latency by 74% compared with IVF indexes.The index models proximity from the query vectors’ viewpoint rather than relying only on key-key closeness.Source: §3.2 · p. 6

3.3 CPU-GPU CO-EXECUTION

  • Prefill attention patterns identify KV vectors kept on GPU, while other vectors are indexed in CPU memory for dynamic retrieval.The implementation uses fixed initial tokens and a sliding window, computes CPU and GPU attention separately, and combines their outputs.Source: §3.3 · pp. 6–7

4 EVALUATION

The evaluation compares RetrievalAttention with full attention and training-free baselines across long-context accuracy and latency tasks. It also tests the attention-aware search index’s recall-versus-scan efficiency.

4.1 EXPERIMENTAL SETUP

  • The evaluation compares RetrievalAttention with full attention and training-free baselines on three long-context LLMs—Llama-3-8B-Instruct-262k, Yi-6B-200K, and Yi-9B-200K—across ∞-Bench, RULER, and Needle-in-a-haystack.Experiments use one NVIDIA RTX4090 GPU with 24GB memory. By default, indexing-based methods retrieve the top 100 nearest key vectors.Source: §4.1 · p. 7

4.2 ACCURACY ON LONG CONTEXT TASKS

  • On ∞-Bench, RetrievalAttention achieves accuracy comparable to full attention and the best KV Retrieval accuracy; on RULER, it remains nearly equal to full attention across context lengths.Retrieving top-2000 keys brings KV Retrieval nearly to full-attention accuracy.Source: Table 2 · p. 8, §4.2 · p. 8, Table 3 · p. 9
  • At the same task accuracy, Flat and IVF scan 100% and 30% of past key vectors, respectively, while RetrievalAttention scans only 1–3%, resulting in much lower decoding latency.Source: §4.2 · p. 8
  • From 4K to 128K contexts, RetrievalAttention focuses on information across positions, while StreamingLLM struggles when critical information lies outside its static patterns.The comparison is shown in Figure 5.Source: §4.2 · p. 8, Figure 5 · p. 9

4.3 DECODING LATENCY

  • At 128K, RetrievalAttention reduces decoding latency 4.9× versus Flat and 1.98× versus IVF; its latency breakdown also reports 91% and 74% lower index-search latency, respectively.It spends 34.0% of end-to-end time on vector search, versus 86.6% for Flat and 67.0% for IVF.Source: §4.3 · p. 9, Table 4 · p. 9, Table 5 · p. 10

4.4 INDEX RECALL VS. SCANNING VECTORS

  • For query-to-key search, RetrievalAttention reaches recall above 0.95 after scanning 1–3% of keys; all indexes perform well for in-distribution key-to-key search.The analysis attributes RetrievalAttention’s query-to-key efficiency to mitigating the query–key distribution gap.Source: §4.4 · pp. 9–10, Figure 6 · p. 10

5 RELATED WORKS

Related approaches exploit attention sparsity or offload KV caches, but their strategies differ in how they select and retrieve relevant vectors. RetrievalAttention uses ANNS indexes for dynamic token retrieval and targets the query–key distribution gap.

5 RELATED WORKS

  • Attention-sparsity methods compress KV caches, but often lose accuracy because attention’s sparse pattern is dynamic.FlexGen and Lamina offload KV caches but retain costly full-attention computation; Quest and InfLLM retrieve blocks using representative vectors.Source: §5 · p. 10
  • RetrievalAttention indexes KV vectors for high-recall, low-cost retrieval, while concurrent MagicPiG and PQCache do not address attention’s OOD issue and may retrieve about 20% of the cache.SparQ, InfiniGen, and LoKi instead reduce channel dimensions to approximate relevant keys.Source: §5 · p. 10

6 CONCLUSION

The conclusion presents RetrievalAttention as a CPU-offloading, vector-search approach to dynamic sparse attention. It reports speedups over retrieval baselines and a 128K-token deployment on one RTX4090 for 8B-level LLMs without compromising accuracy.

6 CONCLUSION

  • RetrievalAttention offloads most KV vectors to CPU memory and uses vector search for dynamic sparse attention to reduce inference cost.Its search approach is tailored to the different query and key vector distributions.Source: §6 · p. 10
  • At a 128K-token context, RetrievalAttention runs 8B-level LLMs on a single 24GB RTX4090 GPU with acceptable latency and without compromising model accuracy.It achieves 4.9× decoding speedup over exact KNN.Source: §6 · pp. 10–11

A.1 MODEL ARCHITECTURE

The evaluated models support grouped query attention, where multiple query heads share one KV head. Their architecture differs in layer and KV-head counts.

A.1 MODEL ARCHITECTURE

  • All three evaluated models support grouped query attention, with multiple query heads sharing one KV head.Yi-9B has more transformer layers, while Llama-3-8B has more KV heads.Source: Table 6 · p. 15, §A.1 · p. 15

A.2 ADDITIONAL RESULTS ON NEEDLE-IN-A-HAYSTACK

On Needle-in-a-haystack, the appendix compares other methods and reports RetrievalAttention’s results on contexts up to one million tokens. It passes all tests from 250K to 1 million tokens.

A.2 ADDITIONAL RESULTS ON NEEDLE-IN-A-HAYSTACK

  • StreamingLLM finds the answer only within the static pattern; InfLLM declines at longer contexts, while SnapKV, Flat, and IVF perform well on this benchmark.The static pattern is 640 tokens: 128 initial tokens plus a 512-token local window.Source: Figure 7 · p. 15, §A.2 · p. 15

A.3 PERFORMANCE IN THE EXTREMELY LONG-CONTEXT INFERENCE

  • RetrievalAttention passes all test cases at context lengths from 250K to 1 million on Llama-3-8B-1048K.The reported result demonstrates robustness of the attention-aware indexes.Source: Figure 8 · p. 16, §A.3 · p. 16

A.4 DECODING LATENCY ON A100

On A100, RetrievalAttention maintains generation accuracy while reducing decoding latency relative to Flat and IVF, with latency increasing only slightly as context grows. It combines GPU- and CPU-computed partial attention outputs, found through prediction and vector retrieval.

A.4 DECODING LATENCY ON A100

  • RetrievalAttention maintains generation accuracy with lower decoding latency than IVF and Flat on A100.The comparison measures per-token generation latency at 128K context length.Source: Table 7 · p. 16, §A.4 · p. 16
  • As context grows from 100K to 1M, RetrievalAttention’s latency increases 8%; vLLM starts triggering OOM above 200K with its GPU KV cache enabled.Full attention with KV-state recomputation has latency that increases quadratically with context size.Source: Table 8 · p. 16, §A.4 · p. 16

B.1 FORMULA OF COMBINING ATTENTION RESULTS FROM THE CPU AND GPU SIDE

  • The method combines GPU- and CPU-computed partial attention outputs so the result equals attention over the selected KV vectors.It partitions KV vectors into predictable GPU-resident vectors and dynamically retrieved CPU-resident vectors, then rescales and combines their partial outputs.Source: §B.1 · p. 17

B.2 OVERALL EXECUTION FLOW

  • At each token, RetrievalAttention predicts active KV vectors for GPU attention while a vector database retrieves relevant vectors for CPU attention, then combines both outputs.The GPU path uses FlashAttention; the CPU search and GPU computation run in parallel.Source: §B.2 · pp. 17–18

C IMPLEMENTATION

The paper describes implementation optimizations, additional baselines, and dynamic retrieval-budget allocation. RetrievalAttention performs near full-attention accuracy against added baselines, while budget allocation has task-dependent effects.

C IMPLEMENTATION

  • Implementation overlaps KV transfer with GPU prefill, parallelizes CPU searches across heads, and shares KV-vector storage across indexes.The authors also report initial 8-bit quantization results with accuracy equivalent to full precision.Source: C IMPLEMENTATION · p. 18

D ADDITIONAL RELATED WORK

  • Prior sparse-attention approaches may require retraining or retain full GPU KV caches; other inference optimizations are described as orthogonal and combinable with RetrievalAttention.The related-work discussion also notes overhead and accuracy concerns for some dynamic-sparsity approaches.Source: D ADDITIONAL RELATED WORK · p. 18

E ADDITIONAL BASELINES

  • RetrievalAttention performs best and achieves nearly full-attention accuracy across two benchmarks, while InfiniGen and Quest show noticeable accuracy drops on RULER.The comparison is reported for 128K context length.Source: Table 9 · p. 19, E ADDITIONAL BASELINES · p. 19

F DYNAMIC RETRIEVAL BUDGET ALLOCATION

  • PyramidKV allocation slightly improves average accuracy over the original RetrievalAttention, improving Retr.KV but slightly reducing En.QA performance.The original method uses a fixed 2,000-token budget for every head and layer; PyramidKV allocates more tokens to lower layers and fewer to higher ones.Source: F DYNAMIC RETRIEVAL BUDGET ALLOCATION · p. 19, Table 10 · p. 19

G PERFORMANCE ON THE LARGER MODEL

On Llama-3-70B-262k, RetrievalAttention nearly matches Flat’s task accuracy, outperforms Quest by 80%, and decodes 3.5× faster than Flat.

G PERFORMANCE ON THE LARGER MODEL

Loading 2409.10516v3…