A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention
Authors: Heejun Lee
Geon Park
Youngwan Lee
Jaduk Suh
Jina Kim
Wonyong Jeong
Bumsik Kim
Hyemin Lee
Myeongjae Jeon
Sung Ju Hwang
Venue: ICLR 2025 Poster (2025)
Keywords: Efficient Attention Mechanism
Long-context LLM Decoding
KV Cache Offloading
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: To response of quadratic scaling (time & GPU memory) problem of LLM decoding framework, we are proposing sub-quadratic time & GPU memory cost HiP attention.

Abstract: In modern large language models (LLMs), increasing the context length is crucial for improving comprehension and coherence in long-context, multi-modal, and retrieval-augmented language generation. 
While many recent transformer models attempt to extend their context length over a million tokens, they remain impractical due to the quadratic time and space complexities.
Although recent works on linear and sparse attention mechanisms can achieve this goal, their real-world applicability is often limited by the need to re-train from scratch and significantly worse performance. In response, we propose a novel approach, Hierarchically Pruned Attention (HiP), which reduces the time complexity of the attention mechanism to $O(T \log T)$ and the space complexity to $O(T)$, where $T$ is the sequence length. 
We notice a pattern in the attention scores of pretrained LLMs where tokens close together tend to have similar scores, which we call "attention locality". Based on this observation, we utilize a novel tree-search-like algorithm that estimates the top-$k$ key tokens for a given query on the fly, which is mathematically guaranteed to have better performance than random attention pruning. In addition to improving the time complexity of the attention mechanism, we further optimize GPU memory usage by implementing KV cache offloading, which stores only $O(\log T)$ tokens on the GPU while maintaining similar decoding throughput. Experiments on benchmarks show that HiP, with its training-free nature, significantly reduces both prefill and decoding latencies, as well as memory usage, while maintaining high-quality generation with minimal degradation.
HiP enables pretrained LLMs to scale up to millions of tokens on commodity GPUs, potentially unlocking long-context LLM applications previously deemed infeasible.

Published as a conference paper at ICLR 2025
A TRAINING-FREE SUB-QUADRATIC COST TRANSFORMER MODEL SERVING FRAMEWORK WITH HIERARCHICALLY PRUNED ATTENTION
Heejun Lee; Geon Park," Youngwan Lee," Jaduk Suh,' Jina Kim
Graduate School of Artificial Intelligence (AI)
Korea Advanced Institute Science and Technology (KAIST)
Scoul, South Korea {ainl, geon park, ywlee88, jaduksuh,
ac.kr
Wonyong Jeong, Bumsik Kim, Hyemin Lee
Myeongjae Jeon Graduate School of AI POSTECH Pohang, South Korea mj. @ejoeeessttccchac..
LLMOps Team
DeepAuto.ai
Scoul, South Korea
(young, 1 aim,, hailey @ddepputo..a .i
Sung Ju Hwang
Graduate School of Al KAIST, DeepAuto.ai
Seoul, South Korea sjhwang@ka st.ac.kr
ABSTRACT
In modern large language models (LLMs), increasing the context length is crucial for improving comprehension and coherence in long-context, multi-modal, and retrieval-uugmented language generation. While many recent transformer models attempt to extend their context length over million tokens, they remain imprac- tical due to the quadratic time and space complexities. Although recent works linear and sparse attention mechanisms applicability is often limited by the need to re-train from scratch and significantly propose novel approach, Hierarchically
achieve this goal, their real-world
worse performance. In response, Pruned Attention (HiP), which reduces the time complexity of the attention mech- anism to O log 7') and the space complexity O(T), where T is the sequence length. We notice pattern in the attention scores of pretrained LLMs where tokens close together tend to have similar scores, which call "attention local- ity". Based on this observation, we utilize a novel rree-search-lit like algorithm that estimates the top-k key tokens for given query on the fly, which is mathemat- ically guaranteed to have better performance than random attention pruning. In addition to improving the time complexity the attention mechanism, we further optimize GPU memory usage by implementing KV cache offloading, which stores only O (log T') tokens on the GPU while maintaining similar decoding throughput. Experiments on benchmarks show that HiP, with its training- free nature, signifi- cantly reduces both prefill and decoding latencies, as well as memory usage, while maintaining high-quality generation with minimal degradation. HiP enables pre- trained LLMs to scale up to millions of tokens on commodity GPUs, potentially unlocking long-context LLM applications previously deemed infeasible.
INTRODUCTION
Large Transformer-bassee generative language models (LLM) trained on huge datasets have recently demonstrated remarkable abilities in various problem domains, such as natural language under- standing (Touvron al., 2023), code generation (Rozière et al., 2024), and multi-modal question answering (Liu et al., 2023a). This is made possible by the effectiveness of the attention mecha- nism, which learns 2 pairwise relationships between all tokens in a sequence of T tokens. Despite their success, the quadratic complexity of the attention mechanism makes it increasingly challenging to meet growing resource demands when processing longer sequences.
"qqull contributors

---

Published as a conference paper at ICLR 2025
Hierarchy of Language
Hierarchically Pruned Attention (HIP) OTT T)
HW-aware Block Sparsity
Information has locality inside of each hierarchy
finding Imporstnt
Taken Hierarchicaly
0
Dinie Sequence into Chunlo uoae
Sequence
Distanoe
Dog
Which calls the first house cal?
Ruogor
Atronn
widdaat
were
earty
human
crttemart
16 Key Broçks
QHIP
Select Regrasentation Token Chunks
10 Bbda
In Beporrrraiin Soure
Demesiicarion en
Measure chunk importance by compar nq
Makasio CS4
Key Taicess
between query and rspresontative tokon
Co
Select Top-k Important Chunk ooe wewnes oely ferems
Blocks Pruned by HP
Corres the inic
Active Blocs
Figure 1: HiP Attention. HiP dynamically prunes block sparse attention depending on given query token in sub-quadratic cost by utilizing the hierarchy and locality of natural language.
Various approaches have been suggested to handle longer sequences efficiently to overcome this limitation. FlashAttention (Dao et al., 2022; Dao, 2023) has reduced the space complexity to O(T) by fusing the component computations its time complexity remains O(T²), making it less applicable to inference tasks with long contexts. Many other methods (Lee et al., 2023; Beltagy et al., 2020; Zaheer et al., 2020b; Tay et al., 2020; Ki- taev et al., 2019; Tay et al., 2021; Liu et al., 2021) tackle the issue by sparsifying the attention matrix or approximate the attention mechanism using kernel methods to reduce its quadratic complexity. However, these works are not widely employed in real-world LLM serving frameworks because they often lead to performance degradation due to drastic changes in the computation flow and are too complex to implement efficiently for actual speedups. Moreover, they often require extensive fine-tuning or even pre- training from scratch, which can be prohibitively expensive and prevent the timely deployment of production-ready pre- pre-trained models.
avoid storing 722 attention scores one time. However,
In this paper, we define and achieve three fundamental objectives for frameworks tailored to long- context transformer serving frameworks: (1) minimizing the algorithmic complexity of attention mechanisms, (2) enhancing GPU compute efficiency, particularly through TensorCore utilization, and (3) maximizing the effective use of limited GPU memory capacity.
First, to serve long sequence in a timely manner, we propose Hicrarchically Pruned Attention (HiP), an efficient training-free attention mechanism reducing the quadratic time complexity O(T log T) by approximating the top-k key tokens in a sequence. HiP exploits "attention locality", where neigh- boring tokens often have similar attention scores, as shown in Figure 1 (Left). Therefore, as shown in Figure 1 (Center), HiP divides the input sequence into 2k chunks, and the center token in each chunk is chosen to represent its neighbors, driven by the attention locality within the chunk. HiP computes the attention scores of these representative tokens to approximate the importance of each chunk for given query. HiP iteratively refines its selection by starting with the top-k most impor- tant chunks and progressively narrowing them down until each chunk contains single token. This hierarchical top-k key estimation takes O log T) time, which used for sparse attention compu- tation that costs O(T), making the overall complexity of our attention mechanism log-linear. We provide mathematical proof demonstrating that our HiP outperforms random selection, supported by empirical evidence from attention score statistics in Section 4.
Second, we introduce hardware- aware optimizations to enhance GPU compute efficiency for our HiP through block-wise key sparsity, illustrated in Figure 1 (Right). Specifically, our top-k approxima- tion is implemented in a tiled manner (Tillet et al., 2019) so that it can fully utilize matrix multiplier units (MMUs; e.g., TensorCores (Nvidia, 2024)) and achieve the highest possible token-proesssing throughput. Additionally, we integrate our attention mechanism into throughput-opimizce LLM serving frameworks, such as vLLM (Kwon al., 2023) and SGlang (Zheng et al., 2024), further enhancing deployment efficiency.
Lastly, to serve extremely long sequences within the limited GPU memory, we propose a KV cache management strategy that stores only Olog T) tokens in GPU memory (HBM) and offloads the re- maining tokens to host memory (DRAM). The Ollo T') tokens stored in GPU memory are the ones accessed most frequently and are meant to provide quick access for the GPU's MMUs. In contrast, other less frequently accessed tokens reside in main memory and are transferred to GPU memory only upon token access misses. With high access hit ratio in HiP, our memory management scheme effectively meets the demand for limited HBM capacity while leveraging the larger DRAM capacity, preventing token access from becoming bottleneck.
We validate HiP on various benchmarks by applying it to Llama3.1-88 (Meta, 2024). In Long- Bench (Bai et al., 2023), HiP maintains 96% of its relative performance while achieving almost 2.7 x speedup in the prefill stage and 16.5x speedup attention computation in the decode stage with 32k context length compared to Flash Attention. Additionally, in passkey retrieval tasks such as
2

---

