---
title: 'DuoAttention: Efficient Long-Context Inference'
url: https://www.emergentmind.com/topics/duoattention
type: topic
---

# DuoAttention: Efficient Long-Context Inference

Searching arXiv for DuoAttention and closely related long-context inference work.
DuoAttention is a long-context inference framework for Transformer large language models 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 [2410.10819]. 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 [2601.17367] [2502.14866].

## 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 [2410.10819].

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 [2410.10819]. 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 [2410.10819].

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 [2410.10819].

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 [2410.10819]. 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 \(\Lambda\)-like mask that permits attention only to attention sinks and a fixed recent window [2410.10819]. Large \(\alpha\) indicates a retrieval head; small \(\alpha\) 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 [2410.10819]. 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 \(\alpha_{i,j}\) are optimized. This makes the optimization lightweight because the number of trainable quantities is only \(N\times H\), with \(N\) layers and \(H\) KV heads per layer [2410.10819].

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 [2410.10819]. The synthetic dataset therefore embeds ten randomly generated passkey sequences of length \(s=32\) 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 [2410.10819].

The distillation loss is given as

$$
\mathcal{L}_{\text{distill}} = \frac{1}{N} \sum_{i=1}^{N} \sum_{j=T-l+1}^{T} \left( h_{\text{full}}^{(i)}[j] - h_{\text{mixed}}^{(i)}[j] \right)^2,
$$

and the sparsity-inducing regularizer is

$$
\mathcal{L}_{\text{reg}} = \sum_{i=1}^{L}\sum_{j=1}^{H} |\alpha_{i,j}|.
$$

The total objective is

$$
\mathcal{L} = \mathcal{L}_{\text{distill}} + \lambda \mathcal{L}_{\text{reg}},
$$

with \(\lambda = 0.05\) in the reported experiments [2410.10819].

After optimization, the gates are binarized using a threshold \(\tau\) chosen from a target sparsity quantile. Heads with \(\alpha_{i,j}>\tau\) use full attention; the remainder use streaming attention [2410.10819]. 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 \(Q\), \(K\), and \(V\) 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 [2410.10819].

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 \(Q,K,V\), 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 [2410.10819]. The streaming-head mask is therefore \(\Lambda\)-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 \(K\). Retrieval heads process chunks causally over the full prefix as usual. Streaming heads compute chunkwise \(K,V\), 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

$$
O(L^2)
$$

to

$$
O(LK),
$$

and prefilling memory improves from

$$
O(L)
$$

to

$$
O(K),
$$

where \(L\) is total sequence length and \(K\) is chunk size [2410.10819].

The implementation uses PyTorch, FlashInfer kernels for RoPE and RMSNorm, and FlashAttention-2 for prefilling [2410.10819]. 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 [2410.10819].

## 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 [2410.10819]. 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 [2410.10819].

The headline efficiency numbers reported in the abstract are:

| Setting | Best reported gain |
|---|---:|
| MHA memory reduction | up to \(2.55\times\) |
| GQA memory reduction | up to \(1.67\times\) |
| MHA decoding speedup | up to \(2.18\times\) |
| GQA decoding speedup | up to \(1.50\times\) |
| MHA prefilling acceleration | up to \(1.73\times\) |
| GQA prefilling acceleration | up to \(1.63\times\) |

These gains are reported with minimal accuracy loss relative to full attention [2410.10819].

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 \(6.4\times\) capacity increase over naive BF16 full-attention deployment [2410.10819].

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**) [2410.10819]. 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**) [2410.10819]. 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** [2410.10819].

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 [2410.10819]. 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 [2601.17367]. 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 [2601.17367]. 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 [2502.14866]. 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 [2502.14866]. 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 [2602.18022]. 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 [1803.09937]. In current LLM literature, “DuoAttention” refers specifically to the retrieval-head/streaming-head framework of [2410.10819].

## 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 [2410.10819]. 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 [2410.10819]. Third, GQA reduces the granularity of head selection because multiple query heads share one KV head, which constrains achievable savings relative to MHA [2410.10819]. Fourth, the method remains approximate whenever any head is assigned streaming attention; the paper does not provide a formal theorem guaranteeing preservation of accuracy [2410.10819].

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 [2410.10819]. 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 [2410.10819].

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 [2410.10819]. 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 [2601.17367] [2502.14866].

Source: https://www.emergentmind.com/topics/duoattention