DuoAttention: Efficient Long-Context Inference

• DuoAttention is a long-context inference framework that partitions attention heads into retrieval heads for full historical context and streaming heads for recent tokens.
• The method leverages a synthetic retrieval task to optimize head classification, reducing KV-cache memory usage and decoding latency during both prefilling and inference.
• Empirical results show up to 2.55× memory reduction and 2.18× decoding speedup across multiple LLMs, with minimal loss in accuracy on long-context benchmarks.

Searching arXiv for DuoAttention and closely related long-context inference work. DuoAttention is a long-context inference framework for Transformer LLMs that exploits a head-level asymmetry in attention behavior: only a subset of heads, termed Retrieval Heads, are critical for accessing information arbitrarily far back in the context, whereas the remaining Streaming Heads primarily attend to recent tokens and attention sinks. The method therefore assigns a full KV cache only to retrieval heads and a constant-length sink-plus-recent KV cache to streaming heads, reducing KV-cache memory and attention latency in both decoding and prefilling while preserving long-context capability more effectively than token-pruning baselines (Xiao et al., 2024). In the long-context inference literature, DuoAttention also serves as a reference point for later hybrid sparse/full head-allocation methods, including systems that make the sparse/full split dynamic or embed it in broader serving stacks (Tang et al., 24 Jan 2026, Yang et al., 20 Feb 2025).

1. Conceptual basis and problem formulation

Long-context deployment of autoregressive LLMs is constrained by the KV cache. During decoding, all past keys and values are typically cached, so memory grows linearly with context length, and each new token attends over the entire cache. During prefilling, standard causal attention over a prompt of length $L$ incurs $O(L^2)$ work. The DuoAttention formulation addresses this bottleneck by asking not which tokens should be retained globally, but which heads actually require access to the full history (Xiao et al., 2024).

The central observation is that attention heads exhibit functional specialization. Retrieval Heads are heads whose outputs change significantly when their access is restricted to recent tokens and attention sinks; these heads are therefore essential for long-range information retrieval. Streaming Heads, by contrast, mostly attend to recent context and attention sinks and can be served with a constant-size cache without materially affecting model outputs (Xiao et al., 2024). The distinction is explicitly output-centric rather than score-centric: a retrieval head is defined by the effect of replacing its full attention with streaming-style attention on the final model behavior, not merely by visually long-range attention patterns (Xiao et al., 2024).

This yields a heterogeneous inference policy within each layer. Retrieval heads preserve a standard growing KV cache and standard causal attention. Streaming heads retain only attention sinks at the beginning of the sequence and a fixed recent sliding window, so their cache size is constant with respect to total context length. A plausible implication is that DuoAttention reframes long-context efficiency as a head-allocation problem rather than a token-selection problem, which distinguishes it from many earlier KV-pruning approaches.

2. Attention decomposition and cache policies

DuoAttention applies to both multi-head attention (MHA) and grouped-query attention (GQA). In MHA, each query head has its own KV head; in GQA, multiple query heads share one KV head, so the method operates at the KV-head or group level. Because GQA uses coarser grouping, the achievable compression ratio is smaller than in MHA (Xiao et al., 2024).

The paper writes standard causal self-attention for a head as

$$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$

where $M$ is the causal mask (Xiao et al., 2024). For head identification, DuoAttention introduces a convex mixture between full attention and streaming attention. For KV head $j$ in layer $i$,

$$attn_{i,j} = \alpha_{i,j} \cdot full\_attn + (1 - \alpha_{i,j}) \cdot streaming\_attn,$$

with $\alpha_{i,j}\in[0,1]$, where $full\_attn$ uses the standard causal mask and $streaming\_attn$ uses a $O(L^2)$0-like mask that permits attention only to attention sinks and a fixed recent window (Xiao et al., 2024). Large $O(L^2)$1 indicates a retrieval head; small $O(L^2)$2 indicates a streaming head.

The deployment policy follows directly from this decomposition. Retrieval heads keep the full history in the KV cache. Streaming heads keep only sink tokens and recent tokens, yielding a constant-length cache. This structure reduces the dominant long-context KV footprint roughly in proportion to the retained retrieval-head fraction, ignoring the constant streaming cache size (Xiao et al., 2024). The paper explicitly characterizes the method as a head-wise heterogeneous attention policy for inference.

3. Retrieval-head identification

DuoAttention has two phases: retrieval-head identification and deployment. In the identification phase, model parameters are frozen and only the gate variables $O(L^2)$3 are optimized. This makes the optimization lightweight because the number of trainable quantities is only $O(L^2)$4, with $O(L^2)$5 layers and $O(L^2)$6 KV heads per layer (Xiao et al., 2024).