Published as a conference paper at ICLR 2025
with Hierarchical Top-k Key Selection
Step Efficient Mask Estimation
T (53.1)
Step 2. Sparse Attention with HIP Attention Mask
O
($3.2)
wKeys
Query Block
Key Block
walk
Values
Keys
Querles
42.0
Attention
Q
23
Estimated Attention Mask .
977
Well you
by the
my
wom-
T
Top
the
10.2
3.5
37.1
Iterations
Querles
Outputs
Well tell
X
Discard Bottom Chunks
X
Keep Top- Chunks
2-44 hy- Select anctions at each horadonn
0 sofumax! M @ (QK V)
walk&'m
Figure 2: Overview of our HiP attention mechanism. In HiP, the model dynamically decides which num- ber of key tokens to attend to for each query by generating a sparse attention mask. The sparse attention mask is generated in a tree search-like manner. At each iteration, the top-k blocks with the largest attention scores selected, and the rest of the branches are discarded. The final mask becomes accurate approximation of the top-k blocks the true attention map. Please refer to Figure 19 for a more detailed illustration.
RULER (Hsich et al., 2024), HiP preserves its original effective context length, while all baselines fail to do so. We also evaluate the effectiveness of the proposed KV cache offloading framework. On a machine capable of serving up to a 16k context length with Flash Attention, our method extends the context length up 64k by offloading the KV cache without significant throughput degradation.
In conclusion, by integrating the three proposed solutions, we present a single long-context serv- ing framework that efficiently manages compute and memory resources while being transparent and easily usable. This extension of serving context length, achieved within the constraints of lim- ited space and compute budgets, delivers substantial benefits for long-context applications, such as question answering with long texts (Kryšciński et al., 2022), multi-agent chatbots (Hu et al., 2024), enhanced retrieval -ugmentt reasoning, and long video data summarization. Furthermore, since our approach is training-fre freee HiP can be seamlessly applied to pretrained LLMs without requiring additional training. As result, we expect our method to be highly practical for wide range of long-context LLM applications.
Our contributions within the proposed framework can be summarized as follows:
We propose a novel, training-free free hierarchically pruned attention mechanism that uses hierarchi- cal score-locality-wwe cost of the attention mechanism to (T log T) time and (T) space complexity (Section 3.
top-k approximation to accelerate LLM serving, reducing the quadratic
We further optimize our HiP mechanism with hardware-awart block-wise tiled optimization using OpenAl Triton, achieving up to speed up to 6.83x speedup in end-to-end decoding for 128k context. (Section 3.2, Table 5)
We implement KV cache offloading to reduce GPU memory efficiency further, increasing serving context from 16k up 64k tokens an RTX 4090 with 8B model (Section 3.3).
RELATED WORKS
Previous studies proposed several attention approximations with linear complexity using either ker- nel methods or sparse attention. Low-rank approximations of softmax attention via kernel methods (Choromanski et al., 2022; Qin al., 2022) achieve faster inference speeds but significantly alter the data flow, leading to performance degradation that is hard to mitigate. In contrast, sparse at- tention methods, which use attention pruning to preserve trained attention scores, allow for simple replacement of pre-trained mechanisms. However, they often require additional fine-tuning to adapt to static attention patterns (Beltagy et al., 2020; Zaheer et al., 2020a; Xiao et al., 2024) or the training of an attention estimator (Lee et al., 2023; Liu et al., 2021). These methods are generally less effi- cient than fused attention techniques (Dao et al., 2022; Dao, 2023) due to their fine-grained sparsity, which prevents optimal MMU utilization. For more details, see Appendix E.I.
3 METHODOLOGY
Given query, key, and value sequences Q, K.K. V Rxxd the conventional single-head attention output O is computed as S = QKT E TXT p= softmax(S) E RT*T O PV e Rxxd where d denotes embedding dimension, and softmax is applied rowe The causal masking and
3

---

Published as a conference paper at ICLR 2025
constant scaling are omitted for brevity. The S and P matrices respectively called the attention scores and probabilities. We focus on the fact that, due to the nature of the softmax function, only the highest attention scores significantly impact the output. Therefore, promising approach to approximating in a sparse format and reducing the complexity from O(T2) is to retain only its top-k elements, detailed in the following equations:
M top.k_mash (QK)) ") (0,1)TxT,
(1)
$ = mask (QK) E TXT
P softmax( E RTTT,
Ô PV E Rxd
(2)
[Sij if Mij =1
where [mask M (S));) = -00 if M.. =0
where top_k_mask( denotes a binary mask which selects the top-k largest elements for each row of the given matrix. Since S is sparse matrix with only kT valid elements, S and Ô in Equation (2) can be computed in (TT) time using sparse matrix operations.
However, obtaining the binary mask M in sub-quadratic time is no easy task. To address this chal- lenging problem, we exploit what we call "attention locality". Observation of attention scores reveal that the scores tend to exhibit local similarity, phenomenon we refer to as attention locality. We ex- ploit this observation by performing tree- based search for the top-k tokens. We divide the sequence into 24 chunks, and then select a representative token from each chunk. Due attention locality, representative token have similar scores to other tokens in its chunk thereby "representing" that chunk. We select the top-k most important chunks based on the attention scores of the representative tokens. By repeating this process, refine the tokens until we details of our method are shown in Section 3.1. We only cover the single-head non-causal case here, but note that method can easily be extended to causal multi-heal attention.
longer divide chunks. Exact
3.1 HIERARCHICAL SCORE-LOCALIT--WAHH
TOP-k ESTIMATION
As shown in Equation (1), our goal is to select the top-k largest elements of each row of pre- trained attention score S without computing the entire matrix. To this end, we use a greedy binary tree search algorithm, as illustrated in the left side of Figure The complete algorithm for mask estimation is presented in Algorithm 1.
For a given query q E Rdd, the first iteration, we divide the key sequence K E Rxxd along the time dimension into ls equal-sized chunks (f11) (1))) (f22) (()))....(())))))))))
12)), where
(1) [j-1)-T] + 1 and 1,0) = ir are the first and last indices of the jth chunk, each.' The J
superscripts denote the iteration number. At each iteration i, we further divide each of the chunks into two equal-sized branches:
B 2 (f)
(6) m 1).
B = (m²) 11))). where my (0) = ( + +1 //2 for j = 1.. k.
(4) is the center key token index for each branch B3 We assume that A representative key index r this representative key represents the entire branch. Thus, among the 21 branches, the top k branches whose representative key's s scores are the highest are chosen for the next iteration:
(( ((+1)) = B!) for j 1 k where {l,,.... [2 == argtope jE(l .. 2k)
(TK K so
(4)
We repeat the above iteration Th = [log2 71 times, i.e., until the length of each branch all becomes 1. In the end, obtain set of indices I [ff(m)e) ..... fk p(mxe)), which is estimation of the top-k indices of K which have the largest attention scores with the query Thus, we obtain m, an estimation of a row of the attention mask 2
m = estimate attn mask: (q. K) == [(),
(5)
In conclusion, this algorithm takes O(T log T') time in total because the total number of iterations is log, T' where each iteration takes constant time Ok), and we do this for each of the queries.
'1.7 denotes rounding to the nearest integer.
21 (z), where A is set, denotes the indicator function: ax 1 if x E A, and otherwise 1 0.
4

---

Published as a conference paper at ICLR 2025
3.2 BLOCK APPROXIMATION OF T ESTIMATION
Despite the log-linear complexity, obtaining competitive latency the state-of-the-ard implementa- tions of dense attention on an accelerator (e.g.. GPU) is difficult. This is because the matrix multi- plier unit (MMU) inside accelerators is optimized for dense attention, where they compute fixed-size blocks of matrix multiplication in a few clock cycles. In contrast, the attention score computation in the top-k estimation of HiP cannot be performed with traditional matrix multiplication because different key matrix is used to compute the dot product for each query vector. To utilize MMU, we use a technique called block approximation during top-k estimation, illustrated in Figure 2 (Right).
In top-k estimation, replace K Rxxd with its tiled version K E RT/bxxbe and Q with its tiled version Q E /2 xd where 6y and by are the size of key block and query block. The top-k estimation iterations are done similarly to before, except that the division and branching of the key sequence done block-wise (using the first dimension of K). Importantly, instead of k k/bb chunks are maintained at each iteration in order to select k tokens, and the score calculation
in Equation (4) is replaced with maXme[Ls,,,.eelll]
[Im.. K (g)) .n:) where q E R xd is the given query block. While this modification enables HiP to reduce the cost further, we internally sample the blocks with stride b.g in the query dimension and 6.k in the key dimension instead of using the full b, by block.
As a result of this optimization, the estimated mask M becomes block-sparse. Therefore, each (bg/bag) x d-block of the query can be matrix- muttii--ul with the same (k/bba) key matrix to obtain (b.//bag) x (bbbs))) elements of S. Thus, b, and bag are critical for the most efficient utilization of the MMU: we can achieve a considerable latency reduction if we set by/bby multiple of 16 or 32, as shown in Appendix E.4. While the choice of by and bet is irrelevant to the MMU utilization, helps reduce the number of top-k estimation iterations.
3.3 KV CACHE OFFLOADING
Decoding Steps with KV offloading
Mask Refresh Step
Cached Mask Step
Thanks to our top-k estimation algorithm, HiP only accesses (kr/bax) log T key states per at- tention head. Moreover, the algorithm's 's mem- ory access pattern exhibits strong temporal lo- cality. Using this fact, we can further enhance efficiency by exploiting the memory hierar- chy: we offload less frequently accessed key- value (KV) states from the GPU to the main memory. This involves caching frequently ac- cessed KV states (hot tokens) by tracking state access patterns of top-k estimation and sparse attention using the estimated HiP mask.
Masking
Sparse Attention
Sparse Krg Attention (5A)
Key
New
New Token
New
Access Log
Mask
Token
Key Tokens for Masking
Olklog
fo lu
GPU
Key-Value Tokens for Sparse Attention ((k) Free makiigg
Road Wiite
Store Every Key Value Tokens O
UNN Request
Offloaded KV Cache
Figure 3: Flow of KV Cache Offloading with HiP.
Our GPU cache that holds the hot tokens consists of two components: token bank containing the actual KV states and a page table with the token-bank index mapping, as shown in Figure 4. One straightforward implementation for the page table would be a vector map: a simple length-?' array of point- ers. While this approach practical for typical sequence lengths (e.g., 128k 1M), its space complexity is O We employ a linear probing hash table to reduce the space complexity, achieving Olog T) space complexity. How- ever, empirical results show that GPU hash map lookups introduce additional latency compared to using a simpler vector-based page table.
Index Translation
Talon Koger
Jecmn Stubcli
Page Table
HashWWpp
Given the distinct memory access patterns in top-k estimation and in sparse attention, we maintain two separate offloading contexts, each containing a page table and set of GPU-resident hot tokens, as illustrated as two separate GPU loaded KV caches in Figure 3. For the top-k estimation stage, Keache == Figure 4: KV Token (ksbbx)) log T key states
Requested Token
held in VRAM, where is hyperparameter Index Translation determining the cache size. For sparse attention, & key and value states are held. In summary, we need to hold (Keache/2 /2 + k) tokens" equivalent of KV states in the GPU. The kernel first queries
5

---

Published as a conference paper at ICLR 2025
the GPU cache when accessing key value tokens. Upon cache miss (which is unavoidable due to the dynamic nature of the attention access pattern), the system attempts to retrieve tokens from the main memory. By using our cache. we can significantly speed up memory access compared to directly accessing CPU memory from the GPU.
In conclusion, we reduce the GPU memory footprint for KV tokens from (T)) Olog T), but this comes with page table overhead that can range between O(T) and O(log T) depending on the data structure used. The overall space complexity is thus determined by the type of page table, allowing for a configurable trade-off between GPU memory efficiency and latency. However, we suggest that users use vector maps in many practical long-context ranges (32-512k) achieve competitive latency compared to Flash attention. Please refer to Section 5.5 for detailed benchmarks.
4 THEORETICAL ANALYSIS
In this section, we justify the design choices of our HiP's approximate top-k key selection algorithm by answering the following questions: (1) Is HiP's key selection algorithm better than the random selection baseline at finding keys with the biggest scores? (2) How should the representative token in each branch be chosen? We answer these questions by providing probabilistic analysis of HiP's key selection algorithm in a simplified setting (k 1), based on the assumption of attention locality.
Observation: keys closer together exhibit higher similarity in attention scores. In each atren- tion head of layer in an LLM, key sequence K" E RTxx is used for computing the attention mechanism. Given a query vector q E Rd, the scores for each key K E RR. can be computed. We investigate how much locality these scores exhibit by studying the correlation between their dis- tance := i - j| and the score difference òa := 8 8y for every i, E [L..T', with sample natural language data. As shown in Figure 6, our empirical observation shows that òλ generally follows normal distribution, whose mean is almost zero and the standard deviation is an increasing function of distance A. More details regarding this observation provided in Appendix A.3.
Analysis. Based on this observa- tion, we assume that we can ap- proximate the difference in attention scores between two keys separated by A tokens as scalar random vari- able 82 ~ (A)))) where (A) an increasing function of A. This can be interpreted as keys that are closer together are more likely to have a similar attention score, which fits well with our observation and at- tention locality assumption. With this assumption, the following Theorem I can be shown.
2.4
6A21
6a15
Mean
6SM
Attention Score Difference
Index difference
3nddxx difference
Figure 6: Score Difference Distribution. We collect the attention
score statistics from the 17th layer and second atlention head of Llama3.1 The left figure shows the raw distribution when A = 500. The middle and right figures show the mean and standard deviation as a function of A.
Theorem 1 (Informal). Consider the case of finding the location of the top-1 key token with the maximum attention score in a context T tokens. Suppose that our locality assumption holds true. We divide the context into two branches with T/2 keys each. Then, the branch whose center token has the bigger attention score is more likely to contain the top-1 key token.
The above shows the effectiveness of one iteration of HiP's key selection algorithm. By recursive application of HiP's key selection iterations, we can intuitively see that the probability of HiP's key selection algorithm finding the location of the top-1 key would be higher than that of uniform random selection as well. Therefore, under the attention locality assumption, on average, HiP's key selection algorithm on average finds the best key tokens more often than random selection. This is also the basis for choosing the center key token as the representative in our algorithm. See Appendix A. for the proof sketch and Appendix A.2 for the formal statement and proof of the theorem.
5
EXPERIMENTS
5.1 EXPERIMENT SETTINGS
Large Language Models (LLMs) are of the most prominent models that utilize the attention mechanism. Thus, we first apply our proposed HiP to Llama3. 1-8B (Touvron et al., 2023), a pre- trained LLM that is reported perform well on various long-context natural language understanding

---

Published as a conference paper at ICLR 2025
Prefill Latency
Decode Latency
PG19 Perplexity
300
80
250
400
70
Crashed
Latency (ms)
200
Latency SU
Perplexity
60
150
40
Crashed
10
200M
100
Frashed
20
50
-00M
8
25
50
75
100
125
50
75
100
125
25
50
100
125
Context length (x1024 tokens)
Context length (x1024 tokens)
Context length (x1024 tokens)
BigBird
StreamingLLM
Hyyee 1yea
Hyper A
Latency and Perplexity Evaluation with Various Context Lengths. We evaluate our proposed HiP and baselines in PG19 (Rae et al., 2019) with various context length on Llama3.1-8B (Meta, 2024). See Appendix D for experiment details.
Figure 7:
Table I:
Passkey Results. We evaluate our proposed Hip and baselines using passkey retrieval which is needle in haystack style context utilization benchmark.
Dense Prefill
Sparse Prefill
Dense
Sparse Decode
Dense Decode
Sparse Decode
936 H
215 alH
31 a!II
2N A
MI AAD
22 Rngga
all
G++ MI-H
XC A
XI GAV
Attention Method
2k
11
128k
100
62.7
64.0
61.6
67.4
80.2
87.2
79.4
82.0
11.1
97.8
8.6
10.5
27.2
27.6
92.9
7.4
9.9
6.2
9.2
19.3
32.7
68.0
66.3
77.3
84.6
90.7
94.0
23.2
20.2
12.7
Context
32k
100
64.5
65.3
82.2
87.1
91.2
98.0
99.7
99.6
16.5
100
26.8
42.
96.9
100
100
14.6
14.8
16.1
38.1
96.3
100
95.3
Length
16k
100
66.2
67.5
90.3
98.0
93.7
98.5
94.7
100
25.3
98.6
43.9
59.5
100
100
100
19.0
13.9
29.8
55.6
98.7
100
91.3
8k
100
69.5
73.5
99.3
95.7
100
100
32.3
61.9
98.4
100
100
30.8
16.9
66.8
95.8
100
100
4k
100
77.3
83.1
100
100
98.5
99.7
100
100
56.8
100
90.4
100
97.8
100
100
56.8
24.7
94.1
100
98.0
100
100
Avg. 100
67.44
70.0
84.8
58.2
90.7
96.9
96.1
95.9
27.5
98.6
41.1
55.2
79.9
84.1
98.8
24.55
15.5
38.3
52.6
77.0
83.5
89.2
Speedup
Prefill
1.0
1.0
1.0
14.6
8.4
8.5
14.6
8:4
8.5
(128k)
Decode
1.0
29.1
14.8
21.8
11.3
29.1
14.5
29.9
15.1
1.0
1.0
1.0
1.0
1.0
1.0
1.0
29.1
45.0
21.8
11.3
29.9
15.1
63.9
tasks up to 128k context tokens, to evaluate the effectiveness of our HiP mechanism. We replace all, but the initial La attention layers with HiP in the pretrained LLM, where L is the total number of layers, and L denotes the remaining dense attention layers. We choose la through an ablation study (Appendix E.5). During LLM decoding, we cache the sparse attention mask from the previous step and refresh it every I'm step to reduce the decoding latency. The latency pertecy--effommacce tradeoff of mm is discussed in Section 5.4. For a detailed description of HiP's decoding process, see Algorithm 2 in the appendix. Further details on the hyperparameters are in Appendix D.
Baselines. We use several sparse attention baselines: A, StreamingLLMM (SLLM) (Xiao et al., 2024), AVD (Jiang et al., 2024; Li et al., 2024), BigBird (Zaheer et al., 2020a), HyperAttention (Han et al., 2024), and H2O (Zhang al., 2023), chosen for their training- free and sub-quadratic properties. Both StreamingLLMM and A use a combination of global sink tokens and sliding window (Beltagy al., 2020), with StreamingI.LM additionally using rolling RoPE indexing (Xiao al., 2024). AVD retains key vertical and diagonal lines in the prefill attention mask based on snapshot scores on top of A. As is a prefill-orientdd method, A is used for decoding. BigBird uses random masking along with the A pattern. MlyperAttention is aakee-clustering-styye (Kitaev et al., 2019) attention mechanism. Finally, H₂O retains the top-k high-scoring KV tokens for the next step's KV cache.
5.2 LANGUAGE MODELING PERFORMANCE EVALUATION
We evaluate HiP on the PG19 (Rae et al., 2019) datasets. We measure latency in two stages: (1) the initial pass (prefill), where the forward pass covers the entire prompt, and (2) subsequent passes (decode), which process one token at a time with a KV cache. In Figure 7, HiP attention is 9.00x faster in prompt latency and 29.99 faster in decoding latency on Llama3.1-8B. with only +0.5348 increase in perplexity on PG19 (8.1151 - 8.6499), Our method leverages block approximation to max- imize MMU efficiency, outperforming quadratic baselines and achieving near-linear decoding la- tency. Further details on experimental settings are in Appendix D.
7

---

Published as a conference paper at ICLR 2025
Table 3: LongBench Results. evaluate HiP and baselines. We measure the speedup of prefill and decode on 32k context length, the maximum context length of LongBench.
Dease Prefll
Sparse Prefill
Sparse Decode
Donse Decnde
Sparse Decade
9S0 O
O
11 puggea
215 Jt
at
CNN
pi
215 anta
31 att
ON 19 HH
A 4
AVD St
AVD S.
BinBind
Alication
Method
mm
215
1
512
215
11
NarativoAA
29.5
15.8
15.3
25.9
26.0
29.0
28.9
20.2
22.0
25.5
23.5
21.4
2399
26.8
27.1
11.4
17.4
18.8
20.6
11.9
19.4
21
25.1
26,9
26.4
Qusper
44.3
19.4
21.6
32.3
35.6
42.1
43.2
28.6
41.4
429
44.3
19,8
45.1
44.2
43.1
7.8
21.4
26.2
26.3
26.9
14,9
41.7
43.6
43.9
43.0
Subset
2WikiMQA
HotpotQA
54.6
39.5
19.3
17.3
20.6
17.8
14.7
33.7
53.9
58.9
54.5
39.5
21.6
34.6
45.6
41.7
53.4
45.0
34.1
51.
45.8
45.1
56.1
53.8
41.2
8.6
22.2
27.1
28.9
39.5
34.5
42.5
25.1
33.0
50.
45.2
440
55.0
53.6
42.4
53.7
39.8
GovRepoet
35.0
26.7
28.2
24.9
27.3
50.3
32.4
33.8
34.2
35.0
34.5
14,0
34.
34.4
34.6
23.2
25.9
25.1
23.1
25.2
27.3
30.1
31
34.9
11.7
MulliNows
27.4
25.1
25.6
26.5
26.8
27.2
27.4
27.1
27.1
27.2
22.6
25.5
22.8
22.5
26.0
26.7
27.9
26.7
Avg. Scores
38.4
2006
21.5
11.3
13.3
56.8
37.6
26.5
34.1
57.7
57.2
12.8
18.01
38.9
57.8
13.9
25.5
26.6
28.1
25.9
20.9
35.8
37.6
38.3
16.9
Rel. Seores (%6)
100
533
56
81
87
96
98
69
89
98
86
99
101
99
36
61
69
68
78
93
98
100
96
Speedup (328)
110
1.0
10
3.2
211
1.7
2.7
32
3.0
1.7
Decode
1.0
7.3
3.7
10.4
5.6
8.5
4.3
1.0
1.0
100
10
1.0
1.0
1.0
1.0
10.6
12.8
12.8
6.4
10.4
5.6
8.5
4.3
4.3
16.5
5.3 LONG CONTEXT PERFORMANCE
In this section, we investigate the performance of our HiP, comparing its latency and accuracy against baselines on various benchmarks. Mainly, we build two kinds of benchmark sets: (1) long-context utilization to verify our method can retrieve the information in given context using needle in a haystack (NIAH) and (2) long-context natural language understanding to show that our method can preserve reasoning and text generation performance of original long-context LLM. We apply the efficient attention method to mimic various deployment settings by replacing prefill, decode, or drefill-dcoode flash attention. We can find our HiP performs robustly in every scenario compared to baselines, by applying efficient attention methods in different phases separately.
Passkey and RULER. First, we analyze the result of long- cong-context utilization performance using passkey retrieval in Tables and 2. Our passkey retrieval test is a simple test to find a five-digit passkey in a repeated haystack sentence. RULER (Hsieh et al., 2024) is a more complex benchmark containing NIAH tests, such as finding multiple passkeys and tracking variable changes inside complicated essay-style haystack sentences. In Table our method is the strongest in every deploy- ment setting. Dense prefill in general scores high in this benchmark because the model has no chance of overlooking the passkey tokens. However, interestingly, AVD shows an almost perfect score with sparse prefill + dense de- code. We think this is because the snapshot heuristic that captures important tokens during prefill is a perfect fit for this benchmark. However, because of this aspect, it performs poorly on more complex tasks such RULER and LongBench. The combination of HiP and AVD slightly increases the performance from regular HiP, achieving 100% accuracy in passkey up 64k context length.
Table 2: RULER Results. We compare the effective context lengths of HiP and baselines with Llama3.1 I-BBB. Accuracies surpassing 80% are marked with bold font.
Dense Prefill
Sparse Prefill
Doacc
Sparse
Dense Decode
Sparse Decode
10 ppppa
AVD
atta
gin x2 d111
49 prggia
HIF
Attention Method
Effective Lengrh
8k
16k
32k
32k
16k
128k
77..0
13.9
38.9
31.3
19.1
52.0
58.2
11:0
8.3
21.1
26.5
79.9
Context Length
32k
87.4
16.9
82.9
58.1
77.0
86.5
89.7
12.5
15.9
27.5
84.2
16k
91.6
27.2
92.4
76.3
89.5
92.1
94.1
24.0
22.6
90.0
93,9
Sk
93.8
54.5
89.5
94.7
94.4
94.6
4k
95.5
55.3
95.9
95.3
95.9
95.9
96.0
87.1
65.1
95.7
95.9
Avg.
65.4
N5.44
72.0
Speodup (128K) Decode
Prefill
11:00
11:00
1.00
4.53
4,801
2.44
1.700
4.53
4,801
2.44
1.70
1.00
5.90
7.05
1.00
1.00
1.00
1.00
5.90
27.01
7,05
7.75
LongBench. We then use the LongBench benchmark (Bai et al., 2023) to evaluate the long context prompt and decoding performance of HiP in Table 3. We believe that this benchmark is the most important because it shows both long context generation performance and knowledge retrieval per- formance, which are critical in many LLM applications. such as multi-turn assistants and in-context learning. Compared to passkey. the dense decode setting scores higher because this benchmark much more decoding-heavy. This means that real-world natural language question answering and long context text generation rely more on decoding accuracy rather than prefill. Therefore, we can see san-decode-frienndly baselines such as StreamingLLM, AVD and A failing to recover long- gegeratenrto performance in GovReport and MultiNews subtasks, which decode 512 tokens. In- terestingly, AVD completely fails on those two subsets while it works moderately well on some QA tasks. We think this is because AVD fails to capture complex reasoning and long-term context due to its restrictive attention mask patterns. In Appendix E.2, we illustrate this long context knowl- edge retrieval ability by using an example from LongBench. HiP outperforms every baseline, and
8

---

Published as a conference paper at ICLR 2025
Table 4: Benchmark Performance on Long-Booksum. Task. We evaluate the book summarization task. 2k tokens are generated for the summary of each book, whose lengths are between 32k-128k tokens. For the 'Hall' and 'Quarter window" settings, the context window size of each sparse attention method is adjusted accordingly. The speedups are measured on the Normal setting.
Decoding context window (tokens)
Iaama,. Instrrrtttttt
Method
KVCache
foorprint (tokens)
Normal Window (x1)
Half Window (x.5)
Quarter Window (x.25)
Decode
ROUGE-1
ROUGE-22
ROUCE-L
ROUGE-1
ROUGE-2 ROUGE-L
ROUGE-1
RODGE-2
kOUGE-L, Speeduu
Un (imial
FlashAttn
41.63%
10.16%
23.58%
41.63%
10.16%
23.58%
41.63%
10.16%
23.58%
1.00x
BigBird 4k
4K
36.05%
7.59%
19.81%
34,20%
7.02%
18.90%
31.97%
6.23%
18.71%
5.90x
Restrieted VRAM Size
FlashArta TRUNC
8K
8K
36.62%
8.33%
20.97%
36.62%
8.33%
20.97%
36.62%
8.33%
20.97%
15.68x
Big 2k TRUNC
SK
4K
35.44%
7.57%
19.65%
34.60%
7.20%
19.32%
32.36%
6.46%
18.62%
5.90x
AVD 8K 8k a 8k
SK
8K
37. 18%
8.28%
21.62%
36.07%
7.72%
20.60%
35.72%
7.67%
20.81%
7.5lx
HiP (Ours)
7K
3K
38.84%
9.11%
21.92%
38.47%
8.57%
21.50%
37.34%
8.52%
21.53%
7.75x
HiP 2 Vak Dik
7K
44
39.98% %
9.61% %
22.61 %
39.44%
9.09
22.05%
38.00% %
8.70%
22.04%
7.05x
Latency Breakdown in Single Layer Decoding
Table End-to-end Decoding Speedup and Quality for each Tm- We show the trade-off between end-to-end decoding speedup and decoding quality using Long- Bench. Metrics are the comparison score Flash At- tention. LongBench, we merge four tasks into 'QA' and merge two tasks into the 'Suummary' column in the table. The latency is measured with Ig=0.
Linear
163.1
303.6
585.3
100
Flach Actention
Latency (us)
Masking
75
Spa Accention
HIP Mask Refreshi
HIn Madk Cached)
Facg Attaintion
25
End-to-Ead Decoding Speedup
LongBench
16
32
64
128
Seq. Length
8k
16k
32k
64k
128k
QA
Surwa
Avg.
Context length (x1024 tokens)
HiP (r==l)
0.99
1.51
2.31
4.05
6.83
95,9
96.1
96.0
Figure& 8: End-to-end Decoding Latency. We show
HiP (rm=2)
1.26
1,93
3.09
5,51
9.57
94.7
95.5
95.0
HiP (rm=4)
1.58
2.44
7.28
12.97
94.2
93.3
93.9
two phases of HiP decoding: mask refreshing step triggered in every I'm decoding step and mask cached sparse attention.
HiP (G==8)
1.65
2.55
4.31
7.89
14.30
93.4
90.5
92.4
with a small amount of fine-tuning with an unrelated dataset, it even recovers the original model's performance ('HiP HEAL"). See Appendix D for more details and discussion about healing.
BookSum. We use the BookSum benchmark (Kryšciński et al., 2022) to assess the long-context and long-response generation capabilities of HiP. We report the average ROUGE Fl-scores (Lin, 2004) for the generated summaries in Table 4. To simulate a realistic long-context decoding scenario and demonstrate the effectiveness of KV cache offloading, we put limit on the GPU KV memory size to 8K tokens. This represents a practical context length on a 24GB GPU with an 8B model without KV offloading. Specifically, for Attentinn and BigBird, we truncate the context to 8K tokens, and AVD uses an 8K token length sliding window. With our method, with KV cache offloading, can expand the effective context length only limited by the main memory's 's capacity, which is much cheaper. HiP outperforms all other baselines in this VRA.M-lmmtted setting while maintaining high decoding speed: over 7x faster than regular Attettnt Although FlashAttention with a truncated context is faster, it suffers from significant performance degradation and, most importantly, breaks the user's expectation that the model can access the entire context. We observe that HiP with a context window of only 512 still outperforms AVD with an 8k window.
5.4 LATENCY BREAKDOWN AND END-TO-END DECODING SPEEDUP
We evaluate the trade-off between attention latency and the model performance with HiP in Fig- ure 7. We observe that our HiP's latency-optimized setting shows about 9.00x speedup of attention decoding latency but only increases the perplexity by 0.5348 in PG19 (Rae et al., 2019), compared to FlashAttention2. In Figure 8. we show the latency breakdown of the HiP-applied transformer model. Our proposed method contains two major components that contribute to the overall latency: (1) top-k estimation iterations and (2) fused sparse attention. We observe that the HiP top-k estima- tion kernel is the only scaling part as the sequence grows; the sparse attention and linear layer shows constant time for each decoding step. Since the top-k estimation iteration results can be cached and reused I'm times, the latency of the HiP method is dominated by fused sparse attention in most practical scenarios, as shown in Figure 8. On the other hand, the rm hyperparameter trades off the generation quality for latency, especially for long decoding, as shown in Table 5. HiP achieves 6.83 times end-to-end decoding speedup with 128k context while maintaining 96.0% relative per- formance in LengBench. We can speed up further to 14.30x when we allow a moderate amount of performance degradation (-3.6%p)

---

Published as a conference paper at ICLR 2025
Decode Throughput and GPU KV Memory
Table 6: Detailed additional data of KV cache offloading performance on RTX 4090 24GB and A100 80GB.
5000
flash wns Offeaa
Decode Throughput (toks) (tol/S)
4000
Rash Adn UVM
IPP win Officed
175
RTX4090 24GB, T=44k
A100 80GB. T-512k
GPU Memory
3000
ISP UVM of Vector Map
HP UVM weo Cache
I4 LIVNSoo on cisch Mep
125
Throughput (tok/s)
Prefill
Decode
VRAM
(MB)
Prefll
Decode
VRAM
(MB)
DOM
FA2
0OM 00M
0OM 00M
00M
0OM
0OM
OOM
2000
HiP offlead
00M
00M
00M
OOM
1000
2 UVM
5678
1.91
N/A
1382
0.18
N/A
HiP UVM
9002
22.91
N/A
8225
20.94
N/A
VecouMMap
6386
95.45
2283
4127
85.11
18404
12
20
20
50
60
HiP HabkMap
359
10.15
2104
48
2.74
10890
Context length Dx1024 tokens)
Figure 9: KV Cache Offloading Performance (left). We measure hatched prefill and decoding throughput (tokens/s) with our novel KV cache offloading framework with an on-device offloading cache. Straight lines show the latency, and dashed lines show the GPU memory usage by KV caches.
5.5 KV CACHE OFFLOADING BENCHMARK
In Figure 9, we evaluate the latency and memory usage of our KV offloading framework. The UVM variants use the CUDA unified virtual memory API to offload the whole KV cache to the main memory. Our HiP has two variants that depend the type of cache implementation. We use Llama3.1 with 16-bit weights, and the KV states are stored in 8-bit floats. We use a single RTX 4090 24GB for the graph on the left, and to additionally test our method up 512k tokens, we also test on a single A100 80GB GPU. We set la = O, and choose the last token for the representative key to reduce the memory access in this test. See Appendix D for details.
As shown in Figure 9, with UVM, both ours and Flash Attention slow down decoding about 5 to 7 times compared to full GPU runtime. However, we could serve until 64k context, while the same machine can serve only 16k at maximum. Since memory access is significantly more costly with UVM, the trend of logarithmic scaling of decode throughput is clearer than when working with pure GPU memory. So, at 64k context length, ours is more than 50 times faster than Flash Attention with UVM. However, UVM slows down both methods too much compared to full GPU runtime.
We test two types of cache implementation: vector map and hash map. A vector map uses a T-sized vector of pointers pointing to the allocated bank to store the mapping between a token index and a bank index. Our GPU-loaded KV offloading cache (Vector Map) shines by achieving 93% decoding throughput compared to no KV offloading at all. Without a significant slowdown, we could extend RTX 4090, which is 4.17x higher decoding throughput compared to HiPUVM and 49.97 49.97x higher decoding throughput compared to Flash Attentionuvv, as shown in Table 6. However, with the vector map, the space complexity is OTT). To reduce the space complexity to Oloo T), we use a linear probing hash map to store the index mapping. This way, we can reduce the GPU memory consumption by 40.8% on 512k context length. However, since the hash map lookup is not friendly to the GPU, it slows down token accesses more than naive UVM.
the serving context from 16k to 64k
We present our KV offloading framework on a standard gaming PC equipped with single RTX 4090. Our experiments confirm that the PCIe 4.0x8 bandwidth is sufficient to manage offloading traffic through KV accesses using UVM. Furthermore, when scaled up to a single A100 80GB, our framework demonstrates its ability to extend serving context length, even server-grade hardware. We anticipate that our HiP's KV offloading framework will effectively increase serviceable context length across a wide range of deployments, from on-device setups cloud-based environments.
6 CONCLUSION
In this study, we present HiP Attention, a novel framework for accelerating pretrained Transformer- based models without any training, with focus on the acceleration of LLMs for long-context tasks. Our proposed HiP rapidly estimates the top-k context keys for computing sparse attention, drasti- cally reducing the computation required for long context inference and fine-tuning from O(T²) to (T T). Our HiP attention a drop-in replacement for the core of any Transformer-basedd model, such as language and multimodal models, and does not require modifying the existing weights This is a practical and meaningful improvement as it allows pre-trained models to be fine-tuned and executed much more efficiently in long sequences without sacrificing quality. We are looking forward to contributing to open-source LLM serving frameworks by combining various efficient decoding strategies with HiP attention.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
We provide every experiment code and kernel code in the attached supplementary file. We also provide detailed instructions on how to run experiments readme markdown files, so please read those files. And we put detailed experiment settings in Appendix D. We will try our best to resolve further reproducibility problems. Inside the HiP library. we have multiple versions of HiP kernlls, all written with OpenAI Triton. The upstream kernel path is hip models attention2 draft prefetch. py. Additionally, you can see the evo- lution of our HiP from the very first HiP implementation hip models hip attention attenti innelioonl.. py; please feel free to enjoy our codebases. We left them all for research purposes when someone needs various settings, such as dynamic retention ratios, that are only supported by old versions. Our main experiment entry file is hip / main / model.eval. py. Please execute --help option gather further information. Our offloading experiment entry file is hip / models / dip.attention / offload runner / offload runner. For Long- bench RULER, we modified the official code to run our method with vLLM. Please refer to Hi PAttennt i cnArgs class to investigate full settings, including every subtle configuration. A, AVD and BigBird are using the same HiP kernel since they are the same block sparse attention. We just modify the block masks that passed to block sparse attention. Streaming is implemented in hip/models / sinkk attentteonnn ion/s sink aankkatte attention py. About HiP-related environment variables of vLLM and SGlang, please refer to HiPAttent ionEnvs in vLLM and SGlang attention backend implementations.
hip attention
ACKNOWLEDGEMENTS
This work was supported by the Institute for Information & communications Technology Planning & Evaluation(IITPP) grant funded by the Korea government (MSIT) (RS-2019-11900775, Artificial In- telligence Graduate School Program (KAIST); No. -0--2019--11999 1906, Artificial Intelligence Grad- uate School Program (POSTECH); No.RS NSS-2024-004597999 Development of ML compiler frame- work for on-device AI; No.RS No.RS-2022-1122 1122-112,0713, 13, Meta-learning Applicable to Real-world Problems), the National Research Foundation of Korea (NRF) grant funded by the Korean government (MSIT) (No. RS-2024-00354944; No. RS-2023-002562599 and NAVER-Inte Co-Lab. The work was conducted by KAIST and reviewed by both NAVER and Intel. Artificial intelligence industrial con- vergence cluster development project funded by the Ministry of Science and ICT (MSIT, Korea) & Gwangju Metropolitan City.
REFERENCES
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023. 2,
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020. URL http: http///awwwoooggab///////oo/
05 50. 2, 7,
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple I1m inference acceleration framework with multiple decoding heads, 2024. 44
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers, 2022. URL http:
//arxiv. orriiv.orrg/bba abw.aaass/00// 11.0099.199.. 14794. 39
Together Computer. Redpajama: an open dataset for training large language models, 2023. URL httm//ttteherccooouuc eeateerre/eeeeeeeee RRePP a jama ammama-ae Data. 38
Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning, 2023.
org. ott.o////aa....0//033 2, 38, 39

---

Published as a conference paper at ICLR 2025
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with I0-awareness, 2022. URL httpttt///xxii..ottt / 3,
Yichao Fu, Peter Bailis, Ion Stoica, and Hao Zhang. Break the sequential dependency of Ilm infer- ence using lookahead decoding. arXiv preprint arXiv:2402.02057, 2024. 44
Google. Our next-generation model: Gemini 1.5, 2024. URL https: /ttttttt//bbloogooooo///
technology/ wwwpa//gggoooooooooogeoo weweimii-eee---eettttiiooo wew.:mmaoeee-rrrrrr------00-
46
Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David Woodruff, and Amir Zandich. Hyperattention: Long-context attention in near-linear time. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: /oop/rppeeeveeeeeeeeeeeeeeeeeeee ew. net. forumoormm iid= EhCOd2BJIM, 7, 38, 39, 40
Dan Hendrycks, Collin Burn,, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding, 2021. TRLLhtt::///xxxt org/ a/a/bb//// ///aabb//009 2009.000.... 03300. 33
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint 7Xiv:2404.006544 2024. 3,
Shengran Hu, Cong Lu, and Jeff Clune. Automated design of agentic systems. arXiv preprint MXXi::2008,08433 2024. 3
Hakan Inan, Kartikeya Upasani, Jianfeng Chi, Rashi Rungta, Krithika Iyer, Yuning Mao, Michael Tontchev, Qing Hu, Brian Fuller, Davide Testuggine, and Madian Khabsa. Llama guard: Llm- based input-output safeguard for human-ai conversations, 2023. 46
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhen- hua Han, Amir Abdi, Dongsheng Li, Chin-Yer Lin, Yuqing Yang, and Lili Qiu. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. arXiv preprint arXiv:2407.0249%, 2024.7
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia- Yuan Chang, Huiyuan Chen, and Xia Hu. Llm maybe longlm: Self-extend Ilm context window without tuning, 2024. 34
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In ICLR 2020, 2019. URL https: http///ppee/eeeeeeeeeeeeeeeeeeee net met/fooummma toet/ 2, 7, 33, 34, 39, 40
Wojeiech Kryšcinski, Nazneen Rajani, Divyansh Agarwal, Caiming Xiong, and Dragomir Radev. BookSum: A collection of datasets for long-form narrative summarization, 2022. URL http: //arxiv.orgg abs 2105. 08209. 3,
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium Operating Systems Principles, 2023.
Heejun Lee, Jina Kim, Jeffrey Willette, and Sung Ju Hwang. SEA: Sparse linear attention with estimated attention mask, 2023. URL http:://aaat org/ abs /abs/233 /21/./21... 01777. 2, 3, 33, 34, 39, 40,44 44
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference Machine Learning, pp. 19274-19286. PMLR, 2023. 44
Bo Li*, Peiyuan Zhang* Kaichen Zhang*, Fanyi Pu*, Xinrun Du Yuhao Dong, Haotian Liu, Yuanhan Zhang, Ge Zhang, Chunyuan Li, and Ziwei Liu. Lmms-eval: Accelerating the development of large multimoal models, March 2024. URL attps://githut com/ Evol Evolvingt.n ving Lab/ Imm=-vall
12

---

Published as a conference paper at ICLR 2025
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before generation. arXiv preprint arXiv:2404.14469 2024.7
Chin-Yew Lin. ROUGE: package for automatic evaluation summaries, In Text Summarization Branches Out, pp. 74-81, Barcelona, Spain, July 2004. Association for Computational Linguis- tics. URL https: / /tw.aaaaaaaooooooo..oooo.ooooola wrg. 20t://04-100000 -0013. 9
Haotian Liu, Chunyuan Li, Yuheng Li, and Yong Jae Lee. Improved baselines with visual instruction tuning, 2023a. URL /ttt:://axiv.... wtt. abs//bb///// 2231 0. 03744.
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. In NeurlPS, 2023b. 33
Haotian Liu, Chunyuan Li, Yuheng Li, Bo Li, Yuanhan Zhang, Sheng Shen, and Yong Jae Lee. Llava-next: Improved reasoning, oer, and world knowledge, January 2024. URL https:// seeva-vi.gitann
io -0001/0000//00/0000000000-111-11
ilava-next/ 33
Liu Liu, Zheng Qu, Zhaodong Chen, Yufei Ding, and Yuan Xie. Transformer acceleration with dynamic sparse attention, 2021, URL http:://aaxxxww.o//////////
2,
Llama Team Al @ Meta. The llama 3 herd of models, 2024. 2407. 2407.21783 21783.1111 7,
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, Chunan Shi, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating large language model serv- ing with tree-based speculative inference and verification. In Proceedings of the 29th ACM In- ternational Conference on Architectural Support for Programning Languages and Operating Systems, Volume 3, ASPLOS '24. ACM, April 2024. doi: 00.1145/320066.3651333 URL
http:://xx.. www/orgg///0 10.
3651335
Nvidia. Tensorcore: Nvidia, 2024. data-center/teesoo- cores/..
URL. https: http://wwwwwamm www.i.iiia...mmm. www.iiiaa.com/euuus.....
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models, 2023. URL httt::::/aaxxi.orrtte htt.b////33.000 00071. 200900001..
34
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosFormer: Rethinking softmax attention, 2022. URL http:// arxiv. orw. abs 220/2222222222222 08791. 3,
Jack W Rac, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL
wwtp:///....aa///////oooo//s///s 9ttp://www...ar/////9111.1...000 7,9
Luka Ribar, Ivan Chelombiev, Luke Hudlass- Galley, Charlie Blake, Carlo Luschi, and Douglas Om: SparQ attention: Bandwidth-efffiient LLM inference, 2013. URL http: /ttp:://arxiii org/ / 2812.04988 32, 43, 44
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Ev- timov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. Code llama: Open foundation models for code, 2024. URL hww.:htt///. rxiv. orgg abs 2308 ..e...2000 12950.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da- Da-Cheng Juan. Sparse sinkhorn attention. In Proceedings of the 37th International Conference on Machine Learning, pp. 9438-9444 94449447 PMLR, 2020. URL https:/ httplp:://ppcceeeee http///pcceeenn...oee mlr, www.sppwwe/eeaaeee html. ISSN: 2640- 3498. 2
13

---

Published as a conference paper at ICLR 2025
Yi Tay, Dara Bahri, Donald Metzler, Da- Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Re- thinking self-attention in transformer models, 2021. URL /tL.ttp:://axxii..o/////. 00743. 2
Gemma Team. Gemma 2: Improving open language models practical size, 2024. URL https:
//arxiv. wtwpiv..rggabb////0000 00118. 34
Philippe Tillet, Hsiang-Tsung Kung, and David D. Cox. Triton: an intermediate language and compiler for tiled neural network computations. Proceedings of the 3rd ACM SIGPLAN Inter- national Workshop on Machine Learning and Programming Languages, 2019. URL https: //api. www.sseaaannnnnaeeeeeeeeeeeeeeee
tinygrad. tinybox, 2024. URL
46
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao. Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nic, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. URL
220..22372222 09288. 1,6
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 Rep- resentations, 2024. URL https:/ hww.///eeeeeeeeeeeeeeeeeeeeeeeee / net/ forum? 1== 7, 32, 39
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, and Amr Ahmed. Big bird: Trans- formers for longer sequences. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 17283-17297. Curran Associates, Inc., 2020a. URL httppss/pprooeeddss...eeeiippppp http://sseeeerpp..pcccppppeeeeee papeer2000 file//ill/// http:///h2.222222222288853aaaaaa
3,7
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Al- berti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Advances in Neu- ral Information Prucessing Systems, volume 33, pp. 17283-17297. Curran Associates, neurips......................... /ccppe//e cttps/cr.ppppeeepeeeeeeee paper/2020/ html.
Inc., 2020b. URL
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. Hąo: Heavy- hitter oracle for efficient generative inference of large language models, 2023. 7, 44
Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E. Gonzalez, Clark Barrett, and Ying Sheng. Sglang: Efficient execution of structured language model programs, 2024. URL //aa//////x.... org/abs// 2212.. 07104. 2
14

---

Published as a conference paper at ICLR 2025
Appendices
A Theoretical Analysis
16
A.1 Proof Sketch
16
A.2 Detailed Proofs
18
A.3 Revisiting Assumptions in Theoretical Analysis
23
B Detailed Methodology Descriptions
31
B.1
Hierarchical Sparse Attention Mask Estimation Algorithm
31
B.2 HiP Decoding Algorithm
B.3 Detailed Flow-diagram of HiP
32
B.4
Additional Optimization Techniques
32
B.5
Training Downstream Tasks with HiP
32
Additional Experimental Results
33
C.1
Large Multimodal Model with HiP
33
C.2 Massive Multitask Language Understanding (MMLU)
33
C33 Comparison with Reformer and SEA
C.4 Context Extention with Self-Extend
34
C.5 Ensemble Hierarchical Top-k Approximation
35
D Detailed Experimental Settings
37
Additional Analysis
39
E.1 More Discussion on Related Works
E,2 Analysis of Summarizing Result between StreamingL.LMM and HiP
40
E.3
Hierarchical Attention Mask Pruning Visualization
41
E,4
Ablation Study on Block Size
E.5
Ablation Study on Dense Layer Choice
43
E.6
Ablation Study Representative Token Location
43
E.7
Discussion about KV Cache Eviction and Compression Strategy
44
Discussion about Speculative Decoding
E.9
Remaining Challenges in HiP and Potential Solutions
44
E.10
Unique GPU Resource Demand Pattern of HiP Compared to Flash Attention
45
Potential Negative Social Impact
46
15

---

Published as a conference paper at ICLR 2025
A THEORETICAL ANALYSIS
A.1 PROOF SKETCH
This section provides sketches of the proofs and derivations for the theorems and lemmas mentioned in the main section of the paper. The full formal proofs are presented in the next subsection.
Using the insight obtained from our observations, we model the random variable وة as follows.
Axiom The atfention score difference between two tokens distance A apart is defined as a random variable d with the following distribution.
8A AOO0 (A))), A>0.
Where the standard deviation (A) is increasing function of A, note that both dΔ and (A) defined only when A > O, and cases when A = will be dealt separately. We also need to define the problem setting precisely in terms of math. Without losing generality, we model the tokens as list of indices of length 2n. For simplicity, we simplify the problem into the case of top-1 selection. The indices of tokens range from 1 to 2n, where indices I belong to the first section, and indices + 1 to 2n belong to the second section. The number k represents the location of the representative token inside each section, and hence, indices è and n + k are the representative tokens of the first and second sections, respectively. We visualize each token position in Figure 10.
Indices
Rep
Representative
Figure 10: Visualization of Problem Setting.
For the ensuing theorems, define the following random variables and events: t is the random variable regarding the index of the maximum token, and is assumed to be a uniform random variable. A and B are the events where the maximum token is located in the first and second section, respectively. X is the event where the first representative token is selected, due to the attention score index being larger than the score at index n + k. The event Y is the exact opposite of X, where the second representative token is selected.
What the probability of a token whose distance from the maximum token a, having a greater attention score than token distance 3 away from the maximum token? Using the random variable da from earlier, since the tokens cannot have larger scores than the maximum token, the scores can be calculated as Smaz - and Smaz jōg(, where Smaz denotes the maximum attention score. Thus, this problem equivalent to calculating the probability P( < (5gl), which can be calculated follows.
Lemma 1.
P( (ō =
( +
(00)) (ax)
=2 dr 1.
detailed derivation of the above lemma can be found in Appendix A.2. For convenience, in the proofs to follow, we will abbreviate P( < =5gl) as (a))))(())) Please note that (a)(a) (B)) is a decreasing function of a, and an increasing function of B. Intuitively, this means that the probability of the token with distance being larger than one with distance 3 gets larger as the first token gets closer and the second token gets further away from the maximum token. This intuition agrees with the assumption of locality.
Using Lemma 1, we can calculate the probabilities for the following four cases - the case of the first or second representative token is selected, either when the maximum token lies in the first or section. The detailed derivation processes are provided in Appendix A.2.
Claim 1.
k-1 Σ do 1), o( i)) +1+ Σ
P(X 2n
(o(ii-),,((n -1))
16

---

Published as a conference paper at ICLR 2025
Claim 2.
i), (k))--))))))))))) 11))
P(YNA)= 2n
Σ Ma(n+k-1),, ()) +0+
Σ
-
Claim 3.
k-1
P( 2n
Σ i ()) +0+ Σ
i k), - k))
i=
Claim 4.
A-
P(YNB) = 2n Σ
vlo(k - oi +1+
+1+ Σ - on R))
1=2+1
From these claims, we now appraise HiP's ability to make the correct choice. In particular, if HiP indeed better than random token selection, then it should select the first representative token with higher probability if the maximum token is in the first section, and vice versa. The following theorems specify the conditions in which HiP performs better than random selection. Similarly, more detailed proof can be found in the appendix section.
Lemma 2. Ifk n/2, then P A) > P(YNA).
Proof (sketch). If :: 2 n/2, then the following three inequalities hold.
c0a(k k --)) bll(+ ok 1)), 1>
Σ - M( 1=l+1
Σ
k o), k))-
i=
Therefore, ifk n/2, (XXNA) P(YNA).
Lemma 3. 1fk S n/2 1, then P n): P(XnB).
Proof (sketch). If < n/2 + 1, then the three inequalities hold.
(aa( -(o(k-1),o(n (n Σ
i)), 1>0, 0,
1=1
Σ i=++1
- k), k(
Σ
i k)).
i=
Therefore, if < n/2 1, P(YBB) > P(XnB).
The final theorem follows directly from the above two lemmas.
Theorem 1. If = n2/2, then P A) > P(Y AA) and P(YNB) > P(ΣпB).
This theorem tell us that if k n/2,, then HiP consistently outperforms random token selection in terms of making the correct choice. Thus, just by setting the representative token as the middle token, the hierarchical structure of HiP consistently outperforms random token selection! We now have the answers for the questions mentioned at the beginning of section Section 4. Through mathematical analysis, we have shown that by using the middle token as the representative token, the hicrarchical approach of HiP consistently guarantees better performance than random token selection. Therefore, throughout our paper, we always use the middle token as the representative token of a section.
17

