DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads
Authors: Guangxuan Xiao
Jiaming Tang
Jingwei Zuo
junxian guo
Shang Yang
Haotian Tang
Yao Fu
Song Han
Venue: ICLR 2025 Poster (2025)
Keywords: Large Language Models; Long Context; Efficiency;
Primary area: foundation or frontier models, including LLMs
TL;DR: We significantly reduce both pre-filling and decoding memory and latency for long-context LLMs without sacrificing their long-context abilities.

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.33 million context length measured on a single A100 GPU. Code is provided in https://github.com/mit-han-lab/duo-attention.

Published as a conference paper at ICLR 2025
DUOATTENTION: EFFICIENT LONG-CONTEXT LLM INFERENCE WITH RETRIEVAL AND STREAMING HEADS
Guangxuan Xiao Jiaming Tang'
Jingwei Zuo? Junxian Guo 1.3
Shang Yang' Haotian Tang*
Yao Fur
Song Han
MIT Tsingguu University 3 SITU University of Edinburgh
5 NVIDIA
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 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 Streaming Heads-do not require full attention. Based on this insight, we introduce Attenton, a framework that only applies 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-base algorithm with synthetic data to identify retrieval heads accurately. Our method significantly reduces long-context inference memory by up to 2.55 for MHA and 1.67 x for GQA models while speeding up decoding by up to 2.18 x and 1.50 and accelerating pre-filling by up to 1.73x and 1.63x for MHA and GQA models, respectively, with minimal accuracy loss compared to full attention. Notably, combined with quantization, DuoAAtention enables Llama-3-8B decoding with 3.3 million context length on a single A100 GPU. Code is provided in the link
INTRODUCTION
Large language models (LLMs) (Tovrron et al. 2223a b OpenAl 2023 Black et al. 2022) are at the
forefront of the AI revolution, powering advanced application such as multi-round dialogues (Schul
man et al. 2022 Taori al.
2023 Chiang et al. (023), long document summarization (Goyal
& Durrett 2020 Zhang et al. 2023a), and tasks involving mixed modalities like visual and video
understanding (Liu et al.. 2023b Lin al. 2023). 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 224x224 image corresponds to 256 tokens (Liu et al. 2023b), and 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-88 (Dubey et al. 2024) model architecture, serving with FP16 KV cache for million tokens would require at least 137 GB of memory-exceeding the capacity of a single 80GB GPU. Additionally, the latencies
Part of the work done during an internship NVIDIA.

---

Published as a conference paper at ICLR 2025
Layeer 15 Head 12 Sicrivaal cepture tolors outoot
Layer 10 Hexd II
100%
Streamay ettenion roont sakens
Fall
80%
Streaming Head First
Accuracy
60%
Retrieval Head Fins
Randomh
40%
20%
0.1
0.2
0.4
0.7
Retrieval Heads
Streaming Heedss
Streaming Attemtion Ratio
Figure I: Visualization of attention maps the Llama-2-7B model for the sentence "The fruit orange. What is the best fruit? Orange." " shows the distinct roles of retrieval heads (c.g., Layer 15, Hcad 12) and streaming heads (c.g. Layer 10, Head 11). 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 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-contex, passkey retrieval accuracy is shown: modifying retrieval heads severely damages performance, while constraining streaming heads has minimal impacts.
of pre-filling and decoding with such large contexts are significant, posing substantial challenges to the effective use of LLMs in Long-contex: 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 trropedd-uuers Attention (GQA)(Ainss) et al., 2023), require model pre-training and fail to reduce computational costs. Linear Attention methods (Gu & Dao 2023 Poli et al. 2023), while less demanding in terms of computation and memory, often underperform in long-context scenarios compared to Transformer models. Approximative attention methods, such H2O (Zhang et al. 2023b). StreamingLLMM (Xiao et al.) 2023b). TOVA (Oren et al.) 2024). and FastGen (Ge et al. 2024 often compromise accuracy in cong-context applications. KV cache quantization (Liu et al.
2024 Hooper et al.) 2024). although useful, does not reduce the computation time of the attention mechanism. System-level optimizations, including FlashAttention (Dao et al. 2022 Dao 2023). FeashDe (Hong et al., 2024). and PagedAttention (Kwon et al.) 2023). 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.
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., (024) and Streaming Heads, as shown in Figure Reirieval 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., 2023b), and can operate effectively with a reduced KV cache that includes only recent tokens and attention sinks.
Building the dichotomy of retrieval and streaming heads, we propose DuoAAtention, 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 con-oompeessible retrieval heads using synthetic datasets. Unlike existing methods that rely on attention pattern profiling (Wu et af.) 2024) Ge et al.) 2024| Tang et al. 2024a). DuoAttention directly measures output deviation resulting from token dropping, achieving higher compression rates and improved deployment efficiency. DuoAAttentinn is designed with simplicity and efficiency in mind: each Transformer layer has two KV caches- full KV cache for crucial retrieval heads and 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.55x for MHA and 1.67 for GQA models while speeding up decoding by up to 2 and 1.50x and accelerating pre-filling by up to 1.73x and 1.63 for MHA and GQA models, respectively, with minimal accuracy loss compared to full attention.
Moreover, DuoAltention 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 single A100 GPU, achieving 6.4x capacity increase compared to standard full attention FP16 deployments. DuoAttention paves the way for deploying LLMs in applications requiring million-level context handling.
2

---

Published as a conference paper at ICLR 2025
Retrieval Head Identification:
Full Attention
Deployment:
Full Attention
X 0Z
a
Yes
x
Q
K
Retrieval Head?
x a)
Streaming Attention
Streaming Attention
Figure 2: Overview of DuoAtfention: (1) the retrieval head identification phase, we assign trainable gate value, a, 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 's output, while simultancously applying regularization loss to encourage lower gate values. This training phase is efficient, requiring only the gate values to be Soinnabll-eevinng all other model parameters frozen- thus allowing to be completed within several hours on an GPU node. (2) During deployment, these gate values are binarized to classify heads as either retrieval or streaming based on threshold r. Retrieval heads, identified by 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.
DUOATTENTION
2.1 RETRIEVAL AND STREAMING HEADS
Retrieval Heads In Transformer-bast based LLMs, attention heads exhibit distinct and consistent patterns, reflecting their specialized functionalities (Clark et al., 2019 Xiao et al.) 2023b Wu et al.. 2024). Figure visualizes visualiiss two types of attention heads in the Llama-2-7B-32--122 7B333K- 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 primarily attends to recent tokens and attention sinks (Xiao et al. 2023b), without highlighting earlier relevant tokens in the context. We refer 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 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 I 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 '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 (1), thereby improving the efficiency of processing long contexts.
2.2 OPTIMIZATION -ASSDD IDENTIFICATION OF RETRIEVAL HEADS
Definition of Retrieval Heads Section 2.1 qualitattivyy defines retrieval and streaming heads, but for precise identification, need concrete and quantitative definition. In this paper, define "retrieval heads" as the attention heads that:
significantly alter model outputs when restricted recent lokens and attention sinks.
3

---

