PaperlayerView exact revision on arXiv

Source-linked AI summary

CAKE: Cascading and Adaptive KV Cache Eviction with Layer Preferences

Ziran Qin, Yuchen Cao, Mingbao Lin, Wen Hu, Shixuan Fan, Ke Cheng, Weiyao Lin, Jianguo Li

arXiv:2503.12491v2cs.CL

TL;DR

Long-context inference creates growing KV-cache memory demands, while existing methods often overlook layer-specific attention patterns. CAKE adaptively allocates cache using spatial and temporal attention preferences, manages budgets cascadingly, and uses a temporally informed eviction indicator. Across LongBench and NeedleBench, it preserves performance with 3.2% of the cache and outperforms existing methods, especially under low memory.

  • Problem

    Existing KV-cache eviction methods often allocate resources uniformly across layers despite differing attention patterns, complicating efficient inference under long-context memory constraints.

  • Method

    CAKE combines preference-prioritized adaptive allocation, cascading cache management, and token eviction based on sustained importance and temporal attention variability.

  • Results

    Across LongBench and NeedleBench, CAKE consistently outperforms baselines across models and memory constraints, preserving performance with only 3.2% of the cache.

  • Takeaways & Limitations

    CAKE provides layer-adaptive KV-cache management for long-context tasks while improving memory and inference efficiency, particularly in low-memory settings.

  • Takeaways & Limitations

    CAKE does not address finer-grained attention dynamics within layers, such as head-level variation.

Abstract

from arXiv · show

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

1 INTRODUCTION

Longer contexts make KV caches grow linearly, while existing eviction methods often allocate cache uniformly across layers despite differing attention requirements. CAKE addresses this with preference-aware allocation, cascading memory management, and temporally informed token eviction.

  • Beyond 128K-token capacities, expanding contexts increase KV-cache memory burdens during inference.
  • Uniform cache sizes can impair performance under memory constraints because layers have different cache requirements.
  • CAKE measures layer preferences using spatial attention dispersion and temporal attention shifts to adapt cache allocation across layers.
  • CAKE uses cascading memory management to adjust layer budgets during prefilling without storing all layers’ KV caches simultaneously.
  • Its eviction indicator combines sustained token importance with attention variability to reduce adverse effects on later decoding steps.

2 BACKGROUND AND RELATED WORKS

KV caching reduces repeated attention computation but grows with sequence length, motivating eviction and more rational allocation under finite memory budgets. CAKE contrasts uniform and fixed-pattern strategies with adaptive allocation based on layer preferences.

  • Uniform allocation assigns equal cache sizes across layers, whereas fixed-pattern and CAKE strategies vary allocation across layers.
  • CAKE’s preference-prioritized adaptive allocation adjusts layer ratios for different contexts and models under specified memory budgets.
  • KV caches are updated by concatenating each newly generated key-value pair, avoiding repeated key-value projections during decoding.
  • Linear KV-cache growth makes efficient management difficult for extremely long input or output sequences within finite cache budgets.
  • Early eviction methods retain sequence-boundary tokens, risking the omission of important tokens located in the middle.

3 INSIGHTS INTO ATTENTION DYNAMICS

The paper analyzes attention weights across spatial and temporal dimensions to characterize how models distribute and shift attention. These dynamics vary substantially across layers, models, and contexts, motivating tailored cache management.

  • Spatial attention dispersion describes how one query token’s attention spreads across other tokens, while temporal attention shift tracks changes in highly attended tokens over time.
  • Spatial Analysis: Higher entropy of an attention row indicates a more even distribution, whereas lower entropy indicates focused attention.
  • Temporal Analysis: Temporal attention shift is measured through column-wise variance, with higher values indicating larger attention changes across positions.
  • LongBench analyses find substantial variation in attention dispersion and shift across layers, models, and contexts.
  • These attention differences indicate a need for cache-management strategies tailored to dynamic layer behavior.

4 METHODOLOGY

