SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget
Authors: Zihao Wang
Bin CUI
Shaoduo Gan
Venue: ICLR 2025 Poster (2025)
Keywords: KV-cache
LLM inference optimization
Primary area: generative models

Abstract: Optimizing the Key-Value (KV) cache of the Large Language Model (LLM) has been considered critical to saving the cost of inference. Most of the existing KV-cache compression algorithms attempted to sparsify the sequence of tokens by taking advantage of the different importance of tokens. However, most of these methods treat all layers equally, allocating the same KV budget to each layer. This approach is suboptimal, as some layers may be less sensitive to input tokens yet still receive the same budget as others. In this work, we found that by identifying the importance of attention layers, we could optimize the KV-cache jointly from two dimensions, i.e., sequence-wise and layer-wise. Based on our observations regarding layer-wise importance in inference, we propose \sys to precisely optimize the allocation of KV-cache budget among layers on-the-fly and then incorporate three representative sequence-wise algorithms to compress the KV-cache for each layer with its very own budget. Specifically, we first measure each layer's importance by calculating the cosine similarity of the input prompt differences before and after the self-attention layers. Based on this similarity, we then categorize the layers into two groups and adjust their KV budgets accordingly. By optimizing the KV-cache from both sequence's and layer's dimensions, \sys achieves around 30\% to 70\% of the memory reductions and up to 2.2 $\times$ of throughput improvements in a wide range of LLMs and benchmarks. The code is available at https://github.com/hetailang/SqueezeAttention.

Published as a conference paper at ICLR 2025
SQUEEZEATTENTION: 2D MANAGEMENT OF KV- CACHE IN LLM INFERENCE VIA LAYER-WISE OPTI- MAL BUDGET
Zihao Wang.' Bin Cuit,
Shaoduo Gan
School of CS & Key Lab of High Confidence Software Technologies (MOE), Peking University Geoming Al
hetai lang0. o@gma com, (bin. cui, shaoduo. waw.gaannn@ppeeeeemmm edd..cn
ABSTRACT
Optimizing the Key- Key-V (KV) cache of the Large Language Model (LLM) has been considered critical saving the cost of inference. Most of the existing KV- cache compression algorithms attempted to sparsify the sequence of tokens by taking advantage the different importance of tokens. However, most of these methods treat all layers equally, allocating the same KV budget to each layer. This approach is suboptimal, as some layers may be less sensitive to input tokens yet still receive the same budget as others. In this work, we found that by identifying the importance of attention layers, we could optimize the V-cachhe jointly from two dimensions, i.e., sequence-wise wiie and layer-wise. Based on observations re- garding layer-wise importance in inference, we propose SQUEEZEATTENTION to precisely optimize the allocation of KV-cache budget among layers on-the-fir fly and then incorporate three representative sequence- wise algorithms to compress the first measure
KV-cache for each layer with its very own budget. Specifically, each layer's importance by calculating the cosine similarity the input prompt differences before and after the self-attention layers. Based on this similarity, we then categorize the layers into two groups and adjust their KV budgets accord- ingly. By optimizing the KV-cache from both sequence's s and layer's dimensions, SQUEEZEATTENTION achieves around 30% to 70% of the memory reductions and up to 2.2 x of throughput improvements in a wide range of LLMs and bench- marks. The code is available at
INTRODUCTION
The remarkable performance achieved by generative large language models (LLM) across wide range of natural language processing (NLP) tasks is making people in the computing industry believe that it has a great potential reshape the way they design their products. The past year has witnessed an unprecedented surge in applications driven by LLMs, such intelligent chatbots, LLM ALM--oweredd search engines, digital personal assistants, automatic programming tools, and so on. Along with the ever- ever-growing LLM upplications, their massive inference cost starts becoming severe challenge that hinders the deployment of LLMs and raises concerns regarding their carbon footprint Faiz et al. 2023
For a decoder-only autoregressive model, which is the most widely adopted LLM architecture, the inefficiencies in inference mainly come from the fact that the model can only generate tokens one by one, and sampling each token requires attending to all previous tokens. In practice, the intermediate key-value embeddings of each layer have been cached incrementally in each iteration avoid fre- quent recomputations. Since the K KV-cache increases linearly with the number of attention layers, context length and batch size, often ends up being multiple times larger than the model itself
"This work is done during the internship at Peking University and Geoming Al "orresponning author

---

Published as a conference paper at ICLR 2025
KV Cache Per Token
KV Cache Per Token
Layess
Layers
KV Cache Token
H2O
SqueezeAtentiin (w/H2O)
Layers
KV Cache Token
KV Cache Per Token
Full Cache
Layer
Layers
tokinn being generated
tokens cache
discarded tolens
nppc S ourput
Sliding Window
SqueezeAteentio (w/Silddin Window)
Figure 1: Demonstrations of KV-cache policies in inference from the view of the sequence and at- tention layer. Full Cache (leftmost column) simply stores the KV embeddings for all the tokens in all the layers. Sequence-wisg sesuncee-wis compression algorithms (middle column) drop tokens in the sequence's 's dimension, where each layer has the same cache budget. SQUEEZEATTENTION (rightmost column) further compresses the KV-cache by adaptively re-allocating the cache budgets in the layer's dimen- sion.
Sheng al. (2023). and therefore, dominating the I/O cost of inference. Recently, optimizing the KV-cache has been broadly considered a critical approach to boost the efficiency of inference. From the perspective of context length, many well-studied algorithms are trying to identify the most "valu- able" tokens in the sequence and evict the unimportant ones to reduce the KV-cache and attention complexity, such as Sliding Window Attention Beltagy et al.] (2020), Heavy-Hitter (H20) Zhang Xiao et al.] 2023), Scissorhands Liu et al. (2024), FastGen Ge et al.
et al. (2024), Streaming (2023) and so on. From the perspective batching, many studies aim to explore how to efficiently manage the memory of KV-cache on batch basis with different sequence lengths Zheng et al. 2002); Kwon et al. (2023). However, the opportunities in the dimension of attention layers have barely been touched by most, if not all, of the existing methods. In other words, all the attention layers have always been treated equally by those KV caching strategies. Therefore, in this paper we ask:
Do all the attention layers that share the same KV caching strategy have to cache the same amount of tokens? If not, how can we precisely allocate the cache budget for each layer such that we can further reduce the KV-cache on top of sequence-wise compressions?
To answer these questions, we need to take closer look at the behaviors of different attention layers during inference. Some inspiring clues could be found in a few existing studies. Early-exiting LLM Del Corro et al. (2023), as widely-adopted inference method, shows that after going through certain number of attention layers, the hidden representations are likely to reach saturation, and therefore, the forward computing can exit early without finishing the entire network and still get reasonable prediction. Besides, very recent work called FastGen Ge et al. (2023) found that attention layers in different positions have different optimal KV caching strategies. For example, attention layers in the very carly part of the model should simply cache all the tokens in the sequence, whereas, some middle layers should apply cache eviction strategies based on the token locality or frequency. Although FastGen could select the optimal eviction strategy for each layer, it is still unclear how to find the optimal cache budgets, instead of pre-defined, unified hyperparameter, for layers that share the same strategy.
Given the fact that attention layers do have different degrees of importance regarding inference, can make reasonable hypothesis that by taking advantage of the layer importance, could
2

