SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget

## Abstract

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 SqueezeAttention 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, SqueezeAttention 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.

## 1 Introduction

The remarkable performance achieved by generative large language models (LLM) across a wide range of natural language processing (NLP) tasks is making people in the computing industry believe that it has a great potential to reshape the way they design their products. The past year has witnessed an unprecedented surge in applications driven by LLMs, such as intelligent chatbots, LLM-powered search engines, digital personal assistants, automatic programming tools, and so on. Along with the ever-growing LLM applications, their massive inference cost starts becoming a severe challenge that hinders the deployment of LLMs and raises concerns regarding their carbon footprint Faiz et al. ().

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 to avoid frequent recomputations. Since the KV-cache increases linearly with the number of attention layers, context length and batch size, it often ends up being multiple times larger than the model itself Sheng et al. (), 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 ”valuable” 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. (), Heavy-Hitter (H2O) Zhang et al. (), StreamingLLM Xiao et al. (), Scissorhands Liu et al. (), FastGen Ge et al. () and so on. From the perspective of batching, many studies aim to explore how to efficiently manage the memory of KV-cache on a batch basis with different sequence lengths Zheng et al. (); Kwon et al. (). 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:

Figure: Figure 1: Demonstrations of KV-cache policies in inference from the view of the sequence and attention layer. Full Cache (leftmost column) simply stores the KV embeddings for all the tokens in all the layers. Sequence-wise compression algorithms (middle column) drop tokens in the sequence’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 dimension.
Refer to caption: extracted/5915379/figure/introduce/intro.png

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 a 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. (), as a widely-adopted inference method, shows that after going through a 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 a reasonable prediction. Besides, a very recent work called FastGen Ge et al. () found that attention layers in different positions have different optimal KV caching strategies. For example, attention layers in the very early 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 a 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, we can make a reasonable hypothesis that by taking advantage of the layer importance, we could 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. (), 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 SqueezeAttention, a 2D KV-cache 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 on the fly by clustering their cosine similarities measured during the prompt prefilling phase. Given a sequence-wise KV-cache eviction policy (like Sliding Window Beltagy et al. () or H2O Zhang et al. ()), and a unified cache budget (like 4096 tokens or 20% of prompt length), SqueezeAttention 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 even better is that SqueezeAttention is orthogonal to all those sequence-wise KV-cache compression algorithms, so it can be smoothly combined with any of them. Figure 1 demonstrates how the SqueezeAttention works jointly with two representative sequence-wise KV-cache eviction algorithms, i.e., H2O Zhang et al. () and Sliding Window Attention Beltagy et al. (). More details about the SqueezeAttention algorithm can be found in Section 4.

To the best of our knowledge, SqueezeAttention is the first algorithm considering the KV-cache budget in a layer-wise way, making it a valuable addition to all those sequence-wise compression algorithms for inference. In our experiments, we integrate SqueezeAttention into 7 popular LLM models ranging from 7B to 70B, i.e., Llama2-7B, Mistral-7B, Falcon-7B, OPT-6.7B, GPT-Neox-20B, Mixtral-8$\times$7B, and Llama2-70B, combining with 3 representative sequence-wise KV-cache compression algorithms, i.e., Heavy-Hitter Oracle (H2O), Sliding Window Attention and StreamingLLM. The results show that SqueezeAttention 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% to 70% of the memory savings and up to 2.2 $\times$ of throughput improvements for inference.

## 2 Preliminaries and Related Work

### 2.1 Anatomy of KV-cache in LLM Inference

For a 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 a time to generate tokens by iterations, and meanwhile, concatenates the newly calculated KV embedding to the KV-cache.

Let $p$ be the length of the prompt, $o$ be the length of the output, $b$ be the batch size, $n_{layer}$ be the total number of attention layers, and $d_{model}$ be the hidden dimension. In the $i$-th layer, denote the model weights regarding attention Key and Value by $\textbf{w}_{K}^{i}$ and $\textbf{w}_{V}^{i}$, where $\textbf{w}_{K}^{i}\in\mathbb{R}^{d_{model}\times d_{model}}$, and $\textbf{w}_{V}^{i}\in\mathbb{R}^{d_{model}\times d_{model}}$.

