FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference
Authors: Xunhao Lai
Jianqiao Lu
Yao Luo
Yiyuan Ma
Xun Zhou
Venue: ICLR 2025 Oral (2025)
Keywords: Large Language Models (LLMs)
LLM inference
Long-context LLMs
Sparse Attention Mechanism
Primary area: generative models
TL;DR: FlexPrefill is a novel sparse attention mechanism for large language models that dynamically adapts attention patterns and computational budgets in real-time to optimize performance for each input and attention head.

Abstract: Large language models (LLMs) encounter computational challenges during long-sequence inference, especially in the attention pre-filling phase, where the complexity grows quadratically with the prompt length. Previous efforts to mitigate these challenges have relied on fixed sparse attention patterns or identifying sparse attention patterns based on limited cases. However, these methods lacked the flexibility to efficiently adapt to varying input demands. In this paper, we introduce FlexPrefill, a Flexible sparse Pre-filling mechanism that dynamically adjusts sparse attention patterns and computational budget in real-time to meet the specific requirements of each input and attention head. The flexibility of our method is demonstrated through two key innovations: 1) Query-Aware Sparse Pattern Determination: By measuring Jensen-Shannon divergence, this component adaptively switches between query-specific diverse attention patterns and predefined attention patterns. 2) Cumulative-Attention Based Index Selection: This component dynamically selects query-key indexes to be computed based on different attention patterns, ensuring the sum of attention scores meets a predefined threshold.
FlexPrefill adaptively optimizes the sparse pattern and sparse ratio of each attention head based on the prompt, enhancing efficiency in long-sequence inference tasks. Experimental results show significant improvements in both speed and accuracy over prior methods, providing a more flexible and efficient solution for LLM inference.

Published as a conference paper at ICLR 2025
FlexPrefill: A CONTEXT-AWARE SPARSE ATTENTION MECHANISM FOR EFFICIENT LONG-SEQUENCE IN- FERENCE
Xunhao Lai'", Lai"",Jianqiie Jianqiao Lu Yao Luo", Yiyuan Ma², Xun Zhou', !ekknng University The University of Hong Kong BByteDance Inc l xunhao@pku. edw.opkkuuuu.... en, jqlu@cs. hku. [luoyao. 0, mayiyuan. unicorn, /houxun//bytedanncs
ABSTRACT
Large language models (LLMs)) encounter computational challenges during long- sequence inference, especially in the attention pre-filling phase, where the com- plexity grows quadratically with the prompt length. Previous efforts to mitigate these challenges have relied on fixed sparse attention patterns or identifying sparse attention patterns based limited cases. However, these methods lacked the flex- ibility to efficiently adapt to varying imput demands. In this paper, we introduce FlexPrefill, Flexible sparse Pre-filling mechanism that dynamically adjusts sparse attention patterns and computational budget in real-time to meet the specific requirements of each input and attention head. The flexibility of our method is demonstrated through two key innovations: 1) Query-Aware Sparse Pattern Deter- mination: By measuring censen-Shannon divergence, this component adaptively switches between query-specific diverse attention patterns and predefined atten- tion patterns. 2) Cumulative-Attention Based Index Selection: This component dynamically selects query-key indexes to be computed based on different atten- tion patterns, ensuring the sum of attention scores meets predefined threshold. FlexPrefill adaptively optimizes the sparse pattern and sparse ratio of each atten- tion head based on the prompt, enhancing efficiency in long-sequence inference tasks. Experimental results show significant improvements in both speed and ac- curacy over prior methods, providing more flexible and efficient solution for LLM inference.
Perfonancc (7)
Peffomanc
A
K-
Raiie
6.2x
5.2x
Threnghput
Thronghpet(()
4.4x
3.4x
3.3x
Perhemuace
Prefoassee
2.2x
1.1x
LLM
64L
1200
38-4k
Threnghput (1)
Thronghpptt (?)
Contect Lengrh
(a) Performance comparison of different inference ap- proaches across LLaMA, GLM, Yi, and Qwen models.
(b) Attention acceleration ratio comparison of differ- approaches various context lengths.
Figure 1: Performance and speed analysis of different inference approaches.
"I Leading co-authors with equal contribution.

---

Published as a conference paper at ICLR 2025
INTRODUCTION
Large language models (LL.Ms) are transforming various domains by enabling sophisticated natural language understanding and generation (Zhao et al. (023). However, as the length of the context supported by these models increases (Zeng et al. 2024; Reid et al. 2024; Dubey et al. 2024 Ab- din et al. 2024 Young et al.] 2024). and as these models are applied to longer and more complex tasks (Bai et al. 2024] Shaham et al., 2023: An et al., 2024 Zhang et al., (024c), the cost of in- ference, particularly during the prefill phase, becomes bottleneck Fu (2024). The computational complexity of full attention mechanisms, which grows quadratically with sequence length, leads to inefficiencies during long-sequence inference tasks.
Considering the sparse nature of attention in LLMs (Deng et al. 2024). sparse attention mechanisms
have been proposed to limit attention calculations selected subset of query-key pairs. Many ex- isting sparse attention methods concentrate on predefined sparse patterns, such as BigBird (Zaheer al. 2020), LongLora (Chen et al.) 2024). and
Window (Jiang et al. 2023). How-
ever, these approaches restrict the flexibility of sparse attention and often necessitate further training or fine-tuning of the model. There also research proposing training -free sparse attention methods for large language models. For instance, Streaming L.M LLM (Xiao et al., 2024b) identifies the attention sink phenomenon and introduces a sparse attention mechanism that combines initial tokens with slid- ing windows. HHan et al. (2024) and Xiao et al. (2024a) further extend the sink attention approach. Additionally, other studies, such as those by Likhosherstov et al. (2021), 1), Liu et al. (202I), and Jiang et al. (2024a), highlight the presence of dynamic sparse patterns in models. Minference (Jiang al., 2024a) analyzes the intrinsic sparse patterns in LLMs and utilizes offline search- based sparse pat- terns along with online search-based sparse indexes. However, this approach relies on predefined sparse patterns and ratios with a limited search space, failing to account for the diverse and dynamic nature of real-world input sequences. Consequently, these attention patterns are often suboptimal, particularly when the complexity of input data varies, as they lack the adaptability required to meet the specific needs of varying input sequences.
To address these limitations, we propose FlexPrefill, a novel flexible sparse prefilling attention mechanism designed to adapt in real-time to the specific needs of each input and attention head. Our approach comprises two key components: (1) Query-Aware Sparse Pattern Determination: We categorize attention heads into Diverse pattern, which requires query-specific estimation, and Struc- tured pattern, which is consistent across queries. Each attention head adaptively switches its sparse pattern based on the input by measuring the Jensen-Shannon divergence between the estimated and true attention score distribution. (2) Cumulative-Attention Based Index Selection: This compo- nent dynamically selects query-key indexes to be computed based on attention patterns, ensuring the sum of attention scores meets predefined threshold. This approach allows for the adaptive allo- cation of computational budgets for different attention heads while maintaining the effectiveness of the model. FlexPrefill's adaptive optimization of sparse patterns and ratios for each attention head based on inputs effectively balances computational efficiency with model performance.
We conduct extensive experiments using state-of-the-art L.MMs, including Meta-Ilama-3.. 1.ma--31--8-- 1-BB- Instruct (Xiong et al., 2024), GLM-4-9B- Bhhh (Zeng et al.] 2024). Yi-9B-200K (Young et al. 2024). and Qwen2-7B- Instruct (Yang al. 2024), on challenging long-context benchmarks such as RULER (Hsich et al. (024) and InfinieBennch (Zhang et al.) (024c). The results demonstrate sig- nificant improvements in both speed and accuracy over prior methods, with FlexPrefill consistently
-B-ntrrutt
preserving or even enhancing model performance across various context lengths and tasks.
2 SPARSE ATTENTION
Problem Setting In the sparse attention mechanism, Q and K represent the query and key ma- trices given sequence length L, respectively. The attention score for each position is computed as the scaled dot product between the query matrix Q and the key matrix K, normalized by the square root of the head dimension d. To improve computational efficiency, sparse attention lim- its computation to subset of query-key pairs whose indexes form the set S U S,, where S, S [(t.)) I < i, j E Z). The sparse attention mechanism is formalized as:
A(Q, K, V, Softmax
(Q.K) K + +ss)
2

---

Published as a conference paper at ICLR 2025
Here, Mg is sparse attention mask based on S defined as:
0, Msli.j] -00
if (1, S,
otherwise.
Goal The goal of the dynamic sparse attention system is to achieve balance between compu- tational efficiency and attention effectiveness. This can be formally expressed as a multi-objective optimization problem:
min
K, -(Q,
(I)
min
||]
where AQQ, K, V) represents full attention result computed and K, represents the sparse attention result. The first objective in Equation Taims minimize the difference between the sparse and full attention results. The second objective seeks to minimize the size of the selected subset S, which directly relates to the computational efficiency.
Attention Sparse Patterns Exhibit Variability In the attention mechanism of LLMs, we observe that attention patterns can vary significantly between attention heads. Figure 2a reveals Diverse pattern for different query positions. This variability suggests that a one-size-fits-ale approach to sparse attention may be suboptimal. Conversely, (2jehibb more consistent and structured pattern across queries, showing a clear Structured pattern similar to that reported by Jiang et al.) (2024a). which also known as Vertical-Slat Slash pattern. attention heads with such a pattern, subset of the attention map may be sufficient to estimate the entire sparse index set.
(a) Attention maps exhibiting Diverse patterns, where attended key tokens scattered across different query positions with independent blocks.
(b) Attention maps exhibiting Structured patterns, e.g. attention is concentrated along vertical and slash- like structures, which is consistent across queries.
Figure 2: Comparison of attention patterns in different attention heads. The Diverse patterms (a) show scattered attention with independent blocks across query positions, while the Structured patterns (b) exhibit attention focused along certain structures.
Varying Samples Require Adaptive Sparsity Ratio Previous work (Jiang et al., 2024a) utilizes an offline search to determine fixed sparsity ratio, applied uniformly across all input cases. How- ever, as shown in Figure 3] the input prompts exhibit varying levels of complexity, necessitating different sparsity ratios for optimal performance. Samples with significant long-range dependencies may benefit from lower sparsity ratio, whereas those with predominantly local dependencies can be efficiently processed with higher ratio.
To address the above problems, we propose method to dynamically adjust each sample's sparse pattern and sparsity ratio, introduced in Section 3
METHOD
In this section, we introduce our method for optimizing the sparse attention mechanism. Our ap-
proach involves key components: (1) Query-Aware Sparse Pattern Determination and (2) Cumulative-Attention Based Index Selection. We then present the overall sparse attention algo- rithm of FlexPrefill in (3) Algorithm.

---