---

Published as a conference paper at ICLR 2025
A.2 DETAILED PROOFS
This section provides more detailed proofs and derivations for the theorems and lemmas mentioned in the main section of the paper.
Lemma 1.
P( 160l) =4 =aʃo Φb
(o(8)
1
[((a))) 2)
$ dr 1.
2n
Proof.
P(-10gl 8a < (0 | /
P((6a 18gl)
P x, y) drdy
-lvl
Assuming that 6,2 and 53 are independent, we can expand the equation as follows.
P(óa = =r,6z = = P(á, = x: P( = y)
:. P(loa < =60l) f / |all
P(o. x))(oa y) ydddy
Myl
From Axiom 1, note that 8n ~ NN(0 (())) and 03 ~ ~N(0,
P = P66 = y) drdy
lal
M//
1
20( dady
lul
(a)
2 (3)²
20(3)
2 dady
2TT (B)
1y v (a)
1
(ul/o(a)
1
2,(
$ drdy
(3)²
2TT
-ul/a(a)
1
2
2Φ
20(01 dy
a(a)
(3)²
Do
[y|
2
+
20( dy
(B)2
Note that the symbol & represents the cumulative distribution function (CDF) of the standard normal distribution.
lys
2
2013y dy
(aa)
(B)
1
4
+
20(0) dy
((a)
(3)
(
o
1
e
dy
(
V2TT
(
1
+b
2
dr
1
a (
2
As mentioned in the main section, for convenience, we will denote P( < (5gl) as ((a(a) (B)). Please remember that (()(a),, (B)) is decreasing function of o, the first argument, and an increas- ing function of 3, the second argument. Note that in our derivation, we assumed that ó, and 6y are independent. More discussion this assumption will be provided in the ensuing subsections.
18

---

Published as a conference paper at ICLR 2025
Claim 1.
P(XNA)
Σ look -000-1,,000 (x= +1+ Σ v(ali - ( (-+1
om
2n
Proof. Recall that uniform random variable denoting the index of maximum token. Since
the number of tokens is 2n, event A is actually the union of events t = 1.... Therefore,
P(KNA) = Σ P(xnt=i- Σ P(=
1 = 2n Σ P i)
1=1
Using Lemma 1, P(X i) can be calculated follows.
P(X = 1) = (b(okk 1), -))
if < k
if = k
(lloi - k), on - i))
if k
A detailed explanation for the above derivation is as follows. When the maximum token index is smaller than the first representative token (i.e. < k), the distances between the maximum token Thus, using Lemma P(X - i)
and the representative tokens are k and (brkk - i). (n + k - i)). Similarly, P(X = (f - k), a [n+ ( token index is larger than the first representative token (i.e. k). When the maximum token the first representative token (i.e. = k), we cannot use Lemma 1 as a ()) is defined only when A >> 0. However, if the maximum token the first representative token, then it will always be larger than
i)) when the maximum
the second representative token hence, P(X|t i) =
Using the above derivation, we can rewrite the original summation as follows.
P(ΣN))
1 t
Σ
2n
Σ (o(k - () +1+
Σ
vllli o(
2n
i=k+1
Claim 2.
P(YNA)
- 2n Σ slo(n + + - okk +0+ Σ k -- oi -k)) (=+11
Proof. The derivation is almost the same as Claim 1, except that the second representative token needs to have larger attention score than the first one. Therefore, the arguments in the 1 function are switched for cases < and i k. When k, the maximum token is the first representative token, P P(X|t i) is zero. Therefore,
o) -)) if < k if = k if >i> k
P(Y -
0
b i), -))
P(YNA) = 2n
P(Y-=1)
2n Σ
(a(n + k- i) i)) + +0+
Σ
+ k))
19

---

Published as a conference paper at ICLR 2025
Claim 3.
P(XBB- = 2n Σ -i) 1)) + +0+
Σ s(a(n - ki, k)) 1
Proof. Recall that is uniform random variable denoting the index of the maximum token. Since the number of tokens 2n, event B8 is actually the union events 1... 2n. Therefore,
2n Σ
2n
2n
P(X PX
= Σ
P(t i)
Σ
P = i)
2n
X
11+t1
im++1
Similarly Claim using Lemma 1, P( - i) derived as follows.
Vill - k), + k - i))
if < n + k
PX 0
if + k
(b(oli k), - n k))
if > n + k
2n
P(NBB= 2n i=++1
Σ P(
n+tk-1
Σ o) +0+
Σ
(i,(i-k),o0( -A))
2n
V=++1
i=++k+1
1
Σ va(nn +i kk +
( +i -i-nni-AAi) -i)) k))
2n
k=
Claim 4.
P(YΠ = 2n voak 1) +i k)) + +++
Σ slali - on -k)) ==
Proof. The derivation is similar to Claim 3, except that the second representative token needs to have larger attention score than the first one. Therefore, the parameters of the v function switched compared to Claim 3. Also, when i + k then the second representative token is the maximum token. Therefore, P(Y i) is derived follows.
( + i), oli k)) 1 - n n (vlllinn
if < n + k if = + k
P(
-n k))
if > n + k
2n
P(YBB) = 1
Σ
P(Y
2n
$=n+1
n
Σ
( - --ik
k)) +1+
Σ
vllli n -o) k))
2n
im++k+1
1
Σ
-( +1+ Σ
vllli - i
2n
Lemma 2. 1f n/2, then P(XNA) P(YNA)
20

---

Published as a conference paper at ICLR 2025
Proof. From Claim and Claim 2,
P(XNA)= 2n
wbo (1 + +1+
Σ
sla(i (o(ii-)),0 (( 1))
i=++1
P)YNA) 1
Σ vlo(n + i) o)) 0+
pla(n -0, 1), k))
k
The following inequalities trivially hold.
R-1
k-1
Σ
(a(( - k 1))> > Σ
(a(n + k- 1), o), - 1))
1>0
For the remaining two terms, the direction of the inequality depends on the relationship between oii k) and in)) In order for P(XNA) > P(YNA), we need to have aii ((i-k) k) S an+ i),
iee. i for all = k ..
i-k n
Vi + 1...
2k
2i
Vi k 1...n
n
Vi = k + 1...n
2
.k 2
Therefore, i 2 n//, then the inequality P( P( A) holds.
Lemma 3. Ifk n/2 + 1, then P(YnB) P(X
Proof. From Claim 3 and Claim 4,
P(XBB= 1
Σ
(o((n wo(n+i-k),odt 1)) +0+ Σ
k))
2n
i=#+1
P(YNB) 2n
Σ - ()) k)) +1+ Σ
- k) k:))
i=
The following incqualities trivially hold.
Mali (n
Σ
k))
For the remaining two terms, the direction of the inequality depends the relationship between
(n+ii-)) k) and (k-i). In order for P(YnB) > PXXBB, we need to have (kk-i) i) S o(n+i-k). l for all 1...k 1.
i.e. h:
A A-isnti- A
Vi = ....
2k 5 n,
Vi = .... 1...k
Si n
Vi = 1...
2
k
+1
2
Therefore, if < n/2++ + 1, then the inequality P(Y P n B) holds.
Theorem 2. (k)o()))( < (k) then P(K A) P V B) is maximized when n1/2
21

