DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads

## Abstract

Abstract Deploying long-context large language models (LLMs) is essential but poses significant computational and memory challenges.
Caching all Key and Value (KV) states across all attention heads consumes substantial memory.
Existing KV cache pruning methods either damage the long-context capabilities of LLMs or offer only limited efficiency improvements.
In this paper, we identify that only a fraction of attention heads, a.k.a, Retrieval Heads , are critical for processing long contexts and require full attention across all tokens.
In contrast, all other heads, which primarily focus on recent tokens and attention sinks–referred to as Streaming Heads –do not require full attention.
Based on this insight, we introduce DuoAttention, a framework that only applies a full KV cache to retrieval heads while using a light-weight, constant-length KV cache for streaming heads, which reduces both LLM’s decoding and pre-filling memory and latency without compromising its long-context abilities.
DuoAttention uses a lightweight, optimization-based algorithm with synthetic data to identify retrieval heads accurately.
Our method significantly reduces long-context inference memory by up to 2.55 × \times × for MHA and 1.67 × \times × for GQA models while speeding up decoding by up to 2.18 × \times × and 1.50 × \times × and accelerating pre-filling by up to 1.73 × \times × and 1.63 × \times × for MHA and GQA models, respectively, with minimal accuracy loss compared to full attention.
Notably, combined with quantization, DuoAttention enables Llama-3-8B decoding with 3.3 million context length on a single A100 GPU.
Code is provided in the link .

## 1 Introduction

Large language models (LLMs) (Touvron et al., ; ; OpenAI, ; Black et al., ) are at the forefront of the AI revolution, powering advanced applications such as multi-round dialogues (Schulman et al., ; Taori et al., ; Chiang et al., ), long document summarization (Goyal & Durrett, ; Zhang et al., ), and tasks involving mixed modalities like visual and video understanding (Liu et al., ; Lin et al., ).
These applications often require processing extensive numbers of contextual tokens; for instance, summarizing the entire Harry Potter series could involve analyzing approximately one million tokens. The challenge intensifies with visual language models (VLMs), where a single 224$\times$224 image corresponds to 256 tokens (Liu et al., ), and a three-minute video at 24 FPS generates around 1.1 million tokens.

A critical issue in deploying LLMs in such applications is the long-context inference problem. The full attention mechanism demands that all tokens attend to every previous token for accurate representation, resulting in linearly increasing decoding and quadratically increasing pre-filling latency as the sequence length grows. Additionally, the Key-Value (KV) Cache technique, which stores keys and values from all preceding tokens, causes memory usage to scale linearly with context length. As sequences lengthen, memory is increasingly consumed by the KV cache, placing a significant computational burden on the attention mechanism. For instance, in the Llama-3-8B (Dubey et al., ) model architecture, serving with FP16 KV cache for 1 million tokens would require at least 137 GB of memory—exceeding the capacity of a single 80GB GPU. Additionally, the latencies of pre-filling and decoding with such large contexts are significant, posing substantial challenges to the effective use of LLMs in long-context scenarios.

Despite numerous efforts to overcome the challenges of attention mechanisms in long-context inference, significant computational and memory issues persist. Architectural modifications, such as Grouped-Query Attention (GQA)(Ainslie et al., ), require model pre-training and fail to reduce computational costs. Linear Attention methods (Gu & Dao, ; Poli et al., ), while less demanding in terms of computation and memory, often underperform in long-context scenarios compared to Transformer models.
Approximative attention methods, such as H2O (Zhang et al., ), StreamingLLM (Xiao et al., ), TOVA (Oren et al., ), and FastGen (Ge et al., ), often compromise accuracy in long-context applications and are incompatible with essential KV cache optimization techniques like GQA.
KV cache quantization (Liu et al., ; Hooper et al., ), although useful, does not reduce the computation time of the attention mechanism. System-level optimizations, including FlashAttention (Dao et al., ; Dao, ), FlashDecoding (Hong et al., ), and PagedAttention (Kwon et al., ), while effective, do not reduce the KV cache size and still require significant computation for extended contexts. These limitations emphasize the need for further advancements to deploy models that handle million-level context lengths.

Figure: Figure 1: Visualization of attention maps in the Llama-2-7B model for the sentence "The best fruit is orange. What is the best fruit? Orange." shows the distinct roles of *retrieval heads* (e.g., Layer 15, Head 12) and *streaming heads* (e.g., Layer 10, Head 1). On the left, retrieval heads capture contextually relevant tokens such as "best," "fruit," and "orange," which are crucial for processing long-context information and, therefore, require a full KV cache. In the middle, streaming heads primarily focus on initial and recent tokens without emphasizing past contextual relevance. On the right, the impact of limiting attention to the sink and recent tokens on long-context passkey retrieval accuracy is shown: modifying retrieval heads severely damages performance, while constraining streaming heads has minimal impacts.
Refer to caption: x1.png

In this paper, we introduce a key observation that attention heads in LLMs can be categorized into two distinct types: Retrieval Heads (Wu et al., ) and Streaming Heads, as shown in Figure 1. *Retrieval Heads*, which represent only a fraction of the total, are crucial for processing long contexts and require full attention across all tokens. In contrast, the majority of attention heads, termed *Streaming Heads*, primarily focus on recent tokens and attention sinks (Xiao et al., ), and can operate effectively with a reduced KV cache that includes only recent tokens and attention sinks.

Building on the dichotomy of retrieval and streaming heads, we propose DuoAttention, a general, straightforward, and easily integrated approach that significantly accelerates both LLM’s decoding and pre-filling and reduces memory footprints, particularly in long-context scenarios.
The core innovation of DuoAttention is a lightweight, optimization-based procedure that identifies non-compressible retrieval heads using synthetic datasets. Unlike existing methods that rely on attention pattern profiling (Wu et al., ; Ge et al., ; Tang et al., ), DuoAttention directly measures output deviation resulting from token dropping, achieving higher compression rates and improved deployment efficiency.
DuoAttention is designed with simplicity and efficiency in mind: each Transformer layer has two KV caches— a full KV cache for crucial retrieval heads and a constant KV cache for streaming heads, which stores only attention sinks and recent tokens. This design allows DuoAttention to dramatically reduce memory usage and improve decoding speed in models like Llama-2/3 and Mistral, achieving up to 2.55$\times$ for MHA and 1.67$\times$ for GQA models while speeding up decoding by up to 2.18$\times$ and 1.50$\times$ and accelerating pre-filling by up to 1.73$\times$ and 1.63$\times$ for MHA and GQA models, respectively, with minimal accuracy loss compared to full attention.

Moreover, DuoAttention is fully compatible with important optimization techniques like GQA and quantization. We show that when combined with 8-bit weight 4-bit KV cache quantization, DuoAttention enables a Llama-3-8B model to handle up to 3.3 million contextual tokens measured on a single A100 GPU, achieving a 6.4$\times$ capacity increase compared to standard full attention FP16 deployments. DuoAttention paves the way for deploying LLMs in applications requiring million-level context handling.

## 2 DuoAttention

Figure: Figure 2: Overview of DuoAttention: (1) In the retrieval head identification phase, we assign a trainable gate value, $\alpha$, to each attention head, which blends the outputs of full attention and streaming attention. The training objective is to optimize these values to minimize the deviation from the full attention model’s output, while simultaneously applying a regularization loss to encourage lower gate values. This training phase is efficient, requiring only the gate values to be trainable—leaving all other model parameters frozen—thus allowing it to be completed within several hours on an 8 GPU node. (2) During deployment, these gate values are binarized to classify heads as either retrieval or streaming based on a threshold $\tau$. Retrieval heads, identified by a gate value above the threshold, use full attention, caching the KV pairs for all tokens. In contrast, streaming heads cache only the KV pairs of recent tokens and attention sinks.
Refer to caption: x2.png