Published as a conference paper at ICLR 2025
Trannn Sampl: This very long stary best:
Luma-2-78
Mistmil-18
Lama---8R
I Llama-3-708
Is lat of leng paragrapnna
0.8
eesate the soutt Lim cete petr pan
10
30-
10
20-
Ie lat loeg pareer anm
6.
15
40
eeaator to. the milli any seke mike sile 32 woods -
le lat long affiall
25-
2% 25
60
0.2
parcerys to the ow le? compute distilation los
1ite
10 20
Heads 30
5
Hoads Hoads
Hoads
Figure 3: Example from the synthetic dataset used identify retrieval heads. We embed ten
Figure 4: Optimized gate values of four LLMs. Llama-2-7B
uses MHA with 32 heads per layer, while Mistral and Llama-3 models use GQA with 8 heads per layer. Retrieval heads have higher scores. MHA models have a lower ratio of retrieval heads compared to GQA models.
32-word passkeys within a long text and ask the model to recall these passkeys. Distillation loss is calculated solely on the passkeys.
We use this criterion distinguish retrieval heads from streaming heads. Note that this definition differs from existing works (Ge et al., 2024, Wu et al., 2024, Tang et al.] (2024a) 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- 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-baseed Identification We employ an optimization- approach to identify retrieval heads, drawing inspiration from prior work in CNN filter pruning (Liu et al.) (017), as illustrated in Figure - First, we assign a gate value arj. 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 0z,j E [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 # #layers& x theads and preventing the impact to the model's 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:
attni,j = ai, full Lattnt + (1 - -a.j)) streamingg_tt attn where the attention calculations are defined as:
full_attn softmax QKT T Meausal V,
st streaming_ate attn softmax QKT M Mstramming) V,
where Measal is the causal attention mask (a lower triangular matrix), and M Mueaming represents A-like mask (Han et al. 2023 Xiao et al. 2023b) that attends only to recent and initial tokens.
Synthetic Dataset for Identifying Retrieval Heads
However, relying solely 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 effectively identify which KV heads can be compressed without compromising the model's performance. As depicted in Figure 3 we create a passkey-retrieva 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 last hidden state of the full attention model (H (aa) and those of the model using DuoAttention (H (Hised)), focusing only the last passkey tokens in the entire inputs with 7 tokens, where N is the batch size:
N
T
Ldistill
1
H
(8) Fall
H (8)
mixed [1)
(1)
N
[+1
4

---

Published as conference paper at ICLR 2025
Retrieval Head's KV Cache
Streaming Head's KV Cache
Decoding Token 4
Attention Sink
0
2
3
4
2
3
4
Chunk
Recent Token
Decoding Token 5
0
2
3
4
5
0
3
4
5
Incoming Token
Decoding Token 6
Chunk +
Unatlended
All tokens
Attention sinks Recent tokens
Figure 5: Decoding (left) and Chunked Pre-filling (right) Processes in DuoAttentions (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 DuoAttention's streaming heads on a 16-token sequence, with one attention sink, two recent tokens, and a chunk size DuoAttention's streaming heads have linear time and constant memory complexity during long sequence pre-filling.
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 be more effective than using natural language modeling alone (see ablation studies in Section 13 We use the LI regularization term (a.k.a, Lasso (Tibshirani) 1996)) encourage sparsity in the gate values:
#layera theads
Lreg
Σ
|
The final training loss combination of the distillation loss and the regularization loss, weighted by hyperparameter A, which set as 0.05 in our experiments:
Caisill
X
(3)
Since the total number of trainable parameters only thousands of floating- -oitin numbers, this optimization process is fairly fast, with only 2,000 steps needed. All training experiments in our paper can be conducted on 8xNVIDIA A100 GPU servers.
2.3 DEPLOYING LLMS WITH DUOATTENTION
Binarizing Attention Implementations At inference time, apply full attention exclusively to the designated retrieval heads, identified using the optimized gate values from the training phase (as shown in 4). We binarize the attention policy for each head based threshold т, determined by a specified sparsity quantile, to differentiate between retrieval heads and streaming heads:
if as streaminggatte attn otherwise
full_attn.
attnu
(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 heads within layer, rather than relying on scattering and gathering operations.
Decoding As shown in Figure S 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 constant size. When token 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 2023) pre-fill the KV caches for both retrieval and streaming heads. In long-conted LLMs, chunked pre-filling is common prac- tice (Agrawal et al.) 2023 Kwon et al.) 2023), dividing the prompt into fixed-length chunks to pre-fill the KV cache. This technique significantly reduces peak memory usage (see Table [0) 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-fillle pre-filling in DuoAttention can be achieved with linear time and constant memory complexity, without requiring specialized kemels. As shown in Figure once layer's KVs computed, the streaming head's

---

Published as a conference paper at ICLR 2025
Llama-2-78-- 33-3Kinnttrct (MHA)
Full
StresmingLUM 25%
Fest DuoAttention
Deoumen Depthhhhhh
67 44
62 14
67
62
89
Full Attention
H20.50%
Streaming! 50%
TOWA
Fast Duodttantion
67 44 22
85 64 22
44 22
22
44
55
Figure DuoAltention provides comparable accuracy as full attention on the Veedle-in-a-Haasttar sayyttck benchmark using 25% full attention ratio on the MHA model and 50% full attention ratio on the GQA model.
StreamingLLMM
DucAttention
DuReader
GovReport
HotpotQA
MultiNews
MutiReldQA-ER
Musique
Uama-2-78-322
30
2D-
20
20
0.5
1.0
0.5
20
3.5 Qasper
0.5 Summ
0.5
1.0
0.5 TREC
L.0
0.5
Lo
Ressageetrieeal--a
40
messagehetrieval.
30
22.5
SamSum
75
TriviaOA
20.0
80
0.
0.5
0.5
0.5
30
0.5
0.5
70
0.5
1.0
DuReader
Ro port
HatpetgA
MultiNewws
MuRIFeldOA:FN
Musique
Liama 5-80 1040K
20
25
20
30
20
10
05
1.
0.5
T.
3.5 Qasper
1.0
0.5
1.0
0.5
1.0
0.5 TREC
1.0
0.5 TriviaOA
L.0
Passgeleerreval-ll
massagenettievva
QM4uum
SamSum
75
16
22.5
20.0
0.
0.5
0.5
1.0
Cache Budget
0.
0.5
1.0
0.5
1.0
0.5
1.0
Figure 7: DuoAttention provides better KV budget and accuracy trade-off LongBench benchmarks.
KV cache is immediately pruned to keep only the sink and recent tokens. The next chunk of incoming tokens will only attend to constant number 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 ( to OLL)) and the memory complexity is reduced from ((L) to (K)).
It's important to note that DuoAttention's 's design is well-suited for batch operations, which can further enhance LLM efficiency in serving scenarios with large batch sizes.
EXPERIMENTS
3.1 SETUPS
Models, Datasets, and Baselines We evaluate DuoAttention on both long-context and short-context benchmarks demonstrate that our method preserves model performance on tasks requiring both long and short contexts while significantly improving efficiency, For long-context evaluations, we Haystack
use the Needle-in-az -a-asssia-- (NIAH) benchmark (Kamradi) (202)) and LongBench (Bai al.) 2023). For short-context evaluations, we assess performance MMLU (Hendrycks et al., 2021). MBPP Austin et al. (02T), and MT-Bench (Zheng al.) (002). We employ state-of-thh-arr open- source models, including Llama-2-7B 2-h-hhttt (Touvron et al. (023b) (and its long-context variant 2-B-a2--7B
(Together) 2023)). 3lama-3-(8,70p) BooB--Inntruee (and its long-contens variant Ba-----nttuut--rratttta
1048k D. and Histral--B-vO.. 1-.3v0..2IInsmree Instruct (Jiang al.
2023) We
compare our method against KV cache compression algorithms, including H2O (Zhang et al. 2023b). TOVA (Oren et al. 2024), FastGen (Ge et al. 2024). and StreamingI.LM (Xiao et al. 2023b).
Implementation Details We implement DuoAttention in PyTorch (Paszke et al. 2019) using RoPE (Su al.) (02T) and RMSNorm kernels from FlashInfer (Ye al. (024). For retrieval head identification, we use a batch size of 1, inserting ten 32- 2o-dord passkeys into the BookSum (Kryšcinski jtt al., (02T) dataset. The identification process uses 128 sink tokens and 256 recent tokens. Training samples are drawn from 50 intervals ranging from 1,000 tokens the model-specific maximum length. Passkeys are randomly inserted at 1000 points within the context. Further details are included
httts://hugging
entaa I.lama- 3- 88- Inst rastruccttGren 104m

---

Published as a conference paper at ICLR 2025
Fdi
H2O
ScrommngLLM
TOYA
Duskitonionn
Table 1: Llama-3-70B results on short benchmarks.
MMLU
Bench
Budget MMLU MBPP
MT-B
9.3
Full
100%
79.38%
47.85%
8.93
H2O
50%
79.26%
32.12%
7.16
TOVA
50%
79.15%
36.09%
7.96
0.2
3.0
SLLM
50%
77.46%
5.57%
5,41
2.5
DuoAttn
50%
79.35%
47.09%
9.14
0.2
0.6
1.0
0.4
KV Cache Budget
Figure Results on short henchmarks. in Appendix Scction[A.1 We optimize gate values using the AdamW (Kingma & Ba 2015) 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-CONTERS CONTEXT BENCHMARKS
We evaluate DuoAAtention using the Needle-in-a- Haysstack (NIAH) benchmark and LongBench (Bai and 1oama-3--B-Insttii Gradient- 1048k. We configure DuoAttentioo with a 25% retrieval head ratio for Llama-2-7B-322 Instruct and a 50% ratio for Llama-3-88 -eaee------Inntiee Iee---nstuut-rrrettee Gradient- 1048k. We compare DuoAttention with H2O, TOVA, and StreamingLLMM 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 FlashAttenionn and simulating decoding for the last 50 tokens of the input, following Tang et al. (0024b). FastGen's 's algorithm does not allow for the specification of the KV compression ratio, as fluctuates with inputs. Therefore, we adjust the attention recovery ratio to ensure the KV cache budget is, on average, above 25% or 50% the experiments shown in Figure 6 Additionally, FastGen's '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 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
et al. 2023). We use two long-context models: 2aamaa--7B-322--mmsss--
Needle-in-a- Haystack (NIAH) is 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 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 result, DuoAttention demonstrates strong performance across all sequence depths, handling lengths up 1048K tokens effectively.
LongBench (Bai et al. 2023) is comprehensive suite of long-context datasets encompassing multiple tasks and natural texts, designed 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 50% KV cache budget for GQA, consistent with the results observed in the needle-in-e aa-n--hayae h-yysttcck benchmark. We compare DuoAAtention with FastGen in Table (and in the Appendix provides full results for all 21 LongBench tasks using the 25% and 50% V cache budget for the two models, showing that AuoAntention consistently outperforms baselines across most tasks and achieves the highest average scores.
6 in the Appendix. Table and
3.3
BENCHMARKS.
To ensure that DuoAAtention does not compromise the model's performance on short-context tasks, we evaluate alongside all baselines on three short-context benchmarks: MMI.U, 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 AutAAtention, 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 g and Table [ DuoAttention consistently outperforms all baselines under the same KV cache budget across various models, including Llama-2-7B. Llama- 3-8B, and Llama-3. 3-700-Instten With a 50% KV cache budget, DuoAttention achieves near-lossless performance most benchmarks, demonstrating that it preserves the model's original capabilities.
7

---

Published as a conference paper at ICLR 2025
Full Attention
DueAttention
Llama-2-78 (MHA 25%) Decoding
Llama-3-88 (GQA 50%) Decoding
80
100
Memory (Ci8) Lafoncy
26
32
45
25
S]
30
25
54
40
120K,
200K
100K
300K
600K.
120
60
32
42
26
20
81
91
110
140
70
52
64
76
125
20K
40K
60K
80K
100K
120K
140K
160K
180K
200K
100K
200K
300K
400K
500K
600K
700K
800K
900K
1M
Context Length
Figure 9: Per-token decoding latency and memory usage of DuoAltention compared to full attention across varying context sizes. DuoAltention uses 25% retrieval head ratio for Llama-2-7B (MHA) and 50% for Llama-3-8B (GQA). Duo achieves up to 2.45% memory reduction for MHA and 1.65 for GQA models, along with up 2.13x latency reduction for MHA and 1.5x for GQA models. These reductions approach the inverse of the retrieval head ratios as context length increases. Out-of- of-momory (OOM) results are linearly extrapolated from measured data.
R Full Attention
DuoAttention
Llama-22 1aama-2-7B (MHA 25%) Pre-filling 100K Context
Llama-3-8B (GQA 50%) Pre-filling 320K Context
Latexcy (6)
19
12
18
14
10
14
19
18
10
1910
160
93
10
10
80
10K
20K
30K
10K
50K
60K
70K
8OK
90K.
100K
32K
61K
96K
128K
160K
192K
221K
256K
258K
320K
Menory (GB)
80
51
65
66
67
68
69
80
58
65
19
75
56
350
20165
40
28
30
30
31
52
36
30K
30K
60K
80K.
96K
128K
192K
256K.
233
320K
Pre-filling Chunk Size
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 1O0K tokens, and a 50% ratio for Llama-3-8B (GQA). pre-filling context of 320K tokens. As the pre-filling chunk size decreases, DuoAAtention achieves up to 1.73x latency reduction for MHA and 1.63 x for GQA models, with memory reductions up to 2.38 for MHA and 1.53x for GQA models.
3.4 EFFICIENCY RESULTS
We evaluate DuoAttention's 's decoding latency and memory usage Llama-2-7B and Llama-3-88 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 16. By employing a retrieval head ratio of 25% for Llama-2-7B and 50% for 1-BB---8BB DuoAttention maintains accuracy while significantly improving efficiency.
Decoding Efficiency As shown in Figure 9 DuoAttention'ss 's decoding speed scales linearly, though with 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 II shows DuoAAteetion's speedup and memory savings across various KV budget settings for 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 Figure IT] AtoAAtention achieves maximum improvements on an A100 GPU: 2.55 x memory reduction for MHA and 1.67 for GQA models, and 2x18x latency reduction for MHA and 1.50x 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 1.73x lateney reduction for MHA and 1.63 x for GQA models, with memory reductions of up to 2.38x for MHA and 1.53x for GQA models.
Combiniation with Quantization To fit more tokens into limited memory, we can integrate weight and KV cache quantization with DuoAttention maximize KV cache capacity. Previous studies
8

---

Published as a conference paper at ICLR 2025
Llama Llama-2-70-32K 2-78-32K nstruct
L-ama-3-88-- Graduutt-Gradien-1044434 w/ Context Length 517000
Context Length 134000
4
80
80
#Tokens (million)
DuoAAtention
3.30
Memory (60)
80
2.58
Labiinyy (mc)
Memary (G8)
60
1.67x
50
3
+ 8-bit Weight
80
2.18x
10
1.50x
2
4-bit KV
1884 1
20
20
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.52
Retrioval Head's Ratio
Figure 11: DuoAttention'ss 's decoding memory and lateney 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.55x memory reduction for MHA and 1.67 x for
Figure 12: Combined with 8-bit weight and 4-bit KV cache quantization, DuoAt- tention can accommodate 3.30 million to- kens single A100-80G 80G GPU for the Llama-3-88 model.
GQA models, along with up 2.18x latency reduction for MHA and 1.50x for GQA models.
Different Retrioval Head Idertifiiation Methods
Different Optimization
Different Deployment
Sink and Recent Sizes
Sink and Recent Sizes
1.00
0.75
Aientionns Lang--- Hodsiig
0.50
Fecont 330.0
Rocelt 121
Syttttiii touro
0.0
Aecrnt 64, 256
0.8
Racert
Mlentione lng
0383
0.55
Langgguige Modeiig
0550
Recint 220.
220
0.50
S.t. Receet 120
Raceet
0.45
en Syttneti tourd
Raceet 256
RV Cache Budget
Figure 13: Ablation studies: (1) Comparison of retrieval head identification methads, 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.
have shown that weight quantization (Xiao et al. 2023a Lin et al. 2024) and 4-bit KV cache quantization (Lin* et al.] 2024) Liu et al., [024]] Hooper al.) (024) do not compromise model performance. We combine DuoAttention with the QServe (Lin et al., 2024) 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 A A100-80G GPU using the Llama-3-8B model, resulting in a 6.4x increase in capacity compared to the naive full attention BF16 deployment.
3.5 ABLATION STUDIES
We conduct ablation studies using the Mistral-7B I1B--nsttuut-O..22 on passkey retrieval and MMLU datasets. For the passkey retrieval task, we embed an 8-word passkey within a 30K-word text and perform linear sweep across 100 insertion depths, reporting exact match accuracies.
Optinization-based vs. Attention Profiling-based Retrieval Head Identification We assess optimization-basee method against attention profiling, as used in FastGen (Ge et al. 2024) and RazorAltention (Tang et al.) 2024a), 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 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 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 Contiguration 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.

