MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention
Authors: Huiqiang Jiang
YUCHENG LI
Chengruidong Zhang
Qianhui Wu
Xufang Luo
Surin Ahn
Zhenhua Han
Amir H. Abdi
Dongsheng Li
Chin-Yew Lin
Yuqing Yang
Lili Qiu
Venue: NeurIPS 2024 spotlight (2024)
Keywords: LLMs Inference
Long-Context LLMs
Dynamic Sparse Attention
Efficient Inference
Primary area: natural_language_processing

Abstract: The computational challenges of Large Language Model (LLM) inference remain a significant barrier to their widespread deployment, especially as prompt lengths continue to increase. Due to the quadratic complexity of the attention computation, it takes 30 minutes for an 8B LLM to process a prompt of 1M tokens (i.e., the pre-filling stage) on a single A100 GPU. Existing methods for speeding up prefilling often fail to maintain acceptable accuracy or efficiency when applied to long-context LLMs. To address this gap, we introduce MInference (Milliontokens Inference), a sparse calculation method designed to accelerate pre-filling of long-sequence processing. Specifically, we identify three unique patterns in long-context attention matrices-the A-shape, Vertical-Slash, and Block-Sparse-that can be leveraged for efficient sparse computation on GPUs. We determine the optimal pattern for each attention head offline and dynamically build sparse
indices based on the assigned pattern during inference. With the pattern and sparse indices, we perform efficient sparse attention calculations via our optimized GPU kernels to significantly reduce the latency in the pre-filling stage of longcontext LLMs. Our proposed technique can be directly applied to existing LLMs without any modifications to the pre-training setup or additional fine-tuning. By
evaluating on a wide range of downstream tasks, including InfiniteBench, RULER, PG-19, and Needle In A Haystack, and models including LLaMA-3-1M, GLM-4-1M, Yi-200K, Phi-3-128K, and Qwen2-128K, we demonstrate that MInference effectively reduces inference latency by up to 10x for pre-filling on an A100, while maintaining accuracy. Our code is available at https://aka.ms/MInference.

MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention
Huiqiang Jiang, Yucheng Li Chengruidong Zhangi Qianhui Wu, Xufang Luo,
Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin- Yew Lin, Yuqing Yang, Lili Qiu
Microsoft Corporation, University of Surrey (hjiang, chengghannggggg yuqyang baggangg/miicooo @ewicooooooo com, yucheng- wwwpeeng...iuuun li@surrey.aaare. ac.
Abstract
The computational challenges of Large Language Model (LLM) inference remain significant barrier their widespread deployment, especially as prompt lengths continue to increase. Due to the quadratic complexity of the attention computation, it takes 30 minutes for an 8B LLM to process a prompt of 1M tokens (i.e., the pre-filling stage) on single A100 GPU. Existing methods for speeding up pre- filling often fail to maintain acceptable accuracy or efficiency when applied long-context LLMs. To address this gap, we introduce MInference (Million- tokens Inference), a sparse calculation method designed accelerate pre-filling of long-sequence processing. Specifically, we identify three unique patterns in long-context attention matrices- the A A-shape. Vertical-Slaste Slash, and Slock-Sparse- that can be leveraged for efficient sparse computation GPUs. We determine the optimal pattern for each attention head offline and dynamically build sparse indices based the assigned pattern during inference. With the pattern and sparse indices, we perform efficient sparse attention calculations via our optimized GPU kernels significantly reduce the latency in the pre-filling stage of long- context LLMs. Our proposed technique can be directly applied to existing LLMs without any modifications the pre-training setup additional fine-tuning. By evaluating on wide range of downstream tasks, including InfiniteBench, RULER, PG-19, and Needle In A Haystack, and models including LLaMA-3-1M. GLM- 4-IM, Yi-200K, Phi-3-128K, and Qwen2- 1aimn---2888 128K, we demonstrate that MInference effectively reduces inference latency by up 10x for pre-filling on an A100, while maintaining accuracy, Our code is available at wtlhttp::// //aps: ms/ msIMInferene
erence
Needle Haystack LLaMA-3-88-11
Context
Depth Percert (%)
22-
-0.8
Minforence
Latencylminn
20
10x
47-
-0.4
6.8x
100
0.0
Tor
7200
mine
10k
300k
500k
Context Windows Size
Context Length
(a) Needle In A Haystack
(b) Latency Specdup
Figure I: Attention weights, especially in long-context I.LMs, exhibit up to 96.8% sparsity in contexts of 128K.
We propose MInference, leveraging dynamic sparse attention accelerate the pre-filling stage of long-context LLM inference. achieves up speedup for IM contexts on a single A100, as shown (b), and matches or surpasses baselines, as demonstrated by Needle A Haystack [Km223 in (a) on LLaMA-3-8B-IMM GGa24]
Equal contribution. Work during internship at Microsoft.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Introduction
Large language models (LLMs) have entered the era of long-context processing, with some of them
supporting context windows ranging from 128K to 10M tokens [Gra24 RST 24 LYZA24 YCL*244
AJA DA24]. These extended context windows enable LLMs to unlock a multitude of complex
real-world applications, such as repository-level code understanding [BSK 24, JYW 23 P3C
long-document guestion-answering [CPG-23 L2), self-play reasoning (Ope24), extreme-label
in-context learning [LZD+24] and long-horizon agent tasks [Wen23).
However, due to the quadratic complexity of attention, it can take several minutes for the model to process the input prompt (i.e., the pre-filling stage) and then start to produce the first token, which leads to unacceptable Time To First Token experience, thus greatly hinders the wide application of long-context LLMs. As shown in Fig. 2a when serving LLaMA-3-8D on single A100 machine, the model would keep users waiting for 3 minutes to finish the pre-filling stage given a prompt of 300K tokens, and this number increases to 30 minutes for prompt of 1M tokens. The overhead of self-attention computation exceeds 90% of the total pre-filling latency, which makes the major bottleneck in long-context processing of LLMs. Previous research has shown that the attention matrices are highly sparse (1-0C+22 DSY24J, which has led to the development of fixed sparse attention methods such Longformer (BPC20 and BigBird [ZGD+ 20 However, prior studies have also noted that attention distributions vary significantly across different inputs [LCW2] LQC+22] This dynamic nature prevents prior sparse methods from being used directly on long-context LLMs without expensive training or fine-tuning. But if the dynamic sparse attention patterns could be efficiently predicted online, the pre-filling latency of long-contes, LLMs could be significantly reduced by calculating only the most important part of the attention weights.
Building upon this idea, present MInference, technique that reduces 95% of FLOPs in the attention computation to significantly accelerate the pre-filling stage of long-context LLM inference via dynamic sparse attention. Unlike existing dynamic sparse attention methods that introduce large computational overhead to estimate attention patterns with low-rank hidden dimensions (LQQC RCHG 24|, our method is designed specifically for long-context scenarios with minimal overhead in estimation. Specifically, we conduct extensive analysis and identify three general patterns of sparse attention in long-context LLMs: A-shape pattern, Vertical-Slash pattern, and Block-Sparse pattern. Based on these findings, we introduce a kernel-aware search method to assign the optimal attention pattern for each head. Importantly, instead of fixed attention masks in prior studies, we perform an efficient online approximation to build dynamic sparse mask for each head according to their assigned pattern and particular inputs. For example, to build a dynamic sparse mask for a specific prompt on one Vertical-Slash head, we use partial of attention weight consisting of the last last_q query and key vectors (i.e. Q ( and K) to estimate the most important indices of the vertical and slash lines globally on the attention matrix. For Block-Sparse heads, we perform mean pooling on both query and key vectors in blocks of 64 and calculate the block-level attention weights to determine the most important blocks and thereby obtain block-sparse dynamic mask. After obtaining the dynamic sparse mask, three optimized GPU kernels are used, which we developed for the above three sparse patterns. These kernels are based on the dynamic sparse compilers PIT 23), Triton [KKC1 and FlashAttention Dao24], which enable extremely efficient computation of dynamic sparse attention.
Extensive experiments are conducted on various Long-context LLMs, including ILLMM-3-8B- IM [Gra24], GLM-4-9B-1M [GZX+24], and Yi-9B-200K [YCL++24) across benchmarks with context lengths over 1M tokens, such as InfiniteBench 24], RULER (HSK+24), Needle In A Haystack [Kmm2]], and PG-19 [RPJ Needle In A Haystack was also tested on Phi-3-Mini- 128K AJA 24) and Qwen-2-7B-128K (BBC*23). Results show that MInference speeds up the pre-filling stage by up to 10x for IM contexts with LLaMA-3-8B on single A100, reducing latency from 30 minutes to minutes per prompt, while maintaining or improving accuracy.
2
Attention Heads: Dynamic, Sparse, and Characteristic
2.1 Attention is Dynamically Sparse
The sparsity of attention weights in pre-trained L.Mss, especially in long-context scenarios, has been well-documented LQC+22 22 RCHG+24 LWD+23 XTC+24 As shown in Fig. 2b for an attention

---

Attention Recall of Top-Kik44096)
96.4%
Attention Recall Top-Kik-4006
81.7%
Total
Latencylming
20
0
16-
12-
0.6
166
12-
0.6
10-
20
0.4
20
24-
-.2
-.2
28-
1-
10k
300K
500k
1M
12 16 20
24
24
32
0.0
8 12
24 28 32
Context Windows Size
Head
Head
(a) Attention incurs heavy cost.
(b) Attention is sparse.
(c) Sparsity of attention is dynamic.
Figure 2: (a) Latency breakdown of the pre-filling stage. (b) How much attention scores top-k (k=4096) columns cover 128k context. (c) Less attention scores are retrieved when reusing the top-k indices from another examples, indicating its dynamic nature. Visualizations are based on LLaMa-3-8B with a single A100.
matrix of size 128k x 128k; retaining only the top 4k columns recalls 96.8% of the total attention. In other words, each token is attending to limit number of tokens despite the long sequence it is processing.
On the other hand, although the sparse nature of attention matrices shared across different inputs, the exact distributions of sparse pattern highly dynamic. That is to say, a token given position only attends to a subset of the sequence in self-attention, and the exact tokens it attends to are highly context-dependent and vary significantly across different prompts. This dynamism has been mathematically demonstrated prior studies [LCW21 LCW23], As depicted in Fig. 2 if we take the top 4k columns found in Fig, 2b and apply it another prompt of 128k, the recall of attention would drop largely to 83.7%.
2.2 Attention Sparsity Exhibits Patterns
Table 1: Comparison of different sparse patterns.
Patterns
A-shape
Vertical-Slase
Block-Sparse
Top-K
Spatial Distribution
Static structured
Dynamic structured
Dynamic structured
Dynamic fine-grained
Lateney GPU
Low
Medium Small
Low
High
Time to build the index
Zero
Small
High
Although the sparsity distribution of attention matrix is dynamic, previous works 2X
HWP+24] have shown that they exhibit certain patterns in the two-dimensional space such as spatial clustering. Through our analysis of long-context prompts of various lengths and tasks, we have
Asshape
Vertical-Slash Block-Sperse
Block-Sparse
Topk Nearest Non-Zero Dist.
20-
Asshape
Black Sparse head
Input
15
Attention Weight Recall
0.92-
0.85
Vertical-S
10
Input
0.85
Aushape head
0.92
Inpur
4
12
16
20
28
32
0.85
10
20
Layer
Dense FLOPsFFLPPs in kemal
(a) Attention patterns
(b) Attention is spatial clustering
(c) Attention pattern recall
Figure 3: (a) Visualization of attention weights from different attention heads. For different prompts and tasks, the pattern of the same head is relatively consistent, the sparse indices are dynamically changing. Distance the top-10 nearest non-zero element the attention matrix. (c) Attention recall distribution using our identified patterns, where FLOPs in the kernel refer to the real FLOPs required for sparse attention computing using on GPUs. Here, 1x64 block size is used for the Verticul-Slan pattern, and 64x64 block size used for others GPUs. All visualization are based LLaMA-3-88 3-nt--3ttttt
Instruct-2621 (Gra24)

---

categorized such attention sparse patterns into the A-shape, Vertical-Slant (VS), and Block-Sparss patterns, as shown in Fig. Baland and Fig. Table I details the characteristies and differences between these three patterns.
A-shape pattern The attention weights of these types of heads concentrated on initial tokens and local windows [XTC 24 HWPP241 exhibiting relatively higher stability.
Vertical-Slash (VS) pattern The attention weights concentrated on specific tokens (vertical lines) [M23] and tokens at fixed intervals (slash lines). The positions vertical and slash lines in this pattern dynamically change with the context content and exhibit certain sparsity, making them difficult to be encompassed by local windows and A-shape patterns.
Block-Sparss pattern This sparsity pattern is the most dynamic, exhibiting more dispersed distribu- tion. Despite its dynamism, the attention weights maintain some characteristies of spatial clustering, which we identify the block-sparse pattern. We analyzed the distances between non-zero attention weights and their top-k nearest non-zero neighbors within a 128k prompt as shown in Fig. 33b The results indicate that across layers and heads, the distances between nearest non-zero values generally concentrated around 5, suggesting a strong spatial clustering of the attention weights.
The point of these three patterns is that we can leverage them to perform highly efficient sparse computing for the attention matrix in long- long-context LLMs. Fig. 3c we test how efficient is our indentified patterns retrieving attention scores with limit computing cost on GPU (FLOPs). First, attention heads are labeled with one of the sparse pattern (detail see $.2)). Then we demonstrate our patterns are significantly more efficient compared to other sparse methods [RCHG+24] XTC+ PPJF24] Specifically, with the same amount of FLOOs, our patterns achieve notable higher recall on attention scores, which can potentially lead to better accuracy. For example, previous Top-K methods [CCHG 244 XT PPJF24] struggle with the Block-Sparse pattern as they focus on specific tokens globally, while our pattern retrieves attention scores more efficiently and accurately. example how we use these patterns on long-context LLMs and how we implement optimized GPU kernels for these patterns in
3 MInference 1.0
Following the analysis in $2 we propose MInference to accelerate the pre-filling stage of long- context LLMs, consisting of three steps: 1) Offline attention pattern identification for each head; 2) Dynamic build of sparse indices w.r.t. the pattem; 3) Sparse attention calculation with optimized GPU kernels.
Dynamic
Sparse Calculation
Approximate by block
Approximate by lest
A-shape head
vertical-slash head
block-sparse head
Figure 4: The three sparse methods in MInference.
3.1 Problem Formulation
When accelerating the pre-filling stage of long-context LLMs with sparse attention computing, the attention matrix can be formulated as follows:
A(M) Softmax QK TKKT - -c(1-M MM),
(I)
where Mr, E (0.1)
represents the dynamic sparse mask for item i, of the attention matrix. Here, a large constant, such le5, ensuring that the less important attention weights for which M.j = have values approaching after the softmax, i.c., A xe 0.

---

