DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads
Under review as a conference paper at ICLR 2025

DUOATTENTION: EFFICIENT LONG-CONTEXT LLM
INFERENCE WITH RETRIEVAL AND STREAMING HEADS

Anonymous authors
Paper under double-blind review

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× for MHA and
1.67× for GQA models while speeding up decoding by up to 2.18× and 1.50×
and accelerating pre-filling by up to 1.73× and 1.63× for MHA and GQA models,
respectively, with minimal accuracy loss compared to full attention. Notably,
combined with quantization, DuoAttention enables Llama-3-8B decoding with
3.3 million context length on a single A100 GPU. Code will be released upon
publication.

1

INTRODUCTION

Large language models (LLMs) (Touvron et al., 2023a;b; OpenAI, 2023; Black et al., 2022) are at the
forefront of the AI revolution, powering advanced applications such as multi-round dialogues (Schul-
man et al., 2022; Taori et al., 2023; Chiang et al., 2023), 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 et 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 224×224 image corresponds to 256 tokens (Liu et al.,
2023b), and a three-minute video at 24 FPS generates around 1.1 million tokens.

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

1

000
001

002
003
004
005
006
007

008
009
010
011
012
013

014
015
016
017
018

019
020
021
022
023
024

025
026
027
028
029
030

031
032
033
034
035
036

037
038
039
040
041

042
043
044
045
046
047

048
049
050
051
052
053

Under review as a conference paper at ICLR 2025

054
055

056
057
058
059
060
061

062
063
064
065
066
067

068
069
070
071
072

073
074
075
076
077
078

079
080
081
082
083
084

085
086
087
088
089
090

091
092
093
094
095

096
097
098
099
100
101

102
103
104
105
106
107

