Research map / Architectures and efficiency / Attention and transformers
RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes RetrievalAttention, a training-free method that offloads most KV vectors to CPU-side ANNS indexes built with query-aware (attention-aware) edges to handle query-key out-of-distribution, enabling 128K-token decoding of 8B LLMs on a single RTX4090.
Problem
Long-context LLM decoding is slow and KV cache GPU memory grows linearly with context, while off-the-shelf ANNS indexes perform poorly for attention retrieval due to out-of-distribution query-key vectors.
Approach
Build the ANNS index over key vectors using query-to-key KNN connections from prefill (projected onto keys, RoarGraph-style), so decoding queries can traverse to critical keys while scanning only 1-3% of keys; combine CPU-retrieved partial attention with a static GPU-resident window via exact log-sum-exp rescaling.
Key results
- On RULER at 128K, RetrievalAttention averages 74.70 vs 78.7 full attention for Llama-3-8B-level setups, close to Flat (74.34) at 128K on Llama-3-8B; on ∞-Bench it matches full attention (48.9 vs 50.4 avg) using 100 retrieved keys.
- At 128K on a single RTX4090, decoding takes 0.188 s/token, a 4.9x reduction vs Flat (0.922 s) and 1.98x vs IVF (0.373 s).
- Needs to scan only 1-3% of key vectors for >0.95 recall on Q-to-K search, versus 30-50% for IVF.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use RetrievalAttention instead of Flat exact KNN retrieval | Llama-3-8B decoding at 128K context · Needle-in-a-haystack | per-token generation latency (s): 0.922 -> 0.188 | Helped |
| use RetrievalAttention instead of IVF retrieval | Llama-3-8B decoding at 128K context · Needle-in-a-haystack | per-token generation latency (s): 0.373 -> 0.188 | Helped |
| use RetrievalAttention rather than full attention | Llama-3-8B · ∞-Bench | average accuracy: 50.4 -> 48.9 | Hurt |
| use RetrievalAttention rather than full attention | Llama-3-8B · RULER | average accuracy: 86.54 -> 84.70 | Hurt |
| use PyramidKV dynamic retrieval-budget allocation (ablation) | RetrievalAttention on Llama-3-8B · ∞-Bench, 128K context | average accuracy: 49.9 -> 50.1 | Helped |
| use PyramidKV dynamic retrieval-budget allocation (ablation) | RetrievalAttention on Llama-3-8B · ∞-Bench, 128K context | Retr.KV accuracy: 14.5 -> 16.0 | Helped |
| use PyramidKV dynamic retrieval-budget allocation (ablation) | RetrievalAttention on Llama-3-8B · ∞-Bench, 128K context | En.QA accuracy: 8.7 -> 8.5 | Hurt |
| select critical tokens dynamically for each query rather than reuse the first query's top-1000 tokens (ablation) | Llama-3-8B attention · 100,000-token prompt | average recovery ratio: 71% -> 89% | Helped |
Surprising findings
- Off-the-shelf ANNS indexes (IVF, HNSW) perform poorly on query-to-key search even though they work well for key-to-key search, showing the OOD gap matters.
- Static StreamingLLM-style patterns drop the top-1000 recovery ratio from 89% to 71% when reused across decoding steps, showing critical tokens change per query.
Limitations
- Accuracy on some tasks (e.g., RULER at 128K for Llama-3-8B-level setups, Retr.N/Retr.KV) still shows small drops; on Llama-3 RULER at 128K the 74.7 average is below full attention's 78.7 in the appendix comparison.
- Latency results are single-batch and rely on CPU retrieval over PCIe; prefill index construction overhead is not the focus.
- Evaluation covers only a few Llama-3 and Yi models with custom static pattern size of 640 tokens.
Experimental setup
Datasets: ∞-Bench RULER Needle-in-a-haystack
Benchmarks: ∞-Bench RULER Needle-in-a-haystack
Metrics: accuracy recall@100 per-token decoding latency scanned key fraction
Baselines: Full attention vLLM StreamingLLM SnapKV InfLLM Quest InfiniGen Flat IVF
Builds on: RoarGraph StreamingLLM FlashAttention Faiss
Closest papers by meaning
- DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads (ICLR 2025 · similarity 0.91)
- LightRetriever: A LLM-based Text Retrieval Architecture with Extremely Faster Query Inference (ICLR 2026 · similarity 0.91)
- RazorAttention: Efficient KV Cache Compression Through Retrieval Heads (ICLR 2025 · similarity 0.91)
- RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression (ICML 2025 · similarity 0.91)
- A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts (ICLR 2025 · similarity 0.90)
- Efficient Low Rank Attention for Long-Context Inference in Large Language Models (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.