RazorAttention: Efficient KV Cache Compression Through Retrieval Heads
Authors: Hanlin Tang
Yang Lin
Jing Lin
Qingsen Han
Danning Ke
Shikuan Hong
Yiwu Yao
Gongyi Wang
Venue: ICLR 2025 Poster (2025)
Keywords: LLMs
KV cache compression
LLM inference acceleration
Primary area: generative models
TL;DR: A novel KV cache compression algorithm that achieves 3X reduction under general cases, compatible with FlashAttention

Abstract: The memory and computational demands of Key-Value (KV) cache present significant challenges for deploying long-context language models. Previous approaches attempt to mitigate this issue by selectively dropping tokens, which irreversibly erases critical information that might be needed for future queries. In this paper, we propose a novel compression technique for KV cache that preserves all token information. Our investigation reveals that: i) Most attention heads primarily focus on the local context; ii) Only a few heads, denoted as retrieval heads, can essentially pay attention to all input tokens. These key observations motivate us to use separate caching strategy for attention heads.Therefore, we propose RazorAttention, a training-free KV cache compression algorithm, which maintains a full cache for these crucial retrieval heads and discards the remote tokens in non-retrieval heads. Furthermore, we introduce a novel mechanism involving a “compensation token” to further recover the information in the dropped tokens. Extensive evaluations across a diverse set of large language models (LLMs) demonstrate that RazorAttention achieves a reduction in KV cache size by over 70% without noticeable impacts on performance. Additionally, RazorAttention is compatible with FlashAttention, rendering it an efficient and plug-and-play solution that enhances LLM inference efficiency without overhead or retraining of the original model.

Published as a conference paper at ICLR 2025
RAZORAFTENTION: EFFICIENT KV CACHE COMPRES- SION THROUGH RETRIEVAL HEADS
Hanlin Tang", Yang Lin', Jing Lin', Qingsen Han', Danning Ke', Shikuan Hong', Yiwu Yao!,
and Gongyi Wang!
Huawei Technologies Co., Ltd
ABSTRACT
The memory and computational demands of Key-Value (KV) cache present sig- nificant challenges for deploying long-context language models. Previous ap- proaches attempt to mitigate this issue by selectively dropping tokens, which irre- versibly erases critical information that might be needed for future queries. In this paper, we propose a novel compression technique for KV cache that preserves all token information. Our investigation reveals that: i) Most attention heads primar- ily focus on the local context; ii) Only a few heads, denoted as retrieval heads, essentially pay attention to all input tokens. These key observations motivate us to use separate caching strategy for attention heads. Therefore, we propose RazorAttention, training-free KV cache compression algorithm, which main- tains a full cache for these crucial retrieval heads and discards the remote tokens in non-retrieval heads. Furthermore, we introduce novel mechanism involv- ing "compensation token" to further recover the information in the dropped to- kens. Extensive evaluations across diverse set of large language models (L.LMs) demonstrate that RazorAttention achieves reduction in KV cache size by over 70% without noticeable impacts on performance. Additionally, RazorAttention is compatible with FlashAttentinn, rendering it an efficient and alug-and-pla solu- tion that enhances LLM inference efficiency without overhead or retraining of the original model.
INTRODUCTION
Long-context large language models (LLMs) have significantly advanced capabilities in natural lan- guage processing across diverse tasks. However, the growth of the Key-Value (KV) cache under increasing imput length has become the major bottleneck for deployment. There are been plenty of previous work designed alleviate this problem by compressing the KV cache size, including quantization (Sheng et al. 2023 Zhao et al. 2024 Lin et al. 2024), token-dropping (Zhang et al.
2023
Xiao et al. 2024), local attention (Jiang al. 2023 Child al. 2019), etc.
Original model
RazorAttention
1.
0.6
Dapth Percent
Depth Percant
22.0
32.0
678
0.460
070
6.40
89
62
D002
3003 31400 Sequence Length
36300
44500 52000 60163
400
1200
31400
2002 3650 40033 52000
Sequence Length
Figure 1: RazorAttention achieves comparable performance to the original model, even with 70% KV cache compressed. To demonstrate this, we tested Llama2-13BB 1amma2-13B-6ae 13B-64K (Fu et al. 2024) on the Needle A Haystack benchmark (gkamradt] (023).
"Corresponding author: Aanghl199464 @anghll994@@gmii.com

---