Figure 1: Visualization of attention maps in the Llama-2-7B model for the sentence "The best fruit is orange.
What is the best fruit? Orange." shows the distinct roles of retrieval heads (e.g., Layer 15, Head 12) and
streaming heads (e.g., Layer 10, Head 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 a full KV cache. In the middle, streaming heads primarily focus on initial and recent tokens without
emphasizing past contextual relevance. On the right, the impact of limiting attention to the sink and recent tokens
on long-context passkey retrieval accuracy is shown: modifying retrieval heads severely damages performance,
while constraining streaming heads has minimal impacts.

Despite numerous efforts to overcome the challenges of attention mechanisms in long-context
inference, significant computational and memory issues persist. Architectural modifications, such
as Grouped-Query Attention (GQA)(Ainslie et al., 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 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. 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),
FlashDecoding (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., 2024) and Streaming Heads, as shown in Figure 1.
Retrieval Heads, which represent only a fraction of the total, are crucial for processing long contexts
and require full attention across all tokens. In contrast, the majority of attention heads, termed
Streaming Heads, primarily focus on recent tokens and attention sinks (Xiao et al., 2023b), and can
operate effectively with a reduced KV cache that includes only recent tokens and attention sinks.

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

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

2

Retrieval Heads0%20%40%60%80%100%00.10.20.30.40.50.60.7FullStreaming Head FirstRandomRetrieval Head FirstStreaming HeadsAccuracyStreaming Attention RatioStreaming heads focus on attention sinks and recent tokens.Retrieval heads capture relevant tokens in the context.Layer 8 Head 4Layer 8 Head 25Layer 9 Head 8Layer 15 Head 12Layer 10 Head 11Layer 6 Head 5Layer 6 Head 27Layer 12 Head 7Under review as a conference paper at ICLR 2025

108
109

110
111
112
113
114
115

116
117
118
119
120
121

122
123
124
125
126

127
128
129
130
131
132

133
134
135
136
137
138

139
140
141
142
143
144

145
146
147
148
149

150
151
152
153
154
155

156
157
158
159
160
161

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

2 DUOATTENTION

2.1 RETRIEVAL AND STREAMING HEADS

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

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

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

2.2 OPTIMIZATION-BASED IDENTIFICATION OF RETRIEVAL HEADS

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

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

3

KT×Retrieval Head Identification:V×α×(1−α)⊙Q⊕×Deployment:Retrieval Head?YesNo⊙Streaming AttentionFull Attentionα>τα≤τKT×VQ×Streaming AttentionFull AttentionUnder review as a conference paper at ICLR 2025

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.

Figure 3: Example from the synthetic dataset
used to identify retrieval heads. We embed ten
32-word passkeys within a long text and ask
the model to recall these passkeys. Distillation
loss is calculated solely on the passkeys.
We use this criterion to 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-context processing, even when they are not
apparent in attention scores. We support this argument with ablation studies presented in Section 3.5.

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

Our optimization-based identification method directly assesses the impact of compressing the KV
cache with only sink and recent tokens for each KV head. We begin by initializing the gate value
αi,j ∈ [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 × #heads 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:

where the attention calculations are defined as:

attni,j = αi,j · full_attn + (1 − αi,j) · streaming_attn

full_attn = softmax(QKT ⊙ Mcausal)V ,

streaming_attn = softmax(QKT ⊙ Mstreaming)V ,
where Mcausal is the causal attention mask (a lower triangular matrix), and Mstreaming 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 on natural language
modeling objectives is insufficient for identifying retrieval heads because the supervision signal in
natural text that requires inference over long spans is sparse, and most tokens can be inferred using
local context. To address this, we design a synthetic dataset specifically aimed at enhancing the
model’s long-context retrieval capabilities, allowing us to effectively identify which KV heads can be
compressed without compromising the model’s performance. As depicted in Figure 3, we create a
passkey-retrieval dataset by embedding ten randomly generated passkey sequences of s tokens in ten
random locations within a very long context (s = 32 in experiments). The model is then tasked with
recalling these ten sequences at the end of the context.

Training and Loss Functions We optimize the distillation loss, which is the L2 difference between
the last hidden state of the full attention model (Hfull) and those of the model using DuoAttention
(Hmixed), focusing only on the last l passkey tokens in the entire inputs with T tokens, where N is
the batch size:

(H (i)

full[j] − H (i)

mixed[j])2

(1)

Ldistill =

1
N

N
(cid:88)

T
(cid:88)

i=1

j=T −l+1
4

162
163

164
165
166
167
168
169

170
171
172
173
174
175

176
177
178
179
180

181
182
183
184
185
186

187
188
189
190
191
192

193
194
195
196
197
198

199
200
201
202
203

204
205
206
207
208
209

210
211
212
213
214
215

Training SampleThis is a very long story book: … [a lot of long paragraphs…] … Remember this sequence of words, it’s the first passkey to the vault: lima zulu … golf papa  … [a lot of long paragraphs…] … Remember this sequence of words, it’s the tenth passkey to the vault: xray echo … mike kilo … [a lot of long paragraphs…] … Based on the content of the book, what are the passkeys to the vault? First Passkey: lima zulu … golf papa … Tenth Passkey: xray echo … mike kilo32 wordscompute distillation loss32 wordsLlama-2-7BLlama-3-70BLlama-3-8BMistral-7BUnder review as a conference paper at ICLR 2025

216
217

218
219
220
221
222
223

224
225
226
227
228
229

230
231
232
233
234

235
236
237
238
239
240

241
242
243
244
245
246

247
248
249
250
251
252

253
254
255
256
257

258
259
260
261
262
263

264
265
266
267
268
269

Figure 5: Decoding (left) and Chunked Pre-filling (right) Processes in DuoAttention: (1) The retrieval
heads’ KV cache stores all tokens, while the streaming heads’ KV cache retains only recent tokens and attention
sinks, ensuring constant memory usage. (2) The chunked pre-filling process of DuoAttention’s streaming heads
on a 16-token sequence, with one attention sink, two recent tokens, and a chunk size of 4. DuoAttention’s
streaming heads have linear time and constant memory complexity during long sequence pre-filling.
Our synthetic dataset ensures that every supervision signal is relevant to the final compression strategy,
making the process lossless in terms of information retrieval accuracy. It proves to be more effective
than using natural language modeling alone (see ablation studies in Section 13). We use the L1
regularization term (a.k.a, Lasso (Tibshirani, 1996)) to encourage sparsity in the gate values:

Lreg =

#layers
(cid:88)

#heads
(cid:88)

i=1

j=1

|αi,j| .

(2)

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

L = Ldistill + λLreg.

(3)

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

2.3 DEPLOYING LLMS WITH DUOATTENTION

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

(cid:40)

attni,j =

full_attn
if αi,j > τ
streaming_attn otherwise

(4)

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

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

Chunked Pre-filling We use FlashAttention-2 (Dao, 2023) to pre-fill the KV caches for both
retrieval and streaming heads.
In long-context LLMs, chunked pre-filling is a 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 10) by
lowering the peak intermediate activation size in linear layers from sequence length to chunk size.
DuoAttention is fully compatible with chunked pre-filling, and the streaming heads’ pre-filling in
DuoAttention can be achieved with linear time and constant memory complexity, without requiring
specialized kernels. As shown in Figure 5, once a layer’s KVs are computed, the streaming head’s

5

Chunk 1Chunk 2Attention SinkIncoming TokenRecent TokenUnattendedChunk 3Chunk 4Chunk 1Chunk 2Attention SinkIncoming TokenRecent TokenUnattendedChunk 3012345Decoding Token 5Retrieval Head’s KV Cache0345Streaming Head’s KV Cache04560123456Decoding Token 601234Decoding Token 40234All tokensAttention sinks + Recent tokensAttention SinkIncoming TokenRecent TokenUnattendedChunk 1Chunk 2Chunk 3Chunk 4Under review as a conference paper at ICLR 2025

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

Figure 7: DuoAttention provides better KV budget and accuracy trade-off on 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 a constant number of contextual tokens during pre-filling. Let L represent
the sequence length and K the chunk size. The pre-filling time complexity for streaming heads is
optimized from O(L2) to O(LK), and the memory complexity is reduced from O(L) to O(K).

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

3 EXPERIMENTS

3.1 SETUPS

Models, Datasets, and Baselines We evaluate DuoAttention on both long-context and short-context
benchmarks to demonstrate that our method preserves model performance on tasks requiring both
long and short contexts while significantly improving efficiency. For long-context evaluations, we
use the Needle-in-a-Haystack (NIAH) benchmark (Kamradt, 2024) and LongBench (Bai et al.,
2023). For short-context evaluations, we assess performance on MMLU (Hendrycks et al., 2021),
MBPP (Austin et al., 2021), and MT-Bench (Zheng et al., 2023). We employ state-of-the-art open-
source models, including Llama-2-7B-chat (Touvron et al., 2023b) (and its long-context variant
Llama-2-7B-32K-Instruct (Together, 2023)), Llama-3-[8,70]B-Instruct (and its long-context variant
Llama-3-8B-Instruct-Gradient-1048k *), and Mistral-7B-v0.2-Instruct (Jiang et 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 StreamingLLM (Xiao et al., 2023b).

Implementation Details We implement DuoAttention in PyTorch (Paszke et al., 2019) using
RoPE (Su et al., 2021) and RMSNorm kernels from FlashInfer (Ye et al., 2024). For retrieval head
identification, we use a batch size of 1, inserting ten 32-word passkeys into the BookSum (Kry´sci´nski
et al., 2021) dataset. The identification process uses 128 sink tokens and 256 recent tokens. Training
samples are drawn from 50 intervals ranging from 1,000 tokens to the model-specific maximum
length. Passkeys are randomly inserted at 1000 points within the context. Further details are included

*https://huggingface.co/gradientai/Llama-3-8B-Instruct-Gradient-1048k

6

270
271

272
273
274
275
276
277

278
279
280
281
282
283

284
285
286
287
288

289
290
291
292
293
294

295
296
297
298
299
300

301
302
303
304
305
306

307
308
309
310
311

312
313
314
315
316
317

318
319
320
321
322
323

2K7K12K17K22K27K32K022446789Document Depth (%)Full Attention2K7K12K17K22K27K32K022446789H2O 25%2K7K12K17K22K27K32K022446789StreamingLLM 25%2K7K12K17K22K27K32K022446789TOVA 25%3K6K9K12K15K18K21K24K022446789FastGen 25%2K7K12K17K22K27K32K022446789DuoAttention 25%80K242K403K564K726K887K1048K022446789Document Depth (%)Full Attention80K242K403K564K726K887K1048K022446789H2O 50%80K242K403K564K726K887K1048K022446789StreamingLLM 50%80K242K403K564K726K887K1048K022446789TOVA 50%4K8K12K16K20K24K28K32K022446789FastGen 50%80K242K403K564K726K887K1048K022446789DuoAttention 50%Llama-2-7B-32K-Instruct (MHA)Llama-3-8B-Instruct-1048K (GQA)0.51.0152025DuReader0.51.02030GovReport0.51.02040HotpotQA0.51.0152025MultiNews0.51.02030MultiFieldQA-EN0.51.02040MultiFieldQA-ZH0.51.01020Musique0.51.02550PassageRetrieval-EN0.51.02040PassageRetrieval-ZH0.51.02030Qasper0.51.017.520.022.5QMSum0.51.03040SamSum0.51.05075TREC0.51.0708090TriviaQA0.51.0102030DuReader0.51.0253035GovReport0.51.03040HotpotQA0.51.02025MultiNews0.51.0304050MultiFieldQA-EN0.51.0304050MultiFieldQA-ZH0.51.0152025Musique0.51.0255075PassageRetrieval-EN0.51.02550PassageRetrieval-ZH0.51.02030Qasper0.51.020.022.5QMSum0.51.037.540.042.5SamSum0.51.04060TREC0.51.08085TriviaQALlama-2-7B-32KLlama-3-8B-1048KKV Cache BudgetFullH2OStreamingLLMTOVADuoAttentionUnder review as a conference paper at ICLR 2025

Table 1: Llama-3-70B results on short benchmarks.

Budget MMLU MBPP MT-B

Full

100% 79.38% 47.85% 8.93

50% 79.26% 32.12% 7.16
H2O
50% 79.15% 36.09% 7.96
TOVA
50% 77.46% 5.57% 5.41
SLLM
DuoAttn 50% 79.35% 47.09% 9.14

Figure 8: Results on short benchmarks.

in Appendix Section 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-CONTEXT BENCHMARKS

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

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

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

3.3 SHORT-CONTEXT BENCHMARKS.

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

7

324
325

326
327
328
329
330
331

332
333
334
335
336
337

338
339
340
341
342

343
344
345
346
347
348

349
350
351
352
353
354

355
356
357
358
359
360

361
362
363
364
365

366
367
368
369
370
371

372
373
374
375
376
377

0.00.1Llama-2-7BMBPP0.30.4MMLU246MT-Bench0.20.40.60.81.00.00.20.4Llama-3-8B0.20.40.60.81.00.40.60.20.40.60.81.02.55.07.5KV Cache BudgetFullH2OStreamingLLMTOVADuoAttentionUnder review as a conference paper at ICLR 2025

378
379

380
381
382
383
384
385

386
387
388
389
390
391

392
393
394
395
396

397
398
399
400
401
402

403
404
405
406
407
408

409
410
411
412
413
414

415
416
417
418
419

420
421
422
423
424
425

426
427
428
429
430
431

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

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

3.4 EFFICIENCY RESULTS

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

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

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

Combiniation with Quantization To fit more tokens into limited memory, we can integrate weight
and KV cache quantization with DuoAttention to maximize KV cache capacity. Previous studies

8

Memory (GB)0408020K40K60K80K100K120K140K160K180K200K3029272624232119181875696357514539322620Full AttentionDuoAttention050100100K200K300K400K500K600K700K800K900K1M555249444036322824219285776962544639312306012020K40K60K80K100K120K140K160K180K200K393634312926242219171101009181716152423222070140100K200K300K400K500K600K700K800K900K1M76706458524640342721137125112100887664523927Context LengthLatency (ms)OOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMMemory (GB)0102010K20K30K40K50K60K70K80K90K100K1917151414141312111119191819191819191919Full AttentionDuoAttention08016032K64K96K128K160K192K224K256K288K320K15614213512512312311010599931541541551541551520408010K20K30K40K50K60K70K80K90K100K34333231303029282726717069686766656463620408032K64K96K128K160K192K224K256K288K320K70666359565249454238757268656158Llama-2-7B (MHA 25%) Pre-filling 100K ContextLlama-3-8B (GQA 50%) Pre-filling 320K ContextLatency (s)OOMOOMOOMOOMOOMOOMOOMOOMLlama-2-7B (MHA 25%) DecodingLlama-3-8B (GQA 50%) DecodingPre-filling Chunk SizeMemory (GB)0408020K40K60K80K100K120K140K160K180K200K3029272624232119181875696357514539322620Full AttentionDuoAttention050100100K200K300K400K500K600K700K800K900K1M555249444036322824219285776962544639312306012020K40K60K80K100K120K140K160K180K200K393634312926242219171101009181716152423222070140100K200K300K400K500K600K700K800K900K1M76706458524640342721137125112100887664523927Context LengthLatency (ms)OOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMOOMMemory (GB)0102010K20K30K40K50K60K70K80K90K100K1917151414141312111119191819191819191919Full AttentionDuoAttention08016032K64K96K128K160K192K224K256K288K320K15614213512512312311010599931541541551541551520408010K20K30K40K50K60K70K80K90K100K34333231303029282726717069686766656463620408032K64K96K128K160K192K224K256K288K320K70666359565249454238757268656158Llama-2-7B (MHA 25%) Pre-filling 100K ContextLlama-3-8B (GQA 50%) Pre-filling 320K ContextLatency (s)OOMOOMOOMOOMOOMOOMOOMOOMLlama-2-7B (MHA 25%) DecodingLlama-3-8B (GQA 50%) DecodingPre-filling Chunk SizeUnder review as a conference paper at ICLR 2025

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

Figure 12: Combined with 8-bit weight
and 4-bit KV cache quantization, DuoAt-
tention can accommodate 3.30 million to-
kens on a single A100-80G GPU for the
Llama-3-8B model.

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

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., 2024; Hooper et al., 2024) 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 A100-80G GPU using the Llama-3-8B model, resulting in a 6.4×
increase in capacity compared to the naive full attention BF16 deployment.

3.5 ABLATION STUDIES

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

Optimization-based vs. Attention Profiling-based Retrieval Head Identification We assess
our optimization-based method against attention profiling, as used in FastGen (Ge et al., 2024) and
RazorAttention (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 in Figure 13 (1), our
approach of using synthetic data to identify retrieval heads produces significantly better results than
traditional language modeling, which computes loss on all tokens in natural data.

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

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

9

432
433

434
435
436
437
438
439

440
441
442
443
444
445

446
447
448
449
450

451
452
453
454
455
456

457
458
459
460
461
462

463
464
465
466
467
468

469
470
471
472
473

474
475
476
477
478
479

480
481
482
483
484
485

1.67×1.50×2.55×2.18×12343.301.840.52+DuoAttention+ 8-bit Weight 4-bit KV#Tokens (million)0.50.60.70.80.91.00.250.500.751.00Passkey RetrievalDifferent RetrievalHead Identification MethodsAttention ProfilingOptimizationw/ Language ModelingOptimizationw/ Synthetic Data (ours)0.40.50.60.70.80.90.00.51.0Different OptimizationSink and Recent SizesSink, Recent = 0, 320Sink, Recent = 320, 0Sink, Recent = 64, 2560.350.400.450.500.550.600.650.700.80.91.0Different DeploymentSink and Recent SizesSink, Recent = 4, 16Sink, Recent = 16, 64Sink, Recent = 32, 128Sink, Recent = 64, 2560.20.40.60.81.00.450.500.55MMLUAttention ProfilingOptimizationw/ Language ModelingOptimizationw/ Synthetic Data (ours)0.40.50.60.70.80.90.5800.585Sink, Recent = 0, 320Sink, Recent = 320, 0Sink, Recent = 64, 2560.20.30.40.50.60.70.80.91.00.500.55Sink, Recent = 4, 16Sink, Recent = 16, 64Sink, Recent = 32, 128Sink, Recent = 64, 256KV Cache BudgetUnder review as a conference paper at ICLR 2025

486
487

488
489
490
491
492
493

494
495
496
497
498
499

500
501
502
503
504

505
506
507
508
509
510

511
512
513
514
515
516

517
518
519
520
521
522

523
524
525
526
527

528
529
530
531
532
533

534
535
536
537
538
539

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)(Shazeer, 2019) and Grouped-Query Attention
(GQA)(Ainslie et al., 2023) reduce the size of the Key-Value (KV) cache by sharing KV heads across
query heads. However, these methods require pre-training with specific architectures and do not
reduce computational costs. Linear attention Transformers (Gu & Dao, 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. StreamingLLM (Xiao et al., 2023b) identifies "attention sinks" and proposes
always retaining initial and recent tokens to maintain constant decoding latency and memory usage,
allowing the model to process significantly more input tokens than the pre-training sequence length.
FastGen (Ge et al., 2024) profiles attention heads to discard tokens during decoding. However, our
experiments show that these methods degrade the long-context abilities of LLMs. Also, methods like
H2O 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
et 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 DuoAttention and can be used
together to further reduce memory usage.

System Optimizations vLLM (Kwon et al., 2023) and FlashAttention (Dao et al., 2022; Dao,
2023) improve attention computation efficiency by optimizing batch processing and utilizing GPU
memory hierarchies. FlashDecoding (Hong et 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-context capabilities. However, their
approach does not compress the KV cache for non-retrieval heads, focusing solely on accuracy.
MInference (Jiang et al., 2024) accelerates pre-filling for long-context LLMs by using sparse attention
patterns but does not optimize KV cache storage or latency during decoding. RazorAttention (Tang
et al., 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-based
approach. Also, RazorAttention doesn’t optimize pre-filling. DuoAttention offers more effective KV
cache management and higher compression rates, leading to better performance for both pre-filling
and decoding in long-context applications.

5 CONCLUSION

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

10

Under review as a conference paper at ICLR 2025

540
541

542
543
544
545
546
547

548
549
550
551
552
553

554
555
556
557
558

559
560
561
562
563
564

565
566
567
568
569
570

571
572
573
574
575
576

577
578
579
580
581

582
583
584
585
586
587

588
589
590
591
592
593

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/
posts/2024-08-01-cold-compress.html.

Amey Agrawal, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S. Gulavani, and Ra-
machandran Ramjee. Sarathi: Efficient llm inference by piggybacking decodes with chunked
prefills, 2023. URL https://arxiv.org/abs/2308.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 arXiv:2108.07732, 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: A 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, Ion Stoica, and Eric P. Xing. Vicuna: An
open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023. URL https:
//lmsys.org/blog/2023-03-30-vicuna/.

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 attention. In Tal Linzen, Grzegorz Chrupała, Yonatan Belinkov, and
Dieuwke Hupkes (eds.), Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and
Interpreting Neural Networks for NLP, pp. 276–286, Florence, Italy, August 2019. Association
for Computational Linguistics. doi: 10.18653/v1/W19-4828. URL https://aclanthology.
org/W19-4828.

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 IO-awareness, 2022. arXiv:2205.14135.

Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha
Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn,
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-Olano, Diego Perino, Dieuwke Hupkes,
Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip

11

Under review as a conference paper at ICLR 2025

594
595

596
597
598
599
600
601

602
603
604
605
606
607

608
609
610
611
612

613
614
615
616
617
618

619
620
621
622
623
624

625
626
627
628
629
630

631
632
633
634
635

636
637
638
639
640
641

642
643
644
645
646
647

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, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov,
Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, 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-Yeary, 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 Bogoychev, 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, Seohyun 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, Tarek
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 Eisenman, 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-Hsiang 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, Duc 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, Jean-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,
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 Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang,
Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam,

12

Under review as a conference paper at ICLR 2025

648
649

650
651
652
653
654
655

656
657
658
659
660
661

662
663
664
665
666

667
668
669
670
671
672

673
674
675
676
677
678

679
680
681
682
683
684

685
686
687
688
689

690
691
692
693
694
695

696
697
698
699
700
701

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, Shuqiang Zhang, Shuqiang 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 DeVito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, and Zhiwei Zhao. The llama 3 herd
of models, 2024. URL https://arxiv.org/abs/2407.21783.

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 https://openreview.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://github.com/mit-han-lab/Block-Sparse-Attention, 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 on 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 llm inference with
kv cache quantization, 2024.

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.06654, 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 https://arxiv.org/abs/2309.
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-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas
Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023.

13

Under review as a conference paper at ICLR 2025

702
703

704
705
706
707
708
709

710
711
712
713
714
715

716
717
718
719
720

721
722
723
724
725
726

727
728
729
730
731
732

733
734
735
736
737
738

739
740
741
742
743

744
745
746
747
748
749

750
751
752
753
754
755

Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua
Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Minference
1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. arXiv preprint
arXiv:2407.02490, 2024.

Greg Kamradt. Llmtest_needleinahaystack: Doing simple retrieval from llm models at vari-
ous context lengths to measure accuracy. https://github.com/gkamradt/LLMTest_
NeedleInAHaystack, 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:
//arxiv.org/abs/1412.6980.

Wojciech Kry´sci´nski, 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 Ion 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: Llm 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-llava: Learning

united visual representation by alignment before projection, 2023.

Ji 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
llm compression and acceleration, 2024.

Yujun Lin*, Haotian Tang*, Shang Yang*, Zhekai Zhang, Guangxuan Xiao, Chuang Gan, and Song
Han. Qserve: W4a8kv4 quantization and system co-design for efficient llm serving. arXiv preprint
arXiv: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: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint
arXiv:2402.02750, 2024.

OpenAI. Gpt-4 technical report, 2023.

Matanel Oren, Michael Hassid, Yossi Adi, and Roy Schwartz. Transformers are multi-state rnns,

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 DeVito, 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 NeurIPS, 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 https://arxiv.org/abs/2302.10866.

14

Under review as a conference paper at ICLR 2025

756
757

758
759
760
761
762
763

764
765
766
767
768
769

770
771
772
773
774

775
776
777
778
779
780

781
782
783
784
785
786

787
788
789
790
791
792

793
794
795
796
797

798
799
800
801
802
803

804
805
806
807
808
809

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. OpenAI 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.09864, 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. URL https:
//arxiv.org/abs/2407.15891.

Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest:

Query-aware sparsity for efficient long-context llm inference, 2024b.

Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy
Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following llama model.
https://github.com/tatsu-lab/stanford_alpaca, 2023.

R. Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical

Society (Series B), 58:267–288, 1996.

Together. Llama-2-7b-32k-instruct — and fine-tuning for llama-2 models with together api, June

2023. URL https://together.ai/blog/llama-2-7b-32k-instruct.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-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.ai/2024/01/08/cascade-inference.html, Jan 2024. URL https:
//flashinfer.ai/2024/01/08/cascade-inference.html. Accessed on 2024-02-
01.

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 NeurIPS, 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 Ion Stoica.
Judging llm-as-a-judge with mt-bench and chatbot arena, 2023.

15

Under review as a conference paper at ICLR 2025

A APPENDIX

A.1 EXPERIMENTAL DETAILS

We use FSDP2 in PyTorch for model training and DeepSpeed Ulysses (Jacobs et al., 2023) sequence
parallelism to support long sequences. During training, we use an efficient block-sparse approximation
of Λ-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 in Table 2.

Table 2: Training Hyperparameters.

Models

Max. Seq. Lengths

Llama-2-7B-chat
Llama-2-7B-32K-Instruct
Llama-3-8B-Instruct
Llama-3-8B-Instruct-1048K
Llama-3-70B-Instruct
Mistral-7B-Instruct-v0.2

4096
32000
8192
32000
8192
32000

Figure 14: Block-sparse approximation of Λ-like attention.

A.2 FULL LONGBENCH RESULTS

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

16

810
811

812
813
814
815
816
817

818
819
820
821
822
823

824
825
826
827
828

829
830
831
832
833
834

835
836
837
838
839
840

841
842
843
844
845
846

847
848
849
850
851

852
853
854
855
856
857

858
859
860
861
862
863

Streaming Mask (Token Granularity)TopK Block Sparse MaskStreaming Mask (Block Granularity)Dense MaskLocal blockSink blockLocal tokenSink tokenUnder review as a conference paper at ICLR 2025

Table 3: Full LongBench results with Llama-3-8B-Instruct-1048K. DuoAttention achieves the best
performance with a 50% KV cache budget on most datasets.

Dataset

Average

2WikiMQA
DuReader (zh)
GovReport
HotpotQA
LCC
LSHT (zh)
MultiNews
MultiFieldQA-en
MultiFieldQA-zh
Musique
NarrativeQA
Passage Count
PassageRetrieval-en
PassageRetrieval-zh
Qasper
QMSum
RepoBench-P
SAMSum
TREC
TriviaQA
VCSUM (zh)

Full

H2O (50%)

SLLM (50%) TOVA (50%) Duo (50%)

40.08

28.78
30.41
34.23
40.37
38.19
38.00
27.73
52.62
50.58
24.22
26.56
1.00
81.00
62.15
29.21
24.52
38.94
42.51
71.50
87.70
11.37

35.76

27.99
24.94
29.44
36.77
43.09
25.00
25.52
38.53
38.25
19.24
25.13
2.05
74.75
52.57
20.65
22.87
39.98
40.78
64.00
85.98
13.45

32.26

29.22
9.41
29.08
39.27
41.94
25.50
24.85
28.11
31.07
20.47
22.06
1.64
49.00
38.90
21.77
22.11
37.60
40.25
67.00
86.11
12.10

35.55

26.93
27.00
30.10
38.45
42.31
24.50
26.32
44.94
40.82
23.07
25.64
1.00
72.00
46.13
23.06
23.16
40.14
40.50
54.00
84.97
11.59

40.21

29.08
29.31
32.72
41.63
44.16
30.00
27.72
51.44
52.40
24.65
24.54
0.00
87.00
62.15
26.93
24.20
46.12
41.83
71.00
87.14
10.46

A.4 NIAH RESULTS ON MISTRAL MODELS

Figure 15: NIAH result on the Mistral-7B-Instruct-
v0.2 model.

Figure 16: NIAH result on the Mistral-7B-Instruct-
v0.3 model.

A.5

IMPLEMENTATION OF FASTGEN ON LONG-CONTEXT BENCHMARKS

Due to the lack of official implementation of the FastGen (Ge et al. (2024)) algorithm, we reproduce
it using a community codebase (Adams et 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 as outlined in the FastGen paper. To quantify sparsity, we
calculated the average KV cache usage across all test cases as the overall measure of sparsity. For the
Llama-2-7B model, we set the recovery ratio to 0.7, ensuring the average KV cache budget was over
25% of the full KV cache. Similarly, for the Llama-3-8B model, we set the recovery ratio to 0.87,
ensuring the average KV cache budget was more than 50% of the full KV cache. Additionally, since
FastGen uses the full attention map of the user-provided prompt to profile the types of different heads,
it results in an O(n2) attention map complexity. Therefore, we are unable to test its performance in
long contexts. For the long context benchmark, we used 8 A100-80G GPUs, achieving sequence
lengths of up to 24k tokens for the Llama-2-7B model and up to 32k tokens for the Llama-3-8B
model. In addition to the needle-in-the-haystack benchmark shown in Figure 6, we also evaluated

17

864
865

866
867
868
869
870
871

872
873
874
875
876
877

878
879
880
881
882

883
884
885
886
887
888

889
890
891
892
893
894

895
896
897
898
899
900

901
902
903
904
905

906
907
908
909
910
911

912
913
914
915
916
917

Under review as a conference paper at ICLR 2025

Table 4: Full LongBench results with Llama-2-7B-Instruct-32K. DuoAttention achieves the best
performance with a 25% KV cache budget on most datasets.

Dataset

Average

2WikiMQA
DuReader (zh)
GovReport
HotpotQA
LCC
LSHT (zh)
MultiNews
MultiFieldQA-en
MultiFieldQA-zh
Musique
NarrativeQA
Passage Count
PassageRetrieval-en
PassageRetrieval-zh
Qasper
QMSum
RepoBench-P
SAMSum
TREC
TriviaQA
VCSUM (zh)

Full

H2O (25%)

SLLM (25%) TOVA (25%) Duo (25%)

37.52

35.59
25.10
31.23
47.98
51.21
34.50
27.11
33.95
45.79
22.97
24.11
0.00
50.92
37.68
33.23
20.79
51.58
42.10
71.50
86.21
14.45

26.84

28.87
15.56
20.66
39.60
45.78
16.50
19.21
21.01
19.81
20.63
19.14
0.53
19.50
11.75
16.84
18.89
45.16
39.73
48.50
85.16
10.71

27.80

29.69
13.96
24.14
40.39
44.25
17.50
20.54
16.69
22.50
20.09
21.13
0.58
19.08
16.77
17.68
20.05
45.25
37.43
56.50
85.24
14.36

29.78

31.18
15.51
22.88
47.45
47.91
18.50
21.41
18.19
24.96
21.00
23.06
0.00
30.17
32.38
20.85
20.16
49.03
36.17
47.00
85.65
11.85

34.49

33.37
23.99
27.98
50.44
48.34
25.50
25.03
25.49
39.23
19.27
20.49
0.33
47.25
40.93
26.59
21.48
48.58
33.10
68.50
86.15
12.35

Table 5: Comparison of FastGen and DuoAttention on a subset of LongBench using the Llama-3-8B-
Instruct-1048K model.

FastGen (>50%) DuoAttention (50%)

Average

2WikiMQA
DuReader (zh)
HotpotQA
LCC
MultiNews
MultiFieldQA-en
MultiFieldQA-zh
Musique
Passage Count
PassageRetrieval-en
PassageRetrieval-zh
Qasper
QMSum
SAMSum
TriviaQA
VCSUM (zh)

32.82

18.61
20.22
33.08
46.50
18.18
44.05
42.15
13.58
0.09
93.12
40.75
26.51
24.03
34.12
69.92
0.23

40.01

29.08
29.31
41.63
44.16
27.72
51.44
52.40
24.65
0.00
87.00
62.15
26.93
24.20
41.83
87.14
10.46

FastGen on LongBench for both models. However, due to the quadratic memory consumption of
FastGen, we only report results for datasets that were feasible to run on 8x A100-80G GPUs using
FastGen. As shown in Table 5 and Table 6, DuoAttention can consistently outperform FastGen on
LongBench datasets.

18

918
919

920
921
922
923
924
925

926
927
928
929
930
931

932
933
934
935
936

937
938
939
940
941
942

943
944
945
946
947
948

949
950
951
952
953
954

955
956
957
958
959

960
961
962
963
964
965

966
967
968
969
970
971

Under review as a conference paper at ICLR 2025

Table 6: Comparison of FastGen and DuoAttention on a subset of LongBench using the Llama-2-7B-
32K-Instruct model.

FastGen (>25%) DuoAttention (25%)

Average

2WikiMQA
MultiNews
MultiFieldQA-en
MultiFieldQA-zh
PassageRetrieval-zh

19.01

28.05
12.60
28.58
22.44
3.38

32.81

33.37
25.03
25.49
39.23
40.93

A.6 COMPARISON WITH RECENT KV CACHE COMPRESSION METHODS (SNAPKV,

PYRAMIDKV)

(a) SnapKV with Simulation Length = 0

(b) SnapKV with Simulation Length = 50

(c) PyramidKV with Simulation Length = 0

(d) PyramidKV with Simulation Length = 50

(e) DuoAttention with Simulation Length = 50

Figure 17: NIAH results for Llama-2-7B-32K-Instruct with a 25% KV cache budget.

SnapKV (Li et al., 2024) and PyramidKV (Cai et al., 2024) are recent KV cache compression methods
that use a local window of observed tokens to determine which KV cache tokens to retain. Both
methods rely on computing attention scores for the last few tokens (typically 8–64) over the entire

19

972
973

974
975
976
977
978
979

980
981
982
983
984
985

986
987
988
989
990

991
992
993
994
995
996

997
998
999
1000
1001
1002

1003
1004
1005
1006
1007
1008

1009
1010
1011
1012
1013

1014
1015
1016
1017
1018
1019

1020
1021
1022
1023
1024
1025

Under review as a conference paper at ICLR 2025

1026
1027

1028
1029
1030
1031
1032
1033

1034
1035
1036
1037
1038
1039

1040
1041
1042
1043
1044

1045
1046
1047
1048
1049
1050

1051
1052
1053
1054
1055
1056

1057
1058
1059
1060
1061
1062

1063
1064
1065
1066
1067

1068
1069
1070
1071
1072
1073

1074
1075
1076
1077
1078
1079

(a) SnapKV with Simulation Length = 0

(b) SnapKV with Simulation Length = 50

(c) PyramidKV with Simulation Length = 0

(d) PyramidKV with Simulation Length = 50

(e) DuoAttention with Simulation Length = 50

Figure 18: NIAH results for Llama-3-8B-Instruct-Gradient-1048k with a 50% KV cache budget.

20

Under review as a conference paper at ICLR 2025

context and pruning tokens based on these scores. This approach performs well on benchmarks like
Needle-in-a-Haystack (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 in real-world scenarios such as multi-turn dialogues or tasks where queries
are positioned earlier in the prompt. This reliance reduces their flexibility and general applicability.

Figures 17 and 18 compare the performance of SnapKV and PyramidKV with DuoAttention under
equivalent KV cache budget constraints (25% for Llama-2-7B-32K-Instruct and 50% for Llama-3-
8B-Instruct-Gradient-1048k). The evaluations include both cases: without simulating the last tokens
as generated tokens (Simulation Length = 0) and with simulation of the last 50 tokens as generated
inputs (Simulation Length = 50, mimicking a second-round dialogue scenario). Details of the testing
procedure are provided in Appendix Section A.3.

As shown, DuoAttention performs comparably or better than SnapKV and PyramidKV when no
simulation is applied. However, when the last 50 tokens are treated as generated inputs, SnapKV
and PyramidKV 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 a 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 19: MInference applied to all attention heads.

Figure 20: DuoAttention + MInference applied to
retrieval heads.

MInference (Jiang et al., 2024) employs sparsity patterns, such as block-sparse and vertical-slash
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 in 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 (Hsieh et al., 2024) is a synthetic dataset designed to rigorously evaluate long-context
language models with configurable sequence lengths and task complexities. It includes 13 tasks
spanning 4 categories, assessing long-context capabilities beyond simple in-context recall.

Table 7 presents the average accuracy of full attention and DuoAttention (50% sparsity) across
different context lengths, using the Llama-3-8B-Instruct-Gradient-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 an average performance increase of 0.05%.

21

1080
1081

1082
1083
1084
1085
1086
1087

1088
1089
1090
1091
1092
1093

1094
1095
1096
1097
1098

1099
1100
1101
1102
1103
1104

1105
1106
1107
1108
1109
1110

1111
1112
1113
1114
1115
1116

1117
1118
1119
1120
1121

1122
1123
1124
1125
1126
1127

1128
1129
1130
1131
1132
1133

Under review as a conference paper at ICLR 2025

Table 7: RULER results comparing full attention and DuoAttention using the Llama-3-8B-Instruct-
Gradient-1048k model.

Context Length

4K

8K

16K

32K

64K 128K Avg.

Full Attention
DuoAttention (50%)

92.78
92.83

90.54
91.17

86.41
85.17

80.59
81.28

76.33
75.81

73.01
73.71

83.28
83.33

These findings validate DuoAttention ’s effectiveness in maintaining strong accuracy on a rigorous
benchmark, even under more challenging long-context 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-Rope quantization and a 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-8B-Instruct-Gradient-1048k model. Notably, both the full
attention model and the DuoAttention-enabled model achieve perfect accuracy when using FP16 KV
caches (refer to Figure 6).

The key results are as follows:

• Baseline (INT4 KV Pre-Rope Quantization): The model achieves an overall accuracy
score of 0.867, demonstrating a 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 of 0.016 in performance relative
to the INT4 KV baseline (Figure 22).

These findings confirm that incorporating DuoAttention (with 50% sparsity) has a negligible impact
on overall accuracy while offering potential computational advantages. This validates the efficacy of
the combined approach in preserving accuracy while optimizing resource efficiency.

A.10

A.11

IMPLEMENTATION DETAILS OF THE NEEDLE-IN-THE-HAYSTACK BENCHMARK

Our implementation follows the setup of the original Needle-in-the-Haystack benchmark Kamradt
(2024). The haystack corpus is constructed by concatenating Paul Graham’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."

22

1134
1135

1136
1137
1138
1139
1140
1141

1142
1143
1144
1145
1146
1147

1148
1149
1150
1151
1152

1153
1154
1155
1156
1157
1158

1159
1160
1161
1162
1163
1164

1165
1166
1167
1168
1169
1170

1171
1172
1173
1174
1175

1176
1177
1178
1179
1180
1181

1182
1183
1184
1185
1186
1187

Under review as a conference paper at ICLR 2025

The corresponding retrieval question is:

"What is the best thing to do in San Francisco?Answer: The best thing to do in San Francisco is"
For evaluation, we calculate a score based on the word-level overlap between the model’s response
and the expected output. Specifically, let model_response denote the model’s response and
expected_answer represent the target output split into individual words, which is:

"eat a sandwich and sit in Dolores Park on a sunny day."
The score is computed as 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
is given by:

score =

|set(model_response) ∩ set(expected_answer)|
|expected_answer|

This approach ensures that the evaluation is robust to minor variations in word order while penalizing
the absence of key words from the expected output.

We perform a linear scan over two dimensions: the insertion depth of 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:

"<|im_start|> This is a very long story book: <book> {context} </book>.
Based on the content of the book, Question: {retrieval_question}Answer:"
Here, {context} denotes the surrounding text from the haystack corpus, and {retrieval_question}
corresponds to the retrieval question.

A.12 EXPERIMENTS ON QUERY POSITIONING

(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-Instruct-Gradient-1048k with a 50% KV cache budget. The query of
the NIAH benchmark is positioned in the middle of the haystack.

To further evaluate DuoAttention’s robustness compared to SnapKV and PyramidKV, we conducted
additional experiments focusing on these methods’ dependency on query positioning within the
context. Specifically, we designed a scenario in which the query is not positioned at the end of the
input context, as SnapKV and PyramidKV typically assume.

In this experiment, the input context was constructed as follows:

23

1188
1189

1190
1191
1192
1193
1194
1195

1196
1197
1198
1199
1200
1201

1202
1203
1204
1205
1206

1207
1208
1209
1210
1211
1212

1213
1214
1215
1216
1217
1218

1219
1220
1221
1222
1223
1224

1225
1226
1227
1228
1229

1230
1231
1232
1233
1234
1235

1236
1237
1238
1239
1240
1241

Under review as a conference paper at ICLR 2025

• An instruction was placed at the beginning of the input: "This is a very long storybook with

a 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 was embedded within the haystack: "A: The best thing to do in San Francisco is

eat a sandwich and sit in Dolores Park on a sunny day."

• At the end of the context, only a partial answer prompt was provided: "Answer: The best" to

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 of this experiment are presented in Figure 23. Each subplot illustrates the performance of
a method under a 50% KV cache budget. The results reveal several key insights:

1. SnapKV and PyramidKV Failures: Both SnapKV and PyramidKV exhibit significant
degradation when the query is not at the end of the context. This highlights their reliance on
specific assumptions about query locations to guide KV cache pruning. As demonstrated in
PyramidKV, even when compressing 32K to 128 with Mistral-7B-Instruct, both SnapKV and
PyramidKV exhibit minimal performance degradation. However, this level of performance
is only attainable when the query is known and used as observation tokens for pruning.
Our updated NIAH results demonstrate that both SnapKV and PyramidKV fail when the
observation tokens are not the query tokens, even at a high retention ratio of 50%.

2. DuoAttention Robustness: DuoAttention achieves accuracy comparable to full attention
in this scenario, underscoring its robustness and general applicability. Unlike SnapKV and
PyramidKV, DuoAttention does not rely on the query’s position, making it 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.

1242
1243

1244
1245
1246
1247
1248
1249

1250
1251
1252
1253
1254
1255

1256
1257
1258
1259
1260

1261
1262
1263
1264
1265
1266

1267
1268
1269
1270
1271
1272

1273
1274
1275
1276
1277
1278

1279
1280
1281
1282
1283

1284
1285
1286
1287
1288
1289

1290
1291
1292
1293
1294
1295

24