---

Published as a conference paper at ICLR 2025
further "squeeze" the amount of KV-cache that has already been compressed by those sequence- wise eviction algorithms, and eventually, achieve even better efficiencies.
To describe the importance of the attention layers quantitatively, we track the cosine similarity, which has been considered a robust metric to reflect the similarity of embeddings in NLP Sidorov et al. (2014), between the hidden representations before and after the self-attention computing in each layer, and then put all layers' data together to demonstrate how an input embedding evolves through the entire model in inference. The intuition is that the more similar the embeddings are after the attention computing (indicated by higher cosine similarity), the less information this attention layer could insert into the embedding. After broad investigations into multiple popular LLM models, e.g., Mistral-7B, Falcon-7B, Llama2-7B. Llama2-70B and so on, as shown in Figure 2 we found some common characteristics. Firstly, the first half of attention layers, in general, contributes more to the output embedding than what second half does. Secondly, some specific layers, typically the first and last few layers, might be more important than other layers, depending on the specific model and dataset.
Based on this simple yet effective indicator, we propose SQUBEZEATTENTION, 2D KV-ccche compression algorithm that prunes KV-cache from not only the sequence's dimension but also the layer's dimension. Since the layer importance is highly dependent on the model and task, SQUEEZEATTENTION categorizes all the layers into groups the fly by clustering their cosine similarities measured during the prompt prefilling phase. Given a sequence-wise wise KV-cache eviction policy (like Sliding Window [Beltagy et al. (2020) H2O ZZang et al. (2024), and unified cache budget (like 4096 tokens or 20% of prompt length), SQUEEZLATTENTE
automatically reallocates the cache budgets among groups of layers such that the important layers could cache more tokens to stabilize the model accuracy and the unimportant layers could drop more tokens to save the I/O cost. What's 's even better is that SQUEEZEATTENTION orthogonal all those sequence- wise KV-cache compression algorithms, so it can be smoothly combined with any of them. Figure I demonstrates how the SQUEEZZAATENTION
works jointly with two representative sequence- wise KV-cache evic- tion algorithms, i.e., H2O /Zhang et al. (2024) and Sliding Window Attention [Beltagy et al. (2020). More details about the SQUEEZEATTENTION algorithm can be found in Section 4.
To the best of our knowledge, SQUELZEATTENTION is the first algorithm considering the KV- cache budget in a layer-wise way, making it a valuable addition to all those sequence-wise com- pression algorithms for inference. In experiments, integrate SQUEEZEATTENTION into 7 popular LLM models ranging from 7B to 70B, i.c., -lama227B, Mistral-7B, Falcon-7B, OPT-6.7B, GPT-Neox-20B, Mixtral-8x7B: and Llama2-70B, combining with 3 representative sequence- wise KV-cache compression a algorithms, i.e., Heavy- Oracle (H20). Sliding Window Attention and StreamingI.LMM The results show that SQUEEZEATTENTIO can achieve better model performance with even lower cache budgets than all three algorithms under a wide range of models and tasks, which lead to approximately 30% 70% of the memory savings and up 2.2 x of throughput improvements for inference.
PRELIMINARIES AND RELATED WORK
2.1 ANATOMY OF KV-CACHE IN LLM INFERENCE
For decoder-only transformer-based model, the inference process typically involves two phases: prefilling and decoding. In prefilling, LLM takes the entire prompt as input to calculate and cache the key-value embeddings of each token in each attention layer. Then the decoding phase takes one embedding at time to generate tokens by iterations, and meanwhile, concatenates the newly calculated KV embedding to the KV-cache.
Let be the length of the prompt, be the length of the output, b be the batch size, 1aayyer be the total number of attention layers, and dynodel be the hidden dimension. In the i-th layer, denote the model weights regarding attention Key and Value by w and w where W R democci and w E dndddd X xdmodet
In the prefilling phase, denote the hidden states of the i-th layer by h prompt where h prompt R pbxpxdmmdell Then the K KV-cache of the i-th layer after prefilling can be formulated as:
3

---

Published as a conference paper at ICLR 2025
Ck prompt Wik CV
h prompt W
In the decoding phase, denote the hidden states of the j-th output token in i-th layer by h (1 j o), where h butput, E [R6x1xdmook Then the KV-cache of the i-th layer after generating j-th tokens can be formulated
Boutput,
C CONCAT Butpput, wwk)
(2)
CV = CONCAT Noutput, wiv)
(3)
As the decoding process goes along, KV-cache growing incrementally until the output sequence is fully finished. Therefore, the maximum number of floats in total of the KV-cache is:
Blayer
bb- (p dmmodell
or simply 2 dmodel Maayer b (p o).
Taking Llama-2-7B in FP16 as an example, where naayer = 32, dmodel = 4096. The entire model weights consumes around 14GB of memory, whereas, the KV-cache takes around 0.5MB per token. In other words, the KV-cache starts to exceed model weights when processing more than 28K tokens, which could be easily made up of batch of 28 requests with 1K content length.
2.2 EXISTING KV-CACHE OPTIMIZATIONS
As analyzed above, the number of layers, batch size, and context length are three critical factors that decide the size of the KV-cache. Therefore, existing optimization studies are likely to seek opportunities from these perspectives.
Sparsifying the context sequence is an effective way to break the linear relationship between the context length and the KV-cache Del Corro et al. (2023); Zhang et al. (0024); Anagnostidis et al. 2023 : Sukhbaatar al. (2019): Rae Razavi (2020). The general intention of these algorithms to find out the unimportant tokens in the sequence and drop the KV-cache of these tokens. For example, Sliding Window Attention MBettay et al. (2020) only caches a certain number of the most recent tokens and drops the rest. StreamingLLM Xiao et al. (2023) found that in addition to the recent tokens, tokens at the beginning of the sequence are also crucial to the output. Heavy-Hitter Zhang al. (2024) and Scissorhands Liu et al. (2024) rank the importance of tokens by comparing their attention scores. As mentioned above, these algorithms treat all the attention layers equally, and therefore, have fixed KV-cache budget for each layer.
Optimizing the KV-cache on a batch basis mainly needs to manage the memory of different requests efficiently. For example, vLLM Kwon et al. (2023) allocates small chunks of memory an on- demand way, instead of a fixed big block for each prompt, to reduce the memory fragmentation in batch. RadixAttention Zheng et al. (2023) manages share the K KV-cache across requests in a batch when they have the same prefix in the prompt.
Last but not least, how relax the linear relationship between the KV-cache and the layers remains largely unexplored compared with the other two dimensions. FastGen Ge et al. (2023) is a very recent work that found layers in different positions may have different optimal sequence-wing KV- cache eviction strategies. It then proposed an algorithm to choose the best eviction strategy from 1) Locality strategy (like Sliding Window), 2) Special Tokens strategy (like StreamingLLMM) 3) Local and Frequency strategy (like Scissorhands, H2O) and so on for each attention head during the inference. However, for all the attention heads that have been assigned the same eviction strategy, FastGen simply gives them unified pre-defined cache budget, like 30% of the sequence length. Therefore, how to adaptively allocate the KV-cache budget to layers with the same sequence-wist token eviction strategy still largely unclear.
4