---

Published as a conference paper at ICLR 2025
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 architec- tures, using approximate attention mechanisms, applying KV cache quantization, and system-level optimizations.
Model Architecture Multi-Query Attention (MQA) (hazeer 2019) and Grouped- Qetrryy-uurry Attention (GQA))Annlii) et al. 2023) reduce the size of the Key- Value (K V) 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 2023) reduce memory usage but tend to underperform on tasks requiring long-context processing.
Approximate Attention
Methods like Sparse Transformer (Child et al. 2019) and Long- Former (Beltagy et al. 2020) use local or block attention patterns to reduce computational complexity. BigBird (Zaheer et al. 2020) 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.) 2023b) and TOVA (Oren et al.
2024) simplify attention by discarding tokens based on query patterns. StreamingI.LMM (Xiao al. 2023b) identifies "attention sinks" and proposes always retaining initial and recent tokens to maintain constant decoding latency and memory usage, allowing the model process significantly more input tokens than the pre-training sequence length. FastGen (Ge et al., 2024) profiles attention heads to discard tokens during decoding. However, our experiments show that these methods degrade the long-contex: abilities of LLMs. Also, methods like H20 and TOVA 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. 2024 Hooper al. 2024 Lin * et al. 2024) reduce the size of KV caches, but they do not address the computational overhead of attention kernels. These methods are complementary to AuoAttention and can be used together to further reduce memory usage.
System Optimizations vLLM (won al. 2023) and FlashAttention (Dao al., 2022 Dao 2023) improve attention computation efficiency by optimizing batch processing and utilizing GPU memory hierarchies. FlashDecoding (Hong al. 2024) and RingAttention (Liu et al. 2023a) 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. (2024)
introduces the concept of retrieval heads to explain LLMs' long-com cong-context capabilities. However, their approach does not compress the KV cache for non-retrieval heads, focusing solely on accuracy. MInference (Tiang al., 2024) accelerates pre-filling for long-context I.LMs by using sparse attention patterns but does not optimize KV cache storage or latency during decoding. RazorAttention (Tang al. 2024a) 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-hant based approach. Also, RazorAttention doesn't optimize pre-fillin. 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.
CONCLUSION
We introduce DuoAttention, a framework that optimizes memory and computational resources in LLMs by distinguishing between Retrieval Heads and Streaming Heads. By applying full KV cache only to retrieval heads, AuoAttention significantly reduces memory usage and latency for both decoding and pre-filling in long-context applications. achieves memory reductions of up to 2.55x for MHA and 1.67 for GQA models, with decoding speed improvements of up to 2.18x for MHA and 1.50x for GQA, and pre-filling accelerations of up to 1.73x and 1.63 respectively, with minimal accuracy loss compared to full attention. When combined with quantization, AuoAttention further boosts KV cache capacity, supporting up to 3.30 million contextual tokens on a single A100 GPU. DuoAttention paves the way for L.LMs handle contexts with millions of tokens.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank MIT-IBM Watson AT 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.
REFERENCES
Griffin Adams, Faisal Ladhak, Hailey Schoelkopf, and Raja Biswas. Cold compress: A toolkit for benchmarking kv cache compression approaches, 8 2024. URL https://www. answer ai/ 03otts/222403- -1l-------o--------------------- compress..e....ce.o html
Amey Agrawal, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav Gulavani, and Ra- machandran Ramjee. Sarathi: Efficient IIm inference by piggybacking decodes with chunked prefills, 2023. URL
16369
Joshua Ainslie, James Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. Program synthesis with large language models. arXiv preprint MXiv:2108.00732. 2021.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020. arXiv:2004.05150.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. GPT-NeoX-20B: An open-source autoregressive language model, 2022. arXiv: 2204.06745.
Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, et al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling. CORR, 2024.
Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng. Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, lon Stoica, and Eric P. Xing. Vicuna: An open-source chatbot impressing gpt-4 with 90%* * chatgpt quality, March 2023. URL https
//
-3-00----00 200vicuua/
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. 2019.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? an analysis of BERT"S s attention. In Tal Linzen, Grzegorz Chrupata, Yonatan Belinkov, and Dieuwke Hupkes (eds.), Proceedings of the 2019 ACL Workshop BlackboxNL.P. Analyzing and Interpreting Neural Networks for NLP, pp. 276-286, Florence, Italy, August 2019, Association for Computational Linguistics. doi: 10.18653/w1/W11-4493 URL httthtttt//ccantiiin
orq kxq/W153 911///99-4488
Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. 2023.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with 10-awareness, 2022. 2it::2205.14333
11