---

Published as a conference paper at ICLR 2025
Proof. From Claim and Claim 4,
P(XNA)+ + P(YNB)
Σ
(( +1+
Σ
2n
k=++1
1
Σ pb((k - (n) +1+
Σ
Vblli - on
-
2n
1=++1
Σ vlo(k - ()) +1+
Σ
2n
i=k+1
1 2n
Ma(k - (() +1+
Σ
valli- (( k), i))
i=k+1
For simplicity, we refer to the first term as 71, and the second term as 72, We now investigate which value of maximizes 71 and 72.
For T2, suppose the value of i changes from l to le + 1.
T2 72x = (i((()) k)) - -o(onnnnnnnnnnn ( k),
can be easily seen that (k(()), - k²) is a decreasing function of k and (o k),, a (k)) an increasing function of k. Therefore, T2 722 is decreasing function of k, and its value goes from a positive value when k 1, and negative value when = n 1. Thus, T2 is maximized when T2k+1 T2, (b((k)) o(n k)) (o k), (k)) 0. The value of k where ((((k)) k)) ( - ok)) = 0 is computed as follows.
3: = n k, . n
2
Therefore, 72 is maximized when k = / For Tl, suppose the value of k changes from &: to k 1.
T4 Th (ka(())) on + k)) v k), k))
Unlike 72, we cannot easily determine whether (k))), o + k)) or vo - ()) o k)) is an increasing or decreasing function of k. Therefore, we turn to the definition (1(((a)) (B)). From Lemma I,
(a(a), -4 / + (o(3) (((a)
e de
V2T
Thus, the positivity of T1 - 71, depends on the characteristics of the function
If
a +n/ (k)) a decreasing function of k. then (kk) ( (n++k) k)) becomes a decreasing function of k, and inversely, (a(nn- - k), o2n k)) becomes an increasing function of k. In this case, similarly to T2, we can show that 71 maximized when = n/2. In order to find the condition where (n + (k) becomes decreasing function of k, we take its derivative in terms of k.
on +k) d'(n + (k) (n + k))
< 0
dk
o(k)
okk)
o' (n d' k)
o'kk)
: (n + k)o(k) < on
kk), ())
(n + k)
o(k)
Thus, that if o (k))/a()) is decreasing function of k then on + (k) also becomes
decreasing function of k. This condition is equivalent to the condition o (k)))((k) < d(k)². The derivation is follows.
d o')) dk o(k)
o" (k()) - d'kk)2
< 0,
<
a(k)²
22

---

Published as a conference paper at ICLR 2025
Wikitext2
Booksum
Longbench
Pg19
0.5
0.4
0.2
0.0
Attention Head
(a) Llama3.1 8B Model.
Wikitext2
Booksum
Longbench
Pg19
0.4
0.3
0.1
25-
232222222 Attention Head
0.0
(b) Qwen2 Model.
Wikitext2
Booksum
Longbench
Pg19
0.5
0.3
0.2
0.1
0.0
Attention Head
(c) Exaone 3.0 7.8B Model.
Logarithmic Function. The above figures show the NRMSE (Normalized Root Mean Squared Error) result fitting the o function as a logarithmic function, as a continuation of Figure 12. For fitting the function, the general formula y log + b) +e was used. We demonstrate the experimental results on various datasets and LLM architectures.
Figure 11: NRMSE of Fitting a (A)
Lk"(())(( < (1))
To summarize, if (k)) < dk))?, then (k)) on + k)) becomes a decreasing function of k and (n (2n-- becomes an increasing function of Therefore, 71 T1, becomes a decreasing function of k, and its value goes from positive to negative as & goes from to n 1. The crossover point where k = n k, and n + 2n k. Coincidentally, this is when = n/2. Therefore, if o" (k) < of ok)))² the value of & which maximizes 7'1 is k n/2.
To summarize, (k)))((k) < db))2, then = n/2 maximizes both and 72. This concludes
the proof.
A.3 REVISITING ASSUMPTIONS THEORETICAL ANALYSIS
In Axiom and Lemma 1, we make two important assumptions: that S follows a normal distribu- tion (0, o (A)²). and that 6, and d3 are independent of each other. However, this justifiable? In this section, we provide analysis and explanation on this issue by providing several empirical results collected across various datasets and LLM architectures, which justify our assumptions.
A.3.1 THE DISTRIBUTION OF AA
First, we justify the assumption d (O, (A)²) by validating three parts that de does follow normal distribution, that (AA is an increasing function of A, and that the mean can be approximated zero.
23

---

Published as a conference paper at ICLR 2025
Head B
Head 16
..00
Standard deviadie
Standard devaation
2.52
2-23
Standard Sanddrdddevvato deviation
4.0
2:75
223
1
2:93
2.0
223
0.07445
200
NIRAEE 022-4922
NAMSE:
400
'eedd 20
Head 24
Head 2g
Head 32
Standard deviation
Standard deviation
Standard deviation
Standard deviation
4.0
8.00
2.2
-.5
4.75
2.0
2-
1.8
244
425
2.0
NRNSE: 5,040E12
NRMSE 55004235
3.0
NIISE: 0.253045
NRMSE 0.078300
difference
Index difference
Index difference
Index difference
(a) Llama3.1 8B, 16th layer.
Head
Head
Head 12
Head
M
Standard devistion
2.2
3.0
Standard devistion
5:25
Standard deviation
2.5
2.6
3.02
Standard
30
2.4
1275
2.
20
125
18
200
400
300
500
1000
200
420
550
800
1300
230
400
500
800
1000
220
400
600
800
1000
Index
Head
edifforonre
Head 28
Tndlax difference
Standard deviation
Standard dessation
Standard deviation
325
8:c0
1.4
275
2:30
2.2
22
1.8
200
230
400
10
238
1030
201
410
500
303
1000
200
400
600
BOD
1006
Index difference
Index difference
Index difference
Qwen2 7B, 12th Layer.
Head
Head 12
Head 16
Standard deviation
devistiion
Standurd deviation
des
Standard deviation
2.0
2.50
2.25
2.5
2.08
2
20
NRMAE O 0366647
1.58
NRMSE: O 264154
1.6
NRMSE: O 036833
NRMISE: 0.097188
Head 20
'Hedd 24
Head 28
Heldd 32
Standard deviation
Standard devation
Standard devatiio
1.0
Standard rd deviation
2.5
2.6
2.8
50
2.4
12
2.4
20
NRSASE: 0.00070
NRMSE: D. D44811
NRMSE: NHNSEE-0220182 050182
NRMSST-.0516677 difference
Index differifee
difference
difference
(c) Exaone 3 7.8B, 16th layer.
Figure 12: Plots of (A) and their Normalized Root Mean Squared Error The above figures show the plots of or (A), and the NRMSE emror for fitting as a logarithmic function across various L.LM models and layers. For the input, we used the Wikitext2 dataset.
Figure 15 shows some examples of the distribution of dة of various models. As can seen in the figure, 8a clearly follows normal distribution. In order to provide statistical evidence that d indeed follows a normal distribution, we compute the KL divergence between òΛ and the normal distribution fitted onto EA- We experiment on all layers for various LLM models and A values, and show some of the results in Figure 16. As can be seen in Figure 16, almost all of the KL divergence values extremely small and close to In fact, with the small exception of the first layer
24

---

Published as a conference paper at ICLR 2025
Head
Head 5
Head 9
Head 13
0.00
10.3
0.15
6.2
Moan
0.22
Mcan
6.13
Mean
0.4
0.00
603
0.
0.22
0.02
-1.
0.33
200
830
600
300
1008
200
400
600
503
1000
230
400
520
300
1003
200
430
500
200
1900
Ind difference
Index difference
Index difference
Index difference
Head
0.0
6.0
0.15
0.15
0.10
Mean
0.15
Mean
0.05
0.05
difference
Index difference
Index difference
Index difference
(a) Llama3.1 8B Model, 17th Layer.
Head
Head
Head
Head
0.825
0.35
0.800
Mean
0.30
-0.25
10000
Index difference
Index difference
Index difference
difference
Head 20
Head 24
Head 28
0.1
6.20
0.0
Mean
0.15
Mean
-2.1
-0.1
005
6.00
272
4.30
500
500
1930
0
237
400
600
300
1000
200
.00
600
830
1006
Index difference
Index difference
Index difference
(b) Qwen2 7B Model, 14th Layer.
Head
Head
0.00
-0.05
Moan
-0.10
-2.15
2.10
-4 691
-0.15
-3575
-030
-0:2
3.00
-025
-025
600
1000
200
1000
230
600
1983
203
KB8
1930
Index difference
Index differance
Indix difference
Index difference
Head 20
Head
Head 28
Head
0.23
Mean
-0.1
Mean
Mean
0.15
0.13
-02
03
-84
200
600
1000
see
200
660
1888
Index difference
Index difference
Index difference
Index difference
(c) Exaone 3.0 7.8B Model, 16th Layer.
Figure 13: Plots of the mean of da. The above figures show the plots of the mean of dλ, function of A for various LLM models and layers. For the input, we used the LongBench dataset.
of the Qwen2 model, all of the KL divergence values are smaller than 0.1, which validates that d follows the normal distribution.
Although the KL divergence values in Figure 16 are very small, it appears that few attention heads stand out extreme outliers in the first layer of the Qwen2 model. In Figure 14, we take closer look at the distribution of 8^ in one of such attention heads. This figure shows that the distribution of da in these outliers still largely resemble Gaussian Distribution. However, the distribution appears
25

---

Published as a conference paper at ICLR 2025
A 400
A
0.10
0.10
0.58
0.08
0.00
0.08
0.06
0.24
0.04
0.04
0.02
0.02
0.00
0.00
-20 -10
-20
-10
20
10 20
Difference
Attention Score Difference
Attention Score Difference
Attention Score Difference
Figure 14: Close Examination of óة for the Qwen2 Model. The above figures plot the distributions of da of the 3rd attention head of the first layer for the Qwen2 model, which was also shown in Figure 16. Close examination on these figures suggest that although the KI. divergence scores are high, da still largely follows Gaussian distribution.
200
400
0.225
0.025
0.525
0.520
0.020
0.015
0.015
0.015
0.010
6010
0.010
0.005
0.065
6.055
0.005
0.000
KL divergence: 0.0030
0.000
divergence: 0.0079
0.000
divergence: 0.0031
0.000
10. divergence: 0,0027
Attention Score Difference
Attention Score Difference
Attention Score Difference
Attention Score Difference
(a) 1lama3.1 8B, 9th attention head of the first layer.
A 20U
4 400
A 600
A=
0025
0.125
0.625
0.22
0.020
0015
0.015
0010
0,010
0025
0.305
0,055
0.005
0032
10. divergencer 5,0073
0.00
KL. divergenon 0,0024
0.000
divergence: 0.0023
0.000
10.
divergence: 0,0021
10
-0
Attention Score Difference
Attention Score Difference
Attention Score Difference
Attention Score Differance
(b) Qwen2 7B, 6th attention head of the 9th layer.
400
0:055
0.025
5.02
0.025
8.200
0.000
0.000
0.22
0.155
5.0%5
0.215
0.100
0.10
8005
0.005
2.025
6.000
divergence:
0.000
divergence:
0,000
divergence: 0.0019
0.000
divergenccc
-5
Attention Difference
Attention Difference
Attention Difference
Attention Score Difference
(c) Exaone 3.0 7.8B, 6th head of the 13th layer.
Figure 15: Distribution of d.. The above figures show the distributions of d across various models for A values 200, 400, 600, and 800, respectively. The Wikitext2 dataset was used as input.
be severely discontinuous and quantized, which is the reason why the KL divergence values were so high. We suspect that this is the result of some internal operations of the Qwen2 model, which we are not aware Although these few occasions stand out as anomalies, we believe that it does not compromise the integrity of our assumption for two reasons. First, it is an extremely that only occurs in the few early layers of the Qwen2 model. Secondly, even though the results are discrete, the distribution of da still resembles normal distribution. Therefore, our assumption of modeling 8A as a normal distribution is valid.
occasion
Next, show that (A) is increasing function of A. We observe an interesting phenomenon, where the overall demeanor with which (A) increases resembles a logarithmic function, as can be seen in Figure 12. Although there may be some amount of minor deviations, the standard deviation
26

---

Published as a conference paper at ICLR 2025
Layer
Layer
Layer
Layer 13
0.10
0.006
17
Layer
29
0.04
0.02
0.00
1 Index Difference
(a) Llama3.1
8B Model.
Layer 1
Layer
Layer
Layer 9
0.10
0.06
Layer 13
Layer 17
Layer 21
Layer 25
00.44
0,00
Index Difference
(b) Qwen2 7B Model.
Layer
Layer 5
Layer
Layer er
0.10
0.08
Layer 17
Layer 21
Layer 25
Layer 29
0.04
Head
Index Difference
(c) Exaone 3.0 Model.
Figure 16: KL Divergence. The above figures express the KL Divergence between da and its fitted normal distributions various layers. The y axis represents attention heads, and the x axis represents A, ranging from 50 950. We the Wikitext2 dataset input for the model.
(A) of da has consistently shown to resemble a logarithmic function for almost all layers and attention heads. In order to prove this claim, show the results for fitting oAA) as logarithmic function for each attention head and layer. As can be seen in Figure 11, except for a few outliers (es- pecially in the early layers), the results converge within NRMSE (Normalized Root Mean Squared
27

---

Published as a conference paper at ICLR 2025
Error) of 0.1, which suggests very good fit. This shows that (AA) resembles logarithmic function and, therefore, can be seen as an increasing function.
The above observation actually opens up the room for discussion about the optimality of HiP. In Theorem 2, (X n A) + P(Y n B3 is the probability of HiP making the right choice, regardless of the location of the maximum token. Therefore, Theorem 2 proves that if o" (a)or( < o' dA)) then not only does HiP perform better than random token selection, but it is also in fact optimal. If (A) resembles a logarithmic function like our observation, then o"(A)
O, and therefore the inequality o" (A) (A) < 0 <= o'( holds. Thus, on top of HiP guaranteeing better performance than random token selection, there also is the possibility of HiP being, in fact, optimal.
However, we are very careful with making this claim because we do not yet have logical explana- tion about why a (A) displays logarithmic increase. In the main paper, based on the assumption of locality, we only claim that (A) is an increasing function of A. However, the assumption of locality does not cover the exact detailed behavior of (A)). There is always the possibility that the logarithmic pattern we see in a (A) may be the result of some model-specific traits, the overall training corpus, or some other reasons that we are not aware of. Therefore, although we find this observation interesting, we only leave it as a discussion topic in the appendix section. We plan to analyze the detailed characteristies of the a (A) function in future work.
Finally, we discuss approximating the mean to zero. Unlike our assumption, the mean of which normal distribution 5n follows does not stay as zero. As can be seen in the following Figure 13, the mean of أs consistently begins from zero when A O, but it moves away from The manner in which the mean gets further away from zero is not consistent. In some cases, the mean decreases negatively away from zero, and in some cases, is the opposite. The exact behavior differs between different layers, or attention heads or the input dataset, so it seems that we cannot deterministically express the mean in terms of math.
A increases.
However, this does not undermine the integrity of our mathematical analysis. Note that in our proof, we use the absolute attention score difference |da]. The main logic of our proof is that since attention scores display locality, the representative token near the maximum token is likely to have a larger score than the other representative token and, hence, is more likely to be chosen by HiP's algorithm. This does not change, even if the mean of 8λ displays the aforementioned nonzero characteristic. Since dn moves away from zero as A increases, |مa| is an increasing function of A. This makes even more likely that the representative token near the maximum token will have larger score than the other representative token compared to when the mean is zero. For this reason, vi (a(()), (3)) is still decreasing function of and increasing function of 3, even if the mean exhibits nonzero characteristics. Thus, our mathematical analysis remains valid.
We approximate the mean zero mainly for three reasons. First, even though the mean gradually moves away from zero, regardless, it is still close to zero and much smaller than the standard de- viation of òλ. Secondly, the behavior of the mean is not consistent, which makes it very difficult to express it in terms of math. Finally, the approximation does not compromise the integrity of the mathematical analysis while making the overall derivation much simpler and easier to understand. These are the reasons why we approximate the mean as zero in this paper.
A.3.2 THE INDEPENDENCE BETWEEN da AND 5g3
Now, we talk about the independence between 8, and dg. In Lemma for simplicity, we have assumed that 8² and óg is independent of each other. However, this is not necessarily true. As can be seen in Figure 17, the correlation coefficient between 8, and 83 is nonzero. From empirical analysis, we observe that the correlation coefficient between 8, and 6.g is actually dependent on la - 31 - the distance between the two tokens. As can be seen in Figure 18, the correlation coefficient decreases from around 0.7 when the distance is short to around 0.4 when the distance is long.
This suggests that d, and ōg have a strong positive correlation when a and 3 are close to each other, and a weak positive correlation when they are far away. In other words, nearby tokens have higher probability of having similar attention scores. This is exactly equivalent to the assumption of attention locality that we discussed in the main section of this paper, and this empirical data is the strongest evidence we have that supports this assumption. To summarize, due to attention locality, dA is not independent, and nearby tokens are less independent than faraway tokens.
28