---

Published as a conference paper at ICLR 2025
3 OBSERVATIONS
Inspired by some previous works that managed to optimize the LLM inference from a layer-wise perspective Del Corro et al. (2023); Ge al. 2023), make hypothesis that attention layers in different positions play distinct roles in terms of importance, and therefore. should have different op- timal KV-cache budgets. To better understand the layer-wise contribution the output embedding, we employ cosine similarity as the metric to quantify the importance of each layer during inference. Specifically. in each attention layer, we track the hidden states of each input embedding two spots, i.e., the embedding before and after the self-attention block. Denote the hidden state before the self- attention by A, and the hidden state after the self-attention by B. The cosine similarity between two embeddings A and B is calculated follows:
n A.B
Σ
A
A
CosineSimilarityA,, B)
(5)
|A|
B
Σ
43
Σ
B
where (A1, A2,..., An) and B (B1, B2.... ality.
Ba)) are the vectors, and n is the dimension-
For any given input embedding, we could get cosine similarity of each attention layer that roughly quantifies how much the embedding changes after the self-attention computing of this layer. Then by comparing the cosine sim- ilarity among layers,
2020
2000
3020
05
could find a way rank attention layers by their importance. Fol- lowing this line of thought, we choose 4 rep- resentative LLMs, i.e., Mistral-7B Jiang et al. (2023. Llama2-7B-32K Touvron al. (2023). Llama2-70B Touvron et al. (2023). Falcon-7B Almazrouei et al. (2023), conduct the ex- periments. Fach model is fed by 200 prompts and we track the cosine similarity of each to- ken in each layer. Figure 2 shows the results after averaging over prompts. Each row of the heatmap demonstrates how input embedding at the given position evolves through all the lay- can
5020
D4
2000
6020
6000
0.
20
25.1
1033
-14
2520
2320
1250
3500
35
30
70
58
00
20
30
35
20
25
Figure 2: Visualization of cosine similarity before and after the self-attention calculation of attention each layer. The layers with higher cosine similar- ity, represented by lighter colors, exert a relatively lower impact the input vectors.
ers. From the brightness of the figures, get some insights as follows: 1) In general, the first half of layers (in darker color) contributes more to the output embeddings than the second half of layers (in lighter color) does; 2) The first and last few layers tend be more critical than other layers, depending on the specific models and tasks; 3) The cosine similarity can effectively depict the layer-wise importance, as the trend it reflects aligns with the previous studies, like Early- exiting Del Corro et al. (2023) and FastGen Ge et al. (2023).
This observation gives us a simple yet effective metric design a new algorithm that is able to optimize the K KV-cahhe from not only the context's dimension, but also the layer's 's dimension.
4 ALGORITHM
In this section, we describe the SQUEEZEATTENTION algorithm inspired by our observations from the last chapter. The most distinct feature of the proposed algorithm is that considers tokens in the KV-cache as a 2D matrix with one dimension of sequence and another dimension of layer, and both dimensions are going to be optimized jointly.
5

---