The goal of the dynamic sparse attention system is to achieve greater speedup with minimal overhead while retaining as much of the attention weights as possible. Formally, this can be expressed as:
min
(A(M)
A is
min Expurse (M) + loverhead [M)
where /sqaree and loverhead represent the time spent on dynamic sparse attention computation and estimation of the approximate dynamic sparse pattern, respectively.
3.2 Speedup of Long-context LLM Inference via Dynamic Sparse Attention
Kernel-Aware Optimal Sparse Pattern Search To
Algorithm 1 Kernel-Aware Sparse Pattern Search
achieve the best accuracy with limited FLOPs budget, we propose an offline Kemel--waac Optimal Sparse Pat- tern Search method. In this step, we determine which sparse pattern will used for each attention head, and the optimal setting for the pattern in real calculation (e.g., the number of vertical/slash lines in VS pattern; or the number of top-k blocks in BS patterns). As shown in Algorithm| [ we first create the search space based on target FLOPs for each pattern, ensuring all potential candidates (i.e., different patterns with different settings) have similar computational cost. Kernel-awear dicates the computational cost reflects the real FLOPs in GPU kernels, instead of conceptual estimations, which crucial to achieve the optimal acceleration.
Input: Q, K, V e RSdd, patterns search space target FLOPs t, initial- ized search space
# Build kernel-aware search space for i 1 to lo| do 4, FLOPs li_ii_nooo (a,
while It; - t > do
σ; + Chan Spacce( pt)
here in-
t, + F1.OPs in (o,)
end while
P <- pUo:
end for
Search for optimal head pattern
Next, we go through the search space with reference ex- ample decide the optimal pattern and setting. Specif- ically, we use recall of the attention output as the objec- tive criterion when searching for the best pattern. This approach leverages FlashAttention [Dao24) to reduce GPU memory overhead and incorporates the informa- tion from the V matrix, enabling end-to-end selection of the best pattern, which further enhances performance.
Phest + 6 y Softmax(QK" /Vd) for +- to |pl
y: Attanionttenioon( / Vd,p.)
Phest argmin(y: W. Preas)
end for return Prest
Sparsity Indices Approximation and Dynamic Sparse Attention Calculation During the infer- ence stage, we will perform an online estimation on the attention matrix to dynamically determine the
Algorithm 2 Pertica--Slashe Head
Algorithm 3 Block-Sparse Head
Input: K,K, V SSxd, Kys kg N
Input: Q, E Sxd,, ky
Approximate verticul and slash pattern (last_g 64)
# Appreximate block-sparse pattern (block size 64)
A + softmax (QI- K" /Vd + Mesuual
Q MeanPooling (Q, bllok_kksisz
K MeanPooling (K block_sizee sizze)s
Indices of top k. vertical line, in verticul
À sofumax (QK' // + massuul
io +- argtopk (sum.(A), Ki
# Indices of top i blocks
Indices of top k. slash line, sum in slash i, argtopk (sum. (A),
<- argtopk A kbb
# Build sparse attention index
Build sparse attention index
sparseformat((.))
sparseformat((jr,..
# Final dynamic sparse attention scores (only index block)
Final dynamic sparse attention scores (only index block)
A softmax (sparse(QKr (QKT, Vo)/Vd)
A softmax (sparse(Q'') /io.)/ Vd
# Sparse mixed scores and values
# Sparse mixed scores and values
y sparse( AV 1, return y
+- sparse( AV, i.s)
return y
5

---

spatial distribution our sparse indices, based on the assigned patterns and the exact input. After that, we conduct the sparse attention computations with our optimized GPU kernels. The implementation details of our kernels can be found in Appendix C.4 Noted that the sparse mask is static for A-shape heads, so there is no overhead in building the dynamic masks, and only sparse calculation is required.
(i) Vertical-Slash head. As shown in Algorithm 2 due to the continuity of vertical and slash lines, we matmul the last query vector [ and key vector K to produce the estimated attention matrix A, which, in turn, is used to determine the indices for the vertical i, and slash i lines. After obtaining the sparse indices for the vertical and slash lines, we convert them into a sparse format tex- Using these sparse indices, we perform block-sparse calculations of the attention weights and attention output.
(ii) Block-Sparse head. Per Algorithm mean pooling applied on Q and K to obtain Q and K, respectively. The two matrices are multiplied to get the estimated block-level attention weights Å. Since the mean pooling and matrix multiplication operations are commutative, the resulting attention weights are approximately equivalent to the actual attention weights after mean pooling. This allows us to approximate the actual attention weights' block-sparse pattern with minimal overhead. Similarly, we build a sparse index és and use it to compute the sparse attention weights and attention output.
4 Experiments
In this section, we investigate two questions: (i) How effective is MInference? We evaluate method on three general long-context benchmarks: InfiniteBench [4C RULER [HSK HSSK+24), and the Needle In A Haystack task [Kam23], as well as the long-context language modeling task [RP]+20] These benchmarks cover long-context QA, multi-hop QA, math reasoning, aggrega- tion tasks, summarization, retrieval tasks, and code debugging, allowing to assess MInference's effectiveness across a wide range of long-context scenarios. (ii) How efficient is MInference?
We delve into the end-to-end latency and its breakdown to evaluate the efficiency of MInference. Additional experimental, latency results, and analysis can be found in Appendix
and F
Implementation Details Our experiments use four state-of-the-art long-context LLMs: LLaMA- 1-8B-Instruct-2224
GLMM--9BB-I [GZX and Yi-9B-
200K 24]. Additionally, we tested Needle in A Haystack [Km22 on Phi-3-Mini-
128K [AJA+2) 244 and Qwen2-7B-128K [BBC*23], as detailed in Appendix D.I To guarantee stable results, we use greedy decoding in all experiments. We provide simple custom implementation of our method in PyTorch, built on FlashAttention [Do224] Triton [TKC19], and the dynamic sparse compiler PIT 23) We set the target FLOPs to 1k global tokens and 4k local windows in the A-shape pattern. We set last_g patterns, respectively. The latency experiments are conducted on a single Nvidia A100 GPU in the b 6 format. More details are provided in Appendix C.2
64 and block_stte size = 64 in the Vertical-Slass and Block- -Saar-
We use the provided metrics and scripts from the following benchmarks for evaluation. More details about dataset can be found in Appendix C.
Dataset & Evaluation Metrics
This benchmark consists of 10 tasks, including retrieval tasks such as PassKey retrieval, Number retrieval, and KV retrieval, as well as representative realistic tasks like guestion-answeringg coding, dialogue, and summarization. The average context length of InfiniteBench is about 214K tokens.
(i) InfiniteBench (24:+
(ii) RULER [HSK* + 24): : A challenging long-conter: benchmark consisting of 4 categories and 13 complex tasks, including retrieval, multi-hop tracing and aggregation, and QA tasks. It contains subsets with different prompt lengths up to 1 128k tokens, allowing us to determine the actual context window size of the model based on the results.
(iii) Needle In A Haystack [Kmm2]] A long-context retrieval benchmark testing L.LMs' performance with context window sizes up to IM tokens where information placed at various positions.
http:///huggingfaacccoooogannnne -aaaa------nstttt----retet Inr-nntrutt-rraient--ee hhtpss//huggiggfaec.ooogrraatett --n-ama--------sttttt-t--- EnstructGrradient- 1048k

---

Table 2: Performance of different methods with different base models on InfiniteBench [ZCH 24
Methods
En. Sum Ea.QA
Ec.MC En. Dia Zh.QA Debu. Math. Find Retr.PaskKey Retr.Num Retr. KV Avg.
LLaMA-3-N0 3-----3--B--- --NB-262K
20.2
12.4
67.3
6.0
12.9
22.1
26.6
100.0
100.0
14.4
38.2
StreamingL.M
21.0
40.2
10.0
10.4
25.9
30.0
5.1
0.8
23.8
StreamingLLM dilated
20.1
44.5
15.5
20.5
27.5
87.5
0.5
24.2
StreumingL.MM w/ strided
17.3
8.2
27.5
14.5
11.2
19.5
27.5
4.0
2.1
1.0
13.3
InfLLM
24.1
7.8
45.0
6.0
11.4
19.5
32.9
100.0
100.0
1.2
34.8
static
43.2
3.5
8.9
20.6
25.1
0.2
31.9
Ours
20.5
12.9
65.9
7.5
12.5
22.3
33.1
100,0
100,0
12.8
38.8
y1-98-200k
8.2
10.6
64,2
1.0
17.3
21.3
23.4
99.8
100.0
28.8
37.5
StreumingLLM
5,4
14.2
38.0
4.0
18.8
18.8
22.3
39.2
6.1
1.6
16.8
Streamingl dilated
StreamingLLM M w/ strided
6.1
4.5
9.8
0.0
16.9
0.0
3.1
1.5
0.0
0.0
4.6
InfLLM
6,3
13.0
45.9
2.5
21.5
20.6
34.6
85.3
88.1
1.4
31.9
w/
48.5
22.9
Ours
7.9
11.2
64,2
1,0
17.9
24.1
23.1
99.5
100.0
27.6
37.7
6J.M-499B-1M
28.3
9.7
68.6
39.5
12.1
29.4
38.9
100.0
100.0
41.0
46.7
StreumingLLM
27.7
6.4
40.2
12.5
10.8
27.7
21.1
97.1
25.6
0.6
27.0
InfLLL
28.0
7.3
45.0
14.0
10.7
27.9
39.4
98.0
100.0
2.6
37.3
Ours
28.8
68.6
38.5
39.1
100.0
100.0
43.0
47.0
(iv) PG-19 IRJ+ 20) Following StreamingLLMM 2TC and H2O ZZZ + we use PG-19 for long-context language modeling tasks with prompts up to 100k tokens.
Baselines We include five training-freee sparse attention approaches our baselines: StreamingILLM [XTC+24), which corresponds to the A-shape pattern. We use 1k global tokens and 4k local windows in all our experiments; 2) StreamingLLLM w/ dilated [BPC20] which sets dilated local windows with intervals in the local windows direction. We use 1k global tokens and 8k dilated attention windows with an interval 1; 3) StreamingLLMM windows while adding dilated attention. We use Ik global tokens, 2k local windows, and 4k dilated attention windows with an interval 1; 4) InfLLM [4 which uses a memory unit to process streaming long sequences. Following the paper, we set 128 global tokens and 8k local windows in all experiments; 5) Ours w/ static, which utilizes static sparse indices in the Vertical-Sie -lashh and Block- Sparse heads. For all basclines, perform sparse computation only during the pre-filling stage, while retaining dense computation during the decoding stage.
strided [CGRST9]. which retains local
InfiniteBench As shown in Table 2 MInference achieves the best overall performance In- finiteBench compared to baseline methods. Remarkably. MInference matches or even slightly surpasses the performance of the original full attention baseline on some tasks, despite the significant acceleration it provided. From the perspective of different tasks, our method not only performs well in natural language tasks such summarization, QA, and code, but also maintains the original model's performance on ectrieval-relate tasks. Baseline methods such as StreamingL on the
Table Performance (%) of different models different methods on RULER HHSK +
evaluated lengths
from 4k to 128k.
Methods
Claimed
Effective
4K
8K
16K
32K
64K
128K
Avg.
LLaMA-3 1-MMAA3--B--200 3-8B-262K
262K
16K
97.2
91.8
87.3
80.8
77.4
72.2
84.4
StreamingL.M.
4K
97.2
38.1
17.2
35.0
StreamingLLMM LLM w/ dilated
<4K
23.4
0.7
1.4
18.8
16.5
15.6
12.7
StreamingL.MM M w/ strided InfLLM
<4K
2.0
0.7
0.6
0.6
0.7
1.3
1.0
4K
89.4
79.8
70.1
55.6
43.0
39.5
62.9
Ours
32K
97.7
91.2
85.0
77.6
87.0
Y2-B-2000
200K
SK
91.9
90.2
78.8
76.3
68.1
62.9
78.1
StreamingL.LMM
4K
37.8
33.9
18.6
12.8
34.3
StreamingLLM w/ dilated
<4K
44.8
42.8
38.5
29.8
26.8
23,9
34.4
StrcamingL.MM w/ strided InfL.LM
<4K
2.6
0.7
0.6
0.6
1.2
0.5
1.1
<4K
80.3
83.9
60.7
45.2
38.6
30.2
56.5
Ours
8K
92.33
89.7
79.0
73.8
56.9
74.7
GLM-4-9B-IAM
64K
93.8
91.6
87.4
88.0
StreamingL.MM
4K
93.8
66.9
58.5
$1.4
45.9
39.1
59.3
InfLLM
8K
94.7
89.5
76.4
66.5
56.8
53,5
72.9
Ours
64K
94.6
93.1
91.0
89.6
85.5
84.0
89.6

---

contrary, struggle with these retrieval tasks. Additionally, on tasks such as dialogue QA, using local attention mechanisms can better handle these tasks, while our performance closer to the original results, indicating that our method is not solely based on local windows. Extending the local windows' intervals StreamingL.LM,, i.e., w/ dilated and w/ strided, has minimal impact on the model's performance.
To further reveal the true potential of our method long-context LLMs, we evaluate MIn- ference with the state-of-the-ard long-context challenge, RULER. As shown in Table MInference effectively maintains the long-context performance even in complex multi-hop or aggregation tasks in RULER. It even outperforms the original full attention for testing lengths beyond 32K, achiev- ing effective context windows of 32K and 64K (context with performance over 85% is considered effective (4) in LLaMA-3-8B-262 and GLM-4-9B-IM.
RULER
Following the approach of StreamingLLM XTC and H2O ZSZ 24] we evaluate our methods against baselines on the language modeling task based on the PG-19 dataset [ As shown in method yields best results compared to other sparse approaches, and exhibits minimal divergence compared to the full attention baseline. For prompts of 100K token, perplexity only 0.2 higher than the full attention, but lower than StreamingLLMM for 0.25 and 0.75 on the Yi-9B-2005 and LLaMA-3-222 --222 models respectively.
Language Modeling
log perplexity
10.0-
Ailayya Boj
9.5
9.0
Preaaing. LMM dtond
Medivenoo cosic Medvrenoe
Smoominll Slaned StoonikollM drdded
Vedeemnne Medemmoce
Swwering LM carded
8.5
1k 10k
30k
50k
75k
100k
10k
30k
sók
75k
100k
Windows
(a) 0an--------ntuut-2222
(b) Yi-9B-200K
Figure 5: Perplexity results on PG-19 RPJ 20
using different models and methods.
Needle In A Haystack Comparing Fig.
Needle in Haystack LLaMA-3-8B-1M Streaming 1M Context
to Fig. 6 our method effectively retains the
ability to process information at different po-
Drpth Percent (%6
22-
08
sitions across various context windows, ranging from 1k IM tokens. In contrast, methods like Streaming LLM and InfLLM (as shown in Ap- pendix D.I. while effective at reducing latency, experience a sharp decline in performance once critical information extends beyond the range of global tokens and local windows.
35-
674
-04
78-
02
100-
140
220e
7000
130
Context Length
Figure 6: Results on Needle In A Haystack of
StreamingLLM [XTC*24 24 in LLaMA-3-8B--IM
To evaluate the contributions of different components in MInference, we in-
A Ablation Study
troduce four variants for the ablation study: (1) Ours w/ static, which uses a static sparse mask in the Vertical-Slash and Block-Sparse patterns; (2) Ours w/ only A-shape, which is equivalent to StreamingLLMM (3) Ours w/ only block- -sarse, which uses only the Block-Sparse pattern in the dynamic sparse calculation. (4) Ours w/ only vertical-slash which uses only the Pertical-Slash pattern the dynamic sparse calculation.
Tables 2 3 and present the ablation results. It first proves that using static indices significantly degrades LLM performance, especially in highly dynamic tasks like KV retrieval, where accuracy nearly drops to zero. This highlight the necessity of our dynamic strategy and the effectiveness of our dynamically built sparse indices. Additionally, remove any pattern from the three leads varying degrees of performance degradation. Specifically, "only A-shape" can only capture information

---

Table 4: Performance finiteBench (ZCH*24 24|.
different ablation
methods
using
LLaMA-3-888 B-8---nsstuutt-222 on In-
Methods
En.Sum En Qn
En.MC En Dia
ZA.QA
Code.Debug Math.Find Retr Passkey Retr.Num Retr
Ours
20.5
12.9
65.9
7.5
12.5
22.3
33.1
100.0
100,0
12.8
38.8
Ours w/ only block-sparse
12.4
3.4
60
12.2
0.0
18.7
Ours w/ only vertical-slush
19.6
12.0
62.1
9.5
11.7
21.6
29.1
100.0
100.0
5.0
37.1
within local windows. The "only block-sparse" variant using only the BS pattern, also results in significant performance declines. On the other hand, "only vertical-slash manages preserve most of the performance due to its balance between dynamicity and the StreamingLLM pattern, but still fall behind the full version of our method.
Latency Fig. lb and TO shows the latency and breakdown of MInference across different context windows single AT00. At 100K, 300K, 500K, and 1M tokens, our method achieves speedups of 1.8x, 4.1 x, 6.8 x, and 10x, respectively. It reduces the pre-filling latency from 30 mins to mins single A10 for prompt of 1M token. By further utilizing tensor parallel (LMZ++2] and context parallel LZA24 JTZ this latency can be reduced to 22 seconds on 8x A 100 GPUs. This significantly lowers the deployment cost of long-context LLMs and enhances user experience. And since our kernel is implemented based on Triton, it can be easily ported to other devices and achieve similar speedups, such as on the H100 MI300X. Additionally, analyzing the latency breakdown, we found about 5 5%-20% of the overhead is spent on dynamic sparse index building, while the remaining time is spent on dynamic sparse calculation.
Integrate with KV cache compression methods We also combined MInference with state-of- the-art KV cache compression method SnapKV [LHY 24], as shown in Table This proves our method is compatible with KV cache compression techniques. For most tasks, performance remains nearly unchanged, with the average score even showing slight increase, which further demonstrates the potential practical value of our method phenomenon is also observed other works, such as ShadowKV (SCB 24).
optimization for serving long-context LLMs. This
Table S: Performance different methods on InfiniteBench ZCH decoding stage,
using SnapKV LIIY 24
Methods
FnnSSum En.QA En.MC Em. Zh.QA Code. Dehug Math.Find Retr.PassKey Retr.Num Retr. /Avg.
LLaMA-3 w/ SnapKV]
18.0
11.8
65.5
2.5
12.0
21.3
26.6
100.0
100.0
1.8
36.0
Ours w/ SnapKV
18.9
11.7
66.4
6.5
12.1
21.8
33.1
100.0
100.0
37.3
Scaling-up on Larger LLMs We also evaluated MInference larger LLMs, such LLaMA- MInference maintains strong performance even larger models.
3-70B-IM As shown in Table Notably, in dynamic tasks such KV retrieval, MInference match even slightly improve performance compared to full attention. contrast, baselines like InfLLM generally struggle with tasks such KV retrieval.
Table Performance different methods using LLaMA-3-700 Instr-nnsttttt22221 InfiniteBench ZZCH +24
Methods
Fn.Sum Fn.QA En.MC Fi. Zh.QA Code Math.Find Retr.Passkey Retr.Num Retr. K Avg.
(LaMA-3-70B-266
20.7
10.3
84.2
9.5
14.0
33.2
61.7
97.0
100.0
34.0
46.5
StreamingLLMM
52.0
10.0
12.6
27.4
21.6
InfLLM
24.1
8.1
57.0
10.0
12.9
27.4
52.3
100.0
100.0
39.2
Ours
20.6
10.1
83.4
10.0
14.1
34.1
61.9
100.0
100.0
39.0
47.3
5 Related Works
Sparse Attention Due to the quadratic complexity the attention mechanism, many previous
works have focused sparse attention to improve the efficiency of Transformers. These meth- ods include static sparse patterns, cluster- based sparse approaches, and dynamic sparse attention.
hhng::/hougginggaae.ooognaaaaiaa 3-7BB
http:///oeoaaoooooooaaooooootnnn Lnsttnctt-rrddien22222

---

Static sparse patterns include techniques such as sliding windows JJMM 23
AJA 24
dilated at-
tention [CGRST9] SGR 21 DMD 23I. and mixed sparse patterns [BCC20
ZGD 20
LCSR21
Cluster-based sparse methods include hash-based [KL.20 and kNN-based RSVVG2
NF.C22
methods. All of the above methods require pre-training the model from scratch, which makes them infeasible to be directly used as plugin for reay-to-use LLMs. Recently, there has been work [DG24] ZAW24] to unify state space models [GGR22] GD24 DG24], and linear atten- tion [KVPF20 2DJ into structured masked attention. Additionally, some works [WZH21] LQC 2CHG*24 leverage the dynamic nature of attention to predict sparse patterns dynami- cally. However, these approaches often focus on low-rank hidden states during the dynamic pattern approximation or use post-statistical methods to obtain the sparse mask, which introduce substantial overhead in the estimation step, making them less useful for long-contess LLMs.
Scaling Context Windows of LLMs Recent research has focused on expanding the context window of pre-trained LLMs, that enables LLMs to handle more complex real-life applications JYW+23
23] These methods can be categorized into: 1) Staged pre-training 23 :P]+ 2)
Modifying or interpolating position embeddings (PSL22 CWCT23 POFS24 DZZ + 24 : 3) Utilizing
external memory modules for context storage BANG23 TSP 23 XZH 24): 4) Expanding com-
putations across multiple devices in a distributed manner [LZA24). However, these methods do not
alleviate the high inference costs in long-context processing.
Long-Come Cont-Co LLM Inference Recent studies [F224] have tackled the high computational cost of attention and substantial KV cache storage in long-context scenarios from two angles: pre-filling and decoding. Pre-filling optimizations are primarily categorized State Space Models [GGR22] GD24], linear attention methods SSDH+ PAA+23], memory-based meth-
ods [MFG24 HBK++41 hybrid methods (LLB*24 2LL, and prompt compression meth- JW+24 2W]+24). However, these approaches require training from scratch or additional overhead and are difficult implement directly in pretrained long-context LLMs. Recently, some studies [MEL24, X7H+24] LCL 24) have focused on using kNN or cluster- based sparse attention to accelerate LLM inference. However, these methods often lead to reduced accuracy, limited speedup, or are restricted to CPU scenarios.
ods (LDGL23 23
In contrast, optimizations for the decoding stage divided into LJW NW+22 +44 1) Reusing attention
KV to reduce KV cache storage [Sha19 23 SDZ 24 DA24
Nr.
2) Static KV cache
dropping XT
HWP+24] 24|; 3) Dynamic KV cache dropping ZZZ 24,
LDL + 24
GZL+24
24
LHY 24
APB 24|; 4) Dynamic KV cache offloading [RCHG + 24 DHJ 24,
TZZ*24
LCL 24
CSY
SCB
24): 5) Methods for restoring performance loss due to KV cache compres-
sion [AAJ 24
24] 2 6) Hierarchical speculative decoding methods [SCY+24] CTS+24); 7) KV
cache quantitation [LY]+2] 24). Nevertheless, these methods do not address the heavy computational
burden of the attention in the pre-filling stage.
6 Conclusion
This paper addresses the expensive computational cost and the unacceptable latency of the attention calculations in the pre-filling stage of long-context LLMs. We propose MInference, a method that accelerates the pre-filling stage by leveraging dynamic sparse attention with spatial aggregation patterns. Specifically, we categorize attention heads into three types: A-shape, Vertical-Slash, and Block-Sparse. Using a kernel-aware optimal sparse pattern search method, we identify the optimal pattern for each head. Subsequently, we utilize a fast approximation approach to build dynamic sparse masks for different inputs, and then apply these mask to perform sparse attention calculations. Experimental results on benchmarks such as InfiniteBench, RULER, language modeling, and Needle In A Haystack demonstrate that our method effectively maintains the long-context capabilities of LLMs while achieving up to a 10x speedup. reducing the latency from 30 minutes to 3 minutes per prompt for million token prompts on a single A100 GPU. Additionally, we have found that similar dynamic sparse attention patterns also exist in both multi-modal LLMs IWWI.+24 and encoder-decoder L.LMs (RSR+20) Using MInference for pre-filling stage inference acceleration holds great promise.
10

