CAKE: Cascading and Adaptive KV Cache Eviction with Layer Preferences

We greatly appreciate your thorough review and detailed suggestions. Addressing your comments has helped us strengthen both the technical content and clarity of our presentation. Below we address your concerns point by point. Corresponding modifications in the paper are highlighted in blue.

W1: The motivation for using spatial dispersion and temporal shift.

Our approach is motivated by a key observation: attention patterns exhibit significant variations across layers, models, and contexts, as demonstrated by our extensive visualizations in Appendix J. This characteristic makes it suboptimal to employ fixed or uniform cache allocation strategies. Through our analysis, we identify two crucial aspects of attention patterns: spatial dispersion and temporal shift. Spatially, we observe that some layers distribute attention broadly across tokens (Fig. 1(a)), while others concentrate on specific tokens (Fig. 1(b)). Layers with broader attention dispersion, when compared with full cache settings, require larger cache sizes to maintain their performance. However, spatial patterns alone cannot fully capture attention dynamics. Temporally, we find that some layers shift their attention focus across different tokens during different steps (Fig. 1(c)), while others maintain fixed attention to tokens (Fig. 1(d)). Layers with dynamic temporal patterns need larger cache allocations to effectively track these changes. Given these observations, we propose that effective cache allocation must consider both spatial dispersion and temporal shift to accurately measure each layer's cache requirements. This comprehensive approach enables CAKE to adaptively allocate resources based on layer preferences, adapting to varying attention patterns in different models and contexts.

Indeed, the above analyses have been provided in 67-75 lines of the submitted paper. Please kindly refer to it. We trust that this explanation adequately addresses your concerns. Should you require additional clarification, please do not hesitate to inform us, as we are more than willing to provide it.

W2: Table 2 shows that the adaptive allocation strategy provides minimal improvement.

To better demonstrate the effectiveness of our adaptive allocation strategy, we have conducted additional experiments across multiple models, comparing three allocation strategies: uniform allocation, pyramid allocation, and our preference-prioritized adaptive allocation (P2A). The table below presents average scores across 16 LongBench datasets with a total budget size of 128L:

Compared with uniform allocation, pyramid allocation only shows modest improvements on the Llama2 model, it actually suffers performance degradation on other models. This demonstrates a key limitation of fixed-pattern allocation strategies: they rely on prior observations that may not generalize across different model architectures due to varying attention patterns. In contrast, our P2A strategy consistently outperforms both uniform and pyramid allocation across all tested models. This consistent improvement stems from P2A's ability to effectively measure layer-specific cache preferences by analyzing both spatial and temporal characteristics of attention patterns, enabling it to adaptively allocate appropriate cache sizes to corresponding layers.

W3: Compatibility with Flash-Attention.

We appreciate the reviewer's concern but want to clarify that CAKE is directly implemented on top of Flash-Attention.

While implementing our KV cache eviction method, we still use Flash-Attention's "_flash_attention_forward" for full attention computations. The only additional computation needed is to obtain attention weights from the observing window $\mathbf{A}[-S_w:,:]$ for calculating preference scores and eviction indicators, which can be efficiently computed via $\mathbf{A}[-S_w:,:] = \text{Softmax}(\frac{\mathbf{Q}[-S_w:,:]\mathbf{K}^T}{\sqrt{D}})$. For a 32K input sequence, this local attention computation introduces negligible overhead (0.1% of full attention). Therefore, CAKE fully preserves the efficiency benefits of Flash-Attention while providing additional memory optimization through adaptive cache management. To better illustrate the compatibility, we provide a detailed PyTorch-style implementation with Flash-Attention in Appendix C, Listing 1.

W4: Evaluation on more backbones.

Following your advice, we have incorporated experiments on two additional backbones, Qwen and Gemma, to further strengthen our analysis. Specifically, we conducted comparisons under both low-budget settings ($B_{\text{total}} = 128L$) and high-budget settings $(B_{\text{total}} = 1024L)$ across 16 datasets on LongBench. The average results are presented as follows, with detailed results provided in Appendix F.2.

As can be seen, CAKE consistently outperforms other baselines across both low-memory and high-memory scenarios, even with the inclusion of the new backbones. Notably, Gemma with $B_{\text{total}} = 1024L$ achieves a performance that surpasses the full-cache baseline (34.18 vs. 34.09). In addition to expanding the range of model architectures, we also conducted experiments with larger model sizes, including 13B, 32B, and 70B. CAKE continues to deliver the best performance in these settings, with detailed experimental results available in Appendix F.2. We believe these results highlight the effectiveness and generalizability of our proposed approach.

W5: Code Reproducibility

We appreciate your valuable feedback regarding code availability and fully understand the importance of open-sourcing for ensuring reproducibility. To address this, we are actively preparing the code and relevant documentation for public release. We will ensure that our work can be fully reproduced by the research community and plan to make the codebase available upon the acceptance of this paper.

W6: Comparison with quantization methods.