Published as a conference paper at ICLR 2025
4.1 SQUEEZEATTENTION
In the sequence's dimension, there are various cache eviction policies that could directly combine with, like Least Recently Used methods (Sliding Window, StreamingLLMD,, Least Frequently Used methods (Scissorhands, H20), and so sequence's 's dimension by Cseg. and its cache budget by binie. Note that all the layers have the same cache budget by default, just like the assumptions made by all these sequence- wise KV-cache compression algorithms.
We denote policy that compresses the KV-cache in
In the layer's dimension, SQUEEZEATTENTION firstly tracks the layer importance with the given prompt in the prefilling phase by collecting the cosine similarities of each layer whenever the self- aftention is conducted. At the end of prefilling, each layer ends up with set of cosine similarities, each of which corresponds to a token that has flowed through this layer. Then we use the averaged value over prompt tokens to represent the layer-wise importance of this layer regarding this task. By clustering the layers into groups based on the layer-wise importance with KMeanS, we reallocate the binit for each layer in way that more budgets are assigned to the more "important" layer groups. Since layers have different cache budgets, in the decoding phase, Caeq works separately with each layer's 's very own budgets. The detailed process described in Algorithm
Algorithm SOUBEZEATTENTION
Require: prompt: input sequence; Ceeqi a KV-cache compressor in sequence dimension; biniti the initial cache budget of each layer; Mayyrr number of attention layers; pt hyperparameter
(0 < p < 1); K(;; KV-cache of the i-th layer;
1: Feed the prompl into the model for prefilling, calculate cos_sim, (6) of the j-th token in the i-th layer by Equation
2: for + 1 to Aaayer do
(ppommp)))) cos_sim (4)
3:
cos cos_sim
[-l
pp(ppr)))))
4: end for
5: G1, G2 G3 KMeans(cossssmmm
> Cluster layers into 3 groups by cos_sim' (t) (1 & i Mayer)), where G3 has the biggest sim average
6: for + 1 to Rioger do
7:
i G3 then
8:
b(9) binit x
9:
else
10:
b(o)
Plager Xhini 1 Ga) Xbina
Iea(Cr)+len(Ga))
11:
end if
12: KO = (promm))
> the prompt's KV is cached after prefilling Di Decoding output tokens one by one
13: end for
14: for < 1 to (onulpp)l do
15:
for
Maaer do
if (e)( ((6) then
(i) Ceey (K( b(0) Compress the KV-cache of this layer by its own budget
18: end if
Finish the Self-attention based on the compressed (())
20: end for
end for
4.2 DISCUSSIONS
4.2.1 How TO DECIDE THE VALUE OF P
SQUEEZEATTENTION involves hyperparameter p to control the percentage of initial budgets that could be removed from the "unimportant" layers. The smaller the is, the more budgets will re- assigned. In experiments, we found 0.3-0.4 is reasonable choice range in most cases. To precisely understand the impact of we have conducted extra experiments to demonstrate how the model accuracy changs with the value of p, please refer to A.2 for more details.

---

Published as a conference paper at ICLR 2025
4.2.2 WHY CLUSTERING INTO THREE GROUPS?
Based on the observations of 7 models we have tried, we found they all have typical pattern (3 groups) with respect to the layer importance. Specifically, Group consists of few special layers (always the first and last few layers) which can be seen never be evicted. Then Group 2 and Group 3 do not have a fixed borderline with each other, but found that Group makes obviously less impact
analogy of special tokens that should
the embeddings, which be analogy of "frequent" and "unfrequent" tokens in the sequence-wisg methods. Therefore, our 3- group policy could be rephrased as: "we firstly identify and prioritize the special layers (Group 1), then classify the rest layers into two groups: important (Group 2) and unimportant (Group 3), then reallocate the cache budget based on the clustering. Even if we cluster the layers into more than groups, we just break down Group and Group into many small groups, but they eventually need be reduced into two classes again, that is, either reducing the budget or increasing the budget. To preserve the model accuracy, we only reduce cache budgets from Group 3, which accounts for around 50% to 70% of total layers.
Table 1: Datasets used in experiments.
Task
Task Type
Eval metric
Avg len
Language
Sample
CNN/ Daily Mail Summarization (3 sentence)
Rouge-2
2,000
EN
1,000
XSUM
Summarization (1 sentence)
Rouge-2
2,000
EN
1,000
SAMSUM
Few shot
Rouge-L
6,258
EN
200
NarrativeQA
Single-doc QA
F1
18,409
EN
200
TriviaQA
Few shot
F1
8,209
EN
200
EXPERIMENTS
5.1 EXPERIMENT SETUP
LLM Models. We choose 7 representative LLMs, with model sizes ranging from 6.7B to 70B and context lengths ranging from 2K to 32K, to evaluate the proposed algorithm: GPT-NeoX-20B,, OPT- 6.7B, Falcon-7B, Mistral-7B, Mixtral-8 7BB LLLmma-7B-324 LLama2-70B.
Datasets. We conduct experiments on datasets: CNN/Daily Mail, XSUM, TriviaQA, SAMSUM, and NarrativeAA. TriviaQA, SAMSUM, and NarrativeQA originate from LongBench Bai et al. 2023), where the data length typically exceeds 8k. CNN/Daily Mail and XSUM have an average length of about 2k. Detailed information about datasets be found in Table
Baselines. 3 sequence-wise sparsification algorithms are chosen as the baselines to integrate into SQUBEZEATTENTION and we assign each algorithm the model that works beat, which we call the beat baseline algorithms:
Heavy- (H20) ZH20) et al. (224)) Identify the important tokens within the sequence by comparing the accumulated attention score of each token.
Sliding Window Attention Beltagy al. (0200) "Local" strategy that only caches the most recent tokens" KV embeddings. This method works especially well with Mistral and Mixtral.
Streaming LLM Xiao et al. 2023): In addition to the most recent tokens, Streamingl always caches the first n tokens in the sequence, as they are identified as "sink tokens". We take n 4 as recommended by the paper.
Hardwares. We conduct all the experiments on the AWS platform (p4d.24xlarge) with Nvidia A A100-40GB GPUs, interconnected by the NVLinks (600 GB/s GPU peer-to-peer bandwidth).

---