---

References
[AAJ*24] Muhammad Adnan, Akhil Arunkumar, Gaurav Jain, Prashant Nair, llya Soloveychik, and Purushotham Kamath. Keyformer: Kv cache reduction through key tokens selection for efficient generative inference. Proceedings of Machine Learning and Systems, 6:1 14- 127, 2024.
[AJA 24] Marah Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat Behl, Alon Ben- haim, Misha Bilenko, Johan Bjorck, Sébastien Bubeck, Martin Cai, Caio César Teodoro Mendes, Weizhu Chen, Vishrav Chaudhary, Parul Chopra, Allie Del Giorno, Gustavo de Rosa, Matthew Dixon, Ronen Eldan, Dan Iter, Amit Garg, Abhishek Goswami, Suriya Gunasekar, Emman Haider, Junheng Hao, Russell J. Hewett, Jamie Huynh, Mojan Javaheripi, Xin Jin, Piero Kauffmann, Nikos Karampatziakis, Dongwoo Kim, Mahoud Khademi, Lev Kurilenko, James R. Lee, Yin Tat Lee, Yuanzhi Li, Chen Liang, Weishung Liu, Eric Lin, Zeqi Lin, Piyush Madan, Arindam Mitra, Hardik Modi, Anh Nguyen, Brandon Norick, Barun Patra, Daniel Perez-Becker, Thomas Portet, Reid Pryzant, Heyang Qin, Marko Radmilac, Corby Rosset, Sambudha Roy, Olatunji Ruwase, Olli Saarikivi, Amin Saied, Adil Salim, Michael Santacroce, Shital Shah, Ning Shang. Hiteshi Sharma, Xia Song. Masahiro Tanaka, Xin Wang, Rachel Ward, Guanhua Wang, Philipp Witte, Michael Wyatt, Can Xu, Jiahang Xu, Sonali Yadav, Fan Yang, Ziyi Yang. Donghan Yu, Chengruidong Zhang, Cyril Zhang, Jianwen Zhang, Li Lyna Zhang, Yi Zhang, Yue Zhang, Yunan Zhang, and Xiren Zhou. Phi-3 technical report: A highly capable language model locally on your phone, 2024.
[ALTa++2]] Joshua Ainslie, James TeerThorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. 2023.
[APB+24] Sotiris Anagnostidis, Dario Pavilo, Luca Biggio, Lorenzo Noci, Aurelien Lucchi, and Thomas Hofmann. Dynamic context pruning for efficient and interpretable autoregres- sive transformers. Advances in Neural Information Processing Systems, 2024.
[BANG23] Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew R. Gormley. Unlimiformer: Long-range transformers with unlimited length input. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[BBC+23] Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, Binyuan Hui, Luo Ji, Mei Li, Junyang Lin, Runji Lin, Dayiheng Liu, Gao Liu, Chengqiang Lu, Keming Lu, Jianxin Ma, Rui Men, Xingzhang Ren. Xuancheng Ren, Chuanqi Tan, Sinan Tan, Jianhong Tu, Peng Wang, Shijie Wang. Wei Wang. Shengguang Wu, Benfeng Xu, Jin Xu, An Yang, Hao Yang, Jian Yang, Shusheng Yang. Yang Yao, Bowen Yu, Hongyi Yuan, Zheng Yuan, Jianwei Zhang, Xingxuan Zhang, Yichang Zhang, Zhenru Zhang, Chang Zhou, Jingren Zhou, Xiaohuan Zhou, and Tianhang Zhu. Qwen technical report. ArXiv preprint, abs/2309. 16609, 2023.
[BPC20] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. ArXiv preprint, abs/2004.05150 2020.
[BSK+24] Ramakrishna Bairi, Atharv Sonwane, Aditya Kanade, Arun Iyer, Suresh Parthasarathy, Sriram Rajamani, B Ashok, and Shashank Shet. Codeplan: Repository-level cod- ing using Iims and planning. Proceedings of the ACM on Software Engineering, ((FSE):6756998, 2024.
[CGRS19] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. ArXiv preprint, abs/1904. 10509, 2019.
[CPG+23] Avi Caciularu, Matthew E Peters, Jacob Goldberger, Ido Dagan, and Arman Co- han. Peek across: Improving multi-document modeling via cross-document question- answering. Proceedings of the 6 Annual Meeting of the Association for Computa- tional Linguistics (Volume 1: Long Papers), pages 1970-1989, 2023.
11

---

[CSY+24] Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye. Yang Zhou, Jianyu Zhang, Niklas Nolte. Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, al. Magicpig: Lsh sampling for efficient Ilm generation. ArXiv preprint, abs/2410. 00./2410.161777 16179, 2024.
[CTS+24) Jian Chen, Vashisth Tiwari, Ranajoy Sadhukhan, Zhuoming Chen, Jinyuan Shi, lan En-Hsu Yen, and Beidi Chen. Magicdec: Breaking the latency-throughput tradeoff for long context generation with speculative decoding. ArXiv, abs/2408. 11049, 2024.
[CWCT23] Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. ArXiv preprint, fbs/2306.15595, 2023.
[CWW+24] Zhe Chen, Jiannan Wu, Wenhai Wang, Weijie Su, Guo Chen, Sen Xing, Muyan Zhong, Qinglong Zhang, Xizhou Zhu, Lewei Lu, et al. Internvl: Scaling up vision foundation models and aligning for generic isuual-inguistic tasks. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 24185-24198, 2024.
[DA24] DeepSeek-AL. Deepseck-v2 v2: A strong, economical, and efficient mixture-of-experts language model, 2024.
[Dao24] Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference Learning Representations, 2024.
[DG24] Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference on Machine Learning, 2024.
[DHJ+24] Jincheng Dai, Zhuowei Huang, Haiyun Jiang, Chen Chen, Deng Cai, Wei Bi, and Shum- ing Shi. Sequence can secretly tell you what to discard. ArXiv preprint, abs/2404. 15949, 2024.
[DMD+23] Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. ArXiv preprint, abs/2307.02444 2023.
[DSY24] Yichuan Deng, Zhao Song, and Chiwun Yang. Attention is naturally sparse with gaussian distributed input. ArXiv preprint, abs/2404.02690, 2024.
[DYZ+24] Harry Dong, Xinyu Yang, Zhenyu Zhang, Zhangyang Wang, Yuejie Chi, and Beidi Chen. Get more with LESS: Synthesizing recurrence with KV cache compression for efficient LLM inference. In Forty-first International Conference on Machine Learning. 2024.
[DZZ*24] Yiran Ding, Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. LongroPE: Extending LLM context window beyond 2 million tokens. In Forty-first International Conference on Machine Learning, 2024.
[F] Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models 128k context. In
Forty-first International Conference on Machine Learning, 2024.
[Fu24] Yao Fu. Challenges in deploying long-context transformers: A theoretical peak perfor- mance analysis. ArXiv preprint, abs/2405.08944, 2024.
[GD24] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024.
[GGR22) Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[Gra24] Gradient. Llama-3 8b instruct gradient 4194k (v0.1), 2024.
12

---

[4] 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.
[GZX*24] Team GLM, Aohan Zeng, Bin Xu, Bowen Wang, Chenhui Zhang, Da Yin, Diego Rojas, Guanyu Feng, Hanlin Zhao, Hanyu Lai, et al. Chatglm: A family of large language models from glm-130b to glm-4 all tools. ArXiv preprini, abs/2406. 12793, 2024.
[4BK+24] Namgyu Ho, Sangmin Bae, Tachyeon Kim, hyunjik.jo, Yireun Kim, Tal Schuster, Adam Fisch, James Thome,, and YouYoung Yun. Block transformer: Global-to-local language modeling for fast inference. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
[HSK*24] Cheng-Ping Hsich, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. RULER: What's s the real context size of your long-context language models? In First Conference on Language Modeling, 2024.
[HWP+24] Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. LM-infinite: Zero-shot extreme length generalization for large language models. In Kevin Duh, Helena Gomez, and Steven Bethard, editors, Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume I: Long Papers), pages 3991- 4008, Mexico City, Mexico, 2024, Association for Computational Linguistics.
[JSM+23] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Deven- dra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. ArXiv preprint, abs/2310.0688 2023.
[JTZ+24] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Reza Yazdani Aminadabi, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. System optimizations for enabling training of extreme long sequence transformer models. In Proceedings of the 43rd ACM Symposium on Principles of Distributed Computing. pages 121-130, 2024.
[JW+24] Huiqiang Jiang. Qianhui Wu, Xufang Luo, Dongsheng Li, Chin- Lin, Yuqing Yang, and Lili Qiu. LongLLMLingua: Accelerating and enhancing LLMs in long context scenarios via prompt compression. Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1658-1677, Bangkok, Thailand, 2024. Association for Computational Linguistics.
Lun-Wet Wei Ku, Andre Martins, and Vivek
[JWL+23] Huiqiang Jiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. LLMLingua: Compressing prompts for accelerated inference of large language models. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pages 13358-1339 13376, Singapore, 2023. Association for Computational Linguistics.
[JYW+23] Carlos E Jimenez, John Yang, Alexander Wettig, Shunyu Yao, Kexin Pei, Ofir Press, and Karthik R Narasimhan. Swe-bench: Can language models resolve real-world github issues? In The Twelfth International Conference on Learning Representations, 2023.
[Kam23] Greg Kamradt. Needle in a haystack pressure testing 11ms, 2023,
[KKL20] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient trans- former. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[KVPF20] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Trans- formers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 1 119 of Proceedings of Machine Learning Research, pages 5156-5165. PMLR, 2020.
13