### 2.1 Retrieval and Streaming Heads

##### Retrieval Heads

In Transformer-based LLMs, attention heads exhibit distinct and consistent patterns, reflecting their specialized functionalities (Clark et al., ; Xiao et al., ; Wu et al., ). Figure 1 visualizes two types of attention heads in the Llama-2-7B-32K-Instruct model using the sentence "The best fruit is orange. What is the best fruit? Orange". The left panel highlights an attention head that emphasizes relevant tokens during decoding; for instance, the first occurrence of "best fruit" is accentuated while decoding the second "best fruit," and the initial "orange" is highlighted when inferring the second "orange." These attention heads, which we term *Retrieval Heads*, are crucial for context processing as they capture contextually relevant tokens. Compressing the KV cache for retrieval heads would lead to the loss of vital contextual information, and thus they require full attention across all tokens.

##### Streaming Heads

In contrast, the attention head depicted in the middle panel of Figure 1 primarily attends to recent tokens and attention sinks (Xiao et al., ), without highlighting earlier relevant tokens in the context. We refer to these as *Streaming Heads*. Compressing the KV cache for Streaming Heads is feasible because dropping the unattended middle tokens does not significantly alter the attention output. Therefore, streaming heads can be optimized by retaining only the KV states of attention sinks and recent tokens, without compromising the model’s ability to manage long contexts.

##### Impact of Token Pruning on Retrieval and Streaming Heads

The right panel of Figure 1 shows a preliminary passkey retrieval experiment, showing that the model’s performance drops significantly when the middle tokens in the KV cache of retrieval heads are pruned, i.e., replaced with streaming attention. In contrast, removing the middle tokens for streaming heads has no significant impact on passkey retrieval accuracy. This observation indicates that we can enhance computational efficiency without sacrificing the model’s long-context capabilities: By dropping middle tokens for streaming heads while keeping full attention for retrieval heads, we reduce the memory demands of streaming heads to $O(1)$, thereby improving the efficiency of processing long contexts.

Figure: Figure 3: Example from the synthetic dataset used to identify retrieval heads. We embed ten 32-word passkeys within a long text and ask the model to recall these passkeys. Distillation loss is calculated solely on the passkeys.
Refer to caption: x3.png

### 2.2 Optimization-Based Identification of Retrieval Heads

##### Definition of Retrieval Heads

Section 2.1 qualitatively defines retrieval and streaming heads, but for precise identification, we need a concrete and quantitative definition. In this paper, we define “retrieval heads” as the attention heads that:

*significantly alter model outputs when restricted to recent tokens and attention sinks.*

We use this criterion to distinguish retrieval heads from streaming heads. Note that this definition differs from existing works (Ge et al., ; Wu et al., ; Tang et al., ) that rely solely on attention scores to identify retrieval heads, which overlook 1) the end-to-end impact of compressing the KV cache for specific attention heads, 2) the role of value states, and 3) the variability of attention distributions across layers and heads. In contrast, our definition directly measures output deviation, allowing us to identify attention heads crucial for long-context processing, *even when they are not apparent in attention scores*. We support this argument with ablation studies presented in Section 3.5.

##### Optimization-based Identification

We employ an optimization-based approach to identify retrieval heads, drawing inspiration from prior work in CNN filter pruning (Liu et al., ), as illustrated in Figure 2. First, we assign a gate value $\alpha_{i,j}$, to each key-value (KV) head in the LLM. This value intuitively represents the importance of the $j$-th KV head in layer $i$ for processing long-context information. Note that in models using GQA, one KV head can be associated with multiple attention heads, and our method accounts for the KV cache compression of an entire group of attention heads.

Our optimization-based identification method directly assesses the impact of compressing the KV cache with only sink and recent tokens for each KV head. We begin by initializing the gate value $\alpha_{i,j}\in[0,1]$ for each head at 1, assuming that all heads initially serve as retrieval heads. These gate values are then optimized, with the LLM’s parameters remaining fixed, limiting the number of trainable parameters to $N\times H$ and preventing the impact to the model’s original abilities.

During the forward pass, we combine the outputs of full and streaming attention (which attends only to sink and recent tokens) for each KV head, using the gate value as the mixing weight:

$$ $\texttt{attn}_{i,j}=\alpha_{i,j}\cdot\texttt{full_attn}+(1-\alpha_{i,j})\cdot \texttt{streaming_attn}$ $$

where the attention calculations are defined as:

$$ full_attn $\displaystyle=\text{softmax}({\bm{Q}}{\bm{K}}^{T}\odot{\bm{M}}_{\text{causal}} ){\bm{V}},$ streaming_attn $\displaystyle=\text{softmax}({\bm{Q}}{\bm{K}}^{T}\odot{\bm{M}}_{\text{ streaming}}){\bm{V}},\vspace{-0.5em}$ $$

where ${\bm{M}}_{\text{causal}}$ is the causal attention mask (a lower triangular matrix), and ${\bm{M}}_{\text{streaming}}$ represents a $\Lambda$-like mask (Han et al., ; Xiao et al., ) that attends only to recent and initial tokens.

##### Synthetic Dataset for Identifying Retrieval Heads

However, relying solely on natural language modeling objectives is insufficient for identifying retrieval heads because the supervision signal in natural text that requires inference over long spans is sparse, and most tokens can be inferred using local context. To address this, we design a synthetic dataset specifically aimed at enhancing the model’s long-context retrieval capabilities, allowing us to effectively identify which KV heads can be compressed without compromising the model’s performance.
As depicted in Figure 3, we create a passkey-retrieval dataset by embedding ten randomly generated passkey sequences of $s$ tokens in ten random locations within a very long context ($s=32$ in experiments). The model is then tasked with recalling these ten sequences at the end of the context.

##### Training and Loss Functions

We optimize the distillation loss, which is the L2 difference between the e last hidden state of the full attention model (${\bm{H}}_{\text{full}}$) and those of the model using DuoAttention (${\bm{H}}_{\text{mixed}}$), focusing only on the last $l$ passkey tokens in the entire inputs with $T$ tokens:

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

Our synthetic dataset ensures that every supervision signal is relevant to the final compression strategy, making the process lossless in terms of information retrieval accuracy. It proves to be more effective than using natural language modeling alone (see ablation studies in Section 13).
We use the L1 regularization term (a.k.a, Lasso (Tibshirani, )) to encourage sparsity in the gate values:

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

The final training loss is a combination of the distillation loss and the regularization loss, weighted by a hyperparameter $\lambda$, which we set as 0.05 in our experiments:

$$ $\mathcal{L}=\mathcal{L}_{\text{distill}}+\lambda\mathcal{L}_{\text{reg}}.$ (3) $$

Since the total number of trainable parameters is only thousands of floating-point numbers, this optimization process is fairly fast, with only 2,000 steps needed. All training experiments in our paper can be conducted on 8$\times$NVIDIA A100 GPU servers.

Figure: Figure 5: Decoding (left) and Chunked Pre-filling (right) Processes in DuoAttention: (1) The retrieval heads’ KV cache stores all tokens, while the streaming heads’ KV cache retains only recent tokens and attention sinks, ensuring constant memory usage. (2) The chunked pre-filling process of DuoAttention’s streaming heads on a 16-token sequence, with one attention sink, two recent tokens, and a chunk size of 4. DuoAttention’s streaming heads have linear time and constant memory complexity during long sequence pre-filling.
Refer to caption: x5.png