In the prefilling phase, denote the hidden states of the $i$-th layer by $\textbf{h}_{prompt}^{i}$, where $\textbf{h}_{prompt}^{i}\in\mathbb{R}^{b\times p\times d_{model}}$. Then the KV-cache of the $i$-th layer after prefilling can be formulated as:

$$ $\displaystyle\textbf{C}_{K}^{i}=\textbf{h}^{i}_{prompt}\cdot\textbf{w}_{K}^{i} ;\textbf{C}_{V}^{i}=\textbf{h}^{i}_{prompt}\cdot\textbf{w}_{V}^{i}$ (1) $$

In the decoding phase, denote the hidden states of the $j$-th output token in $i$-th layer by $\textbf{h}_{output_{j}}^{i}$ ($1\leqslant j\leqslant o$), where $\textbf{h}_{output_{j}}^{i}\in\mathbb{R}^{b\times 1\times d_{model}}$. Then the KV-cache of the $i$-th layer after generating $j$-th tokens can be formulated as:

$$ $\displaystyle\textbf{C}_{K}^{i}=\texttt{CONCAT}(\textbf{C}_{K}^{i},\textbf{h}_ {output_{j}}^{i}\cdot\textbf{w}_{K}^{i})$ (2) $\displaystyle\textbf{C}_{V}^{i}=\texttt{CONCAT}(\textbf{C}_{V}^{i},\textbf{h}_ {output_{j}}^{i}\cdot\textbf{w}_{V}^{i})$ (3) $$

As the decoding process goes along, KV-cache is growing incrementally until the output sequence is fully finished. Therefore, the maximum number of floats in total of the KV-cache is:

$$ $\displaystyle 2\cdot\sum^{n_{layer}}_{i=1}[b\cdot(p+o)\cdot d_{model}]$ (4) $$

or simply $2\cdot d_{model}\cdot n_{layer}\cdot b\cdot(p+o)$.

Taking Llama-2-7B in FP16 as an example, where $n_{layer}=32$, $d_{model}=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 a 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. (); Zhang et al. (); Anagnostidis et al. (); Sukhbaatar et al. (); Rae & Razavi (). The general intention of these algorithms is to find out the unimportant tokens in the sequence and drop the KV-cache of these tokens. For example, Sliding Window Attention Beltagy et al. () only caches a certain number of the most recent tokens and drops the rest. StreamingLLM Xiao et al. () found that in addition to the recent tokens, tokens at the beginning of the sequence are also crucial to the output. Heavy-Hitter Zhang et al. () and Scissorhands Liu et al. () rank the importance of tokens by comparing their attention scores. As mentioned above, these algorithms treat all the attention layers equally, and therefore, have a 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. () allocates small chunks of memory in an on-demand way, instead of a fixed big block for each prompt, to reduce the memory fragmentation in a batch. RadixAttention Zheng et al. () manages to share the KV-cache across requests in a batch when they have the same prefix in the prompt.

Last but not least, how to relax the linear relationship between the KV-cache and the layers remains largely unexplored compared with the other two dimensions. FastGen Ge et al. () is a very recent work that found layers in different positions may have different optimal sequence-wise 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 StreamingLLM), 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 a 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-wise token eviction strategy is still largely unclear.

## 3 Observations

Inspired by some previous works that managed to optimize the LLM inference from a layer-wise perspective Del Corro et al. (); Ge et al. (), we make a hypothesis that attention layers in different positions play distinct roles in terms of importance, and therefore, should have different optimal KV-cache budgets. To better understand the layer-wise contribution to 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 at two spots, i.e., the embedding before and after the self-attention block. Denote the hidden state before the self-attention by $\mathbf{A}$, and the hidden state after the self-attention by $\mathbf{B}$. The cosine similarity between two embeddings $\mathbf{A}$ and $\mathbf{B}$ is calculated as follows:

$$ $\text{CosineSimilarity}(\mathbf{A},\mathbf{B})=\frac{\mathbf{A}\cdot\mathbf{B} }{\|\mathbf{A}\|\cdot\|\mathbf{B}\|}\\ =\frac{\sum\limits_{i=1}^{n}A_{i}\cdot B_{i}}{\sqrt{\sum\limits_{i=1}^{n}A_{i} ^{2}}\cdot\sqrt{\sum\limits_{i=1}^{n}B_{i}^{2}}}$ (5) $$