---

[LC Di Liu, Meng Chen, Baotong Lu. Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chengruidong Zhang, Bailu Ding, Kai Zhang, et al. Retrievalattention: Accelerating long-context Il inference via vector retrieval. ArXiv, 2409. 1051, 2024.
[LCSR21] François Lagunas, Ella Charlaix, Victor Sanh, and Alexander Rush. Block pruning for faster transformers. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10619-10629 Online and Punta Cana, Dominican Republic, 2021. Association for Computational Linguistics.
[LCW21] Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. ArXiv preprint, abs/2106.03764, 2021.
[LCW23] Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive flexibility of self-attention matrices. Proceedings of the AAAI Conference on Artificial Intelligence, 37(7):8773-88) 17818 2023.
[LDGIL23] Yucheng Li, Bo Dong, Frank Guerin, and Chenghua Lin. Compressing context to enhance inference efficiency of large language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 6342-6353, 2023.
[LDL++2] Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for Ilm kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024.
[LHY+22] 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. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[LJW+24] Yucheng Li, Huiqiang Jiang, Qianhui Wu, Xufang Luo, Surin Ahn, Chengruidong Zhang, Amir H Abdi, Dongsheng Li, Jianfeng Gao, Yuqing Yang, et al. Scbench: A kv cache-centric analysis of long-context methods. arXiv preprint 2rXiv:2412.10319, 2024.
[LLB+24] Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. Jamba: A hybrid waasformer-mambl language model. ArXiv preprint, (bs//403..9987, 2024.
[LLWL24] Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. Advances in neural information processing systems, 36, 2024.
[LM Zhiqi Lin, Youshan Miao, Quanlu Zhang, Fan Yang, Yi Zhu, Cheng Li, Saeed Maleki, Xu Cao, Ning Shang, Yilei Yang, Weijiang Xu, Mao Yang, Lintao Zhang, and Lidong Zhou. nnscaler: Constraint-guided parallelization plan generation for deep learning training. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24). USENIX Association, 2024.
[LQC+22] Liu Liu, Zheng Qu, Zhaodong Chen, Fengbin Tu, Yufei Ding, and Yuan Xie. Dynamic sparse attention for scalable transformer acceleration. IEEE Transactions on Computers, 71(12)3116-3177 2022.
[LWD+23] Zichang Liu, Jue Wang. Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, and Beidi Chen. Deja vu: Contextual sparsity for efficient LLMs at inference time. In Andreas Krause, Emma Brunskill, Kyunghyun Cho. Barbara Engelhardt, Sivan Sabato, and Jonathan Scar- lett, editors, Proceedings of the 40th ICML, volume 202 of Proceedings of Machine Learning Research, pages 22137-22176 176 PMLR, 2023.
[LY]+24] Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braver- man, Beidi Chen, and Xia Hu. KIVI: tuning-free asymmetric 2bit quantization for KV cache. In Forty- first International Conference Machine Learning, 2024.
14

---

[LYZA24] Hao Liu, Wilson Yan, Matei Zaharia, and Pieter Abbeel. World model on million-length video and language with ringattention. ArXiv preprint, abs/2402.08268,, 2024.
[LZA24] Hao Liu, Matei Zaharia, and Pieter Abbeel. Ringattention with blockwise transform- ers for near-infinite context. In The Twelfth International Conference on Learning Representations, 2024.
[LZD*24] Tianle Li, Ge Zhang, Quy Duc Do, Xiang Yue, and Wenhu Chen. Long-context Ilms struggle with long in-context learning. ArXiv preprint, abs/2404.0205 020000, 2024.
[MEL24] Yuzhen Mao, Martin Ester, and Ke Li. Iceformer: Accelerated inference with long- sequence transformers on CPUs. In The Twelfth International Conference on Learning Representations, 2024.
[MFG24] Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention. ArXiv preprint, abs/2404.077 07143. 143, 2024.
[MJ23] Amirkeivan Mohtashami and Martin Jaggi. Random-access infinite context length for transformers. In seirty-seventh Conference Neural Information Processing Systems, 2023.
[NEC+24] Piotr Nawrot, Adrian Łańcucki, Marcin Chochowski. David Tarjan, and Edoardo Ponti. Dynamic memory compression: Retrofiuting LLMS for accelerated inference. In Forty-first International Conference on Machine Learning, 2024.
23] Erik Nijkamp. Tian Xie. Hiroaki Hayashi, Bo Pang, Congying Xia, Chen Xing, Jesse Vig, Semih Yavuz. Philippe Laban, Ben Krause, Senthil Purushwalkam, Tong Niu, Wojciech Kryšcinski, Lidiya Murakhovs' kaa Prafulla Kumar Choubey, Alex Fabbri, Ye Liu, Rui Meng, Lifu Tu, Meghana Bhat, Chien-Sheng Wu, Silvio Savarese, Yingho Zhou, Shafiq Joty, and Caiming Xiong. Xgen-7b technical report. ArXiv preprint, abs/2309.0344 034500 2023.
[OHY+24] Matanel Oren, Michael Hassid, Nir Yarden, Yossi Adi, and Roy Schwartz. Transformers multi-state RNNs. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, Proceedings of the 2024 Conference on Empirical Methods Natural Language Processing, pages 18724-18741, Miami, Florida, USA, November 2024. Association for Computational Linguistics.
[Ope24] OpenAl. Learning to reason with 1lms, 2024.
[3] Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Leon Derezynski, Xingjian Du, Matteo Grella, Kranthi Gv, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Jiaju Lin, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Johan Wind, Stanislaw Wožniak, Zhenyuan Zhang, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. RWKV: Reinventing RNNs for the transformer era. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Findings of the ACL: EMNLP 2023, pages 14048-14077, Singapore, 2023. Association for Computational Linguistics.
[POC+23] Joon Sung Park, Joseph C. O'Brien, Carrie J. Cai, Meredith Ringel Morris, Percy Liang, and Michael S. Bernstein. Generative agents: Interactive simulacra of human behavior. Proceedings of the 36th Annual ACM Symposium User Interface Software and Technology, 2023.
[PPJF24] Matteo Pagliardini, Daniele Paliotta, Martin Jaggi, and François Fleuret. Fast atten- tion over long sequences with dynamic sparse flash attention. Advances in Neural Information Processing Systems, 36, 2024.
[PQFS24] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. In The Twelfth International Conference Learning Representations, 2024.
15

---

[PSL22] Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In The Tenth International Con- ference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[PWJ*24] Zhuoshi Pan, Qianhui Wu, Huiqiang Jiang, Menglin Xia, Xufang Luo. Jue Zhang, Qingwei Lin, Victor Ruhle, Yuqing Yang. Chin- Lin, H. Vicky Zhao. Lili Qiu, and Dongmei Zhang. LLMLingua-22 2: Data distillation for efficient and faithful task-agnostic prompt compression. In Lun-Wet Wei Ku, Andre Martins, and Vivek Srikumar, editors, Findings of the ACL 2024, pages 963-981, Bangkok, Thailand and virtual meeting, 2024. Association for Computational Linguistics.
(RCHG+24) Luka Ribar, Ivan Chelombiev, Luke Hudlass-Galley, Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: Bandwidth-efficient LLM inference. In Forty-first International Conference on Machine Learning. 2024.
[24] Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling. ArXiv preprint, abs/2406.07522. 2024.
[RP]+20] Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Compressive transformers for long-range sequence modelling. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[RSR+20] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Mach. Learn. Res.,
2020.
[4S Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap. Jean-baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrit- twieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. ArXiv preprint, abs/2403.05530 2024.
[RSVG21] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content- based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53 0553-68,22 2021.
[SCB+24] Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, and Beidi Chen. Shadowkv: Ky cache in shadows for high-throughput long-context I1m inference, ArXiv preprint, abs/2410.21465. 2024.
[SCY+24] Hanshi Sun, Zhuoming Chen, Xinyu Yang, Yuandong Tian, and Beidi Chen. Tri- force: Lossless acceleration of long sequence generation with hierarchical speculative decoding. First Conference on Language Modeling. 2024.
[3D Yutao Sun, Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang. and Furu Wei. Retentive network: A successor to transformer for large language models. ArXiv preprint, 00//2307.08621.. 2023.
[4] Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang. Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder ar- chitectures for language models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
121 Han Shi, Jiahui Gao, Xiaozhe Ren, Hang Xu, Xiaodan Liang, Zhenguo Li, and James Tin- Yau Kwok. Sparsebert: Rethinking the importance analysis in self-attention. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Con- ference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 9547-9551 PMLR, 2021.
[Sha19] Noam Shazeer. Fast transformer decoding: One write-head is you need. ArXiv preprint, abs/191 2019.
16

---

[TDT+23] Yi Tay, Mostafa Dehghani, Vinh Q. Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Dara Bahri, Tal Schuster, Steven Zheng, Denny Zhou, Neil Houlsby, and Donald Metzler. UL2: Unifying language learning paradigms. In The Eleventh International Conference on Learning Representations, 2023.
[TKC19] Philippe Tillet, Hsiang-Tsung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages. pages 10-19, 2019.
[tri23] Triton implementation of the flash attention v2. Technical report, OpenAI, 2023.
[TSP+23] Szymon Tworkowski, Konrad Staniszewski, Mikolaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Mitoš. Focused transformer: Contrastive training for con- text scaling. In seirth-eeventh Conference on Neural Information Processing Systems, 2023.
[TZZ*24] Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. QUEST: Query-aware sparsity for efficient long-context LLM inference. Forty-first International Conference on Machine Learning, 2024.
[Wen23] Lilian Weng. Llm-powered autonomous agents. illianwengggitubb.i io. 2023.
[WWL 24] Zhongwei Wan, Ziang Wu, Che Liu, Jinfa Huang, Zhihong Zhu, Peng Jin, Longyue Wang, and Li Yuan. LOOK-M: Look-once optimization in KV cache for efficient multimodal long-context inference. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, Findings of the Association for Computational Linguistics: EMNLP 2024, pages 4065-4078, Miami, Florida, USA, November 2024. Association for Computational Linguistics.
[WZH21] Hanrui Wang, Zhekai Zhang, and Song Han. Spatten: Efficient sparse attention archi- tecture with cascade token and head pruming. In 2021 IEEE International Symposium on PrifoPPerforaacce Computer Architecture (HPCA), pages 97-110. IEEE, 2021.
[XTC*24] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Effi cient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2024.
[XZH*24] Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, and Maosong Sun. InfLLM: Training-free long-context extrapolation for LLMs with an efficient context memory. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[YCL++24) Alex Young, Bei Chen, Chao Li, Chengen Huang, Ge Zhang, Guanwei Zhang, Heng Li, Jiangcheng Zhu, Jianqun Chen, Jing Chang, et al. Yi: Open foundation models by 01. ai. ArXiv preprint, abs/2403. 04652, 2024.
[ZAW24] Itamar Zimerman, Ameen Ali, and Lior Wolf. A unified implicit attention formulation for gated-lines linear recurrent sequence models. ArXiv preprint, abs/2405.16500 2024.
[4] Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Hao, Xu Han, Zhen Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. xcBench Extending long context evaluation beyond 100K tokens. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Longg Papers), pages 15262-15277, Bangkok, Thailand, 2024. Association for Computational Linguistics.
[ZGD+20] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12, 2020, virtual, 2020.
17

---

[3]Z*23] Ningxin Zheng, Huiqiang Jiang, Quanlu Zhang, Zhenhua Han, Lingxiao Ma, Yuqing Yang. Fan Yang, Chengruidong Zhang, Lili Qiu, Mao Yang, et al. Pit: Optimization of dynamic sparse deep learning models via permutation invariant transformation. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 331-347, 2023.
[2SZ*24] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy- oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024.
18