CAKE allocates KV-cache budgets according to layer-specific attention preferences, then cascades memory management across prefilling stages while preserving the standard strategy’s eviction result. Its eviction indicator also accounts for sustained importance and attention variability to tolerate shifting attention.

  • Preference-Prioritized Adaptive Allocation: Layer preference scores combine spatial attention dispersion and temporal attention shifts to estimate each layer’s KV-cache requirement.Temperature parameters adjust the influence of dispersion and shift, using a recent prefilling-attention window.
  • Preference-Prioritized Adaptive Allocation: Normalizing layer preference scores allocates the total cache budget adaptively across layers.The resulting cache sizes reflect each layer’s characteristics and the input context.
  • Preference-Guided Cascading Cache Management: Peak memory is reduced to the target by partitioning prefilling into layer-wise stages and redistributing budgets as preference scores become available.The cascading procedure dynamically maintains cache budgets during prefilling and uses constant memory during that stage.
  • Preference-Guided Cascading Cache Management: Cascading eviction produces the same final cache as applying eviction once with the computed layer budget.Theorem 1 establishes equivalence between cascading management and the vanilla preference-prioritized adaptive allocation strategy.
  • Attention-Shift Tolerant Eviction Indicator: The eviction indicator combines sustained importance and attention variability, while assigning the most recent tokens an arbitrarily large preservation score.A pooling layer clusters earlier tokens to maintain context and prevent information fragmentation; TopK then retains the highest-scoring KV pairs.

5 EXPERIMENTATION

CAKE is evaluated across LongBench and NeedleBench, multiple model architectures, memory budgets, and inference-efficiency settings. It generally outperforms competing eviction methods, especially under tight memory, while its adaptive allocation strategy improves existing indicators.

  • Experimental Setup: CAKE is evaluated on LongBench’s 16 datasets and NeedleBench’s retrieval and reasoning tasks across varied cache budgets and model architectures.Experiments cover five open-source LLMs from 7B to 70B parameters and total budgets from 64L to 2048L.
  • LongBench: CAKE outperforms competing methods on LongBench, especially under tight memory, while matching full-cache performance with Btotal > 512L.In the low-memory Llama2 setting, CAKE leads or ranks second on every dataset.
  • NeedleBench: 3.2% of the cache preserves CAKE’s NeedleBench 32K capabilities, including performance surpassing full cache on Single-Needle Retrieval.CAKE also performs particularly strongly on the more complex Multi-Needle Retrieval task.
  • Efficiency: 48.63% lower peak memory usage is achieved by CAKE than full cache at 128K context length.The comparison uses Mistral-7B-Instruct-v0.3 with FlashAttention-2.
  • Efficiency: CAKE maintains relatively stable decoding speed as input length increases, unlike full cache, whose latency grows substantially.Fixed-size KV caching reduces computational and I/O latency bottlenecks during decoding.
  • Compatibility and Ablations: P2A consistently improves nearly all tested tasks when combined with both H2O and SnapKV instead of vanilla uniform allocation.The compatibility study averages performance across six LongBench tasks under Btotal values of 128L and 512L.
  • Compatibility and Ablations: CAKE’s allocation strategy consistently outperforms uniform, pyramid-shaped, and random allocation in the default LongBench ablation.The strategy measures layer-specific attention patterns in both spatial and temporal dimensions.

6 CONCLUSION

CAKE optimizes KV-cache eviction by allocating layer-specific budgets from attention patterns, managing them through cascading updates, and selecting tokens using long-term and temporal importance. Experiments show strong performance and improved inference efficiency across models and memory constraints.

  • Conclusion: CAKE uses a global view of layer-specific attention patterns to dynamically allocate cache sizes and manage them through cascading cache management.The allocation reflects layer preferences rather than a uniform cache size.
  • Conclusion: CAKE’s eviction indicator considers both long-term token influence and temporal variability for more informed token selection.This complements the layer-level allocation strategy.
  • Conclusion: CAKE achieves superior performance across LongBench and NeedleBench models and memory constraints, particularly in low-memory scenarios.The conclusion also reports improved inference efficiency on long-context tasks.