where $\mathbf{A}=(A_{1},A_{2},\ldots,A_{n})$ and $\mathbf{B}=(B_{1},B_{2},\ldots,B_{n})$ are the vectors, and $n$ is the dimensionality.

Figure: Figure 2: Visualization of cosine similarity before and after the self-attention calculation of attention each layer. The layers with higher cosine similarity, represented by lighter colors, exert a relatively lower impact on the input vectors.
Refer to caption: extracted/5915379/figure/heat_map/heat_map_22.png

For any given input embedding, we could get a 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 similarity among layers, we could find a way to rank attention layers by their importance. Following this line of thought, we choose 4 representative LLMs, i.e., Mistral-7B Jiang et al. (), Llama2-7B-32K Touvron et al. (), Llama2-70B Touvron et al. (), Falcon-7B Almazrouei et al. (), to conduct the experiments. Each model is fed by 200 prompts and we track the cosine similarity of each token in each layer. Figure 2 shows the results after averaging over prompts. Each row of the heatmap demonstrates how an input embedding at the given position evolves through all the layers. From the brightness of the figures, we can get some insights as follows: 1) In general, the first half of layers (in a darker color) contributes more to the output embeddings than the second half of layers (in a lighter color) does; 2) The first and last few layers tend to 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. () and FastGen Ge et al. ().

This observation gives us a simple yet effective metric to design a new algorithm that is able to optimize the KV-cache from not only the context’s dimension, but also the layer’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 it 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.

### 4.1 SqueezeAttention

In the sequence’s dimension, there are various cache eviction policies that we could directly combine with, like Least Recently Used methods (Sliding Window, StreamingLLM), Least Frequently Used methods (Scissorhands, H2O), and so on. We denote a policy that compresses the KV-cache in sequence’s dimension by $\bm{C}_{seq}$, and its cache budget by $b_{init}$. 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.

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-attention is conducted. At the end of prefilling, each layer ends up with a 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 $b_{init}$ for each layer in a way that more budgets are assigned to the more ”important” layer groups. Since layers have different cache budgets, in the decoding phase, $\bm{C}_{seq}$ works separately with each layer’s very own budgets. The detailed process is described in Algorithm 1.

Figure: Algorithm 1 SqueezeAttention

### 4.2 Discussions

#### 4.2.1 How to decide the value of p 𝑝 p italic_p

SqueezeAttention involves a hyperparameter $p$ to control the percentage of initial budgets that could be removed from the ”unimportant” layers. The smaller the $p$ is, the more budgets will be reassigned. In experiments, we found 0.3-0.4 is a reasonable choice range in most cases. To precisely understand the impact of $p$, 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.

#### 4.2.2 Why clustering into three groups?

Based on the observations of 7 models we have tried, we found they all have a typical pattern (3 groups) with respect to the layer importance. Specifically, Group 1 consists of a few special layers (always the first and last few layers) which can be seen as an analogy of special tokens that should never be evicted. Then Group 2 and Group 3 do not have a fixed borderline with each other, but we found that Group 3 makes obviously less impact on the embeddings, which can be seen as an analogy of ”frequent” and ”unfrequent” tokens in the sequence-wise 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 3 groups, we just break down Group 2 and Group 3 into many small groups, but they eventually need to 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 our 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 |

## 5 Experiments

### 5.1 Experiment Setup

Figure: Figure 3: Performance of SqueezeAttention, best baselines, and Full Cache under different cache budgets.
Refer to caption: extracted/5915379/figure/result/main_result.png

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$\times$7B, LLama2-7B-32k, LLama2-70B.

Datasets. We conduct experiments on 5 datasets: CNN/Daily Mail, XSUM, TriviaQA, SAMSUM, and NarrativeQA. TriviaQA, SAMSUM, and NarrativeQA originate from LongBench Bai et al. (), where the data length typically exceeds 8k. CNN/Daily Mail and XSUM have an average length of about 2k. Detailed information about datasets can be found in Table 1.