### 2.3 Deploying LLMs with DuoAttention

##### Binarizing Attention Implementations

At inference time, we apply full attention exclusively to the designated retrieval heads, identified using the optimized gate values from the training phase (as shown in Figure 4). We binarize the attention policy for each head based on a threshold $\tau$, determined by a specified sparsity quantile, to differentiate between retrieval heads and streaming heads:

$$ $\vspace{-0.5em}\texttt{attn}_{i,j}=\begin{cases}\texttt{full_attn}&\text{if } \alpha_{i,j}>\tau\\ \texttt{streaming_attn}&\text{otherwise}\end{cases}\vspace{-0.5em}$ (4) $$

##### Reordering Attention Heads

Before deployment, we preprocess the model by reordering the output channels of the Query, Key, and Value projection weights according to the attention head assignments. This reordering groups retrieval heads and streaming heads into two distinct, consecutive clusters, allowing for efficient slicing and concatenation operations when managing the KV cache for these two types of heads within a layer, rather than relying on scattering and gathering operations.

##### Decoding

As shown in Figure 5, we allocate two KV caches for each layer in the LLM during decoding: one for retrieval heads, which stores all past Keys and Values, and another for streaming heads, which stores only attention sinks and recent tokens, maintaining a constant size. When a new token is processed, its query, key, and value vectors are split along the head dimension to compute full attention for retrieval heads and streaming attention for streaming heads. The results are then concatenated along the head dimension for the output projection.

##### Chunked Pre-filling

We use FlashAttention-2 (Dao, ) to pre-fill the KV caches for both retrieval and streaming heads. In long-context LLMs, chunked pre-filling is a common practice (Agrawal et al., ; Kwon et al., ), dividing the prompt into fixed-length chunks to pre-fill the KV cache. This technique significantly reduces peak memory usage (see Table 10) by lowering the peak intermediate activation size in linear layers from sequence length to chunk size.
DuoAttention is fully compatible with chunked pre-filling, and the streaming heads’ pre-filling in DuoAttention can be achieved with linear time and constant memory complexity, *without* requiring specialized kernels. As shown in Figure 5, once a layer’s KVs are computed, the streaming head’s KV cache is immediately pruned to keep only the sink and recent tokens. The next chunk of incoming tokens will only attend to a constant number of contextual tokens during pre-filling. Let $L$ represent the sequence length and $K$ the chunk size. The pre-filling time complexity for streaming heads is optimized from $O(L^{2})$ to $O(LK)$, and the memory complexity is reduced from $O(L)$ to $O(K)$.

It’s important to note that DuoAttention’s design is well-suited for batch operations, which can further enhance LLM efficiency in serving scenarios with large batch sizes.

Figure: Figure 6: DuoAttention provides comparable accuracy as full attention on the Needle-in-a-Haystack benchmark using 25% full attention ratio on the MHA model and 50% full attention ratio on the GQA model.
Refer to caption: x6.png

Figure: Figure 7: DuoAttention provides better KV budget and accuracy trade-off on LongBench benchmarks.
Refer to caption: x7.png

## 3 Experiments

### 3.1 Setups

##### Models, Datasets, and Baselines