---

A Limitations
As the context length decreases, the time required to build the dynamic index becomes more significant attention computation time decreases. For example, with 10k context, the time spent building the index increases from 5% to 30%, resulting in overall end-to-end latency approaching that of FlashAttention. However, this overhead proportion gradually decreases the prompt lengthens. Additionally, when using a higher sparsity rate, the model performance may noticeably decline.
B Broader Impacts
Minference effectively accelerates the inference of long-context LLMs, facilitating their deployment and application. By enabling lower latency, it can reduce the deployment costs of L.LMs, especially for long-context LLMs, helping to democratize access to advanced AI. It also promotes further research and development in related fields.
C Experiment Details
C.1 Dataset Details
InfiniteBench [ZCH+24) includes 10 tasks designed test various aspects of long-context pro- cessing. Specifically, these tasks cover entire novel summarization, open-form question answering based on novels, multiple-choice question answering on novels, question answering on long drama scripts, question answering on Chinese texts, debugging large code repositories, identifying the largest/smallest number arrays, and retrieval tasks with varying pattern lengths. The average token length for these tasks is 214k, and they include 3,992 examples.
is recent synthetic benchmark suite for long-context evaluation with 13 complex tasks across four categories. The retrieval category includes Single Needle-in-a-Haystace (S-NIAH), where a single key-value pair is inserted into noisy text, and the model must retrieve the value. Multi-keys fieedle-in---Haystae (MK-NIAH) involves multiple keys, and the model retrieves one specific value among hard distractors. The Multi-values Needle-in-a-Haystedd (MV- NIAH) task requires retrieving all values associated with single key, while the Multi- Multi-queries Beedle-in-a-Haysttcee (MQ-NIAH) task involves retrieving values for multiple keys. The Multi-hop hop Tracing category includes Variable Tracking (VT), where the model traces and returns variable names pointing to the same value through variable bindings. The aggregation category introduces Common Words Extraction (CWE), where the model identifies the top-K common words from a mixture of common and uncommon words, and Frequent Words Extraction (FWH), where the model identifies the most frequent words from Zeta distribution. The Question Answering (QA) category extends existing short-context QA datasets by adding distracting paragraphs, challenging the model to answer questions based on relevant information surrounded by distractors. These tasks provide comprehensive evaluation of long- long-context modeling capabilities. covering multi-hop reasoning, aggregation, and complex question answering. Following [HSK*24], we test models on 4K, 8K, 16K, 32K, 64K, and 128K context lengths, including 2,600 examples per length.
RULER HHKK* 24
Needle In A Haystack task [Km223] evaluates the performance of sarieval-augmentt generation (RAG) systems by embedding specific, targeted information (the "needle") within a large, complex body of text (the "haystack"). The test assesses a language model's ability to identify and utilize this specific piece of information amidst a vast amount of data. Both RULER and the needle test iterate over various context lengths and document depths (where the ground-truth is placed in the prompt) to measure the cong-context performance. Here we scale the Needle A Haystack task to 1M context length, including 750 examples.
The perplexity on long text is also often used by researchers to evaluate the language modeling performance of long-context LLMs. PG-19 a suitable test set for this task, as includes texts as long as 500K tokens. Perplexity is used as the metric indicating how well model predicts the next token in a sequence. Our experiments are conducted on 1, 1,000 random samples from PG-19 that are longer than 100K tokens.
PG-19 [RJ] 20
19

---