Baselines. 3 sequence-wise sparsification algorithms are chosen as the baselines to integrate into SqueezeAttention and we assign each algorithm to the model that works beat, which we call the beat baseline algorithms:

- •
Heavy-Hitter (H2O) Zhang et al. (): Identify the important tokens within the sequence by comparing the accumulated attention score of each token.
- •
Sliding Window Attention Beltagy et al. (): A ”Local” strategy that only caches the most recent tokens’ KV embeddings. This method works especially well with Mistral and Mixtral.
- •
StreamingLLM Xiao et al. (): In addition to the most recent tokens, StreamingLLM 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 8 Nvidia A100-40GB GPUs, interconnected by the NVLinks (600 GB/s GPU peer-to-peer bandwidth).

### 5.2 End-to-End Result

Figure 3 demonstrates the comparison results of SqueezeAttention and other 3 baseline algorithms 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-wise compression algorithms to different tasks would lead to quite different model accuracies. Therefore, for each task, we choose the best sequence-wise compression algorithm to represent the best case, and then apply SqueezeAttention on top of the best case. As shown in the figure, SqueezeAttention consistently manages to improve the model accuracies under various KV-cache budgets by reallocating the cache budgets among layers. In other words, SqueezeAttention 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 accuracy. We select Mistral-7B (Sliding Window), GPT-NeoX-20B (Heavy-Hitter), and LLama2-70B (StreamingLLM) 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 a single GPU.

Table LABEL:accuary_result demonstrates that in all three settings, by achieving the same model accuracy, SqueezeAttention 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 budget of H2O algorithm. Subsequently, we employ PYTORCH PROFILER to evaluate the diminished memory usage of generating one token during the inference (w/o memory usage of model weights). Figure LABEL:efficiency_result shows that SqueezeAttention 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.

**Table 2: 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-wise sparsification algorithm as the baseline.**
| Model | Size | Dataset | Best Baseline | Performance / Used KV Budget |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Full Cache | w/ SqueezeAttention | w/o SqueezeAttention |  |  |  |  |
| Mistral | 7B | SAMSUM | Sliding Window | 43.53 / 100% | 41.05 / 20% | 40 / 30% |
| GPT-NeoX | 20B | XSUM | Heavy-Hitter Oracle | 0.09 / 100% | 0.09 / 20% | 0.08 / 60% |
| LLama2 | 70B | XSUM | StreamingLLM | 0.18 / 100% | 0.17 / 30% | 0.19 / 40% |

### 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 for Mistral-7B, and 1 to 64 for Llama2. For each task, we select the best baseline algorithm, that is, Sliding Window for Mistral-7B and StreamingLLM for Llama2-70B.

Table LABEL:throughput_result shows the token throughput on 8 A100-40GB GPUs. With the same batch size, SqueezeAttention can enhance throughput by up to 2.2$\times$ for Mistral-7B and 1.4$\times$ 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 SqueezeAttention is reported in A.4.

**Table 3: Generation throughput (token / s) on eight A100 GPUs of Mistral-7B and LLama2-70B with SqueezeAttention and Full Cache. To maintain the model accuracy, SqueezeAttention uses 20% of the cache budget for Mistral-7B and 30% of the cache budget for LLama2-70B. ”OOM” means out-of-memory.**
| Model | Size | prompt len + gen len | Algorithm | Batch size |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1 | 32 | 64 | 128 | 224 |  |  |  |  |
| Mistral | 7B | 512 + 1024 | SqueezeAttention | 20.5 | 496.5 | 682.7 | 824.4 | 892.5 |
| Full Cache | 20.9 | 254.0 | 304.8 | OOM | OOM |  |  |  |
| Model | Size | prompt len + gen len | Algorithm | Batch size |  |  |  |  |
| 1 | 8 | 16 | 32 | 64 |  |  |  |  |
| LLama2 | 70B | 256 + 512 | SqueezeAttention | 5.2 | 37.2 | 71.2 | 116.2 | 170.7 |
| Full Cache | 5.2 | 36.0 | 62.5 | 84.8 | OOM |  |  |  |

