CAKE: Cascading and Adaptive KV Cache Eviction with Layer Preferences

Authors: Ziran Qin, Yuchen Cao, Mingbao Lin, Wen Hu, Shixuan Fan, Ke Cheng, Weiyao Lin, Jianguo Li
Primary area: generative models
Keywords: Large Language Model, Efficient Generative Inference, Key-Value Cache

Abstract:
Large language models (LLMs) excel at processing long sequences, boosting demand for key-value (KV) caching. While recent efforts to evict KV cache have alleviated the inference burden, they often fail to allocate resources rationally across layers with different attention patterns. In this paper, we introduce Cascading and Adaptive KV cache Eviction (CAKE), a novel approach that frames KV cache eviction as a ``cake-slicing problem.''
CAKE assesses layer-specific preferences by considering attention dynamics in both spatial and temporal dimensions, allocates rational cache size for layers accordingly, and manages memory constraints in a cascading manner. This approach enables a global view of cache allocation, adaptively distributing resources across diverse attention mechanisms while maintaining memory budgets.
CAKE also employs a new eviction indicator that considers the shifting importance of tokens over time, addressing limitations in existing methods that overlook temporal dynamics.
Comprehensive experiments on LongBench and NeedleBench show that CAKE maintains model performance with only 3.2\% of the KV cache and consistently outperforms current baselines across various models and memory constraints, particularly in low-memory settings. Additionally, CAKE achieves over 10$\times$ speedup in decoding latency compared to full cache when processing contexts of 128K tokens with FlashAttention-2. Our code is available at https://github.com/antgroup/cakekv.

Decision: accept (poster)
Reviewer scores: 6, 6, 10, 6

Reviews and discussion:
<thread>
creation date:
1730524916908

thread_type:
official_review

<thread_object>
creator: reviewer_zgxb

<content>
summary:
This work introduces CAKE, a method for efficient KV cache management in LLMs that enhances inference by dynamically allocating cache based on each layer’s spatial and temporal attention demands. By framing cache eviction as a cake-slicing problem, CAKE optimally distributes resources across layers and incorporates a novel eviction indicator to account for the shifting importance of tokens over time. Extensive experiments show the potentials of CAKE.

strengths:
- Pioneering Adaptive Memory Budgets: CAKE is the first work to consider adaptive memory budgets for different layers of LLMs. This innovative approach allows for more efficient memory utilization, improving model performance by allocating resources where they are most needed.

- Addressing KV Cache Compression: The paper tackles the timely problem of KV cache compression in LLMs, which is especially relevant for on-device applications. By focusing on this issue, the work makes LLMs more practical and accessible in resource-constrained environments.

- Clear Writing: The overall writing is clear, though the paper structure is a bit chaotic. This clarity facilitates comprehension, replication, and further research based on the paper's findings.

weaknesses:
- The motivation for using spatial dispersion and temporal shift in cache size allocation is unclear. Providing more insights or intuition would help clarify its benefits. Additionally, Table 2 shows that the adaptive allocation strategy provides minimal improvement, suggesting it may not be necessary.

- The adaptive KV compression method is incompatible with flash-attention. Given that flash-attention is widely adopted for efficient training and inference of LLMs, it’s unlikely that practitioners would choose CAKE over flash-attention in practice.

- Only two LLM backbones, Llama and Mistral, were evaluated, which may be insufficient. Consider adding another backbone, such as Phi, Qwen, or Gemma, to strengthen the analysis.

- The team has not open-sourced their code, which could raise concerns about the reproducibility of their work.

- The baseline methods used for comparison are not sufficiently strong. Consider including the following more robust methods (For KVQuant[1] and KIVI[2] should be jointly used with the flash-attention, and GEAR could remain the original settings):

[1] Hooper, Coleman, et al., "Kvquant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization." arXiv preprint arXiv:2401.18079 (2024).