The identification objective is based on synthetic retrieval supervision rather than ordinary language modeling. The paper argues that natural language modeling loss is poorly aligned with long-range retrieval because most tokens can be predicted from local context, so it does not cleanly expose which heads are indispensable for distant retrieval (Xiao et al., 2024). The synthetic dataset therefore embeds ten randomly generated passkey sequences of length $O(L^2)$7 at random locations in a long context and asks the model to recall them at the end. Distillation loss is applied only on the final passkey tokens, so the supervision directly targets the retrieval behavior relevant to KV compression (Xiao et al., 2024).

The distillation loss is given as

$O(L^2)$8

and the sparsity-inducing regularizer is

$O(L^2)$9

The total objective is

$$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$0

with $$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$1 in the reported experiments (Xiao et al., 2024).

After optimization, the gates are binarized using a threshold $$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$2 chosen from a target sparsity quantile. Heads with $$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$3 use full attention; the remainder use streaming attention (Xiao et al., 2024). The paper emphasizes that this procedure directly optimizes the end-to-end effect of replacing a head’s full KV cache by a sink-plus-recent cache, rather than relying on attention-score profiling alone.

4. Inference-time architecture and systems design

At deployment, DuoAttention reorders the output channels of the $$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$4, $$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$5, and $$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$6 projection matrices so retrieval heads and streaming heads are contiguous within each layer. This is a systems optimization intended to avoid scatter/gather overhead and make slicing and concatenation efficient (Xiao et al., 2024).

During decoding, each layer maintains two caches. The retrieval-head cache grows with context length exactly as in full attention. The streaming-head cache retains only attention sinks and a recent sliding window, so its size is constant. For each new token, the layer computes $$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$7, splits heads by type, applies full attention for retrieval heads and streaming attention for streaming heads, concatenates the resulting head outputs, and applies the output projection (Xiao et al., 2024). The streaming-head mask is therefore $$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$8-shaped in the sense used by streaming-attention literature: initial sink tokens are retained, the middle is discarded, and the recent window is preserved.

Prefilling is also modified. DuoAttention supports chunked prefilling with FlashAttention-2, dividing the prompt into chunks of size $$\mathrm{Attn}(Q,K,V;M) = \mathrm{softmax}(QK^\top \odot M)\,V,$$9. Retrieval heads process chunks causally over the full prefix as usual. Streaming heads compute chunkwise $M$0, then immediately prune the cache to sink-plus-recent tokens before the next chunk. The paper states that, for streaming heads, prefilling time complexity improves from

$M$1

to

$M$2

and prefilling memory improves from

$M$3

to

$M$4

where $M$5 is total sequence length and $M$6 is chunk size (Xiao et al., 2024).

The implementation uses PyTorch, FlashInfer kernels for RoPE and RMSNorm, and FlashAttention-2 for prefilling (Xiao et al., 2024). The method is also explicitly compatible with quantization, and the paper combines it with QServe-style 8-bit weights and 4-bit KV-cache quantization (Xiao et al., 2024).

5. Empirical results

The paper evaluates DuoAttention on long-context and short-context benchmarks across multiple model families, including Llama-2-7B-chat, Llama-2-7B-32K-Instruct, Llama-3-8B-Instruct, Llama-3-8B-Instruct-Gradient-1048k, Llama-3-70B-Instruct, and Mistral-7B-Instruct-v0.2 (Xiao et al., 2024). Long-context benchmarks are Needle-in-a-Haystack (NIAH) and LongBench; short-context benchmarks are MMLU, MBPP, and MT-Bench. Baselines include H2O, TOVA, StreamingLLM, and FastGen (Xiao et al., 2024).

The headline efficiency numbers reported in the abstract are:

Setting	Best reported gain
MHA memory reduction	up to $M$7
GQA memory reduction	up to $M$8
MHA decoding speedup	up to $M$9
GQA decoding speedup	up to $j$0
MHA prefilling acceleration	up to $j$1
GQA prefilling acceleration	up to $j$2

These gains are reported with minimal accuracy loss relative to full attention (Xiao et al., 2024).

The paper also reports a deployment milestone: combined with quantization, DuoAttention enables Llama-3-8B decoding with 3.30 million context length on a single A100-80G GPU, described as a $j$3 capacity increase over naive BF16 full-attention deployment (Xiao et al., 2024).