---

Published as a conference paper at ICLR 2025
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, Anirudh Goyal, Anthony Hartshom, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia- Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang. Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Graeme Nail, Gregoire Mialon, Guan Pang. Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Ilyann Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Mísra, Ivan Evtimov, Jade Copet. Jacwon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadcokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong. Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Lauren Rantala- Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin. Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira. Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogaychev, Niladri Chatterji, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng. Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly. Ross Taylor, Ruan Silva, Rui Hou, Rui Wang. Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh. Sean Bell, Scohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarck Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do Vish Vogeti, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang, Xiaoqing Ellen Tan, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos. Aaditya Singh, Aaron Grattafiori, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alex Vaughan, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Franco, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisemman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Changhan Wang. Changkyu Kim, Chao Zhou, Chester Hu, Ching- Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, Danny Wyatt, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang. Due Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Firat Ozgenel, Francesco Caggioni, Francisco Guzmán, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Govind Thattai, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Ibrahim Damlaj. Igor Molybog. Igor Tufanov, Irina- Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Japhet Asher, Dean-Baptiste Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie,
12

---

Published as a conference paper at ICLR 2025
Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Karthik Prasad, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kun Huang, Kunal Chawla, Kushal Lakhotia, Kyle Huang, Lailin Chen, Lakshya Garg. Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Maria Tsimpoukelli, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vattskvv, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang. Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikolay Pavlovich Laptev, Ning Dong, Ning Zhang, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj. Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Rohan Maheswari, Russ Howes, Ruty Rinott, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shiva Shankar, Shugiang Zhang, Shugiang Zhang. Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max. Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad. Sumit Gupta, Sungmin Cho, Sunny Virk, Suraj Subramanian. Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Kohler, Thomas Robinson, Tianhe Li, Tianjun Zhang. Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vítor Albiero, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang. Xiaofang Wang. Xiaojian Wu, Xiaolan Wang, Xide Xia, Xilun Wu, Xinbo Gao, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi. Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yuchen Hao, Yundi Qian, Yuzi He, Zach Rait, Zachary Deiito, 7ef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, and Zhiwei Zhao. The llama 3 herd of models, 2024. URL htt.thttss////rtttt org/ abs 220.//107772283333
Yuan Feng, Junlin Lv, Yukun Cao. Xike Xie, and S. Kevin Zhou Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient 11m inference, 2024. https:/::://axx org/ abs//bb//aa 2407 11550
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive KV cache compression for LLMs. In The Twelfth International Conference on Learning Representations, 2024. URL httth:ts:://// /ttp:s/////peeeeeeeeeee net/ forum?
id=uNrFpDPMyo
Tanya Goyal and Greg Durrett. Evaluating factuality in generation with dependency-level entailment. In Findings of the Association for Computational Linguistics: EMNLP 2020, Online, 2020. Association for Computational Linguistics.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023.
Junxian Guo, Haotian Tang, Shang Yang, Zhekai Zhang, Zhijian Liu, and Song Han. Block Sparse At- tention. https: http:::///tttt/tt/ www.ttu..ooooooo.toooo mi t- 1 ab/Bl BlckoockSppa rse- Attenti on 2024.
Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. LM-Infinite: Simple on-the-fly length generalization for large language models, 2023.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021.
Ke Hong, Guohao Dai, Jiaming Xu, Qiuli Mao, Xiuhong Li, Jun Liu, Kangdi Chen, Yuhan Dong, and Yu Wang. Flashdecoding++: Faster large language model inference gpus, 2024.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization, 2024.
13

