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
Venue: ICLR 2025 Poster (2025)
Keywords: Large Language Model
Efficient Generative Inference
Key-Value Cache
Primary area: generative models
TL;DR: We introduce a dynamic KV cache eviction method that adapts to layer-specific attention patterns, significantly improving LLM inference efficiency while preserving performance on long-context tasks.

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.

Published as a conference paper at ICLR 2025
CAKE: CASCADING AND ADAPTIVE KV CACHE EVICTION WITH LAYER PREFERENCES
Ziran Qin Yuchen Cao², Mingbbo Lin Wen Sui2,hiiuan Fan², Ke Cheng', Weiyao Lin't, Jianguo Lift Shanghai Jiao Tong University, AAn Group, "Independent Researcher
ABSTRACT
Large language models (LLMs) excel at processing long sequences, boosting de- mand for key-value (KV) caching. While recent efforts evict KV cache have al- leviated 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 Seake-slicing problem." CAKE assesses layer-specific preferences by considering attention dynamics in both spatial and temporal dimensions, allo- cates rational cache size for layers accordingly, and manages memory constraints in cascading manner. This approach enables a global view of cache alloca- tion, adaptively distributing resources across diverse attention mechanisms while maintaining memory budgets. CAKE also employs new eviction indicator that considers the shifting importance of tokens over time, addressing limitations 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 set- tings. Additionally, CAKE achieves 10x speedup in decoding latency com- pared to full cache when processing contexts of 128K tokens with FlashAttention- 2. Our code is available
com/ untgroup/cakeky cakekv
INTRODUCTION
Large language models (LLMs) Zhao et al. 2023 Achiam et al. 2023 Dubey et al. 2024 An-
thropic 2024 Al 2024) have enhanced their long-text processing capabilities, improving perfor- mance in multi-turn dialogues (Chiang et al. 2023). document summarization (Zhang et al. 2024a), question answering (Kamalloo et al. 2023), and information retrieval (Liu 22.4420 New mod- els such as GPT-4 (Achiam et al. (023). Claude 3.5 (Anthropic) [024). LLaMA 3.1 (Dubey et al.)
2024] and Mistral Large 2 (AI, 2024) have extended token processing capacities beyond 128K. Expanded contexts necessitate linear increase in key-value (KV) cache size, resulting in heavier inference-time memory burdens. Shazeer (2019); Ainslie et al. (2023) partially address this issue by merging key-value heads during the training phase. However, optimizing key-value cache with- out additional training is crucial for efficient inference of long contexts under memory constraints, particularly in typical deployment scenarios where the model structure is fixed.
One way to maintain a manageable KV cache size on the fly is to remove some KV pairs (Xiao al. 2023 Zhang et al.) 20224b] et al. 2024b). The idea eliminate less important KV pairs based on certain rules. Although recent methods have enhanced pair selection for removal, they typically assign uniform cache sizes across layers, disregarding layer-specific requirements. This approach can impair performance under memory constraints. Recent research is trying to improve this by looking how attention works in different layers (Yang et al. 2024a] Cai et al., (2024). These methods, which depend on prior observations, might not work well for all input types and models. Another idea to adjust the cache size based on the current attention pattern (Wan et al. 2024). This can help, but it may not always optimize overall performance.
"This work was done when Ziran Qin was a research intern at Ant Group. "Correppndiing Authors.

---

Published as a conference paper at ICLR 2025
+Tempurral +
Temmoora
+Temporni
Spatial
Spatial
Spatial
Spotio
High Dispersion
(b) Low Dispersion
(c) High Shift
(d) Low Shift
Figure 1:
Variation in spatial (a, b) and temporal (c, d) characteristics of attention pattern.. We
provide toy examples (left) and real examples from Mistral's different layers (right) for illustration. For more detailed analysis and visualization of attention dynamics, please refer to Appendix J
Allocating optimal cache sizes for different layers with a fixed memory budget akin to "cake- slicing." The crux of this challenge is assessing each layer's affinity for KV cache, which is mir- rored in its attention mechanisms. Our research delves into these mechanisms, examining spatial and temporal aspects across layers. Spatially, attention distribution across tokens can differ markedly be- tween layers. Some layers may spread attention broadly, while others concentrate on specific tokens, as depicted in Figure aa) and Figure Temporally, attention hotspots in certain layers shift over time in Figure cc), whereas in others. they remain constant in Figure The variability of atten- tion pattern necessitates dynamic approach to memory allocation that goes beyond uniform (Zhang et al. 2024b Liu et al.] 2024d, Ren & Zhu 2024 Li et al. 2024b), depicted in Figureßaa) or fixed- pattern methods (Cai et al. (0244). drawn in Figure They lack the flexibility to adapt to the nuanced attention dynamics. To enhance cache efficiency, comprehensive strategy is required, factoring in each layer's attention pattern and KV cache preference. Moreover, conventional meth- ods that evict KV pairs relying solely static spatial attention scores often overlook the temporal evolution of attention, missing critical insights into how token relevance changes in various contexts.
In this paper, introduce Cascading and Adaptive KV cache Eviction (CAKE), enhancing KV adaptive cache allocation strategy cou-
cache eviction by leveraging proposed preference-prioritien pled with innovative token eviction strategy. It is progressive and dynamic, employing novel metric that assesses each layer's cache size preference, taking into account both the dispersion of spatial attention and the shifts in temporal attention, as outlined in Figure To optimize cache allocation from a holistic viewpoint while keeping peak memory usage within desired limits, CAKE further employs cascading memory management method. During the prefilling phase, CAKE dy- namically manages the memory budgets for layers and adjusts the KV cache with the guidance of obtained preference scores, eliminating the need to store all KV cache for all layers. For token eviction, acknowledging the limitations of existing methods, we introduce an eviction indicator that considers sustained importance and attention variability, thereby minimizing the adverse effects of eviction subsequent decoding steps. Extensive experiments have been conducted using vari- ous LLM architectures on the LongBench and NeedleBench benchmarks, encompassing nine major task categories. Results demonstrate CAKE's superior performance across diverse memory scenar- ios within these benchmarks. The allocation strategy in CAKE, which can complement and enhance existing KV cache eviction methods, offers versatile solution for improved cache management. Furthermore, compared to full cache FlashAttention-2 implementation, CAKE significantly reduces memory consumption while simultancously enhancing LLM throughput.
Our contributions are: (1) An analysis of attention dynamics revealing spatial dispersion and tem- poral shifts, leading to a layer-specific metric for cache size requirements. (2) An adaptive cache allocation strategy that optimizes overall cache allocation based on layer preferences. (3) A cascad- ing cache management method that dynamically adjusts the KV cache during the prefilling stage, achieving memory usage efficiency comparable to uniform strategies without sacrificing eviction performance. (4) A new eviction indicator that considers the sustained importance and variability of tokens, enhancing token eviction performance.
BACKGROUND AND RELATED WORKS
Basics of KV Cache Operations. We revisit the fundamentals of KV caching. We focus single attention head, characterized as weight matrices We. WK, Wv E DDD where D denotes the model's hidden dimension. Considering a prompt embedding X E SXD with S representing the sequence length, attention module has two phases: prompt prefilling and token decoding.
2

---

Published as a conference paper at ICLR 2025
liihh
Low
Hid
Eveely
Retalned Cache
Slicin
Fine Pxtrred KV Cache
Any
Memury
(n) Uniferm Allocation
(b) Pived-Pattern Allecation
Loyer Preferences
Wen enuly
Gliing
Contest
Attention Weiph Lavers (Lm4)
Menory Budgex
KV Coche
b Global Layer
Alacstion
Memory Bulge
KV Codhe
B
Attentian Weights Lagern (L-4)
rrefercace-Piortttlaed Adaptive Allocation
Figure 2: Illustration of CAKE compared with existing cache allocation strategies. (a) Uniform cache allocation (Xiao et al. 2023 Zhang et al., 2024b, Li et al. 2024b); (b) Fixed-shape cache
allocation (Cai et al. 2024 Yang et al 2024a); (c) Preference-priorrtizeed adaptive cache allocation used in CAKE. Compared to (a) and (b), CAKE adjusts allocation ratios across different layers with layer preferences, adapting to various contexts and models, with given memory budgets.
Prompt Prefilling: The query, key, and value states are initially calculated as follows:
Q = XWQ,
= XWXW,, V XWy.
(1)
Then, the output of the attention module is determined as:
Attn (Q V) = AV, with attention weights A Softmax QK" in cache, establishing the KV cache.
(2)
E RSSS The key-value states K and V are then stored
vD
Token Decoding: The KV cache is used and updated to produce new tokens. For each decoding step i, let the new token embedding be x; RIID To circumvent repeated key-value projections, the KV cache is concatenated with the newly generated KV pair x; W WK. x; Wv R glxD for the current attention computation and to refresh the KV cache: V CCocat( V, x; WW..
K = Concat(K, x, Wk),
(3)
Though the KV cache alleviates the considerable computational demands of attention mechanisms, its linear growth with sequence length poses a challenge for extremely long input or output se- quences. Efficiently managing the KV cache within a finite cache budget remains formidable issue, especially in contexts that demand longer context lengths.
KV Cache Eviction. KV cache eviction during model inference enhances computational efficiency without altering the attention mechanism. This optimization relies on strategic cache reduction tech- niques. Early approaches to KV cache eviction focus on specific parts of the input sequence. For instance, StreamingLLM (Xiao et al. 2023) and LM-Infinite (Han et al. 2023) prioritize the reten- tion of the first and last tokens. However, this strategy risks ignoring potentially important tokens in the middle of the sequence. Recognizing this limitation, subsequent research introduces more sophisticated indicators to filter unnecessary KV cache entries, such as using cumulative attention scores (Zhang al. 2024b Liu al.] 2024d). last token attention scores (Oren al. (0024). mean attention scores (Ren& & Zhu 2024), or clustering recent attention scores (Li et al.) 2024b). While these methods offer more nuanced approaches cache eviction, they often apply the uniform allo- cation strategy, which may not maximize cache utilization. Recent FastGen (Ge et al. 2023) chooses which tokens to keep based on special markers, but doesn't 't limit the total cache size. PyramidIn- fer (Yang et al. 2024a) and PyramidKV (Cai et al.] 2024) allocate the budget in a pyramid-shaped manner, while D20 (Wan et al.) 2024) adjusts cache size based on the current layer's attention density. While these methods have shown promise, they often rely on predefined cache allocation strategies across all layers or are just based on the local layer attention, which does not fully capture the complex dynamies of attention mechanisms and lacks global perspective on layer preferences for the cache. In contrast, our work addresses these limitations by considering layer-specific cache preferences with comprehensive and global perspective.
3

---

Published as a conference paper at ICLR 2025
Szatas attention diperdarr Layendangle teatmap Mi-hhett
Spatill attertion da ACTORS RyMr
Ulmaa 25-22-
Laran 2n soouat Harr inrmarn
Cample
Samp
Sampls
Temporal attendinn Laper Sumple Teacmpp
Tergurt ottention shift acress Ryyers
Upmot 24 chor
Loras3
wonnet
Metrl ENOTNTTTE
Lamad-
Lamab 140 manvet
(a)
(b)
Figure 3: Analysis of attention dynamics. (a) Heatmap for spatial attention dispersion (upper) and temporal attention shift (lower), red color with high value, while blue for low value. The x-axis represents samples, and the y-axis represents layers. (b) Variation of spatial attention dispersion (upper) and temporal attention shift (lower) across layers and models. The experimental data is derived from the LongBench dataset (Bai et al. 2023).
3 INSIGHTS INTO ATTENTION DYNAMICS
Recent studies (iio et al. | 2023 Zhang et al. 2024b) have shown that only a few KV cache elements crucial during token decoding. Building on this, analysis identifies two pivotal attention traits: spatial attention dispersion, showing how one token's 's attention spreads across others, and temporal attention shift, tracking the evolution of highly-attendee tokens time.
Recall some basics about the attention weights A. Each A shows how much the i-th query token attends to the j-th key token. We analyze the matrix spatially and temporally to explore attention dynamics in decoder-only models. (1) Spatial Analysis: Row i, A [i, shows how the i-th token's 's attention is distributed across tokens in one step. Examining A [i. reveals the attention landscape at that step. This helps reveal how the model prioritizes information at each generation step. (2) Temporal Analysis: Column | f]. shows how attention the j-th token changes over steps. Each row in [[: represents time step, showing the j-th token's 's evolving attention. This dimension is key to tracking the model's 's shifting focus during sequence generation.
In this paper, we quantify the spatial attention dispersion by the entropy of A :]. defined as:
H(A) = Σ A(i, log(AAii:
1=0
where the inner product is conducted between each row of A and its element-wise logarithm. Rows show higher values with even attention distribution and lower values with focused attention. Aggre- gating these values measures the evenness of attention distribution per token.
Accordingly, the temporal attention shift characterized by the variance of A[:.), expressed as:
S-1
V(A)= = Σ Var A (:..))
(5)
1=0
which computes the variance for each column of A. Higher values indicate significant attention shifts across positions, while lower values suggest stable attention. Summing the variances column- wise captures the attention mechanism's temporal dynamics.
We analyze attention dynamics across multiple LLM layers using the LongBench dataset's long- context samples (Bai et al. 2023). [ shows attention dispersion changes across layers, and Figure (lower) shows attention focus shifts. The results show significant variations in attention dispersion and shift across layers, models, and contexts. Despite recent KV cache advances (Xiao ctt al. 2023 Zhang et al.) 2024b), our findings highlight LLM attention complexity, indicating a need for tailored KV cache management strategies effectively handle the dynamic nature of attention.
4

---

Published as a conference paper at ICLR 2025
4 METHODOLOGY
4.1 PREFERENCE-PRRORITTZED ADAPTIVE ALLOCATION
Given the variability in attention mechanisms across layers, models, and contexts, uniform or fixed- pattern cache allocation strategies prove inefficient, especially under limited memory budgets. To address these challenges, considers each layer's unique characteristics and adaptively allocates cache sizes from a global view of attention patterns. We define preference metric for each layer's KV cache requirements, con- sidering both the spatial dispersion and temporal shift of attention:
introduce the ceeference-prooriizee adaptive allocation strategy. It
P= P= , =
H(Al-S 1, -.w)). V= V ;, SSwl), -.:-))),,
(6)
where the preference score for an attention layer's cache size. As dispersed attention (high 7) requires retaining broader contexts, while frequent shifts (high V) demand support for complex tem- poral patterns, layers with high preference score P benefit more from larger KV cache to maintain performance. Accounting for the varying importance of attention dispersion and shift across models and memory constraints, we introduce temperature parameters T1 and 72 to adjust their influence on P. We focus on the submatrix A [- :, - S., of A, representing recent window of size Sw- This focus on the recent window chunk of prefilling attention weights is inspired by recent re- search (Li et al. 2024b Yang et al.) (0224)). These studies have demonstrated that decoding patterns can be effectively captured by analyzing the attention patterns of recent queries.
We implement the infereecce-proritiize adaptive allocation strategy using layer-specific preference scores. These scores adjust dynamically to varying models and contexts, ensuring adaptability, Each layer's cache size set by normalizing the preference scores and allocating the total memory budget accordingly. This results in an adaptive cache size B {B:) calculated follows:
Pt
B1
Bioaal:
2k=O
Pk
where Pi is the preference score for layer 1, and Bunal is the total cache budget. This method optimizes cache size allocation for each layer's unique characteristics and the input context.
4.2 PREFERENCE--GUIDED CASCADING CACHE MANAGEMENT
Our meference-ppioritizee adaptive allocation strategy, though optimal, initially requires compre- hensive view of prefilling attention weights across all layers. This leads to the peak memory usage of O(S L), which may exceed the cache budget B
To overcome this, we further develop the preference-guidee cascading cache management, as illus- trated in Figure dynamically maintains the cache budget during prefilling, effectively reducing peak memory usage to the target. titions the prefilling process into stages, one for each model layer, and cascadingly manages cache memory guided by preference scores. At each stage, as a new layer completes its prefilling compu- tation, the budget redistributed based on the current obtained preference scores (Algorithm lines 3-7), and KV caches are updated among all processed layers according to the redistributed budget.
outlinees our cache management procedure, which par-
This approach ensures constant memory usage during the prefilling stage while maintaining equiv- alent cache distribution and eviction results as the standard strategy. Our algorithm is supported by the following theoretical results (Proposition and Theorem [ proofs provided in Appendix D).
Proposition 1. For any layer e [L], the allocated budget size decreases monotonically from stage to --
B (m++) < <B(m)
E [1, 1].
(
where m)) is the redistributed cache budget for layer at stage m, calculated based on the current obtained preference score P:
P
B (m)
B
(9)
Σ
122
Pk
k=U
5