---

Published as a conference paper at ICLR 2025
Head 1
Head
Head 9
Head 13
10.0
1.5
110.0
Chef
Head 17
Head 21
Head 25
Head 29
2.
1.5
8100
7.5
-10
13.6
Lomrestion CONT
-300
-15
Conalation Co
20
Figure 17: Correlation Coefficient Between 010 and 6100- The above figures show the scatter plots of the joint distribution of 810 and 6100. The nonzero correlation coefficient suggests that there exists some amount of dependency between d1e and 8100- The 16 layer of the Llama3.1 8B model was used for sampling the values. We use Wikitext2 dataset for the input.
Head 1
Head
Head
Head 13
0.20
0.70
0.65
0.65
0.90
0.50
0.60
0.55
0.59
3.5
0.53
0.50
WWY
0.65
0.60
0.40
0.00
Head
100
Head
200
Head
13c
150
20c
Head 29
103
150
Comelation Coefficient
0.90
0.00
0.85
0.55
0.52
0.43
0.45
0.42
0.45
1.0
102
[az - 81
Figure 18: Correlation Coefficient. The above figures plot the relationship of the correlation coef- ficient between & and ōg. and the index difference Ja 3 for different attention heads. In order to obtain a more clearer plot, the value of was fixed to 10. As be seen in the figures, the correlation coefficient consistently decreases from 0.7 to around 0.4, as the distance between a and increases. The values sampled from layer 16 of the Llama3.1 8B model, using the Wikitext2 dataset as input.
If so, then does the dependency between 8,, and 8, disprove our analysis? In order to analyze this issue, we compare p(6a = x, 63 = 8) when ds are independent and dependent. When they are independent, then P(6.. Ir, dg is expressed follows.
1
x²
y
(6 = =x, = y) = =P55. =x)P(h> = v) = 2nd (a)o exp - 20 (
20
29

---

Published as a conference paper at ICLR 2025
When &a and óg are dependent with the correlation coefficient being p. the probability can be pressed using the multivariate normal distribution as follows.
P(é.- x, dz y)
( p exp 2 p²| (0( + y² 20(
1
1
2p ( 1)
We argue that the critical difference between these two equations is the nonlinear term 1 p². When the tokens are far apart, i.e., p 0.4, then the 1- p² 0.85, which is relatively close to 1. Therefore, we can largely approximate the dependent P(ó, = x, dg = y) equation with its independent version without much error. However, when the tokens near each other, i.e., = 0.7, then - p² = 0.51, which is significantly different from 1. In this case, we can no longer justify the approximation of independence.
This means that if the representative tokens are too close to each other, then our mathematical analysis, which proves that HiP performs better than random token selection, might no longer hold. Interestingly, we have actually observed experimental results that support this notion. In the last few iterations of HiP, we have noticed that HiP's binary search algorithm does not perform well, unlike previous iterations. The main reason behind this phenomenon is that in the last few iterations, the sections become too short and too close to each other. Since they are all close to each other, due to attention locality, there longer exists significant attention score difference between each section. Hence, HiP fails to successfully identify the sections with the largest attention scores, and deteriorates almost random token selection level performance.
In order to resolve this issue, we tried to develop a simple patch that acts after the top-k selection process. Recall that in order to foster the efficient usage of matrix multiplier units, HiP was de- veloped with block granularity. Previously, after the hierarchical top-k selection process, only the selected blocks are used for creating the sparse attention mask. Instead, grabbed entire vicinity of tokens around the selected blocks, and used those for the sparse attention mask. The key idea was that due to attention locality, the tokens around the selected blocks would have the highest attention scores in the entire sequence. If HiP cannot choose among those tokens, the easiest workaround would be to use all of them. Therefore, by simply using the entire range around the selected blocks, we thought that we would be able to save the most important tokens from being discarded. During development, we used range sizes of 16 or 32.
The effects of this experimental revision definitely improved the overall performance. As can be seen table Table 7, the relative performance compared to FlashAttention2 was enhanced by 5%. However, as this patch requires the algorithm to access much more key tokens than before, it is sig- nificantly detrimental to the overall latency. Due to this reason, we did not include this experimental patch in the final version of HiP. Yet, believe that with further research, may be able to find a method that can exploit the strengths of this approach while minimizing the additional latency. Thus, reserve this topic as a future work. Table by Patch Results. patch where the by values are modified for the last few iterations. Results show that the patch is beneficial for overall performance yet comes at the cost of latency. The Relative Performance is measured compared to Altetonte The unit of the prefill latency is milliseconds (ms). Also, note that we use 128k context length for experiments.
order to resolve the dependency issue of dΛ, we test an experimental
Prefill
NQA
Qasper
HQA
2WM
GR
MN
Rel. Perf
Latency T-128k
Flash Attention
29.53
44.27
54.61
39.49
34.99
27.39
100.0%
855.2
HiP A: 512
21.44
41.73
50.36
45.17
30.09
26.01
92.40%
105.7
Range Size
25.46
44.12
52.87
42.30
33.40
26.89
97.51 %
181.4
Range Size 32
26.82
44.30
52.85
41.23
33.71
27.18
97.88 %
227.2
To summarize, when the representative tokens are far apart, the dependency between òa is weak, so approximate them as independent. When the representative tokens are nearby. (i.c. last few iterations of HiP), the dependency of 5ة can no longer be ignored, and thus HiP's binary search like algorithm longer suffices. We acknowledge this issue, and have tried to develop experimental
30

---

Published as a conference paper at ICLR 2025
method that solves the problem. We find that while the method definitely helps performance, comes at a large cost of latency. We plan on tackling this problem in future work.
B DETAILED METHODOLOGY DESCRIPTIONS
B.1 HIERARCHICAL SPARSE ATTENTION MASK ESTIMATION ALGORITHM
Algorithm Hierarchical Sparse Attention Mask Estimation
input Queries Q E Ig Keys K E RR Sparsity constant k, Query block size by Key block size by. Top-r approximation constant r. output Estimated attention mask M E 1J,T which is represented by array of indices
I E [1: TTTT/bxk//h
1: Q, K reshapet/M, /h, xh, xd Q) roshapt/// xbs xa K.)
2: Number of iterations = [Tog(T/bb))
for each query block index
T/ do
4: (Fo) ( (L - 1) 1 + + l I for 1.. k.
>> nodes' initial start and end indices
5: for each iteration i 1 nist do
6: for each node index j = 1 k do
7:
m (4)
(1
(10 //2
2)
(1 (
= (1
(1) 1), (m 1.))
8:
9.2 9.23)
m 4j
9:
for each branch index h = 1 2k do
Pick a center index righ from the range ((i)
10:
Compute score sth maxm.n [causal mask
Q Qg.,.
K
) (i)
Pick top-k indices "d.....,te } of the sequence q.1:
(1)
9.2k
Update nodes (F/) (ff(tt)), ,(++1)) = B.)) for j =
4j
Set mask indices Inj = 53 p(tne) for j 1
14:
In Algorithm 1, we describe the complete algorithm used for mask estimation. The peak amount of memory used during the mask estimation process is in OTT), since each iteration, only the immediately previous iteration's 's node indices are needed, and the rest can be discarded.
B.2 HIP DECODING ALGORITHM
Algorithm HiP Decoding Algorithm
input The model M, the number of layers L, the mask estimation period rim number of initial dense layers la-
output Generated sequence y.
1: Initialize y with empty sequence.
2: while generation has not ended do
3: for each layer = 1.. N do
4: if < la then
5: Perform regular dense attention for the current layer.
6: else
7:
if the current generated sequence length is divisible by I'm then For each head, estimate the attention mask with Algorithm 1. Cache the attention masks.
8:
> O(TlogT') time O(T) space O(T) time & space
9:
10:
Perform fused sparse attention using the cached attention masks.
11: Sample the next token and append it to the sequence y.
In Algorithm 2, we show rough sketch of the decoding process with HiP. In particular, note the function the mask estimation period r., and the number of initial dense layers la, as well as the time and space complexities of each component.
31

---

Published as a conference paper at ICLR 2025
Step 0. Initialize Chunk by Uniformly Split Sequence by k/ b_k
Poggan
Step Divide Each Chunk by Half. every chunk not divisable, terminate
1 Chank
Ieratio Churk
teeoter 1 Chunc
henton Crunkia
beatton Churki
headinn Chunk
Step 2. Evaluate importance for each chunk
haraton1 Chank
Banaton Churk2
Chunce
bouson Crunku
acrsson Churks
beraten 1 Chark
Step Select top k/ b_k important chunk
Chunk X- Importance Measure
Kuuii 10
Kry Stock
Key Blick
Step 4. Go back to step with selected chunks.
2 Chank!
Churk2
Memory Mernnrhy
Imprrtinee SC0
Final. Returns approximatly selected k tokens
Scaler
HBM Venter
Detailed Flowchart of Hierarchically Pruned Attention. We illustrate the internal steps of each program of hierarchical attention pruning. We instantiate single program for each attention row due to row-level synchronization on the top-k operation.
Figure 19:
B.3 DETAILED FLOWW-DAGRAM OF HIP
In Figure 19, we illustrate hierarchical attention pruning step-by-sted step.
B.4 ADDITIONAL OPTIMIZATION TECHNIQUES
B.4.1 TOP-R APPROXIMATION
In order to reduce the cost of the mask estimator even further, we take inspiration from SparQ Atten- tion (Ribar al., 2013), where global memory (HBM GDDR) transfer is reduced by selectively fetching only the most relevant components of the key vectors. Specifically, when computing the inequality condition in Equation (4), instead of fetching all components of the key vectors, only fetch r < d most prominent components estimated by the query vector q. Thus, we compute the following as an approximation:
ki =
>
Ap.
kp
(6)
where (P1, ... global memory accesses in the mask estimation stage can be drastically reduced. However, we disable this approximation by default.
p,}
argtop r((q)). By using the top-r approximation, the total number of
B.4.2 BLOCK SPARSE FLASH ATTENTION
We utilize the Flash Attention (Dao et al., 2022) mechanism to reduce the latency of sparse attention and use a small size sliding window to reduce performance degradation on the side of the sparse attention kernel. Following the receipt of StreamingI.LM (Xiao et al., 2024), local sliding window and global sink attention are also added during block sparse flash attention operation.
B.5 TRAINING DOWNSTREAM TASKS WITH HiP
In this section, we describe the HiP attention training strategy for downstream tasks. We discovered that direct fine-tuning after applying HiP could not achieve the performance of the fine-tuned vanilla
32

---

Published as a conference paper at ICLR 2025
attention. Empirically, HiP's highly sparse attention matrices show excellent performance approx- imation during test time but not train gradients. Since our method heavily prunes the attention matrix, the gradient cannot flow through dense attention probabilities. This incomplete and unsta- ble gradient the attention matrix leads to significant training performance degradation because HiP forces the model to have attention patterns similar those of the pretrained model rather than adopting them for the downstream task.
However, we could achieve the same performance as vanilla by using the following two-step training strategy: (1) fine-tuning with vanilla first and (2) further fintuning after applying HiP (healing). First, we train the pretrained model with vanilla attention the downstream task as usual. Then, we load the finetumed weight, apply HiP, and further finetuning with the downstream task dataset with just a few optimizer steps. We call this a healing process because we only finetune a few steps from the finetuned weight. For training details about each experiment, we describe hyperparameter and optimization setups Appendix D.
C ADDITIONAL EXPERIMENTAL RESULTS
C.1 LARGE MULTIMODAL MODEL WITH HIP
Table LMMs-Eval on LLaVA-166-3BB, We apply the efficient attention mechanism in only LLM
backbone of LMM.
MME
MME percep
Method
MMMU
DocVQA GQA TextVQA
Avg.
cog
Flash Attn
324
1575
36%
77%
65%
67%
100.0%
SLLM k=512
242
1094
31%
25%
52%
27%
63.7%
BigBird k=512
320
1526
36%
53%
64%
55%
90.7%
HiP 256
303
1465
37%
68%
64%
63%
94.7%
HiP k=512
286
1500
36%
74%
64%
65%
95.9%
Since large multi-modal models (LMM) (Liu et al., 2023b; 2024) leverage pre-trained LLMs as a backbone for their NLU ability, without changes, except for the tokenizer, also evaluated our (Liu et al., 2024), a large multi-modal model, using LMMs-eval (Li* al., 2024), which provides extensive benchmarks for large multimodal model suites. As we show in Table 8, our method scores 95.9% relative scores, which is similar to the performance recovery ratio of HiP on NLU tasks.
HiP on top of L.L.VA-1.6--
C.2 MASSIVE MULTITASK LANGUAGE UNDERSTANDING (MMLU)
Figure 20: MMLU scores on Llama3.1-8B.
Method
Huma- nities
STEM
Social Sciences
Other
Avg.
FlashAttn 67.35
51.88
73.21
66.52
63.24
SLLM
66.65
52.22
73.98
65.68
63.16
BigBird
66.94
51.06
73.27
66.21
62.81
HiP (Ours)
67.64
51.50
73.19
66.32
63.13
Next, we evaluate HiP on the MMLU benchmark (Hendrycks et al., 2021) to show that our method does not negatively affect the NLU ability of the pretrained model. The results show that our HiP preserves the NLU performance of the original model. All tested methods are able recover original MMLU scores without significant loss here. This is probably due to the nature the MMLU task: the answer is only dependent on the most recent span of tokens rather than the entire prompt (which contains fhw-shot examples in the beginning).
C.3 COMPARISON WITH REFORMER AND SEA
We compare the performance HiP against Reformer (Kitaev et al., 2019) and SEA (Lee et al., 2023) using the Llama2-7b model and show the results in Table 9. Even though HiP requires no
33

---

Published as a conference paper at ICLR 2025
Table 9: Comparison of WikiText2 perplexity against Reformer and SEA. For this comparison, we use the Llama2-7b model. The experimental setting for HiP the same as the main experiment. For Reformer, we use 512 for the bucket size and replace the attention module in all layers except the first three for fair comparison with HiP. For SEA, we use 512 for the attention predictor length and & (the top k sampling factor) 64. We fine-tune both Reformer and SEA models using LoRA with rank 128, the same as HiP HEAL, until convergence.
Method
Perplexity
Vanilla PyTorch
5.5933
Reformer FT (Kitaev et al., 2019)
45.5130 (+39.92)
SEA (Lee al., 2023)
23.1964 (+17.60)
HiP (4512) PP
5.7938 (+0.20)
HiP (e512) HEAL
5.6872 (+0.09)
Table 10: Self-Extend Result with HiP. We apply Self-Extend with HiP k=256 on Gemma2-2B-it We extend the maximum context length of Gemma2 from 8k up to 128k tokens. We measure the perplexity of Wiktiext2.
Self Extend
Sliding Window
Method
T=1k
T=2k
T=4k
T=8k
T=16k
T=32k
T=64k
T=128k
FA2
18.90
14.48
12.35
11.40
46.66
167.40
363.02
FA2
x
18.90
14.48
12.35
30.17
198.84
638.65
1310.26
HiP
/ (x4)
X
18.98
14.77
13.00
12.14
12.00
46.51
277.54
HiP
V (x8)
X
18.98
14,90
13.30
12.57
12.22
12.63
49.26
HiP
/ (x4-8)
X
18.98
14,81
13.15
12.36
12.06
13.62
116.18
HiP
V (x8-16)
18.98
15.00
13.51
12.83
12.53
12.55
14.32
HiP
V (X8)
19.28
14.86
12.87
11.95
11.53
11.36
11.69
HiP
(x16)
19.28
14.98
13.11
12.25
11.86
11.69
11.54
11.68
fine-tuning, unlike Reformer and SEA, which need fine-tuning, our HiP attention is far superior in language modeling performance compared these two baselines.
For fair comparison, we fine-tune Reformer and SEA using LoRA (Low-rank adapters), which have the same rank as the healed version HiP. Due this, the Reformer and SEA's performance converges to much-degraded value compared to the values reported in their respective original papers. We conclude that both methods need much modification to the original weights in order to perform well, whereas since HiP performs well even without any modification to the weights whatsoever, a small amount of LoRA training can even close this small gap.
C.4 CONTEXT EXTENTION WITH SELF-EXTEND
Since our method has a considerable advantage long- decoding, we need the pre-trained long-context model. Unfortunately, however, all pretrained transformer models support long contexts. Therefore, many previous studies (Peng et al., 2023; Jin et al., 2024) try to extend max- imum position embeddings of the trained model extend the context size of LLM. We adopt the SelfExtend (Jin et al., 2024) method into our method because it is also a training-free context ex- tension method. We picked Gemma2 (Team, 2024) to target LLM and extend the context model because the model is hybrid of sliding windows and dense attention. Therefore, it will have the advantage of long context model with HiP by saving the KV cache of sliding window layers. The Gemma2 repeats the attention layer by repeating the stack of different attention blocks: sliding win- dow and dense attention. To evaluate the effectiveness of the combination of HiP and SelfExtend, apply them to attention layers.
We can observe Gemma2 explode after its pretrained context length, which is 8192 (First row of the Table 10). We can see the model fails after the sliding window context length, which is 4096 for the sliding window layer (Second row of the Table 10). Therefore, we know that treating sliding windows especially is quite essential for performance. In the third and fourth rows of the Table 10, we apply the same Self-Extend group size for every attention layer, including the layer that was orig- inally sliding window, before replacing with HiP. We could observe the settings struggling to
34

---

Published as a conference paper at ICLR 2025
recover performance right after Self- Group Size x Sliding Window Size, we apply twice the larger Self Extend group size for the HiP layers originally sliding window. The modified group size application is in the fifth and sixth rows of the Table 10. We could observe that can extend the context window as expected, not bounded by sliding window size. However, the above replace- ments are quite restarted setting because they are not even allowed use sliding window attention, which is usually very efficient in practical LLM frameworks, Therefore, we replace only dense at- tention to HiP, and we could observe a significant performance boost from all layer replacements while extending context length up to 1 128k, as shown in the last two rows of the Table 10.
C.5 ENSEMBLE HIERARCHICAL ToP-k APPROXIMATION
For the first attempt to address the challenges in HiP described in Appendix E.9, we perform an ensemble of HiP masks by introducing slight randomness to the branching decisions within a given iteration. Our empirical results on Wikitext2 indicate that the ensemble of HiP with no dense layers (la 0) achieves comparable perplexity to HiP with dense layers (le 3) across context lengths from 4k to 16k. This highlights how considering different branching strategies within a given itera- tion can improve inference performance and suggests the potential for replacing the dense layers in HiP with comparable performance.
Methodology. First, we generate ne HiP attention mask samples by adding randomness to the branch direction when the middle iteration branches out to the next iteration. As re slightly adjusts the selected node in the iteration, each sample can take a slightly different path during tree branching, leading to a diverse construction of attention masks.
Second, we count the number of agreements of the indices per query and retain the indices that exceed or equal to the agreement threshold vyte. Therefore, v = 1 functions as a union operator and Owore n. an intersection operator. By adjusting O from I to ne. we perform an operation that lies between union and intersection, prioritizing indices with a higher number of votes. To prevent the union-like Oose increasing the number of active attention elements much, we truncate the number of the final selected indices by original A when T E {0, is 1, prioritizing the indices having more votes.
Lastly, we perform the introduced ensemble method for the first L. layers. like we do with la-
Experiments. We first provide experimental evidence how ensemble enables end-to-end sub- quadratic complexity and then give details on our hyperparameter tuning process. The experiments are conducted with the LLaMA2-7B model.
PPL. Lengths
PPL over Rel. Retention Ratios
6.85
fncembo
5.30
Sueel (y
Spg (r
5.03
Crusmble
5.70
5.5
5.00
5.5
7456
7-124
1254
1,559
1,442
1.577
Seq. Langth
Rel. Ratention Ratio 1
Figure 21: Perplexity Evaluation on Long Context. (Left) Perplexity result in Wikitex12 with different ensemble settings (8vote: 7) where re 5.0, 4 = all. (Right) Perplexity comparison between full HiP (la = 0) and ensemble (f = 1, T = 0, re = 5.0, Le = all) using same relative retention ratio in each sequence length.
Performance Comparison with Original HiP. To show that ensemble enables end-to-end sub- quadratic complexity with comparable performance to our default HiP (la = 3), we compare the full HiP (lg = 0), the default HiP (la 3), and the ensemble with la = 0. We fix r. 5.0, le all that gave the best performance in T = 4096, as shown in Figure 22. The result indicates that ensemble with yooe 1, 0 outperforms both full HiP and default HiP, shown in Figure 21 (Left), and therefore this suggests that ensemble could not only improve the performance but also replace the dense layers with comparable performance.
Moreover, we provide a comparison with full HiP (l. = 0) at the same level of sparsity as the ensemble to demonstrate that the improvement is not solety due to the increased number of selected
35

---

Published as a conference paper at ICLR 2025
Effect of 9at and ra in PPL
Where to Ensemble
5.0
2.84
5.04
2.33
5.55
3.2
5.02
3.87
Recet (1m-2)
5.00
279
Raw S22
5.79
Randomness Vta)
*f Ensemble Enabled Layer (1.)
Figure 22: Effect of Randomness in HiP sampling and Eassemble-Enabled Layers l.. Perplexity in Wikitext2 with T = 4k. (Left) We adjust randomness r. with fixed L. = all. (Right) We adjust number of L. with fixed Ovate = 1, re = 5.0. The dashed horizontal line shows the performance of HiP (le = 0, La = 3).
indices resulting from our ensemble method. As shown in Figure 21 (Right), our ensemble method is Pareto frontier compared to HiP, with performance measured against the retention ratio.
Latency of Ensemble. Since we sample multiple HiP masks by ne and perform voting operations across R. x k number of indices, the ensemble costs a few times more than the original HiP. How- ever, since the cost of dense attention grows quadratically, the ensemble will become more efficient compared to the dense attention as the context length increases. Therefore, we think that the use of the ensemble method could be particularly advantageous in extremely long contexts.
Rel. Retention Ratio After Ensemble
Retention Ratio Changes over
Retention Ratio Changes over 1
Suyel (r 3)
Sug E1
125
Ret. R...eentiin Raito
1.4
11 1D
Rel. Ratention Rado
1.23
Rel. Racention Ratoo
1.20
1.15
1.15
1.2
1.10
6.1
1.85
5.0
7.5
12.0
20
Seq. Lengtth (k)
Randornness (ra)
Ensemble Enabled Layer
Figure 23: Relative Retention Ratio and Ensemble Factors. Relative retention ratio after mask ensemble. The ratio is computed by dividing the number of active indices after the ensemble by before. (Left) We adjust the sequence length I' with fixed Fe = 5.0, Le = all. (Center) We adjust the randomness re in T = 4k with fixed le = all. (Right) We adjust the number of ensemble enabled layer le in T = 4k with fixed vuce = 1, re = 5.0. Ensemble disabled layers are computed as 1.0.
Ablation Study of Ensemble Hyperparameter: @oote T, Tes Le. As shown in Figure 22 (Left), Ovale = with - = O, and Te >= 5.0 gives the highest score. This indicates that performing the union operation with large randomness in sampling gives the highest performance. Also, in Figure 22 (Right), with &vose 1, we show that le = all works the best, therefore if we want to replace the few layers of vanilla attention, then we would have to apply ensemble in every layer.
Correlation Between the Relative Retention Ratio and Ensemble Factors. In Figure 23, we show that the relative retention ratio shows no correlation with sequence length, while randomness shows a positive correlation. Moreover, when measuring the relative retention ratio changes over le with &ver 1, since we treat the ensemble disabled layer as having a 1.0 relative retention ratio, more ensemble-enabled layers lead to higher relative retention ratio.
Analysis with Visualization. In Figure 24, we provide a visual analysis to show how the ensemble selects indices that HiP missed to fill up a complete attention pattern and how it enables dynamic sparsity. We can see how the ensemble catches missed important indices such as diagonal, vertical, and stride attention patterns in (a). (b), and (c) of Figure 24, Moreover, compared to HiP (left), the union operation (center) enables dynamic sparsity per head. Especially in (c), we can see that the ensemble is effective for filling missed indices in long sequence while providing dynamic sparsity for each row (red pixels are gradually frequent in the bottom). Lastly, in Figure 24 (center, right), show how selects indices that receive more votes compared to those selected by 0.
36

---

Published as a conference paper at ICLR 2025
HP Altertion Mask Visudizztion in LLaWA2 79 Ruyer Head 23)
[a]
HP Altention Mark Visulitattoon in LLaVA3 70 3layer31, Heed
HP 5123
Ensemble EI
Ensemble
10M
[d]
HU $1B)
Ensenblee
Ensemble (r
Figure 24: Attention Mask Ensemble Visualization. Visualization of attention mask in = 16k for (Left) HiP (la = 0). (Center) ensemble (0 1, T= = 0). (Right) ensemble (6vote 1, 1). Red indicates indices added by our ensemble method, yellow indicates indices from HiP, and green indicates where attention will not be computed.
D DETAILED EXPERIMENTAL SETTINGS
Computation Resources. We use machines We use this local development machine. Most of the micro-benchmark and kernel optimization done with this machine. (2x RTX 4090 24GB, 128GB DDR5-3600, Ryzen 7950x), (2) A100 Machine. We use this AWS cloud node as the main computation horse. All training and most long context benchmarks are measured with this machine. We different GPU architectures for kernel development because we could not get an H100 AWS node due to lack of quota. Therefore, our kernel's s CUDA hyper-parameters, such as the number of warps per program and block sizes, are not optimal in the A100 machine. To overcome these mismatches between the development machine and the computation horse, we used t ton. autotune as much as possible. However, for the above reasons, the optimization of the GPU kernel may not be optimal for every GPU architecture.
experiments and training. (1) 4090 Machine.
Experiment Settings. By default, for HiP experiment, we use by = 32, by 2, = 512, = 3 if 7B else 4, r'm 8. For StreamingLLMM, we use num_sink 4.
We show overall experiment settings, such as the number of GPUs and model IDs to which the experiment is introduced (c.g., the caption of the figure and table). To reference, leave the huggingface model path Table 11. We used an instruction model from the same provider for instruction following ability-required tasks such passkey, LongBench, and RULER.
Experiment Details. We measure every latency measures with a single NVIDIA RTX 4090. PCLe 4.0 x8, 128GB DDR5-5600, and Ryzen 7950x. The batch size is 32 for decoding and I for prefilling. The official implementation of HyperAttention is not available for decoding; therefore, we did not measure the decoding latency and decoding benchmark.
37

---

Published as a conference paper at ICLR 2025
Table 11: Model IDs on Huggingface. We use large language models trained on long context inputs order demonstrate our method's effectiveness long context lengths.
Maximum context length
Model
Huggingface ID
LLaMA2-7B
32K
LLaMA2-13B
-ft
32K
Qwen1.5-2 1.57B
Oxer/0wen1.55-78-00
32K
Qwen1.5-14B 1.5-14
Oxen/Owenl.5-14g) 142 Chat
32K
Luxia-21 21.BB
32K
Llama3. 1-8B
neta--Laamm//ettttlllaaas......
128K
Llama3. 1-8B-Instruw.
atta--1aamm//ett--_aaa-333----2 met
128K
Gemma2-2BB-4
2 2-2b-it
8K
Gemma2-9B-it
google/gemna 2 9b it
8K
Exaone3-7.88
0-7.....0-7....
4K
For FlashAttention, we use backend, such as vLLM, SGlang, and HyperAttention.
tt===.66666 6.3 for every baseline that requires FlashAttention
We do not use HyperAttention in another experiment because it fails to recover the perplexity of PG19, the most basic metric of language modeling. For HyperAttention (Han et al., 2024) we used 1sh projs=7, bl block. s ze=64, sample s ze=1024, min seq len=32. We select mi ninneeq_ees to match the size the MMU's block size (32) rather than 4096, which is suggested by the authors' code repository. Since in sortL.H (Han et al., 2024) it processes shorter block size of min_seq.len with vanilla attention. Therefore, we have to reduce its size to a smaller size than 4096.
Since StreamingLLM does not use block-wise computation just like HiP's block approximation, cannot utilize TensorCore in GPU. This downside degrades the throughput significantly compared to HW-aware algorithms such as ours and Flash Attention (Dao, 2023). However, since the method requires different RoPE index for every query-key dot product, cannot easily adopt block sparsity on their method. This will slow down attention computation more than twice because the RoPE re-computation costs twice much as attention score computation.
In the Figure 7, the latency is measured with our latency measure machine (RTX 4090). StreamingLLM shows OOM 32k context due to the overhead of the COO sparse matrix. Hy- perAttention shows invalid kernel parameters over 32k context due to heavily nested tensors in long context. It uses nigh-dimensional striding to perform reformer-style token clustering, but the backbone flash attention kernel does not support that sigh--imensional striding with a larger tensor.
In the Figure the latency is measured with our latency measure machine (RTX 4090). The ma- chine has about 4GB of VRAM available for the KV cache, excluding model weight and temporary buffers. We limit the size of the CPU offloaded KV cache 32GB. The tested model is Llama3.1 8B.
Training Details (Healing HiP in Arxiv). The HiP healing in Table done as follows. For the Llama3.1 1-8B model, after applying HiP, we fine-tune the pretrained model on the Arxiv dataset for 500 steps with AdamW optimizer with lcarning rate Ic-5, and with batch size 32, LoRA rank 256, and HiP's hyperparameters set to by 2, b, 64, baq 2, buk 1, k 512, L = 3. We use the Arxiv dataset Redpajama (Computer, 2023). The inputs are truncated to a maximum of 8192 tokens to speed up the process.
The purpose of this fine-tuning, which we call healing, is to make the model adapt to the slight differences in the activations that appear when the original dense attention layers are replaced with sparse HiP attention. As shown in Table 3, the healed model performs slightly better than the plug-an houg-and-play (unhealed) model on LongBench. However, emphasize that HiP is meant to be training- free, and healing is just an additional option for extra performance, as our method already works perreeeffetty without training.
38

---

Published as a conference paper at ICLR 2025
Complexity
Is Training free?
Is Dynamic Attention?
Time
Space
HiP (Ours)
Log-linear
Linear
SEA (Lee et al., 2023)
Linear
Linear
X
FlashAttention (Dao. 2023)
Quadratic
Linear
/
StreamingI.I.M (Xiao al., 2024)
Linear
Linear
HyperAttention (sortI.HH (Han et al., 2024)
Near-linear
Reformer (LSHH) (Kitaev et al., 2019)
Log-lincar
Log-linear
X
Performer (Choromanski et al., 2022)
Linear
Linear
X
Cosformer (Qin et al., 2022)
Linear
Linear
X
Longformer (Beltagy et al., 2020)
Linear
Linear
X
Table 12: Comparison of time and space complexities of various efficient attention methods.
Prefill
Decode
Prefill
Decode
A
H2O
StreamingL.LMM
HiP (Ours)
AVD
Hip VD (Ours)
BigBird
Sink & Sliding Window
V (Vertical)
(Diagunal)
Random mask
"Heavy-Hirters"
HiP's Top-k Fstimated
Figure 25: Visualization of sparse attention baselines' attention masks.
E ADDITIONAL ANALYSIS
E. MORE DISCUSSION ON RELATED WORKS
In Table 12, we compare efficient attention methods that we do not include in our main benchmarks. For each method, show the time and space complexities, whether is training-free, and whether it uses dynamic attention. Dynamic attention methods can attend to each other content dynamically rather than using static attention patterns such as sliding window. Besides HiP, HyperAttention is the only method that satisfies all four criteria, but HyperAttention suffers from substantial per- formance degradation (see Figure 7). Figure 25, we conceptually visualize the various sparse attention baselines' attention patterns for extra clarity.
StreamingLLM (Xiao al., 2024). StreamingLLM uses sliding window attention with an tention sink, which processes the input sequence in linear complexity without resetting the KV cache; they call this process 'streaming." StreamingI_LM introduces the attention sink, which is similar to the global attention token in Longformer (Beltagy et al., 2020), and streams the KV cache using RoPE indexing. However, due to the sliding window, the method cannot perform long-context knowledge retrieval. Therefore, this method cannot utilize the full context, and they do not extend the context window of the model by any amount. Since the method loses the key-value memory time passes, it cannot take advantage of the Transformer's strength: its powerful past knowl- edge retrieval ability. Furthermore, since they use a different RoPE indexing for every query-key doodpuo, they cannot utilize MMU, which is critical speedup factor in modern accelerators.
39

---

Published as a conference paper at ICLR 2025
Atteetnttntion (Han et al., 2024). HyperAttention introduces sortLSH, improved version of LSH (Kitaev et al., 2019), to work as plug-and-play The method uses block spansity to utilize MMU. It is training-free free, has sub-quadratic time complexity access to every past key token, much like our method. However, PlyperAttentinn struggles to recover vanilla performance when replacing most of the layers in the trained model in training-free manner (see Figure 7).
and has the ability to potentially
Sparse Linear Attention with Estimated Attention Mask (SEA) (Lee et al., 2023). Inspired by SEA's framework, which introduces linear complexity attention estimation and sparse matrix interpolation, we aimed to improve its efficiency. SEA estimates each query's 's attention probabilities over the keys with fixed-size vector, turns it into sparse mask by selecting the top-k elements, and resizes it; this process is done with linear complexity. However, the method is difficult to implement efficiently due its extra modules, mainly the estimator and sparse matrix interpolation. Furthermore, the method does not support block sparsity; thus, it cannot utilize the MMU. We were motivated to improve this work drastically by introducing a fused and train-free attention mask estimator, HiP.
E.2 ANALYSIS OF SUMMARIZING RESULT BETWEEN STREAMINGLLM AND HIP Example of Generated Summary (StreamingLLM, Human, Hip)
StreamingLLM
Human Ground Truth
HiP (Ours)
procurment Bitilin
press 1000)
the
gom or Baok Defense
Depariment
coluior colao,
on lnghtes Jude
Danpuree
MEl slow
sproohh enrmaning
years. tyrray
of Mil preammer
feds the
reduping
aporzech
formel
ha grom fariearyy wht losd valu or Tralel
sccrrl percent Uneer amiull canrattr B0S ta and of Under
statcoory conmance The repert highigg
Mardiiaa key
compared with contrcting The BBd The Cosel though
tacon comtistigg Wil sueen nuch incroved
resed cas gige
ranong
the
compattion Howeer the lenges riman, indlcing need
whhecr hadrg to
conmect
each year after
dedclly and ind.una
anate The report Ragpecs congree my conider thase
reanaal for untonnded conanquaren uerte The report a mojr opotens predde rmiganc approsd for the miyy interrtee the movirance Will menmiing
ONE pormaness manee goominng nn Bill commentn
(Tirth dy for Compared nored romt under pimposed
Dasaromen
ranged lace than fo move thon
controciigg with the sumorrty to Bid for ownnn maye cheigg by the Coom Guard or Bel ware sod aed cocle programs
angorig
no uotoo and improveeent
with may and moner baing
Sic to 109 In desevv verfy beca cost Wil
report oi cuuses the and the poosiaall exparding ther kuose croating parmann mermnnen for black buy cormint ng. repon als the
Lack of Keywords:
cone settig emiar MDl that pomis 600
har wothh picurnnnt of
given afe the
had to mamoie cermic ecten
arcarement prog time reppriee crnteggt bustalary
Black Buy Contracting (BBC)
800 wa opomen with tie ower
IBUSC 23058
There requiemenn foo
oter spproar aronc Progress loing condleed
Retrieval Failed!
becs
manes governig the Bid
than the years plrend arcarmments Econorri crzer and parch component
the bone be precure under the cormer iarn ng the you ther not nome botnrosch dtatat sutomant dan, buses tkey incldd carcalation indude wheter
The
Compare to Human Ground Truth
BBC eathorie beaane welilles 500 authonty
Compare to Human Ground Truth
festure BBC cortress
ROUGE-1:
25.47 %
ROUGE-1:
48.33 %
and BRE fiza more
ROUGE-2:
ROUGE
16.45
cureerd une whether to creae permanere matine given the usd and wlathe the Cousi Courd alould bagn
ROUGE-L:
14.18 %
ROUGE-L:
22.05 %
Keyword Occurrence Histogram (Stacked)
Some keywords are outside of the accessible key window a single StreamingLLM layer.
10
MYP
BBC
DOD
10 U.S.S. 2309b.
512 Token Window
1024 Token Window
10000
20000
30000
40000
50000
Document Location (Characters)
Potential Accessible Key Window Single HiP Layer
Accessible Key Window Single StrearmingLLM Layer
Figure 26: Summarizing Example of GovReport Dataset from LongBench. We sample random examples from GovReport summarization results with Qwen1.5-14B.
In Figure 26, we analyze one example generation result from GovReport in LongBench. We pick four important keywords from the human ground truth summary: MYP, BBC, DOD, and 10 U.S.C
40

---

Published as a conference paper at ICLR 2025
2306b. We pick two keywords (MYP, DOD) that appear in every summary and two keywords (BBC, 10 U.S.C 2336b) that appear only in the ground truth and HiP result. The result clearly shows that StreamingLLM is struggling gather information beyond its single-layer window size, & = 1024. Streamingl..A should be accessible to a much longer distance than & because information keeps ex- changing across time dimensions in each layer, like Mistral. In contrast StreamingI.LMM our pro- posed HiP attention shows successful knowledge retrieval in summary from long-range in-context documents, with the same plug-and-plan manner. Also, quantitatively, ROUGE-* scores show that the summary generated by HiPP is much better in coherence with ground truth than StreamingLLM.
E.3 HIERARCHICAL ATTENTION MASK PRUNING VISUALIZATION
Hierarchical Attention Mask Pruning
Iteration
Iteration
Group Size -
Group Size -
1230
1506
1500
Iteration
Iteration
Group Size
Group Size #1
1U0
1500
1750
500
1500
Figure 27: Visualization of Hierarchical Attention Mask Pruning. Yellow indicates non-zero entry of the attention matrix, and green indicates an empty entry of the attention matrix. We use k 512, by = 32, by 2, T 4k.
In Figure 27, we demonstrate real-world examples of hierarchical attention mask pruning. We sam- ple the Q. K, and V tensors from the first layer of LLaMA2-7B with random text sample from Wikitext-2. Note that each attention mask in the masking iteration is not the final attention mask. The final attention mask generated by this process is from iteration sparsity of the mask is low because the group size of blocks is very large (8), so the 8 2 key values are treated as single groups The attention score of that group is represented by the attention score between the query and the group's 's first block (b²).
In an earlier iteration, the
E.4 ABLATION STUDY ON BLOCK SIZE
Table 13: Ablation Study on Block Size: Perplexity on WikiText2. We use LLaMA-2-7B-- model, k = 512, di = 3, T 12k. Average shows the average perplexity of each row and column. Red means a bad perplexity score, and green means good perplexity score.
bg
PPL..
2
16
Avg.
5.5210
5.5172
5,4996
5.4952
5,5064
5.5331
5.5121
2
5.6032
5.6060
5,5863
5.5599
5.5577
5.5901
5,5839
by
4
5.6316
5.6280
5,5971
5.5872
5.5838
5.5870
5.6024
8
5.5807
5.5928
5,5553
5.5238
5.5100
5.5212
5.5473
Avg.
5.5841
5,5860
5.5596
5.5415
5.5395
5.5578
5.5614
We perform ablation study on block sizes (b,, by) using our method. Block size b, determines how many queries are grouped into the block during the masking iteration and sparse attention. And block size by determines how many keys grouped. Block size is really important factor in utilizing MMU (e.g.. NVIDIA TensorCore) in modern accelerators. MMU enables matrix multiplication and tensor operations be performed in single fewer cycles rather than processing one by one
41

---

Published as a conference paper at ICLR 2025
Table 14: Ablation Study on Block Size: Attention latency on decoding phase with single query. We use same setting with Table 13. Batch size is 96 and I'm 1 and sequence length is 12k. Average shows the average latency of each row and column. Red means bad latency, and green means good latency. The unit of latency is milliseconds. We use RTX-4090 to measure the latency of attention operation. Pytorch shows 20.7215 ms and FlashAttention2 shows 20.2527 ms.
by
Lat.
16
Avg.
1.6416
1.6397
1.6397
1.6398
1.6397
1.6401
16601
2
1.3543
1.3542
1,3530
1,3533
1.3535
1.3534
1,3536
by
4
1,2312
1.2311
1,2296
1,23015
1.2322
1,2306
1.2309
8
1,2836
1.2842
1,2830
1.2827
1,2830
1,2835
1.2833
Avg.
1.3777
1.3773
1.3763
1.3766
1.3771
1.3769
1.3770
Table 15: Ablation Study on Block Size: Attention latency on decoding phase with multiple queries (Speculative Decoding). We use same setting with Table 14. We use 32 query tokens to mimic speculative decoding scenarios. PyTorch shows 31.0632 ms and FlashAttention2 shows
20.3378 ms.
by
Lai.
2
4
8
16
32
Avg.
152.4423
77.2624
39.2077
20.1341
10.4925
5.5461
50.8475
2
85.3140
43.8207
22.7708
12.1814
6.4987
3,5056
29.0152
bg
4
54.2530
27.7025
14,3960
7.9660
4.5839
2.6156
18.5862
41.5079
21,3036
11,5369
7.4481
5.0634
3.0337
14.9823
Avg.
83.3793
42.5223
21.9779
11.9324
6.6596
3.6753
28.3578
using floating point multiplication and addition. This kind of accelerator trend leads to mismatching of wall-clock latency and FLOPs in modern hardware. Therefore, we check the performance and latency trade-off among grouping queries and keys by block size by. by.
In Table 13, show that perplexity gets better as b, increases while it gets worse as by increases. is not intuitive that increasing by shows better perplexity than before because they lose the resolution across the query dimension attention mask estimation. However, the result shows that more block size (more averaging) across the query (time) dimension shows better performance. In contrast to this observation, by works as expected, like that less resolution key (past knowledge or memory) dimension leads to worse performance.
This phenomenon makes our method speed up without any performance loss, even achieving bet- ter performance. In Table 14, we measure the micro latency benchmark of our attention operation during the decoding stage, which feeds single query into the attention operator. With a single query, cannot utilize the MMU fully because, during sparse attention and attention score esti- mation in masking iteration, we cannot matrix multiply between the Q group and K group. We have single query vector; therefore, multiplication, which is the main key feature of MMU. However, in Table 15, we measure the miero latency benchmark of our attention operation during the decoding stage with a speculative decoding strategy, which feeds multiple queries into the attention operator. We feed 32 query vectors within query dimension in the input tensor; therefore, now we can utilize a matrix-matrix multiplier in an MMU. With these multiple queries and MMU utilization, our method could achieve 10.23 times speedup on 12k sequence length compared PyTorch naive implementation (using bmm)
need a vector-matrix multiplier instead of matrix-matrix
We use by = 32, by = 2 by default, according to the ablation study across the block sizes. We choose by 32 because increasing by leads to better latency and perplexity. However, we stopped increas- ing by by 32 because the current modern GPU, especially the NVIDIA Ampere series, usually does not support matrix-matrix multiplication larger than 32. And maybe in the future, some variants will support larger matrix multiplication, just like Google TPU. However, larger blocks need more register allocation for block masking and address calculation. Therefore, considering implementa- tion limitations, we think there is benefit to increasing by infinitely. Also, from performance perspective, we do not think this trend will keep over b, 32. We choose by = because latency speedup from by = 1 to by = 2 is huge respect to perplexity loss.
42

---

could minimize Published as a conference paper at ICLR 2025
Additionally, measure the latency with I'm 1. which means without mask caching. Therefore, this speedup will be amplified with rm in a practical setting.
E.5 ABLATION STUDY ON DENSE LAYER CHOICE
Perplexity Wikitext2 Number of Dense Layers (4)
7-4k
5.8
7-8k
5.7
7=12k
PPL.]
5.4
1 Default
1 Full Dense
FullHHP
20
25
30
of Dense Layers
Figure 28: How Many Layers Should be Remained Dense Layer lg? We use la 3 as the default value. The Y-axis means how many first layers the Transformer model are kept as dense attention layers rather than replaced with HiP. We use Llama2 7B 32k for PPL evaluation Wikitext2
We do not replace the first few layers (la) of the Transformer because the first few layers tend to have dense attention probabilities rather than sparse attention probabilities. This phenomenon is well described in previous results (Ribar et al., 2013). The first few layers exchange information globally and uniformly across the tokens.
Therefore, in Figure 28, perform ablation study how many first layers should remain as dense attention. We observe that the first layers are kept as dense attention and then more perplexity. In other words, if replace the original transformer block with HiP attention, the performance degradation. However, for maximum practical speedup. we want to minimize the number of dense layers for the experiment. Therefore, we run the ablations study on different numbers of dense layers, and we choose 3. For 2 to 3, the performance (perplexity) improvement maximized. In conclusion, the practical effect dense layers on latency is minimal because the number of dense layers (e.g.. 3, 4) is small compared to the number of HiP layers (e.g., 29, 96). We show that we could achieve practical end-to-end speedup compared baselines in Figure 8.
E.6 ABLATION STUDY ON REPRESENTATIVE TOKEN LOCATION
In the theoretical analysis section, we have mathematically proved that selecting the middle token as the representative token of section guarantees better performance than random selection. In this subsection, we perform a brief ablation study on the location of the representative token to see this claim is true. In the following experiment, we measured the perplexity values of the Llama3. I 8B model for various representative token locations. We used the PG19 dataset as the input for the model.
'abble 16: Ablation Study of Representative Token Location. Perplexity is measured on the first 3.2M tokens the PG19 dataset with 128k context length. We Llama3, 1-8B for evaluation.
Position
First
0.2
0.4
Middle
0.6
0.8
Last
PG19 Perplexity
8.851
8.752
8.691
8.685
8.818
8.934
8.812
Table 16 shows the experimental results. As can be seen in the table, the perplexity value is min- imized when the middle token is selected as the representative token. This closely matches our expectations, and it asserts that our theoretical analysis of the representative token location is valid. Therefore, this ablation study justifics the selection of the middle token as the representative token.
43

---

Published as a conference paper at ICLR 2025
E.7 DISCUSSION ABOUT KV CACHE EVICTION AND COMPRESSION STRATEGY
We think the KV eviction and compression strategy is an orthogonal method to our proposed HiP method, and we can cooperate with KV cache strategies. Users can use sparse linear attention methods like ours with a KV eviction strategy. If the KV eviction strategy is careful enough, our method should retain the same performance as vanilla attention.
Also, the typical retention ratio (612/32000 = 1.6%) of our method is much more extreme than state-of-art eviction strategies (10 to 20% (Zhang et al., 2023)). Moreover, the KV eviction strategy loses information permanently, which should be problem. We think we can solve the memory pressure from the KV cache should be solved with the memory hierarchy the computer system. NVMe storage should be large enough to store everything. We think KV eviction has limitations because we cannot estimate which information will be important in the future. Therefore, we should store every important piece of knowledge somewhere in memory. During the storage of the KV cache, we can utilize a partial KV cache eviction strategy.
We believe KV cache offloading is the future direction to tackle the memory limitation of the atten- tion mechanism, as we proposed in the main section.
E.8 DISCUSSION ABOUT SPECULATIVE DECODING
We think that HiP can cooperate with many other speculative decoding strategies orthogo- nal (Leviathan al., 2023; Miao al., 2024; Fu et al., 2024; Cai al.. 2024) because they are working with the output of LLM. which is logits. Also, the speculative decoding method tries to decode multiple queries simultaneously to verify the speculative generation candidates. This char- acteristic of speculative decoding will take advantage of additional speedup with the large batches in HiP.
E.9 REMAINING CHALLENGES IN HIP AND POTENTIAL SOLUTIONS
Although HiP successfully replaces the existing vanilla attention, there is room for improvement as follows:
While HiP outperforms the baselines and performs similarly with Flash Attention in long context evaluation, HiP still underperforms Flash Attention with smaller & (lower compute budget).
As shown in Appendix E.5, HiP uses few layers (la) of quadratic dense attention, which could lead to higher cost as the context length increases.
Since HiP enforces every row in the attention matrix to have the same sparsity, this is not optimal handle dynamic sparsity of attention (Ribar al., 2013; Lee et al., 2023).
As shown in Figure 2, because HiP discards the bottom chunks for unselected branches, it is impossi- ble to select tokens from the actual top-k set if they happen to be within those discarded chunks. We refer this 'branch early termination' in HiP, and addressing this issue could help resolve above improvement points. Therefore, we propose two possible research directions that tackle the branch early termination problem while also enabling dynamic sparsity: (1) an ensemble hierarchical top-k approximation and (2) an improved tree traverse strategy.
First, for the ensemble hierarchical top-k approximation, jecting randomness into branching decisions during a specific iteration and create a final ensemble mask by aggregating indices from these masks. The ensemble method demonstrates that applying different branching in given iteration can enhance inference accuracy and indicates the potential to replace the dense layers la in HiP with comparable performance, as shown in Appendix C.5. Moreover, Figure 24 illustrates how the ensemble method enables dynamic sparsity across layers and heads, addressing the limitation of uniform sparsity in HiP.
generate multiple HiP masks by in-
Second, we could explore more diverse traversal methods, rather than strictly relying binary branching and selecting the top-k tokens at each iteration. In Appendix C.5, we examine the ef- fectiveness of applying ensemble techniques to HiP masks with slight randomness. However, this approach incurs additional computational costs due to the oversampling of the mask, which can be quite expensive. Therefore, achieve similar effect, could diversify the tree branching beyond
44

---

Published as a conference paper at ICLR 2025
the binary structure, similar to an n-beam search. Another potential solution specifically tailored to HiP is to apply bralciiban in a certain iteration and oversample the chunks in subsequent iterations, maintaining multiple paths until the end. By doing so, the final iteration would include more than k candidates, resolving the branch early termination issue and allowing us to decide how many tokens to select for dynamic sparsity across the layers.
E. UNIQUE GPU RESOURCE DEMAND PATTERN OF HIP COMPARED TO FLASH ATTENTION
Our novel HiP attention is quite different from any other attention implementations. We do not heavily rely ALUs (floating operations) like Flash Attention. Also, we do not heavily rely on memory bandwidth like previous sparse attention methods like H₂O (it has to store attention scores in global memory). The HiP relies on thread resources because of the top-k operator in between HiP iterations. Moreover, we have highly complex algorithms compared to traditional GPU applications like graphics shaders. Due to this highly complex algorithm, we heavily rely on thread resources, even if we are heavily optimized with MMU for attention score sampling.
Thanks to the reduced complexity of HiP, O log we are winning every configuration with long context prefill and decoding. Since the decoding phase is highly memory-dependent for Flash Attention, we also always win in most practical context lengths. However, we sometimes lose if Flash Attention is way too much faster because of the trillions of floating operation specifications of high-end server-grade GPUs. Moreover, Flash Attention 2 and utilize special floating point valuation resources in GPU, especially on H100; FA3 is way too much faster in another setting. Therefore, we are starting to lose in short context (T=32kk to FA2 and FA3 because of the speedup of Flash Attention on H100.
This phenomenon is disappointing to us. Therefore, we try to investigate why HiP in H100 is so slower than others, even compared to consumer-grade GPUs such RTX 4090. We think the high demand for CUDA thread resources is due to our internal sorting algorithm. Since we must remain top-k blocks in every iteration, we must perform O log k) cost sorting operation Ollog T)) times. Therefore, as k grows, we are staving to allocate worker thread for score comparison.
Table 17: Prefill Latency Speedup of HiP by Removing Sorting on RTX 4090 and H100 with different h and T. The time unit is milliseconds.
Device
4090
H1000
Context Length
32k
128k w/ Scot
Profill Latency
wio Sect
w/ Sort
Speeclup
wlo Sort
w/ Sort
Speedup
wlo Sort
w/ Sort
Speedup
wlo Sort
Speedup
Flash Attention
56.0
855.2
24.05
430.79
HiP k=512
16.48
19.33
1.173
85.62
1.204
19
86.65
108
1.251
HiP
26.96
33.70
1,250
150.90
190.98
1.266
27.56
36.
1.317
153.41
210.78
1.374
HiP k=2018
44.12
62.27
1.11
263.60
386.97
1.468
43.74
69.3
1,590
260.22
434.30
1.669
Table 18: Technical Specifications of 4090 and H100. We put comparison of prefill latency
compared to RTX 4090.
Rel. CUDA Core
Rel. TFLOPs
Rel. Men. Bandwidth Rel. Clock Speed
Rel. HiP Speedup
Rel. FA2 Specdup
RTX 4090 11100
1.00
1.00
1.00
1.00
1.00
1.00
1.00
3.66
3.32
0.71
0.89
1.99
We want to show that the elimination of the sorting operation will speed up our top-k estimation. To do so, we replace sorting with an identity function. So, in this version, we always select the first half blocks to pass the next HiP iteration. As shown in Table 17, eliminating sorting speed up our HiP significantly. In 4090, we could observe 46.8% speedup, and in H100, we could observe 66.9%. We can see the high relation between (CUDA cores + relative clock speed) and HiP speed as shown in Table 18. So, we will try to investigate removing the sorting and replacing it with some approximations for more practicality of HiP.
This characteristic is quite good for cost-efffectiveness. Nvidia does not usually reduce CUDA cores on sesumerr-rrdee GPUs; therefore, we could achieve the same speed as H100 while reducing some great cost-effective
GPU costs more than ten times. Even in server-grade GPUs, there alternatives. For example, L40s has more ALU than 4090 and the same amount of CUDA core. Therefore, L400s will offer A A100-level linear layer computation while offering H100-level attention,
45

---

Published as a conference paper at ICLR 2025
which is way after than flash attention L40s. In the L40s, flash attention will be extremely slow, like in A100 and 4090, because they have similar FLOPs due to the price tag. We wanted to test L40s during submission, but unfortunately, could not find any possible option to get L40s.
The lower-grade GPUs often struggle with the small size of VRAM. However, the tiny memory of lower-grade GPU is not problem with method due the powerful KV cache offloading feature without decoding throughput degradation. We have already shown that we can serve 64K context length with a single RTX 4090 card, and if you put 8 of them together, then can serve around 512K context length with high decoding throughput. For example, the tinygrad company offers 8x 4090 workstations with only 40,000S (tiny grad, 2024) (we are not them, just for clarification). The price is almost similar to single H100 card, but you can serve 512K context length with more than twice TFLOPs! This means that if you have two nodes of that machine, you can actually run Google Gemini class (Google, 2024) long context L.M in the home. And if the tensor parallelism is linearly scaled with two nodes, you can decode I 1,527 tokens with 64k context length. Since our method is almost logarithm scale with context length during decoding, tokens per second with a one million context length. So, we are really excited to introduce our KV cache offloading feature with HiP in many practical aspects.
can expect to decode around 1K
F POTENTIAL NEGATIVE SOCIAL IMPACT
In this paper, do perform careful investigation LLM alignment performance with HiP. There is the potential that HiP might break the LLM safety guard. However, as far as we observed during the experiment, the HiP could preserve most of the behavior of trained LLM. Furthermore, we could always adopt the third-party LLM safety guard model such as LLaMA Guard (Inan et al., 2023).
46


Reviews:
[{"review_id": "cjBjirJ8Fn", "reviewer": "Reviewer_RsGC", "summary": "This paper proposes Hierarchically Pruned Attention (HiP) to reduce the time complexity of attention to O(T logT) and space complexity to O(T) where T is the sequence length. By exploiting the continuity of token sequence (tokens close together tend to have similar scores), for each query, the HiP use a tree-search like algorithm to approximately search the top k key tokens that yield large attention weights. Further, the author developed a KV cache offloading scheme to offload KV cache to host memory and reduce the GPU memory usage.", "questions": "1. It would be easier to understand to have a figure illustration showing the tree search (or improve figure 2, the figure 2 step 1 is a bit confusing) for section 3.1. \n2. The algorithm divides the sequence into k segments, and one token will be selected in each segment. Is it correct? If so, then the top k tokens must be distributed in the sequence, and cannot be concentrated on certain regions. Why do you make this design choice? \n3. I am aware of some literatures that also use iterative refinement to dynamically calculate attention for efficiency. Have the authors tried to compare to these literatures?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method uses iterative refinement to dynamically and approximately locate top k tokens, which is interesting. \n2. HiP shows promising efficiency improvement with only small performance drop. \n3. The appendix provides a lot of ablation study to study the behavior of the proposed method. \n4. The method is training free.", "weaknesses": "see question section", "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": "1YyC4qmBza", "reviewer": "Reviewer_T2wK", "summary": "The paper introduces Hierarchically Pruned Attention (HiP) which reduces time and space complexity of the attention mechanism. HiP exploits \"attention locality\" to estimate top-k key tokens (w/ theoretical justifications as insights) for a given query and does so in a hardware aware way. Moreover, the paper introduces a further KV cache offloading steps which reduces space complexity further.", "questions": "* Gemma 2 uses this mix between sliding window and normal attention. it would be great to understand if there are any degradation on non-llama architectures.\n* The method does show that MMLU is not degraded when using this method. It would be interesting to see this for a broader set of metrics if possible.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* The method being training-free means it can be used as a drop-in to already trained models.\n* The paper does careful complexity analysis of its claims but strikes a balance on introducing information in a way to aid presentation (thinking of the informal theorem) while still being rigorous later. \n* It is extremely valuable to have code examples and implementation released.\n* The paper is extremely comprehensive when understanding its metrics across different hardware and comparing with the many version of flash attention.", "weaknesses": "* It would be great to have the long context benchmarks also for different models -- gemma and mistral are both open source and around similar sizes.", "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": "rj8qen2ghx", "reviewer": "Reviewer_NLBm", "summary": "This paper proposes a novel, training-free attention mechanism called Hierarchically Pruned Attention (HiP) to accelerate the serving of pre-trained Transformer-based large language models (LLMs) for long-context tasks. HiP addresses the computational challenges posed by the quadratic time and space complexity of standard attention mechanisms in handling long sequences.", "questions": "1. How does the choice of chunk sizes in HiP's hierarchical pruning affect the trade-off between accuracy and efficiency?\n2. What strategies can be employed to further optimize KV cache offloading, such as using different memory tiers or compression techniques?\n3. How well does HiP integrate with other efficiency techniques, such as quantization or model pruning, to further improve serving efficiency?\n4. Can HiP be effectively applied to other Transformer architectures beyond the specific LLM model used in the paper?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. HiP's training-free nature eliminates the need for costly retraining of large LLMs, making it readily applicable to existing models.\n2. HiP's hierarchical pruning significantly reduces the computational complexity of attention from quadratic to log-linear, enabling efficient handling of long sequences.\n3. KV cache offloading effectively addresses the memory limitations of GPUs, allowing HiP to scale to much longer context lengths.\n4. The paper provides thorough experimental results on diverse benchmarks, showcasing HiP's effectiveness in terms of speedup, performance preservation, and context length extension.\n5. The paper includes a theoretical analysis of HiP's hierarchical pruning algorithm, providing insights into its superior performance compared to random key selection.", "weaknesses": "1. HiP's effectiveness relies on the assumption of attention locality. While this assumption generally holds, there might be cases where it's violated, potentially impacting performance.\n2. HiP enforces the same sparsity across all rows of the attention matrix. Exploring mechanisms for dynamic sparsity, where sparsity varies based on the input, could further enhance efficiency and performance.\n3. The implementation and optimization of HiP, particularly the KV cache offloading, are tailored for specific hardware platforms (e.g., RTX 4090). Performance and optimal configurations might vary across different hardware accelerators.\n4. The paper doesn't specifically address potential LLM alignment issues that might arise from applying HiP. Further investigation is needed to ensure HiP's safety and robustness in practical deployments.", "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": "xB2qXvQDys", "reviewer": "Reviewer_KprR", "summary": "The paper presents HiP (Hierarchically Pruned Attention), a novel approach aimed at reducing the time and space complexity of the attention mechanism in Large Language Models (LLMs). HiP leverages the observation that tokens close together tend to have similar attention scores to estimate the top-k key tokens for a given query on the fly. This results in sub-quadratic time complexity (O(T log T)) and linear space complexity (O(T)), where T is the sequence length.  Experimental results show that HiP significantly reduces both prefill and decoding latencies while maintaining high performance on benchmarks such as LongBench.", "questions": "See the weakness for details.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. HiP does not require retraining, making it easy to apply to pre-trained models. In addition, by reducing the time complexity to O(T log T) and space complexity to O(T), HiP enables the use of longer context lengths without the associated quadratic cost.\n2. Through KV-cache offloading, HiP optimizes GPU memory usage, which is especially beneficial for large models.\n3. HiP can dynamically adjust to different sequence lengths, making it suitable for a variety of tasks that involve long contexts.", "weaknesses": "1. The effectiveness of HiP is contingent upon the presence of \"attention localities,\" which might vary across different LLM architectures or tasks. The robustness of this operation deserves more discussions.\n2. While HiP reduces latency in the decoding phase, the iterative pruning process might introduce overhead in the initial stages. How does the proposed method balance these two parts.\n3. The authors have already shown improvements for certain sequence lengths, while it is expected to thoroughly explore how HiP scales to extremely long sequences or very large models. The current configuration on the sequence length (128k) and the model size (8B) are relatively small for a paper studies on the efficient decoding. Alternatively speaking, the scalability of this method deserves more experiments to support.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)