We evaluate DuoAttention on both long-context and short-context benchmarks to demonstrate that our method preserves model performance on tasks requiring both long and short contexts while significantly improving efficiency. For long-context evaluations, we use the Needle-in-a-Haystack (NIAH) benchmark (Kamradt, ) and LongBench (Bai et al., ). For short-context evaluations, we assess performance on MMLU (Hendrycks et al., ), MBPP (Austin et al., ), and MT-Bench (Zheng et al., ). We employ state-of-the-art open-source models, including Llama-2-7B-chat (Touvron et al., ) (and its long-context variant Llama-2-7B-32K-Instruct (Together, )), Llama-3-[8,70]B-Instruct (and its long-context variant Llama-3-8B-Instruct-Gradient-1048k (^*^**[https://huggingface.co/gradientai/Llama-3-8B-Instruct-Gradient-1048k](https://huggingface.co/gradientai/Llama-3-8B-Instruct-Gradient-1048k))), and Mistral-7B-v0.2-Instruct (Jiang et al., ). We compare our method against KV cache compression algorithms, including H2O (Zhang et al., ), TOVA (Oren et al., ), FastGen (Ge et al., ), and StreamingLLM (Xiao et al., ).

##### Implementation Details

We implement DuoAttention in PyTorch (Paszke et al., ) using RoPE (Su et al., ) and RMSNorm kernels from FlashInfer (Ye et al., ). For retrieval head identification, we use a batch size of 1, inserting ten 32-word passkeys into the BookSum (Kryściński et al., ) dataset. The identification process uses 128 sink tokens and 256 recent tokens. Training samples are drawn from 50 intervals ranging from 1,000 tokens to the model-specific maximum length. Passkeys are randomly inserted at 1000 points within the context. Further details are included in Appendix Section A.1.
We optimize gate values using the AdamW (Kingma & Ba, ) optimizer, starting with a learning rate of 0.02, warming up from 0.002 in the first 400 steps, and reducing back to 0.002 in the final 400 steps. All experiments run for 2,000 steps on NVIDIA A100 GPUs.

### 3.2 Long-Context Benchmarks

We evaluate DuoAttention using the Needle-in-a-Haystack (NIAH) benchmark and LongBench (Bai et al., ). We use two long-context models: Llama-2-7B-32K-Instruct and Llama-3-8B-Instruct-Gradient-1048k. We configure DuoAttention with a 25% retrieval head ratio for Llama-2-7B-32K-Instruct and a 50% ratio for Llama-3-8B-Instruct-Gradient-1048k. We compare DuoAttention with H2O, TOVA, and StreamingLLM using the same KV cache budget. We use 64 sink, 256 recent tokens, and 32,000 pre-filling chunk size for DuoAttention. Since the original designs of H2O and TOVA do not support long contexts, we modify their algorithms by replacing the pre-filling stage with FlashAttention and simulating decoding for the last 50 tokens of the input, following Tang et al. (). FastGen’s algorithm does not allow for the specification of the KV compression ratio, as it fluctuates with inputs. Therefore, we adjust the attention recovery ratio to ensure the KV cache budget is, on average, above 25% or 50% in the experiments shown in Figure 6. Additionally, FastGen’s quadratic memory cost during the attention profiling phase limits its ability to handle long-context samples. We measure FastGen’s performance on NIAH for Llama-2-7B up to a 24K context and for Llama-3-8B up to a 32K context; beyond these sizes, it results in out-of-memory errors. Detailed baseline implementations and justifications are provided in Appendix Section A.3 and Section A.5.

Needle-in-a-Haystack (NIAH) is a challenging pressure test designed to assess the ability of models to accurate identify and retrieve relevant information from lengthy context. As shown in Figure 6, all baseline methods fail to retrieve correct answers from the various depths of the long sequence, as they discard the KV cache containing the necessary information during generation. In contrast, DuoAttention retains all KV caches in the retrieval heads while discarding only those in the streaming heads, preserving the model’s retrieval capability. As a result, DuoAttention demonstrates strong performance across all sequence depths, handling lengths up to 1048K tokens effectively.

LongBench (Bai et al., ) is a comprehensive suite of long-context datasets encompassing multiple tasks and natural texts, designed to assess long-context understanding capabilities more thoroughly. Figure 7 shows the performance on 14 LongBench tasks, comparing different methods based on their KV cache budgets. DuoAttention shows a superior trade-off between KV budget and accuracy on most tasks, underscoring its generalizability. Notably, DuoAttention achieves performance comparable to full attention on most tasks, using a 25% KV cache budget for MHA and a 50% KV cache budget for GQA, consistent with the results observed in the needle-in-a-haystack benchmark. We compare DuoAttention with FastGen in Table 5 and 6 in the Appendix.
Table 3 and 4 in the Appendix provides full results for all 21 LongBench tasks using the 25% and 50% KV cache budget for the two models, showing that DuoAttention consistently outperforms baselines across most tasks and achieves the highest average scores.

Figure: Figure 8: Results on short benchmarks.
Refer to caption: x8.png

### 3.3 Short-Context Benchmarks.

To ensure that DuoAttention does not compromise the model’s performance on short-context tasks, we evaluate it alongside all baselines on three short-context benchmarks: MMLU, MBPP, and MT-Bench. These benchmarks assess the model’s knowledge, coding abilities, and helpfulness.
We use one-shot prompting for MMLU and zero-shot prompting for MBPP and MT-Bench. For DuoAttention, we configure 32 sink tokens and 128 recent tokens on MMLU, and 16 sink tokens and 64 recent tokens on MBPP and MT-Bench.
As shown in Figure 8 and Table 1, DuoAttention consistently outperforms all baselines under the same KV cache budget across various models, including Llama-2-7B, Llama-3-8B, and Llama-3-70B-Instruct. With a 50% KV cache budget, DuoAttention achieves near-lossless performance on most benchmarks, demonstrating that it preserves the model’s original capabilities.

Figure: Figure 9: Per-token decoding latency and memory usage of DuoAttention compared to full attention across varying context sizes. DuoAttention uses a 25% retrieval head ratio for Llama-2-7B (MHA) and 50% for Llama-3-8B (GQA). DuoAttention achieves up to 2.45$\times$ memory reduction for MHA and 1.65$\times$ for GQA models, along with up to 2.13$\times$ latency reduction for MHA and 1.5$\times$ for GQA models. These reductions approach the inverse of the retrieval head ratios as context length increases. Out-of-memory (OOM) results are linearly extrapolated from measured data.
Refer to caption: x9.png

Figure: Figure 10: Pre-filling latency and memory usage of DuoAttention compared to full attention across varying pre-filling chunk sizes. DuoAttention uses a 25% retrieval head ratio for Llama-2-7B (MHA), pre-filling a context of 100K tokens, and a 50% ratio for Llama-3-8B (GQA), pre-filling a context of 320K tokens. As the pre-filling chunk size decreases, DuoAttention achieves up to 1.73$\times$ latency reduction for MHA and 1.63$\times$ for GQA models, with memory reductions up to 2.38$\times$ for MHA and 1.53$\times$ for GQA models.
Refer to caption: x10.png

Figure: Figure 11: DuoAttention’s decoding memory and latency *vs*. KV budget with a fixed context length. Memory and latency are reduced linearly when the ratio of retrieval heads is reduced. DuoAttention achieves up to 2.55$\times$ memory reduction for MHA and 1.67$\times$ for GQA models, along with up to 2.18$\times$ latency reduction for MHA and 1.50$\times$ for GQA models.
Refer to caption: x11.png

### 3.4 Efficiency Results

We evaluate DuoAttention’s decoding latency and memory usage on Llama-2-7B and Llama-3-8B models on a single NVIDIA A100 GPU. We pre-allocate the KV cache for the entire benchmark sequence to prevent the extra overheads of dynamic memory allocations. The default number format for weights and activations is BFloat16. By employing a retrieval head ratio of 25% for Llama-2-7B and 50% for Llama-3-8B, DuoAttention maintains accuracy while significantly improving efficiency.

##### Decoding Efficiency

As shown in Figure 9, DuoAttention’s decoding speed scales linearly, though with a flatter slope compared to full attention, reflecting the chosen retrieval head ratio. This efficient scaling leads to significant reductions in memory usage and notable improvements in decoding speed. These improvements approach the inverse of the retrieval head ratios as context length increases.
Figure 12 shows DuoAttention’s speedup and memory savings across various KV budget settings for a fixed context size. Both decoding latency and memory usage decrease linearly as the ratio of retrieval heads is reduced in the deployment configuration. Under the settings in Figure 12, DuoAttention achieves maximum improvements on an A100 GPU: 2.55$\times$ memory reduction for MHA and 1.67$\times$ for GQA models, and 2.18$\times$ latency reduction for MHA and 1.50$\times$ for GQA models.

##### Pre-filling Efficiency

DuoAttention also accelerates long-context pre-filling for LLMs, as discussed in Section 2.3.
Figure 10 shows that DuoAttention significantly reduces both pre-filling latency and memory usage, with these savings increasing as the pre-filling chunk size decreases. This is because the time and memory complexity for the streaming heads are reduced with smaller chunk sizes. DuoAttention achieves up to 1.73$\times$ latency reduction for MHA and 1.63$\times$ for GQA models, with memory reductions of up to 2.38$\times$ for MHA and 1.53$\times$ for GQA models.

##### Combiniation with Quantization

To fit more tokens into limited memory, we can integrate weight and KV cache quantization with DuoAttention to maximize KV cache capacity. Previous studies have shown that weight quantization (Xiao et al., ; Lin et al., ) and 4-bit KV cache quantization (Lin* et al., ; Liu et al., ; Hooper et al., ) do not compromise model performance. We combine DuoAttention with the QServe (Lin* et al., ) quantization method and kernels to enable 8-bit weight and 4-bit KV cache LLM inference.
Measured results are shown in Figure 12. Combining quantization techniques with DuoAttention allows us to accommodate up to 3.30 million tokens on a single A100-80G GPU using the Llama-3-8B model, resulting in a 6.4$\times$ increase in capacity compared to the naive full attention BF16 deployment.

### 3.5 Ablation Studies

Figure: Figure 13: Ablation studies: (1) Comparison of retrieval head identification methods, showing the superiority of our optimization-based approach with synthetic data over attention profiling and language modeling. (2) Analysis of start and recent token sizes shows that combining sink and recent attention optimally identifies retrieval heads. (3) Deployment performance indicates 16 attention sinks and 64 recent tokens are optimal, with minimal gains beyond these values.
Refer to caption: x13.png

We conduct ablation studies using the Mistral-7B-Instruct-v0.2 on passkey retrieval and MMLU datasets. For the passkey retrieval task, we embed an 8-word passkey within a 30K-word text and perform a linear sweep across 100 insertion depths, reporting exact match accuracies.

##### Optimization-based vs. Attention Profiling-based Retrieval Head Identification

We assess our optimization-based method against attention profiling, as used in FastGen (Ge et al., ) and RazorAttention (Tang et al., ), utilizing the same synthetic passkey dataset for both. Results in Figure 13 (1) show our method significantly outperforms attention profiling, which struggles to identify retrieval heads, affecting model optimization accurately.

##### Optimizing with Synthetic Data vs. Language Modeling

As illustrated in Figure 13 (1), our approach of using synthetic data to identify retrieval heads produces significantly better results than traditional language modeling, which computes loss on all tokens in natural data.

##### Necessity of Sink+Recent Attention in Optimization

Figure 13 (2) highlights the importance of combining sink and recent attention during the optimization phase. Exclusive reliance on either starting or recent token attention is inadequate for effective retrieval head identification.

##### Deployment Phase Configuration

We analyze the deployment configuration for attention sinks and recent tokens within streaming heads. Our findings indicate that performance plateaus at 16 sink tokens and 64 recent tokens (Figure 13 (3)). Further increases yield marginal improvements.

## 4 Related Work

Various approaches have been developed to scale up LLMs and improve their efficiency in handling long contexts. These methods can be grouped into four main categories: optimizing model architectures, using approximate attention mechanisms, applying KV cache quantization, and system-level optimizations.

##### Model Architecture

Multi-Query Attention (MQA)(Shazeer, ) and Grouped-Query Attention (GQA)(Ainslie et al., ) reduce the size of the Key-Value (KV) cache by sharing KV heads across query heads. However, these methods require pre-training with specific architectures and do not reduce computational costs. Linear attention Transformers (Gu & Dao, ) reduce memory usage but tend to underperform on tasks requiring long-context processing.

##### Approximate Attention

Methods like Sparse Transformer (Child et al., ) and LongFormer (Beltagy et al., ) use local or block attention patterns to reduce computational complexity. BigBird (Zaheer et al., ) achieves linear complexity by combining local and global attention, but many of these methods require custom GPU kernels or retraining, limiting their practicality. H2O (Zhang et al., ) and TOVA (Oren et al., ) simplify attention by discarding tokens based on query patterns. StreamingLLM (Xiao et al., ) identifies "attention sinks" and proposes always retaining initial and recent tokens to maintain constant decoding latency and memory usage, allowing the model to process significantly more input tokens than the pre-training sequence length. FastGen (Ge et al., ) profiles attention heads to discard tokens during decoding. However, our experiments show that these methods degrade the long-context abilities of LLMs. Also, these methods cannot reduce the pre-filling cost of long-context LLMs.

##### KV Cache Quantization

Techniques such as 8-bit and 4-bit quantization (Liu et al., ; Hooper et al., ; Lin* et al., ) reduce the size of KV caches, but they do not address the computational overhead of attention kernels. These methods are complementary to DuoAttention and can be used together to further reduce memory usage.

##### System Optimizations

vLLM (Kwon et al., ) and FlashAttention (Dao et al., ; Dao, ) improve attention computation efficiency by optimizing batch processing and utilizing GPU memory hierarchies. FlashDecoding (Hong et al., ) and RingAttention (Liu et al., ) introduce further improvements in decoding speed and sequence-level parallelism. While these methods enhance computational performance, they do not address KV cache size reduction, making them complementary to DuoAttention for additional speed and memory optimization.

##### Recent Works

Several recent works share similar ideas with DuoAttention. Wu et al. () introduces the concept of retrieval heads to explain LLMs’ long-context capabilities. However, their approach does not compress the KV cache for non-retrieval heads, focusing solely on accuracy. MInference (Jiang et al., ) accelerates pre-filling for long-context LLMs by using sparse attention patterns but does not optimize KV cache storage or latency during decoding. RazorAttention (Tang et al., ) also divides attention heads into retrieval and non-retrieval categories but relies on attention profiling, which, as our experiments show, is less accurate than our optimization-based approach. Also, RazorAttention doesn’t optimize pre-filling. DuoAttention offers more effective KV cache management and higher compression rates, leading to better performance for both pre-filling and decoding in long-context applications.

## 5 Conclusion

We introduce DuoAttention, a framework that optimizes memory and computational resources in LLMs by distinguishing between *Retrieval Heads* and *Streaming Heads*. By applying a full KV cache only to retrieval heads, DuoAttention significantly reduces memory usage and latency for both decoding and pre-filling in long-context applications. It achieves memory reductions of up to 2.55$\times$ for MHA and 1.67$\times$ for GQA models, with decoding speed improvements of up to 2.18$\times$ for MHA and 1.50$\times$ for GQA, and pre-filling accelerations of up to 1.73$\times$ and 1.63$\times$, respectively, with minimal accuracy loss compared to full attention.
When combined with quantization, DuoAttention further boosts KV cache capacity, supporting up to 3.30 million contextual tokens on a single A100 GPU. DuoAttention paves the way for LLMs to handle contexts with millions of tokens.

#### Acknowledgments

We thank MIT-IBM Watson AI Lab, MIT and Amazon Science Hub, MIT AI Hardware Program, National Science Foundation, Hyundai and Samsung for supporting this research. We thank NVIDIA for donating the DGX server.

## Appendix A Appendix

### A.1 Experimental Details

We use FSDP2 in PyTorch for model training and DeepSpeed Ulysses (Jacobs et al., ) sequence parallelism to support long sequences. During training, we use an efficient block-sparse approximation of $\Lambda$-like attention for streaming attention, as implemented in Guo et al. () and illustrated in Figure 14. Maximum sequence lengths vary across models, as detailed in Table 2.

**Table 2: Training Hyperparameters.**
| Models | Max. Seq. Lengths |
| --- | --- |
| Llama-2-7B-chat | 4096 |
| Llama-2-7B-32K-Instruct | 32000 |
| Llama-3-8B-Instruct | 8192 |
| Llama-3-8B-Instruct-1048K | 32000 |
| Llama-3-70B-Instruct | 8192 |
| Mistral-7B-Instruct-v0.2 | 32000 |

Figure: Figure 14: Block-sparse approximation of $\Lambda$-like attention.
Refer to caption: x14.png

### A.2 Full LongBench Results

**Table 3: Full LongBench results with Llama-3-8B-Instruct-1048K. DuoAttention achieves the best performance with a 50% KV cache budget on most datasets.**
| Dataset | Full | H2O (50%) | SLLM (50%) | TOVA (50%) | Duo (50%) |
| --- | --- | --- | --- | --- | --- |
| Average | 40.08 | 35.76 | 32.26 | 35.55 | 40.21 |
| 2WikiMQA | 28.78 | 27.99 | 29.22 | 26.93 | 29.08 |
| DuReader (zh) | 30.41 | 24.94 | 9.41 | 27.00 | 29.31 |
| GovReport | 34.23 | 29.44 | 29.08 | 30.10 | 32.72 |
| HotpotQA | 40.37 | 36.77 | 39.27 | 38.45 | 41.63 |
| LCC | 38.19 | 43.09 | 41.94 | 42.31 | 44.16 |
| LSHT (zh) | 38.00 | 25.00 | 25.50 | 24.50 | 30.00 |
| MultiNews | 27.73 | 25.52 | 24.85 | 26.32 | 27.72 |
| MultiFieldQA-en | 52.62 | 38.53 | 28.11 | 44.94 | 51.44 |
| MultiFieldQA-zh | 50.58 | 38.25 | 31.07 | 40.82 | 52.40 |
| Musique | 24.22 | 19.24 | 20.47 | 23.07 | 24.65 |
| NarrativeQA | 26.56 | 25.13 | 22.06 | 25.64 | 24.54 |
| Passage Count | 1.00 | 2.05 | 1.64 | 1.00 | 0.00 |
| PassageRetrieval-en | 81.00 | 74.75 | 49.00 | 72.00 | 87.00 |
| PassageRetrieval-zh | 62.15 | 52.57 | 38.90 | 46.13 | 62.15 |
| Qasper | 29.21 | 20.65 | 21.77 | 23.06 | 26.93 |
| QMSum | 24.52 | 22.87 | 22.11 | 23.16 | 24.20 |
| RepoBench-P | 38.94 | 39.98 | 37.60 | 40.14 | 46.12 |
| SAMSum | 42.51 | 40.78 | 40.25 | 40.50 | 41.83 |
| TREC | 71.50 | 64.00 | 67.00 | 54.00 | 71.00 |
| TriviaQA | 87.70 | 85.98 | 86.11 | 84.97 | 87.14 |
| VCSUM (zh) | 11.37 | 13.45 | 12.10 | 11.59 | 10.46 |

**Table 4: Full LongBench results with Llama-2-7B-Instruct-32K. DuoAttention achieves the best performance with a 25% KV cache budget on most datasets.**
| Dataset | Full | H2O (25%) | SLLM (25%) | TOVA (25%) | Duo (25%) |
| --- | --- | --- | --- | --- | --- |
| Average | 37.52 | 26.84 | 27.80 | 29.78 | 34.49 |
| 2WikiMQA | 35.59 | 28.87 | 29.69 | 31.18 | 33.37 |
| DuReader (zh) | 25.10 | 15.56 | 13.96 | 15.51 | 23.99 |
| GovReport | 31.23 | 20.66 | 24.14 | 22.88 | 27.98 |
| HotpotQA | 47.98 | 39.60 | 40.39 | 47.45 | 50.44 |
| LCC | 51.21 | 45.78 | 44.25 | 47.91 | 48.34 |
| LSHT (zh) | 34.50 | 16.50 | 17.50 | 18.50 | 25.50 |
| MultiNews | 27.11 | 19.21 | 20.54 | 21.41 | 25.03 |
| MultiFieldQA-en | 33.95 | 21.01 | 16.69 | 18.19 | 25.49 |
| MultiFieldQA-zh | 45.79 | 19.81 | 22.50 | 24.96 | 39.23 |
| Musique | 22.97 | 20.63 | 20.09 | 21.00 | 19.27 |
| NarrativeQA | 24.11 | 19.14 | 21.13 | 23.06 | 20.49 |
| Passage Count | 0.00 | 0.53 | 0.58 | 0.00 | 0.33 |
| PassageRetrieval-en | 50.92 | 19.50 | 19.08 | 30.17 | 47.25 |
| PassageRetrieval-zh | 37.68 | 11.75 | 16.77 | 32.38 | 40.93 |
| Qasper | 33.23 | 16.84 | 17.68 | 20.85 | 26.59 |
| QMSum | 20.79 | 18.89 | 20.05 | 20.16 | 21.48 |
| RepoBench-P | 51.58 | 45.16 | 45.25 | 49.03 | 48.58 |
| SAMSum | 42.10 | 39.73 | 37.43 | 36.17 | 33.10 |
| TREC | 71.50 | 48.50 | 56.50 | 47.00 | 68.50 |
| TriviaQA | 86.21 | 85.16 | 85.24 | 85.65 | 86.15 |
| VCSUM (zh) | 14.45 | 10.71 | 14.36 | 11.85 | 12.35 |

**Table 5: Comparison of FastGen and DuoAttention on a subset of LongBench using the Llama-3-8B-Instruct-1048K model.**
|  | FastGen (>50%) | DuoAttention (50%) |
| --- | --- | --- |
| Average | 32.82 | 40.01 |
| 2WikiMQA | 18.61 | 29.08 |
| DuReader (zh) | 20.22 | 29.31 |
| HotpotQA | 33.08 | 41.63 |
| LCC | 46.50 | 44.16 |
| MultiNews | 18.18 | 27.72 |
| MultiFieldQA-en | 44.05 | 51.44 |
| MultiFieldQA-zh | 42.15 | 52.40 |
| Musique | 13.58 | 24.65 |
| Passage Count | 0.09 | 0.00 |
| PassageRetrieval-en | 93.12 | 87.00 |
| PassageRetrieval-zh | 40.75 | 62.15 |
| Qasper | 26.51 | 26.93 |
| QMSum | 24.03 | 24.20 |
| SAMSum | 34.12 | 41.83 |
| TriviaQA | 69.92 | 87.14 |
| VCSUM (zh) | 0.23 | 10.46 |

**Table 6: Comparison of FastGen and DuoAttention on a subset of LongBench using the Llama-2-7B-32K-Instruct model.**
|  | FastGen (>25%) | DuoAttention (25%) |
| --- | --- | --- |
| Average | 19.01 | 32.81 |
| 2WikiMQA | 28.05 | 33.37 |
| MultiNews | 12.60 | 25.03 |
| MultiFieldQA-en | 28.58 | 25.49 |
| MultiFieldQA-zh | 22.44 | 39.23 |
| PassageRetrieval-zh | 3.38 | 40.93 |

### A.3 Implementation of H2O and TOVA on Long-Context Benchmarks

The original designs of the H2O and TOVA algorithms are not compatible with FlashAttention during pre-filling, as they rely on attention scores to perform token eviction. Since attention scores in FlashAttention are never materialized, these algorithms cannot be used in pre-filling, which is one of their main flaws. Therefore, it’s not possible to evaluate these algorithms in long-context settings like needle-in-the-haystack and LongBench, as they cause OOM during context pre-filling. To compare with these strategies, we modified the algorithms: during pre-filling, we used FlashAttention for exact calculations. During the decoding stage, we perform token eviction based on the generated tokens’ attention scores to contextual tokens. This modification improves performance compared to the original design since pre-filling is exact and token eviction occurs only during decoding. In extreme scenarios, if there is only one generated token in the answer (e.g. multiple-choice tasks), our implementation of H2O and TOVA will be exact with full attention, unlike their true accuracy. To approach their true performance, we simulate the last 50 tokens in long input benchmarks (needle-in-the-haystack and LongBench) as generated tokens to perform their token eviction policy long enough, as well as our algorithm. This experimental setting is also used by  Tang et al. (). Experimental results show our method can pass this pressure test, while H2O and TOVA cannot.

### A.4 NIAH results on Mistral models

Figure: Figure 15: NIAH result on the Mistral-7B-Instruct-v0.2 model.
Refer to caption: extracted/5908680/figure/Mistral-7B-Instruct-v0.2_duo_attn-attn_pattern=lr=0.02-reg=0.1-ctx=1000_32000-multi_passkey4-sparsity=0.5.png

### A.5 Implementation of FastGen on Long-Context Benchmarks

Due to the lack of official implementation of the FastGen (Ge et al. ()) algorithm, we reproduce it using a community codebase (Adams et al. ()), which is referenced by FastGen’s official repository. In the FastGen algorithm, the pruning ratio cannot be directly configurable; instead, the recovery ratio $T$ is used to control sparsity as outlined in the FastGen paper. To quantify sparsity, we calculated the average KV cache usage across all test cases as the overall measure of sparsity. For the Llama-2-7B model, we set the recovery ratio to $0.7$, ensuring the average KV cache budget was over 25% of the full KV cache. Similarly, for the Llama-3-8B model, we set the recovery ratio to $0.87$, ensuring the average KV cache budget was more than 50% of the full KV cache. Additionally, since FastGen uses the full attention map of the user-provided prompt to profile the types of different heads, it results in an $O(n^{2})$ attention map complexity. Therefore, we are unable to test its performance in long contexts. For the long context benchmark, we used 8 A100-80G GPUs, achieving sequence lengths of up to 24k tokens for the Llama-2-7B model and up to 32k tokens for the Llama-3-8B model. In addition to the needle-in-the-haystack benchmark shown in Figure 6, we also evaluated FastGen on LongBench for both models. However, due to the quadratic memory consumption of FastGen, we only report results for datasets that were feasible to run on 8x A100-80G GPUs using FastGen. As shown in Table 5 and Table 6, DuoAttention can consistently outperform FastGen on LongBench datasets.

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 8
presentation: 3
questions: One of the concurrent works, RazorAttention, also utilizes a very similar recipe. The authors compared it in Figure 13(1), but I find the reading to be very different from RazorAttention's own reporting. Any insight of why?

---

## Post-rebuttal update

The latest response from the authors is satisfactory for the most part: 

* It is a reasonable choice to feature Llama-2-7B-32K-Instruct if the goal is to include a long context-capable MHA model.
* The added Llama 3.1 results demonstrate that DuoAttention remains performant on one of the most mainstream long context models.
* I appreciate the authors for recognizing the connection between their optimization-based approach and prior gating network studies.
* The budget discussion regarding RazorAttention and NIAH is sound. However, I must note that if the authors did not explore echo heads — a key element of RazorAttention — in the ablation studies on "attention profiling," this omission should be explicitly noted.
* Similarly, if Figure 19 is not done with an aligned budget, it is not fair to directly compare Duo's NIAH performance with MInference by citing Figure 19. That said, using it as a baseline for Figure 20 is reasonable. I also have some reservations about claiming that the "sandwich in a SF park" needle is harder than the magic city/passkey retrieval-like needle, as the former is known to be very prompt-sensitive. However, this is a minor distinction, and as long as the setting is explicitly noted, both can be considered solid needle setups.


I am improving my rating to 8 if the authors can supply full $\infty$Bench report in their revision.

(In addition, I also add that I don't believe DuoAttention's incomptability with MQA models is a big deal. Yes, this limitation should be clearly noted; but it is rare to find powerful MQA models, and the community adopts GQA as a middle ground between MQA and MHA for good measures.)
soundness: 4
strengths: 1. The proposed method is extremely performant. It is rare to see a prefilled compressible token dropping method capable of doing NIAH. 

2. The efficiency optimization is solid, and it is supported by a thorough efficiency evaluation.

3. The optimization-based head selection approach offers a different avenue for attention head pattern matching.

4. The paper is nicely written.
summary: DuoAttention leverages finding around retrieval heads. By allocating more KV budget to such heads and compressing the rests, the method claims to have excellent long context performance.
weaknesses: 1. The method mainly involves retrieval head + token dropping (in this case, StreamingLLM), which is slim in novelty. I don't think it is much of an issue given the impressive performance boost achieved while remaining practical, but this is worth mentioning.

2. The featured compared methods — SteamingLLM, H2O, and TOVA — are very dated. I'd like to see how it stands off with more modern KV cache compassion/sparse inference methods, e.g., SnapKV and MInference.

3. The NIAH setting is unclear. What does the needle look like? What background is used? The author should better clarify this, as different settings can lead to very different results. 

4. Following #3, it looks like LongBench is the only real long context test evaluated. I'd like to see more coverage on other long context datasets, e.g., $\infty$Bench and RULER.

I believe the proposed method marks a breakthrough in eviction-based methods. Its execution is solid and the presentation is clear. I am willing to boost my score should the coverage/clarity of performance evaluation be improved.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Please address the questions in the weaknesses.
soundness: 3
strengths: 1.	It saves the GPU memory and enhances the inference efficiency for both MHA and GQA.
summary: This paper introduces a lightweight, optimization-based method DuoAttention that only applies a full KV cache to retrieval heads while using a constant-length KV cache for streaming heads. The proposed method achieves good performance in both MHA and GQA.
weaknesses: 1.	There are some related works identifying the different patterns of retrieval and non-retrieval heads, such as RazorAttention, Retrieval Head Mechanistically Explains Long-Context Factuality, MInference, etc. It would be better to clarify the novelty of the proposed method.
2.	It can only be applied to multiple KV cache scenarios. For only single KV cache model structure, such as YOCO, it cannot be applied.
3.	The baselines are not comprehensive. It would be beneficial if the authors could compare their results with more advanced baselines.
4.	It seems the proposed method cannot achieve the best performance in all the benchmarks in Figure 7. It would be beneficial to give detailed analysis on the worse cases.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: * How effective is this method for reasoning-based tasks? Do the retrieval heads remain the same for question-answering type tasks as well as reasoning type tasks?
* Rather than merely illustrating the distinction between retrieval and streaming heads through specific examples, it would be more compelling to present statistically significant evidence demonstrating how the identified retrieval heads generalize across various sequences and tasks.
* Is there any way to extend this method to multi-GPU setting?
* To what extent can the memory savings lead to increased throughput? (Just a suggestion)
soundness: 3
strengths: * The paper is well-structured and the claims are easy to follow.
* This work addresses a significant research problem in LLM inference - KV cache optimization. 
* The authors perform comprehensive ablation studies to motivate each aspect of their proposed approach. The advantage of their proposed novel retrieval head selection method over the existing methods is well-supported by empirical results.
summary: This work presents methods to reduce memory footprint and computational demands of Large Language Model's (LLM) attention layers to enhance inference speed and to accommodate longer context lengths. The authors build upon the observation that only a few attention heads within an LLM, also known as retrieval heads, have considerable impact on long-context performance. They propose a novel, optimization-based approach to identify these less critical heads, unlike previous methods that relied on heuristic techniques. This optimization process utilizes synthetic data designed specifically to evaluate long-context capabilities. Their results indicate superior model performance compared to current state-of-the-art retrieval head allocation techniques and certain KV sparsification methods, while maintaining comparable KV cache budget.
weaknesses: * While tensor parallelism has proven effective in reducing latency and memory usage per GPU, this paper omits any discussion on how its proposed methods could be adapted to such model-parallel settings. For example, in an 8-way tensor-parallel configuration, the retrieval heads might become the performance bottleneck, potentially negating the gains from this work.

* Some state-of-the-art KV sparsification methods are missing in the comparative analysis. For instance, SnapKV, PyramidKV, AdaKV are some notable omissions.
   1. Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation, 2024d. *Arxiv /abs/2404.14469*
   2. Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, and Wen Xiao. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024. *Arxiv /abs/2406.02069*
   3. Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, and S. Kevin Zhou. Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference, 2024. *Arxiv /abs/2407.11550*

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1) Is the proportion of Retrieval Heads the same for different models? Why choose 25% Retrieval Heads for MHA and 50% for GQA?
soundness: 3
strengths: 1) The paper is well written and well motivated.
2) DuoAttention is a plug-and-play solution compatible with FlashAttention.
3) DuoAttention can accelerate inference during both the prefill and decoding stages.
summary: To tackle KV Cache problem, the paper introduces a new approach called DuoAttention, which classifies attention heads in LLMs into two types: Retrieval Heads and Streaming Heads. Retrieval Heads are essential for processing long contexts and require full KV cache storage, while Streaming Heads focus on recent tokens and can operate with reduced KV caches. DuoAttention uses a lightweight optimization process to identify non-compressible Retrieval Heads, allowing for efficient memory usage and faster processing. This method integrates easily with existing optimization techniques and significantly reduces the memory footprint and decoding time. When combined with quantization, DuoAttention enables models like Llama-3-8B to handle up to 3.33 million tokens on a single GPU, achieving a 6.4× increase in capacity compared to standard deployments.
weaknesses: 1) Lack the experiment with the accuracy results after combining DuoAttention and KV quantization.
2) Lack the comparisons with new KV Cache compression technologies such as Minference.
3) Using different compression rates for different heads may lead to uneven computation across cards during parallel inference, potentially affecting performance.