---

Published as a conference paper at ICLR 2025
Algorithm Preference-guiie
Cascading Cache
lager
Prefilling Stage 1
Management
Freentt Laper Poz2
18 PP 1
Meonaring
Redistrbutee Budget
(0.74.220)
Input:
weights Aml-Sw 1, :::] E R ma-
Total cache budget Bunal local attention
Pull muchir
trices Km, Vn. E RSxD number of layers L
Bvicy
Layen
Output: Retained
cache
set
C
Previous KV Cache (Suge 0)
Updated KV Cache (Stage 1)
[[K [(K((-1)
(1(11))
E [T]]
1: Initialize (()) , (
B p(u) <-
B
Curmert Lager
Prefilling Stage
2: for stage m 0 to
do
3: C(m) f C C ((Km, Vm
Bia) (0)%) S. _p 1
Trocund (22)
Mennaring
Redistmbute Budget (Pail cache)
(0.74.220) - -...0.....
4: Compute preference score Pm for layer m
using Eq. 6
Lagar
Layer
5: Compute indicator In (L6))) using local atten-
tion weights from chosen eviction method
6: pim) <- p(m-1) U (Pm).
Previous
Updated (Stage
((m) (
(Stage
7: B gim) [B] (m) : [m]], where B(m)) is
Prefilling Stoge 3
computed using Eq.
g (16 2 a p
8:
// Manage current KV Caches
Mazsuring
Redatmbute Bidget
45400444444 2000000000000
9:
for layer t m
Secorred Lapen (02)
YcPall mocker
10:
(K,, V.) <- C
Laper
Sayers
B <- ( [T] I ((m))
Lagar
ExctX
2
11:
#/ Update cache for layer
12:
(K, V.) f((m)
60
EVICT((Ki Vi)(m),
B(()) I.(m)
Previnus
lppdated
13: C(
(Ki, Vi) Vz) (m+1)
(Stage
(Stage
(R, Ve)(m), S s(m ( (00)
Lipec
14: end for
Figure 4: Illustration of preference-guided
15: end for
cascading cache management algorithm dur- 16: return Retained cache set ing three stages.
{( (1.1)) (0- ::mm E [Z]]
Proposition 1 guarantees that the cache budget allocated to each layer decreases monotonically as stages process, ensuring that the KV cache (K,, V, (mm+))) is always proper subset of the previous stage's cache (K,, V.)(m), regardless of the eviction method used.
Before presenting Theorem let's define some operational details related to our algorithm. Given indicator vector 1 E SR it measures token importance for layer l. The specific calculation of our eviction indicator will be detailed the subsequent section. For each stage m E [L], the eviction operation, denoted EVICT(-), retains KV pairs ((("), (((m) 1 == from the current cache
(m) ()) 1 xD corresponding to the top-B( (m) positions with the highest scores in T(m))
E TRB m(- The indicator vector T(m)) is also updated to i(m) B( retaining only the elements corresponding the preserved positions. The updated cache and indicator (Ki, Vi) (m))
i((m) will serve as the basis (Ki, V, (m++)) and fm+1)) for further updates in the next stage m + 1.
Theorem 1. For any layer E [L]. the KV cache (K,, (.)( obtained through cascading evic- tion from stage i stage L is equivalent the result of applying the eviction operation once the full KV cache (K1, V1} E Sxb with the cache budget B₁ computed in Eq. 7
[K, (L-11)) EVICT( V.)(m) (m)) (mm)
(m) L-1 (,, VI), B I.)
(10)
Theorem demonstrates that our sefterence-guideed cascading cache management achieves equiv- alent KV cache eviction results to the vanilla weefrence--piorittin adaptive allocation strategy, despite its cascading operation. This method is theoretically compatible with existing KV eviction techniques, claim that will be empirically validated in Section 5.5

---

Published as a conference paper at ICLR 2025
4.3 ATTENTION-SHIFF TOLERANT EVICTION INDICATOR
Current top eviction indicators, such as accumulated or mean attention scores (Zhang et al. 2024b Liu et al., 2024d Ren & Zhu 2024), simplify tokens' attention profiles into single values, effec- tively identifying important tokens. However, this approach may overlook tokens whose importance fluctuates, as it fails capture the dynamics of shifting attention. This could result in prematurely evicting tokens from the cache, affecting the model's ' future ability to retrieve relevant information.
We propose robust eviction strategy tolerant of attention shifts. We use a multi-faceted indicator considering sustained importance and attention variability. For layer I, the eviction indicator I E R$ computed, where each element [|[n] signifies the importance of token n. computed as:
Mcan [Maan(Ail- (Arl-S n, + Var(Ai i,n],, if S S Sw otherwise,
[m] 12
(11)
where Mean(.) and Var(.) measure sustained importance and attention variability, respectively, with adjusting their influence. Inspired by recent research (Li et al. [222bb] Yang et al. (024a). we assign an arbitrarily large value to ensure the preservation of the most recent W tokens. Also, a Sw] pooling layer clusters l:[: S S.] to maintain context and prevent information fragmentation.
The eviction operation EVICT((K, V.). Bt, It) is then executed, retaining the KV pairs (K₃, Vi) that correspond to the top- B, scores in Iy:
K = D/.,,
V = :..), D, = TIp B.),
(12)
where TopK selects the indices D2 E TR "h of the top Bi scores in I,
5 EXPERIMENTATION
5.1 EXPERIMENTAL. SETTINGS
Backbone LLMs. Our experiments involve five major open-source LLMs (7B-70B parameters) capable of handling 4k-128k token contexts. These models feature two representative attention structures: (1) Multi-head attention: including Llama2-Chat (Touvron et al., 2023) and Gemma- Instruct (Team et al.] 2024). (2) Grouped-query attention: including -lama3--nnstuee (Dubey et al.)
2024), Mistral-v0.3 (Jiang et al. 2023), and Qwen2.5-Instruet (Team 2024).
Baseline Methods. We assess CAKE's performance against five baselines: (1) Uniform Allocation: StreamingL.LMM (Xiao et al., (023) balances initial and recent tokens, H20 (Zhang et al., (024b) cumulative attention, TOVA (Oren al. 2024) adopts last-token attention, and SnapKV (Li ctt al. 22024b) clusters recent attention. (2) Non- Uniform Allocation: PyramidKV (Cai et al. 2024) employs fixed pyramid-shaped allocation strategy with SnapK eviction indicator.
Evaluating Tasks. To evaluate CAKE's performance across various memory budgets, we use two carefully designed benchmarks: (1) LongBench (Bai et al.) (023): Focuses long-context under- standing, encompassing 16 datasets in six categories: Mongle/Mult-DDcummee QA, Summarization, Few-shot Learning, Synthetic Tasks, and Code Completion. (2) NeedleBench (Li et al. 2024a): Tests retrieval and reasoning in complex contexts through three subtasks: Single-Needle Retrieval, Multi-Needle Retrieval, and Multi-Needle Reasoning.
Implementation Details. CAKE is compared to baseline methods across different total memory budgets ranging from 64L to 2048L. For uniform allocation, each layer has fixed KV cache size. Non-uniform methods like PyramidKV and CAKE vary cache sizes per layer while keeping total memory constant. Tokens are evicted during both prefilling and decoding maintain memory usage. Experiments are run on NVIDIA 1100 80GB GPUs. For specifics, see the Appendix
5.2 EVALUATIONS ON LONGBENCH DATASET
We benchmark CAKE against other methods on 16 datasets. Figure displays the performance of various methods under different memory constraints across three models. SnapKV surpasses legacy methods like StreamingLLM and H2O by using an observation window for memory reten- tion. TOVA, relying solely the last token, falls short by ignoring the broader context. Despite
7

---

Published as a conference paper at ICLR 2025
Uama2-7B-Chat
Mistral78-irsstuccto0
Caree
SocemenylLM
Snapty
DonamingguM
Seaphy Prumaay
Snapky
H20
Pyreresty
264
Figure Average score among 16 datasets of LongBench under different cache budgets.
Table I: Performance comparison over 16 datasets of LongBench. The best result is highlighted in bold, the second best in underline.
Doeemnnnn QA
Muli-Documann QA
Summerizaion
Fow-slot Lewwning
Synthetic
Clode
Method
NAvDA
Queper
MF-on
HosporZA
ZWILMOA
Musqqe
GevRRport
@MSam
MultNews
TREC
TrivinAA
SAMSim
PCaun
PR-cn
RB-P
Ang.
Liama2 7B
Chat, Da.
125Z
StreaminLLLM
13.57
13.33
22.39
20.16
17.12
6.38
16.72
19.58
15.37
51.50
74,84
33.00
1.57
4.00
45.21
40:45
23.35
120
14.56
14.59
26.52
20.05
28.75
5.00
16.20
20.30
20.34
38,00
72.75
37.10
2.33
3.30
49.68
47.37
25.06
TOVA
13.56
13.46
20.28
22.22
25.0T
4.79
15.72
19:11
15.55
42.50
79.58
37.06
3.45
3.61
46.34
39.70
25.31
SnppKV
13.70
16.27
32.52
22.76
28.59
7.56
18.87
19.96
20.15
43.50
79.90
36.74
2.79
7.00
51.86
47.41
28.09
PynmidKV
13.45
16.61
31.21
23.74
28.43
5.89
18.89
20.56
19.81
44.50
50.32
37,82
2.20
8.00
51.19
47.56
28.20
CARE
15.15
16.38
32.17
25.25
31,09
7.00
19.47
21,07
28.36
48.50
$0.66
37.77
4.42
9.00
$1.48.
48.90
29.29
Lluma2: BB-Cltt 2yu
10244
SMresmingL.LM
13.23
16.75
26
25.97
25.67
5.87
22:21
19.27
24.07
61.00
$0.88
40.90
2.25
1.00
56.71
48.79
29.86
H2O
16.58
17.67
32.04
26.01
26.75
7.81
22.87
20.99
25.15
62.00
53.02
40.06
2.50
3.50
57.57
52.02
31,05
TOYA
15,00
17.32
29.57
27.48
31.71
7.05
21.36
20.54
24.75
63.50
$3.35
41.78
2.63
8.00
56.09
52.11
51.46
SuppKV
18.05
19.73
38.07
27.57
30.86
8.41
23.41
30.79
25:22
63.50
83.13
40.08
3.53
7.50
58.20
53.23
82.50
PymmidKV
17.94
30.68
38.85
27.25
29.99
7.51
25.17
21.55
25.45
63,50
82.97
40.79
3.57
7.50
57.72
51.93
32.55
CAKE
18.35
30.95
38.49,
27.99
31.15
8.59
24.29
21.06
25.40
63.50
33.33
41 30
3.11
9.50
57.82
52.93
32.98
1. 1-88--B--BBB0 -iB--lssttre, Busi
125C
StreamingLLLM
23.01
22.50
31.00
46.79
40.52
24.75
18.28
20.89
16.96
40.50
$6.00
38.55
7.00
99.50
51.06
47.16
38.79
H2O
26.18
34.33
44.08
51.33
43.29
27.78
211.79
22.36
20.69
41.50
99.85
40.94
6.20
99.00
58.12
49.82
41.67
TOYA
29.44
27.10
45.20)
54.15
44.97
28.87
22.21
20.54
20.52
53.50
92.27
40.99
6.17
99.50
51.14
43.15
42.47
Snapky
26.40
27.67
47:93
53,40
24.20
29 65
211.56
23.11
20.13
45.50
59.36
39.59
6.33
99:00
57.69
50.69
42.13
PymmidKV
25.57
27:40
46 16
$3,14
44:51
77.16
20.74
20.26
19.90
45.50
87.00
40.13
6.50
99.50
56.98
48: 26
41.94
CAKE
27.41
32,01
49.58
$2.99
45.76
28.32
21.85
23.15
21.56
47.50
89.61
40.00
6.33
99.50
57.02
40.83
43.39
Limma.. Instt-nstract, Fi..
10247
StreamingLLM
26.64
30.77
35.59
47.31
42.03
24.17
25.81
21.31
25.66
63.50
58,
42.76
6.59
88.00
61.36
53.47
42.75
H2O
29.57
36.15
45.94
54.43
44.81
29.04
27.64
23.31
36.47
62.00
91.83
45.14
6.36
99.00
62.74
55.39
46.11
TOVA
50.65
40.95
51.09
54.58
46.51
50.62
28.12
23.61
26.24
68.00
91.49
43,80
5.92
99.50
60.73
52.64
47.15
SuppkV
30.95
44.74
52.58
55.09
26.83
30.37
27.87
24.57
25.99
08.081
22.113
42.60
6.50
19.50
63.00
56. $6.50
47.45
PynmidKV
50.54
43.64
52.73
55.20
45.29
31.28
27.53
24.50
26.00
58.00
92.09
41.75
6.05
99.50
62.35
55
47.69
CAKE
50.68
44.95
52.
55,49
46.99
30.87
28.48
34.91
26.
69,00
91.3
42.60
6.00
99.50
62.65
56.89
48.13
Mistral 7B-Lasttuu---..3 Mu
1282.
StremmingLLMM
16.91
21.51
24:85
54.14
26.99
16.64
15.67
18.61
14.40
43,50
53.25
37.00
0.00
23.00
41.63
42.12
28.76
H2O
21.25
36.66
35.13
38,82
29.00
18.85
21.00
19.50
18.63
41,08
57.64
38.22
1.50
67.00)
44.35
47.29
34.79
TOYA
22.47
34.26
37.22
42.26
25.85
19.97
19:40
18.70
17.86
63,00
S9.MS
37.31
0.50
50.50
31.42
37.22
34.52
SupkVV
21.02
27.36
41:35
45.15
29.25
22.75
30.47
20.17
17.75
43.50
87.3
36.01
0.50
00.90
44 48
45
35.81
PynmiuKV
21.73
36,60
41 46
43.20
29.32
21.47
20.23
19.82
17:46
40.00
67.64
37.11
1.055
02.00
42.55
43.98
35.14
CAKE
22.31
29.15
43:51
$1.51
30.36
22,85
21.56
20.47
18.96
$7.00
60.00
39.36
1,00
76.50
44.96
16.
37,33
Mittral-8--nsstait--00
Feaa
10247.
StreamingLLM
20.96
28.15
3003
37.06
27.56
16.03
24.03
19:07
22.79
67.00
57.61
40.96
1.59
21.50
48.05
48.87
33.82
HID
23.76
31.63
41.31
43.24
31.07
30.43
26:74
30.41
23.93
67.50
59.84
42.62
1.50
72.00
50.40
49.87
99.72
TOYA
26.97
34.51
45.58
44.32
32.58
22.83
26.91
30.75
23.49
75.00
53
43.17
1.00
91.08)
47.51
43.06
41.96
SnapKV
26.61
35.78
48.11
45.75
32.20
23.37
26.71
21.84
23.18
70.50
$3.60
41.37
0.50
83.00
00.40
$1.79
62.18.
Pynmidky
25.51
56.02
47.72
24.74
33.16
33.91
55.55
21.83
23.27
70.50
68.4
40.94
1.00
87.00
50.17
50.91
41.96
CAKE
26.09
36.34
48.11
45.07
32.37 39
23.40
27.56
21.45
24.03
73.50
88.6
42.71
0.00
91.50
51.06
$1.25
42.69
similarities to SnapKV, PyramidKV underperforms in many scenarios, possibly due to its limited adaptability. CAKE, however, outperforms others by dynamically allocating memory to each layer based on its specific needs, especially under tight memory conditions. Notably, CAKE achieves per- formance comparable to full cache models while only utilizing small cache budget (Bioal > 512[).
For detailed analysis, see Table for results across two memory scenarios: a low setting (Busaa 1287) and high setting (Bunal 1024L). Full results under different constraints are detailed Appendix(F.] CAKE consistently ranks among the top performers or closely matches them across various tasks. In the low memory scenario with the Llama2 model, CAKE leads finishes second in all datasets. CAKE's adaptive allocation strategy and robust eviction indicator consistently demonstrate superior and balanced performance, its generalizability extends well to various model architectures with ranging model sizes. We evaluate this on two more architectures, Qwen2.5 and Gemma, well larger models (13B-70B), with results presented in Appendix F2 and F3
8

---

Published as a conference paper at ICLR 2025
Matral TB Instrct +03 NeedieSenc 33K
Jamab.1 1 03 retruct NeecheSanch 3IK
Muil
On0 Scare
col
125
Vil
cndt f22
yora
oagt 4oo
dn 2912
thi
ago soees iose sive.
Figure 6: Performance comparison on NeedleBench 32K.
5.3 EVALUATIONS ON NEEDLEBENCH DATASET
We also test CAKE on retrieval and reasoning tasks using the NeedleBench dataset, with the results depicted in Figure Consistent with the LongBench outcomes, CAKE shows superior overall per- formance. While utilizing only 3.2% of the cache size, CAKE preserves the model's capabilities for processing 32K-token-lonn contexts and even surpasses the full cache in the Single-Needle Retrieval task. This suggests that pruning specific cache information may eliminate superfluous details, allow- ing generation to focus on the most salient aspects and enhance its overall efficacy. Notably, CAKE outperforms existing methods in more complex tasks akin real-world scenarios, particularly in Multi-Needle Retrieval. This can be attributed to CAKE's design, which balances long-term sig- nificance and short-term relevance, maintaining more comprehensive and balanced representation of contextual information while avoiding premature discarding of information that could be vital for subsequent complex retrievals. For a more in-depth examination of the experimental outcomes, including detailed statistics and additional visualization graphs, please refer to Appendix G
5.4 EVALUATION OF MEMORY AND THROUGHPUT
We assess the effectiveness of our method in reducing memory consumption and enhancing time efficiency during L.M inference by analyzing peak memory usage and decoding latency Mistral- 7B-Inntrruct-0033 implemented with FlashAttention-22 (Dao. 2023). Our comparison includes full cache, and three KV cache eviction methods, all with Biotal 1024T: SnapKV (uniform allocation strategy), PyramidKV (fixed-pattern allocation strategy), and our proposed CAKE method.
Peak Memory Usage. As depicted in the left panel of Figure 7 CAKE exhibits substantial memory- saving capabilities, comparable to other KV cache eviction methods such as uniform allocation (SnapKV) and fixed- sredd-pttern non-uniform allocation (PyramidKV). All these approaches maintain fixed-size KV cache. When compared to the full cache implementation, CAKE achieves an impres- sive reduction peak memory usage of approximate 48.63% with 128K context length.
Throughput Analysis. Regarding decoding latency, the right panel of Figure shows that both uniform and non-uniform eviction methods exhibit comparable inference performance. As input length increases, decoding latency for the full cache method grows significantly due to escalating computational demands and I/O latency bottlenecks. In contrast, CAKE maintains a relatively stable decoding speed by preserving fixed amount of KV cache, resulting in significantly lower latency compared to the full cache, particularly for longer sequences. It is noteworthy that CAKE demon- strates remarkable efficiency, achieving over 10x speedup in decoding latency compared to the full cache approach when processing sequences with 128K context length.
5.5 COMPATIBILITYY WITH EXISTING KV EVICTION METHODS
Our preference-priorttin adaptive allocation strategy demonstrates strong compatibility with ex- isting eviction indicators. We evaluate its performance using two representative methods: H2O, which serves as classical indicator, and SnapKV, which represents new advanced indicator. The experiments are conducted on LongBench datasets using Llama2-7B-Chie under Buaat of 128 and 512L. We report the average performance across six tasks. In Figure when compared with vanilla uniform cache allocation, methods equipped with our allocation strategy consistently improve per- formance across nearly all tasks. Importantly, they achieve significant overall performance gains. This comprehensive improvement across different eviction methods and tasks demonstrates the ver- satility and effectiveness of our allocation approach. It stresses the potential of our strategy as a generalizable framework for optimizing KV cache eviction, enhancing various existing methods.

---

Published as a conference paper at ICLR 2025
Table 2:
Performance compari-
Snapky
Snapit/
Paramidky
Pyomiayy
son with different allocation strate- gies. "P2A" denotes our preference- prioritized adaptive allocation.
CAKE
CAKE
Peak Momoryy
Fall Cache Fiashlttt22
Fut Cadhe w. F
Full Cacre DOM (NAA)
Ful Cache GDM (UA)
Strategy
Avg.
4k
16k
32k
64k
256k
4k
8k
16k
32k
64k 126k
Uniform
28.36
Context Length
Context Length
Pyramid
28.69
Figure 7: Peak memory usage and decoding latency on A100 80GB.
Random
28.3:1:.2
P2A
29.29
Serventatioo Budget
Summa
Fea shttleriing
Teals jeg v/t dailly 12:3
Table 3: Performance comparison with different eviction indicators.
Orge Dxc
Syrthitic
Singe 2ox
Synthetc
Indicator
Avg.
20 QA
Code Cotp oc
Vur: Cec
Cada Compleze
Mcan
28.82
Var
28.68
Figure 8: Performance methods w/ or w/o preference- prioritized cache allocation strategy, abbreviated as "P2A."
Mean
Var
28.6
Mean
Var
29.29
5.6 ABLATION STUDIES
In this section, we present series of ablation studies on Longßench evaluate the effectiveness of our proposed allocation strategy and eviction indicator. We use Llama2-7B-Che with cache budget B 128L as the default setting. Additional ablation studies are provided in Appendix[]
Effectiveness of Proposed Allocation Strategy. To assess our preffreecee-rioritizee adaptive allo- cation strategy, we compare it with various cache allocation strategies, including uniform, pyramid- shaped, and random allocation. As shown in Table the pyramid-shape allocation strategy performs better than the uniform allocation strategy in this case. However, it may become ineffective in many scenarios, as we verified in the previous experimental section. Notably. the random method even outperforms the uniform strategy in some cases. This indicates that the conservative allocation strat- egy of uniform distribution fails to fully utilize the limited memory budget. Our allocation strategy consistently outperforms all baselines, as our method comprehensively measures the KV cache pref- erence by capturing complex layer-specific attention patterns both spatially and temporally.
Effectiveness of Proposed Eviction Indicator. We further investigate effective eviction indicators to identify the most crucial tokens for preserving model performance. We compare different met- rics for the eviction indicator: mean only, variance only, and combinations mean and variance (multiplication and addition). The results in & that, compared to variance only, tokens selected by mean attention are more conducive to maintaining performance. While the multiplica- tive combination of mean and variance slightly underperforms the individual metrics, the additive combination achieves the best performance overall. As we analyze, using mean attention effectively captures tokens with long-term importance, ensuring the retention of consistently relevant informa- tion. On the other hand, incorporating variance helps identify positions with the most significant changes, thus aiding in maintaining the attention distribution. The additive combination optimally balances these two aspects, leading to more informed eviction decisions.
6 CONCLUSION
In this paper, we propose Cascading and Adaptive KV cache Eviction (CAKE), a novel approach for optimizing KV cache evicting in L.MMs. CAKE dynamically allocates cache sizes by leveraging global view of layer-peecific attention patterns, using cascading cache management guided by layer preferences. Additionally, CAKE introduces new eviction indicator that accounts for both the long-term influence and temporal variability of token importance, allowing for more informed token selection. Experiments on the LongBench and NeedleBench benchmarks highlight CAKE's superior performance across different models and memory constraints, especially in low-lewee mwmmmerry scenarios. Our method enhances both L.MM performance on long-context tasks and inference efficiency.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
The paper is supported in part by the National Natural Science Foundation of China 62..19255 U21B2013) and Ant Group.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.0877 230:0087438874 2023.
Mistral AL
Large enough, 2024.
URL batts::/mistral.aa/
mist ral Large- 2407/
Joshua Ainslie, James Lee-Theee Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check-
points. arXiv preprint arXiv:2305. 13245, 2023.
Silas Alberti, Niclas Dern, Laura Thesing, and Gitta Kutyniok. Sumformer: Universal approxi- mation for efficient transformers. In Topological, Algebruic and Geometric Learning Workshops 2023, pp. 72-86. PMLR, 2023.
Anthropic. The claude 3 model family: Opus, sonnet, haiku, March 2024. URL https:// http:://w.............//.a...... Model Card Claude 3. 3.pdf
www-edn.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang. Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint 2xXx:2004.005500 2020.
Zefan Cai, Yichi Zhang, Bofci Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling. arXiv preprint 2.XX::2406002099 2024.
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Uni- fying sparse and low-rank attention. Advances in Neural Information Processing Systems, 34: 17413-17426, 2021.
Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang. Yonghao Zhuang, Joseph E. Gonzalez, et al. Vicuma: An open-source chatbot impressing GPT-4 with 90%* * ChatGPT quality, March 2023. URL httpt:tt//// htt//mmmyss./// .rgg b 0log/20033003300- cura/
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprini arXiv: 10509, 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, 2023. URL. https://arxii: org/
Harry Dong, Xinyu Yang, Zhenyu Zhang, Zhangyang Wang, Yuejie Chi, and Beidi Chen. Get more with less: Synthesizing recurrence with kv cache compression for efficient 1lm inference. arXiv preprint S10.:::202.099990 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Leuman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
11

---

Published as a conference paper at ICLR 2025
Qichen Fu, Minsik Cho, Thomas Merth, Sachin Mehta, Mohammad Rastegari, and Mahyar Na- jibi. Lazyllm: Dynamic token pruning for efficient long context Ilm inference. arXiv preprint arXiv:2407.14057, 2024.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprinl arXiv:2310.01801, 2023.
Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint vrXiv:2308.6617, 2023.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.0682) 2023.
Ehsan Kamalloo, Nouha Dziri, Charles LA Clarke, and Davood Rafiii. Evaluating open-domain question answering in the era of large language models. arXiv preprint arXiv:2305.00998. 2023.
Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong, Zaoxing Liu, Tushar Krishna, and Tuo
Zhao. Gear: An efficient kv cache compression recipefor near-lossless generative inference of Ilm. arXiv preprint arXiv:2403. 005527, 2024.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers rnns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165. PMLR, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint D.xXiv:2001.044550 2020.
Mo Li. Songyang Zhang, Yunxin Liu, and Kai Chen. Needlebench: Can Ilms do retrieval and reasoning in million context window? arXiv preprint 4 2024a.
Yuhong Li, Yingbing Huang, Bowen Yang. Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation. arXiv preprint a4Xiv:2404. 14469, 2024b.
Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al. Deepseek-v2: A strong, economical, and efficient mixture- of-experts language model. arXiv preprint arXiv:2405.04444 2024a.
Akide Liu, Jing Liu, Zizheng Pan, Yefei He, Gholamreza Haffari, and Bohan Zhuang. Mini- cache: Kv cache compression in depth dimension for large language models. arXiv preprint arXiv:2405. 14366, 2024b.
Nelson Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157 2024c.
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang. Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for Ilm kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024d.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: A tuning-free asymmetric 2bit quantization for ky cache. arXiv preprint arXiv:2402.02750, 2024e.
Piotr Nawrot, Adrian Lańcucki, Marcin Chochowski, David Tarjan, and Edoardo M Ponti. Dy- namic memory compression: Retrofitting Ilms for accelcrated inference. arXiv preprint arXiv:2403.09636, 2024.
Matanel Oren, Michael Hassid, Yossi Adi, and Roy Schwartz. Transformers are multi-state mns. arXiv preprint 4rXiv:2401.06. 004004 2024.
12

---

Published as a conference paper at ICLR 2025
Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovec, Dale Schuurmans, and Bo Dai. Combiner: Full attention transformer with sparse computation cost. Advances in Neural Information Processing Systems, 34:22470-2222 22482, 2021.
Siyu Ren and Kenny Q Zhu. On the efficacy of eviction policy for key-value constrained generative language model inference. arXiv preprini urXiv:2402.06200 2024.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguis- fics, 9:53-68, 2021.
Noam Shazeer. Fast transformer decoding: One write-head all you need. arXiv preprint arXiv:1911.02150, 2019.
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. arXiv preprint arXiv:2405.0525. 2024.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, al. Gemma: Open models based on gemini research and technology. arXiv preprint 2430v22403.08295, 2024.
Qwen Team. Qwen2.5: A party of foundation models, September 2024. URL httph:h//pe:tttttt cotiubbiiooloo///eo /wwpo////wee255.mmm 5/
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.0922 0998899 2023.
Zhongwei Wan, Xinjian Wu, Yu Zhang, Yi Xin, Chaofan Tao, Zhihong Zhu, Xin Wang. Siqi Luo, Jing Xiong, and Mi Zhang. D2o: Dynamic discriminative operations for efficient generative inference of large language models. arXiv preprint arXiv:2406. /3..224466.3335 2024.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.00 04768, 2020.
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Song Han, and Maosong Sun. Infllm: Unveiling the intrinsic capacity of ulms for under- standing extremely long sequences with training-free memory. arXiv preprint arXiv:240..06617, 2024.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453 2023.
Yuhui Xu, Zhanming Jie, Hanze Dong. Lei Wang, Xudong Lu, Aojun Zhou, Amrita Saba, Caiming Xiong, and Doyen Sahoo. Think: Thinner key cache by query-driven pruning. arXiv preprint arXiv:2407.21018. 2024.
Dongjie Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. Pyramidinfer: Pyra- mid kv cache compression for high-throughput IIm inference. arXiv preprint arXiv:2405. 12532, 2024a.
June Yong Yang, Byeongwook Kim, Jeongin Bae, Beomseok Kwon, Gunho Park, Eunho Yang. Se Jung Kwon, and Dongsoo Lee. No token left behind: Reliable ky cache compression via importance-awart mixed precision quantization. arXiv preprint arXi::2402.18096 2024b.
Yuxuan Yue, Zhihang Yuan, Haojie Duanmu, Sifan Zhou, Jianlong Wu, and Liqiang Nie. Wkvquant: Quantizing weight and key/value cache for large language models gains more. arXiv preprint arXiv:2402.12065. 2024.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283--1297 2020.
13

---

Published as a conference paper at ICLR 2025
Tianyi Zhang, Faisal Ladhak, Esin Durmus, Percy Liang, Kathleen McKeown, and Tatsunori B Hashimoto. Benchmarking large language models for news summarization. Transactions of the Association for Computational Linguistics, 12:39-57, 2024a.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient gen- erative inference of large language models. Advances Neural Information Processing Systems, 36, 2024b.
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yinggian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et survey of large language models. arXiv preprint arXiv:2303.18222 2023.
14

---

Published as a conference paper at ICLR 2025
A ADDITIONAL RELATED WORKS
Efficient Attention. Numerous strategies have been devised to tackle the high time and space de- mands of the attention layer. These strategies include low-rank techniques to transform the original attention mechanism into one with linear or near-linear complexity (Choromanski et al. [020 Wang
al. 2020 Katharopoulos et al. 2020 Alberti al. 2023), sparse atfention confine the attention
to specific. predetermined or adaptively learned patterns (Child et al., 2019 Kitaev et al.
2020 Roy
al.
2021). ). and their integration for more comprehensive solution (Beltagy et al. [ Zaheer
et al. 2020, Chen et al., 202T] Ren et al.] 202I). Another group minimizes KV cache size. This
achieved by either re-engineering the architecture of multi-head attention or by innovating the
update process of the KV cache. Multi-Query Attention (MQA) (Shazeer (219)) and Group-Query Attention (GQA) (Ainslic et al., 2023) reduce KV heads by sharing KV pairs. Multi-head Latent Attention (MLA) (Liu et al. (024a) furthers efficiency by compressing keys and values into latent vectors, while (Su et 20024) employs cross-attentics to recycle the shared KV cache. Dy- namic Memory Compression (DMC) (Nawrot et al.) 2024) dynamically consolidates the KV cache during the training process, thereby optimizing resource utilization.
Orthogonal KV Cache Compression. Beyond the eviction of the KV cache, numerous strategies have been investigated for its compression. Specifically, quantization methods (Liu et al. 2024e)
Yue et al.] 2024; Kang et al., 2024) have been applied to condense the KV cache into a lower-
bitwidth format. Moreover, a fusion of KV cache eviction with quantization (Yang et al. 2024b
(2024), low-rank factorization (Dong al. 2024), or cache merging (Liu et al. 2024b, Wan et al.
pruning Xu al.
2024) has been explored to achieve more compact form of the cache, preserving
some of the discarded data. Additionally, there are initiatives transfer the evicted KV cache to
auxiliary memory resources for potential redeployment in computation (Xiao et al.) 2024 Fu et al.) 2024). It crucial to acknowledge that these compression techniques are compatible with KV cache eviction and, when integrated, can amplify their individual benefits, leading to more robust solutions for easing the KV cache load.
B LIMITATION AND FUTURE WORK
CAKE introduces a novel perspective on KV cache management by considering the model's vary- ing attention patterns across layers and comprehensively measuring layer-specific preferences for KV cache allocation. This approach enables effective adaptation to different models and input con- texts, optimizing the use of limited memory resources. Despite these strengths, several limitations and areas for future work remain. While CAKE improves cache management at the layer level, does not address finer- meee-graieed dynamics within layers. Future work could explore integrating head- level attention patterns with layer-specific variations, potentially offering more nuanced control over cache allocation. Additionally, although we demonstrate CAKE's compatibility with existing evic- tion methods (Section 5.5) and quantization methods (Appendix (.5). there is potential to further enhance memory efficiency by combining it with other KV cache optimization techniques such as cache merging (Liu et al. 2024b) and pruning (Xu et al. 2024). This integration could significantly expand CAKE's applicability and resource-savine capabilities. In the future, we will continue to explore strategies for efficient LLMs, focusing on both memory efficiency and computational opti- mization to further push the boundaries of resource-constrained inference.
C MORE IMPLEMENTATION DETAILS
In this section, we provide more detailed implementation information for CAKE. Our algorithm is divided into two phases: prompt prefilling and token decoding. During prompt prefilling, we utilize preference-guides cascading cache management (Section 4.2 based on the preference-priorrtizeee adaptive allocation strategy (Section[4.1 to allocate appropriate cache sizes for different layers. During preference-guider cascading cache management, we manage the KV cache by updating using eviction operation based on the attention-shifi tolerant eviction indicator (Section For the preference -rrioriiiee adaptive allocation strategy, the temperature parameters Ti and T2 are set to 0.2 ~ 2 and 0.4 ~ respectively, optimized through grid search. For our experimental models, aggregate each quantized value to represent its layer characteristics. For the attention-shhe tolerant evietion indicator, we fix 7 200. Following SnapKV (Li et al. 2024b), we an observation
15

---

Published as a conference paper at ICLR 2025
window of size S. 32 and utilize pooling layer clustering information. In the decoding stage, once the attention mechanisms for layers have been well established, we fix the budget size for each layer according to the allocation results obtained from the ceeference-prioriiin adaptive allocation. Listing provides Pyforch-style pseudo-code for CAKE's implementation with Fltshint
Listing I: Implementation of CAKE in pseudo Py PyTorh style.
class CakeCache(Cache):
def .-initt..(seff) self peff.pref..cee scores =
indet.indiicaorss [1
self. ladds.
Layer 0 L-1 def arrention.forwar hidden.states, past.key.value (tt..eyy.aalee::pionnl[CCCC [Ckeannl[[a None):
10
Compute query states. key states, value stares
bsz, q.len. hidden.states. size))
12
13
Compute local allention to the current layer
14
local_ate
compute.local_attenttiii (query.,qqquyytttttttte [t, self. wtt.:/se....wwwiii........ window_size:..l key_states, value value_states)
15
observed if prefill:
local attal:.:. window size:.:- self.. window size
16
17
# Calculate preference seore
18
dispersion salcullate.ttt..ispeeee.
(bbserved.attn))
19
shift = calculate shift (observed
past.key.value. pref_seores [layer] (dispersion ddsspesionn**(//// (+I/tannl) shift=e(1/uu22)) # Calculate eviction indicator
21
mean oeservbeer atin)
23
var = var(ubserved attln)
24
indicator= attn_mean self. galmagamma- attn-var indicator inddit(in
26
a Update
past_key pref_scores append ( (pref_score) pref.score. score. indicator)
28
past kast.key www.ee.eee...eeeeeeeeeeaeee
kast_key. peiikey_vvauee value cake key. value, q.len. self. layer
30
else:
31
indicator=
paxt.key-value update. (corecindicatoo
33
Compute attention output using fasb_attention
35
atta output = flash attention forward (query states, key states. value states.
17
return attn_output, past_key_value
38
39
def cake(self, past key kast.key
value, seq len, surrent-layer idx): if seq_len<=self cache cacheesize-see size- window.size:
40
return past_key_var values
42
Cache budget management
prefill:
44
pref_scores past_key. values. pref.scores.
45
layer. = [pref.score /sum prrmmppree.oooe s total. for pref.score in pref scores]
46
# The following loop can be exccuted in parallel for cach layer layer-idx budgei_size
caumerate(layer.bddee budget,size min( (iiddegei,siie. seq.len self. window,size)
48
past.key_value evict_layer.kvcache (aa.teee.a.tttt.aatllll.e layer_idx, budget_size)
50
past.key.value
budget_size
5i
else:
bide kast_key. valuevalus... budget past.key.values = self. wew....vvi...yy..kvaa.... past.key.values. current.layer.idk. budget.size)
53
54
return pasl.key.values
56
def evict_ larettllleee.eeeeeeeeee kveache (self. gest.key_values layer_idx. budget.size):
57
bsz, num key value heads.. seq len head dim = past key values. key cache layer idx I. shape wum.key.value.gron
self num_hcads num_key.value.he indicater = -asluusy-values.. evict_indicaton [ayer_idx] indices indicator. lopk (budget. dims-l).indices
59
61
Update the eviction indicators pest.key.values.. evit.iiirvet eiet..iiiiiiiaooooolor..ttt [layer_idx] = sompress.kv indices, self. window size)
62
65
indices indices. unsqueeze(-l).
-1, -1. head dim)
Update the KV cache
65
past.key-values. key-e layerzidx compress.ky stytstates. indices self. window
past.key. values. value-cache (laeer_idx compress. (value states. indices, self. window_sine size)
67
return past.key.valos
16

---

Published as a conference paper at ICLR 2025
D PROOFS
Proof of Proposition Given the budget calculation formula in Eq.
we have:
B(m)
Pt
Biotal,
(13)
Σ m
Pk
k=
B
B(m+1)
P.
-
(14)
k=U
Pk
And the inequality Eq. (8) is equivalent to:
P m++
P
(15)
1k=0
Pk
2k=
PE
Since preference scores are always positive, this inequality holds. Therefore, we complete the proof that the allocated budget size monotonically decreases.
Proof of Theorem Recall that the eviction operation on layer I essentially retains KV pairs based on top-B, selection indices D, on the values I, where B₁ the target budget size, Therefore, proving Eq.(IO) is equivalent to demonstrating:
TopK(I)")) (I(") _((")) - TI, Bt).
(16)
The left-hand side of the above equation can be expanded as a cascading process:
D (1) (I ) D.T R
((+1) Topk(I! (1+1) Bi B((t1)), D (((11) e R
D
(17)
D(1-1) TopK(I! B((-1)), D E RB!T
For the vanilla method, Dj is obtained by single TopK operation:
Dj TopK(I, Bt). Given these definitions, our goal simplifies to proving that:
(18)
D (1 D;-
(19)
To prove this equality, we first show that these two sets have the same number of elements. Note that (( B ( and DDi = Bt. At the final stage, both methods have access to the preference scores of all layers, and the budget calculation for layer 1 identical both cases:
B( B P
B
(20)
Σ1 Pa
where Pi is the preference score for layer I, and Brotal is the total cache budget. Thus, (D(T-1) [D/]- Next, we further prove that elements in D Df2-1) and D, are identical. Given that B/ (m)) is
decreasing increases, for any stage [7, L 1), have:
"ppK(""",,(""" TI, B(m+1)) TopK(I!m) (m)
T B.(m-1))..
(21)
Therefore, we have:
D 3 D (((1) .... ((1-1) = Topk(I(L-1), BB B(b-1)) TI, Bt).
(22)
In the final stage, set D (L-1)) contains exactly the (((-11) B, highest-scoring elements from I which is identical to D,. Thus, we have proven that D (L-1) Dr, which completes the proof of Theorem
17

---

Published as a conference paper at ICLR 2025
33.25
33.00
LongBench Avg. Score
32.75
32.50
32.25
32.00
Full Cache (16 bit)
31.75
CAKE
KCVC (2 bit)
24:08
KCVC (4 bit)
KIVI (2 bit)
23.75
KIVI (4 bit)
KIVI (4 bit) CAKE
23.50
10%
20%
30%
40%
50%
Compression Ratio
Figure 9: Comparison of CAKE, KV cache quantization methods (KIVI and KCVC), and their combination (KIVI with CAKE) on Llama2-7B-Chen using LongBench dataset.
E COMPARISON AND COMPATIBILITY WITH QUANTIZATION METHODS
While CAKE focuses on selectively dropping less important KV pairs to reduce memory footprint, quantization approaches aim to compress the cache by reducing the bit precision of stored values. These two strategies are orthogonal and potentially complementary in nature.
In this section, evaluate CAKE against two quantization approaches: KIVI (Liu al.) 2024e), state-of- of-the-art method using asymmetric KCVT quantization (quantizes Key cache per-channel and Value cache per-token), and computationally efficient channel-wise quantization variant for both key and value caches (denoted KCVC), As is shown in Figure under matched compres- sion ratios, CAKE outperforms full cache (33.23 vs. 33.07) and both quantization methods 25% and 12.5% compression ratios. Notably, when combining KIVI-INT4 with CAKE, it demonstrates significant advantages in maintaining performance at high compression ratios. For instance, KIVI- INT4 with CAKE achieves 32.51 12.5% compression ratio, outperforming KIVI-INT2 (32.17) while using the same storage overhead. Even at very aggressive 6.25% compression ratio, the combined approach still maintains strong performance (32.48), showing the effectiveness of inte- grating eviction and quantization strategies.
F EXTENDED EXPERIMENTAL RESULTS ON LONGBENCH
In this section. we provide comprehensive experimental results on LongBench (Bai ct al. 2023
benchmark focused on long-context understanding with input lengths ranging from 1,235 18,409 tokens. We present detailed performance evaluations across cache budgets from 64I to 2048L for three base models: Llama2-7B-Che (Touvron et al.) 2023). Llama3.1-88 -8B-Instruut (Dubey et al.) 2024), and Mistral--BB-Inttruccl
(Jiang et al. 2023 (Appendix F.I.. To demonsirate CAKE'S generalizability, we conduct additional experiments on two more LLM architectures: Qwen2.5-7B-
Instruct (Team] 2024) and Gemma-77 7B-Instrucl (Team et al. 2024) (Appendix F.2). Furthermore,
we evaluate CAKE'S scalability on larger models ranging from 13B to 70B parameters, including Llama2-1 -3-Chht, Swen2.5-33BB-Inse and Llama3-70B. 1anaa-70B-Instruct (Appendix
H3
18

---

Published as a conference paper at ICLR 2025
Table 4: LongBench Task Type.
Task Type
#English Task
#Code Task
Multi-document QA
Single-document QA
Summarization
Few-shot learning
Synthetie Tasks
Code Completion
2
LongBench is a long context dataset comprised of six task categories covering key long-text applica- tion scenarios: single-document QA, multi-document QA, summarization, few-shot learning, syn- thetic tasks and code completion. Our experiments are conducted on all of its 14 English subtasks and 2 code subtasks, spanning throughout all six categories (Table)4), with input lengths ranging from 1,235 to 18,409 tokens.
F.1 DETAILED PERFORMANCE ANALYSIS ACROSS CACHE BUDGETS
Tables 5 and present detailed LongBench results of CAKE and comparative methods applied to Llama2- 7B Chat, Llama3. 1aama3.1-8B 1-88-Instuuct. and Mistral-7B M respectively.
As shown in Table CAKE consistently outperforms other KV eviction methods in terms of average score. Additionally, CAKE achieves top or second across each individual subtask and at all cache sizes. Notice that when using with 512L or more total budget size, CAKE achieves evaluation scores that comparable with full cache across all subtasks.
Results on Llama2-7B-CHR
Results on Llama3.1-88 ---ttruuttt The evaluation results are detailed in Table The Llama3.1 series models have all incorporated the Grouped Query Attention (GQA) mechanism, and are able to support longer context length. The results show that CAKE continues to perform outstandingly across all subtasks in models based on the GQA structure.
1iin------ntttt-v.... Around all 16 subtasks, CAKE shows consistent advantage throughout all the cache size budgets we have tested. Especially when the budget is limited, say 64L, CAKE, equipped with the robust indicator and strategy, expresses relatively strong resistance score decline.
Results on Mistrrl-7B--Isstrutt--...
Table contains the results of
The results as a whole demonstrate that CAKE, compared with other methods, consistently outper- forms across all task types in LongBench when applied to the tested models and under various cache size budgets ranging from 64L to 2048L. This demonstrates CAKE's effectiveness and broad ap- plicability KV cache-efficient long-context processing across various domains and across widely used open-source LLMs.
19

---

Published as a conference paper at ICLR 2025
Table 5: Performance comparison over 16 datasets of LongBench on Llama2-7B- 7B-Chat for cache budgets from 64L to 2048I. The best result is highlighted in bold, the second best in underline.
Single Document QA
Malii DocumenD QA
Summarication
Few Learning
Symthetic
Code
Method
Nrn-A
Qhaper
Mf-
Hospo AA
2WiKiMQA
Musique
GovRespat
QMSum
MutiNews
TREC
TriviaQA
SAMSom
PCoumi
Pk -K-
Avg.
Lac
RB-1
Llama2:7P CB-Chat Bust
Full
Full
20.68
20.19
36.09
27.83
31,45
8.21
27.09
20.71
26.24
64.0
83.09
41.41
2.94
7.75
58.51
$2.29 33.07
Hlama2-7B-Chail Puus
647.
Streaming...M7...
7.35
16.62
6.84
10.73
14.68
2,84
13.34
18.42
12.76
17.50
19.38
11.44
0.25
4.50
20.87
14.91
12.03
H20
12.77
13.76
21.00
19.16
26.53
4,47
14.09
20.01
18.49
32,00
70.92
34.79
2.75
3.00
42.09
41.36
23.57
TOVA
8.64
14.06
12.61
13.71
13.56
3,26
11.92
16.59
13.73
36.00
74,88
31.88
2.15
4,00
36.72
31.94
20.35
SmapKV
1263
14.95
19.93
15.48
22.94
4,78
15.54
19.50
15.10
34.50
72,02
34,10 16
2,00
4,12
40.03
36,22 29 22.75
PyrmmiKVV
10.4
15:53
28.43
23.53
28.82
5.10
15.33
19.79
15.29
34.50
76.16
2.88
CAKE
12622
16.64
28.04
21.46
28.00
5.19
17.20
20.16
17.23
35.00
74,30
35.02
3.33
4.50
45.86
42.33
25.43
Llama2:7B. Chat, Bl
128L.
SL
13:57
13.33
22.39
20.16
17.12
6.38
16.72
19.58
15.37
31.50
74.84
33,00
1.67
4,00
43.21
40.45
23.33
14.59
26.52
20.05
28.73
5.00
16.20
20.30
72.75
37.10
49.68
26.06
TOVA
13.56
15:46
20.28
22.22
25.01
4,79
15.72
19.11
15.58
42,50
79.58
37.06
3:45
4.61
46.34
39,70
25.31
SnapKV
13.70
16.27
32.52
22.76
28.59
7.56
18.87
51.86
28.09
PyramidKV
13:45
16.61
31.21
23.74
28.43
6.89
18.88
20.56
19.81
44.50
$0.32
37.82
2.29
8.00
51.19
47.56
28.20
CAKE
15.15
16.38
32.17
25.25
31.09
2,00
19.47
21.07
20.36
48.50
$0.66
37.77
4.42
9.00
51.48.
48.90
29.29
Llama2-78 7B-Chat, Bou 2562
StreamingLLM
13.20
13:45
24.36
21.50
24.08
5.74
17.28
18.97
18.12
44.00
78.43
37.55
1.58
4.50
53.85
47.03
26,48
H20
15.82
14.72
27.26
21.07
27.00
6.32
19.75
20.29
22.23
45.50
79.64
37.93
2.93
4.00
53.91
50.33
28.04
TOVA
13.78
14.69
23.81
25.92
28.93
7.24
18,00
19.48
19.04
58.00
82.33
39.33
2.69
7.00
51. 42
47.16
28.68
SnupKV
15.47
15.94
34.51
25.17
30.
9.30
20.44
20.89
21.78
39.22
55.93
30.52
PyrmmidKV
15.13
15.86
33.93
25.58
29.51
9.23
20.35
21.22
22.01
58.00
82.39
38.47
2.15
7.50
55.81
49.53
30.42
15:93
16.80
26.26
30.54
9.01
20.98
39.01
56.97
31.21
Llama2-77 2- 2B-Chat, Byy 5122
StreamingLLM 1261
24.21
5.40
19.19
55.57
H20
15.86
15.89
30.14
22.14
28.89
7.01
21.40
20.89
23.54
52.50 81.81
40.55
3.19
5.00
55.6
68
51.07
29.72
TOVA
13663
15.47
26.26
26.55
31,25
7.63
19.88
83.10
40.63
SnupKV
16.42
17.19
36.52
26.32
31,15
9.09
21.33
21.11
23.58
64.00
81.82
39.29
3.33
8.50
58.00
52.46
31.88
PyramidKV
16.91
18:30
36,02
26.74
29.53
9.13
21.53
21.22
23.72
63.50
$1,88
39.76
3.00
8.00
58,02
51.14
31.77
CAKE
16.36
19.31
37.82
26.51
31.34
8.10
21.77
21.42
24.36
64.00
$2.63
39.62
3.33
8.00
57.73
52.33,
32.16
Alama2-78-Chia Buss
1024Z
StreamingLLM
13.73
16.75
26.28
25.97
25.67
5.87
22,21
19.27
24.07
61.00
30.88
40.90
2:28
1.00
56.71
48.79
29.16
H20
16.58
17.67
32.04
26.01
28.73
7,81
22,87
20.99
25.05
60.00
83.02
40.06
2.90
3.50
57.57
52.02
31.05
TOVA
15.00
17.32
29.57
27.48
31.71
7.04
20.54
41.78
56.69
SnupKV
18:05
19.73
38.07
27.57
30,86
8.41
23.41
20.79
25.22
63.50 82.13
40.98
1.33
7.50
58.20
52.23
32.50
PyramidKV
17.94
38.85
27.25
29.99
7.91
21.55
25.45
57.72
32.55
CAKE
18.35
20.93
38.49
27.99
31.15
8.59
24.29
21.06
25.40
63.50
83.33
41,30
3.11
9.50
57.82
52.93
32.98
Llama2-7B-C 7B-Cha
20/8L
Streaming L 1695
17.98
32.25
28.15
30.31
7.66
23.79
19.97
25.75
63.50
83.17
40.36
2.05
5.17
57.
37.36
$1.30
31.61
H20
18:35
20.26
36.71
28.65
30.85
8.33
25.05
21.20
26.31
63.00
83.32
40.78
2.60
7.00
58.
52.92
32.72
TOVA
18.3
20.52
35.38
28.10
30.91
8.87
24.86
20.87
26.22
64.00
83.13
41.99
2.75
7.00
58.07
52.77
32,76
SnapKV
18.50
22.18
38.18
27.72
30.92
8.37
25.87
21.35
25.92
64.00
83.13
40,95
3.32
8.50
58.18
52.23
33.10
PyramidKY
18.51
22.68
37.04
27.76
27.84
7.95
26.42
20.95
26.08
64.00
83.26
40.72
3:05
7.75
58.80
52.55
32.83
CAKE
19.28
22. 16
38.70
27.72
31.58
8.91
26.19
20.70
16.19
64.00
83.25
41.02
3,86
9.00
$7.92
52.76
33.33
20

---

Published as a conference paper at ICLR 2025
Table 6: Performance comparison over 16 datasets of LongBench on Llama3. 1lama3.1-8B-... 1-1---88Btttttt for cache budgets from 64L to 2048L. The best result is highlighted in bold, the second best in underline.
Single
Single-Document Document QA
Multi Document QA
Summarication
Few-shot Learning
Synthetic
Code
Method
Niv-A
Qusper
Hospar MP--0
2WisiMQA
Masique
GoeRspon
QMSum
MattNews
TREE
TrivkaQA
SAMSum
PCaun
PR
Lac
RB-1
Avg.
Llama3.1-80 I-sttnerr Bust
Full
Full
31.06
45.43
53.78
55.04
47.34
31.29
34,87
25.33
27.49
72.5
91.65
43.81
6.0
99.5
63,36
56.65 49.06
Llama3.1 In--ruttt Buss
64L
SMreamingLLMMM
22.52
30.62
46.45
40.02
23.96
16.92
20.95
15.18
38,50
82.71
34.97
6.50
99.50
52.01
45.07
37.38
H20
27.35
24.95
40,70
51.75
43.57
26.72
20.23
22,06
19,35
39,00 88. 30
38.69
6.10
98.00
54,48
46.41
40.48
TOVA
27.14
41.82
53554
27.42
16.67
39.12
6.50
47.41
SnupKV
26.09
24.06
41.22
51.32
45,11
25.71
17.38
21.47
15.81
38.50
84.98
37.36
6.50
99.50
53.16
46.09
39.64
PyramidKV
24.71
25.58
17.14
15.96
37.27
6.50
98.50
52.38
45.54
39.03
CAKE
27.03
26.56
42,96
52.42
45.00
27.73
21.29
18:57
18.26
40.50 88. 42
39.11
6.07
99.00
54,
47.35
40.93
1llma3.1-8B.
128L
Streamingl.I.M
23.01
22.50
31.00
46.79
40.52
24.76
18.38
20.89
16.96
40.50
86.00
38.55
7,00
99.50
57.
06
47.16
38.79
H2O
26.18
21.83
44,08
51.33
43.29
27.78
20.79
22.36
20.69
41.50 89.88
10.91
6.20
99:00
58,
38.12
49.82
41.67
TOVA
29.44
27.03
45,20
54.15
44.87
28.87
22.21
20.54
20.52
53.50 92.27
40.99
6.17
99.50
51.14
43.15
42/47
SnupKV
26.40
27.67
47.83,
53.40
44.20
28.68
20.56
23.11
20.13
45.50 39.36
39. 39.98
6.33
99.00
57.6
$7.69
50.69
42.53
PyramidKV
25.57
56.98
CAKE
27.41
32.01
49.58
$2.99
45.76
28.32
21.85
23.15
21.56
47,50 89.61
40.90
6.33
99.50
57.92
49.83
43.39
Llama3.
1-8B--B-Bs-ueeee Bos 2561
StreamingL.LMM
23.35
23.57
24.73
30.31
46.42
39.19
24.23
20.61
21.10
19.46
45.50
87, ..
41.01
6.50
99.50
59.55
48. 70 39.87
H20
25.93
27.16
43.86
52.23
28:28
22.10
22.75
47,50 23
41.
6.28
61.02
42.93
TOVA
29.70
30.85
48.80
54.36
46.71
30.00
23.21
22,18
22.88
57,50
91.83
42.77
6.05
99.50
54.83
46.15 44.21
SnapKV
27.22
35.25
50.96
54.50
45.73
28.56
22.75
23,80
22.83
55,50
91.28.
41,42
6.03
99.50
60.64
52,80
44.92
PyrmmidVV
28.73
36.14
51.01
54.57
29.58
90:37 56.00
41.09
6.25
60.03
51.45
CAKE
29.75
38.74
52.27
54.52
45.39
29.57
24.15
24.24
23.77
$7.00 90:97
41.82
6.10
99.50
61.01
51.90
45.67
Llama3. 1-8B--8BBBiiii Instructt Brsi
5122
StreamingL.MMM5
25.64
27.48
33.30
47.36
40.06
24.80
23,16
20.80
22.85
57,50 87.69
42.08
6.50
97.00
60.51
51.28
41.75
27.76
29.22
24.71
24.56
91.38
42.10
6.36
62.30
44.39
TOVA
30.11
35.77
49,88
54.58
45.86
29.73
25.45
22,73
25.00
63.50 11.90
43.33
6.00
99.50
58.38
48.58
45.64
SnupKV
30.76
42.03
52.13
54.15
46.14
30.51
24.98
24.24
24.65
64.00 92.05
42.04
6.08
99.50
62.62
54,90
46.92
PyramidKV
30.47
42.15
52.17
54.67
45.25
30.00 60
25.00
24.33
24.51
62.50
91.24
41.67
5.95
99.50
61.58
53.89
46.59
CAKE
31.82
42.99
51.65
54.37
46,89
30.73
26.36
24.94
25.27
63,50
91.54
42.52
6.33
99.50
62.31
$4.30
47.19
Llama3. 1- 1-8--nsstruc,, Busa
1021L
Streamingl IIM 26.64
30.77
35.59
47.31
42.03
24.17
25.81
21,31
25.66
63.50 38. 84.84
42,76
6.50
88.00)
61.36
53.47
42,73
29.57
36.15
54.43
29.04
26.47
6.36
55.30
TOVA
30.66
40.95
51.09
54.58
46.51
30.62
28.12
23.61
26.21
68.00 91.19
43,80
5.92
99.50
60.73
$2.64
47.15
SnapKV
30.95
41.74
52.58
55.09
16.83
27.87
21.57
25.99
92.03
42.60
6.50
63.00
PymmiiKV
30.54
43.64
52.73
55.29
46.29
31.28
27.53
24.50
26.00
68.00
92.09
41.75
6.05
99.50
62,
02.35
55.44
47.69
CAKE
30,88
44.95
52.38
55.49
46.99
30.82
28.68
24.91
26.39
69.00
94
42.60
6.00
99.50
62.65
$6.89
48.13
Llama3.1
Ins-uuctutt, Brs.
20/8L
Streaming 27.40
36.91
37.85
49.23
44.66
24.31
28.57
21.67
27.12
67.50 90.98
42.49
6.12
87.00
63.06
55.39
44.39
H2O
29.65
39.53
48.64
54.23
46.50
29.28
29.97
23.68
27.21
68.00 91. .48
43.06
6.11
99.50
63.
56.91
47.30
TOVA
30.30
43,40
52.78
55.12
46.38
30.62
30,80
24.10
27.09
70,50 91 48
43.63
6,00
99.50
62.76
55.32
48.11
SnupKV
55.25
30.24
27.32
91.48
63.23
PyramidKV
31.13
45.06
53,80
55.78
46.59
30.89
30.25
24,8
27.35
71.00 91.65
42.62
6.00
99.50
63.27,
56.44
48.51
CAKE
30.79
45.83
46.60
30.47
31,12
27.16
43
6.00
99.50
63.28
48.55
21

---

Published as a conference paper at ICLR 2025
Table 7: Performance comparison over 16 datasets of LongBench on Mistral-7B-Insteee for cache budgets from 64L to 2048I. The best result is highlighted in bold, the second best in underline.
Single
Single-Document Document QA
Multi Document QA
Summarication
Few-shot Learning
Synthetic
Code
Method
Niv-A
Quaper
Hospar
2WitiMQA
Masique
GoeRspon
QMSum
MattNews
TREE
TrivkaQA
SAMSum
PCaun
PR
Avg.
Lac
RB-1
Boaaral-7-Inssmatt-e Best
Fuli
Full
28.26
39.33
19.17
45:29
33.21
21.57
32.25
22.16
24.58
75.5
88.36
43.78
0.5
83.0
19.79
51.3 43.21
Mistral-7B 7B-itaal----ttttaat--0.... Instract-v0.33 Bou
64.
L
17.59
21.81
23.39
34.93
26.73
15.14
13.11
18.29
12.59
37,50
82.27
33.97
0.00
28.00
36.66
36.79
27.42
H20
21.04
26.57
34.31
39.34
28.84
17.95
20.11
19.25
16.64
36,50
85.91
36.79
2,00
70.00
41.53
41.28
33.63
TOVA
21.87
41.60
28.35
19:46
89.16
0.00
32.78
SnupKV
20.63
24.03
34.69
41.22
29.94
18.04
15,42
18.82
12.62
36.50
86.44
36.70
1.00
46.50
37.77
40.65
31.31
PyramidKV
20.33
21.96
15.34
18.46
12:55
36.65
0.00
38
39.30
CAKE
21.83
26,85
35.19
39.70
28 89
19.70
20.12
19.19
17.21
38,00
86. 49
37.02
0.50
75.00
41.17
42.18
34.11
Mistral- Mistral-7B-nnssna 78.nnmmattx00.33 Bysi
128L
Streamingl.IMM
16.91
21.51
24.85
34.14
26.99
16.64
15.67
18.61
14.40
43.50
83.26
37.00
0.00
23.00
41.63
42.12
28.76
H2O
21.25
26.66
35.13
38.82
29.80
18.88
21.00
19.50
18.63
41.00
87.6 87.64
38.25
1,50
67.00
44.35
47.29
34.79
TOVA
22.47
24.26
37,22
42.26
28.85
19.97
19,10
18.70
17.86
63.00
88.98
37.71
0.50
56.50
37.42
37.22
34.52
SnupKV
21.02
27.26
41.25
45.15
29.23
22.75
20.47
20.17
17.75
42,50
87.28
38.01
0.50
69.50
44 48
45.69
35.81
PyramidKV
21.73
41.46
43.20
37.11
1.05
CAKE
22.31
29.15
43.51
44.51
30.36
22,85
21.56
20.47
18.96
47,00 88. 60
39.36
1,00
76.50
44,96
46.19 37.33
Mistral-7 7B 78--nxmaatt-x.33 Besu 256L
StreamingL.LM
10
23.41
25.99
34.55
26.78
16:25
17.97
18,82
16.56
52,50
84.26
39.42
1.00
16.50
44,
14.99
50.12
H20
22.36
27.14
30.25
22.35
20,06
39.49
1.50
47.
35.89
TOVA
22.52
27.29
41.32
44.49
31.08
22.36
21.99
18.96
20.00
68.00
88.58
39.76.
0.50
66.00
40.71
41.35 37.18
SnapKV
23.06
29.43
44.61
45.05
30.39
22.89
22.32
20.77
20.19
56.00
88.
39.34
0.00
83.00
48.09
48.72 38.88
PyrmmidKV
22.66
28.85
30.43
$3.00
38
38.88
1.05
46.
38.59
CAKE
24.01
30.78
45.88
45.06
31.45
23.43
22.81
20,87
20.92
58.10 38.29
40.16
01.50
87.00
47.76.
48.43
39.74
7B-
I0.. Bysa
512E
19.93
25.83
28.25
36.21
27.60
15.41
21.23
18.67
19.50
64.00
36.45
39.75
1.00
16.50
46.
70
47.57
32,16
22.34
40.12
30.25
24.37
21.80
2,00
TOVA
25.06
31.69
44.22
44.90
32.48
23.44
24.33
19.81
21.77
71.50
88. 18.46
42.39
0.00
77.
44. 54
45.66
39.86
SnupKV
24.88
32,92
47.31
45.23
31.78
23.08
24.60
21,80
21.52
65.50
88.61
40.43
0.50
88.50
50.59
50.22
41.09
PyramidKV
24.15
33.10
46.62
44.16
34.14
23.37
24.22
21.15
21.29
68.00, 88
40.27
1.55
87.50
48.74
49.17 40.87
CAKE
25.82
34.24
46,93
45.34
31,71
23.38
25.81
21,47
22.41
68.00
88,
41.27
0.50
90.00
49.04
49.67
41.50
Mistral-7B
7B------7--tttt--..... Buxi
1021L
Streamingl I.M 20.96
28.05
30.03
37.06
27.56
16.03
24.03
19.07
22.79
67.00 87 61
40.96
1.50
21
48.05
48.87
33.82
23.78
43:24
31.077
20.41
23.93
88.84
1.50
50.60
TOVA
26.97
31.51
45.58
41.32
32.58
22.83
26.91
20.75
23.19
75.00
88. 66
43.17
1.00
91.00
17.51
47.06
41.96
SnapKV
26.63
45.75
32.20
26.71
23.18
41.37
0.50
50.60
51.79
42.18
PyrmiiKV
25.51
36,02
47.72
44.74
33.16
23.91
26.55
21.83
23.27
70.50
88 41
40.94
1.00
87.00
50.17
50.93
41.98
CAKE
26.09
36.34
48.11
45.97
32.39
23.49
27.56
21.45
24,03
72.50 88.61
42.71
0.00
91.50
$1,06
51.25
42.69
Mistral Ins.ruct-v0.3,, Busi
2018E
Streaming
32.48
33.26
40.53
29.57
15:69
26.64
19.70
23.64
69.00 $8.41.
42.44
2,00
28.50
48.69
50.85
35.88
H2O
25.76
34.28
44,42
44.41
32.17
20.14
28.79
21.23
24.38
71.00 41
43.49
1.50
83.00
50.51
50.39
41,43
TOVA
27.48
37.20
47
45.25
33.14
23.73
29.37
21.73
24.52
75.00 88.66
43.33
1,000
62.00)
49.08
50.18
42.45
SnupKV
26.43
48.27
45.37
33.08
24.40
49.62
51.28
PyramidKV
26.00
37.63
48.25
45.82
37.13
24.00
28.90
22.13
24.35
72,50
88.36
42.28
1.00
84.50
50.06
$1.07
42.75
CAKE
25.43
37.84
48.30
45.60 60
32:21
30,00
24.35
43.38
(1.50
89.00
50.68
51.48
42.98
22

---

Published as a conference paper at ICLR 2025
Table 8: Performance comparison over 16 datasets of LongBench on Qwen2.5-7B-Inse and Gemma-7Bi Conma-7B-Instruett The best result is highlighted in bold, the second best in underline.
Single Single-Deementt QA
Mult-DDxument QA
Summarization
Fow-shot Learning
Synthetie
Code
Method
Nnt QA
Qesper
MF-4n
HotporQA
ZWiLIMQA
Masique
GoeReport
QMSum
MotoNews
TRBC
TrivisQA
SAMSSum
PCount
PRos
Avg.
Lor
RB-
Qwen2.5.77 5.2.2..BBBunn Insstruct, Bost
Fult
29.05
30.24
31.78
23.64
23,96
72.5
100.0
59.61
Qwen2.5-7B Instrcct., Bias
128L
StreamingLLM
18.80
26.71
16.24
15.70
18.01
H2O
22.73
24.52
31.04
43.89
40.19
18.92
18.28
19.82
16,03
42,00
83.11
41.88
8.50
96.50
50.39
05
38.30
TOVA
21.01
23.98
33,65
47.65
38.97
19.42
19.84
19.41
15.54
51.50
85.43
42.49
8.50
93.50
42,60
44.48
38.00
SmppV
23.85
27.19
42,43
51.31
42.59
24.38
18.19
19.76
14.54
42,50
84.24
41.10
9.00
97.00
49.45
53.95
40.09
PyramidKV
19.56
23.39
37.72
48.33
38.33
19.53
15.10
19.32
12.05
42.00
83.29
38.84
9.00
84,50
47.87
50.39
36.83
CAKE
25.99
46.22
25.59
19.51
20.95
15.82
50.62
55.27
41.68
Dwwn2.5-7B- 31B-Insruct, Busa
102
Streamingt
22.72
17.99
24.33
19.47
87.53
34,000
55.
37.38
H20
26.45
34.94
40.49
48:63
42.02
22:27
25.67
20.90
22.41
59.00
87.83
45.07
8,50
98.48
59.77
63.88
44.14
TOVA
26.36
37.97
47.88
55.32
45.72
28. 38
26.15
21.64
22.70
69.50
88.64
44.23
8.50
100,00
56.66
60.07
46.23
SnupKV
29.34
41.61
45,50
29.30
25.63
23.06
65.50
100.00
58.
47.27
PyramidKV
29.34
38.60
50.17
55.67
45.12
27.82
23.26
22.16
20.55
62.50 86.85
43.26
8.50
100.00
S7.76
61.99
45.85
CAKE
29.47
42.71
52,12
56.11
46.41
29.13
26.86
23.12
22.72
62.50 89.23
45.46
8.50
100.00
59.11
64.79 47.70
Gemma 7B Instruct, Brss Full
Full
15:33
34.27
47.43
28.2
22.69
7.48
26.64
19.67
23.84
69.0
79.41
32.59
1.0
44.5
-17.39
45.96
34.09
Gemma 7B Instruct. Bus 128L
S
10.65
20.76
27,41
24.74
20.06
5.91
13.53
16.98
14.30
47.00
74.51
29.16
4,00
9.50
-66
45.70
25.68
12.88
23.58
17.70
18.17
18.23
79.12
31.84
47.28
30.49
TOVA
13:73
27.38
42.38
28.75
23.34
7.80
17.06
17.76
17.61
57.50
79.65
31:35
1.50
49.00
41.85
43.22
31.24
SnapKV
12.88
25.45
17.14
17.65
48.50
47.
20
PyramidKV
12.51
28.02
42.24
28.18
23.33
6.66
16.82
17.94
17.11
46.50
77.94
30.94
2.00
46.00
45.2
48.35
30.61
CAKE
14,09
30.33
44.33
28.63
24.47
6.70
18.20
18.83
18.98
53,00
79.54
32.14
1.50
48.00
48.71
59,59
32.38
Gemma In-nnntttt, Bus
1024E
StramingLLM
13:23
27.01
29.08
25.39
22.15
5.67
20.91
17.74
21.97
67.50 79.28
33.64
1.50
11.50
49.50
49.25
29.73
13.39
21/11
22.85
18.68
22.73
79.78
1.57
TOVA
13.58
32.92
45.60
28.97
24.50
7.96
22.19
18.95
22.86
73.00
79.94
32.61
1.50
44.50
48.
47.29
34.02
SnapKV
13.72
33.76
46.00
29.19
23.85
7.49
22.44
19.35
22.73
71.50
79.96
33.35
1.50
43.50
18.61
17.86
34.05
PyramidKV
21.66
19.37
CAKE
18.16
35.12
47.67
28.24
23.12
6.59
23.28
19.72
23.40
68.50
79.47
34.07
1.50
43,00
48.84
49.20
34.18
F.2 EXTENDED EVALUATION ON ADDITIONAL MODEL ARCHITECTURES
To demonstrate the generalizability CAKE across different model architectures, we conduct ad- ditional experiments on Qwen2.5-7B-Innse Instruct and Gemma-7B- BB-nstruut. The experiments are con- ducted on two memory scenarios: a low setting (Brotal = 128L) and a high setting (Brull = 1024L). As shown in Table similar to results on Llama and Mistral, CAKE consistently outperforms base- line methods across both low-memory and high-memory scenarios on Qwen and Gemma architec- tures. CAKE demonstrates significant advantages in low-memory settings by rationally allocating cache sizes and evicting KV pairs with more robust indicators. Under the high-budget setting, CAKE achieves superior performance, even exceeding the full-cache baseline Gemma (34.18 34.09). These results further validate CAKE's 's adaptability across different model architectures, maintaining its performance advantages regardless of the underlying model design.
23

---

Published as a conference paper at ICLR 2025
Table 9: Performance comparison over 16 datasets of LongBench on different models with sizes from 13B to 70B. The best result is highlighted in bold, the second best in underline.
DocuunnD QA
Mult--Doumeni QA
Summaatcation
Fow-shot Learning
Synthetic
Cooe
Method
NrtQQ
Qusper
Mileei
HogusQQ
ZWIKMOAA
Masique
GiveRoprrt
QMSum
MuhNNews
TREC
TrivisAA
SAMSum
PCont
PR-cn
Loe
RB-1
Avg.
Llama2-1 13B-Chat Bust
Full
Full
14.39
17.07
27.52
27.52
20.89
26.61
68.5
87.75
2.27
15:25
48.27
49.87
29.95
Liama2-15 13B Chat. Brsi
128L
Streaming
10.48
13:44
19.16
12.86
13.81
4.90
15.63
18.91
16.26
40.00
81.66
34.00
2.29
5.00
40.38
35.16
22.75
120
12:74
17.63
25.71
14.12
13.29
3.84
19.35
19.71
21.51
40.00
50.62
39.37
2.03
10.50
40.75
41.31
25.15
TOVA
11.18
17.46
18.17
16.51
58.00
7.50
25.01
SuapKV
13.33
14.73
24.31
13:17
15.10
1.68.
20.37
20.01
20:44
12.00
86.71
38.35
3
11.50
43.07
12.28
25.85
PyramidKV
13.04
13.91
25.26
14.37
15.24
5.68
20.01
19.74
20.52
43.50
85.37
38.83
2.80
13,00
43.57
41.31
25.95
CAKE
12.19
17.17
27.28
13:12
14.73
4.33
20.04
20.14
20.74
42,00
87.75
40.60
2.96
12.50
45.52
43.95
26.56
Llama2-13B--C2 Busi
10244
Streamingt
12.47
12.47
15.03
19.85
11:39
13.68
3:90
23.48
19.33
24.88
63.00
86.
40.12
4.40
6.50
47.
45.56
27.32
H2O
14.76
19.89
29.10
14.84
14.05
4.41
23.38
20.18
25.03
62.50
82.
41.85
1.13
12.50
47.34
47.07
28.76
TOVA
13.69
24.41
23.81
20.22
24.99
69.00
88.15
42.18
3.19
13.25
46. 01.71-
29.10
S
14.39
24.65
20.74
25.08
86.47
42.26
2.62
14,00
48
PyramidKV
14.15
18:30
28.35
12.98
14.42
5.10
25.05
20.84
25.62
68.50
88.21
42.52
2.28
15.00
47.99
48.
29.84
CAKE
15.33
17.84
29.05
13.99
14.30
5.46
24.43
20.60
25.29
68.00
87.42
42.58
3.14
15.50
48.40
48.41
29.98
Qwen2 5-32B-InstruccC Bus
Full
29.25
45.24
51.9
63.31
61.5
38.15
30.33
23.38
23.01
73.5
87.86
45.36
12.0
100.0
51.3
38.17
48.39
Owee22.-32BB Instract, Breu
128L
Streamingl J.M 16.26
22.62
25.95
44.14
47.57
23.37
15.66
17.50
13.22
41.50 81.47
38.03
12.00
58.04
43.51
32.46
33.33
H2O
24.03
30.84
17.00
19.67
16.27
83.64
12.22
45.85
TOVA
25.53
27.73
38.07
$7.08
54.92
32:92
19.28
17.47
15.08
60.50
77.55
41.02
11.81
97.58
38.05
26.75
40.08
SnpkKV
21.28
26.68
10.43
56.74
55.61
33.75
17.15
19.65
14.55
48.50
86.46
10.86
12.50
95.58
45.75
33.55
40.56
PyramidKV
18.31
24.62
36,85
55.50 30
55,40
32.38
15.33
18.02
13.15
46.50
84
39.62
11.50
88.83
44.05
31.76
38.51
24.03
31.85
18.54
20.42
50.00
97.42
Qwee2.5-522
2aw.22.5552--mmeee 335tt-nttrut,, Busi
1024L
StreamingLMM21.
211444
29.11
31,07
46.44
47.31
26.17
23.76
19:07
21.68
63.50
86.12
42.61
12.50
63.08
49.84
36.13
38.74
H2O
25.96
32.84
40.35
56.46
55.57
32.91
25.06
21.02
21.94
60.50 86.63
14.85
12.50
97.58
50.96
37.15
43.89
TOVA
29.38
60.91
34.77
24.86
20.69
22.03
73,00
87.34
45.09
12.00
100,00
49,
35.90
46.47
S
29.93
42.29
50,00
61.97
61,37
37.22
25.01
22,47
21.94
69:50
87.31
44.42
12.50
100,00
51.54
37.32
47.17
PyramidKV
29.39
41.56
49.31
63.12
61.52
38.25
22.47
21.87
18.30
70.50
87.78
44.01
11.83
100.00
50.27
37
46.73
CAKE
30.77
50.73
61.59
60.60
37.
25.86
22.77
22.18
22.50
87.78
12.12
100.00
51.13
37.71
Mamma--70B-Insteee Rlyaa Full
Full
26.08
47.26
49.73
49.61
54.84
27.51
31.01
22.48
27.99
73.5
92.88
45.21
11.0
68.5
37.72
67.28 45.79
Lluma3-708- Ios-tuuuu Busi 1281.
Streamingt
22.33
28.13
28,07
42.85
50.73
22.
17.75
20.02
44.00
82.95
40.05
11.50
67.50
37.53
H20
24.92
29.15
36.50
45.98
50.40
24.52
19.38.
20.61
21.40
43.50
91.78
42.90
10.55,
70.15
46.92
63.45
40.13
TOVA
20.61
25.79
30.59
43:67
47.05
24.33
6.18
17.45
18:22
58.50
91.52
37:04
7.46
34.62
41.71
51.35
34,76
SnupKV
25:29
33.37
44.25
48.01
52.93
20.01
18.90
21.53
20.30
48.50
91.63.
41.53
7.50
68.00
45.87
64.16
41.20
PyramidKV
25.87
34.16
19.17
21.20
21.02
91.43
46.52
CAKE
25:21
39.26
44.39
47:47
54.48
27.22
21.26
22.37
22.25
$7.00
91.28
43.69
10.00
67.50
44.41
64.16
42.62
LJama3-700 Instruct, Bwai
1021L
Streamingl
25.1
35.93
32.89
45.57
49.49
21.53
24.42
20.28
26.03
66.00
91.32
43.53
11.00
68.00
42.61
65.87
41.86
H2O
26.56
26.26
26.05
21.52
26.75
67.00
92.88
11.00
TOVA
27.03
45.73
47.61
49.77
55.54
27.02
26.04
21.59
26.58
73.00
92.95
46.24
11.00
68.00
39.
62.50
45.03
SmpKKV
26.63
45.96
47.66
48.57
$4,92
27.60
26.31
22.34
26.75
72.00
92.88
45.10
10.00
68.50
40.72
67.22
45.20
PyramidKV
26.99
48.01
49.04
54.17
27.37
26.08
22.45
26.22
73.00
91.95
45.10
10.50
68.50
38.44
67.99
45.08
CAKE
26.95
46.94
26.39
73.50 92.44
11.00
44.23
F.3 EXTENDED EVALUATION ON LARGER-SCALE MODELS
We further extend experiments to models ranging from 13B to 70B parameters, including 1lama2-13B-Chal Qwen2.5-325 Instruct, and Llama3-70B. Instruct, under two settings: B 128L and Binaa 1024L. In Table CAKE outperforms baseline methods across all model sizes. Notably, under the high memory setting (Brual = 1024T), CAKE achieves even better performance than full-cache settings for both Llama2-13B (29.9 (29.98 vs. 29.95) and Llama3-70B (45.83 vs. 45.79), showing that our method scales well to larger models while maintaining its efficiency advantages.
24

---

Published as a conference paper at ICLR 2025
G EXTENDED EXPERIMENTAL RESULTS ON NEEDLEBENCH
In this section, we provide more detailed experimental results on NeedleBench (Li al.. 2024a)
using Distral-7B-Inntrrutt-h// (Jiang et al. (023) and 1-mma3.--8-Insseeee (Dubey et al. 2024) as backbone L.Mss. We examine three distinct tasks on NeedleBench: (1) Single-Needle Retrieval (S-RT): Assessing precision in locating individual details within extensive texts. (2) Multi-Needee Retrieval (M-RT): Evaluating the retrieval multiple related information pieces dispersed through- out lengthy texts. (3) Multi-Needle Reasoning (M-RS): Measuring complex reasoning abilities by extracting multiple information elements from extended contexts to answer questions.
All experiments conducted rigorously following the original NeedleBench protocol: Levenshtein distance is used to measure the similarity between predictions and references. Each case is repeated ten times ensure stable scores, and the results weighted-averaged to obtain overall score, providing balanced representation of each task.
Results on - The results presented in Table and Table demonstrate that CAKE consistently outperforms other KV cache eviction methods across all evaluated tasks. CAKE exhibits the smallest decrease in performance scores when compared with the full cache baseline. The robustness of CAKE is evident across all tested cache sizes, with particularly no- table improvements in Multi-Needle Retrieval tasks (Figure IO and Figure This superior per- formance indicates CAKE's enhanced ability to tolerate attention-shifte in long contexts, which can be attributed to eviction strategy that incorporates both sustained importance and attention vari- ability. The proposed indicator and strategy used in CAKE proves to be effective in maintaining performance even under constrained cache conditions.
Results on 1-a.a..-----Inttuuee We further evaluate CAKE's performance Llama3. Instruct using NeedleBench 32K. The results presented in Table Tiilustateee CAKE's comprehensive advantages across all three tasks, and Figure showcases CAKE's outstanding performance multibench tasks, confirming its effectiveness on models with GQA architecture as well.
On the whole, CAKE achieves better overall performance across the two models and various cache sizes, especially for Multi-Needle Retrieval tasks. Yet, given extremely small cache size budgets, Multi-Needle Retrieval tasks performance experiences sharp decline in performance. This may be attributed to the stringent cache limitations are inadequate to support the information demanded by multiple needles. Nevertheless, compared alternative approaches, CAKE demonstrates superior capability in mitigating this issue. Future research should focus on exploring more sophisticated methods for retaining KV pairs, this remains a promising avenue for further improvement.
Table 10: Performance comparison over three subtasks of NeedleBench 32K benchmark on Mistral- v03. The best highlighted in bold, the second best underline.
Singlo-Retrieval
Melti-Retrieval
Multi-Reasoning
Method
Overall
ZH
Overall
ZH
EN
Overall
ZH
EN
Overall
Mistral-T7 IB-Instruct-vvl Mos
Full
Full
91.72
83.01
87.37
90.36
96.95
93.66
55.64
45.89
50.77
78.28
Mistral Instruct-vO... v0.33 Buna 1024/.
H2O
59.20
39.75
49.48
39.50
32.91
36.20
36.95
46.58
41.76
43.18
SnapKV
91.73
90.25
90.99
32.27
81.59
56.93
48.87 87
44.34
46.61
67.46
PyramidKV
91.94
90.69
91.32
25.23
70.55
47,89
47.68
45.32
46.50
64.84
CAKE
89.43
89.89
89.66
51.68
90.32
71.00
52.50
43.57
49.04
71.87
Mistral-7B-
77--Innruuutt-OO..33 Bucl
5121
H2O
60.87
38.44
49.66
4.50
12.68
8.59
36.21
46.15
41.18
34.79
SnapKV
91.21
90.49
90.85
4.95
33.32
19.14
43.88
45.20
44.54
$5.44
PyramidKV
90.06
3.50
17.91
10.70
42.43
45.77
44.10
CAKE
89.97
92.006
91.01
26.18
70.91
48.55
47.60
46.28
46.94
65.05
Mistral 7B Instruct v0.3, Buod 2567.
H2O
66.54
53.61
0.23
4.64
34.61
47.74
41.17
34.53
SnapKV
89.27
89.10
89.18
0.59
2.95
1.77
35.67
47.10
41.39
48.62
PyramidKV
81.51
84.87
83.19
0.18
4.27
2.23
46.90
45.99
CAKE
90.33
89.83
90.08
1.27
7.50
4.39
39.03
47.21
43.12
50.29
25

---

Published as a conference paper at ICLR 2025
Table 11: Performance comparison over three subtasks of NeedleBench 8K benchmark Mistral- 7B- Instruuct-vO... The best is highlighted in bold, the second best is in underline.
Single-Retrieval
Multi-Retrievl
Multi- Maati-Reasoning
Method
Overall
ZH
EN
Overal
ZH
EN
Overall
ZH
EN
Overall
Mistral- 78-Instruct-vO.3 Hona
Full
98.33
84,96
91.65
98.25
96.85
97.55
63.59
57.14
60.37
84.03
Mistral 7B Instruct-vO.33 v0.3, Buot 1024L
H2O
77.93
54.86
66.39
5.15
42.20
23.67
43.55
51.39
50.47
48.80
SnapKV
97.66
89.92
93.79
12.25
50.65
31.45
57.38
57.52
57.45.
64.19
PyramidKV
98.68
87.55
93.11
9.95
37.65
23.80
55.59
56.77
56. 18
61.24
CAKE
90.27
94.24
53.80
87.25
70.53
57.75
57.57
$7.66
76.15
Bus
5126
H2O
72.73
50.71
61.72
0.85
2.75
1.80
39.79
57.5 86
48, 82
39.87
SnapKV
98.11
90.58
94.35
2.25
3.05
52.71
57.78
55.25
$5.23
PyramidKV
97.45
87.87
92.66
1.25
3.00
2.12
49.23
58.52
53.87
$3.86
CAKE
45.25
Mistral 7B Instruct-v0.3. Baya 256L
H2O
74.35
48.40
61.37
0.20
0,20
0.20
37.38
51.87
47.62
38.90
SnapKV
91.89
PyramidKV
92.02
88.24
90.13
0.30
0.55
0.43
39.88
57.02
48.45
50.71
CAKE
97.22
91.93
94.57
1.00
1.22
44.21
58.73
51.47
53.64
Mistral 78 -Instruct---- v0.3 Cache Baseline
Mistral 78 Instruct---tt-- v0. SnapkV Cache Size 1024L
Multi-eeelle AEN 81:59
Widd 21.8
Percerc:5)
Depih Deccentt)
Avcrage
ago Dupth
Token Lenigth
Taken Length
(a) Mistral-7B I--Intrruct--O.3 full cache bascline
(b) SnapK V: cache budget Buni 1024L
Mistral - 78 Instruct v0.3 CAKE Cache Size 1024L
Mistral 78 Inntstssttt--- v0.3 PyramidKV Cache Size 1024L
Multi-Needle Retrieval--iiieeeeee -EN-vva------------ 90.32
Mult--Needie-ee 70.55
56-12
Depth Purcent90)
Deeih Percentta
Average Depth Score
Average Depth Score
Length
(c) PyramidKV: V: cache budget Piotal 1024T.
(d) CAKE: cache budget Brotal 1024T.
Figure 10: Performance comparison of Mistral-7B 1-ii-------ntttttO.... on NeedleBench 32K Multi-Needle Retrieval Task (EN) using different KV eviction methods under B 1024T.
26

---

Published as conference paper at ICLR 2025
Mistral -7B Instruct v0.3 H20 Cache Size 5122
Mistral SnapKV
Mull-Needle-- ReulB--eeede--Reiieil---h-- Setrieeval-E--225 12.68
Retrieval-EN-338
Depth Parcente))
Depth Srow
512L Depth
Length
Token Leingt
(a) H2O: cache budget Brona 512L
SnapKy V: cache budget Brous 512L
Mistral 78 Instruct- PyramidKY Cache Size 512L
Mistral Multi-Needle MeaiiNeeell-Retreeaa.EEEEMM
Multi-Needic Retreeeeeeeevee........ 17.91
65.42
Depth Percerett))
Depth Parcentti
$8.46
13.29
Taken Length
Taken Length
(c) PyramidKV: cache budget Butal 512T
(d) CAKE; cache budget Bual = 512T
Figure 11: Performance comparison of Mistral-7B-Instinnn In--runsttuu---0.33 NeedleBench 32K Multi-Needle
Retrieval Task (EN) using different KV eviction methods under Bunal = 512T.
Table 12: Performance comparison three subtasks of NeedleBench 32K benchmark Llama3.1 1-BB Instruct. The best is highlighted in bold, the second best is in underline.
Single-Retriee
Multi ReitieReetrevall
Multi- Relti--Re
Method
Overall
3aamm3.1 1 Instruct, Buoul Full
Full
100.00
82.99
91.49
99.64
99.77
99.70
72.50
77.70
75.10
89.04
Llama3.1 1-8B Instruct, Busi
1024L
H2O
89.51
66.59
78.05
89.45
91.77
90.61
54.52
71.91
63.21
77.37
SVapKV
99.90
84.21
92.06
86.68
98.23
92.45
71.11
73.29
72,20
86.22
PyramidKV
100.00
85.53
92.76
56.77
88.55
72.66
71.59
76.16
73.87
81.07
CAKE
100.00
84.02
92.01
89.41
98,91
94.16
69.29
73.42
71.35
86.46
Llama3.
Instruct, Busa
5121
H2O
88.14
66.08
77.11
30.05
80.14
$5.09
50.97
73.20
62.08
66.00
SnapK'V
99.90
84,30
92.10
36.09
78.05
$7.07
69.20
73.41
71.31
75.35
PyramidK'V
100.00
86.
93.09
11.00
49,36
30.18
68.72
77.13
72.93
68.17
CAKE
100.00
82.07
91.04
47.36
85.68
66.52
67.67
71.15
69.41
77.19
Llama3.11
Baya
256.
H2O
88.54
69.61
79.08
4.95
6.23
5.59
50.11
72.86
61.49
$1.75
SnapK'V
99.90
82.90
91.4n
4.05
8.45
6.25
66.68
74.48
70.58
59.61
PyramidKV
100.00
85.07
92.53
1.91
6.09
4,00
64.55
76.14
70.34
59.32
CAKE
99.81
81.92
90.86
6.36
27.64
17.00
64.57
68.61
66.59
61.42
27

---

Published as a conference paper at ICLR 2025
Llama3.1 -BB Instruct Full Cache Baseline
Llama3.1 Retrieval
Multi Meeeee--Reeeeeeeeee Ret- 99.77
1924 1922 10Ee
9922
Depth Percentee)
18
Smow
Depth
Length
Langth
(a) Llama3. 1-8B- -n-trutttct full cache baseline
SnapK V: cache budget Brous 512L
Llama3.1 Instruct PyramidKV Cache
Llama3.1
Multi Needle Retneval/EN EN-33K 49.36
Nueti-Needle Ret.rleeeeeeeeeeeeeee
00.45
Depth Parcentti
5.
29.59
Tiken Length
Taken Length
(c) PyramidKV: cache budget Butal 512T
(d) CAKE; cache budget Buaa = 512T
Figure 12: Performance comparison of Llama3.1 1-taa...----BM- -Bstuutt NeedleBench 32K Multi-Needle Retrieval Task (EN) using different KV eviction methods under Bunal = 512T.
Table 13: Prompt prefilling and token decoding latency comparison.
Prefill-Deodding
Method
Overall Generation Time
Prompe Prefill Time (8)
Decoding Time (6)
Decoding
Length
Time (ms))
SnapKV
33.74
0.71
33.03
32.25
PyramidKV
32.29
0.78
31.50
30.77
CAKE
31.77
0.81
30.96
30.24
Cache
SnapKV
34.78
1.66
33.12
32.34
PyramiaVV
33.02
1.73
31.29
30.56
15360+1024
CAKE
32.96
1.77
31.19
30.46
Full Cache
50.71
1.66
49.06
47.91
SnapKV
PynmidKV
35.78
4.20
31.58
30.84
31744+1024
CAKE
36.42
32.12
31.37
Full Cache
86.42
4.14
82.18
80.35
SnapKV
44.29
32,96
PynmiiKV
43.11
11.48
31.63
30,89
64512+1024
CAKE
43.10
11.47
31.63
Full Cache
176.95
11.42
165.53
161.65
SnapKV
68.98
35.38
33.60
32.81
PyraiiKVV
67.61
35.51
32.10
31.35
130048+1024
66.83
Full Cache
364.15
35.86
328.28
320.59
28

---

Published as a conference paper at ICLR 2025
H ADDITIONAL EXPERIMENTS AND ANALYSIS ON EFFICIENCY
In this section, we present a detailed time breakdown during both prompt prefilling and token decod- ing to better assess CAKE's effectiveness across different inference stages compared to other KV eviction methods, including SVapKV and PyramidKV with various allocation strategies. All meth- ods are implemented as extensions of FlashAttention-2 (Dao 2023) using Mistral-7B wan.------ntru--O.... and their performance was benchmarked against a full cache baseline utilizing FlashAttention. To measure computational time more precisely throughout the entire text generation phase, we have separately evaluated the prompt prefilling time and the duration of autoregressive decoding. By fix- ing the generation token count at 1024, the experiment simulates typical long-context processing scenarios in LLMs, which are characterized by lengthy inputs and concise outputs.
As shown in T3 while dramatically lowering latency in the decoding phase, CAKE, SnapKV, and P achieve comparable prompt prefilling times. For SnapKV and PyramidKV with fixed allocation strategies, the prompt prefilling stage requires waiting for KV caching and per- forming one eviction operation per layer, resulting in L eviction operations. However, compared to the computation-intensiin
prefill stage, the eviction operation is relatively inexpensive and can be considered negligible. Notably, while CAKE utilizes ment, it achieves similar timing to SnapKV and PyramidKV with fixed allocation strategies. This is attributed to the parallelizable execution of KV cache eviction between layers (Algorithm [ lines 04),), which allows the dynamic update process to incorporated into single eviction operation. Therefore, the total execution time ultimately equates to that of performing L eviction operations.
seteerence-guidee cascading cache manage-
I ADDITIONAL ABLATION STUDIES
I.1 ABLATION STUDY ON METHOD COMPONENTS
In this section, we conduct ablation studies on LongBench to analyze the contribution of each com- ponent in our method. These studies highlight the importance of our design choices, focusing on two key aspects: (1) the preference-prioritized adaptive allocation strategy, which combines atten- tion dispersion (H) for spatial influence and attention shift (V) for temporal dynamics, and (2) the attention-shiff tolerant eviction indicator, which integrates mean attention score (Mean) for identi- fying sustained important tokens and attention variance (Var) to capture dynamic attention changes.
As shown in Table 14 incorporating each additional component leads to a noticeable performance improvement, demonstrating the synergistic effect of our method. The sefference-prioritizee adap- tive allocation strategy improves performance from 28. 14 to 28.82 when both 7. and V are included, underscoring the importance of capturing spatial and temporal attention dynamics. This highlights how our approach goes beyond traditional uniform strategies by tailoring cache allocation based on complex layer-specific attention patters.. Furthermore, incorporating attention variance (Var) in the eviction indicator significantly enhances performance by accounting for attention fluctuations, further improving average performance to 29.29. This confirms that addressing both spatial distri- bution and temporal evolution of attention is crucial for effective KV cache management, especially in long-context scenarios where dynamic attention patterns heavily influence performance. Our comprehensive approach successfully balances these factors, yielding superior results compared to methods that overlook such nuanced attention characteristics.
Table 14: Ablation study on different components of our method. The experiments are conducted on Llama2-7b ---Chat with total cache budget B 1287.
Allocation strategy
Eviction indicator
Avg.
H.
V
Mean
Var
28.07
28.14
28.82
29.29
29

---

Published as a conference paper at ICLR 2025
1.2 INFLUENCE OF TEMPERATURE PARAMETER
In this section, examine the impact of temperature parameters and T2. which modulate the influence attention dispersion and shift in our preference sefference-prioitizzedd adaptive allocation strategy. We evaluated various settings I 1-8B-Inttutt and a cache budget of Buaal = 1287. Our approach is compared against SnapKV, a state-of-the-ar KV eviction method employing uniform cache allocation, which serves our baseline. As shown in Figure [3] our method consistently outperforms the baseline across different configurations for Biirral-BB-Inttuctt---
models, using
achieves peak performance of 37.33 (baseline: 35.81), while LLama3.1 --8--Instruct reached 43.39 (baseline: 42.53), demonstrating the robustness and ef- fectiveness of our strategy. While our approach is inherently robust, we found that simple, low-cost adjustments to and can further enhance performance. These adjustments allow for better cap- ture of model-specific attention patterns, optimizing cache allocation without expensive retraining extensive hyperparameter searches. This characteristic is particularly valuable in industrial de- ployments where specific models need to be optimized under constrained memory budgets, offering practical solution for resource-efficient inference in production environments.
both models.
(37 35
37.11
Aug
Avg
36.15
42.63
Boceline: 42.53
Baseline:
0.2
0.2
1.7
0.9
49
0.4
(a) 1-8--.33--B--nstrrue
(b) Mistral-8B-IIntruuc
Figure 13: Performance comparison of different тı and 72 configurations for Mistral-7B-Inss Instruct-- v0.3 and 1--ma3.--8-IInstuue models. The plots show the impact of varying тj and 72 model performance, with the baseline (SnapKV) and best scores indicated.
J ADDITIONAL DETAILED ANALYSIS AND VISUALIZATION OF ATTENTION DYNAMICS
In this section, we provide a more intuitive and comprehensive analysis of the differences in attention patterns by presenting detailed visualizations of attention weights across various dimensions. To illustrate the variations between different layers, we visualize the model's attention weights across all layers by aggregating the attention weights of each attention head. We also highlight the differences in attention between different models by conducting experiments on both Llama3 and Mistral. To demonstrate the attention differences across various input contexts, we visualize attention weights for different task types on Llama3, including single-document QA, summarization, and code tasks. Additionally, we visualize attention weights for the same task types but with varying contexts on Mistral.
Attention Differences across Layers. As shown Figure 14] and Figure attention patterns ex- hibit significant differences across layers, as mentioned in the main text. Regardless of the model or input context, we can observe varying degrees of attention dispersion and attention shift across different layers. For example, ure or lower dispersion (e.g., layer layer 24 in Figure (c)) or lower shift (e.g., layer 35 in T5 These observations demon-
layers with higher attention dispersion (e.g., layer 0 in Fig-
(())), as well as layers with higher shift (e.g.,
30

---

Published as a conference paper at ICLR 2025
strate the importance of carefully considering the differences in attention mechanisms and allocating appropriate cache resources accordingly.
Attention Differences across Models. As shown in Figure 14 and Figure 15| these models have notable differences attention patterns. For instance, Mistral often exhibits higher attention dis- persion the beginning and end of its layer stack. In contrast, Llama3 frequently displays high attention dispersion primarily in the early layers. These distinct patterns underscore the importance of model-specific considerations when designing caching strategies, as different architectures may have varying attention characteristics across their layer stacks.
Attention Differences across Contexts. As evident in Vigure[44a))-4 [))))cc), there striking differ- ences in attention patterns across different task types. Even within the same task type, while there may some similarities attention dispersion, attention shift patterns show marked differences. This can lead to different layer preferences for KV cache (comparing layers 10-16 16 in Figure with layers 10-16 in Fig)) Therefore, whether the task types are the same or different, attention patterns are consistently dynamic and varied.
Given the dynamic nature of attention across layers, models, and even contexts, adopting uniform allocation strategy, while safe, does not effectively utilize memory. On the other hand, using fixed pattern allocation strategy fails generalize effectively across multiple models and task scenarios. In contrast, our CAKE thoroughly considers layer preferences for KV cache, formulating reason- able cache size allocation scheme from a global perspective. It takes into account the substantial impact of attention dynamics, allowing it dynamically adapt to different models and contexts. This approach ensures that CAKE can effectively respond to the varied and changing attention pat- terns observed across different scenarios.
31

---

Published as a conference paper at ICLR 2025
Layer 0
Layer 1
Layer 2
Layer 3
Layer
Layer
Layer
Layer 7
0.05
Layer 8
Layer 9
Layer 10
Layer 11
Layer 12
Layer 13
Layer 14
Layer 15
0.04
Layer 16
Layer
Layer 18
Layer 19
Layer 20
Layer 21
Layer 22
Layer 23
9.02
0.01
Layer 24
Layer 25
Layer 26
Layer 27
Layer 28
Layer 29
Layer 30
Layer 31
0,00
(a) Visualization of attention weights Qasper.
Leye
Love:
lsse
kaye
Lowr
Lise/22
Mpp 21
14p
18pa
Laver
Lowr 24
Lopr 24
Use
143
1250 11
(b) Visualization of attention weights on Mutinews.
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
0.05
Layer 8
Layer 9
Layer 10
Layer
Layer
Layer
Layer 14
Layer 15
0.04
Layer 16
Laver 17
Layer 18
Layer
Layer
Layer
Layer 22
Layer 23
0..02
0.91
Layer 24
Layer 25
Layer 26
Layer 27
Layer
Layer
Layer 30
Layer 31
(c) Visualization attention weights Lec.
Figure 14: Visualization results of attention weights for Llama3, 1-8B- -B-Instruct across datasets from different tasks.
32

---

Published as conference paper at ICLR 2025
Layer 0
Layer
Layer
Layer 3
Layer
Layer
Layer 6
Layer
Layer
0.05
Layer 9
Layer
Layer 11
Layer 12
Layer
Layer 14
Layer 15
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
-0.02
0.00
27
Layer
30
33
0.00
(a) Visualization of attention weights on NarrativeQA.
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer 9
Layer 10
Layer
Layer 12
Layer
Layer
Layer 15
Layer 16
Layer 17
0.03
Layer 18
Layer
Layer 20
Layer 21
Layer
Layer
Layer 24
Layer
Layer 26
0.02
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
(b) Visualization of attention weights Qasper.
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
-0.04
0.00
Layer 18
Layer
Layer 20
Layer 21
Layer
Layer 23
Layer 24
Layer 25
Layer 26
Layer 27
Layer 28
Layer 29
Layer 30
Layer 31
Layer
Layer 33
Layer 34
Layer
0.00
(c) Visualization of attention weights MultiFicldQA-en.
Figure 15: Visualization results of attention weights for Mistral-7B. 1 across datasets
from Single dinglendo QA task.
33


Reviews:
[{"review_id": "AoCGZ1SwWK", "reviewer": "Reviewer_rCQd", "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.", "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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "(1)\tThis 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.\n\n(2)\tThe 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.\n\n(3)\tThe 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.\n\n(4)\tThe extensive experiments on several open-source LLM benchmarks truly validate the effectiveness of the proposed algorithms.\n\n(5)\tOverall, the proposed algorithm CAKE model is novel, practical and efficient. The corresponding experimental results are extensive and sound.", "weaknesses": "(1)\tThe 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.\n\n(2)\tIn 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”.\n\n(3)\tEquation (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.\n\n(4)\tI 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.\n\n(5)\tIn 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "N1kTmIMX1W", "reviewer": "Reviewer_LJXW", "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.", "questions": "Please see weakness aspects above.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\nThe 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.\n\nEmpirical 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.\n\nThe empirical improvements over existing baselines are relatively modest, which could suggest limited practical advantages in some cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9JsK0xnV2P", "reviewer": "Reviewer_BrWg", "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.", "questions": "Q1. I'm just curious how this technique could scale to larger models (no need to verify it empirically)\n\nQ2. Could this technique potentially help transformers stay competitive against RNN-based models like Mamba?", "limitations": "", "rating": 10, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "S1. Identified and visualized the attention dynamics in spatial and temporal axis\n\nS2. Strong empirical analysis on multiple datasets\n\nS3. Actual improvement in inference latency\n\nS4. Excellent presentation (great paper flow and beautiful figures)", "weaknesses": "I do not identify any major weaknesses of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TKeGdKwWkp", "reviewer": "Reviewer_zgXb", "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.", "questions": "Please refer to the weakness", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n- The team has not open-sourced their code, which could raise concerns about the reproducibility of their work.\n\n- 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):\n\n[1] Hooper, Coleman, et al., \"Kvquant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization.\" arXiv preprint arXiv:2401.18079 (2024).\n\n[2] Liu, Zirui, et al., \"KIVI: Plug-and-Play 2bit KV Cache Quantization with Streaming Asymmetric Quantization.\" (2024).\n\n[3] GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLM.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)