Published as a conference paper at ICLR 2025
Altertion
Altention Ooverage
Attention Goverage 0.
0.8
0.8
lyprr
16
0.6
0.6
0.6
20
32 28
72
16 head
20
24
28
32
0.0
8
12
16 head
20
24
28
0.0
12
16 20
24
28
0.0
head
(a) longer prompt, task A
(b) longer prompl, task B
(c) shorter prompl, task B
Figure 3: Adaptive sparsity ratios across different attention heads layers for varying sample complexities. Each heatmap shows the sparsity rate different context lengths and task types given a fixed attention score coverage, where darker colors indicate more attention calculations. The sparsity distribution of different atten- tion heads varies with different sample types (a, h) and different context lengths (b,
Query-Aware Sparse Pattern Determination Our observations of Diverse (Figure 2a and Struc- tured patterns (Figure 2b) motivate an adaptive approach sparse attention estimation. We classify attention heads into two types: the Query-Aware head, which is designed to capture variable sparse patterns contingent upon query positions, and the Vertical-Slast head, which intended handle commonly occurring structured sparse attention patterns in L.Mss. We propose a method that dy- namically determines whether to use Query-Aware attention estimates or a more consistent pattern, balancing computational efficiency with accuracy.
Given the computational challenges calculating the full global attention map, we adopt a prag- matic approach:
Representative Query Selection: We select the last block_size query vectors, denoted as Q, representative subset for subsequent estimation.
2. Block- Wise Attention Estimation: By applying a pooling operation on the sequence, we compute two block-wiss attention distributions:
a softmax aavgpool avggool /V////)
à sumpool softmax (QKT/Va)
where the kernel size of the pooling operation is equal to block_size, à represents the estimated distribution and à represents the true distribution.
3. Distribution Comparison: We quantify the discrepancy between these distributions using the square root of the Jensen-Shannon divergence:
DJs(ā, a) VJSD(al()) V/2 (Dxi(allm) + Dx (ā
where = [a (ā) is the mean distribution, and Dx (-ll-) denotes the Kullback- divergence. We adopt the censen-Shanoon divergence for its symmetry and boundedness properties, which ensure more stable and reliable pattern determination.
4. Adaptive Decision:
We compare DJs predefined threshold If DJS <T, the block- wise attention estimation adequately approximates the true distribution. Consequently, we calculate for all queries, enabling query-aware sparse index selection. Otherwise, the approximation is inadequate, we fall back to a more conservative approach. Specifically, we utilize only subset of queries for subsequent sparse index searching and extend the global index set using vertical-slast pattern.
This adaptive method allows to leverage Query-Aware attention patterns when appropriate, po- tentially capturing more diverse patterns. Meanwhile, provides a fallback mechanism to a more consistent Vertical-Slash pattern when the attention estimation is unreliable.
4

---

Published as a conference paper at ICLR 2025
Cumulative-Attention Based Index Selection After deciding whether an attention head belongs to Query-Aware or Vertical-Slash, we attempt to minimize the computation while ensuring effec- tiveness. Our sparse attention mechanism aims select the smallest possible subsets S; for each query position while ensuring that the sum of normalized attention scores within each subset meets predefined cumulative attention threshold y. This objective can be formally expressed as:
exp KJ /Vd) (L,))S Ey exp(Q: K ///Va)
min S:
S
subject to
Σ
2 Vi [n]]
where S U,S, and |Sa| denotes the size of chosen index subset, Q; and K, are the query and key vectors at position i, respectively. The detailed rationale behind optimizing Equation (3 is provided in Appendix
To achieve this objective, we employ different strategies depending on the attention head type:
1. Query-Aware head:
Estimate the block-wise attention scores by first wveage-ppooinng the queries and keys, then computing the attention.
Sort the blocks in descending order and select blocks in order until the sum of their normalized attention scores exceeds the threshold y, forming the sparse index set S.
2. Vertical-Slash head:
Select a representative query subset Q and compute attention scores, then calculate average scores for the vertical and slash lines.
Sort these lines and select vertical and slash lines in descending order until their cu- mulative normalized attention scores exceed ?
Extend selected lines to the entire attention matrix, forming the sparse index set S.
This adaptive approach tailors the selection process to the attention patterns each attention head, allowing us to efficiently determine the sparse index set S for sparse attention computation. By doing so, we balance computational efficiency with attention accuracy, ensuring that we focus com- putation the most relevant token interactions while maintaining predefined level of cumulative attention scores.
Algorithm Algorithm 1 presents the overall procedure of our proposed method for efficient sparse attention computation. The algorithm takes the query matrix Q, key matrix K, value matrix V, sparse pattern threshold 7 and cumulative attention threshold y as input. It is divided into the fol- lowing three parts:
(i) Sparse Pattern Determination: Algorithm determines whether to use Query-Aware pattern or fall back to the Vertical-Slac pattern for each attention head.
Based on the attention patterns obtained in (i) and the given cumulative attention threshold 7, the sparse index set S that needs to be computed for each attention head is obtained by Algorithm (Query-Aware) Algorithm (Vertical-Slash)
(ii) Sparse Index Selection:
(iii) Sparse Attention Computation: The algorithm performs sparse attention computation for each attention head using the obtained sparse index and returns the final attention result.
EXPERIMENT
4.1 SETTINGS
In this section, we provide a comprehensive overview of the models and datasets employed in our experiments to assess the performance and capabilities of various long-context language models.
5

---

Published as a conference paper at ICLR 2025
FlexPrefill
Algorithm 1 Sparse Attention
Algorithm Sparse Pattern Search
Input: Q, K, V E RSxd E [0, 1}. Y E (0,1)
Input: Q.K.
Take a representative query subset
select Q Q ([[bbockksize))
Determine the sparse pattern based on the threshold T
Compute estimated plock-poola attention a and true block block-pooled attention à
pattern + Sparse Pattern Scarch (Q, K,T) r)
Decide the sparse index set S based on pattern an if pattern query then
ā + softmax (pool (pool KK] lvd/
à + pool (softmax (QKT// \vaa
S
--Quuyy Aware Index (Q, K,
Compute -hansenn divergence d + [JSD(al(a)
else pattern vertical_ slash then
S +
Vertical Slash Index (Q, (,
Determine whether to use Query-Speciti tention paltern if ds < then pattern queer_useeciiie
end if
Compute the final sparse attention output t A(Q, K, K.V. return
else
pattern + vertical slash
end if return pattern
Algorithm 3 Vertical Slash Index Search
Algorithm Query Aware Index Search
Input: Q, K,y
Input: K.
# Compute subset of the full attention map
Compute estimated attention scores using
Å softmax (QK*) /vV) where Q C Q
pooled Q and
Q + pool(Q), + pool(K)
# Sum and normalize attention scores along the
Ā <- softmax
(QR /va
vertical and slash directions
ag <- sum (ereieerrtaaa)tt (A)/ Ex, A[i,J]
Flatten and normalize attention map
a₂ + sum_slash(A) Ei,, A
A + flatten( A/ //,,, Ali
# Sort vertical and slash attention scores
Sort attention scores
I. + argsort(an)
Ia argsort(A)
I. argsorl(az)
# Obtain the minimum computational budget
# Obtain the minimum computational budget making the sum of the scores exceeds 7
making the sum of the scores exceeds
K +- min(k: EiETa(LA) A 7}
K. min(k:
EiEL.[IN [ 2 7
Get final index
K. min(k: EiEr, a. 7
S Ta : K] return
# Select vertical and slash index
S. I.[1: K.],, S. [,[1 K.]
< S. U. return S
Models We utilize four state-of-the-arr LLMs known for their proficiency in handling long-
tasks: (i) Meta-Llama-3, 1-8B-IIsstrrut-1225 (Dubey et al. (024) (LLaMA) (ii) GLM-4-9B-Chat
1024k ZZng et al., 2024) (GLM) (iii) Yi-9B-200K (Young et al.] 2024 (Yi) (iv) Qwen2-7B-
instruct-128k (Yang el al. 2024) (Qwen), where Yi is a pre-trained model and all others are instruct
models. We use the default chat template for all instruct models the subsequent experiments.
Datasets
We evaluate the models on two datasets, each offering unique challenges in long-contex. understanding: (i) RULER (Hsieh al. 2024): synthetic benchmark dataset created to evaluate

---

Published as a conference paper at ICLR 2025
long-context LLMs with customizable sequence lengths and task complexities. It extends the basic seedle-in-a-haysee test as well as introduces new task categories such as multi-hop tracing and aggregation. (ii) Infinite Bench (Zhang et al. 2024c): a benchmark dataset designed to test LLMs' understanding of long dependencies within extensive contexts, with an average token count of 214k. It includes 10 synthetic and real-world tasks across various domains.
Implementation Details Our experiments are conducted in computing environment equipped with a single NVIDIA A100 GPU with 80GB memory. For our experiments, we implement a custom pipeline using PyTorch, building FlashAttention (Dao 2024 to ensure efficient atten- tion mechanisms over long-contex imputs. Our implementation leverages Triton (Tillet et al. 2019) for optimizing the performance of GPU-accelerated computations and uses block_size 128 in all experiments. We choose the last block_siee _izeeseeeeeeeeeeeeeeeee query vector the representative query set Q for both sparse pattern determination and Vertical-Slash sparse index selection, which allows the represen- tative attention score be computed only once and reused in both attention pattern determination and sparse index selection. The sparse pattern threshold T is set to 0.1 for all models. Additionally, by adjusting the parameter y we assign tailored computational budget each attention head in model,
real-time for each input scenario. We set 0.9 for Yi-9B-2000 and Qwen2-7B-Insiiee and 0.95 for the other models. To ensure that all attention heads work normally. we retain the first and last key blocks of each query block, while also requiring each attention head to compute at least 1024 tokens. All experiments were conducted using greedy decoding to maintain consistency across results.
Baselines We evaluate our method against three strong baseline approaches to highlight its ef- ficiency and performance in long-context tasks: 1) FlashAttention (Dao 2024): highly opti- mized attention mechanism that exploits underlying hardware optimizations and efficient memory access patterns, which significantly reduces the computational and memory overhead of attention operations. 2) Streaming (Xiao et al.) 2024b): A sparse attention mechanism that combines global attention sink tokens with dilated local attention windows. figuration with interval of 1, encompassing 1,000 global tokens and 8,000 local windows. 3) MInference (Jiang et al., 2024a): An efficient sparse attention method for handling long sequences, employing three distinct sparse attention patterns. According to the recommended settings from the original paper, we utilize offline search attention patterns and dynamically generated indexes to compute each attention head separately. Additionally, we include more baseline comparisons for comprehensive evaluation in Appendix D
use a dilated window con-
All baselines use sparse computation during pre-filling, switching to dense computation during the decoding stage. This ensures fair comparison across different methods and highlights the balance between latency and performance.
4.2
MAIN RESULT
This section presents a comparative analysis of various attention mechanisms. The evaluation in- cludes Full Attention, Miniference, Streaming LLM, and our proposed method. The comparison aims to demonstrate our approach's efficiency and effectiveness against existing methods.
RULER demonstrate the potential our approach for long-context large language models, we conduct evaluations using the RULER benchmark. Table illustrates that Streaming LLM's per- formance significantly deteriorates as context length increases, while MInference shows suboptimal performance on certain models. In contrast, method consistently preserves various models' per- formance across multiple context lengths. Notably, our approach can enhance models' capabilities in some scenarios while accelerating the computational process. Detailed latency comparisons are provided in Appendix(C
Infinite Bench
To further demonstrate our approach's potential for long-contex: LLMs, we con- Bench challenge. Table 2 shows that our method
duct evaluations using the state-of saate-of-the-aee Infinite preserves most of the model's performance in retrieval and QA tasks, while maintaining efficacy in complex mathematical and coding tasks.
7

---

Published as a conference paper at ICLR 2025
Table I: Performance comparison of different methods various models and sequence lengths RULER. The best results are highlighted bold.
Models
Methods
4k
8k
16k
32k
64k
128k
Avg
Full-attn
95.67
93.75
93.03
87.26
84.37
78.13
88.70
Streaming LLM
95.43
93,99
74.76
48.56
26,20
30.77
61.62
LLaMA
Minference
95.67
93.99
93.27
86.54
84.86
58.17
85.42
Ours
95.43
93.51
94.71
89.42
82.93
79.09
89.18
Full-attn
93.75
93.03
89.66
90.63
85.34
81.97
89.06
Streaming LLM
93.75
92.79
76.92
56.97
40.14
34.86
65.91
GLM
MInference
93.75
93.03
90.38
89.90
85,10
82.93
89.18
Ours
93.51
91.83
89.90
91.35
86.06
83.41
89.34
Full-attn
92.79
86.06
85.3
76.93
69.47
66.35
79.49
Streaming LLM
93.03
83.41
65.38
51.68
39.18
34.61
61.22
Yi
Minference
92.79
85.58
82.69
73.32
63.94
57.69
76.00
Ours
93.27
86.78
83.89
72.36
64.66
56.97
76.32
Full-attn
89.90
88.70
80.77
79.33
56.49
17.79
68.83
Streaming LLM
90.14
88.94
57.93
43.03
25.48
12.26
52.96
Qwen
MInference
89.90
88.70
79.33
78.61
49.04
10.82
66.07
Ours
90.39
89.91
83.17
81.25
59.14
20.67
70.75
Table 2: Performance comparison of different methods on various models and tasks on Infinite Bench. The bold/underlined numbers indicate the first/second highest value in each column.
Methods
En.Sum
Ra.QA
Zl.QA
Code.Dehog
Math:Find
Retr.PassKeey
Rete:Numberr
RetrKV
Axg
Full-atte
1LaMA
Minference
Streaming
31.04
22.00
63.76
14.50
28.17 20
5.33
27.43
77.12
14.00
34.05
Ouis
31.82
24.82
69.43
19.50
35.46
16.75
31.14
98.64
99.83
44.00
47.14
Pull-atta
28.00
20.70
42.36
35.00
15.69
32.99
26.29
99.15
100.00
19.00
41.95
GLM
Seroming LLM
Minferstee
28.1 28.15
27.37
19.42
11.58
44.98
31.00
29.50
18.50
15.83
13:77
36.29
26.65
20.20
23.71
100.00
15.08
100.00
100.00
48,60
6.50
44.55 56
27. 16
Oues
28,90
20.27
12.36
29.50
13.55
33.50
23.71
99.15
100.00
$5.40
44,83
Pull aum
8.32
11.47
65.50
1.50
17.49
20.51
23.14
100000
99,66
29.00
37.69
Yi
Streaming LLM
Minference
6.01
6.06
10:27
11.31
62.45
42,70
3,50
2.00
17.15
17.74
24.37
10.30 54
24.86
22.86
108.00
10.17
97.63
36.
29,00
3.50
37.44
17.29
Ours
6,62
11.89
65.07
2.50
19.54
24,00
99.83
91.53
24.00
36.
Full
4.65
5.57
34.50
9.00
11.27
24.87
24.57
95.42
75.42
0.00
29.53
18.54
10.94
23.24
Minferatee
7.45
58.47
23.65
13.27
20.81
30.07
Performance vs. Latency Our method leverages online search characteristics, enabling seamless application to various models and flexible adjustments between computational speed and perfor- mance. Tuning the > parameter balances model efficacy and output latency: decreasing acceler- ates processing while increasing 7 preserves model quality. Figure compares approach with MInference, demonstrating that our method achieves better performance with lower latency. A more detailed latency analysis can be found in Appendix G
(a) Average
(b) 128K context
(c) 32K context
90
90
RULER Score
RULER Score
RULER Score
86
20
86
84
84
82
80
50
15
AVG Attention Latoncy (ms)
128K Attention Latency (ms)
Attention Latency (ms)
Full Attention
Minference
Fixed Budge Vertical-Slash
FlexPrefill
Figure 4: Comparison of method with MInference and Fixed budget Vertical-Slash attention, showing the trade-off between model performance and attention latency. Our method consistently outperforms MInferenee and Fixed Budget Vertical- Slash across different attention latencies. More details are provided Appendix
8

---

Published as a conference paper at ICLR 2025
4.3 ABLATIONS
We conduct various ablation studies to analyze the components of our proposed method and under- stand the impact of each configuration, with some details provided in the Appendix due to space constraints. We begin by validating the necessity of dynamically computing budget allocations. We explore the significance of the Query-Aware attention head and the selection of threshold т. We also investigate the impact of various minimum computational budget allocations the robustness of the sparse attention mechanism.
Furthermore, we find that the block size of sparse attention does not significantly impact model effectiveness (Appendix F22 We observe that different representative query set Q minimally affect the Query-Aware pattern determination but substantially influence the index selection of the Vertical- Slash head (Appendix This observation justifies the use of the last block size query vectors as
Q. We evaluate alternative index construction implementation for Algorithm| observe that it has negligible impact on model performance (Appendix[F4) (F4). Finally, we study the effect of limiting the maximum computational budget of the attention heads and discover that the performance of some models saturates with increased computation, indicating the potential for further acceleration (Appendix [.6).
Fixed Budget vs. Dynamic Budget We experiment with fixed and dynamic budget allocations. Fixed budget uses the same computational budget for all attention heads across all layers. Results in Figure 4 show that dynamic settings lead to improved performance and a better balance between inference speed and model effectiveness compared to static budget allocation.
Query-Aware head threshold We evaluate model performance with and without Query-Aware attention pattern and examine different 7 impacts on performance. Figure 5 demonstraaee that en- abling Query-Aware attention head at an appropriate т strengthens the model without increasing computation. However, if is set too large, some attention heads with inaccurate attention estima- tion are recognized as Query-Aware heads, potentially degrading model performance. More detailed results can be found in Appendix F.
LLaMA
GLM
Yi
Qwen
70.5
$5.0
18.5
70.2
68.0
68.1
56-2
(tau)
(tau)
((auu)
Figure 5: Comparison of our method under different thresholds 7 indicates that an appropriate threshold improves model performance.
Minimum Budget Limit We perform ablation experiments to analyze the minimum computa- tional budgets required for each attention head. Our findings reveal that setting minimum budget threshold enhances performance and prevents the collapse of attention heads with extremely high sparsity ratios. More detailed results can be found in Appendix [.5 and Table 10
To further analyze the components of our approach, we visualize the distribution of sparse patterns between different layers (Appendix]B H), the sparse masks of Query-Aware and Vertical-Slast heads in approach (Appendix [1). and the distribution of sparsity ratios for different layers and heads (Appendix D). respectively. This visualizes the flexibility of our approach and the significance of each component.

---

Published as a conference paper at ICLR 2025
5 RELATED WORKS
Ongoing advancements in engineering and algorithmic design continue to push the boundaries of long contexts large language models (LLMs), enabling them to tackle in- creasingly complex tasks. Some methods extend the context length of the mode by gathering exten- sive long-text datasets and persistently pre-training or fine-tuning the models (Fu et al. 2024b Chen al. 2024; Xiong al. 2024). Given that many modern L.LMs employ RoPE-based positional embeddings (Su et al. 2021), various innovative positional embedding techniques have been intro- duced to extend model lengths, exemplified by Peng al. (2024), Ding al. (2024), and al. (022d).. Moreover, approaches leveraging external memory or retrivval augmentation have been proposed to enhance long context processing capabilities (Mohtashami & Jaggi 2023 Twokowski
Long Context LLMs
et al. 2023
Xu et al. 2024).
LLM Inference Acceleration Given the quadratic time complexity of self-attention relative to sequence length, accelerating inference for long contexts is crucial for LLMs. Some strategies optimize the original attention computation through algorithms synergized with hardware features.
including FlashAttention (Dao et al. 2022 Dao 2024, Shah et al. 2024 and RingAttentionn (Liu
al., 2023; [Bandon et al. 2023). Additional methods facilitate attention acceleration by reducing
context length, such as retrieval sariivaal-aaeee based techniques (Mohtashami & Jaggi 2023 Wang et al. [203 Xu
al. |020|22022aaa Munkhdalai al., 2024) and compression- based approaches (Li et al. al.3]233|Jianng
Jiang et a
2024b Pan et al. (2024). Strategies addressing the quadratic complexity of attention computations
include employing recurrent mechanisms to aggregate information (Zhang et al. 2024b Bulatov
et al.
2023 Martins et al. 2022), using State Space Models (Gu et al.] [022
Gu & Dao 2023
Lieber et al. 2024). and exploring innovative model architectures (Sun et al. 2023 Peng et al,
2023 Beck et al.] 2024). Given the inherent sparsity in attention mechanisms, numerous sparse
attention methods have been proposed. Many focus on fixed attention patterns, such as shifted sparse aftention (Chen al. (2024), sink attention (Xiao et al.)
(024b), and other approaches (Child al.) 2021: Ding et al.33 Modern LLMs like 2023) and phi3 (Abdin et al.) (2024) also employ fixed sparse attention patterns, including sliding window attention. Han et al. (2024) and Xiao et al. (2024a) further extend the sink attention approach. Additionally, other studies, such as those by [ikhosherstov et al. (202T), Liu
| Bellagy et al. 20200 Zaheer et 22220 Shi et al.
mistral (Jiang et al.
et al. (0 and Jiang et al. (2024a), highlight the presence of dynamic sparse patterns in models.
For the decoding phase, KV caching techniques that store past keys and values are prevalent, intro- ducing unique challenges for I.M inference acceleration. Some methods optimize computational processes or KV caches management more efficiently, such as FlashDecoding (Dao et al., 2023] Hong et al. (023) and PatedAttention (Kwon et al. (023). Other approaches aim to reduce KV cache sizes through quantization (Liu et al. [0024a] Kang et al. 2024, Liu et al. 2024b), token merging (Nawrot et al., 2024 Li al. (022).
KV discarding (Zhang al. 2023 Dong al.)
2024]. Furthermore, sparse attention methods are widely applied in the decoding stage. Some tech- niques retain the entire KV cache but utilize only subset for sparse attention computations Ribar eet al.] (2024): [Tang et al.] (2024); Zhhng et al. (2024a), while others diminish the KV cache via specific sparse attention patterns Ge et al. (004 Xiao et al. (2024b); Zhang et al. 2023)
6 CONCLUSION
In this paper, we introduced FlexPrefill, a novel flexible sparse attention mechanism designed to adapt in real-time for efficient long-sequence pre-filling in LLMs. Our approach, comprising Query- Aware Sparse Pattern Determination and Aamulative-Atteniion Based Index Selection, addresses the limitations of previous methods by dynamically optimizing sparse patterns and ratios for each at- tention head based on inputs. Extensive experiments across sfate-of-the-art LLMS and challenging long-context benchmarks demonstrate that FlexPrefil consistently preserves or enhances model per- formance while significantly improving computational efficiency. The adaptive nature of our method allows for better balance between speed and accuracy, outperforming prior methods across vari- ous scenarios. As LLMs continue to evolve and tackle increasingly complex long-contes tasks, FlexPrefill offers promising solution for maintaining high performance while managing compu- tational resources effectively. Future work could explore further optimizations and applications of this adaptive approach in diverse model architectures.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Marah I Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat S. Behl, Alon Benhaim, 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 Mi- tra, 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, Yuman Zhang, and Xiren Zhou. Phi-3 technical report: A highly capable language model lo- cally your phone. CoRR, abs/2404. 14219, 2024. doi: 10.48550/RXXV.2444 14219, URL ht.t https:///oot wor.ioorg/.... 48550 /0.5550/axxiv.. 2wtp:/2ww.w............ 14219
Chenxin An, Shansan Gong, Ming Zhong, Xingjian Zhao, Mukai Li, Jun Zhang, Lingpeng Kong, and Xipeng Qiu. L-eval: Instituting standardized evaluation for long context language mod- els. In Lun-We Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd An- nual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 1-16, 2024, pp. 4338-14411. Association for Computational Linguisties, 2024. doi: 10. 18653/v 13653/v1/2024 acl 1101--1000 long. 776
URL https: orr/10.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilin- gual, multitask benchmark for long context understanding. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Com- putational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11 11-16, Association for Computational Linguistics, 2024. doi: 10.18653/V1/2024 www.:oooorrg/00 wttp://....165////1 aoc.o//ooooo.....ooo.........oo 172
2024, pp. ACL-LONG. 172. URL
Maximilian Beck, Korbinian Pöppel, Markus Spanring. Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. CoRR, http:/22550045111 2024. doi: 10.4550/ARXXV.2445.0555 144544 URL
https: /t:p:://dd0 . /0. 10. 48550 a arxiv 2405. 04517
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020. httth:tt///aatt
org/ 2004. 05150
William Brandon, Aniruddha Nrusimha, Kevin Qian, Zachary Ankner, Tian Jin, Zhiye Song, and Jonathan Regan-Kelley. Striped attention: Faster ring attention for causal transformers. CoRR, abs/231 1. 004311 2023. doi: 10.48550//RXIV.231..000 48550/arXiv. 09431
URL https::/p//....
Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Scaling transformer 1m tokens and beyond with RMT. CORR, abs/2304. 1062, 2023. doi: 10.48550/ARXIV2234..100 URL https:
//doi.org/// 48550/8551i 2304. 11062]
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Lon- glora: Efficient fine-tuning of long-context large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenRe- view.net, 2024. URL URL
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse 904. 10509
transformers. CoRR, abs/1904. 10509, 2019. URL

---

Published as a conference paper at ICLR 2025
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL htttp:://pperre//// .et forum?id- mZr 22yh
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Sanmi Koyejo, S. Mo- hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 9, htt.aappee_fllleeeerrrr222222220 Abstrratt-sstttttt-ta Contrracct--ooffeeeneehhttt
2022, 2022. URL http::://appes...pattee 312677222220000a77333200.......
Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flash-decoding for long- context inference, October 2023. URL wttps///ern..sstfod..ee...
baashdecoding.htne html Accessed: 2024-9-29.
Yichuan Deng, Zhao Song, and Chiwun Yang. Attention is naturally sparse with gaussian distributed input. CoRR, abs/24044 2024. doi: A005555/ARRXV.24440000000 Utthtts::///od 02690
brg/10. 4ww54///05....222
Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers to I, 000, 000, 000 tokens. CoRR. abs/2307.02488 024466 2023. doi: 10.48550/ARXX1.2337.000 02486. URL fhttps://dol org/10.
48550/arXiv 2307 02486
Yiran Ding, Li 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, ICML 2024, Vienna, Austria, net/ fttt/formmmioo
July 21-27, 2024. OpenReview.net, 2024. URL
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 infer- ence. In Forty-firsi International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL httph::/ http:s////peeevviieeee ew. net/ forum? did=
uhHDDVKEW MW
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshom, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Rozière, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits. David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia- Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle I Lee, Georgia Lewis Ander- son, Graeme Nail, Grégoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Ko- revaar, Hu Xu. Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel M. Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Lec, Jan Geffert, Jana Vranes, Jason Park, Jay Ma- hadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong. Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Al- wala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, and et al. The llaa 3 herd of models. CoRR, hts//200.221783 2024. doi: 1044850/ARXXI.240..000 217.21837.......... URL //doi.
https: doi. /10 1111110.8855555 48540/5544000 www... 21783
Tianyu Fu, Haofeng Huang, Xuefei Ning, Genghan Zhang, Boju Chen, Tianqi Wu, Hongyi Wang, Zixiao Huang, Shiyao Li, Shengen Yan, Guohao Dai, Huazhong Yang, and Yu Wang. Moa: Mix- ture of sparse attention for automatic large language model compression. CoRR, abs/2406.1499
12

---

Published as a conference paper at ICLR 2025
2024a. doi: 10.48550/ARXIV2246611822 /ttp:/t//////...oo//000088855550
2406. 14909
Yao Fu. Challenges in deploying long-context transformers: A theoretical peak performance analysis. CoRR, abs/2405 089444, 2024. doi: 10.48550/ARXIV2400.0066 URL https:
10.
2405 08944
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models to 128k context. In Forty-first International Con- ference Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, net/forum? Fewfooommdd-Taageeeeeee Ta
2024b. URL /tttpp::///peneeviiie
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 Leurning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenRe- net wwwormmidduuurppDDD.
netw 2024. URL
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. CoRR, abs/2312.00752, 2023. doi: 1048550/ARXXV.2312.0077 URL htttp:://oi.crrg
00752
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The Tenth International Conference on Learming Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL httpt::t///..m / /pwpreeveeeeeeeeeeeemmomee net/
Eorum?id=uYLPozvv2
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 Gómez- Adomo, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Vol- ume 1: Long Papers), NAACI. 2024, Mexico City, Mexico, June 16-21, 2024, pp. 3991-4008. Association for Computational Linguisties, 2024. doi: 10.18653/112004.AAA...
URL https: / http:///oooo.ogg1//// 10. 1tteg//1..1165/1112222
naacl long. .22
Ke Hong, Guohao Dai, Jiaming Xu, Qiuli Mao, Xiuhong Li, Jun Liu, Kangdi Chen, Yuhan Dong, and Yu Wang. Flashdecoding++: Faster large language model inference on gpus. CoRR, abs/2311.01282, 2023. doi: 10.4550/ARXX.2331.0000 18550/arXis.
URL hattp:://do..rg//
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. RULER: what's the real context size of your long-context language doi: 10448550ARXX12244400666 //
models? CoRR, abs/2404.06664, 2024.
URL https:
48550//r21........// 2404. 00..444.6666444
Albert Q. Jiang. Alexandre Sablayrolles. Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b. CoRR, abs/2310.06825 2023. doi: 10.44550/ARXI.233000088 URL htttt::///o// 0rg/10.49550/rxxiv 2310.
06825
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhen- hua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Minfer- ence 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. CoRR, abs/2407.02499. 2024a. doi: 1044850/ARXX..240.00044 URL pettpp::/ddi.org./ 48550/arXiv. 2407 02490
Huiqiang Jiang, Qianhui Wu, Xufang Luo, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Longllmlingua: Accelerating and enhancing Ilms in long context scenarios via prompt compression. In Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Pa- pers), ACI. 2024, Bangkok, Thailand, August 11 11-16, 2024, pp. 1658-1677. Association for
13

---

Published as a conference paper at ICLR 2025
Computational Linguisties, 2024b. doi: 10.1663WV1200444C--LOGGGG URL https: /4bo...rr/10..1155/////22000 ott.::0066662/1222222....o-11110
Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong, Zaoxing Liu, Tushar Krishna, and Tuo Zhao. GEAR: efficient KV cache compression recipe for near-lossless generative inference of LLM. CoRR, abs/2403.05527 2024. doi: 1048550/ARXXY2203305557,,,
URL https:
7/ //doi. wrwpo..g/....... 48550/ 2403. 2403.05527
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Jason Flinn, Margo I. Seltzer, Peter Druschel, Antoine Kaufmann, and Jonathan Mace (eds.), Proceedings of the 29th Symposium on Operating Systems Principles, SOSP 2023, Koblenz. Germany, October 23-26, 26, 2023, pp. 61 626. ACM, 2023. doi: 10.1145/ 3600006.3613165 URLhtttss:// http::://doi.coooo orq/ 10 110111115/00000..1 3613165
Hcejun Lee, Geon Park, Youngwan Lee, Jina Kim, Wonyoung Jeong, Mycongjae Jeon, and Sung Ju Hwang. Hip attention: Sparse sub-quadratic attention with hierarchical attention pruning. CoRR, abs/2406.09827, 2024. doi: 10.48550/ARXIV.2406666 48550/arXiv 1.406.08822
URL httph::t///// / //doi. org/10.
Yucheng Li, Bo Dong. Frank Guerin, and Chenghua Lin. Compressing context to enhance inference efficiency of large language models. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023. pp. 6342-6333 Association for Computational Linguistics, 2023. doi: (0.18653/V1/2223.EMM 10. 10.18653/v1/2022 emwp://222.mmmmm..mm..mmmmmm main.391
URL https: //toi.:://doooro
Yuhong Li, Yingbing Huang. Bowen Yang. Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: LLM knows what you looking for before generation. CoRR, abs/2404. 14469, 2024. doi: 10.48550/ARXIV.2404.4 14469.. URL https:
//doi.org//l/ 485444555//530 2404 14469
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Mcirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, and Yoav Shoham. Jamba: A hybrid transformer-mamba language model. CoRR, abs/2403. 19887, 2024. doi: 1044550/ARXY.2400...855
URL
19887
Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self- attention matrices. CoRR, abs/2106.03764, 2021.
03764
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near- infinite context. CoRR, abs/2310.01889, 2023. doi: 1048550/ARXI.2231000082
BRLhttps:
//doi. o0l...o.g/1..44... 4ttg://0...50////..v..22. http://axxxx.220...0. 01889
Liu Liu, Zheng Qu, Zhaodong Chen, Yufei Ding, and Yuan Xie. Transformer acceleration with dynamic sparse attention. CoRR, abs/2110.11299 2021. URL httttts//axxii.... wwtp:/wwwwrrraaoooooooo abs/
1110.11291
Ruikang Liu, Haoli Bai, Haokun Lin, Yuening Li, Han Gao, Zhengzhuo Xu, Lu Hou, Jun Yao, and Chun Yuan. Intactkv: Improving large language model quantization by keeping pivot tokens intact. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics, ACL 2024, Bangkok, Thailand and virtual meet- ing. August 11-16, 2024, pp. 7716-7741, Association for Computational Linguistics, 2024a. doi: 1116////1/22222AIIIIGSS 2024. f..4ffindinge www.diinnnnnal....
URL https:// /ttp:://ddo...o wwwpiioorg/00.. 1tt.:/1....111/1111 18653/v1
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. KIVI: A tuning-free asymmetric 2bit quantization for KV cache. In Forty-first International Conference on Machine Learning. ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024b. hRLhttts:: htlhtttsss/oppereviieeee net wtttt//oormmmdddfioe ottp://foooummmmo15522222222
14

---

Published as a conference paper at ICLR 2025
Pedro Henrique Martins, Zita Marinho, and André F. T. Martins. co-former: Infinite mem- ory transformer. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Pro- ceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Vol- ume 1: Long Papers), ACI. 2022, Dublin, Ireland, May 22-27, 2022, pp. 5468-5485. Asso- ciation for Computational Linguistics, 2022. doi: 10.1665/112222AA--1OONG.. http:////oo..rrg11.1113533
URL
Long 375
Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. CoRR, abs/2305.166000 2023, doi: 1044850AARRXI.22355.66000 URL
https::/dbi.orgg// 48550/ http:aaxiiiv..
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini- CoRR, abs/2404.07143,, 2024. doi: 10.48550/
ARXIV.2404.0714 143. URL
2404. 07143
Piotr Nawrot, Adrian Lancucki, Marcin Chochowski, David Tarjan, and Edoardo M. Ponti. Dy- namic memory compression: Retrofitting Ilms for accelerated inference. In Forty-first Interna- fional Conference Machine Learning, ICML, 2024, Vienna, Austria, July 21-27, 2024. Open- Review.net, 2024. URL
net http://rrre...eeeoooommmooaaaaaa
Zhuoshi Pan, Qianhui Wu, Huiqiang Jiang, Menglin Xia, Xufang Luo, Jue Zhang, Qingwei Lin, Vic- for Rühle, Yuging Yang, Chin- Lin, H. Vicky Zhao, Lili Qiu, and Dongmei Zhang. LImlingua- 2: Data distillation for efficient and faithful task-agnostic prompt compression. In Lun-Wat Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Lin- guistics, ACL 2024, Bangkok, Thailand and virtual meeting, August 11-16, 2024, pp. 963-981. Association for Computational Linguisties, 2024. doi: f ndings acl.5. 157
URL httptttp://// /topp:://do...o org 100 1et.//.11111111// 18w53/16611e/
Leonid Pekelis, Michael Feil, Forrest Moret, Mark Huang, and Tiffany Peng. Llama 3 gra- dient: A series of long context models, 2024. URL https::/gradientt.... ai/glog//////ooeooooooooooom/pba for long- cont Languageggggeeee L models
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Leon Derczynski, Xingjian Du Matteo Grella, Kran- thi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong. Bart- lomiej Koptyra, Hayden Lau, Jiaju Lin, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Johan S. Wind, Stanislaw Wozniak, 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 (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, Singapore, December 6- 10, 2023, pp. 100048814077... Association for Computational Linguistics, 2023. doi: 10.18653/V/2002335NDDD
EMNLP.936. URL htttps::/dbi.oor//000 100.01..1666///0 2023 www.::..iiiinnnnnemnneee.e 936
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context win- dow extension of large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https:://ppeneviies
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy P. Lillicrap, Jean- Baptiste Alayrac Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, loannis Antonoglou, Rohan Anil, Sebastian Borgeaud, Andrew M. Dai, Katie Millican, Ethan Dyer, Mia Glaese, Thibault Sottiaux, Benjamin Lee, Fabio Viola, Malcolm Reynolds, Yuanzhong Xu, James Molloy, Jilin Chen, Michael Isard, Paul Barham, Tom Hennigan, Ross Mcllroy, Melvin Johnson, Johan Schalkwyk, Eli Collins, Eliza Rutherford, Erica Moreira, Kareem Ay- oub, Megha Goel, Clemens Meyer, Gregory Thornton, Zhen Yang. Henryk Michalewski, Za- heer Abbas, Nathan Schucher, Ankesh Anand, Richard Ives, James Keeling, Karel Lenc, Salem Haykal, Siamak Shakeri, Pranav Shyam, Aakanksha Chowdhery. Roman Ring, Stephen Spencer, Eren Sezener, and et al. Gemini 1.5: Unlocking multimodal understanding across millions of doi: 10.4850/ARXX.2240.555552 URL https://ddo .ww.iioogg11..0 48550/.8555/aa0 240..
tokens of context. CoRR, abs/2403.05552 2024.
15

---

Published as a conference paper at ICLR 2025
Luka Ribar, Ivan Chelombiev, Luke Hudlass- Gallla,,,Glley,, Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: Bandwidth- LLM inference. In Forty-fit first International Conference on Machine Learning. ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL ew. net forrm/ormm/// 1==55ddxmmtt
Jay Shah, Ganesh Bikshandi, Ying Zhang. Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. CoRR, abs/2407.0860x 2024. doi: 10.48550ARXIV.2407..00 48550/arXiv www. 08608
URL httpps://ddt htttp::/ddo..oz/ICC. org/10.
Uri Shaham, Maor Ivgi, Avia Efrat, Jonathan Berant, and Omer Levy. Zeroscrolls: A zero- shot benchmark for long text understanding. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Findings of the Association for Compulational Linguistics: EMNLP 2023, Singa- pore, December 6-10, 2023, pp. 7977-7989. Association for Computational Linguistics, 2023. URL https://doi. org/o0g.g/... 18653/
doi: 10. www.iiisnnssm...................
v v1/2023. Findings
Han Shi, Jiahui Gao, Xiaozhe Ren, Hang Xu, Xiaodan Liang, Zhenguo Li, and James Kwok. Sparsebert: Rethinking the importance analysis in self-attention. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learn- ing Research, pp. 9547-9557. PMLR, 2021. URL http://proceeddnns wttpeeiiiinnsm........... press/ v v139/s wtt9/shi21.. html
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. CoRR, (bs/2104.09864, 2021. URL https: http:///axxvv.oo org/
abs/ /2s 220 2104. 09864
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. CORR, abs/2307.08621, 2023. doi: 1044550AAXXV.23000066611 URL https:://oll org/10. 48550/arXiva 08661
Jiaming Tang. Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. QUEST: query-aware sparsity for efficient long-context LLM inference. In Forty-first International Con- ference Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL httlhttss://opereevvii
Philippe Tillet, Hsiang-Tsune Kung, and David D. Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Tim Mattson, Abdullah Muzahid, and Ar- mando Solar-Lezama (eds.), Proceedings of the 3rd ACM SIGPLAN International Workshop Machine Learning and Programming Languages, MAPL@ PLDI 2019, Phoenix, AZ, USA,
June 22, 2019, pp. 10-19. ACM, 2019. doi: 10.1145/3115508.332999 7/doi. 10. 11
URL https:
3329973
Szymon Tworkowski, Konrad Staniszewski, Mikolaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Milos. Focused transformer: Contrastive training for context scaling. Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural In- formation Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 fi
16, 2023, 2023. http:htttp///ppee http:://pppppaaaa..cc..co/eaee/p -btp:a/000.......aoaaaaaaat-----
851
html]
Weizhi Wang, Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. Augmenting language models with long-term memory. In Alice Oh, Tristan Nau- mann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Informa- tion Processing Systems 2023, NeurIPS 2023. New Orleans, LA, USA, December 10 16, n wttp:eewww.pps....eaaeeeeeeeeeae ct/paperrcppppeeeeeeeeeeeeeeee http:///w...pppo..aaaaaaa///////
2023, 2023. URL oot.ra/eeeeetttaoaaaaaaaoooo--to html 2bd82705f44793bbn www.eeabdeeaeeeeeeeeeeeee
16

---

Published as a conference paper at ICLR 2025
Chaojun Xiao, Pengle Zhang. Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Song Han, and Maosong Sun. Infllm: Unveiling the intrinsic capacity of Iims for under- standing extremely long sequenees with training-free memory. CoRR, abs/2402.04617,, 2024a. doi: 10445500AAXXI.224000666 URI. hut.https://ddo httt:::://ol..rgg10..6555////// /rt.b//0orrr00....5////11111111 2402. 04617
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024b. URL, https://cpenreval
net/ Eorummidddi-o............. NG
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Mar- tin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan. Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. Effective long-context scaling of foundation mod- els. In Kevin Duh, Helena Gómez-Adorno, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: LLnng Papers), NAACL 2024, Mexico City, Mex- ico, June 16-21, 2024, pp. 4643-4661 Association for Computational Linguistics, 2024. doi: URL hRRLhhtss://do// http:///ooo..oggg.... 10...//00//...13/////0000000
10. 2At.://11/222.NAALLLLL. naacl-long 260
Peng Xu, Wei Ping. Xianchao Wu, Lawrence McAfee, Chen Zhu, Zihan Liu, Sandeep Sub- ramanian, Evelina Bakhturina, Mohammad Shoeybi, and Bryan Catanzaro. Retrieval meets long context large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL htttps://openrews ew net/ forummrrmmmmmmmm www.::mdxdwwwwwxwwmmmmm
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jianxin Yang, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang. Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Xuejing Liu, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yugiong Liu, Zeyu Cui, Zhenru Zhang, Zhifang Guo. and Zhihao Fan. Qwen2 technical report. CoRR, abs/2407. 10671, 2024. doi: 10.48550/ARXXI.24477.. 10671. URL. phttps://do/ 10671
48150...4555/aiiii.. 2407.
Alex Young, Bei Chen, Chao Li, Chengen Huang, Ge Zhang, Guanwei Zhang, Heng Li, Jiangcheng Zhu, Jianqun Chen, Jing Chang, Kaidong Yu, Peng Liu, Qiang Liu, Shawn Yue, Senbin Yang, Shiming Yang. Tao Yu, Wen Xie, Wenhao Huang, Xiaohui Hu, Xiaoyi Ren, Xinyao Niu, Pengcheng Nic, Yuchi Xu, Yudong Liu, Yue Wang, Yuxuan Cai, Zhenyu Gu, Zhiyuan Liu, and Zonghong Dai. Yi: Open foundation models by Ol.ai. CoRR, 0as/24033044552 2024. doi: 10. 0S550/ARXIV.2403.0466 URL http:: /doi. org/10.../..ooo 48550/ arxiv 2403. 04652
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: Trans- formers for longer sequences. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, De- cember 12, 2020, virtual, 2020.
Aohan Zeng, Bin Xu, Bowen Wang, Chenhui Zhang, Da Yin, Diego Rojas, Guanyu Feng, Hanlin Zhao, Hanyu Lai, Hao Yu, Hongning Wang, Jiadai Sun, Jiajie Zhang, Jiale Cheng, Jiayi Gui, Jie Tang, Jing Zhang, Juanzi Li, Lei Zhao, Lindong Wu, Lucen Zhong, Mingdao Liu, Minlie Huang, Peng Zhang, Qinkai Zheng, Rui Lu, Shuaiqi Duan, Shudan Zhang, Shulin Cao, Shuxun Yang, Weng Lam Tam, Wenyi Zhao, Xiao Liu. Xiao Xia, Xiaohan Zhang, Xiaotao Gu, Xin Lv, Xinghan Liu, Xinyi Liu, Xinyue Yang, Xixuan Song, Xunkai Zhang, Yifan An, Yifan Xu, Yilin
17

---

Published as a conference paper at ICLR 2025
Niu, Yuantao Yang. Yueyan Li, Yushi Bai, Yuxiao Dong. Zehan Qi, Zhaoyu Wang, Zhen Yang, Zhengxiao Du, Zhenyu Hou, and Zihan Wang. Chatglm: A family of large language models from GLM-130B to GLM-4 all tools. CoRR, abs/2406. 1os//240..1299,, 2024. doi: 10.48550/ARXIV24006. 12793. URL htt.pp//./.l...rr000...55///////
2406. 12793
Hailin Zhang, Xiaodong Ji, Yilin Chen, Fangcheng Fu, Xupeng Miao, Xiaonan Nie, Weipeng Chen, and Bin Cui. Pqcache: Product quantization kveache for long context LLM inference. CoRR, abs/2407. 2024a doi: 10.48550/ARXIV22407..0 12820. CtLhttts::/// //o./.oro///////o/o/ooooooo. org/
10. 4w55448555iivvv.. 2407. 2820
Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, and Zhicheng Dou. Soaring from 4k to 400k: Extending Ilm's 's context with activation beacon. CoRR, abs/2401 003462, 2024b. doi: 10. 48550/ARXIV.2401.055 ...446. URL https: /dii. org/10......oo...ooooo.. 48550/ arwiv. 2wt.:/.....11.34.000
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. xcbench: Extending long context evaluation beyond 100k tokens. CoRR, abs/2402. 13718, 2024c. doi: 10.48550ARRXV224022 2402 13718
13718. URL httph:h/toot http://tio.ooro.ooee/e 111110.455//aa.
Zhenyu Zhang, Ying Sheng. Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song. Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen. H20: heavy-hitter oracle for efficient generative inference large language models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural In- formation Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. 6eee2fac0155777881777eccceb222 ott.r//cceeecccaaaooaaattterreee Abstract
html
Zhenyu Zhang, Runjin Chen. Shiwei Liu, Zhewei Yao, Olatunji Ruwase, Beidi Chen, Xiaoxia Wu, and Zhangyang Wang. Found in the middle: How language models use long contexts better via plug-and-ple positional encoding. CoRR, abs/2403.04777 04797, 2024d. doi: 0044855/ARXXI.2400.. 04797. httlhhtps:////t : www.g/g........ 48550 / 10/8555///iiii..// aww.i.ivv...... 2403. 04797)
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, Yifan Du, Chen Yang. Yushuo Chen, Zhipeng Chen, Jinhao Jiang, Ruiyang Ren, Yifan Li, Xinyu Tang, Zikang Liu, Peiyu Liu, Jian- Yun Nie, and Ji- Rong Wen. A survey of large language models. CoRR, abs/2303. 18223, 2023. doi: 10.48550/ ARXIV.2303. 18223. URL /tKLL:tteee//oo..or//////5/////. 230
18

---

Published as a conference paper at ICLR 2025
A RATIONALE
In practice, we can achieve the multi-objective optimization goal (Equation (T) by setting a toler- ance rate "Is for the size of the selected subset |.S Given this constraint, the goal to minimize the attention difference IA(Q, K, V) A(Q, K, S)S to maintain the effectiveness of the sparse attention mechanism while reducing its computational cost. This can be formulated as the following constrained optimization problem:
min S
subject to
|S| S 7s
SC ((i,)) 1 S S (jjj,j
We now provide a deeper understanding of how optimizing the sparse attention objective (Equa- tion 3) leads to achieving the goal (Equation F. For brevity, let us focus on the optimization problem for single query position, omit the subscript i, and consider the 7th query position
in the following analysis. The standard attention computation for this query position is given e
by: AQQ, K, V) Σin
and the sparse attention computation for this subset is: E. Ep-re
AQQ, K, V, = EjEs Eyese"
Let the attention scores be defined as 0j FL
E The
difference between the full attention and sparse attention computations be expressed as:
AQQ, K,Q,K,
Σ
en
Σ jeS
E
Eesse"
1y==1
j=1
1
Σ
agty
(5)
aj j'ES
Σ ajpy - jES
(E
Σ jlES
2z
Let as Ejes aj. Then, as S ag for S, and as 11 as for k & S. Substituting these bounds into Equation 5 obtain upper bound for the absolute difference between full and sparse attention:
-(Q,KKV)A(QQKK, -AL1V~-AQQAA,) A(Q, - Caslust +
1)
Σ
-
a5
JES
jES
1
< (1 as)
l0yl
1
as
(6)
as
JES
JES
I.
(1 as)
J=1
To minimize the upper bound of the error, we can maximize the sum of normalized attention scores as in the selected subset, subject to constraint on the subset size:
exp(x;)
MAX S
subject to
|S| S 7s
JES 2 exp(x"))
By applying duality principles (Appendix B for detailed derivation), objective (Equation (7) into its dual form, which minimizes the subset size subject to a constraint on the sum of normalized attention scores. Given tolerance 0 aa 1, the dual optimization
transform the primal
19

---

Published as a conference paper at ICLR 2025
objective for single query position becomes:
exp(x,) JES 2j'=1 (xp(x.).)
min s
S),
subject
2 Ta
The general objective for all query positions, shown in Equation 3 can be obtained by aggre-
gating the objectives for individual query positions (Equation (8)) across all E [n]]
B DUAL FORM: BALANCING SUBSET SIZE AND ATTENTION SCORES
To better understand the relationship between the original objective and its dual form, we derive the dual objective step by step. For brevity, let us focus on the optimization problem for a single query position, omit the subscript i, and consider the 7th query position in the following analysis.
Step 1: Binary Variables
If 1f E S and 2j otherwise, the primal objective is:
I.
exp(x;))
max (0,(
Σ
2j Ey'=1 exp(x3)))
subject to
Σ
S 78
Step Relax Binary Variables
Relaxing to continuous variables:
exp(x;)
max
subject to
Σ
SS
=1=1 exp(x,,))
==1
1=1
Step 3: Lagrangian Formulation
Let's form the Lagrangian for the primal problem:
exp(2z)
L(z,A,,uu) = 2
+ A
2g 78)
Σ n3 + Σ U5(2g
E exp(xt))
2 0 is the multiplier for size constraint (4j 2 0 are multipliers for 2j 2 0 and vj 2 0 are multipliers for 2j S 1.
Step 4: KKT Conditions
The KKT conditions give:
Stationarity:
OL
exp(x;)
A 11j Vj
02j
- exp(113)
Complementary slackness:
L
A
2j - 78)
j=1
Ví - - =0
Step 5: Optimal Solution Structure
From stationarity condition:
if exp(2,)
A
exp(2gg)
0 if exp(2,) 1
Ey= exp(21))
Step 6: Duality The dual problem becomes finding minimum |S] (or equivalently Σ/= 2q) that achieves required attention mass %a- This gives us:
exp(x,) %o
min 053
zj subject to
2j
Ej= exp(xx)
j=1
20

---

Published as a conference paper at ICLR 2025
Understanding the Duality Transformation The transformation from primal to dual form arises naturally from the optimality conditions of the Lagrangian. Given the primal problem exp(æ,) max;; 2j 2j Eي.م exp(z,)) subject to Σ) 2j S 7s. the Lagrange multiplier À associated with the size constraint effectively becomes threshold on attention scores. At optimality, positions
selected (i.e., 2j 1) if and only if their normalized attention score -/ exp(2,,) exceeds X. This exp(x,)
exp(x;)
Ta. where To corresponds to the optimal attention mass achieved in the primal problem. The equivalence holds by strong duality, as the primal problem's 's convex relaxation ensures zero duality gap. The dual objective minimizes the size of the selected subset S while ensuring that the sum of normalized attention scores in the subset is greater than or equal threshold )a- This is precisely the same objective defined in Equation (3)
naturally leads to the dual formulation min Σ 2y subject to Σj ビ
2a.
exp(x,)) "xp(z,,),
By deriving the dual form and showing that is equivalent to our main objective, we establish
strong connection between the Equation (7) and Equation (3) formulations. This connection high- lights the inherent trade-off between the subset size and the attention scores, and it provides solid theoretical foundation for our optimization goal.
C DETAILED LATENCY OF DIFFERENT METHODS
To highlight the efficiency and performance advantages of our method, we present detailed atten- tion computation latency comparisons on the RULER dataset. Specifically, we evaluate the average latency of a single attention function call for long sequences (64k and 128k tokens) and the overall average latency across all sequence lengths. As shown in Table 3 our FlexPrefill method outper- forms competing approaches by achieving better performance while maintaining lower latency.
Table 3: Latency comparison of different methods various models and sequence lengths RULER.
Models
Methods
64k Time (ms)
128k Time (ms)
Avg Time (ms)
Avg Score
Full-attn Streaming LLM
157.88
658.83
144.97
88.70
88.41
180.40
56.08
LLaMA
Minference
152.56
215.53
116.63
Ours (? 0.9)
61.92
185.75
50.66
Ours (7 0.95)
85.4
271.07
70.61
89.18
Full-attn Streaming I.M.M
267.26
1066.21
237.6 65
89.06
88.1
179.95
55.72
65.91
GLM
Minference
295.95
493.12
226.71
(7 0.9)
279.82
71.89
88.58
Ours (y 0.95)
120.07
408.74
100.75
89.34
Full-attn Streaming LLM
243.17
1018.07
223.53
79.49
107.45
218.96
68.40
MInference
171.48
305.69
138.86
76.00
Ours (? 0.85)
81.77
244.90
53.85
74.28
Ours ( 0.9)
102.32
319.43
83.30
76.32
Full-attn Streaming L.M
137.53
$69.86
125.67
68.83
77.62
157.51
49.46
52.96
Qwen
Minference
220.51
306.47
162.25
66.07
(7 0.85)
133.19
38.00
70.39
Ours (Y 0.9)
61.51
173
47.62
70.75
D COMPARISON WITH ADDITIONAL BASELINES
To further assess the efficiency of our proposed method, we provide performance comparisons against additional baselines. Specifically, we evaluated LM-Infinite (Han et al.] 2024). In- fLLM (Xiao et al.) (024a)) and MoA (Fu et al. 2024a) using the -i--a----B-Innttuctt
(Pekelis) 3.B-
ct al. 2024) models. Additionally, we assessed HIP (Lee et al. 2024) using the performed the RULER dataset.
Instruct- 128k (Dubey al. 2024 model. All evaluations
21

---

Published as a conference paper at ICLR 2025
As shown in Table 4 inference speed.
method consistently outperforms these baselines in both performance and
Table 4: Performance comparison of additional methods on various models and sequence lengths on RULER
Models
Methods
4k
8k
16k
32k
64k
128k
Avg
128k speedup
Full-attn
91.11
89.18 18
88.94
84.37
81.01
77.88
85.42
Minference
90.87
89.90
90.38
82.21
81.73
78.13
85.54
1,89x
MoA
91.35
88.22
85.58
83.9
81.97
76.44
84.58
1.31x
3LaMA-3
InfLLM
89.4
79.8
43
39.5
62.9
1.51x
LM-Infinite
91.11
54.81
38.99 94
21.39
14.66
OOM
35.82
Ours 0.9)
89.18
89.18
89.66
83.65
79.33
77.64
84.78
3.47x
Ours (2 0.95)
90.87
89.9
90.38
82.21
81.25
79.81
85.74
2.43x
Full-attn
95.67
93.75
93.03
87.26
84.37
78.13
88.70
MInference
95.67
93.99
93.27
86.54
84.86
58.17
79.09
3,05x
LLaMA-3.1
HiP
96.0
94.7
79.9
58.2
85.4
1.70x
Ours 0.9)
95.43
93.27
94.23
87.98
81.49
77.89
88.38
3.49x
Ours 0.95)
95.43
93.51
94.71
89.42
82.93
79.09
89.18
2.43x
E PERFORMANGE-LATENCY TRADE-OFF WITH DIFFERENT y
In Figure 4 we plot model performance against the average prefill time for a single attention head. The average latency of the full attention model is reported follows: 4.5 ms (averaged across all input lengths), 20.58 ms (for an input length of 128k), and 1.19 ms (for an input length of 32k). Our proposed method demonstrates varying latency depending the parameter % while maintaining comparable performance. We provide a comprehensive comparison for different values of 7, facilitating empirical selection process. As shown in Table reducing ? results in faster speeds but with a slight trade-off in performance. In practice, ? can be flexibly adjusted to meet specific performance speed requirements.
Table Performance comparison different y for Llama-3.1-8B model
4k
8k
16k
32k
64k
128k
Avg
128k speedup
(Full-Attn)
95.67
93.75
93.03
87.26
84.37
78.13
88.70
0.97
95.43
93.99
94.47
89.66
83.41
79.09
89.34
1.89x
0.95
95.43
93.51
94.71
89.42
82.93
79.09
89.18
2.43x
0.9
95,43
93.27
94.23
87.98
81.49
77.89
88.38
3,49x
0.85
95.91
93.51
92.31
86.54
80.05
74.28
87.10
4.60x
0.8
96.39
92.79
92.31
87.74
80.05
71.39
86.78
5,20x
0.75
95.19
93.51
92.31
86.06
78.13
70.43
85.94
6.75x
0.7
95.67
93.03
91.35
84.62
76.44
71.15
85.38
7.75
0.65
95.91
92.55
91.35
83.65
74.04
67.07
84.09
8.79x
0.6
95.67
93.27
91.83
82.21
74.52
62.02
83.25
9.81x
F ADDITIONAL ABLATION STUDY
F.1
QUERY-AWARE HEAD THRESHOLD
We evaluate the impact of different т on the model performance and provide the detailed perfor- mance of the model under different context lengths in Table
F.2 TRITON BLOCK SIZE
We explore the impact of adjusting Triton block sizes, specifically comparing block sizes of 64 and 128. The results in Table 7 show that the block size does not have significant effect the performance of the model, so different block sizes can be flexibly selected according to different hardware.
22

---

Published as a conference paper at ICLR 2025
Table 6; Performance comparison of different models using different Query-Aware head threshold the RULER dataset.
model
4k
8k
16k
32k
64k
128k
avg
95.43
93.75
94.71
89.67
82.45
78.61
89.10
0.1
95.43
93.51
94.71
89.42
82,93
79.09
89.18
LLaMA
0.2
95.19
93.75
94.95
89.42
83.17
77.88
89.06
0.3
95.19
93.27
93.99
89.66
82.21
75.96
88.38
93.75
91.35
90.14
91.83
86.78
81.49
89.22
0.1
93.51
91.83
89.90
91.35
86.06
83.41
89.34
GLM
0.2
93.99
92.31
90.38
91.11
86.78
83.66
89.70
0.3
93.51
92.07
89.66
90.14
87,26
82.69
89.22
93.51
87.26
83.41
71.88
62.74
56.49
75.88
0.1
93.27
86.78
83.89
72.36
64,66
56.97
76.32
Yi
0.2
93.75
86.54
83.41
73.80
62.50
57.93
76.32
0.3
93.27
85.82
84.37
73.56
62.02
57.45
76.08
89.42
87.50
82.21
81.73
56.97
18.99
69.47
0.1
90.39
89.91
83.17
81.25
59.14
20.67
70.75
Qwen
0.2
90.63
88.22
82,21
80.29
57.93
19.47
69.79
0.3
91.83
88.46
83.41
80.05
58.41
19.47
70.27
Tahle 7: Performance comparison of different models using different Triton Mack_size the RULER dataset
model
block size
4k
8k
16k
32k
64k
128k
avg
128
93.51
91.83 83
89,90
91.35
86.06
83.41
89.34
GLM
94.23
92.07
89.90
91.11
86.06
82.69
89.34
128
93.27
86.78
83.89
72.36
64.66
56.97
76.32
Yi
64
93.03
86.06
83.89
75.48
61.06
$5.05
75.76
128
90.39
89,91
83.17
81.25
59,14
20.67
70.75
Qwen
64
91.11
89.90
83.65
82.93
60.10
20.43
71.35
F.3 REPRESENTATIVE QUERY SUBSET
We perform ablation experiments on the representative query subset used in our method. We replace the position of the subset used sparse pattern determination and vertical-slash index selection from the sequence end to the middle. Results in Table [hsw that replacing the subset for sparse pattern determination has no significant effect on model performance. However, replacing the subset for vertical slash sparse index selection significantly reduces performance. This demonstrates the rationale for choosing the last block _siee query vectors as the representative query subset.
Table 8: Performance comparison of the Llama-3.1 1--maa-..--B-----seeee --Struutttt model the RULER dataset using different representative query set, where last denotes using the last block_size query vectors, and middle means using the middle _ize query vectors.
model
Q (sparse pattern)
Q (vertical slash)
4k
8k
16k
32k
64k
128k
avg
last
last
95.43
93.51
94.71
89.42
82.93
79.09
89.18
LLaMA
middle last
last
95.43
93.27
94.71
89.91
82.93
77.88
89.02
middle
68.27
51.68
38.46
40.87
19.23
10.58
38.18
F.4 ALTERNATIVE IMPLEMENTATION FOR QUERY-AWARE INDEX SEARCH
In Algorithm we flatten and normalize the estimated attention map, then identify the minimum computational budget required for the sum scores to exceed a given threshold 7. We also explore alternative implementation that performs query-wise index selection. In this approach, the cumu- lative attention scores of selected key blocks exceed y for each query block. As shown in Table l this altemnative achieves comparable performance. We chose the global approach with the flatten
23

---

Published as a conference paper at ICLR 2025
operation primarily for its implementation efficiency, as it simplifies the attention mechanism while maintaining performance parity.
Table 9: Performance comparison of different implementations for Query-Aware index search
Implementation
LLaMA
GLM
Yi
Qwen
w/ flatten
89.18
89.34
76.32
70.05
wo/ flatten
89.78
76.08
70.55
F.5 MINIMUM BUDGET LIMIT
We conduct ablation experiments to analyze the effect of the minimum computational budgets, Ta- ble[oodemonssrraees that when y high, the model effectively captures most the important tokens, making minimum budget limit unnecessary. However, with smaller the number of tokens lected by some attention heads may insufficient, causing them to malfunction. In this scenario, implementing minimum budget threshold significantly enhances the model's performance. More- over, this minimum budget limit does not need to be excessive, as increasing the budget further does not lead to better results.
Table 10: Performance comparison of 1.ama-3.1-8B- I-sttuutttttt and GLM-4-9B-CT model using different min- imum budget on the RULER dataset.
model
budget
8k
16k
32k
64k
128k
avg
0.95
94.23
93.27
93.03
92.31
85.34
78.37
89 .2.42
0.95
1024
95.43
93.51
94.71
89.42 42
82.93
79.09
89 18
0.95
2048
95.19
93.27
94.23
89.90
82.93
77.64
88.86
LLaMA
0.9
91.35
88.70
89.66
85.82
77.64
75.48
84.78
1024
95.43
93.27
94.23
87.98
81.49
77.89
88.38
2048
95.67
94.47
94.23
89.90
83.41
75.24
88.82
0.95
92.07
90.15
90.87
90.63
86.06
83.17
88.82
0.95
1024
93.51
91.83
89.90
91.35
86.06
83.41
89.34
0.95
2048
93.51
92.55
89.90
91.59
86.54
83.17
89.54
GLM
0.9
86.78
83.17
87.50
89.66
83.17
81.25
85.26
0.9
1024
93.75
91.59
89.42
90.63
83,90
82.21
88.58
0.9
2048
93.51
92.07
89.66
91.35
84.62
81.73
88.82
F.6 MAXIMUM BUDGET LIMIT
We perform ablation experiments to analyze the effect of the maximum computational budgets for each attention head on model performance and latency. Table shows that capping the maximum computational budget affects different models variously. Imposing maximum budget constraint on models like LLaMA leads to negative impacts. Conversely, for models like GLM, such a con- straint maintains or even enhances performance. These findings highlight that different models are optimized for processing different context lengths.
G LATENCY BREAKDOWN
The computational complexity of FlexPrefill comprises the following components:
Representative Attention Score Computation: approximately O(bdd), where is the
block size, d is the hidden dimension, and n is the sequence length. This includes the computation of attention scores between the representative query set Rgx and all key vectors K E Raxd
approximately O(bn), involving block-pooles attention score estimation and the calculation of the Jensen-Siee Shanmnnn divergence between the estimated and true at- tention distributions.
Pattern Search:
24

---

Published as a conference paper at ICLR 2025
Table 11: dataset.
Performance comparison of different models using different maximum budgets the RULER
model
max budget
4k
8k
16k
32k
64k
128k
avg
95.43
93.51
94.71
89.42
82.93
79.09
89.18
32k
95.43
93.51
94.71
89.42
84.14
75.72
88.82
LLaMA
16k
95.43
93.51
94.71
89.42
84.37
72.12
88.26
8k
95.43
93.51
94.95
87.98
84.13
63.87
86.98
93.51
91.83
91.35
86.06
83.41
89.34
93.51
91.83
91.35
86.54
84.14
89.54
GLM
93.51
91.83
90.38
87.02
82.93
89.26
93.51
91.83
91.11
90.87
86.06
82.45
89.30
93.27
86.78
83.89
72.36
64.66
56.97
76.32
32k
93.27
86.78
83,89
72.84
63.46
58.17
76.40
16k
93.27
86.78
83.65
71.63
62.74
57.69
75.96
8k
93.27
87.02
84.62
73.56
62.50
54.57
75.92
90.39
89,91
83.17
81.25
59.14
20.67
70.75
90.39
89.91
83.17
81.25
57.69
20.19
70.43
Qwen
90.39
89,91
83.17
81.97
59.86
21.15
71.07
90.39
89,91
82.45
80.77
59.62
19.71
70.47
Sparse Index Construction: approximately On n), required for sorting the represen- tative attention scores and constructing the sparse attention indices.
Sparse Attention Computation: approximately O(an²d), where a is the sparsity factor, representing the fraction of computations performed relative to dense attention.
In contrast, the standard dense attention mechanism has computational complexity of O(n'd) d). FlexPrefill introduces modest overhead (approximately (an'd + O(n log n) + ((nd)), which significantly offset by the computational savings achieved through sparsity.
Figure 6 presents the practical latency measurements for each component of FlexPrefill, including representative attention score computation, attention pattern search, sparse index construction, and sparse attention computation. At shorter input lengths, non-attention computation overheads are higher. As input length increases, index search construction time grows, but its percentage gradually decreases.
Sparse Attention Latency Breakdown Comparison
Sparte Attention
Index Search
Pattern Search
Representative Attention
8k-64k Latency
126k-5 52k Lalency
1.75
60
1.50
Latency (mm)
1.25
40
1.00
30
0.75
0.50
10-
0.00
128k
512k
Context Length
Context Length
Figure 6: Sparse Attention Latency Breakdown Comparison the contribution of different components (Sparse Attention, Index Search, Pattern Search, Representative Attention) to the overall latency for various input lengths. As input length increases, the proportion of time spent on sparse attention computation grows, while other components' relative contributions decrease.
different context lengths. The graph shows
25

---

Published as a conference paper at ICLR 2025
H SPARSE PATTERN DISTRIBUTION
We analyze the Jensen-Shannon distance and sparse attention pattern distribution, focusing on Query-Aware head configurations. Figure 7 indicates that block-wise attention score estimation varies with task type and context length. Figure g shows most attention heads use Vertical-Sie patterns, with fewer Query-Aware patterns mainly the model's first layer. These ablation studies provide insights into our method's s components and configurations, enabling informed decisions to optimize sparse attention mechanism performance and efficiency.
Distance
Distance
4
4-
8
8
12
16
16
20
20
24
24
28
28
12
16
20
24
28-
12
16
20
24
28-
head
head
(a) context, task
(b) 128k context.
Js Distance
Js Distance
8
8
12
12
16
16
20
20
24
24
28
28
28
28
head
head
(c) 32k context, task A
(d) 32k context, task B
Figure 7: Shasen-Shannon (JS) distance heatmaps comparing sparse attention pattern distributions across dif- ferent attention heads and layers. The comparison shown for different context lengths (128k vs. 32k) and task types (lask A vs. task B). Darker colors indicate lower JS distance, suggesting more accurate attention estimation of Query-Aware pattern.
SPARSE ATTENTION MASK
The sparse patterns searched by proposed algorithm for different attention heads highly dy- namic, and we visualize them on the Llama-3. 1-B--...--BB--Ittue model. Figure[]showe typical Vertical- Slash and Query-Aware attention heads, and demonstrates that there are large differences in the sparsity rates required for different attention heads.
J SPARSITY RATIO
We visualize the sparsity ratio of different samples on various attention heads in the Llama-3.1 1- 8B-Instruct model. [O that samples with different difficulties require varying sparsity rates and have inconsistent sparsity distributions across attention heads. Additionally, different input lengths exhibit varying sparsity rates, with longer inputs showing higher sparsity ratios.
26

---

Published as a conference paper at ICLR 2025
Qurrydare Heak Cape
dayy Heak Cuaet
Qu-ny Latinheesi
Qunry hwwpeaateeeeaattc Couet
Qunty sare Werick Sab
Quily faar Vittil 5043
Ouety A43 Veriaal
Ganry tol Weick Sas.
15
2
25
10
10
(a) Task A, 128k, 7 ==0.1 (b) Task B, 128k, 7 0.1 (c) Task A, 32k, T 0.1
(d) Task B, 32k, 7
Dany sar
Dany 1ar
Nim
N4m Hogds
Nmm
(c) Task A, 128k, 0.2 (f) Task B, 128k, 0.2 (g) Task A, 32k, 7 0.2
(h) Task B, 32k, 0.2
Figure 8: Comparison of the distribution of the number of different attention patterns between layers for differ- ent context lengths (128k vs. 32k) and task types (lask A vs. task B).
(a) Vertical-Slar Slash head
(b) Query-Aware head (1)
(c) Query-Aware head (2)
Figure 9: Visualization of sparse masks for different attention heads the Llama-3.1 1-amma--.--8B-mmae -BB---ruttttt model. (a) shows the sparse mask of Vertical-Slan heads. (b) shows the sparse mask Query-Aware head, where there are many diverse blocks that stray from specific pattern. (c) shows that Query-Aware head may still cxhibit the Vertical-Slast pattern.
total 5p3 why 79.83
1.0
total sparsity: 85.72
1.0
total spansiy; 91.27
1.0
total sparsity 91.08
1.0
0.8
0.8
0.8
0.8
0.8
0.6
0.6
06
0.4
0.0
12
(a) 64k context, task A
(b) 64k context, task B (c) 256k context, task A
(d) 256k context, task B
model. The heatmaps show varying sparsity distributions for different sample types (task vs. task B) and context lengths (64k vs. 256k). Darker colors indicate lower sparsity. Longer inputs (c. d) exhibit higher overall sparsity ratios compared shorter inputs b).
Figure 10: Visualization sparsity ratios across different attention heads the 1aama-311-8B-Isstte
27


Reviews:
[{"review_id": "xVjzQ4rDpF", "reviewer": "Reviewer_35xX", "summary": "This paper introduces a pre-filling mechanism for LLMs, FlexPrefill, which dynamically adjusts sparse attention patterns and computational budget with respect to different input examples. There are two key innovations:\n- Question aware sparse pattern determination, which leverage Jensen-Shannon divergence to determine whether an attention head should apply the \"question-aware\" pattern with respect to the given query. If not, it will apply the \"vertical-slash\" pattern similar to MInference.\n- Cumulative-attention based index selection, which selects the sparse index by ensuring the sum of the normalized attention scores meets a predefined cumulative attention threshold, rather than selecting the sparse index to meet a given budget.\n\nExperimental results indicate that FlexPrefill achieves comparable results with previous SOTA methods with GLM and Yi, and outperforms previous SOTA methods with a large margin with Llama and Qwen.", "questions": "(a) How is the subset of queries $\\hat{Q}$ selected for representative query selection in Query-Aware sparse pattern determination and Vertical-Slash head index selection? Is it randomly sampled?\n\n(b) How was Figure 8 generated? Was it based on a specific case study or aggregated over all examples from one of the benchmarks?\n\n(c) In Algorithm 4, the sparse index is selected globally (indicated by the flatten operation). What if perform query-wise index selection?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The key idea that the complexity of input data varies so we should dynamically adjust sparse attention patterns and computational budget with respect to different input examples is interesting and reasonable.\n- Experimental results indicate that FlexPrefill achieves comparable results with previous SOTA methods with GLM and Yi, and outperforms previous SOTA methods with a large margin with Llama and Qwen.", "weaknesses": "1. Figure 8 indicates that only a small portion of attention heads use Query-Aware patterns, whereas the fallback Vertical/Slash pattern is also applied in the baseline method, MInference. Given this observation, further illustration of the original attention patterns and an analysis of why FlexPrefill demonstrates significant performance gains over MInference are necessary.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CdCQBsYxbv", "reviewer": "Reviewer_9EQr", "summary": "The paper introduces FlexPrefill, a novel sparse attention mechanism aimed at improving the efficiency of long-sequence inference for large language models (LLMs). The proposed approach addresses the computational challenges associated with long-sequence processing by dynamically adjusting sparse attention patterns and computational budgets based on input characteristics. FlexPrefill consists of two key components: Query-Aware Sparse Pattern Determination, which selects between diverse and structured attention patterns using Jensen-Shannon divergence, and Cumulative-Attention Based Index Selection, which ensures efficient attention computation by dynamically selecting query-key indices based on attention scores. Experimental results demonstrate  improvements in both speed and accuracy compared to existing methods, such as Streaming LLM and Minference.", "questions": "1. Can you provide a detailed computational complexity analysis on the proposed sparse attention mechanism?\n\n2. Given the complexity of the multi-step decision-making process, how feasible is it to integrate FlexPrefill into existing LLM frameworks without significant modifications?\n\n3. In Figure 4, what is the average attention latency of the Full Attention model?\n\n4. Can you include a comparison of how much computation is saved by using FlexPrefill compared to the baselines in Tables 1 and 2?\n\n5. Have you considered including direct comparisons with other recent related works, such as Han et al. (2024) and Xiao et al. (2024a)? How does FlexPrefill compare to these methods?\n\n6. How were the values of the thresholds (e.g., sparse pattern threshold and cumulative attention threshold) chosen? Were these values empirically determined, or based on specific criteria?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Originality: The paper presents an adaptive and flexible approach to sparse attention that addresses the limitations of fixed sparse attention mechanisms. The real-time adjustment of sparse attention patterns is a notable advancement in handling varying input needs during long-sequence inference, contributing to the growing body of research focused on optimizing LLM efficiency.\n\n2. Quality: The quality of the experimental setup is high, with comprehensive comparisons to baseline methods like Streaming LLM and Minference. The authors provide detailed analyses that demonstrate improvements in both speed and accuracy, supporting the efficacy of the proposed method. The ablation studies further strengthen the validity of the claims by examining the contributions of each component of the approach.\n\n3. Clarity: The paper is well-structured and clearly written, making it accessible to a broad audience. The authors provide a thorough explanation of the motivation, methodology, and experimental results.", "weaknesses": "1. Complexity of the Proposed Sparse Attention Mechanism: The multi-step decision-making process, involving both Query-Aware Sparse Pattern Determination and Cumulative-Attention Based Index Selection, introduces additional complexity to the attention mechanism. A detailed computational complexity analysis is missing, which raises questions about how the overhead introduced by this process is justified by the computational savings achieved through sparse attention.\n\n2. Lack of Detailed Comparison of Computational Savings: The paper lacks a comprehensive comparison of how much computation is saved by using FlexPrefill compared to the baselines presented in Tables 1 and 2. This information would be valuable in quantifying the benefits of the proposed approach.\n\n3. Limited Comparison with Recent Related Works: The paper does not include comparisons with other recent related works, such as Han et al. (2024) and Xiao et al. (2024a). Including these comparisons would help to better position the contributions of FlexPrefill within the context of current research.\n\n4. Mixed Performance Benefits: The reported performance benefits compared to the Full Attention model on long-context benchmarks are mixed, as shown in Tables 1 and 2. Additionally, the computational savings compared to the Full Attention model and other baselines are not adequately reported, which limits the understanding of the trade-offs involved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OdJQtr9jzt", "reviewer": "Reviewer_3fJR", "summary": "This paper addressed the flexibility of the sparse attention pattern to various input length and task currently applied in LLM. The authors split the attention heads into two types: the Query-Aware heads that capture variable sparse pattern upon query, and the Vertical-Slash heads that capture structured and invariant sparse patterns observed in pervious works. The authors decide the sparse pattern for each head by evaluating the discrepancy between the estimated and true attention distribution, and select the sparse pattern by summing up the largest attention scores to a threshold. The proposed method balances the computational efficiency with attention accuracy. The authors evaluate their approach on RULER and Infinite Bench and show that their method preserves high performance across multiple context length and tasks.", "questions": "- Q1: Figure4 shows the attention latency and performance on RULER. However, this figure is not comprehensive enough. Why there is only one point for MInference? Why not adjust the computation budget of MInference and plot a curve like \"Ours\"? A curve for \"Fixed Budget Query-Aware\" is also missing.\n- Q2: In Algorithm 4, the attention map is normalized among all the attention scores, but Algorithm 3 only computes the mean values on each vertical and slash lines. The mean operation in Algorithm 3 can not assure the sum of $a_v$ and $a_s$ to be 1, thus the criterion of Algorithm 3 and 4 may be different even with the same $\\gamma$. Could the authors explain why there is no normalization in Algorithm 3?\n- Q3: in Appendix A, Step 7, the equivalent form of the dual objective is not that intuitive. How can we derive this from the equation in Step 6?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper propose a more dynamic and flexible approach for sparse attention calculation. Their method performs sparse pattern search and sparse index selection both in a dynamic, query-aware manner. The flexibility and dynamism distinct it from MInference, and contribute to the performance improvement.\n- The method is based on solid observations. \n- The paper gives a concrete and simple algorithm to construct these sparse attention patterns.\n- The \"Cumulative-Attention Based Index Selection\" method is convincing with the details discussed in Appendix A. The authors derive the algorithm from the original  multi-objective optimization problem in details, providing the rationale in summing up the largest attention scores.", "weaknesses": "Methodology:\n\n- **Insufficient Comparison with Baselines**: There are also other sparse attention method in recent works such as SampleAttention [1] and HiP-Attention [2]. The drawbacks of StreamingLLM in handling longer contexts is obvious in other papers, it would be better to compare to other more recent methods.\n- **Insufficient evaluation of latency**: Although \"Cumulative-Attention Based Index Selection\" improves performance, it may result in unbalanced computation allocation between different attention heads. I think this could lead to an overall increase in latency, as MInference explicitly avoids this by setting a target FLOPs for each pattern. I think it is necessary to show to what extend the dynamic, head-distinctive computation allocation increases the end-to-end latency compared to the fixed budget setting within the same FLOPs.\n- **Determination of attention patterns**: The algorithm select the vertical-slash pattern when the distribution of estimated and true attention is not that close. However, in my opinion, the distribution discrepancy only indicates that the result of block-wise sparse attention deviates the full attention, but doesn't lead to using the vertical-slash patterns. Is there any theoretical or empirical guarantee that vertical flash patterns are better than block-wise pattern in this scenario?\n\nPresentation:\n\n- \"eacct\" -> \"exact\" in Line 1,043.\n- The caption of Table 5 should be above the table.\n- In Figure 4, it would be more clear to show the values of $\\gamma$ for each point, which will be helpful to provide an empirical selection of $\\gamma$. \n- The authors states \"accelerating the computational process\" in Line 357 but Table 1 does not include any result concerning efficiency. It would be better to give a reference to figure 4 here. \n\n[1] Zhu, Q., Duan, J., Chen, C., Liu, S., Li, X., Feng, G., ... & Yang, C. (2024). Near-lossless acceleration of long context llm inference with adaptive structured sparse attention. *arXiv preprint arXiv:2406.15486*.\n\n[2] Lee, H., Park, G., Lee, Y., Kim, J., Jeong, W., Jeon, M., & Hwang, S. J. (2024). HiP Attention: Sparse Sub-Quadratic Attention with Hierarchical Attention Pruning. *arXiv preprint arXiv:2406.09827*.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GQwNvBwmns", "reviewer": "Reviewer_mNvJ", "summary": "This work investigates the LLM long sequence inference efficiency challenge and proposes a novel and flexible method FlexPrefill to solve it. The proposed FlexPrefill featured with two components: Query-Aware Sparse Pattern Determination and Cumulative-Attention Based Index Selection, of which the former one adaptively optimize the usage of the sparse attensions while the latter one ensures the quality of selected sparse attentions. Extensive experiments elaborate the effectiveness and superiority of FlexPrefill.", "questions": "a) Could the authors elaborate the intuition behind the usage of JS divergence? Would the W-Distence be more appropriate since W-Distence help measures the change costs of given two attentions? Also, how the JS divergence can help FlexPrefill to be effective and flexible, would it introduce addtional calculation costs?\n\nb) Do the authors have a vision on how the model structure can impact on the performance of FlexPrefill?\n\nc) For the baseline selections, concerns may arise since the FlexPrefill poses an automatic selection process which makes FlexPrefill more flexible while the baselines do have. It would be better to see the comparisons with methods equipped automatic selections such as Moa[1].\n\nd) In formula (2), should the JS distance be the calculation on a||b and b||a given two distributions a, b. Can the authors elaborate the usage of a\\hat{-}||m and a\\hat{^}||m?\n\n\n**Reference**\n\n[1] Moa: Mixture of sparse attention for automatic large language model compression.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "a) The proposed Query-Aware Sparse Pattern Determination strategy step further on existing work to provide flexibility to sparse attention usage.\n\nb) Promising improvments on extensive experiment compared with SOTA.", "weaknesses": "a) Unclear costs analysis of pre-calculation on JS divergence.\n\nb) The baseline selections are insufficient.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)