Published as a conference paper at ICLR 2025
CHNltally Mal, CPTM
CHNMibllly Nal, Faltume3
XCSL/%, CPT--Mwak--00
0.00
0.12
0.09
0.08
0.10
0.08
20UGE-22
0.07
0.00
0.07
0.06
0.06
0.06
Heaoyuettter Orarle Scuerzelttent Cache
Sneeminggi
Hesoy-litter Oracle
0.04
Squeese
0.05
Scuesgelctent
Cache
Coche
0.03
0.00
+
Budqes
daet
Budqus
OP7-6.78
SAVIUUM Mixtrel 9-78
Wivisii,, Mistra-e-78
0.10
0.08
46
0.06
0.00
44
Heary-titter Oracle
Säiing Window
Siring Window
0.02
Scuezzoktennt on hull Cache
Sq adn.cctttttiionn
SqusencAtsniion
Pull Cache
Full Cache
150
150
KV Budget
Sudget
SAMSUM Mistrai-70
Thxleg,, Histral-70
NArroticOA, Mistool
40
22.5
35
20.0
17.5
40
10.0
Siiding Window
Siddnng Windew
Siiiig Window
DauoccosMtsst ntion
Squeccekttaniion Pull Coche
Pull
20
Cache
5.0
0
KV dget
KV BL dget
<V Budget SAMOUM L1am0--7T-1
XSUM, 1lammo-Ton
45
0.32
RGUGE-L
ROUGE
0.08
42
0.14
0.06
StrermingLM Cache
0.12
0.04
Full
Full
39
P
100
&
KV Rudg
Ru. dgat
Suuteer
Figure 3: Performance of SQUEEZE best sequence-wise wise compression baselines, and Full Cache under different cache budgets.
5.2 END-TO-END RESULT
Figure demonstrates the comparison results of SQUEEZEATTENTION and other baseline algo- rithms over all 7 models and 5 datasets. Full Cache (dashed line) means all tokens' KV embeddings are fully cached during the inference, therefore, it represents a relatively good model accuracy. The blue and orange curves in each subfigure illustrate how the model accuracy changes with the KV- cache budgets ranging from 10% to 100% of the total sequence length. Note that applying different sequence-win wise compression algorithms to different tasks would lead quite different model accura- cies. Therefore, for each task, we choose the best sequence-wiss wise compression algorithm to represent the best case, and then apply SQUEEZEATTENTION on top of the best case. As shown in the fig- ure, SQUREZEATTENTION consistently manages to improve the model accuracies under various KV-cache budgets by reallocating the cache budgets among layers. In other words. SQUEEZEAT- TENTION can achieve similar inference accuracies with much less KV-cache in total.
5.3 MEMORY CONSUMPTION
Now we evaluate the memory consumption of the proposed algorithm. We choose three settings to compare how much GPU memory it needs to run the inference without degradation of model
8

---