---

confidence: 5
contribution: 3
initial_rating: 8
presentation: 4
questions: 1. the offline method to find the retrieval heads is more important for practical application
2. the compression ratio of KV cache should be improved
soundness: 3
strengths: 1. the experiments is enough and confident
2. the writing and clarity are clear
3. the concept of retrieval heads is significant for static KV cache compression
summary: The DuoAttention uses the retrieval heads and local streaming heads, realizing the KV cache compression and acceleration:
1) the retrieval heads preserve the major information of the long-context, which can not be compressed;
2) the non-retrieval heads are compressed by locality, which is accelerated with StreamingLLM;
3) the retrieval heads are detected with importance learnable weights;
4) Finally, the prefill is accelerated by chunked prefilling, and the decoding is accelerated with compressed KV cache.
weaknesses: 1. the detection of retrieval heads is based on the fine-tuning of importance weights, the usability should be enhanced
2. the compression ratio is limited, if the retrieval heads preserve the full context information

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. The core of this paper focuses on prioritizing the retention of the KV cache in important attention heads while attempting to discard less important cache in other heads, based on offline detection results. Some follow-up works on SnapKV seems to align closely with this approach[3][4]. For example, [4] employs a similar strategy by identifying important attention heads through online analysis of "altering model outputs," subsequently allocating more budget to these key heads and reducing the budget for others. What do you think is the relationship between these budget allocation strategies and the detection of retrieval heads?