---

Published as a conference paper at ICLR 2025
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.066544 2024.
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models, 2023. URL
14509
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anm Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo. Surin Ahn, Zhenhua Han, Amir Abdi, Dongsheng Li, Chin- Lin, Yuqing Yang, and Lili Qiu. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. arXiv preprint arXiv:2407.02490, 2024.
Greg Kamradt. Limtest_ seeeesseeeeeinnhastteekk Doing simple retrieval from Il models at vari- ous context lengths to measure accuracy. hhtpp:: htt///ggttuba hub. com/ gkamradt LLLMeest GeedleInAHaystacce 2024. Accessed: 2024-05-23.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http:
faaaxii.org/bba......... al111122 6980
Wojeiech Kryšciński, Nazneen Rajani, Divyansh Agarwal, Caiming Xiong, and Dragomir Radev. Booksum: A collection of datasets for long-form narrative summarization. 2021.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and lon Stoica. Efficient memory management for large language model serving with pagedattention, 2023.
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before generation. arXiv preprint arXiv:2404. 14469, 2024.
Bin Lin, Yang Ye, Bin Zhu, Jiaxi Cui, Munan Ning. Peng Jin, and Li Yuan. Video-Ilava: Learning united visual representation by alignment before projection, 2023.
Lin, Jiaming Tang, Haotian Tang, Shang Yang. Wei-Ming Chen, Wei-Chen Wang, Guangxuan Xiao, Xingyu Dang, Chuang Gan, and Song Han. Awq: Activation-aware weight quantization for II compression and acceleration, 2024.
Yujun Lin*, Haotian Tang*, Shang Yang*, R Zhekai Zhang, Guangxuan Xiao, Chuang Gan, and Song Han. Qserve: W4 quantization and system co-design for efficient I1m serving. arXiv preprint arXhv:2405.04532. 2024.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near- infinite context, 2023a.
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning, 2023b.
Zhuang Liu, Jianguo Li. Zhiqiang Shen, Gao Huang, Shoumeng Yan, and Changshui Zhang. Learning efficient convolutional networks through network slimming. In ICCV, 2017.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu, Kivi: tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.027500 2024.
OpenAI. Gpt-4 technical report, 2023.
14

---

Published as a conference paper at ICLR 2025
Matanel Oren, Michael Hassid, Yossi Adi, and Roy Schwartz. Transformers multi-state mns,, 2024.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zachary Vito,, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. In NeurlPS, pp. 8024-8035, 2019.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models, 2023. URL
10866
John Schulman, Barret Zoph, Christina Kim, Jacob Hilton, Jacob Menick, Jiayi Weng, Juan Fe- lipe Ceron Uribe, Liam Fedus, Luke Metz, Michael Pokorny, et al. Chatgpt: Optimizing language models for dialogue. OpenAl blog, 2022.
Noam Shazeer. Fast transformer decoding: One write-head is all you need, 2019.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.0980 2021.
Hanlin Tang, Yang Lin, Jing Lin. Qingsen Han, Shikuan Hong. Yiwu Yao, and Gongyi Wang. Razorattention: Efficient kv cache compression through retrieval heads, 2024a. |RLhttps //arxiv. wwg. abs/ 2407 15891
Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Query-aware sparsity for efficient long-context IIm inference, 2024b.
Quest:
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: lab/ stanford attpaa/staffoodaapaaaaaaaaaaaaaa aal_aalpaaaaa 2023.
instruction-folloiing Ilama model.
/tttps://qiet thgg..hb.... com/
R. Tibshirani.
Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society (Series B), 58:267-288 1996.
Together. Llamaa2-7b--2--inntttn - and fine-tuning for llama-2 models with together api, June http:///tetteeeaaaaaall//////aaa 2 7b- 32k inst ruct
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Manie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023b.
Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanistically explains long-context factuality, 2024.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. SmoothQuant: Accurate and efficient post-training quantization for large language models. In Proceedings of the 40th International Conference on Machine Learning, 2023a.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv, 2023b.
Zihao Ye, Ruihang Lai, Roy Lu, Chien- Yu Lin, Size Zheng, Lequn Chen, Tianqi Chen, and Luis Ceze. Cascade inference: Memory bandwidth efficient shared prefix batch decoding. https: ///
flashinfer wtt.:/eer...//2/////// wttp://///..///occcceeeeeeieeeee //flashinfer. ai/ httpsssinerr.../2222////////////
ctto///ccccca inwrraaddiifeeeeeeeee hamml Accessed on 2024-02-
html Jan 2024. URL https:
01.
15

---

Published as a conference paper at ICLR 2025
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, and Amr Ahmed. Big Bird: Transformers for longer sequences. In Proc. of NeurlPS, volume 33, 2020.
Tianyi Zhang, Faisal Ladhak, Esin Durmus, Percy Liang, Kathleen McKeown, and Tatsunori B. Hashimoto. Benchmarking large language models for news summarization, 2023a.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. H2o: Heavy- hitter oracle for efficient generative inference of large language models, 2023b.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric P. Xing. Hao Zhang, Joseph E. Gonzalez, and lon Stoica. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023.
16

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 EXPERIMENTAL DETAILS
We use FSDP2 in Py Torch for model training and DeepSpeed Ulysses (Jacobs et al. 2023) sequence parallelism support long sequences. During training, we use efficient block-sparse approximation of A-like attention for streaming attention, as implemented in Guo et al. (2024) and illustrated in
Figure 14 Maximum sequence lengths vary across models, as detailed Table 2
Table 2: Training Hyperparameters.
Models
Max. Seq. Lengths
-lama-2-7B-chh
4096
Llama-2-78 1aae----------sstiiiiii
32000
Llama-3-SB-Insrruet
8192
32000
Llama-3-70B-Instruct
8192
Distral-BB-Instructta
32000
Local block
Sink block
Streaming Mask (Block Granularity)
Figure 14: Block-sparse approximation of A-like attention.
A.2 FULL LONGBENCH RESULTS
Table 3 and Table show the full LongBench results of DuoAttention and baselines.
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, they rely attention scores to perform token eviction. Since attention scores in FlashAttention are never materialized, these algorithms cannot be used in pre-filling, which one of their main flaws. Therefore, it's possible evaluate these algorithms in long-context settings like needle-in-the-haysted 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 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 only one generated token 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 hachhaystcck and LongBench) generated tokens to perform their token eviction policy long enough, as well as our algorithm. This experimental setting also used by Tang et al. (2024b). Experimental results show method can pass this pressure test, while H20 and TOVA cannot.
17