On LongBench, the method preserves long-context quality better than pruning baselines at the same KV budget. For Llama-3-8B-Instruct-1048K at 50% budget, the average score is 40.21 for DuoAttention versus 40.08 for full attention, exceeding H2O (35.76), StreamingLLM (32.26), and TOVA (35.55) (Xiao et al., 2024). For Llama-2-7B-32K at 25% budget, DuoAttention obtains 34.49 versus 37.52 for full attention, again well above H2O (26.84), StreamingLLM (27.80), and TOVA (29.78) (Xiao et al., 2024). On short-context evaluation, the paper reports near-lossless behavior; for example, on Llama-3-70B at 50% budget, DuoAttention yields MMLU 79.35% versus 79.38% for full attention, MBPP 47.09% versus 47.85%, and MT-Bench 9.14 versus 8.93 (Xiao et al., 2024).

The retrieval-head ratio is architecture-dependent. The paper’s main deployment settings use roughly 25% retrieval heads for Llama-2-7B MHA and 50% retrieval heads for Llama-3 and Mistral GQA models (Xiao et al., 2024). The layerwise gate visualizations indicate that retrieval heads are not uniformly distributed across layers.

6. Relation to adjacent methods and subsequent developments

DuoAttention occupies a specific position in the long-context efficiency literature: it is a static hybrid-attention method that mixes full-attention and streaming-attention heads inside the same layer. Later work uses it explicitly as a baseline or subsystem.

Elastic Attention describes DuoAttention as a representative method that combines full attention and Streaming Sparse Attention within individual layers by assigning different heads to different modes, but does so with a fixed topology or pre-defined ratio prior to inference (Tang et al., 24 Jan 2026). Elastic Attention keeps the hybrid-head idea but replaces the static assignment with input-conditioned per-head routing, arguing that downstream tasks vary in sparsity sensitivity (Tang et al., 24 Jan 2026). This suggests that DuoAttention’s static split is best understood as a strong fixed-policy baseline rather than a universally optimal allocation.

LServe is even more directly connected. It explicitly states that, inspired by DuoAttention, it converts half of the attention heads into streaming heads and adopts DuoAttention’s optimization-based classification to label heads as retrieval or streaming (Yang et al., 20 Feb 2025). LServe then embeds that split inside a unified block-sparse serving framework, adds dynamic query-centric KV-page selection, and combines static and dynamic sparsity for multiplicative serving speedups (Yang et al., 20 Feb 2025). A plausible implication is that DuoAttention’s lasting significance lies not only in its own standalone gains but also in providing a reusable abstraction—retrieval heads versus streaming heads—for later long-context serving systems.

The term “DuoAttention” should be distinguished from unrelated “dual attention” usage in other areas. For example, “Dual-Channel Attention Guidance” manipulates Key and Value channels jointly inside diffusion-transformer attention for image editing, but it is an inference-time guidance method rather than a head-partitioned long-context LLM architecture (Li et al., 20 Feb 2026). Likewise, “Dual Attention Matching network” in person re-identification uses intra-sequence and inter-sequence attention for feature refinement and alignment, not KV-cache heterogeneity in LLM inference (Si et al., 2018). In current LLM literature, “DuoAttention” refers specifically to the retrieval-head/streaming-head framework of (Xiao et al., 2024).

7. Limitations, caveats, and interpretation

The paper identifies several caveats. First, the method presumes a reasonably clean retrieval/streaming dichotomy among heads; if a model uses many heads in more mixed ways, the binary separation may be less accurate and quality may degrade (Xiao et al., 2024). Second, the offline head classification is calibrated on a synthetic passkey retrieval task. If deployment workloads diverge substantially from that calibration signal, the selected retrieval heads may not be optimal (Xiao et al., 2024). Third, GQA reduces the granularity of head selection because multiple query heads share one KV head, which constrains achievable savings relative to MHA (Xiao et al., 2024). Fourth, the method remains approximate whenever any head is assigned streaming attention; the paper does not provide a formal theorem guaranteeing preservation of accuracy (Xiao et al., 2024).

A common misconception is to view DuoAttention as a token-pruning method. The paper argues against that framing. Token pruning discards context positions globally, which can destroy hidden but necessary retrieval anchors. DuoAttention instead preserves the full history selectively for the heads that need it, while allowing other heads to operate with a cheap sink-plus-recent cache (Xiao et al., 2024). Another misconception is that all heads contribute equally to long-context ability; the method is premised on the opposite claim and offers empirical evidence that long-range retrieval is concentrated in a minority of heads (Xiao et al., 2024).

The practical interpretation is therefore narrow but consequential: DuoAttention is not a general sparse-attention theorem, nor a retrained backbone architecture. It is an inference-time framework that exploits head specialization to reduce KV memory and latency while maintaining long-context behavior much more effectively than naive cache pruning (Xiao et al., 2024). In the subsequent literature, it functions both as a deployable method and as a conceptual template for hybrid full/sparse head allocation in long-context LLM serving (Tang et al., 24 Jan 2026, Yang et al., 20 Feb 2025).