2. In Equation 2, why is the index i set to $L$? Shouldn't it be $N$ instead?

3. How do you control the cache budget to a specifical ratio like 50% in the experiments? It seems challenging to precisely manage the cache budget within this approach.

[1] Li, Y., Huang, Y., Yang, B., Venkitesh, B., Locatelli, A., Ye, H., ... & Chen, D. (2024). Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404.14469.

[2] Zhang, Y., Gao, B., Liu, T., Lu, K., Xiong, W., Dong, Y., ... & Xiao, W. (2024). PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling. arXiv preprint arXiv:2406.02069.

[3] Rehg, I. (2024). KV-Compress: Paged KV-Cache Compression with Variable Compression Rates per Attention Head. arXiv preprint arXiv:2410.00161.

[4] Feng, Y., Lv, J., Cao, Y., Xie, X., & Zhou, S.K. (2024). Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference. ArXiv, abs/2407.11550.

While this paper is engaging and provides extensive evaluations, several limitations hold me back from giving it a higher score.  If the authors can address these concerns during rebuttal phase, I believe it would greatly enhance the paper’s quality, and I’d be glad to reconsider my score.
soundness: 3
strengths: 1. The paper is well-written and easy to follow, with a novel methodology.
2. The work is solid, having been evaluated on multiple benchmarks. It also combines the orthogonal technique of quantization, further enhancing its results.
3. The evaluations demonstrate promising results.
summary: Overall, this is an interesting and solid work. It categorizes all attention heads of LLMs into two types—Retrieval Heads and Streaming Heads—based on an offline analysis. By pruning the KV cache of Streaming Heads, this study achieves more efficient inference.
weaknesses: 1. **Choice of Baseline**: DuoAttention is essentially a KV cache compression algorithm, and all the performance gains come from reducing the amount of KV cache that was previously stored, whether in the prefilling stage or the decoding stage. While the authors compare their method to H2O, TOVA, and StreamingLLM, these approaches are no longer state-of-the-art . For instance, SnapKV[1], which was published in April and is now widely regarded as a more powerful KV cache compression method, should have been considered. In the commonly used LongBench evaluation, SnapKV outperforms H2O with only one-quarter of the cache budget[1]. Additionally, in needle-in-a-haystack tasks, SnapKV benefits from its pooling operations and significantly improves accuracy compared to H2O. When compressing an 8K sequence to 128 tokens, SnapKV incurs minimal performance loss, outperforming H2O by a large margin, as shown in [2]. In the experimental section, the authors highlight a limitation of H2O: “Since the original designs of H2O and TOVA do not support long contexts, we modify their algorithms by replacing the pre-filling stage with FlashAttention and simulating decoding for the last 50 tokens of the input, following Tang et al. (2024b).” This limitation arises because H2O requires the accumulation of global attention weights, which is incompatible with FlashAttention. However, SnapKV and its successors only observe attention scores within a small observation window for compression, making them highly compatible with FlashAttention. This allows recalculating only a small portion of attention weights  for compression when combined with chunked prefill. Therefore, SnapKV would be a more appropriate baseline for this paper.