### 5.5 Overhead of the Algorithm

The computational overhead of SqueezeAttention comes from two operations: Cosine similarity and Kmeans. The execution of all these computations only happens during the prefilling phase. Therefore, the cost of SqueezeAttention is a 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 SqueezeAttention to compare the wall-clock time. The prompt length is up to 8k tokens. As shown in Table 4, the SqueezeAttention only caused 6.3% increasement of prefilling time. Note that if we 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 SqueezeAttention is basically neglectable. The more detailed overhead of SqueezeAttention can be found in 5.

**Table 4: Overhead of SqueezeAttention (Prefilling Time in seconds on one Nvidia A100-40GB)**
| Model | w/o SqueezeAttention | 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 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 inference. Besides, it can optimize the user experience of LLM applications. Since the proposed algorithm 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 a 2D KV-cache compression algorithm called SqueezeAttention. 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 observation, SqueezeAttention reallocates the KV-cache budgets over attention layers to further reduce the memory cost of inference. Experiments over a wide range of models and tasks show that SqueezeAttention can achieve better model accuracies with lower memory consumption compared with state-of-the-art algorithms that only compress KV-cache in a sequence-wise way.

## Appendix A Appendix

### A.1 detailed overhead of squeezeattention

In this section, we further evaluate the overhead introduced by SqueezeAttention, we broke down the time taken for two primary operations: cosine similarity computation and K-means clustering. The experiment setup follows the same setup as 5.5. 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 8000×4096, repeated 32 times (since Mistral has 32 layers). Additionally, K-means clusters 32 numbers into 3 classes. Therefore, the total time consumption can be calculated as 0.00068×32+0.001=0.02276 seconds. It is noteworthy that this overhead is incurred only once, regardless of the number of tokens processed.

### A.2 the function of p 𝑝 p italic_p

The hyperparameter $p$ is crucial in the SqueezeAttention mechanism, as it directly determines the final allocation of the KV cache. To illustrate this with a concrete example: suppose we have a model with 32 layers, where 18 layers are deemed important and the remaining 14 are considered less important. Each layer initially has a budget of 1000 tokens. If we set $p$ to 0.3, we will take 70% of the budget from the less important layers and redistribute it equally among the important layers.

Specifically, the budget for the less important layers will be reduced to 1000×0.3=300. Meanwhile, the budget for the important layers will be calculated as follows: (1000×18+1000×0.7×14)/18=1544. In this way, the total budget remains unchanged.

To evaluate the sensitivity of the hyperparameter $p$, we tested the Mistral-7B model on 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 budget will be transferred to other layers, which can yield better results up to a certain point. When $p$ decreases to 10% or even lower, the performance decreases as the less important layers’ KV budget becomes severely inadequate. We can clearly observe an accuracy improvement when adjusting $p$ while keeping the overall budget unchanged. This demonstrates the impact of $p$ on model performance, 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.

**Table 7: Mistral-7B**
| Dataset | Samsum | TriviaQA | LCC |
| --- | --- | --- | --- |
| Important | 17 | 18 | 19 |
| Unimportant | 15 | 14 | 13 |