[2] Liu, Zirui, et al., "KIVI: Plug-and-Play 2bit KV Cache Quantization with Streaming Asymmetric Quantization." (2024).

[3] GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLM.

questions:
Please refer to the weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730547812505

thread_type:
official_review

<thread_object>
creator: reviewer_brwg

<content>
summary:
A novel technique that solves KV cache management to improve computational efficiency considering spatial and temporal attention dynamics. 10× faster decoding for extended sequences. Layer-wise memory budget allocation.

strengths:
S1. Identified and visualized the attention dynamics in spatial and temporal axis

S2. Strong empirical analysis on multiple datasets

S3. Actual improvement in inference latency

S4. Excellent presentation (great paper flow and beautiful figures)

weaknesses:
I do not identify any major weaknesses of the paper.

questions:
Q1. I'm just curious how this technique could scale to larger models (no need to verify it empirically)

Q2. Could this technique potentially help transformers stay competitive against RNN-based models like Mamba?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730576639544

thread_type:
official_review

<thread_object>
creator: reviewer_ljxw

<content>
summary:
This paper introduces a method for optimizing KV cache eviction through a cache allocation strategy to enhance LLM inference efficiency. The proposed cache allocation adapts to layer preferences, adjusting KV cache injection to improve efficiency while maintaining satisfactory performance. Extensive experiments are conducted to demonstrate the method’s effectiveness.

strengths:
The paper is well-written, with clear motivations for the proposed method and pipeline illustrations. Incorporating layer-wise preference modeling to guide KV caching strategies is intuitive, given the insights from attention dynamics analysis.

The proposed method is straightforward and compatible with existing KV caching strategies, making it easy to integrate while achieving decent efficiency. The authors provide ample experiments to substantiate this point.

Empirical analysis is comprehensive, covering 16 tasks with various LLMs of different specifications, offering a comprehensive evaluation.

weaknesses:
Most empirical analyses focus on smaller LLMs with 7B-8B parameters, which may limit the generalizability of this approach for much larger LLMs. Specifically, it would be valuable to see how computational costs and performance are impacted across different LLM sizes.

The empirical improvements over existing baselines are relatively modest, which could suggest limited practical advantages in some cases.

questions:
Please see weakness aspects above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731287506036

thread_type:
official_review

<thread_object>
creator: reviewer_rcqd

<content>
summary:
This paper proposes Cascading and Adaptive Key-value cache Eviction (CAKE) method for optimizing Key-Value cache evicting in large language models. Specifically, CAKE assesses each layer’s KV cache needs by considering attention dynamics in both spatial and temporal dimensions. During the prompt prefilling, CAKE allocates rational cache size for layers by analyzing layer-specific KV cache preferences and manages the memory budgets with the guidance of these preferences in a cascading manner. Besides, CAKE introduces a novel eviction indicator that accounts for both the long-term influence and temporal variability of token importance. Extensive experiments demonstrate CAKE’s superior performance across different models and memory constraints, especially in low-memory scenarios.

strengths:
(1)	This paper provides a novel and practical key-value cache eviction approach to enhance LLM’s proficiency in handling long sequences, based on layer-specific KV cache preferences in a cascading manner.

(2)	The paper is well-organized, and the writing is clear. In particular, Figure 2 clearly points out the differences between the proposed CAKE and other existing models, and I find almost no typos in this paper.

(3)	The theoretical analysis (e.g. Theorem 2) rigorously demonstrates the equivalent KV cache eviction results of the proposed preference-guided cascading cache management to the vanilla preference-prioritized adaptive allocation strategy.

(4)	The extensive experiments on several open-source LLM benchmarks truly validate the effectiveness of the proposed algorithms.

(5)	Overall, the proposed algorithm CAKE model is novel, practical and efficient. The corresponding experimental results are extensive and sound.

weaknesses:
(1)	The abbreviation “KV” in the title and abstract (e.g. line 12, or the second line of abstract) should be clearly written as “Key-value”, as this word appears for the first time in the whole paper.