2. **Inappropriate statements:**

a.  _“Moreover, DuoAttention is **fully compatible** with important optimization techniques like GQA and quantization.” “Approximative attention methods, such as H2O (Zhang et al., 2023b), StreamingLLM (Xiao et al., 2023b), TOVA (Oren et al., 2024), and FastGen (Ge et al., 2024), often compromise accuracy in long-context applications and are **incompatible** with essential KV cache optimization techniques like GQA"_ 

I believe the compatibility of DuoAttention with GQA is not significantly different from other methods. DuoAttention forcedly classifies a group of attention heads under GQA into the same category through offline analysis, thus upporting GQA. However, other methods could achieve the same effect by accumulating weights for a group of kv cache during cache eviction. Therefore, GQA compatibility seems more like an implementation detail rather than a unique feature of the algorithm. In fact, StreamingLLM does not encounter any GQA compatibility issues. The lack of GQA compatibility in earlier works, in my opinion, stems from the fact that GQA was not widely adopted when those methods were initially proposed. Later methods, for the sake of comparison, maintained this approach without further integrating GQA, which could have been easily addressed at the code implementation level. A recent study[3] applying cache eviction in paged attention has demonstrated that this compatibility can be easily achieved in practice.

b._“Also, these methods (H2O, TOVA, StreamingLLM) cannot reduce the prefilling cost of long-context LLMs.”_

 Previous evaluations of cache compression often compress the KV cache after the prefilling stage to ensure comparability. Since the KV cache for each layer can be compressed immediately after the completion of that layer's prefilling computation, this already substantially reduces the peak memory usage during the prefilling process. If further reduction in computational cost during the prefilling stage or a more significant decrease in peak memory usage for very long input texts is desired, combining these methods with chunked prefilling can provide additional acceleration—a straightforward solution. This is particularly applicable to StreamingLLM, which is same with the same Streaming Head in this paper, and thus faces no obstacles in applying such methods. If one argues that the H2O method combined with chunked prefilling may require additional accumulation of global attention weights, the additional baseline SnapKV can effectively address this issue.

Meta-review: This paper presents DuoAttention, a method designed to make long-context inference in large language models (LLMs) more efficient. The approach reduces memory use and processing time by focusing full key-value caching on a subset of attention heads—referred to as retrieval heads—while applying a streamlined cache strategy to other heads. The authors propose an optimization-based method for identifying these critical attention heads, achieving significant improvements in efficiency without compromising performance on long-context tasks. The experiments demonstrate up to 2.55x memory savings and 2.18x faster processing for specific models.

The paper is well-written and addresses an important challenge in deploying LLMs with long-context capabilities. Its strengths include a clear motivation for the proposed method, strong empirical results, and a thoughtful evaluation across multiple benchmarks. However, some limitations remain. The work could benefit from more extensive testing on real-world datasets and a broader range of baseline comparisons. Additionally, while the optimization strategy is effective, the paper would be stronger if it provided deeper insights into how the method generalizes across diverse tasks.

This paper is recommended for acceptance due to its potential to advance the efficiency of LLMs in long-context scenarios. The strong experimental results and practical implications outweigh the identified gaps.

Decision: Accept (Poster)