---

Published as a conference paper at ICLR 2025
Table Full LongBench results with performance with 50% KV cache budget on most datasets.
DuoAAttnntin achieves the best
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
MultiFicldQA-en
52.62
38.53
28.11
44.94
51.44
MultiFicldQA-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
Pu
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
A.4
NIAH RESULTS ON MISTRAL MODELS
14
Figure 15: NIAH result on the Mistral-7B-Insttee v0.2 model.
Figure 16: NIAH result on the -n-ssructtt v0.3 model.
A.5 IMPLEMENTATION OF FASTGEN ON LONG-CONTENT BENCIMARKS
Due to the lack of official implementation of the FastGen Ge et al. (024)) algorithm, we reproduce it using a community codebase (Adams al. (2024)) ), 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 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 0.7, ensuring the average KV cache budget was over 25% of the full KV cache. Similarly. for the Llama-3-88 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 an O(n²) attention map complexity. Therefore, we are unable to test its performance in long contexts. For the long context benchmark, used A 100-80G GPUs, achieving sequence lengths of up to 24k tokens for the Llama-2- 2- model and up to 32k tokens for the Llama-3-88 model. In addition to the seelle-i--the-hasstee benchmark shown in Figure 6 we also evaluated
18

---

Published as a conference paper at ICLR 2025
Table 4: Full LongBench results with LLama-2-7B-instrut--222 DuoAAtention achieves the best performance with 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
MultiFicldQA-en
33.95
21.01
16.69
18.19
25.49
MultiFicldQA-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
PussageRetrieval-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 LongBench using the Llama-3-8B. Instruct- 1eatt-1048 1048K model.
FastGen (>50%)
tuoAttention (50 %
Average
32.82
40.01
2WikiMQA DuReader (zh)
18.61
29.08
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
MultiFicldQA-zhh
42.15
52.40
Musique
13.58
24.65
Passage Count
0.09
0.00
RassageRetrievallenn
93.12
87.00
Passageeetrreval-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
FastGen on LongBench for both models. However, due to the quadratic memory consumption of FastGen, only report results for datasets that were feasible to 8x A100-80G GPUs using FastGen. As shown in Table 5 and T DuoAttention can consistently outperform FastGen on LongBench datasets.
19

---

Published as a conference paper at ICLR 2025
Table 6: Comparison of FastGen and DuoAltention on subset of LongBench using the Llama-2-7B.
32K-In -l-nsstruct model.
FastGen (>25%)
DuoAttention (25%)
Average
19.01
32.81
2WikiMQA
28.05
33.37
MultiNews
12.60
25.03
MultiFielQQA-en
28.58
25.49
MultifielQQA-zh
22.44
39.23
PassageRetrieval-zhh
3.38
40.93
A.6 COMPARISON WITH RECENT KV CACHE COMPRESSION METHODS (SNAPKV, PYRAMIDKV)
(a) SnapKV with Simulation Length
(b) SnapKV with Simulation Length 50
(c) PyramidKV with Simulation Length =0
(d) PyramidKV with Simulation Length
312
232
36
2000
Male
(c) DuoAAtention with Simulation Length = 50
Figure 17: NIAH results for Llama-2. 2-7B-33K--Instrre with a 25% KV cache budget.
SnapKV eL al. 2024h and PyramidKV Caa al. 2024) recent KV cache compression methods that use a local window of observed tokens to determine which KV cache tokens to retain. Both methods rely computing attention scores for the last few tokens (typically 8 8-64) over the entire
20

---

Published as a conference paper at ICLR 2025
(a) SnapKV with Simulation Length 0
(b) SnapKV with Simulation Length 50
(c) PyramidKV with Simulation Length =0
(d) PyramidKV with Simulation Length 50
1An Usmoo Seweeeeeeeeeeeeaaaaeaa 1154 Duckterton Spariiye03
0.00
302
123
(e) AutAttention with Simulation Length = 50
Figure 18: NIAH results for Llama-3-88 B--B--Isttuec-rrraeeeee 1048k with 50% KV cache budget.
21

---