B LIMITATION AND FUTURE WORK

CAKE adapts KV-cache allocation at the layer level but does not model finer-grained dynamics within layers. Future work may incorporate head-level patterns and combine CAKE with additional cache-optimization techniques.

  • Limitation: CAKE does not address finer-grained attention dynamics within layers.The authors identify head-level attention patterns as a potential direction for more nuanced allocation control.
  • Future Work: Combining CAKE with cache merging or pruning remains a future opportunity for further memory-efficiency improvements.The paper demonstrates compatibility with existing eviction and quantization methods but not these additional integrations.

C MORE IMPLEMENTATION DETAILS

CAKE separates inference into prompt prefilling and token decoding. It allocates layer budgets during prefilling using preference-guided cascading management, then fixes those budgets during decoding.

  • Prompt Prefilling: During prompt prefilling, CAKE uses preference-guided cascading management to allocate cache sizes across layers.This process is guided by the preference-prioritized adaptive allocation strategy.
  • Token Decoding: During decoding, CAKE fixes each layer’s cache budget using the allocation results established during prefilling.The implementation uses an observation window of size Sw = 32 and pooling-based layer-clustering information.
  • Implementation: The CakeCache implementation stores preference scores, eviction indicators, and per-layer budgets for attention computation.The attention-forward routine processes hidden states together with optional past key-value state.

D PROOFS

The proofs establish that CAKE’s cascading cache management preserves the target eviction result while progressively reducing allocated budgets. The section also situates CAKE alongside complementary cache quantization methods and LongBench evaluation settings.

  • Budget monotonicity: Positive preference scores ensure that the allocated budget for each layer decreases monotonically as cascading stages advance.This monotonicity supports progressively smaller cache allocations during the staged process.
  • Cascading eviction proof: CAKE’s proof reduces equivalence to showing that cascading and vanilla eviction retain identical top-B_l selection indices.The eviction operation retains KV pairs selected from each layer’s indicator values using the target budget B_l.
  • Cascading eviction proof: At the final cascading stage, both methods use the same layer preference scores and budget calculation, yielding equal retained-set sizes.The proof states that both sets contain B_l elements and that their final-stage budget calculation is identical.
  • Quantization comparison: CAKE and cache quantization are complementary: CAKE drops less important KV pairs, whereas quantization lowers stored-value precision.The paper evaluates CAKE with KIVI and KCVC under matched compression ratios on LongBench.
  • Evaluation setting: LongBench evaluation covers six long-context task categories, including question answering, summarization, few-shot learning, synthetic tasks, and code completion.The experiments span 16 subtasks with input lengths from 1,235 to 18,409 tokens.

F.1 DETAILED PERFORMANCE ANALYSIS ACROSS CACHE BUDGETS

Across LongBench cache budgets and model architectures, CAKE consistently matches or outperforms competing KV eviction methods. Its advantages remain especially visible under constrained memory and extend to larger models.

  • Llama2-7B-Chat: With 512L or more on Llama2-7B-Chat, CAKE achieves scores comparable with full cache across all subtasks.This result is reported for every evaluated subtask at the larger cache budgets.
  • Llama3.1-8B-Instruct: CAKE maintains strong performance across all subtasks on Llama3.1-8B-Instruct, including models using grouped query attention.The results are reported in the 64L–2048L LongBench budget range.
  • Mistral-7B-Instruct-v0.3: On Mistral-7B-Instruct-v0.3, CAKE shows a consistent advantage across all 16 subtasks and cache budgets, with resistance to score decline at 64L.The strongest robustness is reported under limited cache budgets.
  • Additional architectures: CAKE consistently outperforms baselines on Qwen2.5-7B-Instruct and Gemma-7B-Instruct under both low and high memory settings.On Gemma, it exceeds full-cache performance at Btotal = 1024L, scoring 34.18 vs. 34.09.
  • Larger-scale models: Across models from 13B to 70B, CAKE outperforms baselines and exceeds full-cache scores for Llama2-13B and Llama3-70B at Btotal = 1024L.The reported comparisons are 29.98 vs. 29.95 and 45.83 vs. 45.79, respectively.