We appreciate the reviewer's suggestion. It's important to note that CAKE focuses on KV cache eviction through dropping unimportant KV pairs, which is orthogonal to KV cache quantization methods that aim to reduce storage overhead through bit reduction. Both CAKE and quantization methods can be jointly used with flash-attention. More importantly, CAKE is compatible with quantization methods to pursue more efficient KV cache storage as we evaluate in the following part. We have conducted additional experiments on Llama2-7B-Chat comparing CAKE with two typical quantization methods: 1) KIVI[1], a state-of-the-art KV cache quantization method that adopts asymmetric KCVT quantization (quantizes Key cache per-channel and Value cache per-token, similar patterns are also adopted by KVQuant [2] and GEAR [3]), and 2) KCVC, which quantizes KV cache both per-channel for efficiency. Due to time constraints during rebuttal, we compare with KIVI as a representative method, since it shares similar quantization schemes with KVQuant and GEAR. The experimental results on LongBench are summarized as follows:

For a fair comparison, we evaluate different methods under the same compression ratios. CAKE achieves better performance (33.23) than full cache (33.07), at a compression ratio of 50%; At 25% and 12.5% compression ratios, CAKE consistently outperforms both KIVI and KCVC; Most importantly, combining CAKE with KIVI-INT4 achieves better results at 12.5% and 6.25% compression ratios compared to KIVI-INT2 alone. These results validate that: (a) CAKE is effective as a standalone method, (b) CAKE can work synergistically with quantization approaches, (c) The combination enables even higher compression ratios while maintaining performance. We have added this discussion in Appendix E.

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

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

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

We are deeply grateful for your thorough review and strong endorsement of our work's technical novelty, presentation quality, and empirical analysis. We appreciate your insightful questions, which we have carefully addressed below.

Q1: How this technique could scale to larger models.

This question is of such high quality that it motivates us to delve into the potential of larger models, even though it was not a requirement from the reviewer.

Although attention patterns vary significantly across different model architectures and sizes, our method can universally analyze their spatial and temporal attention characteristics and adaptively accommodate different model scales and architectures. This adaptability enables our approach to demonstrate robust performance across various model sizes and architectures. To empirically validate this generalizability, we have conducted additional experiments on larger models including Llama2-13B-Chat, Qwen2.5-32B-Instruct, and Llama3-70B-Instruct. We evaluated both low-budget (128L) and high-budget ($1024L$) settings on Longbench. The table below presents the average scores across 16 datasets (detailed results are provided in Appendix F.3).

Across different model sizes, CAKE consistently outperforms other methods, and notably, under the 1024L setting, CAKE even achieves better performance than full-cache settings for both Llama2-13B and Llama3-70B. We have further validated our method's effectiveness on additional model architectures, including Qwen and Gemma, with detailed results available in Appendix F.2.

Q2: Could this technique potentially help transformers stay competitive against RNN-based models like Mamba?

Yes, CAKE could potentially help transformers remain competitive against RNN-based models like Mamba by addressing several critical challenges:

1. Memory Efficiency: While Mamba achieves linear memory scaling, CAKE significantly reduces the Transformer's memory footprint by maintaining a fixed KV cache budget without sacrificing performance. Additionally, CAKE is compatible with efficient attention mechanisms such as FlashAttention, further improving inference efficiency.

2. Decoding Speed: CAKE enables faster decoding for long sequences (up to 10x speedup for 128K sequences) through optimized cache management, narrowing the speed gap with Mamba.

3. Attention Capabilities: A key advantage of transformer-based models over Mamba lies in their stronger ability to handle complex contexts. CAKE preserves this strength while enhancing efficiency.

While CAKE makes LLMs more efficient for generation, our analysis reveals that attention mechanisms can exhibit redundancy, as only a subset of information is required for effective inference. This insight suggests the potential for combining the strengths of both architectures to develop a more efficient hybrid framework.

We sincerely appreciate your positive assessment of our work. Your valuable feedback has helped us improve the quality and completeness of our paper. We have carefully addressed each of your comments below, with the corresponding modifications highlighted in blue in the revised manuscript.

W1: Most empirical analyses focus on smaller LLMs with 7B-8B parameters, which may limit the generalizability of this approach for much larger LLMs.

To demonstrate the generalizability of our approach across different model sizes, we have conducted additional experiments on larger models including Llama2-13B-Chat, Qwen2.5-32B-Instruct, and Llama3-70B-Instruct. We evaluated both low-budget ( cache size = $128L$) and high-budget (cache size = $1024L$) settings on Longbench. The table below presents the average scores across 16 datasets (detailed results are provided in Appendix F.3).

As shown in the results, CAKE consistently outperforms baseline methods across different model sizes. Under constrained memory conditions (cache size = $128L$), CAKE demonstrates significant advantages over other methods. These advantages are maintained with larger cache sizes ($1024L$), where CAKE even achieves slightly better performance than full-cache settings for some models (Llama2-13B: 29.98 vs 29.95, and Llama3-70B: 45.83 vs 45.79). These results indicate that our approach scales effectively to larger models while maintaining its efficiency advantages.