Published as a conference paper at ICLR 2025
context and pruning tokens based on these scores. This approach performs well on benchmarks like Needle-in-at a---astacckk (NIAH) and LongBench, where queries appear at the end of the prompt.
However, these methods assume that critical query information is located at the end of the context, which is not always valid real-world scenarios such as multi-turn dialogues or tasks where queries positioned earlier in the prompt. This reliance reduces their flexibility and general applicabilityy
Figures 17 and 18 compare the performance of SnapKV and PyramidKV with Duo-Attention under and 50% for Llama-3- MB-Instruct-Gradiint -radd-et--eeieeetkkk 1048k). The evaluations include both cases: without simulating the last tokens generated tokens (Simulation Length 0) and with simulation of the last 50 tokens generated inputs (Simulation Length 50, mimicking recond-roumg dialogue scenario). Details of the testing procedure are provided Appendix Section A.3
equivalent KV cache budget constraints (25% for Llama-2-7B-33 2-----22----ae-
As shown, DuoAttention performs comparably or better than SnapKV and PyramidKV when no simulation is applied. However, when the last 50 tokens are treated generated inputs. SnapKV and PyramidKY experience severe accuracy drops, even under large KV cache budgets. This failure occurs because these methods rely on observing the final tokens to guide pruning, which breaks under these conditions. In contrast, DuoAttention maintains robust accuracy under the same stress test.
These results highlight DuoAttention as more general and robust KV cache compression method, capable of adapting to diverse real-world scenarios without relying on assumptions about token positions within the context.
A.7 COMBINATION WITH PRE-FILLING ACCELERATION METHODS (MINFERENCE)
Figure 20: DuoAttention MInference applied to retrieval heads.
Figure 19: MInference applied all attention heads.
MInference (Jiang et al. 2024) employs sparsity patterns, such block-sparse and vertical-slan patterns, observed within token windows to accelerate pre-filling. However, it is limited to the pre-filling stage and does not improve decoding speed or reduce the KV cache size.
We demonstrate that MInference is an orthogonal method that can complement DuoAttention by further accelerating the pre-filling stage of retrieval heads. As shown Figures 19 and 20 applying MInference alone on our NIAH benchmark results in some accuracy degradation compared to full attention or pure DuoAttention (refer to Figure 6).
By combining MInference with DuoAttention, we replace half of the attention heads in LLMS with streaming heads. This approach maintains comparable accuracy while achieving significant reductions in both the KV cache size (nearly halved) and decoding overhead. These results highlight the compatibility and efficiency of combining DuoAttention with MInference.
A.8 RESULTS ON RULER
RULER (Hsich al. 2024) is synthetic dataset designed to rigorously evaluate long-context language models with configurable sequence lengths and task complexities. includes 13 tasks spanning 4 categories, assessing long-context capabilities beyond simple in-context recall.
Table presents the average accuracy of full attention and DuoAttennion (50% sparsity) across different context lengths, using the 1iama-3-8B-Instiie -rrarcettrrdiett--- 1048k model for sequences up to 128K. The results demonstrate that DuoAttention achieves accuracy scores comparable to full attention across all context lengths, with even average performance increase of 0.05%.
22

---

Published as a conference paper at ICLR 2025
Table RULER results comparing full attention and DuoAttention using the 1irma----B-Instruue
Gradient- 1048k model.
Context Length
4K
8K
16K
32K
64K
128K
Avg.
Full Attention
92.78
90.54
86.41
80.59
83.28 76.33
73.01
DuoAttention (50%)
92.83
91.17
85.17
81.28
75.81
73.71
83.33
These findings validate DuoAttention 's effectiveness in maintaining strong accuracy on a rigorous benchmark, even under more challenging long-contess evaluation settings.
A.9 ACCURACY RESULTS WHEN COMBINING WITH QUANTIZATION
Figure 21: Full Attention with INT4 KV Cache
Figure 22: DuoAttention with INT4 KV Cache
We conducted experiments to evaluate the performance of combining DuoAttention with KV quanti- zation. Specifically, we examined two configurations:
1. Baseline: The original model with INT4 KV Pre-Rot Pre-Rope quantization and group size of 128, as proposed in KIVI (Liu et al. 2024) (see Figure 21
2. Proposed Combination: The model incorporating DuoAttention with 50% sparsity along- side the same INT4 KV Pre-Rope quantization (see Figure 22
For this study, we utilized the Llama-3-aa 3-B------BB--- attention model and the iet.Atttttooneaallee model achieve perfect accuracy when using FP16 KV caches (refer to Figure
1048k model. Notably, both the full
The key results are as follows:
Baseline (INT4 KV Pre-Rope Quantization): The model achieves overall accuracy score of 0.867, demonstrating slight accuracy drop compared with using the FP16 KV
cache (Figure 21
DuoAttention + INT4 KV Quantization: The combined approach achieves an overall accuracy score of 0.851, reflecting only a minor reduction 0.016 in performance relative to the INT4 KV baseline (Figure 22
These findings confirm that incorporating DuoAttention (with 50% sparsity) has negligible impact on overall accuracy while offering potential computational advantages. This validates the efficacy of combined approach preserving accuracy while optimizing resource efficiency.
A.10 RESULTS ON THE LLAMA-3.1 MODEL
Table 8 shows the LongBench results the 1amma-3.1-BB -Isttrutt model. The trends consis- tent with Llama-3-8B. 3aama---B-Insttucttee Instruct-Gradienn -eettt-00885 1048k used in the main text, with DuoAttention achieving accuracy comparable full attention and outperforming baselines.
23

---

Published as a conference paper at ICLR 2025
Table 8: LongBench results with Llama-3-8.1 DuoAttention achieves accuracy comparable to full attention and outperforms baselines.
Dataset
Full
H2O (50%)
SLLM (50%)
TOVA (50%)
Duo (50 %)
Average
39.01
35.61
31.32
36.18
38.91
2WikiMQA
16.37
13.91
13.25
14.22
16.20
DuReader (zh)
29.30
21.53
12.95
22.07
31.31
GovReport
34.53
30.56
30.47
30.78
32.87
HotpotQA
17.23
17.31
15.78
16.29
19.53
LCC
52.39
53.08
52.90
52.39
53.31
LSHT (zh)
46.00
39.00
36.00
42.50
45.00
MultiNews
26.91
25.52
24.97
25.14
26.29
MultiFicldQA-en
28.44
21.89
16.05
21.59
27.77
MultiFicldQA-zh
20.19
14.87
15.92
16.55
21.98
Musique
11.82
10.15
10.19
9.64
12.97
NarrativeQA
31.99
31.09
24.15
31.56
29.12
Passage Count
6.26
5.40
4.75
6.68
6.31
PassageRetrieval-en
97.95
89.86
52.11
97.44
98.59
Pu
77.54
69.73
35.14
71.81
75.37
Qasper
25.14
16.96
23.56
20.75
21.12
QMSum
23.63
22.54
21.48
22.82
23.89
RepoBench-P
49.46
49.51
49.95
49.36
53.74
SAMSum
43.69
42.56
43.32
42.28
43.40
TREC
72.50
66.50
69.50
58.00
73.00
TriviaQA
91.65
90.07
90.06
91.73
89.60
VCSUM (zh)
16.26
15.80
15.17
16.09
15.83
A.11 LONGBENCH RESULTS COMPARING WITH SNAPKV, PYRAMIDKY, AND ADAKV
Table presents a detailed comparison of DuoAttention, SnapKV al. (2024). PyramidKV Cai et al. (2024), and AdaK V Feng et al. (2024) under consistent 50% KV cache budget using the 1aama----8--Instuut--rraaaa
(048k model, on LongBench.
DuoAttention achieves the highest average performance (40.21), outperforming SnapK PyramidKV, and AdaKV on most datasets. While SnapKV, PyramidKV, and AdaKV rely on an observation window to determine relevant KV cache entries, making them sensitive to query positioning, Duo does not depend on this heuristic, This allows it to perform robustly across various scenarios, including continuous pre-filling and multi-round dialogue, where queries not always positioned at the end of the context. Furthermore. DuoAttention demonstrates strong generalization across different tasks, maintaining higher accuracy under constrained KV cache budgets. These results highlight its applicability to real-world retrieval and reasoning tasks.
A.12 IMPLEMENTATION DETAILS OF THE NEEDLEE-NN-HHEE CHEIIIN-TTE--AAE HAYSTACK THASS--TAAT-H--AAAACEC BENCHMARK
Our implementation follows the setup of the original Needle-in in-the-Hayee benchmark Kamradt (2024). The haystack corpus is constructed by concatenating Paul Graham's 's essays. The "needle" inserted into this haystack is the text:
"Remember, the best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day."
The corresponding retrieval question is:
"What is the best thing do in San PranciscoAAnswer The best thing to in San Francisco is"
For evaluation, calculate and the expected output. Specifically, let model aidel_ree relr__ee_pp_e resonsee denote the model's response and expected_ answer represent the target output split into individual words, which is:
based the word-level overlap between the model's response
"eat sandwich and sit in Dolores Park sunny day."
24

---

Published as a conference paper at ICLR 2025
Table Comparison of DuoAttention, SnapKV, PyramidKV, and VdaK under a 50% KV cache budget using the aceuracy across tasks.
1044le model. DuoAttention achieves the highest
Dataset
Full
SnapKV
PyramidKV
AdaKV
DuoAttention
Average
40.08
38.47
38.39
38.67
40.21
2WikiMQA
28.78
29.00
28.12
28.97
29.08
DuReader (zh)
30.41
24.04
26.63
22.65
29.31
GovReport
34.23
26.84
27.59
24.22
32.72
HotpotQA
40.37
40.86
41.56
40.23
41.63
LCC
38.19
38.83
37.59
39.67
44.16
LSHT (zh)
38.00
38.00
38.50
36.50
30.00
MultiNews
27.73
22.84
22.93
21.81
27.72
MultiFieldQA-en
52.62
51.96
52.54
52.99
51.44
MultiFieldQA-zh
50.58
50.74
49.85
50.59
52.40
Musique
24.22
24.86
24.63
24.68
24.65
Narrative
26.56
26.63
26.17
27.36
24.54
Passage Count
1.00
1.00
1.00
1.00
0.00
PassageRetrievallen
81.00
80.50
80.00
80.50
87.00
PassageRetrieval-zhh
62.15
58.53
54.56
61.92
62.15
Qasper
29.21
26.00
23.63
27.02
26.93
QMSum
24.52
24.90
24.45
24.65
24.20
RepoBench-P
38.94
38.20
37.48
38.50
46.12
SAMSum
42.51
40.90
40.90
41.38
41.83
TREC
71.50
66.00
70.00
71.00
71.00
TriviaQA
87.70
87.30
87.20
86.80
87.14
VCSUM (zh)
11.37
9.91
10.80
9.62
10.46
The score is computed the ratio of the number of unique words shared between the model's response and the expected answer to the total number of words in the expected answer. Formally, this given by:
score |set(mode1 aec(model_responsee) response) n [agpooeeee)(((xpeeeeeesassss)))) _nsserr)
lexpected_ _answer]
This approach ensures that the evaluation robust to minor variations word order while penalizing the absence of key words from the expected output.
We perform a linear scan over two dimensions: the insertion depth the needle and the context size presented to the model. Insertion depth varies across 10 levels: 0%, 11%, ..., 100% of the corpus length. Context size varies across 13 context sizes as visualized in our paper.
The context provided to the model is formatted as follows:
"<lim_starte This is very long story book: <book> (context) </book>.
Based on the content of the book, Question: cerrivval_question Answerr."
Here, [context] denotes the surrounding text from the haystack corpus, and (retrieval__uestion)) corresponds to the retrieval question.
A.13 EXPERIMENTS ON QUERY POSITIONING
To further evaluate DuoAttention's 's robustness compared to SnapKV and PyramidKV, conducted additional experiments focusing on these methods' dependency on query positioning within the context. Specifically, designed scenario in which the query is not positioned the end of the input context, as SnapKV and PyramidKV typically assume.
In this experiment, the input context constructed follows:
25

---

At the end of the context, only a partial answer prompt was provided: "Answer: The best" to Published as a conference paper at ICLR 2025
(a) Full Attention
(b) DuoAttention w/ 50% KV Budget
(c) SnapKV w/ 50% KV Budget
(d) PyramidKV w/ 50% KV Budget
Figure 23: NIAH results for Llama-3-8B ---B-Instruct- Gradient- 1048k with a 50% KV cache budget. The query of the NIAH benchmark is positioned in the middle of the haystack.
An instruction was placed at the beginning of the input: "This a very long storybook with question embedded. Please answer the embedded question at the end of the book."
The query. "Q:
What is the best thing to do in San Francisco?", was positioned immediately before the needle in the middle of the haystack.
The needle embedded within the haystack: "A: The best thing to do in San Francisco is eat sandwich and sil in Dolores Park on sunny day."
elicit the model's response.
We evaluated SnapKV, PyramidKV, and DuoAttention on the NIAH benchmark using this context. For this experiment, no simulation of the last tokens was applied: the entire input context (instruction, query, haystack, and partial answer) was provided before KV cache compression.
The results this experiment are presented in Figure Each subplot illustrates the performance of a method under a 50% KV cache budget. The results reveal several key insights:
1. SnapKV and Py Failures: Both SnapKV and PyramidKV exhibit significant degradation when the query is not at the end of the context. This highlights their reliance specific assumptions about query locations to guide KV cache pruning. As demonstrated in PyramidKV, even when compressing 32K 128 with Misttal-7B-IIstrruct, both SnapKV and PyramidKV exhibit minimal performance degradation. However, this level of performance is only attainable when the query known and used as observation tokens for pruning. Our updated NIAH results demonstrate that both SnapK and PyramidK V fail when the observation tokens are not the query tokens, even a high retention ratio of 50%.
2. DuoAttention Robustness: DuoAttention achieves accuracy comparable full attention in this scenario, underscoring its robustness and general applicability. Unlike SnapKV V and PyramidKV. DuoAAttention does not rely on the query's 's position, making suitable for real-world tasks where query positions are not fixed or predictable.
These findings reinforce the conclusion that DuoAttention offers a more reliable and versatile approach for KV cache compression, particularly in scenarios with diverse query positions.
26


Reviews:
[{"review_id": "UjEk4zfHPT", "reviewer": "Reviewer_4AyC", "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.", "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?\n\n---\n\n## Post-rebuttal update\n\nThe latest response from the authors is satisfactory for the most part: \n\n* It is a reasonable choice to feature Llama-2-7B-32K-Instruct if the goal is to include a long context-capable MHA model.\n* The added Llama 3.1 results demonstrate that DuoAttention remains performant on one of the most mainstream long context models.\n* I appreciate the authors for recognizing the connection between their optimization-based approach and prior gating network studies.\n* 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.\n* 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.\n\n\nI am improving my rating to 8 if the authors can supply full $\\infty$Bench report in their revision.\n\n(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.)", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is extremely performant. It is rare to see a prefilled compressible token dropping method capable of doing NIAH. \n\n2. The efficiency optimization is solid, and it is supported by a thorough efficiency evaluation.\n\n3. The optimization-based head selection approach offers a different avenue for attention head pattern matching.\n\n4. The paper is nicely written.", "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.\n\n2. 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.\n\n3. 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. \n\n4. 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.\n\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pWFqdu35i9", "reviewer": "Reviewer_2WGg", "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.", "questions": "Please address the questions in the weaknesses.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1.\tIt saves the GPU memory and enhances the inference efficiency for both MHA and GQA.", "weaknesses": "1.\tThere 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.\n2.\tIt can only be applied to multiple KV cache scenarios. For only single KV cache model structure, such as YOCO, it cannot be applied.\n3.\tThe baselines are not comprehensive. It would be beneficial if the authors could compare their results with more advanced baselines.\n4.\tIt 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XjKo6dEayj", "reviewer": "Reviewer_oer2", "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.", "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?\n* 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.\n* Is there any way to extend this method to multi-GPU setting?\n* To what extent can the memory savings lead to increased throughput? (Just a suggestion)", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "* The paper is well-structured and the claims are easy to follow.\n* This work addresses a significant research problem in LLM inference - KV cache optimization. \n* 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.", "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.\n\n* Some state-of-the-art KV sparsification methods are missing in the comparative analysis. For instance, SnapKV, PyramidKV, AdaKV are some notable omissions.\n   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*\n   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*\n   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*", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XwC5VZjmnV", "reviewer": "Reviewer_jov7", "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.", "questions": "1) Is the proportion of Retrieval Heads the same for different models? Why choose 25% Retrieval Heads for MHA and 50% for GQA?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) The paper is well written and well motivated.\n2) DuoAttention is a plug-and-play solution compatible with FlashAttention.\n3) DuoAttention can accelerate inference during both the prefill and decoding stages.", "weaknesses": "1) Lack the experiment with the accuracy results after combining DuoAttention and KV quantization.\n2) Lack the comparisons with new KV Cache compression technologies such as Minference.\n3) Using different compression rates for different heads may lead to uneven computation across cards during parallel inference, potentially affecting performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sqyAWlkIxt", "reviewer": "Reviewer_n7ZA", "summary": "The DuoAttention uses the retrieval heads and local streaming heads, realizing the KV cache compression and acceleration:\n1) the retrieval heads preserve the major information of the long-context, which can not be compressed;\n2) the non-retrieval heads are compressed by locality, which is accelerated with StreamingLLM;\n3) the retrieval heads are detected with importance learnable weights;\n4) Finally, the prefill is accelerated by chunked prefilling, and the decoding is accelerated with compressed KV cache.", "questions": "1. the offline method to find the retrieval heads is more important for practical application\n2. the compression ratio of KV cache should be improved", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. the experiments is enough and confident\n2. the writing and clarity are clear\n3. the concept of retrieval heads is significant for static KV cache compression", "weaknesses": "1. the detection of retrieval heads is based on the fine-tuning of importance weights, the usability should be enhanced\n2. the compression ratio is limited, if the retrieval heads preserve the full context information", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Fw00E5jYii", "reviewer": "Reviewer_JVPb", "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.", "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?\n\n2. In Equation 2, why is the index i set to $L$? Shouldn't it be $N$ instead?\n\n3. 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.\n\n[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.\n\n[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.\n\n[3] Rehg, I. (2024). KV-Compress: Paged KV-Cache Compression with Variable Compression Rates per Attention Head. arXiv preprint arXiv:2410.00161.\n\n[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.\n\nWhile 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.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow, with a novel methodology.\n2. The work is solid, having been evaluated on multiple benchmarks. It also combines the orthogonal technique of quantization, further enhancing its results.\n3. The evaluations demonstrate promising results.", "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.\n\n2. **Inappropriate statements:**\n\na.  _“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\"_ \n\nI 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.\n\nb._“Also, these methods (H2O, TOVA, StreamingLLM) cannot reduce the prefilling cost of long-context LLMs.”_\n\n 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)