Tabel8 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 8 shows the distribution of important layers for the LLama2-70B model across three different datasets: Xsum (Summarization), Samsum (Few shot), and LCC (Code, Python/C#/Java).

From these tables, we can observe that there is a 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 are also some layers that are always important / unimportant. A detailed analysis of this phenomenon could be an interesting extension of this work. However, we would still recommend an adaptive way since it can precisely capture the importance of layers.

### A.4 throughput comparison between best baseline algorithms and SqueezeAttention

**Table 9: 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, SqueezeAttention uses 20% of the cache budget for Mistral-7B and 40% of the cache budget for LLama2-7B while Sliding Window uses 30% of the cache budget for Mistral-7B and StreamingLLM use 60% of the cache budget for LLama2-7B. ”OOM” means out-of-memory.**
| Model | Size | prompt len + gen len | Algorithm | Batch size |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1 | 32 | 64 | 128 | 224 |  |  |  |  |
| Mistral | 7B | 512 + 1024 | SqueezeAttention | 20.5 | 496.5 | 682.7 | 824.4 | 892.5 |
| Sliding Window | 20.6 | 404.5 | 512.2 | 587.8 | OOM |  |  |  |
| Model | Size | prompt len + gen len | Algorithm | Batch size |  |  |  |  |
| 1 | 32 | 64 | 128 | 256 |  |  |  |  |
| LLama2 | 7B | 512 + 1024 | SqueezeAttention | 20.0 | 143.0 | 150.4 | 144.9 | OOM |
| StreamLLM | 20.4 | 113.7 | 102.4 | OOM | OOM |  |  |  |

We also conducted experiments to compare throughput of SqueezeAttention with best baseline under a set of batch sizes. Both experiments used an input length of 512 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 compared with those SOTA algorithms that only compress KV-cache from the sequence’s dimension.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. What inference framework is used for the memory and throughput experiments? Is SqueezeAttention compatible with current inference memory optimization like vllm[1]?

[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).
soundness: 3
strengths: 1. The method can augment other KV-Cache eviction policies, which will benefit the research community.
2. The algorithm is clearly presented and the method's effectiveness has strong experiment evidence.
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.
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.
2. 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?

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: - Do G1,G2,G3 changes frequently across different samples? otherwise we can assign the layer-wise budget through a offline process.
soundness: 2
strengths: - The proposed methods is evaluated with multiple LLMs on various downstream tasks, demonstrates non-trivial improvements against previous baselines.

- The manscript is clearly organized with several illustration figures and equations. It's easy to understand the main method of this work.

- Both perfomance comparison and end-to-end memory/thoughput comparison are reported.
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.
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].

- 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.

- 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.
 

[1] https://arxiv.org/abs/2312.17276

[2] https://proceedings.neurips.cc/paper_files/paper/2023/file/fde1a69a5b6e554b2f1f727197d2651d-Paper-Conference.pdf

[3] https://arxiv.org/pdf/2202.08625

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 2
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.
soundness: 3
strengths: The observations of comparing the inputs and outputs of attention modules are good.
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.
weaknesses: 1.	The 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.
2.	The 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.
3.	The 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.
4.	The 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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: See the discussion of weaknesses and kindly address them.
soundness: 4
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.
Significant Performance Improvement: The proposed method improves memory consumption and throughput by reallocating cache budgets based on layer importance.
Robust Experimental Validation: The authors test their approach on multiple models (ranging from 7B to 70B parameters) and datasets, demonstrating its generalizability and efficiency.
Compatibility with Other Methods: SqueezeAttention integrates smoothly with various sequence-wise compression techniques, enhancing its versatility.
Energy Efficiency: The memory and throughput improvements have practical implications, potentially reducing the environmental impact of LLM deployment.
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.
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.
Potential 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.
Limited 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.
Fixed 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.
Risk 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.
In 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.

Meta-review: The paper presents a practical approach of layer-wise dynamic allocation for KV compression,  based on their importance determined by the cosine similarity of embeddings before and after self-attention layers. The method is compatible with other sequence-based compression algorithms, augmenting their performance by optimizing layer-level cache budgets. In their experiment, the method achieves 30%-70% memory reduction and up to 2.2× throughput improvement across diverse models (i.e., Llama2, Mistral, Falcon, OPT, GPT-Neox, etc) combined with 3 representative sequence-wise compression algorithms (i.e. H2O, Sliding window and Streaming LLM).

Reviewers generally praised the paper for its comprehensive evaluation and practical benefits. There is initially a missing comparison with layer-adaptive compression methods like PyramidKV which weakens the experimental design, but the authors conducted experiments during the rebuttal period that provided a comparison. Another shortcoming comes from fixed three-group clustering and fixed design of compression ratio which in principle may limit its adaptability for broader applications. However, the concern is alleviated by the extensive experiments that demonstrated the strong performance of their method. 

Overall, the paper represents a solid study of a layer-adaptive KV compression technique, which could be of great practical interest. The paper is relatively weak in terms of novelty as layer-adaptive strategies have been explored before in PyramidKV, and has a relatively limited scope since it does not have much broader implications beyond KV compression. Therefore I recommend an acceptance for a poster presentation.

Decision: Accept (Poster)