Additionally, we have validated CAKE's effectiveness on two more model architectures, including Qwen and Gemma, further demonstrating the generalizability of our proposed method. Detailed results can be found in Appendix F.2.

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

Our initial submission might show modest improvements in some aspects. However, by incorporating our new experiments to address W1, we wish to emphasize that CAKE offers significant practical benefits in four critical areas:

1. CAKE consistently outperforms baselines across different model architectures (Llama, Mistral, Qwen, Gemma) ranging from 7B to 70B parameters on LongBench across all memory settings (Appendix F).

2. CAKE shows significant advantages in memory-constrained settings. For example, with cache size $64L$ on Mistral, CAKE achieves 34.31 average score on LongBench versus SnapKV's 31.31 and PyramidKV's 30.50 (Table 7, lines 1147-1149), highlighting its efficiency in resource-limited scenarios. Similar cases can be also found in other models.

3. CAKE not only narrows the gap with full cache but sometimes surpasses it while maintaining minimal memory, as demonstrated on Gemma-7B, Llama2-13B, and Llama3-70B with only $1024L$ cache size.

4. CAKE significantly outperforms existing methods on challenging Multi-Retrieval tasks. For example, with $1024$ cache size on Mistral, CAKE achieves 71.00 versus SnapKV's 56.93 (Table 10, lines 1340-1342) and maintains 48.55 accuracy versus SnapKV's 19.14 (Table 10, lines 1344-1346) with $512$ cache size. More similar cases can be found in Appendix G.

We sincerely appreciate your thorough review and positive assessment of our work. Your constructive feedback is invaluable in helping us improve both the technical content and presentation clarity of our paper. We have carefully addressed each of your comments below, with the corresponding modifications highlighted in blue in the revised manuscript.

W1&W2: Suggestions on Terminology and Wording

We appreciate your valuable feedback on our writing. We agree with both suggestions and will revise "KV" to "Key-value (KV)" on its first appearance in the abstract. We also concur that replacing "optimally" with "adaptively" more accurately reflects our method's capabilities. Thank you for helping us improve the precision and clarity of our paper. These changes have been incorporated in the revised version of our paper.

W3: Clarification on Equation (4).

We regret the error in our previous statement. The correct operation involves transposing $\text{log}\mathbf{A}[i, :]$ and then computing the inner product with $\mathbf{A}[i, :]$ by using the expression $\mathbf{A}[i, :]\text{log}(\mathbf{A}[i,:])^T$. We have made the necessary correction in our paper and are grateful for your meticulous review.

W4: Clarification on Theorem 1.

Thank you for your careful review of the mathematical presentation. We have changed "Theorem 1" to "Proposition 1" given its straightforward proof. In Proposition 1, "For layer $l∈[L]$" means the allocated budget size decreases monotonically from stage $l$ to $L-1$ for any layer with index in $[0,1,...,L-1]$. We have modified this to "For any layer $l∈[L]$" for better clarity.

W5 & Q1: Could you give more analysis in which kind of benchmark datasets can CAKE obtain better experimental results than the existing baselines?

We appreciate your insightful question. This is an excellent point that helps us better articulate the comprehensive strengths of CAKE across diverse benchmarks. Our analysis demonstrates CAKE's robust performance and specific advantages in different types of evaluations:

On LongBench tasks (detailed in Table 1 and Appendix F), CAKE demonstrates consistent performance improvements across various cache sizes and task types. While other methods occasionally show marginal advantages in specific cases, none exhibits consistent superiority across all conditions, highlighting CAKE's robust general performance.

For NeedleBench (detailed in Appendix G), CAKE's advantages become more pronounced in complex tasks, particularly in Multi-Needle Retrieval. For instance, compared to the previous SOTA SnapKV, CAKE achieves significantly better results on Mistral: 71.00 vs. 56.93 ($B_\text{total}=1024L$, lines 1340-1342) and 48.55 vs. 19.14 ($B_\text{total}=512L$，lines 1344-1346) as shown in Table 10 of Appendix G. Such advantage becomes increasingly pronounced as the cache budget decreases. CAKE's superior performance in these challenging scenarios stems from its design, which considers both long-term significance and short-term fluctuations in information relevance. In contrast to existing methods, CAKE employs a more nuanced approach, maintaining a more comprehensive and balanced representation of contextual information, instead of depending solely on static attention score or fixed cache allocation strategies, thus avoiding premature discarding of information that could be vital for subsequent complex retrievals.

Nevertheless, as an eviction strategy, under extremely constrained budgets, CAKE's performance inevitably experiences some degradation. To address this, we suggest combining KV Cache quantization with CAKE eviction. Our experiments illustrate that this hybrid approach proves adept at maintaining performance under severe memory limitations, as detailed in Appendix E.

Dear reviewer rCQd,

Could you please respond to authors' rebuttal and see if you would like to update your review? Thanks very much!

AC