Research map / Architectures and efficiency / Attention and transformers
DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads
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 DuoAttention, which splits attention heads into full-KV retrieval heads and constant-KV streaming heads, identified by optimizing per-head gate values on synthetic passkey data, to cut long-context decoding and pre-filling cost and memory.
Problem
Long-context LLM inference is limited by KV cache memory growing linearly with context and by attention compute, while existing KV cache pruning methods either hurt long-context ability or give limited efficiency gains.
Approach
Only a fraction of heads (retrieval heads) need full attention; the rest (streaming heads) can attend only to attention sinks and recent tokens. Head roles are identified by optimizing per-head gates that blend full and streaming attention, using a distillation loss on a synthetic passkey dataset plus L1 regularization, rather than by attention-score profiling.
Key results
- Reduces long-context decoding memory by up to 2.55x for MHA and 1.67x for GQA models, with decoding speedup up to 2.18x (MHA) and 1.50x (GQA) on a single A100.
- Accelerates pre-filling by up to 1.73x (MHA) and 1.63x (GQA) with memory reductions up to 2.38x and 1.53x.
- Combined with 8-bit weight and 4-bit KV cache quantization, enables Llama-3-8B to hold 3.30M tokens on a single A100-80G GPU, a 6.4x capacity increase over FP16 full attention.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DuoAttention instead of H2O at the 50% KV budget | Llama-3-8B-Instruct-1048K · LongBench Average | score: 35.76 -> 40.21 | Helped |
| Use DuoAttention instead of FastGen | Llama-3-8B-Instruct-1048K · LongBench Average | score: 32.82 -> 40.01 | Helped |
| Use DuoAttention instead of H2O at the 25% KV budget | Llama-2-7B-Instruct-32K · LongBench Average | score: 26.84 -> 34.49 | Helped |
| Use DuoAttention instead of FastGen | Llama-2-7B-Instruct-32K · LongBench Average | score: 19.01 -> 32.81 | Helped |
| Add DuoAttention to full attention | Llama-2-7B-Instruct-32K · LongBench Average | score: 37.52 -> 34.49 | Hurt |
| Use DuoAttention instead of full attention | Llama-3-8B-Instruct-1048K · LongBench Average | score: 40.08 -> 40.21 | No clear effect |
| Add optimization-based retrieval-head identification rather than attention profiling (ablation) | DuoAttention · Mistral-7B-Instruct-v0.2 passkey retrieval | exact match accuracy | Helped |
| Add synthetic passkey data for retrieval-head identification rather than language-modeling data (ablation) | DuoAttention · Mistral-7B-Instruct-v0.2 passkey retrieval | exact match accuracy | Helped |
| Use both sink and recent tokens during retrieval-head optimization rather than only one token group (ablation) | DuoAttention · Mistral-7B-Instruct-v0.2 passkey retrieval | exact match accuracy | Helped |
| Increase deployment sink and recent-token counts beyond 16 and 64 (ablation) | DuoAttention · Mistral-7B-Instruct-v0.2 passkey retrieval and MMLU | performance | No clear effect |
Surprising findings
- Attention-score profiling (FastGen/RazorAttention-style) fails to identify retrieval heads well compared with the output-deviation optimization.
- Pruning middle tokens of retrieval heads severely damages passkey retrieval, while pruning streaming heads has minimal impact.
Limitations
- FastGen could only be evaluated on shorter contexts (24K-32K) due to quadratic memory, so comparisons at long lengths are limited.
- H2O and TOVA were modified (simulated decoding of last 50 tokens) to run in long-context settings, so baseline numbers may differ from originals.
- Retrieval head ratio is a fixed hyperparameter chosen per model (25% or 50%), and identification requires an optimization run per model.
Experimental setup
Datasets: BookSum synthetic passkey retrieval dataset
Benchmarks: Needle-in-a-Haystack LongBench MMLU MBPP MT-Bench
Metrics: accuracy decoding latency memory usage pre-filling latency speedup
Baselines: H2O TOVA FastGen StreamingLLM RazorAttention
Builds on: StreamingLLM FlashAttention-2 Needle-in-a-Haystack CNN filter pruning
Closest papers by meaning
- A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts (ICLR 2025 · similarity 0.92)
- SCBench: A KV Cache-Centric Analysis of Long-Context Methods (ICLR 2025 · similarity 0.91)
- RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval (NeurIPS 2025 · similarity 0.91)
- CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration (ICML 2025 · similarity 0.91)
- RazorAttention: Efficient KV Cache Compression Through Retrieval Heads (ICLR 2025 · similarity 0.91)
- QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference (ICML 2024 · similarity 0.91)
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.