G EXTENDED EXPERIMENTAL RESULTS ON NEEDLEBENCH

NeedleBench results show that CAKE retains an advantage across models, tasks, and cache sizes, especially for multi-needle retrieval. Its indicator and allocation strategy reduce degradation under constrained caches, although extremely small budgets remain challenging.

  • Evaluation protocol: NeedleBench scores use Levenshtein similarity, ten repetitions per case, and weighted averaging across tasks.These procedures are used to obtain stable, balanced overall scores.
  • Mistral-7B-Instruct-v0.3: CAKE consistently outperforms other eviction methods on Mistral-7B-Instruct-v0.3 across all evaluated NeedleBench tasks and cache sizes.It exhibits the smallest performance decrease relative to the full-cache baseline.
  • Mistral-7B-Instruct-v0.3: CAKE shows particularly notable improvements on Multi-Needle Retrieval, where its strategy incorporates sustained token importance and attention variability.The paper links this behavior to greater tolerance of attention shifts in long contexts.
  • Llama3.1-8B-Instruct: CAKE demonstrates advantages across all three NeedleBench tasks on Llama3.1-8B-Instruct, including models with grouped query attention.Figure 12 highlights its performance on multi-needle tasks.
  • Limitations: Under extremely small cache budgets, Multi-Needle Retrieval performance declines sharply because the cache cannot support the information demands of multiple needles.CAKE mitigates this degradation relative to alternative approaches but does not eliminate it.

I.1 ABLATION STUDY ON METHOD COMPONENTS

The ablation studies examine how adaptive allocation and attention-shift-tolerant eviction components affect CAKE’s performance. Additional analysis evaluates temperature settings for allocation on two instruction-tuned models against uniform SnapKV allocation.

  • Component ablations: The ablation isolates two design aspects: preference-prioritized allocation using H and V, and eviction using Mean and attention variance Var.These components represent spatial dispersion, temporal shifts, sustained importance, and changing attention.
  • Component ablations: 28.14 to 28.82: including both attention dispersion H and attention shift V improves adaptive allocation performance.This compares the allocation strategy with both spatial and temporal factors included.
  • Component ablations: 29.29 average performance: adding attention variance Var to the eviction indicator further improves the ablated method.Var captures dynamic attention changes alongside Mean, which identifies sustained important tokens.
  • Temperature analysis: The temperature analysis evaluates τ1 and τ2 for Llama3.1-8B-Instruct and Mistral-7B-Instruct-v0.3 with Btotal = 128L, using SnapKV as the baseline.The parameters modulate the influence of attention dispersion and shift in adaptive allocation.

J ADDITIONAL DETAILED ANALYSIS AND VISUALIZATION OF ATTENTION DYNAMICS

Additional visualizations show that attention dispersion and temporal shifts vary across layers, models, and contexts. These differences motivate cache allocation that adapts to layer preferences rather than relying on uniform or fixed patterns.

  • Attention differences across layers: Attention patterns differ substantially across layers, including layers with high or low attention dispersion and varying attention shifts.The visualizations aggregate attention weights across heads to compare layers.
  • Attention differences across layers: Appropriate cache resources should reflect layer-specific attention mechanisms because dispersion and shift vary across layers.The analysis explicitly connects these observed differences to cache-resource allocation.
  • Attention differences across models: Mistral often has higher attention dispersion near the beginning and end of its layer stack, whereas Llama3 often concentrates it in early layers.These model-level differences support model-specific caching considerations.
  • Attention differences across contexts: Attention patterns vary across task types and contexts, producing different layer preferences even within the same task type.The comparison includes differing attention-shift patterns across corresponding layers.
  • Adaptive allocation implications: Uniform allocation underuses memory, while fixed-pattern allocation generalizes poorly across models and task scenarios.CAKE instead formulates cache sizes from a global view of layer preferences.
Loading 2503.12491v2…