C.2 Additional Implementation Details
Our experiments are based number of state-of-the-art long-context LLMs: 1) LLaMA-3-8B Instruct-26244 is LLaMA-3 variant with further -TK-aware interpolation and minimal fine-tuning with Ring Attention, which achieved SOTA results on long-context assessments such as the Needle --a----------ntuut-----
In A Haystack test; 2) but supports context lengths up to 1M tokens; 3) Yi-9B-200K
is similar to ALaMA-3-8DB 3-8B-Insttuutt2622
24] is SOTA LLM that balances long-context performance with general capabilities; 4) Phi-3-Mini- Mii--- [AJA* a small but powerful language model that offers capabiliities equivalent models ten times its size with up to 128K context window powered by LongRoPE, ID7Z.*24 5) Bwvn2-7B-128K [BBC+23 recently release update of Qwen series model with up to 128K context window that achieve superior or comparable performance compared to LLaMA-3; 6) GLM-4-9B-IM [GZX+24] has been improved from its predecessor terms of 1M context window, performance on downstream tasks and inference efficiency. To guarantee stable results, we use greedy decoding in all tests. Our kernel implementations are developed and optimized based on the dynamic sparse compiler PIT [ZJZ 23 in the Triton language (TKC19). The latency experiments are done single Nvidia A100 GPU using 66. We provide a simple custom implementation of attention in PyTorch, building on FlashAttention and Triton.
We set the target FLOPs to be the same as 1k global tokens and 4k local window tokens in the A A-shape pattern. The step size of ChangeSpace is set to 50, with the corresponding search space shown in Table [7] Additionally, we use only one sample as our validation set from KV retrieval synthetic data with 30k token inputs, which exhibits strong generalization and stability across different lengths and domains. The search time approximately
minutes on a single A100. Additionally, model
we use the same optimal sparse pattern configuration for both the and the L.aMMA-3--B-Insruue IM model. The specific distribution is shown in Fig.
Table 7: Kernal-aware optimal head pattern search space. In this context, A A-saae represents the global tokens and local window number, Vertical- Slash represents the Top-K K number of vertical and diagonal lines, and Block-Sparse represents the Top-K number of blocks retained.
Patterns
Search Space
A-shape
{11024, 4096))
Vertical-Slasi ((30, 2048), (100, 1800), (500, 1500), (3000, 200) (100)
Block-Sparse
C.3 Single A100 Implementation Details
The original PyTorch implementatior of the LLaMA model causes an out-of-memory error on a single A A100 (80G) when the prompt exceeds 50k tokens. To enable running 1M prompt inference on a single A 100, we implemented the following optimizations:
1. Tensor Splitting: We split the Attention by head and the MLP by sequence dimension. In Long-context scenarios, where computation is the bottleneck, this splitting keeps GPU utilization at 100%, and the overhead of splitting is negligible;
2. Reduction of Intermediate Variables: We minimized intermediate variable allocation by removing the attention mask and implementing causal mask logic directly within the kernel;
3. Elimination of Unnecessary Computations: In long-contex: scenarios, only the logits corresponding to the last token in the prompt phase are meaningful. Thus, only retain the computation of the LM Head Linear layer for the last token.
"htts://huggingfacccco Aag.
http://wwggngggacooooooooaimiiaa colgradienaai/llama-- hanaaaaaaaaa---------aa---nninnn 1048k hhupss//hugiingfaee https///w....gooooeeeeee//////me Ilama.py
20

---

C.4 Kernel Implementation
C.4.1 Block-Sparse Flash Attention
Our MlokkSSprse kernel implementation is based on the Triton version of the FlashAttention ker- nel [n13]]. With the selected block index as an additional input, each thread block loops through the top-K blocks in a row. As discussed in FlashAttention [Dao24), the lateney of the block-sparse FlashAttention kernel is linearly related to the number of blocks, and the speedup ratio (compared to dense FlashAttention kernel) approximately as,
S
8p
2B Kb
C.4.2 Vertical-Slash Attention
The Vertical-Slat Slash attention includes custom kernels: the Vertical-Slash sparse index kernel and the Pertical-Slash sparse FlashAttention kernel.
Figure 7: The dynamic sparse mask for the vertical-slash pattern using LLLMA-3-8B in the summarization task IZCH 24]. Yellow areas indicate the computed parts. Slash lines use 64 blocks, while vertical lines use blocks.
The Vertical-Slast sparse index kernel in Algorithm builds the index for each row of blocks. Since slash line segment be masked by square block, attention mask is mix of blocks and columns, as shown in Fig. [7] We apply point-range two-way merge algorithm where vertical indexes treated points and slash indexes
converted to ranges given the index. The output consists of two parts: merged ranges and separate column indexes, where the ranges are represented by block indexes. The time complexity to build index for row O k..)
The Vertical-Slast sparse FlashAttention kernel in Algorithm 5 iis a mix of the block block-sparse attention kernel and the PIT [3] sparse attention kernel. PIT is technology that loads sparse data into dense compute blocks via Permutation Invariant Transformation. A thread block first loops through the block indexes described in the previous section (block part) and then loops through the column indexes grouped by block size (PIT part). The latency of this hybrid kernel is linearly related to the total area of blocks and columns.
D Additional Experiment Results
Needle in Haystack LLaMA-3-6B-IN w/ InFLLM 1M Context
D.1 Needle In Haystack
1.0
Dopth Percent (86)
22+
03
In addition to the Needle In A Haystack re-
33-
sults for LLaMA-3-Instructt-l- shown in $4
56-
we also present the
using
67-
0A
InfLLM results Fig. [8] and results for GLM-
78.
4-9B- 1M, Yi-9B- 200K, Phi-3-Mini--122
and
100-
00
Qwen2-7B-128K, shown in Fig. Compared to Full Attention, using MInference has minimal impact on the ability to understand semantic in- formation across different context windows and
220e
2000
5700
72
Context Length
Figure 8: Results on Needle In A Haystack using In-
[LLM in LLaMA-3-SD B----B-rettt-t111
21

---

needle depths. There is even a slight performance improvement around the 100k context length using Yi-9B-200K and Phi-3-Mini- 128K.
Needle Hayslack GLM-4-98-1M 1M Context
1.0
Noodie
0-
Haysttakk GLM 4.444.90 1M w/ Minference 1M Context
1.0
Depth Percert (%6
33-
44-
33
44-
0.6
67-
89-
78-
84-
02
100
7DC
2000
STow
Prok
soor
3000
S0
Syroe
500r.
Context Length
Co Context Length
(a) GLM-4-9B-IM
(h) GLM--99-IMM w/ MInference
Needle in A Heystack M-99- 200K Context
Needle in Haystack Yi-98-200K w/ Minference 200K Context
0-
1.
11-
-08
Death Rerent
33-
44-
0.6
46
04
56-
89
-0
80-
02
100-
100
204
300
400
2300
sook
2004
20r
30.
40x
6oc
7or
Tooe"
Zoor"
soor.
Context Length
Context Length
(c) Ti-9B-2008
(d) 9B-9K- w/ MInference
0-
Needle in Haystack Phi-3-Min Mini 128K 128K Context
1.0
Neodle in A Haystack Phi Mini 128K Mini-112aK w/ Minference 128K Context
11-
Depth Percent
33-
44-
-06
23
48
06
56.
84-
80-
-02
100
00.0
100
201r
120r
120c
10c
15c
52
201r
2iter
120r
Context Length
Context Length
(e) 1--Miiii2288
(f) Phi-3-Mini-11 128K w/ Minference
Haystack Qwen2-78-122
Needle in Haystack Owen2: 78 12-1 w/ Minference 128K Context
Dejth Percent (%6
22-
11-
+0.8
56-
67-
-0.4
0A
02
100
..0
100
0.0
200k
lor
2,00
2.2
Contoxt Longth
Contoxt Longth
(g)
Qwen2- Bren2-7B-128K
(h) Qwen2-7B-1128 MInference
Figure 9: Needle In A Haystack Kam23 results using GLMM4-9BB-M GZX 24). Yi-9B-200K [YCI+24)
Phi-3-Mini-128K AJA 24 and Qwen2-7B-128K (BBC*23)
D.2 Latency Breakdown
Fig. 10 shows the miiro-benchmaat results of the three attention patterns proposed in this paper, as well as FlashAttention. can be seen that Vertical-Slash the slowest among the three patterns, but it still achieves a 13x speedup compared to FlashAttention under 1M context windows. A- shape is slightly faster than Vertical-Slash, but IM, A-shape is 50% slower than Vertical-Slan Block- Block-Sparse is the fastest, achieving 30x speedup over FlashAttention under IM context windows.
22

---

A-shape
Block SpakkSparse
Flashättention
Vertical-Slash
Block-Sparse Index
InTLLM
VerticalSlash Index
1500
500
CANTIN 13
100
80-
Latancy(ms)
60-
40-
20+
100K
600K 700K Windows
Figure 10: The latency breakdown of a single attention kernel for three patterns and FlashAttention [ao24 across different context windows single A A100, including the index time for dynamic sparse approximation building dynamic sparsity. 10k tokens, latency the four kernels is very close and all are less than Ims. At 1M tokens, the latency for A-shape is 164ms.
The estimation and index- index-building time for the dynamic sparse pattern accounts for approximately 5%-15% and 25% of the total time for Vertical-Slash and Block-Sparse patterns, respectively. The index-building overhead higher for Block Block-Sparse mainly due to the time-consuming MeanPooling and block-level matmul computations. Additionally, the memory overhead for sparse indexing is relatively small, remaining within 160MB for a LLaMA-3-88 model in 1M context.
D.3 Additional Ablation Study
Table 8: Performance of different ablation methods using --0------8----tttu--66266 on In- finiteBench (4) It is important to note that due to kernel limitations, we must retain least one vertical and slash. Therefore, "ours w/ only vertical" retains the top-1 slash, and "ours w/ only slash" retains the top-1 vertical.
Methods
[En.Sum En EA.QA Ma En.
QA Code Debug Math.Find Retr.PassKey Retr.Num Retr KV Avg.
Ours
20.5
12.9
65.9
7.5
12.5
22.3
33.1
100.0
100.0
12.8
38.8
Ours w/ only vertical
13.7
6,2
30.1
2.0
7.9
65.4
52.7
18.6
Ours only slash
18.4
11.5
60.1
3.0
114
22.1
28.4
100.0
100.0
35.9
To further analyze the role of dynamic vertical and slash lines in the Vertical-Slash pattern for sparse computation, we introduce new set of ablation studies as follows: 1) Ours w/ only vertical, which only vertical lines and the top-1 slash line in Vertical-Slash pattern. 2) Ours w/ only slash, which only uses slash lines and the top-1 vertical line Vertical-Slash pattern. The corresponding top-K quantities are set after converting based on FLOPs in kernel.
As shown in Table X using only vertical lines results significant performance drop, especially in retrieval tasks, where performance is similar to only using block-sparse. In contrast, using only slash lines retains most of the performance, but in highly dynamic tasks such as KV retrieval, performance further decreases, with an average performance drop of 2.9% compared to Ours.
E Pattern Distribution
Fig. II shows the distribution of the optimal head configuration obtained through our search. Firstly, most of the patterns are the Vertical-Slas pattern (>90%). However, according to the ablation study, using only the Vertical-Slash pattern significantly impacts performance in highly dynamic tasks like KV retrieval. Secondly, the Block-Sparse pattern is primarily distributed in several intermediate to later layers, while the A-shape pattern is found in the middle layers. Although the optimal patterns vary slightly across different models, they generally align with these observations.
Additionally, used the same configuration for two versions of LLaMA in our experiments, and the results show that the IM model also performs very well, with nearly perfect results in the Needle In A Haystack task. This demonstrates the generalizability of the optimal sparse pattern.
23

---

A-shape
Vertical Slash
Dock-Sparee
A A-shope
Vertical Slash
Block Sparse
30
30
Head Number
Head Number
20
20
13 2309222332555
7 1013161922223 34
Layer
Layer
(a) LLaMA-3 A-M-MA--------- 3-88--instmett-662KN
(b) Yi-9B-200K
Figure Distribution three sparse head patterns different models. use the same optimal sparse pattern configuration for both LLaMA-3-8DB LLLMMA---S--nnstuut-222222 and LLaMA-3-8B- Ins-B---tttutt-MM...
Sparsity in Kernel Distribution
0.95
0.9
Sparsity in Kernel
0.8
0.7
0.6
0.5
A-shape
Vertical-Slash
0.4
Block-Sparse
10k100k
300k
500k
1M
Context Windows Size
Figure 12: The distribution of sparsity in the kernel across different context windows refers to the proportion of the kernel that is actually computed after block coverage, compared to the sparsity rate when using FlashAttention with causal mask.
As shown in Fig. 12 the sparsity distribution of the three patterns during the actual kernel computation process is displayed. It can seen that when the context windows exceed 200k, the actual sparsity of all three patterns surpasses 90%. Even considering 20% index-building overhead, this ensures that the kernel achieves speedup of over 8x. Furthermore, when the context windows exceed 500k, the sparsity relative to FlashAttention exceeds 95%, with theoretical speedup of over 15x.
G Does This Dynamic Sparse Attention Pattern Exist Only in Auto-Regressive LLMs or RoPE-Based LLMs?
Similar vertical and slash line sparse patterns have been discovered in BERT |SGR and multi- modal LLMs /WWL*22 Additionally, shown in Fig. 13 analyzed the distribution of attention patterns T5 across different heads. It evident that there are vertical and slash sparse patterns even in bidirectional attention.
Recent studies [WWL++24) have analyzed sparse attention patterns multi-modal LLMs, re- vealing the presence of vertical and slash patterns in models like LLaVA [LIWI.24] and In- te (CWW+24). Using MInference for pre-filling stage inference acceleration holds great promise.
24

---

Figure 13: The sparse pattern in TS-style Encoder Attention using Flan-UL2 TDT*23 the Summarization
dataset ZCH
H
Case Study
Table presentss a comparison of the generation performance for various methods on the EN.SUM task (200K input length) from InfiniteBench based the LLaMA-3-8B-2622 model. The orig- inal summary provides a comprehensive and coherent narrative, detailing the Bronwyn family's s trip to the Kindergarten and touching themes such nostalgia, loss, and the passage of time. StreamingLLMM's summary, although looks coherent, introduces elements that are not present the original story, leading to serious factual errors. For example, it mentions boat trip to school for boys and specific details like fishermen, sandwiches, and a spot where men were drowned. These details deviate from the original story, which is about the Bronwyn family preparing for trip the Kindergarten. In addition, the summaries generated by Streamingl with dilated and strided tech- niques largely incoherent, consisting primarily of repetitive and nonsensical characters, indicating a failure to produce meaningful content. In stark contrast, the summary generated by our proposed method offers detailed and coherent narrative, comparable to the original, with clear depiction of the story's 's main events and themes. This includes the preparation of the Bronwyn family for their trip, the characterization of family members and guests, and the exploration of deeper themes such as love, marriage, and the search for meaning. The results demonstrate the superiority of our proposed method in generating high qughttyalty, numan-like summaries over the bascline methods.
Table compares the performance of various methods on the Retrieve.KY task (200K input length) using the LLaMAA---BB-6222 model. The original method demonstrates perfect retrieval, correctly predicting the exact strings of the ground truth for both examples. StreamingL.LM again, generates predictions that looks coherent and real, but factually incorrect. In addition, StreamingLLM with dilated and strided techniques, and our method with static pattern, fail significantly, producing outputs that are either repetitive sequences of characters nonsensical strings, indicating their inability to accurately retrieve the required key-value pairs. Our method, however, performs on par with the original, accurately retrieving and predicting the exact key-value pairs for both examples. This demonstrates the superior capability of our method in handling KV retrieval tasks, providing precise and reliable outputs consistent with the ground truth. The results highlight our method's effectiveness and robustness compared to the baselines, making it reliable choice for such tasks.
25

---

Table 9: Comparison of generated results by different methods using LLaMA-3-888 ------------nttt--22222222 in summariza- tion task IZCH*24
Methods
Summary
Original
The story follows the Bronwyn family and their guests as they prepare for a day trip to the Kindergarten, mysterious destination that holds great significance for them. The story is told from multiple perspectives, including that of Pete Lucy, who staying with the family and is struggling find her place in their lives, As the day progresses, the family members prepare for their trip, and Pete reflects on her own past experiences with the Bronwyns. The story ends with the family setting off on their journey. leaving Pete alone on the lawn to contemplate her own feelings and the meaning of life. The Kindergarten is symbol of the family's past and their memories, and the story explores themes of nostalgia, loss, and the passage of time. Ultimately, the story suggests that even though the past cannot be changed, it can still be celebrated and remembered.
StreamingLLMM
The story follows a family on a boat trip to the Kindergarten, school for boys. The father, Mr Bronwyn, is reading book and seems to be in trance. The children, Seb and Annalisa, They pass by group of fishermen Mr Bronwyn shares out sandwiches. They come across a spot where three men were drowned and Bronwyn takes out his watch, making a calculation. They reach the Kindergarten and Mr Bronwyn jumps out of the boat, holding parcel. The story ends with Pete Lucy, who has been watching the family from the shore, feeling relieved that the trip that Bronwyn has reached his destination. The story told from multiple perspectives, including Pete Lucy's, 's who watching the family from the shore and reflecting on his own life and relationships.
excited about the trip but also little scared.
and
StreamingLLM w/ dilated
StreamingI.IMM w/ strided
Graphics
Ours w/ static
def solverome2 def solverome2 def solverome2 def solverome2 def solverome2 def solverome2 def solverome2 def solverome2 def solverome2 def solverome2 def solverome2 def solverome2 def
Ours
The story follows the Bronwyn family and their guests as they prepare trip to the Kindergarten, a remote island off the coast of Scotland. The family is made up of Mr. Bronwyn, philosopher, and his wife. Mrs. Bronwyn, who is known for her beauty and her ability to surround and protect her family. The guests include Kevin Emma, young man who is studying for a fellowship, and Castor Kiki, botanist who is also widower. The story explores themes of love, marriage, and the search for meaning in life. It also touches on the idea of the other, or the person who is different from oneself and cannot be understood or judged conventional standards. Ultimately, the story is about the search for truth and the desire connect with others on a deeper level.
26

---

Algorithm 4 Vertical-Slash Index
Algorithm Vertical-Slash Flash Attention
Input: vertical indexes i, E Nkv, slash indexes &, E Nk.
Input: Q,K, V E Rxdd, block count Cbik NN block index ésu E NN*k, col-
umn count Gool NN, column index iuo N xkv
Sort vertical and slash indexes
by f anremmnnaaSorf (i.)
4, DescendingSort (é.)
Scale 7
Calculate block number (block size B)
Initialize O - (0) E Sxds
N + -
# Parallelized in GPU
Initialize outputs
for i 4- 1 N do
block count Cos E N,,, block index dan E Nokk, column count Coall column index beal E N xck
Load Qctip Q Q(xB:(44)x R Rxdd
Initialize Osep (0) Bx) E Rxxd,
Parallelized GPU
Initialize m + inf) E I
for + 10-100 to N do Ju +-1
Initialize I + (0) IR
# Loop through h block indexes: block sparse flash attention
Find the first slash line that crosses the row
j. <- biset lefi( i.,ix x B)
for to Ook do
Define the range by slash index
Block start <- bikk Load K ship &- Ks E Bxde
Taan + (i 1) B 1
Fend
B 4
Load Vaup Vis+B E R Bxd,,
S - qaiip Khhip
# Merge points (vertical indexes) and ranges (slash indexes)
S mask(S)
m m + max(m' rowmax(S)) E R
while 8y S kig do
if jo S ki and 2 < radd then
S S mnew
P exp(S)
# Record the point if not the range
hew rowsum(S))
th Fant then
a t exp(mm minew)
dat -- ccal + 1, ci it!""
I' t- al' Inew
end
Ohiip aOO PVaip
j. -- j
end
else
+ su +1
# Loop through column indexes PIT sparse flash attention
# Update raagee
j +
if (i - -1) x B 1/ > rend then
while < c/u do
Record last range
cols +- Voo 3. NB
+ Paua
Load Ketip + Kcols E Bxdd,
while < Tend do
Load Vehip +- Veola E B Bxd,
Calk +- clul
S + aqui KTip
s.cia <-
bik
S - mask(S).
f- + B
m max(m' rowmax R
end while
S < S mhew
# Calculate the new range
P exp(S)
Tsuat
+ (i - 1) B 4/.
I + rowsum ())
exp(m" m)ew))
Fen
else
+ al' -new
chip +- a PV
Extend the range
j j+ B
Ten &- Vend
end while
end
end
# Write outputs
end while
O stip + diag(!') 'OOpip
# Record the last range
Save O, O
end for
Tstan while s Fend do
Cik &- Ciu +
..cbit --
B
end while
end for
return Colk, dosks Cool, dcol
27

---

Table 10: Comparison of generated results by different methods using LLaMM-2--BB--nssrut--22282 in KV Retrieval task [ZCH* 24
Methods
Summary
Original
Ground Truth: 5
Prediction: 256677990--104400333300022222220
2. Ground Truth: Tuuul
Prediction: "Iccshab--to77-44555554444111133 string of length 32).
StreamingLLM
Ground Truth: huttm:/cccccooooou------23333320
2b1 01110000000000--222--4-111111113
Ground hwu.hhhhhhhhh--22144446222222222
Prediction: "6852ee11-6639-47744444477737223
StreamingLLM w/ dilated
1. Ground Truth: Tuuphh/bbbbbbo.oo000000000033330
Prediction: 5 ((CCCCCCCCCCCCC
Ground Truth: floc3dd8--667-44555-55---- 9
Prediction: "def solverome2 def solverome2 def solverome2 def solverome2 def
solverome2 def solverome2 def solverome2 def solverome2 def solverome2 def solverome2
StreamingLLM strided
Ground Truth: 5c6b7690-71 1100-4953-990 3c96b2cadbi2
Prediction: "def solverome2 def solverome2 def solverome2 def solverome2 def solverome2
2. Ground Truth: fees:all---i7-444455555-0t00t000 Prediction: "O (CcdiciiooOO(((((((((((((((CCCCC ((C((((((((((CC))
Ours w/ static
1. Ground Truth:
Prediction: "def solverome2 def solverome2 def solverome2 def solverome2 def
Ground Truth: 1
Prediction: "def solverome2 def solverome2 def solverome2 def solverome2 def
Ours
1. Truth: hht.h:/ccchhho00000000030003-3-0
Prediction: 00667770001100------00--00000020
Ground Truth: hat.:ehhcclssse-4444a4aaaaaaaaa0
Prediction:
(a string of length 32).
28

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We discuss the contributions and scope in detail and accurately in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of MInference in Appendix A
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]

---

The methods discussed and proposed in this paper are based on empirical observations and do not involve theoretical proof.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide detailed descriptions of our experiments and implementation in $4 and Appendix C Additionally, we have released our code at lttps: //aka. /akaa.a///////
MI MInference
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
30

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We provide detailed descriptions of our experiments and implementation
in $4 and Appendix C
Additionally, we have released our code at https: //aka/aaaaaaaaa.a
MInference
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
healiii/ggies///Ce Cawelee//ooSSSbbbiicaa SorooddSumisssiooPP... wwiiiisssiiiicyoyioo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide detailed deseriptions of experiments and implementation in $4 and Appendix C
Additionally, we have released our code at hhtpsa // //aka. ms/
MInf erencel
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: We follow previous works by conducting tests on public benchmarks, without including error bars similar information.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
31

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide detailed descriptions of our experiments and implementation details including computation resources
$4 and Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics Ethics/https:
Answer: [Yes]
Justification:
We follow the guidelines of the NeurlPS Code Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there special consid- eration due to laws or regulations in their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the broader impacts of MInference in Appendix B
Guidelines:
The answer NA means that there is no societal impact of the work performed.
32

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not involve the release of new models or data.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The LLMs, datasets, and codebase used in our work comply with open-source licenses and can be used for scientific research.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
33

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We have released our code at batttss//aks. ms/ ms/MInference
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve data annotation or crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
34

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
35


Reviews:
[{"review_id": "eDQC5VR9Vb", "reviewer": "Reviewer_gjr6", "summary": "This paper proposes MInference, a method to accelerate the pre-filling stage for long-context LLM generation. The key method leveraged by MInference is dynamic sparsification, which consists of three sparse patterns observed in attention matrices: the A-shape pattern, the Vertical-Slash pattern, and the Block-Sparse pattern. A sparse pattern search method is developed to minimize the sparsification error for each attention. MInference has been tested across various LLMs and benchmarks, demonstrating that it significantly accelerates LLM generation in long-context settings with no performance drop.", "questions": "- How does MInference scale up to LLMs with larger sizes, e.g., 70B? And how does MInference perform for MoE-based LLMs, e.g., Mixtral 8x7B?\n- How general are the observed sparse patterns across LLMs?\n- Is MInference compatible with other inference engines, e.g., vLLM?\n- Are the Vertical-Slash Head and Block-Sparse Head part of CUDA kernels in MInference, or does MInference leverage other implementations of those sparse head kernels?", "limitations": "The limitations of the proposed method have been thoroughly discussed. And there is no potential negative societal impact of this work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and well-motivated.\n- Studying how to accelerate LLM generation for long-context settings is of great practical importance.\n- The experimental results of the proposed MInference method are promising.", "weaknesses": "- The experimented model scales are relatively small (e.g., LLaMA-3-8B and Yi-9B) and are only constrained to dense LLMs (i.e., not MoE).\n- It is not clear how general the observed sparse patterns are across LLMs of various scales.\n- The implementation details on how the proposed MInference interacts with existing CUDA kernels (e.g., flashattention) are not very clear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1. _\"For larger models and MoE\"_\n\nThank you for your suggestion. We have added experimental results for **LLaMA-3-70B-1M**, as shown in Table 1, where MInference maintains excellent performance in larger models, significantly surpassing StreamingLLM and InfLLM[1], and nearly matching full attention. Due to resource constraints, we could not provide MoE results during the rebuttal period, but based on our experience, modifications to FFN do not affect the sparse attention patterns. We will include MoE results in future versions.\n\n|Methods |  En.Sum | En.QA | En.MC | En.Dia | Zh.QA | Code.Debug | Math.Find | Retr.PassKey | Retr.Number | Retr.KV |  Avg. |\n|-|-|-|-|-|-|-|-|-|-|-|-|\n|***LLaMA-3-70B-1M*** | 20.7 | 10.3 | 84.2 | 9.5 | 14.0 | 33.2 | 61.7 | 97.0 | 100.0 | 34.0 | 46.5 |\n|StreamingLLM | 20.5 | 8.5 | 52.0 | **10.0** | 12.6 | 27.4 | 61.1 | 14.0 | 10.0 | 0.0 | 21.6|\n|InfLLM | **24.1** | 8.1 | 57.0 | **10.0** | 12.9 | 27.4 | 52.3 | **100.0** | **100.0** | 0.0 | 39.2|\n|**MInference** | 20.6 | **10.1** | **83.4** | **10.0** | **14.1** | **34.1** | **61.9** | **100.0** | **100.0** | **39.0** | **47.3** |\n\nTable 1. Performance of different methods with different base models on InfiniteBench.\n\n[1] InfLLM: Unveiling the Intrinsic Capacity of LLMs for Understanding Extremely Long Sequences with Training-Free Memory, 2024.\n\n2. _\"Generality ability across models\"_\n\nWe tested MInference on a variety of long-context LLMs, including LLaMA-3-8B-1M, Yi-9B-200K, GLM-4-9B-1M, Phi-3-mini-128K, and Qwen2-7B-128K, all of which maintained performance close to full attention (**see general response PDF**).\nAdditionally, we observed that three sparse attention patterns, especially the vertical and slash patterns, are present in BERT[2], T5 (both Encoder and Decoder), and MLLM[3]. We also found that \"induction heads\" exhibit a pattern similar to the \"vertical and slash\" pattern (see [4]).\nIn conclusion, we believe MInference is highly generalizable across different structures, sizes, pre-training data, training pipelines, RoPE structures, and extended long-context methods.\n\n[2] SparseBERT: Rethinking the Importance Analysis in Self-Attention, ICML 2021. \n[3] LOOK-M: Look-Once Optimization in KV Cache for Efficient Multimodal Long-Context Inference, 2024.\n[4] A Mathematical Framework for Transformer Circuits, 2021.\n\n3. _\"Compatible with other inference engines, e.g., vLLM\"_\n\nYes, we have implemented MInference within vLLM. Our method only requires replacing the corresponding FlashAttention kernel and does not affect the scheduling of vLLM PageAttention or similar features. We will open-source the implementation after the paper review.\n\n4. _\"CUDA Kernel\"_\n\nWe have utilized Dynamic Sparse Compiler PIT, Triton, and FlashAttention to implement CUDA Kernels for Vertical and Slash Head and Block-sparse Head, ensuring transferability and support for long-context LLMs inference across different GPUs and CUDA versions. Details can be found in Appendix C.4. We have submitted the relevant code with our submission and will open-source the implementation after the review.", "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": "1w3LSlfJya", "reviewer": "Reviewer_9cfV", "summary": "Targeting at the time-consuming profiling stage of long contexts, this paper proposed an efficient prefilling stage spare attention mechanism. It's based on the observation to common attention patterns. The proposed method can be integrated into most existing LLMs, such as LLama3 and Yi-9B. It achieves good performance compared to the original model while significantly reduce the computation times.", "questions": "1. Is the slash-vertical pattern distinguishable from lambda-shaped patterns? By my understanding, lambda-shaped patterns are just a special case of the slash-vertical pattern. Perhaps if you replace all lambda-shaped patterns with the slash-vertical pattern, it would still work well.\n2. Could you please do a comparison between the proposed MInference and H2O[1]? From my understanding, without considering the real efficiency, it’s hard for a method relying on attention scores/similarity to obtain significantly better performance than H2O.\n3. From my own experience, when the tasks become more challenging, the performance of efficient attention methods might show quite different effectiveness. For example, many long context methods cannot pass the passkey retrieval test with a 100-digit passkey, while with a 5-digit passkey, most methods work perfectly. Could you please construct experiments on more challenging tasks? Especially on those most recently proposed ones, for example: [2][3].\n4. Also, Llama-3 is much more robust to perturbation and error accumulation. Could you please construct experiments with some weaker LLMs, such as Mistralv0.2, Llama2, and gemma-1.1?\n\nReferences:\n[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models\n[2] KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches\n[3] BABILong: Testing the Limits of LLMs with Long Context Reasoning-in-a-Haystack", "limitations": "Yes", "rating": 7, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1: The motivation is clear and the paper is easy to understand. \n\n2: The proposed method seems simple yet effective.\n\n3: The experiments are comprehensive.", "weaknesses": "1: It’s tricky to claim that “we found the three patterns”. Those patterns have been found many years ago, since the years of Bert / transformer models were proposed. Even the idea of leveraging such sparsity of attention has been invented many times for different transformer based models. \n\nSome other major concerns can be found in the \"Question\" section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review, and we apologize for any confusion caused.\n\n1. _\"Claim that 'we found the three patterns'\"_\n\nThank you for your critique. We will revise the wording accordingly. Indeed, we discussed the importance of sparse attention in related works and how it inspired our research. We want to emphasize our understanding of the dynamic nature of these sparse patterns, which led us to propose the training-free MInference method.\n\n2. _\"Lambda is part of vs?\"_\n\nYes, the A-shape pattern is a specific type of \"vertical and slash\" pattern. However, the A-shape head exhibits a more fixed sparsity pattern across different tasks and examples. Due to its static nature, it does not require online sparse index building and allows for more extreme sparse kernel optimization, resulting in better inference acceleration. We differentiate these three patterns based on kernel speedup and spatial locality.\n\n3. _\"Compare with H20\"_\n\nIn fact, MInference and H2O optimize two different problems in long-context LLMs inference: prefilling stage latency and KV cache storage, respectively. These methods are orthogonal. Furthermore, we applied MInference with SnapKV in our paper (Table 5 in paper), which is a SoTA KV cache compression method superior to H20. MInference can be used in conjunction with KV cache compression methods to achieve even better performance.\n\n4. _\"More challenging tasks\"_\n\nThank you for your suggestion. However, I would like to clarify that we have tested on very challenging long-context benchmarks, including RULER, InfiniteBench, Needle In A Haystack, and Language Model, with lengths ranging from 128k to 1M tokens.\n- RULER includes subtasks such as Multi-Needle and Multi-hop Tracing and is one of the most challenging long-context benchmarks, sharing a similar tasks with BABILong and a similar benchmark ranking.\n- We have additionally conducted evaluations on the KV cache compression benchmark[5], including LongBench (refer to Table 2(b) below) and Needle In A Haystack (see Figure 8 in the paper and Figure 1 in the general response PDF), showing that MInference does not negatively impact performance.\n- InfiniteBench includes many challenging long-context tasks. For example, KV retrieval task requires LLMs to recall a 36-bit value from several random 36-bit Key-Value pairs by retrieving the value corresponding to a random 36-bit key, which is not as simple as the 5-digit passkey task mentioned in the comments.\n\nOverall, our tests have included similar or even more challenging long-context tasks, effectively reflecting the minimal impact of MInference on the capabilities of long-context LLMs.\n\n|**LongBench** | SingleDoc | MultiDoc | Summ. | FewShot | Synth. | Code | Avg. |\n|-|-|-|-|-|-|-|-|\n|*LLaMA-3-8B-262K* | 33.5 | 28.3 | 29.4 | 66.9 | 43.0 | 42.4 | 40.6 |\n|StreamingLLM | 29.7 | 23.7 | 28.8 | 65.8 | 19.3 | **42.8** | 35.0 |\n|InfLLM |  31.6 | 25.8 | 29.1 | 66.3 | 36.3 | 41.8 | 38.5 |\n|**MInference** | **34.0** | **28.0** | **29.7** | **66.8** | **42.8** | 42.2 | **40.6**|\n\nTable 2(b). Performance of different methods on LongBench.\n\n[5] KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches.\n\n5. _\"Weaker models\"_\n\nBesides testing on the powerful LLaMA-3-8B-1M, we also conducted tests on the Yi-9B-200K model, which has an effective context size of only 8k in RULER, showing that MInference can achieve performance close to Full Attention without changing the effective context size. Moreover, we have supplemented our results with GLM-4-9B-1M, Phi-mini-128K, and Qwen2-7B-128K (**see general response PDF**). Notably, the latter two models have weaker long-context capabilities, with effective context sizes of 4k and 8k in RULER, respectively. They show a significant performance drop in the Needle In A Haystack task when exceeding 100K tokens, but using MInference, they can achieve close or even better performance. We will include these results in future versions. Additionally, the weak LLMs mentioned by the reviewer have context windows smaller than 32K and cannot be tested on benchmarks like InfiniteBench. We will adapt more long-context LLMs in the future.", "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": "PFRXsqpy3z", "reviewer": "Reviewer_MtKP", "summary": "A key challenge for LLM inference with processing long context lengths is time-to-first token for long prompts. This paper introduces a sparse attention method designed to accelerate prefill with long context lengths. They utilize a strategy that incorporates three different types of sparsity (A-shape, vertical-slash, and block-sparse). They calibrate for which sparsity pattern to apply to different heads, thereby adapting to the varying sparsity patterns across LLM heads. Their method is training-free and obtains significant latency gains in terms of time to first token.", "questions": "- Do you have any analysis that shows that the attention head patterns are the same regardless of the input task (or that their search converges to the same solution for the same head across different tasks)?", "limitations": "Yes", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Time to first token is a significant problem that is under-addressed in existing KV cache compression works, as many applications have long input prompts and short generation lengths\n- Their approach is able to adapt dynamically to different inputs for certain attention patterns (which have different sparsity patterns), as well as to handle differing behaviors across attention heads\n- They provide efficient online kernel implementations for each sparsity pattern (both for constructing the sparsity pattern for the two dynamic patterns as well as for sparse computation)\n- They observe significant prefill speedups attained with minimal accuracy degradation\n- They provide detailed evaluation across a range of long context tasks, as well as ablation for each type of sparsity pattern (and additional analysis showing the distribution of sparsity patterns across attention heads, etc.)", "weaknesses": "- Their approach is heuristic-based (as it is based on fixed types of attention patterns), which may not generalize to new models that are released (some of these patterns could be specific to models that employ particular positional encodings or other model architecture features, which may limit generalizability)\n- The offline kernel search assumes that each head will play the same role regardless of the target task (or alternatively that calibration data will always be available for the target task in order to pre-determine which general attention pattern to use and what budget to allocate to that head)\n- The writing in some sections is poor (eg. end of Section 2)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1. _\"...which may not generalize...\"_\n\n- To demonstrate the generalizability of MInference, we tested it on most open-source long-context LLMs, including LLaMA-3-8B/70B-1M, Yi-9B-200K, GLM-4-9B-1M, Phi-3-mini-128K, and Qwen2-7B-128K (**see general response PDF**). MInference consistently achieves good performance and acceleration across models with different pre-training data, training pipelines, RoPE structures, extended long-context methods, and sizes.\n- Additionally, we observed that three sparse attention patterns, especially the vertical and slash patterns, are not only present in GPT-like LLMs but also in BERT[2], T5 (both Encoder and Decoder), and MLLM[3]. We also found that \"induction heads\" exhibit a pattern similar to the \"vertical and slash\" pattern (see [4]).\n- By design, MInference accounts for different dynamic sparse attention patterns, which lends it a certain degree of generalization capability.\n\n[2] SparseBERT: Rethinking the Importance Analysis in Self-Attention, ICML 2021. \n[3] LOOK-M: Look-Once Optimization in KV Cache for Efficient Multimodal Long-Context Inference, 2024.\n[4] A Mathematical Framework for Transformer Circuits, 2021.\n\n2. _\"... assumes that each head will play the same role regardless of the target task...\"_\n\nIndeed, based on our experiments and observations, these dynamic sparse patterns are relatively fixed across different tasks and examples. This is corroborated by our experimental results, where a configuration searched from a single example can achieve performance close to full attention across various scenarios and benchmarks.\n\n3. _\"The writing in some sections is poor (e.g., end of Section 2)\"_\n\nWe apologize for any confusion caused by our writing and will rewrite the relevant content, particularly Section 2.2, to improve readability.\n\n4. _\"...analysis that shows that the attention head patterns are the same regardless of the input task...\"_\n\nThank you for your suggestion. We explain the relative stability of head-level dynamic sparse attention patterns from several perspectives:\n- Visualization: Fig.3(a) shows the visualization of the sparse attention pattern distribution of the same head across different tasks and examples, demonstrating good consistency.\n- Attention Recall: Fig.3(c) shows the average recall rate of different tasks and examples within the searched pattern as the compression rate changes. The searched pattern achieves a higher recall rate across different tasks.\n- Downstream tasks performance: Using a configuration searched from a single example, we achieve performance close to full attention across different scenarios and benchmarks, further indicating the stability of dynamic sparse attention patterns across various tasks and examples.\nWe will rewrite the paper to highlight this point and thank you again for your helpful and insightful comments.\n\n|**RULER** | Effective | 4K | 8K | 16K | 32K | 64K | 128K | Avg. |\n|-|-|-|-|-|-|-|-|-|\n|*GLM-4-9B-1M* | 64K | 93.8 | 91.6 | 89.3 | 87.4 | 85.2 | 80.8 | 88.0 |\n|StreamingLLM | 4K | 93.8 | 66.9 | 58.5 | 51.4 | 45.9 | 39.1  | 59.3 |\n|InfLLM |  8K  | **94.7** | 89.5 | 76.4 | 66.5 | 56.8 | 53.5 | 72.9 |\n|**MInference**  | 64K | 94.6 | **93.1** | **91.0** | **89.6** | **85.5** | **84.0** | **89.6**|\n\nTable 2(a). Performance of different methods on RULER.", "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": "gfN4AV4f0E", "reviewer": "Reviewer_nzY9", "summary": "This paper presents a sparse calculation technique for the attention mechanism in long-context large language models during the pre-filling stage. Specifically, the technique builds on the authors' observation of three patterns in the attention map: the A-shape pattern, the Vertical-Slash pattern, and the Block-Sparse pattern. In the proposed method, the authors first determine the optimal pattern for each attention head offline and then dynamically decide the hyperparameters for each pattern (e.g., sparse indices) on-the-fly. It is worth noting that the A-shape pattern is more regular in terms of the distribution of non-zero indices compared to the Vertical-Slash and Block-Sparse patterns. To provide a low-cost estimation of the non-zero indices in the Vertical-Slash and Block-Sparse patterns, the authors predict these indices using the matrix multiplication results between the last query vector and the key matrix and the matrix multiplication results between a mean-pooled query matrix and the key matrix, respectively. Results on three different pre-trained LLMs and four datasets indicate that the proposed method effectively maintains information from long prompts while reducing pre-filling latency compared to current state-of-the-art methods.", "questions": "1. What is the dataset of input data used for the analysis in Section 2.2? How was it chosen?\n\n2. In Figure 3(c), why does the block-sparse pattern perform so poorly in attention compared to the A-shape and Vertical-Slash patterns? Based on my understanding, the block-sparse pattern is a more general and dynamic case for the A-shape and Vertical-Slash patterns, as shown in Figure 4. Could the block size be an important factor causing this issue?\n\n3. In Figure 3(c) and the corresponding description, I suggest explicitly mentioning which Top-K method is used. Based on my understanding, the Top-K method involves selecting the top-K tokens for all heads, while the other methods select the top-K tokens for each head. Presenting them together may confuse readers into thinking that Top-K means the top-K for each head.\n\n4. What does \"real FLOPs\" mean in Line 137? FLOPs itself is a \"conceptual estimation\" compared to real measured latency. I believe the authors may want to refer to the FLOPs after discarding the zero values in the attention pattern. If so, please be more specific.\n\n5. The LLaMA-3-8B-Instruct-262k is linked to LLaMA-3-70B-Instruct-262k in Line 176. Please clarify which model is used here.\n\n6. There are minor consistency issues, such as StreamLLLM and StreamingLLM.", "limitations": "The authors discussed the limitation and social impact of their work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Impressive results: Based on the reported experimental results, the proposed method can reduce the pre-filling latency for a 1M context by ten times without compromising accuracy.\n\n2. Well-motivated: I appreciate the analysis and high-quality figures in Section 2, which explicitly illustrate the key takeaways for readers and effectively motivate the proposed method.\n\n3. Comprehensive summary in the related works section.", "weaknesses": "1. How to ensure that the observations on the three LLMs in this paper are generalizable to different LLMs? Since the entire method is built on the observation of the attention patterns in these three LLMs, the generalizability of this observation is crucial for the robustness of this work. For example, the block-sparse pattern is the most dynamic among the three. What if there are even more dynamic attention maps in other LLMs with specific inputs? How is the search space determined, and how general is it? The offline Kernel-Aware Optimal Sparse Pattern Search determines the pattern for each head offline, and the pattern does not change during inference regardless of the input. Is it always true that the attention pattern for a specific head remains nearly the same?\n\n\n2. Lack of evaluation on larger models: For longer contexts, larger models seem to be a better option to ensure quality. However, the models evaluated in this paper are relatively small (<10B parameters). Thus, it is unclear whether the improvements still hold for larger models.\n\n3. Lack of discussion on usage for the generation/decoding stage: Pre-filling is important because it determines the first token latency, but generation/decoding is also important because it determines the throughput. It would be beneficial to add some discussion on whether the observations and the proposed method are still applicable for the generation/decoding stage. For example, while the A-shape pattern might still exist, the sparse indices approximation in the vertical-slash pattern might not work for generation/decoding.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1. _\"...generalizable... built on the observation of the attention patterns...\"_\n\nThank you for your question. We address it from two angles: the generalization of MInference and the relative stability of dynamic sparse attention patterns across different examples.\n\nFor the generalization of MInference,\n- To demonstrate the generalizability of MInference, we tested it on most open-source long-context LLMs, including LLaMA-3-8B/70B-1M, Yi-9B-200K, GLM-4-9B-1M, Phi-3-mini-128K, and Qwen2-7B-128K (**see general response PDF**). MInference consistently achieves good performance and acceleration across models with different pre-training data, training pipelines, RoPE structures, extended long-context methods, and sizes.\n- Additionally, we observed that three sparse attention patterns, especially the vertical and slash patterns, are not only present in GPT-like LLMs but also in BERT[2], T5 (both Encoder and Decoder), and MLLM[3]. We also found that \"induction heads\" exhibit a pattern similar to the \"vertical and slash\" pattern (see [4]).\n- By design, MInference accounts for different dynamic sparse attention patterns, which lends it a certain degree of generalization capability.\n\n[2] SparseBERT: Rethinking the Importance Analysis in Self-Attention, ICML 2021. \n[3] LOOK-M: Look-Once Optimization in KV Cache for Efficient Multimodal Long-Context Inference, 2024.\n[4] A Mathematical Framework for Transformer Circuits, 2021.\n\nFor the relative stability of dynamic sparse attention patterns across different examples,\n- Visualization: Fig.3(a) shows the visualization of the sparse attention pattern distribution of the same head across different tasks and examples, demonstrating good consistency.\n- Attention Recall: Fig.3(c) shows the average recall rate of different tasks and examples within the searched pattern as the compression rate changes. The searched pattern achieves a higher recall rate across different tasks.\n- Downstream tasks performance: Using a configuration searched from a single example, we achieve performance close to full attention across different scenarios and benchmarks, further indicating the stability of dynamic sparse attention patterns across various tasks and examples.\n\n2. _\"What if there are even more dynamic attention maps in other LLMs with specific inputs\"_\n\nGiven the extreme sparsity of attention in long-context scenarios, where the significant attention weights are concentrated in a few blocks, even if the graph differs from the vertical and slash pattern and the A-shape, it can still be covered by the block-sparse pattern.\nAdditionally, our experiments in Question 1 demonstrate the ubiquity of this dynamic sparse pattern in current long-context LLMs and the effective of MInference.\n\n3. _\"How is the search space determined, and how general is it?\"_\n\nThe search space is determined by a specific sparsity rate and adjusted according to the actual FLOPs required for different patterns in the kernel. Based on testing, the search space obtained is quite consistent across different tasks and examples.\n\n4. _\"Lack of evaluation on larger models\"_\n\nThank you for your suggestion. We have included results for **LLaMA-3-70B-1M** in the supplementary PDF. MInference achieves performance close to or even better than full attention, especially compared to baselines.\n\n|Methods |  En.Sum | En.QA | En.MC | En.Dia | Zh.QA | Code.Debug | Math.Find | Retr.PassKey | Retr.Number | Retr.KV |  Avg. |\n|-|-|-|-|-|-|-|-|-|-|-|-|\n|***LLaMA-3-70B-1M*** | 20.7 | 10.3 | 84.2 | 9.5 | 14.0 | 33.2 | 61.7 | 97.0 | 100.0 | 34.0 | 46.5 |\n|StreamingLLM | 20.5 | 8.5 | 52.0 | **10.0** | 12.6 | 27.4 | 61.1 | 14.0 | 10.0 | 0.0 | 21.6|\n|InfLLM | **24.1** | 8.1 | 57.0 | **10.0** | 12.9 | 27.4 | 52.3 | **100.0** | **100.0** | 0.0 | 39.2|\n|**MInference** | 20.6 | **10.1** | **83.4** | **10.0** | **14.1** | **34.1** | **61.9** | **100.0** | **100.0** | **39.0** | **47.3** |\n\nTable 1. Performance of different methods with different base models on InfiniteBench.\n\n5. _\"Decoding\"_\n\nThank you for your suggestion. However, extending to decoding would require significant system work for CPU offloading, which we have earmarked for future work. Nonetheless, our preliminary experiments indicate that the extrapolation generality of these patterns is also very good, as shown in:\n\n|Methods|En.Sum|En.QA|En.MC|Zh.QA|Code.Debug|Math.Find|Retr.Number|\n|-|-|-|-|-|-|-|-|\n|GLM-4-9B-1M|28.3|9.7|68.6|12.1|29.4|38.9|100.0|\n|MInference|28.8|9.6|68.6|12.0|30.7|39.1|100.0|\n|MInference in prefilling and decoding|28.0|9.4|67.3|11.9|31.7|39.1|100.0|\n\nTable 3. Performance of different methods on InfiniteBench using GLM-4-9B-1M.\n\n6. _\"Dataset used in Sec 2.2\"_\n\nIn Section 2.2, we used data from InfiniteBench to create Fig.3, including summarization, QA, Math, and retrieval tasks. Both (b) and (c) are results averaged after randomly selecting 10 examples from each subset.\n\n7. _\"Why does the block-sparse pattern perform so poorly?\"_\n\nThis is mainly because different patterns use different block sizes in the kernel. For example, the Vertical and Slash pattern can use a 1x64 block to process vertical lines instead of 64x64 in block-sparse. If this head pattern is processed using block-sparse, it would result in significant computational waste, leading to a lower attention recall at the same sparsity rate.\n\n8. _\"I suggest explicitly mentioning which Top-K method is used\"_\n\nWe apologize for any confusion caused by our writing. Yes, your understanding is correct; the Top-K mentioned in Fig.3(c) refers to token-level Top-K. We will add all corresponding granularities of TopK and rewrite Section 2.2 for better comprehension in the next version.\n\n9. _\"Real FLOPs\"_\n\nReal FLOPs refer to the FLOP values in the kernel after excluding non-zero calculations. Thank you for your suggestion; we will clarify this meaning in the next version.\n\n10. _Error link and typos_\n\nThank you for your detailed review. We will correct these issues in the next version of our paper.", "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:
We are grateful for the diligent efforts and insightful comments from the reviewers. Your suggestions have been incredibly valuable to our work. We will address and resolve these issues in our responses and in subsequent versions of our paper. Here we respond to some common questions and have included additional experimental results (**see attached PDF**), specifically:

1. _"For larger models"_

We have supplemented the results for **LLaMA-3-70B-1M** on InfiniteBench, demonstrating that MInference continues to perform exceptionally well in larger models, outperforming SoTA baselines such as StreamingLLM and InfLLM[1], and matching the performance of full attention.

|Methods |  En.Sum | En.QA | En.MC | En.Dia | Zh.QA | Code.Debug | Math.Find | Retr.PassKey | Retr.Number | Retr.KV |  Avg. |
|-|-|-|-|-|-|-|-|-|-|-|-|
|*GLM-4-9B-1M* | 28.3 | 9.7 | 68.6 | 39.5 | 12.1 | 29.4 | 38.9 | 100.0 | 100.0 | 41.0 | 46.7 |
|StreamingLLM | 27.7 | 6.4 | 40.2 | 12.5 | 10.8 | 27.7 | 21.1 | 97.1 | 25.6 | 0.6 | 27.0|
|InfLLM | 28.0 | 7.3 | 45.0 | 14.0 | 10.7 | 27.9 | **39.4** | 98.0 | **100.0** | 2.6 | 37.3|
|**MInference** | **28.8** | **9.6** | **68.6** | **38.5** | **12.0** | **30.7** | 39.1 | **100.0** | **100.0** | **43.0** | **47.0** |
|***LLaMA-3-70B-1M*** | 20.7 | 10.3 | 84.2 | 9.5 | 14.0 | 33.2 | 61.7 | 97.0 | 100.0 | 34.0 | 46.5 |
|StreamingLLM | 20.5 | 8.5 | 52.0 | **10.0** | 12.6 | 27.4 | 61.1 | 14.0 | 10.0 | 0.0 | 21.6|
|InfLLM | **24.1** | 8.1 | 57.0 | **10.0** | 12.9 | 27.4 | 52.3 | **100.0** | **100.0** | 0.0 | 39.2|
|**MInference** | 20.6 | **10.1** | **83.4** | **10.0** | **14.1** | **34.1** | **61.9** | **100.0** | **100.0** | **39.0** | **47.3** |

Table 1. Performance of different methods with different base models on InfiniteBench.

[1] InfLLM: Unveiling the Intrinsic Capacity of LLMs for Understanding Extremely Long Sequences with Training-Free Memory, 2024.

2. _"Generality ability across models"_

- To demonstrate the generality of MInference, we tested it on various open-source long-context LLMs, including but not limited to LLaMA-3-8B-1M, Yi-9B-200K, GLM-4-1M, Phi-3-mini-128K, and Qwen2-7B-128K, compared with additional baseline InfLLM. Our method consistently achieves good performance and acceleration across models with different pre-training data, training pipelines, RoPE structures, extended long-context methods, and sizes.
- We have also included results from LongBench to further substantiate the generality of MInference in scenarios around 32K tokens.
- Moreover, we observed that three sparse attention patterns, particularly the vertical and slash patterns, are not only present in GPT-like LLMs but also in BERT[2], T5 (both Encoder and Decoder), and MLLM[3]. We also found that "induction heads" exhibit a pattern similar to the "vertical and slash" pattern (see [4]).
- In summary, we believe MInference exhibits strong generality across various model architectures.

[2] SparseBERT: Rethinking the Importance Analysis in Self-Attention, ICML 2021. 
[3] LOOK-M: Look-Once Optimization in KV Cache for Efficient Multimodal Long-Context Inference, 2024.
[4] A Mathematical Framework for Transformer Circuits, 2021.

|**RULER** | Effective | 4K | 8K | 16K | 32K | 64K | 128K | Avg. |
|-|-|-|-|-|-|-|-|-|
|*GLM-4-9B-1M* | 64K | 93.8 | 91.6 | 89.3 | 87.4 | 85.2 | 80.8 | 88.0 |
|StreamingLLM | 4K | 93.8 | 66.9 | 58.5 | 51.4 | 45.9 | 39.1  | 59.3 |
|InfLLM |  8K  | **94.7** | 89.5 | 76.4 | 66.5 | 56.8 | 53.5 | 72.9 |
|**MInference**  | 64K | 94.6 | **93.1** | **91.0** | **89.6** | **85.5** | **84.0** | **89.6**|

Table 2(a). Performance of different methods on RULER.

|**LongBench** | SingleDoc | MultiDoc | Summ. | FewShot | Synth. | Code | Avg. |
|-|-|-|-|-|-|-|-|
|*LLaMA-3-8B-262K* | 33.5 | 28.3 | 29.4 | 66.9 | 43.0 | 42.4 | 40.6 |
|StreamingLLM | 29.7 | 23.7 | 28.8 | 65.8 | 19.3 | **42.8** | 35.0 |
|InfLLM |  31.6 | 25.8 | 29.1 | 66.3 | 36.3 | 41.8 | 38.5 |
|**MInference** | **34.0** | **28.0** | **29.7** | **66.8** | **42.8** | 42.2 | **40.6**|

Table 2(b). Performance of different methods on LongBench.

Meta-review: This paper discussed an important problem in improving the efficiency of long context inference. The solution is solid with good novelty, and the experimental section is well designed to verify the design of the proposed solution.
Decision: Accept (spotlight)