Published as a conference paper at ICLR 2025
Table Comparisons of the required KV-cache budget to achieve the best accuracy. Three models are selected to represent the small (7B), medium (20B), and large models (70B). For each task, we choose the best existing sequence-war wise sparsification algorithm as the baseline.
Performance Used KV Budget
Model
Size
Dataset
Best Baseline
Full Cache
w/ SQUaaZEATTENTION w/o SQUEEZEATTENT TENTION
Mistral
7B
SAMSUM
Sliding Window
43.53 100%
41.05 220
40 30%
GPT-NeoX
20B
XSUM
Heavy-Hitter Oracle
0.09 / 100%
0.09 / 20%
0.08 / 60%
LLama2
70B
XSUM
StreamingL..M
0.18/ 100%
0.17/ 30%
0.19 40%
5,73
Memory(Gat
2:16
10
1.74
15-
SoeaningLUM
Cach
Ses tasing Sinebn
Tull Cach
Sau GOLIOOTTE 1900
100 Algortthn
PMI PATC
Algorttim
Aigorthm
(a) LLama2-70B -OB
(b) Mistral-7B
GPT-NeoX-20BB
Figure 4: Comparisons of per-token decoding memory usage (exclude model parameters) among the Full cache, SOUREZEATTENTION and best baselines in order to achieve the same accuracy, as shown in Table
accuracy. We select Mistral-7B (Sliding Window), GPT-NeoX-20BB (Heavy-Hitter),, and LLama2- 70B (StreamingL.LM) to cover all three baseline algorithms and models in small, middle, and large sizes. We utilize multi-GPU inference if the model and KV-cache do not fit into single GPU.
that in all three settings, by achieving the same model accuracy, SQUEEZEAT- TENTION consumes the least amount of KV-cache budgets compared with the algorithms that only compress from the sequence's dimension. In some cases, it only takes one-third of the cache bud- get of H2O algorithm. Subsequently, we employ PYTORCH PROFILER to evaluate the dimin- ished memory usage of generating one token during the inference (w/o memory usage of model weights). Figure 4 shows that SQUREZEATTENTION can save 70% to 80% of memory usage per token compared with Full Cache method, and 25% to 66% of memory usage compared with baseline algorithms.
5.4 THROUGHPUT OF TOKEN GENERATION
Since SQUEEZEATTENTION manages to save the memory cost of inference, as shown in the previous sections, we want to explore how these memory reductions can be interpreted into improvements of token throughput. We choose 2 models, i.e., Mistral-7B and Llama-70B, to represent models in small and large sizes. With a fixed content length, we increase the batch size from 1 to 224 Mistral-7B, and 1 to 64 for Llama2. For each task, select the best baseline algorithm, that is, Sliding Window for Mistral-7B and StreamingLLM for Llama2-70B.
Table B shows the token throughput 8 A100-40GB GPUs. With the same batch size, SOUEEZEATTENTIO can enhance throughput by up to 2.2x for Mistral-7B and 1.4x for Llama2- 70B compared to the Full Cache. Besides, SQUEEZEATTENTION also enables batch sizes up to 224 and 64 for two models, which would cause out-of-memory for the Full Cache method. The throughput comparison between best baseline and SQUEFZEATTENTION reported in A.4
5.5 OVERHEAD OF THE ALGORITHM
The computational overhead SQUEEZEATTENTION comes from two operations: Cosine similar- ity and Kmeans. The execution of all these computations only happens during the prefilling phase.

---

Published as a conference paper at ICLR 2025
Table 3: Generation throughput (token / on eight A100 GPUs of Mistral-7B and LLama2-70B with SQUEEZEATTENTIN and Full Cache. To maintain the model accuracy, SQUEEZEATTENTIE uses 20% of the cache budget for Mistral-7B and 30% of the cache budget for LLama2-70B. "OOM" means out-of-memory. More results can be found in A.4
Batch size
Model
Size
prompt len gen len
Algorithm
SQUEEZEATTENTION
20.5
496.5
682.7
824.4
892.5
Mistral
512 + 1024
Full Cache
20.9
254.0
304.8
0OM
OOM
Batch size
Model
Size
prompt len gen len
Algorithm
8
16
32
64
SQUEEZEATTENTIO
5,2
37.2
71.2
116.2
170.7
LLama2
70B
256 512
Full Cache
5.2
36.0
62.5
84.8
OOM
Therefore, the cost of SQUEEZEATTENTION is one-time price, which is much more cost-efficient than those algorithms (like H2O) that require extra calculations in each iteration of token generation.
We profiled the entire prefilling phase of Mistral-7B with/without the wall-clock time. The prompt length up to 8k tokens. As shown Table 4 the SQUEEZEAT- TENTION only caused 6.3% increasement of prefilling time. Note that if take the decoding time into consideration, this ratio is going to be much smaller, depending on the actual number of tokens generated. Therefore, the overhead of ATTENTTNT detailed overhead of SQUEEZEATTENTION can be found in A.I
ATTENNIOO compare
basically neglectable. The more
Table Overhead of SQUEEZE (Prefilling Time in seconds on one Nvidia A100-40GB)
Model w/o SQUEEZEAATENTIION w/ SQUEEZEATTENTION
Overhead Ratio
Mistral-7B
0.636
0.676
6.3%
5.6 LIMITATIONS AND BROADER IMPACTS
SQUEEZEATTENTION works jointly with a sequence-wise sparsification policy, so the assumption we make is that for a given model and dataset, there exists a sequence-wise KV-cache eviction policy that won't 't hurt the model accuracy under a certain tolerance. However, the generalizability of these sequence-wise algorithms is still an active research topic. If this assumption cannot be met, SQUEEZEATTENTION might not work as expected.
The positive social impacts include reducing the energy cost and carbon emissions of LLM infer- ence. Besides, can optimize the user experience of LLM applications. Since the proposed algo- rithm could accelerate the LLM inference, the potential negative impact is that it might exaggerate the improper usage of LLM, like generating harmful information.
6 CONCLUSION
In this paper, we propose 2D KV-cache compression algorithm called SQUNEZEATTENTIO,, By tracking the cosine similarity of each attention layer, we found that layers in different positions have distinct degrees of importance regarding the output embedding. Inspired by this observa- tion, SQUEEZEATTENTION reallocates the KV-cache budgets over attention layers to further re- duce the memory cost of inference. Experiments over wide range of models and tasks show that SQUEEZEATTENTION can achieve better model accuracies with lower memory consumption com- pared with state-of-the-art algorithms that only compress KV-cache in sequence- wise way.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This work is supported by National Natural Science Foundation of China (U23B2048, U22B2037), the Fund of Kunpeng and Ascend Center of Excellence (Peking University), and High-performance Computing Platform of Peking University.
REFERENCES
Ebtesam Almazrouei, Hamza Alobeidli, Abdulaziz Alshamsi, Alessandro Cappelli, Ruxandra Co- jocaru, Mérouane Debbah, Étienne Goffinet, Daniel Hesslow, Julien Launay, Quentin Malartic, Danicle Mazzotta, Badreddine Noune, Baptiste Pamnier, and Guilherme Penedo. The falcon series of open language models, 2023. URL htttts://aaxiiiii orgg abs 2211 11511..1 11..6666
Sotiris Anagnostidis, Dario Pavllo, Luca Biggio, Lorenzo Noci, Aurelien Lucchi, and Thomas Hoff- mann. Dynamic context pruning for efficient and interpretable autoregressive transformers. arXiv preprint 2rXiv:2300..5500 2023.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint S.4iv:2300.145000 2023.
Iz. Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
Luciano Del Corro, Allie Del Giorno, Sahaj Agarwal, Bin Yu, Ahmed Awadallah, and Subhabrata Mukherjee. Skipdecode: Autoregressive skip decoding with batching and caching for efficient Ilm inference. arXiv preprint arXiv:2307.02622 2023.
Ahmad Faiz, Sotaro Kaneda, Ruhan Wang. Rita Osi, Parteck Sharma, Fan Chen, and Lei Jiang. Limcarbon: Modeling the end-to-end carbon footprint of large language models. arXiv preprint arXiv:2309. 14399, 2023.
Suyu Ge. Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for 1lms. arXiv preprint arXiv:2310.01801, 2023.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, al. Mistral 7b. arXiv preprint arXiv:2310.06825. 2023.
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 on Operating Systems Principles, 2023.
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Vietor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for I1m kv cache compression test time. Advances Neural Information Processing Systems, 36, 2024.
Jack W Rae and Ali Razavi. Do transformers need deep long-range memory. arXiv:2007.03356, 2020.
arXiv preprint
Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li. Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, Ion Stoica, and Ce Zhang. Flexgen: High-throughput generative inference of large language models with single gpu. In International Conference Machine Learning, pp- 31094-31116, PMLR, 2023.
Grigori Sidorov, Alexander Gelbukh. Helena Gómez-Adorno, and David Pinto. soft cosine measure: Similarity of features vector space model. Computación y Sistemas, 18 (3):4915504, 2014.
Soft similarity and
11

---

Published as a conference paper at ICLR 2025
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. arXiv preprint arXiv:1905.07799 2019.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. 10099.7744 /7453, 2023.
Zhenyu Zhang, Ying Sheng. Tianyi Zhou, Tianlong Chen, Lianmin Zheng. Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, al. H2o: Heavy-hitter oracle for efficient gen- erative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024.
Liammin Zheng, Liangsheng Yin, Zhiqiang Xie, Jeff Huang, Chuyue Sun, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E Gonzalez, et al. Efficiently programming large language models using sglang. arXiv preprint arXiv:2312.07104, 2023.
12

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 DETAILED OVERHEAD OF SQUEEZEATTENTION
In this section, we further evaluate the overhead introduced by SQUBEZEATTENTION, we broke down the time taken for two primary operations: cosine similarity computation and K-means clus- tering. The experiment setup follows the same setup as This experiment was conducted using a single Nvidia A100-40GB GPU with prompt lengths of up to 8k tokens.
Table 5: detailed onverhead
Cosine similarity
K-means
Total time
0.00068s
0.001s
0.02276s
The operation of computing cosine similarity involves calculating between two arrays of size 8000x4096, repeated 32 times (since Mistral has 32 layers). Additionally, K-means clus- ters 32 numbers into classes. Therefore, the total time consumption be calculated as (00068x2++000110000222
seconds. It is noteworthy that this overhead is incurred only once, regardless of the number of tokens processed.
A.2 THE FUNCTION OF
The hyperparameter is crucial in the SqueezeAttention mechanism, as it directly determines the final allocation the KV cache. To illustrate this with a concrete example: suppose we have a model with 32 layers, where 18 layers less important. Each layer initially has a budget of 1000 tokens. If we set to 0.3, we will take 70% of the budget from the less important layers and redistribute it equally among the important layers.
deemed important and the remaining 14 are considered
Specifically, the budget for the less important layers will be reduced to (000x003=3000 Meanwhile, the budget for the important layers will be calculated as follows: In this way, the total budget remains unchanged.
To evaluate the sensitivity of the hyperparameter p, we tested the Mistral-7B model the Samsum dataset with p values ranging from 0.1 to 1.0. The total KV budget was set to 20% of the prompt length. The results are presented in the table below:
Table 6: The precision changing with p
ROUGE-L
15.7
34.01
37.26
37.69
27.41
26.29
11.11
10.48
8.72
9.07
p
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
100% means we do not alter the structure of the model's KV cache. As p decreases, more KV bud- get will be transferred to other layers, which can yield better results up to a certain point. When decreases to 10% or even lower, the performance decreases the less important layers' KV budget becomes severely inadequate. We can clearly observe an accuracy improvement when adjusting while keeping the overall budget unchanged. This demonstrates the impact of p on model perfor- mance, highlighting its importance in optimizing KV budget allocation.
A.3 LAYER IMPORTANCE ACROSS DIFFERENT TASK
We conducted an extra experiment using two models and various datasets to determine whether the importance of different layers is an intrinsic property of the model.
13

---

Published as a conference paper at ICLR 2025
Table Mistral-7B
Table 8: LLama2-70B
Dataset
Samsum
TriviaQA 18
LCC
Dataset
Xsum
Samsum
LCC
Important
17
19
Important
17
21
18
Unimportant 15
14
13
Unimportant
63
59
62
Tabe 7 displays the distribution of important layers for the Mistral model across three different datasets: Samsum (Few shot), TriviaQA (Single-document QA), and LCC (Code, Python/C#/Java) and table [ shows the distribution of important layers for the LLama2-70B model across three dif- ferent datasets: Xsum (Summarization), Samsum (Few shot), and LCC (Code, Python/C#/lava)
From these tables, can observe that there is rough pattern regarding the layer's group with task- specific fluctuations. We believe there exist some task-sensitive layers that may be classified into different groups with different tasks. Similarly, there also some layers that are always important / unimportant. A detailed analysis of this phenomenon could be an interesting extension this work. However, would still recommend adaptive way since it can precisely capture the importance of layers.
A.4 THROUGHPUT COMPARISON BETWEEN BEST BASELINE ALGORITHMS AND SQUEEZEATTENTION
Table Generation throughput (token / s) on eight A100 GPUs of Mistral-7B and LLama2-7B with SQUEEZEATTENTION and best baseline algorithms. To maintain the model accuracy, SQUEEZEAT- TENTION uses 20% of the cache budget for Mistral-7B and 40% of the eache budget for LLama2-77B while Sliding Window uses 30% of the cache budget for Mistral-7B and StreamingL use 60% of the cache budget for LLama2-7B. "OOM" means out-of-memory.
Batch size
Model
Size
prompt len + gen len
Algorithm
32
64
128
224
SQUEEZEAFTENTIO
20.5
496.5
682.7
824.4
892.5
Mistral
7B
512 + 1024
Sliding Window
20.6
404.5
512.2
587.8
OOM
Batch size
Model
Size
prompt len gen len
Algorithm
32
64
128
256
SQUEEZEATTENTION
20.0
143.0
150.4
144.9
LLama2
512 + 1024
StreamLLM
20.4
113.7
102.4
OOM
OOM
We also conducted experiments to compare throughput of SQUEFZEATTENTION with best baseline under a set of batch sizes. Both experiments used an input length of 12 and an output length of 1024. We choose the compression hyperparameters for each algorithm such that they could all achieve the best mode accuracy. The result show that our algorithm can obviously increase the throughput com- pared with those SOTA algorithms that only compress KV-cache from the sequence's 'ss dimension.
14


Reviews:
[{"review_id": "OMSHRCiw9S", "reviewer": "Reviewer_ARNg", "summary": "This paper proposes SqueezeAttention, a KV-Cache management algorithm that can be combined with KV-Cache eviction policies to further reduce memory footprint and improve throughput. SqueezeAttention allocates size budgets for the KV-Cache of different layers by utilizing statistics on the importance of the attention layers. Specifically, SqueezeAttention first computes the cosine similarity between the activations before and after each attention layer.  Based on this similarity, the layers are then categorized into two groups and their KV budgets adjusted accordingly. SqueezeAttention achieves around 30% to 70% memory reductions and up to 2.2 × of throughput improvements in a wide range of LLMs and benchmarks.", "questions": "1. What inference framework is used for the memory and throughput experiments? Is SqueezeAttention compatible with current inference memory optimization like vllm[1]?\n\n[1] Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C.H., Gonzalez, J., Zhang, H. and Stoica, I., 2023, October. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles (pp. 611-626).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The method can augment other KV-Cache eviction policies, which will benefit the research community.\n2. The algorithm is clearly presented and the method's effectiveness has strong experiment evidence.", "weaknesses": "1. There's little analysis of the reason for performance improvement as shown in Figure 3. Some hypothesis or statistics analyses could give readers a deeper understanding of the algorithm.\n2. The memory usage of Figure 4 is not clearly explained. What tensors are counted in the PyTorch Profiler? Besides, why does LLama2-70B consume a similar amount of memory to Mistral-7B?", "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": "zW4KQLWTAG", "reviewer": "Reviewer_2BQ8", "summary": "This work proposed a layer-wise KV cache compression method that reduce the overhead during decoding stage of LLM inference. The proposed squeezeattention use cosine similarity of embeddings before and after attention block to identify the redundancy of kv cache with respect to specific layer. Then more redundant layers will then be assigned with smaller kv cache budget. For each layer, squeezeattention based on previous methods to remove redundant kv pairs, such as H2O, SteamingLLM and Sliding windows.", "questions": "- Do G1,G2,G3 changes frequently across different samples? otherwise we can assign the layer-wise budget through a offline process.", "limitations": "", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The proposed methods is evaluated with multiple LLMs on various downstream tasks, demonstrates non-trivial improvements against previous baselines.\n\n- The manscript is clearly organized with several illustration figures and equations. It's easy to understand the main method of this work.\n\n- Both perfomance comparison and end-to-end memory/thoughput comparison are reported.", "weaknesses": "- The main observation that the cosine similarity of embeddings changes across layers while the first and last layers tend to have more diverse embeddigns, is not very new. Several works have showed similar results[1-3].\n\n- It would be helpful to consider more recent kv cache compression methods, like SnapKV, PyramidKV, KIVI, etc. As the layer-wise strategy seems can be used in either KV cache pruning/quantization/low-rank decomposition methods, etc.\n\n- In Table 3, it's a little bit unfair to compare the thoughput only with the full cache, since the KV cache evicted method is not the contribution of this work while the part of the thoughput improvements is achieved by the kv eviction, rather than the layer-wise strategy.\n \n\n[1] https://arxiv.org/abs/2312.17276\n\n[2] https://proceedings.neurips.cc/paper_files/paper/2023/file/fde1a69a5b6e554b2f1f727197d2651d-Paper-Conference.pdf\n\n[3] https://arxiv.org/pdf/2202.08625", "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": "CXVtpgC2s0", "reviewer": "Reviewer_fGAX", "summary": "This paper identifies the importance of different attention layers, and proposes a layer-wise strategy named as SqueezeAttention to allocate different KV cache size for each layer. However, the proposed method still have several significant issues.", "questions": "It is unclear why the authors use ROUGE-2 for CNN/Daily Mail and XSUM, but ROUGE-L for SAMSUM. ROUGE-L is generally considered a more accurate metric for summarization tasks and could be applied consistently across all datasets. The authors could either evaluate all datasets with ROUGE-L for consistency or provide a rationale for choosing different metrics for each dataset.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "The observations of comparing the inputs and outputs of attention modules are good.", "weaknesses": "1.\tThe proposed method is designed only for the prefilling stage and does not allow for dynamic adjustment of the KV cache size during the decoding stage. To improve applicability, it would be helpful if the authors discussed potential ways to extend the method to the decoding stage, or provided a rationale explaining why it may not be feasible in that context.\n2.\tThe reduced KV cache size is controlled by the hyperparameter \\( p \\), with values in the range of 0.3-0.4 based on a single model and task. This approach lacks generality. To improve robustness, the authors could conduct experiments across multiple models and tasks to determine if this \\( p \\) value range holds more broadly. Alternatively, they could propose a method for automatically selecting \\( p \\) to adapt to different scenarios.\n3.\tThe method uses a fixed number of clusters, specifically 3, which may limit its generalizability. To strengthen the justification for this choice, the authors could either provide a rationale for using 3 clusters or experiment with different numbers of clusters to determine the optimal setting across various scenarios.\n4.\tThe experiments appear incomplete. While Figure 3 includes four baselines, such as the full KV cache, each experiment only presents one baseline alongside the proposed method for comparison. Including all baselines in each experiment would allow for a more comprehensive evaluation. If certain baselines were omitted, the authors should explain why.", "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": "fbn3v25fzt", "reviewer": "Reviewer_Pb9Y", "summary": "The paper proposes SqueezeAttention, a novel 2D Key-Value (KV) cache management algorithm designed to optimize memory usage and processing efficiency during Large Language Model (LLM) inference. The motivation behind this work is that existing KV-cache compression strategies handle all attention layers equally, which is suboptimal. Instead, SqueezeAttention dynamically allocates the KV-cache budget based on each layer's importance, determined by the cosine similarity of embeddings before and after each self-attention layer. By combining sequence-wise and layer-wise cache optimization, SqueezeAttention provides substantial memory savings (30%-70%) and throughput improvements (up to 2.2×) across various LLM models, including Mistral-7B, Falcon-7B, and Llama2-70B. The experimental results show significant performance gains in memory efficiency and token generation speed.", "questions": "See the discussion of weaknesses and kindly address them.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Novel Layer-Wise Approach: This paper introduces a layer-wise approach to KV-cache optimization, differentiating it from existing sequence-based compression methods. This work fills a gap in current LLM efficiency research.\nSignificant Performance Improvement: The proposed method improves memory consumption and throughput by reallocating cache budgets based on layer importance.\nRobust Experimental Validation: The authors test their approach on multiple models (ranging from 7B to 70B parameters) and datasets, demonstrating its generalizability and efficiency.\nCompatibility with Other Methods: SqueezeAttention integrates smoothly with various sequence-wise compression techniques, enhancing its versatility.\nEnergy Efficiency: The memory and throughput improvements have practical implications, potentially reducing the environmental impact of LLM deployment.", "weaknesses": "Dependency on Sequence-Wise Algorithms: The effectiveness of SqueezeAttention relies on combining it with existing sequence-wise compression methods, which limits its standalone applicability.\nPotential Task-Specific Tuning: Although the layer importance measurement is automated, there may be task-specific variations, suggesting possible limitations in generalizing to unseen tasks without fine-tuning.\nLimited Analysis of Computational Overheads: Although the paper claims that SqueezeAttention adds a negligible overhead, more analysis on computation costs, particularly for real-time applications, would strengthen the results.\nFixed Group Clustering: The choice of clustering layers into three fixed groups may oversimplify the optimization for some models or tasks where layer importance does not align neatly with this structure.\nRisk of Reduced Accuracy: The method risks performance degradation for certain parameter values by under-allocating cache to less \"important\" layers, which might be essential for specific tasks or models.\nIn conclusion, the paper presents a promising contribution to LLM inference optimization with its innovative, adaptive KV-cache management strategy. However, further exploration into standalone performance and task-specific tuning would enhance the robustness of SqueezeAttention.", "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)