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

Authors: Zihao Wang, Bin CUI, Shaoduo Gan
Primary area: generative models
Keywords: KV-cache, LLM inference optimization

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.

Decision: accept (poster)
Reviewer scores: 6, 3, 5, 8

Reviews and discussion:
<thread>
creation date:
1730196894636

thread_type:
official_review

<thread_object>
creator: reviewer_pb9y

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

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.

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.

questions:
See the discussion of weaknesses and kindly address them.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730658396694

thread_type:
official_review

<thread_object>
creator: reviewer_fgax

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

strengths:
The observations of comparing the inputs and outputs of attention modules are good.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730869314315

thread_type:
official_review

<thread_object>
creator: reviewer_2bq8

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

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.

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

questions:
- Do G1,G2,G3 changes frequently across different samples? otherwise we can assign the layer-wise budget through a offline process.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731266940229

thread_type:
official_review

<thread_object>
creator: reviewer_arng

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

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.

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?

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).
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.