Published as a conference paper at ICLR 2025
Input context: "DOD's MILCON appropriations are used fund the acquisition, construction, installation...
Mary's favorite number 34251 Bob's favorite number is 7690... reviewing project cost estimates."
01: 'hhat is Mary's 's favourite number?"
Original model: "Mary's ' favorite number is 34251."
H20: "Mary's favorite number is not explicitly mentioned in the text provided." x
V "Mary's favorite number 34251.'
RazorA teetton::
"Mary's favorite numberis 34251
Q2: Whtt Bob's favourite number?" Original model: "Rob's favorite number 7690." H20: !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
SnapKV: "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
RazorA ttention: 'Bob's favorite number is 7690."
Figure 2: importance-based token-dropping methods cannot work when querying the less relevant information to the main theme. Here, we use an 8K document from LongBench (Bai et al., 2023b) and add two sentences that are not relevant to the main theme. In this case, H20 discards tokens that are less relevant to the main theme, leading to failures in both QI and Q2. SnapK discards tokens based on the first query, making effective for QI but failing in subsequent queries like Q2. Only RazorAttention successfully outputs the exact information from the lengthy input even when we compress 70% of the KV cache.
One major direction for KV cache compression is to directly drop tokens deemed unimportant so far (Zhang et al. 2023 Xiao et al. 2024 Liu et al. [003bb Li et al. 2024). These methods inherently assume that tokens considered unimportant will not be needed in future queries, which does not hold in practical scenarios. For instance, user might request information that is not directly aligned with the main theme of the processed text, or engage in a multi-round conversation querying different segments from the context. In these cases, the importance-baset token-dropping methods can lead to significant performance degradation since the actual information required by the query might be discarded if considered unimportant (see our example on Qwenl.5-7B-CC (Bai et al. 2023a) in Figure D) This leads us to pose critical question:
"Can we find a way to reduce the KV cache size without losing semantic information?"
In this work we address this problem from a novel perspective. Our investigation reveals that there exists "retrieve and process" mechanism in LLMs when processing long context. More specifi- cally, LLMs can accurately recall the queried information from a lengthy input through certain group of attention heads, which we denote as "retrieval heads" (see Section for definition). These heads capable of concentrating most of their attention weights the relevant information (w.r.t. the queries) and increasing the output probability for those words. Another important finding is that non-retrieval heads primarily focus local context the attention sink (Xiao al. 2024), which means these heads cannot effectively utilize all the semantic information from the input. Based on these important findings, we hypothesize that LLM runs the reasoning procedure on "retrieve and process" basis. That says, the model first uses the retrieval heads to gather relevant information, and then non-retrieval heads to process the retrieved information and generate the final response. This motivates us to design separate caching strategies for different heads: For retrieval heads, we keep the KV cache unaltered; for the rest heads, we only cache recent tokens and attention sinks.
Beyond this, we notice that there still exists a certain accuracy gap when we directly discard all the remote tokens in the non-retrieval heads. Therefore for these non-retrieval heads, we designed a "compensation token" for compressing the dropped cache into one token, and proved that the accuracy degradation due to the truncated KV cache gets further improved with this compensa- tion token. With retrieval heads and compensation tokens, we prove that our algorithm, namely RazorAttemtion, can successfully compress 70% of the KV cache without noticeable performance degradation as illustrated in Figure [ RazorAttention sequences, achieving nearly lossless precision in the Needle in a Haystack task after compressing the KV cache by 50%, shown in Figure
support the compression of 1024K
2

---

Published as a conference paper at ICLR 2025
Last but not least, previous importance-baser token-dropping methods cannot be combined with FlashAttention due to their reliance on the attention weights to compute the importance score, mak- ing them impracticable for implementation since FlashAttention is one of the most important com- ponents in long-context inference. RazorAttention addresses this problem since it does not use the attention map as the metrie. The head-wise pruning criterion is totally compatible with FlashAtten- tion, and the computation overhead of the compensation token is negligible. Therefore RazorAtten- tion could achieve substantial inference speedup when compared to previous methods.
To the best of our knowledge, RazorAttention is the first training-free token reduction algorithm that achieves a nearly lossless 3X KV cache reduction. We evaluated RazorAttention on models including Qwen (Bai et al. 2023a), Llama-2 (Touvron et al. (202)), Llama-3 (AI@ Metal (2024) and Baichuan (Baichuan (023) on long-context tasks to prove its effectiveness. Our contribution can be
summarized as follows:
We systematically analyze the attention dynamic of Transformers under lengthy inputs. Our work reveals that only few retrieval heads can essentially recall information from the whole input while the rest heads mainly focus on the local context.
We introduce novel algorithm, namely RazorAttention, that is capable of reducing the KV cache size by 70% under minimal impact performance for contexts ranging from 8K to 100K tokens. We designed an accurate and dlta--free metric for allocating all the retrieval heads, together with error compensation strategy for compensating the information loss due to the truncated KV cache.
RazorAttention introduces negligible overhead in compression and compatible with FlashAttention, rendering it an efficient and plug-and-play solution that enhances LLM in- ference efficiency without training or significant overhead. Extensive experiments demon- strate that RazorAttention can be effectively applied to various models and tasks.
2 RELATED WORK
As the sequence length increases, the memory consumption of KV cache rapidly expands, poten- tially surpassing the size of the model parameters themselves. This leads to an urgent need for KV cache compression, particularly in scenarios with limited GPU memory. One direction is non- Transformer architecture design, such as Mamba (Gu & Dao 2024), Mamba2 (Dao & Gu 2024) Infini-Transfome Transformer (Munkhdalai al. 2024), RWKV (Peng al. [023)and Griffin (De et al. 2024 However, in this paper we focus on KV cache reduction for typical Transformers, which is the most widely used model structure. Below we introduce several approaches for KV cache compression.
Quantization Quantization is a classic yet effective approach to neural network compression. In the field of LLM Quantization, while the outlier challenge attracts great attention (Xiao al. 2023 Wei et al.) [202]] (023) to tackle, the application of which on KV cache is often seen as a by -y-orudtct of activation quantization. Nevertheless, there are several noteworthy works demonstrating the value of KV cache quantization. FlexGen, Atom and QServe (Sheng et al. 2023 Zhao et al.) 2024]]
Lin et al., (024) carefully designed guantization pipelines that utilize KV cache compression to boost the overall inference throughput. KVQuant (Hooper et al.
2024)
integrates several techniques
to minimize KV quantization error and KIVI (Zirui Liu et al.
2023 pushed the limit towards 2-
bits. Besides the post-training methods, LLM-QAI (Liu et al.]
2023a
offers data-free distillation
process that further recovers the performance of the model.
Token-dropping Token- loken-dropping methods assume that not all key-value pairs are essential in self- attention computations, so memory usage can be saved by identifying and removing unimportant KV cache. StreamingLLMM (Xiao af., (024) utilizes sliding window technology, preserving only the KV pairs of attention sink tokens and those within the sliding window, thereby reducing memory footprint and stabilizing model performance. H2O (Zhang et al. 2023) is one the pioneers that use the attention scores to evaluate the importance of each token, followed by eviction strategy that greedily selects cache with higher scores. Scissorhands (Liu et al. 2023b and one of the latest work SnapKV (Li et al. 2024)
similar ideas by narrowing the computation range to consider attention scores related to recent information. Built on that, PyramidKV V and PyramidInfer (Cai. et al. 2024 Yang et al. 2024) analyze the attention concentration patterns and further reduce KV
3

---

Published as a conference paper at ICLR 2025
cache in later layers. Moreover, research efforts have been made to understand KV cache from different perspectives: FastGen (Ge et al. 2024) paid attention to special tokens and punctuation, SubGen (Zndich et al., 2024) investigated the clusterability of key embedding and CORM (Dai
et al. 2024) discovered strong correlation amongst tokens of near neighbors.
Non-MHA Attention Another category focuses reducing KV cache by sharing cache across attention heads. MQA (Shazeer (019) aggressively uses single KV head for all heads, whereas GQA (inslie et al. 2023) suggests an intermediate number of heads to balance the trade-off be- tween inference speed and output quality. Furthermore, MLA (DeepSeek-AL et al. (2024) presents a novel caching method by low-ranking KV cache of all heads into single latent space.
Our algorithm is motivated by the idea from Olsson et al. 2022), where the authors noticed that there are certain groups of attention heads, denoted as the induction heads, that can effectively recall the queried information from the input. Recent study (Wu et al. (2024) also validated this property under extended inputs. This is the first work that proposes a head-wise pruning criterion for KV cache compression based on the interpretability of the attention mechanism.
3 METHODOLOGG
In this section, we introduce the key components of RazorAttention. We firstly apply RazorAtten- tion models using ALiBi (Press al. 2022) positional embedding (denoted ALiBi models) to provide an intuitive understanding of the retrieval and non-retrieval heads. Afterwards, we demon- strate that models using RoPE (Su et al., 2023) positional embedding (denoted as RoPE models) also exhibit this crucial characteristic, which reveal that KV cache within RoPE models can also be efficiently compressed under minimal loss of accuracy.
3.1 RAZORATTENTION FOR ALIBI MODELS
For ALiBi models, its h-th attention head computes the attention score according to
(I) where qm is the query tensor at the m-th position, k, is the key tensor at the n-th position, 4a is the head-specific slope, Sm- (q: k) is the attention score. Notice that (m 2 n) is guaranteed by the causality of attention.
Sm-n (q: k) km l1( n))
In the scenario where (m - n) significantly dominates km the attention between qu and kn would decay zero, meaning that the contribution of any tokens positioned further than becomes negligible for the output at position m. The following theorem formalizes this observation.
Theorem Given an attention head that calculates the attention score as per equation for any E (0, 1), the attention weight from qm to kn can be upper bounded by:
exp (5m-n (q; k))
Attnm (q: k)
Σm exp (Sm (q: kq: Se
Vn < m Co,
LA = 21IWQ, WK, K. lla (11 + [lb log(c)
(2)
li
Here WQn and WK. are the query and key matrices of the h-th attention head, y and b are the weight and bias for the LayerNorm layer before attention (b = 0 for RMSNorm (Zhang & Sennrich 2019)) and ||2 denotes the 1z-norm of the matrix. LA can be viewed as the vision scope of the head. The detailed proof can be found in Appendix A.
Theorem equation indicates that when the distance between Jm and k exceeds Co, the attention weight between these two tokens falls below When e is sufficiently small (e.g.. 0.1%) remote tokens impose minimal influence on the final output and can thus be discarded. Building on this principle, ALiBi models dynamically adjust the KV cache size for each head. We first compute the effective attention scope L,, and keep only the recent L, tokens in the KV cache, since any token further than LA impose attention weight more than e, can safely discard them for compression. Therefore, for ALiBi models, the retrieval heads are the ones with a larger Ln, while the non-retrieval heads has smaller attention vision T.p-
4

---

Published as a conference paper at ICLR 2025
3.2 RAZORATTENTION FOR RoPE MODELS
For RoPE models, each attention head computes the attention score according to kn R
Sm-n (q: k) 4m qm R.q.
(3)
Where Om and kn are the query and key states after the rotary transformation, and Rm and Rr are the rotary matrices at positions and n (see(Su al. (2023) for details). Although RoPE embedding does not inherently suggest long-range decaying attention, our empirical findings indicate that only about 15% of the heads, which we term "retrieval heads," are capable of effectively utilizing long- range information to ensure the model's reasoning performance, while the remaining heads primarily focus on local context. As shown in Table 1, when all heads are set to focus only on local context (Protect 0% heads), the accuracy significantly decreases by an average of 7.3%; when 15% of the heads are randomly selected to attend to long-range information (Protect 15% heads Random) while the rest focus on local context, the accuracy significantly decreases by an average of 7. %%; however, when retrieval heads are selected to focus on long-rang information (Protect 15% heads Retrieval) while the remaining heads focus on local context, the accuracy is comparable the baseline, with only a decline of 0.3%.
Protection heads
MultiFieldQAAen
Hotpotaqa
2Wikimqa
Baseline Protect 0% heads
46.94%
50.96%
36.36%
37.36%
42.36%
31.33%
Protect 15% heads (Random)
38.33%
42.46%
31.97%
Protect 15% heads (Retrieval)
46.66%
50.49%
36.12%
Table I: We protected the KV cache within different groups of attention heads while keeping only the recent 4K tokens in the rest with Qwenl.5-7B. The results indicate that both unprotected heads and randomly protected heads yield similar poor performance, while only protecting the retrieval heads effectively retains most of the LLM's performance. This clearly shows that most attention heads rely solely on local context, and only retrieval heads can fully utilize all contextual informa- tion.
Based on the findings above, we directly decrease the KV cache for all non-retrieval heads. The performance of the model is mostly retained as shown in Table Tooweer, notable accuracy gap remains, indicating that some information is still being lost. Moreover, the test result Needle a Haystack shows clear performance degradation even when we protect the KV cache re- trieval heads (see our ablation result in Figure 8 To further improve performance, designed lightweight and effective way to compress the information in the dropped token into "compensa- tion token". The compensation token is defined
Σ ME(D)
Km:
1
Σ Un-
Na
Nd mE(D)
Here k,, are the compensation tokens for the dropped KV cache, (D) contains the indices of the dropped tokens and Na is the number of the dropped tokens. Afterward, we discard the dropped tokens and augment the KV cache with the compensation token & and è where {K, V3 are the KV cache of the remaining token after rotary transformation. Denoting the compressed KV cache as {K, , k. and (V, 6}. the attention output of the current token follows
exp (mmkk)) 8 + Eng(ry exp (qmk)) un
Attn(gm: {K, [K, (V,
(5)
Na exp (mmkT) Eng(Dy exp (qmkk)
In Figure 3(a) we provide an illustrative example of RazorAttention for RoPE models. With com- pensation tokens, the accuracy is further improved, making RazorAttention almost lossless dropping 70% of the KV cache in the non-retrieval heads. Below we introduce how we determine the retrieval heads group.
5

---

Published as a conference paper at ICLR 2025
[arkk]
Idropped cache)
troling cachi)
Non-retrieval
(compressed cache)
Retrieval head
(ful cache)
Induction head:
(a)
(b)
Figure 3: In Figure 3(a) we present the illustration of how RazorAttention compress the KV cache. For retrieval heads, we maintain full cache for retaining all the tokens' information. For non- retrieval heads, we directly discard remote tokens and compress the discarded tokens into a com- pensation token whose KV cache is denoted {k, o}. In Figure 3(b) provide illustration example the echo head and induction head. The current token "B" and the generated token is "C". In this case, the echo head would mainly attend to token "B" while the induction head mainly attend to token "C" in previous context.
pppre-parrmeter Buffer length
Settings max(4000, N/5) top 14%
Induction head protection
Echo head protection
top 1 % 4
Sink token num
Table General parrame- settings for experiments in the paper, which leads to 3. 125x com- pression of KV cache under long context input.
3.3 IDENTIFICATION RETRIEVAL HEADS
For ALiBi models, the attention scope can be directly determined via equation 2 and KV cache can dropped accordingly. However, for RoPE models, the retrieval heads need to be identified in more sophisticated way. Our investigation reveals that two groups of heads are essential in processing long context, so both of them should be included as retrieval heads as stated below.
Echo head: The head tends attends back previous token (referred as echo token) identical to the current token.
The head tends to attend to the previous token (namely induction token) that immediately succeeded by the current token. Basically it attends the coming token that also exists in previous context.
Induction head:
In Figure 3(b) we present an illustrative example explaining the echo heads and induction heads. In order identify the retrieval heads, we generate (for example, K 2500) random tokens, repeat these tokens 4 times, and then use it as the input of the model. This design minimizes semantic dependencies among tokens, thereby allowing clearer observation of the behavior of echo and induction heads.
Subsequently, we calculated the echo score (attention weight to the echo token) and induction score (attention weight to the induction token) of all words across all heads. The selection of retrieval heads involves the top-14% attention heads with the highest induction score and top- of attention heads with the highest echo score (see Table 2 Notice that although we only much fewer echo heads than retrieval heads, our investigation indicates that both heads are crucial for the retrieving performance for LLMs (see Section 4.2 for ablation results).
With the retrieval heads being identified, we hereby introduce RazorAttention for RoPE Models in Al Algorithm

---

Published as a conference paper at ICLR 2025
Algorithm 1 RazorAttention for RoPE Models
Input:
Non-retrieval headset {H), original KV cache (after rotary transformation) (K, V), com- pression ratio C, compression threshold So, sink token num No-
1: for non-retrieval head h E } do
2: Compute the buffer length La = max (So, N here N is the number tokens in the head.
3: Keeping only the recent T,,, tokens near output and first No sink tokens, discarding the re- maining tokens and compress them into a compensation token according to equation
4: end for
5: Non-retrieval heads compute attention according to equation while retrieval heads follow the
original attention.
Output: Generated output tokens.
4 EXPERIMENTS
A variety of recent-released LLMs are selected to validate our proposals, including Owen Bai al. 2023a), Llama2 (Touvron et al. 2023). Llama3 (etaMeta, 2024) and Baichuan (Baichuan, The selected models are evaluated on Longbench (Bai et al. 2023b) and Needle In A Haystack (gkamradi) 2023) to demonstrate their capabilities long-context circumstances. The experiments are conducted on NVIDIA GeForce RTX 4090 (24GB). We will first validate the ef- fectiveness of our proposal on various tasks, followed by the ablation study each component in RazorAttention with the syper-parameters as in Table 2 We use H2O (Zhang et al. (023) and StreamingLLM (Xiao et al. (024) for compari- son. Notice that do not include SnapKV (Li et al. 2024) as the baseline because assumes that the query is known before compression, which does not hold in general cases or in a multi-round conversation where the user might query different information from the context (as discussed in Section
2023
algorithm design. Unless explicitly stated,
4.1 ACCURACY EVALUATION
4.1.1 LONGBENCH EVALUATION
In Table 3 present the results of different algorithms LongBench (Bai al. 2023b). which provides comprehensive assessment to evaluate long-context related abilities of LLMs. We use Qwenl.5-7B and Qwenl.5-72B for testing since they are RoPE models with a context length of 32K. We also include Llama3-8B to validate the performance of RazorAttention on GQA models. The compression settings are shown in Table 2 fectiveness of RazorAttention on ALiBi models. is set to 0.001, the contribution of the current token in attention map less than 0.001 under the Alibi encoding. The distance position calculated according to formula (2). It can be seen that RazorAttention achieved superior perfor- mance across all models compared to StreamingLLM and H2O. The compelling outcomes indicate that
We choose Baichuan2-138 to demonstrate the ef-
can achieve comparable performance the uncompressed baseline, even under 3X compression ratio.
Moreover, test Llama3-88 8B- Instructt GQA instance where every 4 attention heads share sin- gle set of KV cache. Hence, we consider the attention heads in a group as all retrieval if one or more heads satisfy inductive echoing property. The results in Table cleaarl prove that RazorAttention still work for GQA models.
4.1.2 NEEDLE IN A HAYSTACK EVALUATION
In Figure we present the results on Needle A Haystack. We use Llama2-7B-806 from Fu et al. 2024) since the context length of this model is 80K. Unlike H2O, whose performance is severely degraded under long inputs, RazorAttention can still accurately recall the queried information. This strong evidence proving that RazorAttention the original context, while importance-batt based methods inevitably discard information that might be useful in future queries. The compression settings shown in Table 2
retain all the semantic information within

---

Published as a conference paper at ICLR 2025
Original model
HacocMiaationn
H20
Sento
Figure 4: Performance comparison of RazorAttention and other compression algorithms on Llama2- 7b-80K, Needle In A Haystack. Notice that H20 is incompatible with FlashAttention so we get OOM errors when tested on longer sequences, and its performance has already become unusable this case.
LLMs
NAYQA
Queper
Mi--0
HelporDA
SWOAMGA
Moniqpio
GurRoport
QNi
MuitiNines
VCSCM
TREE.
LISTT
Average
Oheen 5-
12.58
43.10
06.94
50/46
36.36
27.86
15.24
24.12
13.91
17.04
83.77
36,03
Stramngli
H20
16.5
38.15
00.22
51.46
50.19
35.09
27.12
25.43
22.00
22.70
14.03
18.25
83.72
16.4
47.54
34.16
RA
16,33
43.1
46.46
61.18
50.49
36.1
28.79
26:06
22.99
33.96
13.83
20.87
83333
15.00
47.85
3587
Chu
KB
20.33
06.75
65.75
65.41
55.91
46.23
56.56
37.47
12.09
14.00
15.01
71.0
91.15
66.5
Stmamnng AM
9.57
25.35
15.26
14.05
15.51
13422
10133
4.11
3.70
21.1
3.74
45.0
45.72
311.
55.0
22.30
37.08
41.45
3461
7000
64.29
Limel--0
KV
21.84
52.04
45.07
$2.50 54
44.63
21.38
33.04
25.18
36.54
16.26
14.41
35.00
90.17
35.44
H20
21.14
14:1
47.13
27.33
21.31
21.85
16.36
21.3
14.53
86.17
30.26
37.99
RA
21.16
36.22
42.88
$1.03
44.17
26.50
22,03
26.56
23.86
25.83
15.09
$5.63
90.35
3486
30.15
13.00
34.58
52.5
$6.61
5.02
12:44
23.55
32:52
16.03
16.08
5.15
5.53
1.03
5.6
3.94
42.22
30.15
7332
35.42
16.27
120
12.81
20.80
43.74
49.55 54
37.02
91:91
13.54
24,8
18.96
33.51
15.11
67.41
55.25
76.86
$4.5
65.69
13.6
36.55
Table Performance comparison of RazorAttention and other compression algorithms across vari- ous LLMs on LongBench. Notice that the performance of Llama3- ----BB--Inttct on TREC and LSHT are not applicable (close to 0), hence we do not include their result Llama3-8B.
In Figure 5 verified that the RazorAttention algorithm support ultra-long sequences by testing the Needle in Haystack task using GLM-9B-IM (GLM et al. 2024). It is observed that RazorAttention achieves nearly lossless accuracy the 1024K sequence length. To ensure effective retrieval of needles in ultra-long sequences, we retained a higher number of retrieval heads (28% induction heads 2% echo heads).
In Figure 6 we compare the performance of RazorAttention with SnapKV (Li al. 2024). SnapKV competitive approach for KV cache compression, particularly effective when user queries are known in advance. We present SnapKV's performance results in both query-awart and query- agnostic settings. query-agnostic scenarios, we selected the observation window while disre- garding the query tokens, so the user query is not pre-defined, and SnapK performance deterio- rates significantly. This performance degradation attributed to SnapK reliance on predefined queries. In contrast, RazorAttention demonstrates superior resilience due to its training- free design and headwise sparse pattern, which minimize information loss across a variety of query types.
Original model
Razoráttention
Depth Percent
66
Depth Percant
6.8
42
62
94.
52
2204
350g
502k
2558
520k
630e
500k
5024
Sequenoe Length
Sequence Length
Figure 5: GLM-9B-IM the 1024K Needle In A Haystack task.
8

---

Published as a conference paper at ICLR 2025
g40
Sequana Lengi
Serquancs Length 60%
Sequan Lengh
(a)
(b)
(c)
Figure 6: Comparison of Attention and SnapKV on the Llama 3.1-70B model for Needle-in-a- Haystack task accuracy. SnapKV shows significant accuracy differences between query-aware and query-agnostic settings, as shown Figure 6(a) and Figure 6(a) while RazorAttention consistently
maintains higher accuracy, as shown in Figure 6(C)
Protection scheme
Score
1% Echo + 5% Induction Head
69.54%
1% Echo + 8% Induction Head
78.40%
1% Echo 11% Induction Head
84.55%
1% Echo + 14% Induction Head 86.59%
Baseline
87.05%
Table 4: Owenl.5-7B-Chr. using RazorAttention with different numbers of heads protected, tested on Needle A Haystack.
4.2 ABLATION STUDIES
Below we present the ablation results of RazorAttention, and prove that the algorithm design and configuration are optimally chosen to achieve higher compression ration with acceptable perfor- mance degradation.
4.2.1
IMPORTANCE OF ECHO HRADS
Although we only include 1% echo heads in RazorAttention, we notice that this group of heads quite essential in retrieving information under long context as shown in Figure 7 One possible explanation is that the induction heads depend the existence of echo heads discussed in Olsson cct al. (002)).
Induction Head + Echo Head
Induction Head Only
Depth Percent
Depth Percent
66.0
6.1
60.
0.5
0469
87.
19000
43000
52000
61600
60
100
53000 Sequence Length
43600
Sequence Length
Figure 7: Adding 1% of the echo heads can significantly enhances the retrieving performance of RazorAttention Llama2-7B-88
4.2.2 NUMBER OF INDUCTION HEADS
To determine the optimal number of induction heads use in RazorAttention, in Table we present the accuracy of RazorAttention under various numbers of induction heads. The results show that the accuracy improves continuously with increasing number of induction heads. We decide to include 14% of the induction heads order to achieve an optimal balance between the compression ratio and model performance.

---

Published as a conference paper at ICLR 2025
RazorAttention
RazorAttention Compensation Token
Deptt Percent
22.0
64
Depth Parcent
12.0
6.1
44.0
09.
64
00.0
64
4003
19720
25550
47000
52200
64200
75500
0030
12307 15600
20001
24003
20020
1200
Sequence Length
Sequence Length
Figure 8: The compensation token is critical for recovering the information loss introduced by the truncated KV cache.
4.2.3 IMPORTANCE OF THE COMPENSATION TOKEN
In Figure 8 RazorAttention. The compensation tokens successfully compressed most of the information from the dropped tokens,thereby, maintaining high accuracy even with significant KV cache reduction.
it clearly demonstrated that compensation tokens critical for the performance of
5 CONCLUSION
In this paper, propose RazorAttention, novel KV cache compression algorithm, which success- fully achieves a 3X compression ratio for models use RoPE or ALiBi embeddings. Unlike previous importance-basee token-dropping methods which inevitably discard semantic information, RazorAt- tention preserves all semantic information within retrieval heads. We demonstrate that remote tokens can be effectively compressed into compensation tokens within non-retrieval heads. Furthermore, our head-wise pruning criterion fully compatible with FlashAttention, making RazorAttention a plug-an huug-and-llay compression method that accelerates the inference of LLMs under extended context. Our experiments demonstrate that Razor Attention can achieve comparable performance with the original model and surpasses previous methods in both accuracy and efficiency.
6 LIMITATION
However, there are still certain limitations of our work. The first question is why attention heads in LLMs behave so differently and how retrieval heads operate under lengthy inputs. The second challenge lies in achieving higher compression ratio. Although we have successfully reduced the KV cache by 70%, we believe this number can be further improved. Moreover, although we have tested our algorithm meaning that we might need more or less retrieval heads under different cases. These topics are quite important and we will keep investigating them in the future work.
several models, the optimal configuration
other models might be different,
REFERENCES
AI@Met Meta. Llama model card. 2024.
URL https::/github.ccm cow.///ttttoo.mmmtttm.mmoo 1lama/
1aama3/blb//mm maroom/maing M/MODEL, CARD. mc
Joshua Ainslie, James Lee- Thorph Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points, 2023.
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Han, Fei Huang, Binyuan Hui, Luo Ji, Mei Li, Junyang Lin, Runji Lin, Dayiheng Liu, Gao Liu, Chengqiang Lu, Keming Lu, Jianxin Ma, Rui Men, Xingzhang Ren, Xuancheng Ren, Chuanqi Tan, Sinan Tan, Jianhong Tu, Peng Wang, Shijie Wang, Wei Wang, Shengguang Wu, Benfeng Xu, Jin Xu, An Yang. Hao Yang, Jian Yang, Shusheng Yang, Yang Yao, Bowen Yu, Hongyi Yuan, Zheng Yuan, Jianwei Zhang, Xingxuan Zhang, Yichang Zhang, Zhenru Zhang, Chang Zhou, Jingren Zhou, Xiaohuan Zhou, and Tianhang Zhu. Qwen technical report, 2023a.
10

---

Published as a conference paper at ICLR 2025
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, 2023b.
Baichuan. Baichuan 2: Open large-scale language models. arXiv preprint arXiv:230 2023.
URL https: //ap::///aa /ttp///aavvvv/aa////v///aa/aaaav org/ abs/2305 10305
Zefan Cai., Yichi Zhang, Bofei Gao, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, and Wen Xiao. Pyramidkv: Dynamic kv cache compression based on pyrami- dal information funneling, 2024.
Rewon Child, Scott Gray, Alec Radford, and lly Sutskever. Generating long sequences with sparse transformers, 2019.
Jincheng Dai, Zhuowei Huang, Haiyun Jiang, Chen Chen, Deng Cai, Wei Bi, and Shuming Shi. Sequence can secretly tell you what to discard, 2024.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024.
Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Des- jardins, Arnaud Doucet, David Budden, Yee Whye Tch, Razvan Pascanu, Nando De Freitas, and Caglar Gulcehre. Griffin: Mixing gated linear recurrences with local attention for efficient lan- guage models, 2024.
DeepSeek-AL, Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Hanwei Xu, Hao Yang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, J. L. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li, Jin Chen, Jingyang Yuan, Junjie Qiu, Junxiao Song, Kai Dong, Kaige Gao, Kang Guan, Lean Wang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Liyue Zhang, Meng Li, Miaojun Wang, Mingehuan Zhang, Minghua Zhang, Minghui Tang, Mingming Li, Ning Tian, Panpan Huang, Peiyi Wang, Peng Zhang, Qihao Zhu, Qinyu Chen, Qiushi Du R. J. Chen, R. I.. Jin, Ruiqi Ge, Ruizhe Pan, Runxin Xu, Ruyi Chen, S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shirong Ma, Shiyu Wang, Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Size Zheng. T. Wang, Tian Pei. Tian Yuan, Tianyu Sun, W. L. Xiao, Wangding Zeng, Wei An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang, Xiao Bi, Xiaodong Liu, Xiaohan Wang, Xiaojin Shen, Xiaokang Chen. Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang, Xin Liu, Xin Xie, Xingkai Yu, Xinnan Song, Xinyi Zhou, Xinyu Yang. Xuan Lu, Xuecheng Su, Wu, Y. K. Li, Y. X. Wei, Y. X. Zhu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yaohui Wang, Zheng, Yichao Zhang, Yiliang Xiong, Yilong Zhao, Ying He, Ying Tang, Yishi Piao, Yixin Dong, Yixuan Tan, Yiyuan Liu, Yongji Wang. Yongqiang Guo. Yuchen Zhu, Yuduan Wang, Yuheng Zou, Yukun Zha, Yunxian Ma, Yuting Yan, Yuxiang You, Yuxuan Liu, Z. Z. Ren. Zehui Ren, Zhangli Sha, Zhe Fu, Zhen Huang, Zhen Zhang, Zhenda Xie, Zhewen Hao, Zhihong Shao, Zhiniu Wen, Zhipeng Xu, Zhongyu Zhang, Zhuoshu Li, Zihan Wang, Zihui Gu, Zilin Li, and Ziwei Xie. Deepseek-v2 v2: A strong, economical, and efficient mixture-of-expeee language model, 2024.
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models to 128k context, 2024.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms, 2024.
gkamradt. Needle In A Haystack Pressure Testing LLMs, 2023. URL/tttps://githitm com/ gkamradt w.t.dd.................... Ltw.LiMMeetteeeeeeeeeeeee T InAHaystack
Team GLM, Aohan Zeng, Bin Xu, Bowen Wang, Chenhui Zhang, Da Yin, Diego Rojas, Guanyu Feng, Hanlin Zhao, Hanyu Lai, Hao Yu, Hongning Wang, Jiadai Sun, Jiajie Zhang, Jiale Cheng, Jiayi Gui, Jie Tang, Jing Zhang, Juanzi Li, Lei Zhao, Lindong Wu, Lucen Zhong, Mingdao Liu,

---

Published as a conference paper at ICLR 2025
Minlie Huang, Peng Zhang, Qinkai Zheng, Rui Lu, Shuaiqi Duan, Shudan Zhang, Shulin Cao, Shuxun Yang. Weng Lam Tam, Wenyi Zhao, Xiao Liu, Xiao Xia, Xiaohan Zhang. Xiaotao Gu, Xin Lv, Xinghan Liu, Xinyi Liu, Xinyue Yang, Xixuan Song, Xunkai Zhang. Yifan An, Yifan Xu, Yilin NIu, Yuantao Yang, Yueyan Li, Yushi Bai, Yuxiao Dong, Zehan Qi, Zhaoyu Wang, Zhen Yang. Zhengxiao Du, Zhenyu Hou, and Zihan Wang. Chatglm: A family of large language models from glm-1300 130b to glm-4 all tools, 2024.
Albert Gu and Tri Dao. Mamba: Linear- time sequence modeling with selective state spaces, 2024.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadehh Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization, 2024.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models?, 2024. URL https://////
/ 22.4444.66666
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chap- lot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Ann Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023.
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before generation, 2024.
Yujun Lin, Haotian Tang, Shang Yang, Zhekai Zhang, Guangxuan Xiao, Chuang Gan, and Song Han. Qserve: W4a8kv4 quantization and system co-design for efficient Ilm serving, 2024.
Zechun Liu, Barlas Oguz, Changsheng Zhao, Emie Chang, Pierre Stock, Yashar Mehdad, Yangyang Shi, Raghuraman Krishnamoorthi, and Vikas Chandra. Llm-qat: Pree-free quantization aware training for large language models. arXiv preprint arXi: 2023a.
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for I1m kv cache compression at test time, 2023b.
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: infinite context transformers with infini-attention, 2024.
Efficient
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song. Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. Rwkv: Reinventing mns for the transformer era, 2023.
Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation, 2022.
Noam Shazeer. Fast transformer decoding: One write-head all you need, 2019.
Ying Sheng, Lianmin Zheng. Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Re, Ion Stoica, and Ce Zhang. FlexGen: High-throughput generative inference of large language models with a single GPU. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th
12

---

Published as a conference paper at ICLR 2025
International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research. pp. 31094-311 31116. PMLR, 23-29 Jul 2023. URL
.mlr.
ppress/v /202/ www.ghaa.nnhaeeeeeeeeeaeeee html
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En-
hanced transformer with rotary position embedding, 2023.
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 Nie, 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.
Xiuying Wei, Yunchen Zhang, Xiangguo Zhang, Ruihao Gong, Shanghang Zhang, Qi Zhang, Fengwei Yu, and Xianglong Liu. Outlier suppression: Pushing the limit of low-bit transformer language models. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Process- ing Systems, volume 35, pp. 17402-17414. Curran Associates, Inc., 2022. URL htt.://rrrppssccpprrrerrrieeeaaa
https: [fcogpooeesff1aaaaaaaaaaaaa0--00
Xiuying Wei, Yunchen Zhang, Yuhang Li, Xiangguo Zhang, Ruihao Gong, Jinyang Guo, and Xian- glong Liu. Outlier suppression+: Accurate quantization of large language models by equivalent and optimal shifting and scaling, 2023.
Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanisti- cally explains long-context factuality, 2024.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. SmoothQuant: Accurate and efficient post-training quantization for large language models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 38087-38099, PMLR, 23-29 Jul 2023. URL mlr wttppppeess/////eee.eeepp
html
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks, 2024.
Dongjie Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. Pyramidinfer: Pyramid kv cache compression for throuhhroughput IIm inference, 2024.
Amir Zandieh, Insu Han, Vahab Mirrokni, and Amin Karbasi. Subgen: Token generation in sublin- ear time and memory, 2024.
Biao Zhang and Rico Sennrich. Root mean square layer normalization, 2019.
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. xcbench: Extending long context evaluation beyond 100k tokens, 2024. URL TtLLattps:// http//iivvvv. org. http://b//2222211 13718
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song. Yuandong Tian, Christopher Re, Clark Barrett, Zhangyang Wang, and Beidi Chen. H2o: Heavy- hitter oracle for efficient generative inference of large language models. In Thirty-eevnnth Confer- ence on Neural Information Processing Systems, 2023. URL forum? d=RkRrPP7GKO
13

---

Published as a conference paper at ICLR 2025
Yilong Zhao, Chien-Yo Yu Lin, Kan Zhu, Zihao Ye, Lequn Chen, Size Zheng, Luis Ceze, Arvind Krishnamurthy, Tianqi Chen, and Baris Kasikci. Atom: Low-bit quantization for efficient and accurate Ilm serving, 2024.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi Plug-and-play 2bit kv cache quantization with streaming asymmetric quantization. 2023. doi: 10. 10.11114 37282. URL https: https::/rgdoi. net/10, 13140/RG. 28
A APPENDIX
A.1 PROOF OF THEOREM 1
Below we first give an upper bound for the product of the queries and keys, and then show that the attention weight would decay bound. Since we have Wo, and k = WK, x where æ is the input of the Attention block, this leads to
when the positional bias is significantly larger than that upper
gk = æWQ, WK. x & We WK, 122/211/
(6)
Since is attained after LayerNorm, which means
11
x:
b
1
1
Σ i2 -1)²
i=1
i=1
Here & is the input of LayerNorm, d its dimension and &; is the i-th dimension of 2. The equation above leads
x 2
7 O
1 b
0 ② p
<2
+ 2|1b112
0
$ 2/ + b/22
(7)
Combining equation 6 and equation 7 we get
qk S [WQ, WK, lla (21| + 22(b) In order to give upper bound for the attention weight, have
(8)
exp (Sm-on (g: k))
m- (q; k)
In= exp (Sm (q: k))
exp (qkT 4, (m n2)
On= exp (S (Sm-an -m (g: k))
177
exp (qkt - (m n2)
exp (S (q: k))
exp (qkt 1n (m 11))
exp (q
S Sexp (qkT - Ip (m n)))
exp (qk.T
exp (m (m - n))
Therefore to ensure Att (g: k) s e, which is equivalent as log (Atth (g: < log(e),
need
log (Atto kq: Sgkt 4m n) < log(c)
14

---

Published as a conference paper at ICLR 2025
Taking equation 8 into the equation above, we get
I| Wxalla (21lv]l² 21]b1l - (m - n) - log(e),
which gives us
W WK, 1l2 (l>)1 + [1b113) log(()
4n
In this case, we have
Attmm (g: S
A.2 OTHER ACCURACY EVALUATION
We conducted additional evaluations of RazorAttention the RULER (Hsieh al. 2024) and InfiniteBench (Zhang et al., 2024) datasets using the its effectiveness. In experimental setup, retained 30% of retrieval heads(28% induction heads and 2% echo head) and implemented compression mechanism for non-retrieval heads when sequences exceeded 4k tokens: we preserved an attention window of size 4 and local window covering 20% the sequence length, with the remaining tokens being directly compressed.
model to demonstrate
Task
Original Model
RA Model
Task
Original Model
RA Model
codedebug
21.57
21.57
Ruler-166 92.07
92.15
ensum endia enqa enmc
30.71
30.58
19.50
19.50
Ruler-32k
84.93
85.90
29.09
29.50
63.31
63.31
Table RULER Benchmark Results
Table 6: InfiniteBench Results
A.3 EFFICIENCY EVALUATION
We evaluate the decoding latency and throughput of RazorAttention on the GLM-9B-IM model using Ascend 910B NPUs, considering different input lengths for both prefill and decoding, as shown in Table 7 Additionally, we achieve a maximum throughput improvement of 64.2% and 71.81% % for input lengths of 128k and 256k, respectively, due to the ability utilize larger batch size following KV cache compression.
Input Length
Batch Size
Prefill Speedup
Decoding Speedup
128k
1.2%
3.1%
4
5,66%
15,9%
8
8.03%
26.67%
10
7.05%
29.47%
18
From OOM to feasible inference
256k
1
6.58%
9.2%
8.76%
25.47%
5
9.93%
29.01% %
9
From OOM to feasible inference
Table 7: The performance acceleration of the GLM-9B-IM model Ascend 910B NPUs.
15


Reviews:
[{"review_id": "ZvHZVXehAS", "reviewer": "Reviewer_jpBx", "summary": "This work proposed a method advocating keeping a full KV cache for retrieval/induction heads while adopting a StreamingLLM-like cache pattern for the rests. The proposed method claims to have strong long context performance and many efficiency perks.", "questions": "DuoAttention and Not All Heads Matter are two concurrent works of RazerAttention. I'd appreciate it if the author could discuss their similarities and differences. In particular, I am interested in Figure 13(1) of DuoAttention — which supposedly features RazorAttention — but the needle performance is very different from your own reporting. I wonder why?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The proposed method demonstrates decent performance, providing a significant improvement over many token dropping-based baselines.\n2. The evaluation is thorough in terms of model coverages.\n3. As recognized by many recent arts, being FlashAttention compatible is a prerequisite of being a pratical long context serving method, and this work fulfills this prerequisite.", "weaknesses": "1. The baseline methods (H2O and StreamingLLM) are dated and do not reflect the current SOTA of KV cache compression advancement. I understand if the authors do not intend to feature methods like KIVI or TOVA — as they are with different schools of approaches — but compare with sparse inference techniques like SnapKV or MInference are necessary.\n\nThe authors' discussion around L358 regarding SnapKV is faithful, as it corresponds to the findings of SharedContextBench (also submitted to ICLR here). However, SnapKV — and, in essence, methods like LoCoCo and LESS — are still worth featuring since we can almost always find weaknesses in any method. The authors should consider adding some quantifiable evaluations to reflect the weakness of SnapKV-like methods, if they intend to highlight the superiority of RazerAttention in this regard.\n\n2. The needle setting is unclear. It is often very easy to do the \"eating sandwich in a park\" needle with the \"grass sky sun\" repetition background. But it is hard with the \"magic city number\" needle with a noisy background (e.g., Paul Graham's essay). The authors should clarify this setting and consider adding more standardized needle evaluations, like the RULER.\n\n3. Similar to #2, the LongBench truncation settings of Table 3 can also use some clarifications.\n\n4. RazorAttention is proposed as an efficiency work. However, the efficiency evaluation is very limited. I'd like to see latency and throughput evaluation under different context length/batch size workloads.\n\n5. The core design of the work heavily relies on the findings of Induction Heads and Retrieval Heads, making its novelty limited. While I don't think this is much of an issue due to the practical performance gain the proposed method achieved, the relationship of such works deserves much better highlight. The mere mentioning around L167 is not enough.\n\n6. I'd like to see more dataset coverage beyond LongBench — as it is no longer very long w.r.t. models like Llama 3.1. Evaluation of Llama 3.1 on $\\infty$Bench would fill this gap.\n\n\nIn summary, I think this is a decent work that cleverly leveraged some observations of retrieval/induction ideas. I am open to increasing my rating should the author provide the requested clarifications and evaluations. I also apologize for the slightly delayed review — I will try to be more engaging during the discussion period.", "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": "9ReWuUhT0u", "reviewer": "Reviewer_TBcp", "summary": "The paper proposed a novel KV cache mechanism basing on retrieval heads. To be more specific, the paper presents the observation of some retrieval heads that is essential for handling long context workload, then proposed a token-dropping based method to compress the to KV cache of the non-retrieval heads with compensation tokens, while keeping the retrieval head integrated, leading to substantial accuracy improvement compare to previous proposed method.", "questions": "See the weakness", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper observes the phenomenon of \"retrieval heads\" aligning with observations from other works.\n\nBasing on the retrieval heads, the paper presents a straightforward token dropping based KV compression mechanism.\n\nExperiments results show improvement compare with the previous token dropping based baselines.", "weaknesses": "Lack of experimental results on compression ratio:\nThe paper compares against StreamLLM and H2O. However, StreamLLM has a very limited compression ratio, and the performance of different compression ratio in H2O also will be different. The current manuscript lack the study of how different compression ratio in RazorAttention will affect the performance.\n\nLack of overhead evaluation on efficiency\nThe paper claims that RazorAttention can enhance LLM inference efficiency without overhead. However, there is no efficiency evaluation (both peak memory usage as well as the inference latency analysis). It is essential to have such real efficiency gain evaluation.\n\nThe novelty of \"retrieval heads\"\nPrevious paper [1] has demonstrated the existence of \"retrieval heads\" and its importance in serving long context. Do the authors aware of this and is there a difference between the two \"retrieval heads\" proposed in the two paper?\n\nMiss of clarification:\n\"Retrieve and Process\" in line 84: The manuscript claims that some heads are used for retrieval while some are used for process. However, lacks the analysis of \"process\"\n\"Protect\" in line 234: what is the specific protection mechanism? \n\nPotential typos:\n115 as follows:\n\n[1] Retrieval Head Mechanistically Explains Long-Context Factuality", "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": "zg4zSTqjb2", "reviewer": "Reviewer_hKDU", "summary": "The paper has made an observation: induction and echo heads are retrieval heads and others are local processing heads. They then keep complete KV Cache for retrieval heads and drop non-local KV for local processing heads. Following their hyper parameters, this leads to 70% asymptotic KV Compression.", "questions": "Please address the weaknesses above.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The observation is neat.  And experiments support the hypothesis quite well. \n2. The recipe to identify the heads is very lightweight (although is this a contribution or is it previously known from anthropic paper?)", "weaknesses": "1. Experiments\n\n     A. why are some datasets from longbench missing like passage_retrieval etc\n\n     B. Can you perform experiments on infinity benchmark for comparison.\n\n     C. what is the exact KV compression in the longbench datasets? The 70% number is misleading since longbench generally has small context lengths and authors have used a context thold of 4000.  Please add exact KV Cache compression numbers for all the methods to the table. \n\n2. Novelty: Can you please elaborate on the novelty of defining / identifying the echo and induction heads. what is already known and what is figured out?\n\n3. Writing errors: Please revise the manuscript for text/language errors. e.g. : \n    \n     A. \"There are been plenty ...\" (page 1)  \n    \n     B. \" and proved that accuracy ..\" ( page 2) {you showed it empirically .. and also 70% is not what you show in the experiments.}.\n\n     C. C_0 in equation 2. Do you mean L_h ?", "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": "fjsxqg2ELJ", "reviewer": "Reviewer_2S3S", "summary": "The paper introduces RazorAttention, a training-free algorithm for compressing KV cach in LLMs. This method selectively retains information in specific “retrieval heads” of the multi-head attention mechanism while discarding remote tokens in other heads. A notable aspect of RazorAttention is using a \"compensation token\" to summarize discarded information, minimizing data loss. RazorAttention shows a 70% reduction in KV cache size with minimal performance impact, supporting compatibility with FlashAttention for efficient large-context inference without modifying the base model.", "questions": "- What are the latency and throughput improvements, both absolute and relative to the baseline?\n- What are the comparative performances at different KV cache compression ratios?\n- Has RazorAttention been tested on other models like Mistral-7B-Instruct-v0.2 or longchat-7b-v1.5-32k?\n- Can you provide a breakdown or approximation of the distribution between retrieval and non-retrieval heads?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- RazorAttention’s head-specific strategy, distinguishing retrieval heads from non-retrieval heads, is a novel approach to KV cache compression that aligns with the attention heads’ different roles.\n- This method is compatible with FlashAttention and offers practical efficiency gains without training or model alterations.\n- RazorAttention is tested across various architectures, including ALiBi and RoPE models, proving effective on different LLM types, with compression ratios of up to 70%.", "weaknesses": "- The paper lacks detailed efficiency reports, particularly on latency and throughput improvements. A breakdown of the actual performance gains would give better insight into its practical benefits.\n- RazorAttention introduces operations like compensation token creation and slicing, yet the associated computational overhead is not profiled and discussed. A detailed overhead analysis and implementation would help.\n- RazorAttention’s comparison is primarily with StreamingLLM and H2O, omitting comparison and discussion of other types of KV cache compression methods. (like what are others’ performances under similar compression rates? can this method be combined with others?)\n- The results focus on one compression rate without experiments on varying compression ratios compared with others.", "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": "NHmn1X2a36", "reviewer": "Reviewer_gP7k", "summary": "This paper introduce a KV cache compression algorithm in LLMs that only maintains a full cache for retrieval heads and reduces all the remote tokens in non-retrieval heads to single compensation token per head. The proposed method can achieve 3X KV cache compression ratio without noticeable impacts on performance.", "questions": "1.\tWhat is the compression ratio for GQA in Llama3-8B-Instruct model? What is the compression ratio for 1024K sequence length?\n2.\tHow to decide the 1% echo heads and 14% induction heads configuration? Have the authors tried different combination, e.g. more echo heads and less induction heads?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1.\tIt is pluggable solution that enhances the inference efficiency of a diverse set of large language models without noticeable impacts on performance.", "weaknesses": "1.\tThere are some related works identifying the different patterns of retrieval and non-retrieval heads, such as DuoAttention, Retrieval Head Mechanistically Explains Long-Context Factuality, MInference, etc. It would be better to clarify the novelty of the proposed method.\n2.\tIt can only be applied to multiple KV cache scenarios. For only single KV cache model structure, such as YOCO, it cannot be applied.\n3.\tThe baselines are not comprehensive. It would be beneficial if the authors could compare their results with more advanced baselines.\n4.\tIt seems the proposed method cannot achieve the best performance in all the benchmarks. It would be beneficial to give detailed analysis of the worse cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5AEpa81MDR", "reviewer": "Reviewer_CNZf", "summary": "The paper introduces RAZORATTENTION, motivated by the goal of reducing the size of the key-value cache as much as possible without losing semantic information. The rationale behind this approach is the retrieval heads in LLMs. For these retrieval heads, RAZORATTENTION retains the KV cache in its original state; for the other heads, it only caches recent tokens and attention sinks. Additionally, the paper designs a compensation strategy for compensating the information loss. Experiments indicate that RazorAttention can compress up to 70% of the KV cache without causing a noticeable decline in performance.", "questions": "1. As mentioned in the weaknesses, it is beneficial to compare with DuoAttention in terms of effectiveness.\n2. As a training-free method, a potential issue is that different models may have different optimal hyperparameters. For example, how can we quickly determine the suitable proportion of Induction Heads and Echo Heads for different models?\n3. If a general theoretical analysis can be found explaining why retrieval heads are widely present in various models, that would be highly significant.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The research problem addressed in this paper is significant. Reducing the key-value  cache in large language models is a hot topic, and improvements in this area can have a substantial impact on the real-world deployment of LLMs.\n2. Based on the phenomenon of retrieval heads present in large language models and in conjunction with the token dropping method, this paper explores an training-free way to reduce the key-value cache. The results from downstream tasks and NIH tests have both demonstrated the effectiveness of this approach.\n3. The writing of this article is easy for readers to understand.", "weaknesses": "1. I agree that this work is the first training-free token reduction algorithm based on retrieval heads. However, to my knowledge, the phenomenon of retrieval heads being prevalent in LLMs was first articulated in [1]. [1] provides a detailed discussion on the properties of retrieval heads, which diminishes the novelty of this paper. Although the paper provides a theoretical analysis of AliBi-based models, which is commendable, it does not cover Rope-based models.\n2. A recent work [2] is also motivated by the presence of retrieval heads. The main idea of [2] is to combine streaming heads with retrieval heads, which is very similar to the approach taken in this paper. But [2] is not training-free. It would be beneficial to elaborate on the advantages of this paper over [2]. \n\n[1] Retrieval head mechanistically explains long-context factuality.\n[2] DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads", "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": "2dHBiNzINI", "reviewer": "Reviewer_ZhVH", "summary": "This paper investigates the attention mechanisms within transformer models, observing that most attention heads primarily focus on local context while only a few, termed \"retrieval heads,\" are capable of attending to all input tokens. To optimize memory and computational efficiency, the authors propose maintaining a full key-value (kv) cache for retrieval heads and only keeping a recent kv cache for other attention heads. The identification of these retrieval heads is based on observed patterns known as the \"Echo head\" and \"Induction head.\"", "questions": "The paper lacks detailed experimental setup descriptions. When comparing with baseline methods, was the total kv budget (i.e., the combined kv cache size for all heads) kept constant?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written and straightforward to follow.", "weaknesses": "1. The authors identify retrieval heads as echo heads and induction heads, demonstrated through an example using a synthetic repeated sequence. However, this example appears specifically designed to highlight these two types of heads. It raises the question of whether this pattern holds in more generalized query settings. Intuitively, it does not seem that in every type of query, the most crucial information is located solely in echo or induction tokens. Would these patterns still be dominant in more diverse and complex input data?\n\n2. The authors state that they did not compare their method with SnapKV because SnapKV relies on the query for selecting the kv cache. However, this reasoning is weak, as SnapKV is a well-known and accepted method in the field. A comparison would provide more comprehensive insights.\n\n3. I'm confusing that whether the selection of retrieval heads is model-based or  also query-specific. Are these retrieval heads consistent across different queries, or do different heads have different functions across different queries?\n\n4. The authors note that reducing the kv cache for all non-retrieval heads introduces a notable accuracy gap and that compensation tokens can mitigate this gap. However, it is unclear whether the observed improvement comes purely from the use of compensation tokens or from the strategic allocation of the kv cache itself. An experiment comparing the combination of compensation tokens with other baseline models could clarify the extent of the improvement attributable to each factor.", "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)