(2)	In line 23, the abstract section says “this approach allows for a global view of cache size allocation, distributing resources OPTIMALLY”. The word “optimally” is somewhat controversial, until you can demonstrate rigorously from a theoretical point that the proposed resources distribution method is optimal (with respect to certain theoretical property). Therefore, I would suggest to use less controversial word, like “adaptively”.

(3)	Equation (4) is a little confusing for me. Since $A[i,:]$ is a row vector, $\log{A[i,:]}$ is also a row vector, then is $ A[i,:] \log{A[i,:]}$ the inner product of these two vectors? Should give more clear explanations.

(4)	I read the proofs line by line, and according to my experience, the proof is sound. However, since the proof of Theorem 1 is truly basic and short, I believe it will be better to describe “Theorem 1” as a proposition. Besides, in Theorem 1, “For layer $l \in [L]$”, does it mean for any (fixed) layer, or mean there exists a layer? More explanations should be clarified.

(5)	In Table 1 of the experimental results, the proposed CAKE obtains SOTA results in most benchmark datasets. But in some benchmarks, some existing models like TOVA and SnapKV can achieve better results. Could you give more analysis in which kind of benchmark datasets can CAKE obtain better experimental results than the existing baselines?

questions:
In Table 1, CAKE could not outperform existing methods on some benchmarks. Could you give more analysis in which kind of benchmark datasets can CAKE obtain better experimental results than the existing baselines? Or in what situations, CAKE will fail to achieve good results?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732370828417

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary and general reply to the reviewers

comment:
We sincerely appreciate all reviewers' time and efforts in reviewing our paper. Your constructive feedback has substantially helped improve the quality of our work.
We are particularly encouraged that the reviewers have recognized our key contributions in several aspects:

* Novel and Practical Contribution (Reviewer-rCQd, BrWg, zgXb)

* Strong Technical Foundation (Reviewer-rCQd, LJXW, BrWg)
* Comprehensive Empirical Validation (Reviewer-rCQd, LJXW, BrWg)
* Clear Presentation and Organization (Reviewer-rCQd, LJXW, BrWg, zgXb)
  
In response to the reviewers' suggestions, we have made the following major improvements:

* Extended evaluation to **additional LLM architectures** (Qwen and Gemma), with detailed results provided in **Appendix F.2**. (Reviewer-zgXb)
* Conducted experiments on **larger models** ranging from 13B to 70B parameters (Llama2-13B, Qwen2.5-32B, and Llama3-70B), with comprehensive results presented in **Appendix F.3**. (Reviewer-LJXW, BrWg)
* Added discussion on orthogonal KV cache quantization methods in **Appendix E**, demonstrating that our work is **compatible with and complementary** to these approaches. (Reviewer-zgXb)
* Addressed all other clarification requests from the reviewers. (Reviewer-rCQd, zgXb)

All major modifications have been **highlighted in blue** in the revised manuscript for easy reference. We believe these changes have significantly strengthened our paper and addressed the reviewers' concerns.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732539413626

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Kind Request for Discussions and Feedback for Paper 9488

comment:
Dear Reviewers,

We are deeply grateful for your thorough reviews and valuable feedback on our paper.

As the discussion period nears its end, we hope our responses have effectively addressed the points you raised. Your insights have been instrumental in strengthening our work, and we welcome your continued engagement.

Thank you for your dedication and expertise.

Best regards,

Authors of Paper 9488
</content>
</thread_object>

</thread>


Metareview:
metareview:
All reviewers agreed the paper proposed a useful contribution to KV cache eviction strategy to speed up decoding process of LLM. 
Strength:
1. The proposal method appeared to be novel and useful.
2. Good experimental results

Weakness:
1. Results were only on small models, limiting its impact on practical use with larger models (more results came in rebuttal period).
2. Baseline for comparisons were not particulartly strong. (Improvements were made during rebuttal.)