Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference
Authors: Yuan Feng
Junlin Lv
Yukun Cao
Xike Xie
S Kevin Zhou
Venue: NeurIPS 2025 poster ()
Keywords: Efficient AI
Large Language Model
LLM Inference
KV Cache
Primary area: deep_learning
TL;DR: Optimize KV cache eviction by adaptively allocating budgets across different attention heads for efficient LLM inference

Abstract: Large Language Models have excelled in various domains but face efficiency challenges due to the growing Key-Value (KV) cache required for long-sequence inference. Recent efforts aim to reduce KV cache size by evicting vast non-critical cache elements during runtime while preserving generation quality. However, these methods typically allocate compression budgets uniformly across all attention heads, ignoring the unique attention patterns of each head. In this paper, we establish a theoretical loss upper bound between pre- and post-eviction attention output, explaining the optimization target of prior cache eviction methods, while guiding the optimization of adaptive budget allocation. Base on this, we propose {\it Ada-KV}, the first head-wise adaptive budget allocation strategy. It offers plug-and-play benefits, enabling seamless integration with prior cache eviction methods. Extensive evaluations on 13 datasets from Ruler and 16 datasets from LongBench, all conducted under both question-aware and question-agnostic scenarios, demonstrate substantial quality improvements over existing methods. Our code is available at https://github.com/FFY0/AdaKV.

Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference
Yuan Fengl.3t 133 Junlin Lyl.3t Yukun Cao Xike Xie S. Kevin Zhou
School of Computer Science, University of Science and Technology China (USTC) Schooo of Biomedical Engineering. USTC
3 Data Darkness Lab, MIRACLE Center, Suzhou Institute for Advanced Research
Abstract
Large Language Models have excelled in various domains but face efficiency challenges due the growing Key- Value (KV) cache required for long-sequence inference. Recent efforts aim to reduce KV cache size by evicting vast non-critical cache elements during runtime while preserving generation quality. However, these methods typically allocate compression budgets uniformly across all attention heads, ignoring the unique attention patterns of each head. In this paper, we estab- lish a theoretical loss upper bound between pre- and post-eviction attention output, explaining the optimization target of prior cache eviction methods, while guiding the optimization of adaptive budget allocation. Base on this, propose Ada-KV, the first head-wise adaptive budget allocation strategy. offers plug-and-plas benefits, enabling seamless integration with prior cache eviction methods. Exten- sive evaluations on datasets from Ruler and 16 datasets from LongBench, all scenarios, demonstrate
conducted under both question-aware and substantial quality improvements over existing methods. Our code is available at httpp:://githbb http://tt...ccoo////aaaaaaaa
Introduction
Autoregressive Large Language Models (LLMs) have achieved significant success and are widely utilized across diverse natural language processing applications, including dialogue systems M. document summarization 22, and code generation [3] The widespread deployments of LLMs have propelled the development of their capacities to process extended sequences. For instance, GPT supports sequences up to 128K [4). Claude3 up to 200K (5), and Gemini-Pro-L.5 [6] up 2M tokens. However, this growth in token length introduces significant challenges, particularly the rapid expansion of cache size during inference. For 8B L.M,, handling single sequence of 2M tokens can require up to 256GB of cache, severely impacting both GPU memory efficiency and computational runtime efficiency.
In particular, the inference process, for each multi-head self-attention layer, consists two phases: prefilling and decoding. In the prefilling phase, LLMs compute and store all Value (KV) cache elements for the tokens
input prompt. In the decoding phase, model autoregressively uses the most recently generated token to retrieve information from the stored cache, producing the next output iteratively. The large KV cache size introduces two important efficiency challenges: First, severely affects GPU memory efficiency, making it increasingly difficult scale with longer inputs. Second, during decoding, increased I/O latency duo to cache access results in substantial delays, degrading runtime efficiency.
Equal Contribution, "Corresponding Author: Xike
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Aggregated Altention Weight
Uniform Ruddget Allncatin
Adaptive Budger Allocation with Sloe
0.62
0.75
1.00
Tayer
1830 moooX
Head Index/
Budeet
Bedeet
5/4
32
0%
1006
20%
32%
40%
5a%6
0%
1006
20%
30%
40%
50%
Aggregated Top Ratio
Aggregatedd
Top
Ratio
(a) Varied Attention Concentration Across Heads.
(b) From Uniform to Adaptive Budget Allocation
Figure 1: Adaptive budget allocation accommodates varying attention concentration across heads. Left: Analysis using Llama-3. 1-8B-Instruct shows most heads retain nearly all attention weights with a small cache (e.g. top 5%), while dispersed heads require larger cache proportions. Right: Adaptive allocation, which shifts budgets from sparse to dispersed heads, increases the aggregated retained attention weights (from 2.26 to 2.48) and reduces eviction loss compared to uniform allocation.
To address the challenges posed by large KV cache sizes, various cache eviction methods have been developed 7 8 10 These methods constrain the cache size to a predefined budget by retaining only a subset of elements in each attention head and evicting the rest. They reduces memory usage and accelerate decoding, facilitating efficient long-sequence inference. Most of these methods rely on a Top-k selection based on attention weights, to identify and prioritize critical cache elements, deciding which to retain and which to evict.
However, existing methods uniformly allocate cache budgets across attention failing to account for the unique characteristics of each head. As shown in Figure [a] attention heads exhibit diverse concentration patterns: some focus narrowly (or sparsely) on a small portion of the cache, while others distribute their attention more broadly | This uniform allocation results in inffficiencies-eitheee wasting cache budgets on heads with sparse concentration or incurring significant eviction losses in heads with dispersed distribution. Such imbalances degrade the trade-off between overall budget utilization and the quality of post-eviction generation. To address this issuc, we analyze the impact of existing Top-k eviction methods on attention outputs and propose the first adaptive budget allocation strategy, called Ada-KV, as illustrated in Figure [6] This strategy dynamically reallocates budgets from heads with sparse concentration to those with more dispersed attention, improving the quality of post-eviction generation.
Our study begins by revisiting Top-k eviction methods, demonstrating that, under specific budget allocation, these methods are equivalent to minimizing an upper bound of eviction loss between pre- and post-eviction attention o Based on this, we propose Ada-KV, simple yet effective adaptive allocation strategy designed to integrate and enhance existing Top-k eviction methods in a plug-an and- -nd--aadyyaay play manner. Guided by minimizing the theoretical upper bound of eviction loss, Ada-KV adapts to the varying concentration patterns across attention heads, significantly reducing practical eviction loss.
By integrating the Ada-KV into two state-of-the-an (SOTA) methods, SnapKV and Pyramid 9 10]. we have developed two integration cases: Ada-SnapKV and Ada-Pyramid" These methods are evaluated using two comprehensive benchmarks, Ruler and LongBench, which include 13 and 16 datasets, respectively. Experimental results demonstrate that A could effectively improves performance across various budget settings under the widely studied question-aware compression scenario [TI 10 { Furthermore, in the more challenging question-agnoste [12 scenario-where compression is applicd without leveraging guestion-specife information, an important scenario that has been largely overlooked- Ada-KV demonstrates even greater advantages. The main contributions are summarized as follows.
Adaptive Budget Allocation. We identify a critical limitation in current KV cache eviction methods: their uniform budget allocation overlooks the unique attention patterns of individ-
1 Mooe visualizations can found in Appendix A.11
For simplicity, we use L1 distance to quantity eviction loss; extensions to 4p norms are possible but orthogonal to this work.
'The plug-and-plas nature of Ada-KV enables broad applicability to existing methods and future develop- ments, extending well beyond the two cases presented in this paper. Section provides further examples of its adoption in subsequent work.

---

ual heads. To address this, we propose Ada-KV, the first adaptive budget allocation strategy, which enhances budget utilization across individual heads, leading to more efficient cache eviction methods.
Theoretical Insights. We establish a theoretical framework for cache eviction by defining eviction loss and deriving its upper bound. This framework not only explains the optimization target of prior methods but also guides the design of Ada-KV, enabling principled and adaptive budget allocation.
Empirical Advances. With efficient CUDA kernel implementations, Ada-KV offers plug- and-play compatibility, enabling seamless adoption and performance enhancements in existing methods. We evaluate V by integrating it into SOTA cache eviction methods, and demonstrate significant improvements across 29 datasets from Ruler and LongBench, all under both question-awar aware and question-agnostic scenarios.
2 Related Works
2.1 Cache Eviction Methods
In the long-sequence inference, the vast scale of the KV cache elements leads to memory-bound situation, causing significant memory burden and I/O latency [33. Numerous studies have sought
to mitigate by reducing the cache size, notably through the eviction of non-critical cache elements") These methods are primarily divided into two categories: sliding window eviction and Top-k eviction methods. The sliding window eviction methods 14 15 16], exemplified by StreamingI.LM (16), simply retain several initial cache elements and those within a sliding window, while evicting others. However, the undiscriminating sliding eviction of cache elements results in significant reduction in generation quality. In contrast, Top-k eviction methods [7] [7 18
S 9 TO [T] identifyy and retain a
selected set of is critical cache elements based on attention weights, for enhancing the post-eviction generation quality. The adaptive budget allocation presented this paper, tailored for Top-k Eviction Methods, enhances post-eviction generation quality within the same overall budget.
In the study of Top-k eviction methods, early work like FastGen 7 searches and combines multiple strategies, such as maintaining caches of special elements, punctuation elements, recent elements, and Top-k selected elements, based on the characteristics of attention heads. H2O, as the representation (8] I7. develops Top-k based eviction scheme that leverages the query states of all tokens to identify critical cache elements. However, under unidirectional mask in LLM computations, aggregating
attention weights across all query states often causes recent KV cache elements to be mistakenly evicted, degrading the quality of subsequent generations. Recent works, such as SnapKV V [TT] address this issue by using query states within an observation window to identify critical elements, thereby achieving SOTA performance. Subsequent methods, such as Pyramid (9] [0]. further introduce budget allocation across different layers. However, existing Top-k eviction methods typically assign the overall budget uniformly across different heads. In contrast, Ada-KV, which has demonstrated superior theoretical and empirical results through adaptive budget allocation, provides novel strategy to optimizing these methods.
2.2 Sparse Attention Methods
Sparse attention methods are conceptually related to KV cache eviction, but differ fundamentally in their approach 19 20 21 23 24). The key difference is that KV cache eviction retains only subset of the KV cache, while sparse attention methods keep all entries but selectively utilize only a critical subset during computation. Consequently, sparse attention methods do not reduce the memory footprint of the KV cache, thus often necessitating cache offloading to CPU memory 22 (3) or disk storage (24). Additionally, the two technique lines are, in fact, orthogonal. Future research could explore first employing KV cache eviction to compress the cache to a certain proportion and then applying sparse attention for further acceleration. Ada-KV presented in this paper leverages head-wise adaptive allocation to enhance cache eviction methods, and this approach can similarly be applied in future work to enhance dynamic sparse attention methods.
For additional related works, see Appendix A.I
3

---

3 Methodology
We begin by providing formal description of multi-head self-attention layer (Section B.1p. Building on this, we theoretically revisit the foundational principles of existing Top-k eviction methods by introducing an L1 eviction loss metric (Section 3.2). Inspired by theoretical findings, we propose simple yet effective strategy for adaptive budget allocation, which is proven to outperform traditional uniform budget allocation (Section B.3). We further demonstrate its compatibility with existing Top-k eviction methods (Section 3.4) and present an efficient implementation (Section 3.5
3.1 Preliminaries
LLMs operate through an autoregressive generation process, where each step relies on the last token to predict the next one. Let X E gnxd represent the embedding matrix containing all tokens in the sequence, and let x 1 be the embedding of the last token used as input at the current time step. Using the notation system from [25] with h attention heads, we provide a formal description of one multi-head self-attention layer. For each head i E [1, h]. the transformation matrices W.O, W.". W.Y- E idxda map token embeddings to their respective Query, Key, and Value states, while the final output matrix 1 E Rd, xd transforms the intermediate result to the output hidden states. At each time step, the previous KV cache elements of head are initialized as: K; XW. Vi2 XW.. Then, the embedding of input token x is mapped to its respective Query, Key, and Value states for each head, and the previous KV cache is updated accordingly:
qi W hu rW: 1/4 zW. and K, Cat(K; &.), V. Cat(Vi m]
(1)
Finally, the output y E R 1 is computed using the attention weights A; R as follows
Σ
A.V., where A, = softmax K.)
(2)
SEDA
3.2 Theoretical Foundation: Revisiting Top-k Methods with Bounded Eviction Loss
Top-k eviction methods typically presuppose the stability of critical cache elements during future generation process, to facilitate cache eviction [8] [T [10] SOTA methods [[] [0] commonly leverage query states of series of tokens within an observation window calculate observed attention weights with past KV cache elements. These weights, combined with Top-k selections, approximate the identification of cache elements critical in subsequent generations.
For ease of presentation, we assume a window size of 1, implying the eviction procedure relies solely on the attention weights A between the past cache elements and the query state of the last token for
the detection of critical cache elements. A set of indicator variables {I: E R Rixa 16 represent the eviction decision, with allocated budgets {B,} for all heads {i E [1, h]]. where LiE[I,N] B, is the overall budget for one attention layer:
Eviction Decision: I [1
Retain K- and VtJ Evict Ki and V?
where IÍ indicates whether the jth E [1, KV cache element in K,, V22 E Rxd is evicted for head i. Thus, only budget size B; of cache elements is retained for head d: ΣjE(1m) I B;. Then, the post-eviction output ú of multi-head self-attention mechanism is:
Σ Á.V.W.D, where A = softmax C (1 (1-1)) + kgKKK) FE[LA]]
(4)
and @ denotes element-wise multiplication.
The degradation in generation quality after cache eviction stems from changes in the attention output. We quantify the eviction loss as the L1 distance between the pre- and post-eviction outputs of the self-attention mechanisms:
L, Eviction Loss = lly 0ll2
(5)
"The scaling factor Vdn omitted for simplification.
'Given that the first dimension of Iz is 1, I is used to simplify the notation for I,( Similarly, AQ follows the same notation.

---

Utilizing the row norm of the matrix, we derive an upper bound e for the L1 Eviction Loss in Theorem For detailed proof, please refer to Appendix (A.8]
3.1
Theorem 3.1. The Lı eviction loss can be bounded by E:
L1 Eviction Loss S c =2C 2C
E T'A' SE
(6)
where C = Max [IV,WW ] a constant number, representing the max row norm.
Essentially, existing cache eviction methods are equivalent to minimizing the upper bound e, as defined in Theorem33.] One such method, Top-k cache eviction method 8 [I 9 [0] stands out as prominent approach, designed to selectively retain the most critical cache elements by prioritizing those with the highest attention weights. Specifically, the Top-k eviction decision is formulated as
[1 if A E (Ap, k B.). [0 Evict Kj and V."
Top-k Eviction Decision T
This approach ensures that each head retains only the top B, critical cache elements based on attention weights A/, evicting the rest. In Theorem 3.2 we prove that the Top-k eviction decision {I,} } achieves the tightest upper bound e", establishing its effectiveness, under given budget allocation [..). For a detailed proof, please refer to Appendix A.9
Theorem 3.2. The Top-k cache eviction resulting in e
} minimizes the upper bound of L1 eviction loss,
A
Top-k eviction decision [I,) arg mine and mine c 2hC -C (2.)
Σ SE(1,6]
Σ JELL
A (7)
{2;}
Theorems 3.1 and 3.2 establish the theoretical basis for Top-k eviction methods under any given budget allocation B J. In the sequel, we show how an adaptive strategy can further minimize c*,
outperforming uniform budget allocation (i.e., (B, = B/h)).
3.3 Optimizing Top-k Methods with Adaptive Budget Allocation
Algorithm Ada-KV: Adaptive Budget Allocation
Input: Total budget: B; Attention weights for head {Ai}
Output: Allocated budgets {B²}
1: Concatenate attention weights across heads A = Cat ({A+)
2: Select top B weights from Top-k(A, k B)
3: Count number of selected weights for each head i: {fi}
4: Set the allocated budgets as { {B; = fi}
Return allocated budgets {B,}
Here, we propose the first adaptive budget allocation strategy for Top-k eviction methods in Algorithm I. designed to further minimize the upper bound e* in Theorem 3.2 It begins by identifying the B largest attention weights across all heads within a single layer. Based on the frequency of selection in each head, the strategy dynamically determines the budget allocation { Thereby, attention sparse heads, where most attention weights are concentrated on only a few cache elements, are assigned a smaller budget. The saved budget is reallocated to attention-disperse heads with more widespread concentration patterns. Thus it effectively adapts to the distinct attention patterns inherent to each head. Below, we further demonstrate the superiority of this strategy theoretically and empirically in subsequent sections.
Theoretical Advantages of Adaptive Allocation. Given the adaptive budget allocation { J. the upper bound associated with the Top-k eviction decision {I;) expressed as e e**
A (oA,
2hC 2C
Σ
Σ
A
(8)
(E[L,M]
]E
5

---

Relathe Evictio Lose Io It-Sh
0.3
intter Sucpee Alccattion
Algorithm 2
SSSapkV/Ada-Pyraamaaa in One Layer
Adscttive Dudger Alccaton
Input: total budget B, tokens in observation window X xwin E R Ruinndd cache in observation window ( VV,un J, cache out- side observation window {K,, Vi}
02
0.-
Output: retained cache Ki, 1/3
50
100
150
200
I: for + 1 to do
Samples
2: Q = = W.
B 20%
cache size
Recatve Evittin L--th
A. sat so bmaaw Q. K.)
Usifor Budget Alocation
4: A; = Aamarpooligggdii 1) meman(din = 0)
0:15
5: end for
0.10
6: B B winsize x h
Derive budget allocation {Bī) using Algorithm 1 B, (A.)) a) x (B7h)
0.05
8: Safeguard {B; = a x {B:
+ }+
9: Determine the Top-k eviction decision (I based on {
200
10: Select K., k.) from {K,, V/3) according to {II}
40% cache size
11: { kr, 1.}- = Car( K.., 1. [K"
Figure 2: Adaptive allocation yields
lower practical eviction losses (Llama- 3.1-8B-Instrus on Qasper Dataset).
Return retained cache K,,
Theorem 3.3
demonstrates that the adaptive budget allocation achieves the minimum upper bound c" compared to any other allocation strategy, including the commonly used uniform allocation. A detailed proof is provided in Appendix AA.0
Theorem 3.3.
The adaptive budget allocation (B)), derived from Algorithm Algorithn Jachivvss the minimal
upper bound e** for loss associated with Top-k eviction strategies: c* (B.) c*
Although a lower upper bound does not strictly guarantee reduced eviction loss, is generally key objective in algorithm design, as optimizing tighter bound typically leads to better results.
Empirical Evidence for Adaptive Allocation. Empirically, the adaptive budget allocation described in Algorithm|| [ assigns larger budgets to dispersed heads and conservatively adjust budgets for sparse heads. This approach aligns well with the significant variation in attention concentration across heads, as shown in Figure Ta By adjusting the head-wise budget, this strategy effectively reduces practical eviction loss. Figure provides detailed visualization, showing that adaptive budget allocation consistently reduces practical eviction loss across most samples under the same cache budget. These results demonstrate the effectiveness of the proposed method in real-worke scenarios.
3.4 Integration into Existing Cache Eviction Methods
We demonstrate the Plug-and-Plas compatibility of Ada-KV strategy with two SOTA methods, SnapKV and Pyramid, by integrating into their existing Top-k eviction frameworks to create enhanced versions: Ada-SnapKV and Ada-Pyramid. Both SnapKV and Pyramid utilize tokens Xwin E Ruuinsszed from a recent observation window (typically size 32) identify and evict the less crucial elements in past KV cache. SnapKV excels under larger budget scenarios, while Pyramid optimized for constrained budget conditions. This is because Pyramid uses a pyramidal budget distribution across attention layers via pre-set hyperparameters, prioritizing shallower layers. Thus, for a given layer with a total budget of B. their eviction algorithms are identical. However, both methods allocate budgets evenly across all heads within one layer.
Incorporating our adaptive allocation, as outlined in Algorithm 2 we modify these methods to better manage budget allocation at the head level. This integration occurs prior to the eviction process in each layer, where our strategy adaptively adjusts budget allocations based on the observed attention weights among heads, as shown in Line 7. Overall, they first calculate the observed attention weights A, of past cache elements using the query states within the observation window. max pooling layer processes these weights to preserve essential information [II] followed by Top-k selection of past cache elements outside the observation window. These selected elements, along with others within the observation window, are retained, while the rest are evicted. Moreover, we introduce safeguard hypper-paramete, a (defaulted to 0.2), to prevent the allocation of excessively small budgets to highly sparse heads, thereby enhancing fault tolerance for presupposed critical stability [T8]

---

3.5 Implementation of Computation under Adaptive Budget Allocation
Variable-length Attention with Variable-sized Cache Elements. Adaptive allocation improves cache eviction quality but introduces challenges for efficient computation due to variable-sized cache elements across attention heads. We found that the variable-length FlashAttention technique [26, 27], widely adopted in many inference frameworks for continuous batching [28], could support efficient computation under adaptive allocation. To further enable this technique, we implement a flattened cache storage layout that concatenates the caches of all attention heads within layer into a single tensor structure. This layout, combined with a custom CUDA kernel, enables efficient cache update operations. As demonstrated in Section 4.4 these components work in synergy to maintain computational efficiency under adaptive allocation, comparable to that of conventional FlashAttention.
Compatibility with Group Query Attention. Existing SOTA LLMs like Llama [29] and Mis- tral [30, have employed Group Query Attention (GQA) [31] technique to reduce KV cache sizes. However, existing cache eviction methods, such SnapKV and Pyramid, lack GQA compatibility, redundantly replicating grouped KV caches across heads. We implement a simple GQA-compatible mechanism that uses the mean attention weight within each group as the selection criterion, elimi- nating redundancy. This enables SnapK V, Pyramid, and their adaptive Sla-Pyramid-- to achieve significant cache size reductions, such 4x reduction in .lama-3.1-88
Saantts-Ad-Snaappp V and
Experiments
4.1
Settings
Base Models. We employ open-source base models: Llama-3. 1-8---.-r-----sttccc [29] and Mistral-7B- instruct-v0.2 [30] These models, widely adopted due to moderate parameter sizes and impressive performance on long-sequence tasks, both leverage the GQA 3] technique, which reduces the KV cache size to one-quarter of its original.
Datasets. We conduct evaluations using two popular benchmarks: Ruler [32 and LongBench [33] Ruler includes 13 unn-sequence tasks, primarily variations of the Meedle-in-a-Haysstal test, offering range of difficulty levels for comprehensive model assessment. According the official evalua-
tion [32], the maximum effective length for Mistral-7B Balll-----istuuce
and Llama-3.1-83 1aa....--B-Instrure with
the full KV Cache is 16K. Therefore, select the 16K Ruler Benchmark the primary benchmark for evaluation. Additionally, we also conduct evaluations using LongBench, including 16 datasets covering multiple task domains. Detailed dataset information is provided in Appendix A.5 and A.7
Baselines. We select the SnapKV[II] and Pyramid 9 [0] as the primary baselines, given that they SOTA methods and foundational bases for our A and Ada-Pyramid methods. All these methods implement GQA compatibility as stated in Section 3.5 reducing cache size to a quarter of previous naive implementations. Additionally, StreamingLLLM 16 is also included representative of Sliding Window Eviction Methods for reference.
Parameters. Cache eviction methods can be applied whenever KV cache compression is required. Following standard practices prior studies [T] 9 10), we perform cache eviction after the prefilling phase of each layer for consistent comparison. In all experiments, the hyperparameter for adaptive budget allocation fixed 0.2. Both Ada-SnapKV and Pyramid, as well as SnapKV and Pyramid, follow the configuration settings outlined in [1I], including an observation window size of 32 and a max pooling kernel size 7.
4.2 Ruler Benchmark
We evaluate quality scores for each method using budgets set to 20%, 40%, 60%, and 80% of the original cache size. Beyond the commonly studied question-aware compression scenario- questions are known in advance-as we also consider the more challenging question-agnostic set- ting [12), in which questions scenarios, such as prompt caching [34, 335, where numerous question-independen prefix prompts need compression, multi-turn dialogue [I]. where compression occurs without future questions.
revealed only after compression. This mimics more challenging
Overall Results. Figure 3 presents the average scores across 13 datasets in the Ruler Benchmark. Overall, existing Top-k eviction methods exhibit significant performance drops in the question-
7

---

Ada-SnapKV
Ada- -arraid
SnapKV
Pyramid
StreamingLLM
Scors
Aatny Scor
8
80
2
10
2008
50%
50%
45%
20%
1029
50%
5.6
42%
20%
1029
50%
40%
20%
200%
50%
6070
40%
20%
Budget
Sudge:
Subget
Cacha Sudget
(a) Question-awand
(b) Question-aware, Mistral-7B-Instruut-ve
(c) Question-agnostic. Bamma-3.1-BB-Inaee
(d) Question-agnostic,
Mistral-7B-Inntrucce
Figure Average Score on Ruler Among Datasets.
Ada-SnapKV
Ada- syrami
SnapKV
Pyramid
StreamingLLM
(a) CWE
(h) FWE
(c) S-NTAH-1
(d) S-NTAH-2
(c) S-NTAH-3
MK-NIAH-I
(g) MK-NIAH-2 (h) MK 3 (i) MV-NIAH
MI-NIAH
(k) VT
(I) S-QA
Figure 4: Subtask Analysis on Ruler (Question-agnosiic Llama-3.1 1-BB 1-8B-Instruct)
Table 1: Task Analysis for Llama-3. 3.1m-3.118B (Question-agnostic)
Full Cache
10%
20%
40%
Task Domains
SnapKV
Ada-SnapKY
SnapKV
Ada- Sna-S
SnapKV
Ada-SnapKV
Single-Doc. QA
43.10
23.38
23.97
28.78
31.39
35.27
36,63
Multi-Doc. QA
46.49
26.61
29.06
33.51
34.90
40.50
41.36
Summarization
28.97
21.98
22.43
23.82
24.29
26.11
26.66
Few-shol
69.45
$7.89
60.40
61.95
63.70
65.10
66.43
Synthetic
53.73
36.50
34.88
48.19
50.39
53.17
53.00
Code
57.86
59.72
60.91
60.05
61.14
60.49
60.30
Ave.
49.20
36,38
37.45
41.29
42,87
45.52
46.24
agnostic scenario. In contrast, our Ada-SnapKV and Ada-Pyramic consistently outperform the original SnapKV and Pyramid across both question-aware and question-agnostic scenarios on the two LLMs. Using the Llama-3. 1-8-3.11---nItttue model as an example: in the simpler question-aware scenario (Figure Bap. both Ada-SnapKV and Ada-Pyramii significantly reduce quality loss under small compression budgets, such as 40% and 20% cache sizes, compared to the original SnapKV and Pyramid. In the more challenging Pyramid substantially reduce quality loss across all cache budget settings. For instance, leveraging the adaptive allocation strategy, Ada-SnapKV improves SnapK V's scores at 80% and 20% cache sizes from 87.59 and 44.02 to 92.67 and 53.29, respectively.
grestion-agnostte scenario (Figure 3c Ada-SnapKV and Ada-
Subtask Analysis. Figure presents the results for 12 datasets using Llama-3. 1i-ma-3.1-8BII-sse -nnstructtt in the challenging question-agnostie scenario, highlighting the improvements brought by the adaptive budget allocation strategy. More results are provided in Appendix A.2] where our method also shows significant improvements. Overall, our adaptive budget allocation strategy significantly enhances the performance of existing methods. For example, at 80% cache budget, the original SnapK demonstrates noticeable degradation on many tasks, whereas Ada-SnapKV maintains near-lossless performance across most tasks. Particularly in difficult Boadle-in-a-Haysttae tasks like S-NIAH-3 and MK-NIAN-2 with 80% cache budget, Ada-SnapK increases SnapK scores from 62.4 and 85.2 to 97.6 and 99.6, as shown Figures 4e and 4g These results emphasize the effectiveness of incorporating adaptive budget allocation for enhancement, especially in difficult tasks.

---

Ada-SnapKV
Ada-Py adaaPyramid
SnapKV
Pyramid
StreamingLLM
30
Farmmy Score
34-
5.32
30-
10.0
42
27.5
Cache Sudget
Cocha Sudget
Cacle budget
Cochs Budget
(a) Question-aware
(b) Question-awee aware
(C) ceesstion-agnoste (d) Question-
Llama-3. 3.1 88-IIBtruce
Mistral-7FB 7----BBBtnn--teii
Llama-3 3.B Instructt
Mistral-7B Bearal-7B-Instreee Instruct-v0.2
Figure 5: Average Score on LongBench among 16 Datasets under Fixed Budgets.
4.3 LongRench Benchmark
To align with previous evaluations values of {128, 256, 512, 1024, 2048}. We further extend from standard question-awan aware compression to aggesiin-agnostic scenarios by separating the questions Longbench, thereby better reflecting real-world performance, as suggested by 12 For more details, please refer to Appendix A.7
we first the average budget for each head fixed
across 16 datasets under fixed budget constraints with the question-aware aware compression. / SnapKV and Pyramid, both utilizing Top-k selection within an observation window, achieve closely matched performance and outperform the previous sliding window eviction method, BreemmnngLLM. Moreover, our Ada-SnapK V and Ada-Pyramid methods consistently improve generation quality across different budgets, alternately outperforming their respective base methods and, under a fixed budget of 2048, approaching lossless performance. These consistent gains highlight the effectiveness of adaptive budget allocation.
Question-aware Compression: Figures Sn 5 5b present the average
From auare to Question-agnostie Compression: Figures 5 and [c further present the average scores with the auestion-agnosee compression. Our enhanced Ada-SnapKV and Ada- Pyramid methods continue to outperform the original SnapKV V and Pyramid methods. However, a comparison between question-aware and question-agnostic settings reveals notable performance drop across all methods. For example, with Llama and a 2048 budget, the average score of SnapK decreases from 49.09 42.86 while transfer from question-aware to question-agnosti settings. This indicates that commonly used question-aware settings do not adequately capture the limitations of cache eviction methods in more realistic compression scenario. We therefore suggest that future evaluations of cache eviction techniques should place greater emphasis on question -agnosic compression to improve robustness in practical applications.
More Evaluation in Question-agnostic Compression: To more accurately assess
compression, we apply ratio-based cache budget compression, which better reflects real-world perfor- mance on Longbench's highly imbalanced sample lengths. Table provides a detailed breakdown of performance across task domains for Llama-3. 3.8-BB. Overall, Code tasks are largely insensitive to cache compression, and in some cases, even see improved performance after compression. In con- trast, other tasks experience significant degradation, particularly in Summarization and QA. Notably, our Ada-SnapKV effectively mitigates these losses and improves overall performance. Across 18 domain cases and three cache budgets, Ada-Snann Ada-SnapKV delivers quality improvements in domains, demonstrating its robust benefits. Table 2 further reports results on the larger 1.1 model, showing similarly strong gains from our approach. Except for the Code domain, which remains insensitive to compression, significant improvements are observed in all other tasks.
4.4 Computation Efficiency Under Adaptive Allocation
Cache eviction methods aim to enhance the computational (memory and runtime) efficiency by selectively evicting KV cache elements. We demonstrate that, supported by our flattened cache layout and custom CUDA kernels, the computational efficiency of cache eviction under adaptive
Detailedd scores be found in Appendix A.4
"A shown Table 12 (Appendix (A.7) J Longbench sample lengths are highly imbalanced for example, the Code task Lec averages 1,235 tokens, while the QA task NarrativeQA averages 18,409. Thus, applying tixed budget, like 2048, across all datasets may introduce bias when evaluating cache eviction across tasks.

---

Table 2: Task Analysis for Llama-3.1-70B (Question-agoostic).
Stapec
b 20%
b 40%
Domain
Full Cache
SnapKV
Ada-SnapKV
SnapKV Ada-Snape ada-SnapKY
Contect Lange
Single-Doc. QA
47.15
32.54
36.05
39.35
42.21
Memory
Multi-Doc. QA
60.07
46,98
48.45
54.07
55.03
Summarization
28.69
24.40
24.81
26.21
26.57
Cache
Few-shot
72.07
66.71
67.80)
69.29
70.29
Docuainn
SnappV
Synthetic
58.25
51.25
51.50
56.00
56.75
Code
47.82
52.44
51,51
51.90
49.23
Comoott Length
Runtime
Ave.
52.25
44.95
46.08
48.91
49.64
Figure 6: Efficiency (All with FlashAttention-22)
Table 3: Performance gains from applying the V strategy to follow-up methods. All results are from the Llama-3.1-8D model the Longbench benchmark under question-agnostic settings.
DuoAttn Training based
HeadKV Training based
CriticalKV
CriticalKV
DefensiveKV w/o. Ada-KV
DefensiveKV w/. Ada-KV
Cache
w/o. Ada-KV
w/. Ada-KV
20%
39.52 48.17
42.64 47.23
42.99 47.29
43.77
43.78
46.68
40%
48.00
47.76
49.21
allocation, based on the variable-length FlashAttention technique, matches that of traditional uniform eviction methods within the same cache budget constraints. As shown in Figure 6 with a fixed budget of 1024, Ada-SnapKV achieves peak memory usage and decoding latency comparable the original SnapK and both significantly outperform the full cache case. These results demonstrate that Ada-K' KV strategy effectively enhance post-eviction generation quality while maintaining strong computational efficiency.
5 Broad Benefits of the Adaptive Budget Allocation Strategy
Due its alug-and-play design, Ada-KV broadly applicable beyond the two cases presented in this paper. As of publication, many follow-up works have adopted Ada-KV strategy, yielding wide range of enhanced applications [36]]]
37 38 especially CriticalKV (39] and DefensiveKV [40]. Beyond direct integration, concurrent methods have explored budget allocation through training- -baied profil- ing. For instance, AuoAtteetion 4T categorizes heads into "full attention" and "streaming" types, implementation. We
while HeadKV 42 enables finer-graine allocation building present the results these follow-up studies in Table to demonstrate the broad benefits of adaptive budget allocation strategy. When combined with Critical and DefensiveK Ada- KV consistently improves performance, which in turn scales with the strength of the base method. Notably, even the plug-and-play CriticalKV and DefensiveKV methods, when enhanced with the Ada-KV strategy, surpass the performance of training- based approaches like DuoAttention and HeadK These findings demonstrate that Ada-KV retains substantial value and offers significant enhancement potential, alongside more recent and advanced methods.
Conclusion
In this study, we revisit cache eviction strategies for efficient LLM inference and uncover a key overlooked factor: adaptive budget allocation across attention heads. Guided by theoretical analysis of the loss upper bound, we propose Ada-KV, the first adaptive budget allocation strategy for optimizing KV cache eviction methods. To demonstrate its plug-and-play benefits, we seamlessly integrate Ada-KV into two existing SOTA methods, introducing Ada-SnapKV and Ada-Pyramid, two adaptive eviction methods. Beyond the commonly studied question- aware compression, we also evaluate these methods in the more challenging and less explored question-agnostic compression scenario. Using the Ruler and LongBench benchmarks, we demonstrate the effectiveness of adaptive budget allocation in both settings. Our results not only expose the limitations of current cache eviction strategies but also highlight the potential of adaptive allocation to enhance cache eviction.
10

---

Acknowledgments
We would like thank the anonymous reviewers for their valuable feedback and constructive comments. This work was supported by the National Natural Science Foundation of China (NSFC) under Grants 62472400 and 622714665 1465, the National Key R&D Program of China under Grant 2025YFC3408300, and the Suzhou Basic Research Program under Grant SYG202338.
References
[1] Zihao Yi, Jiarui Ouyang, Yuwen Liu, Tianhao Liao, Zhe Xu, and Ying Shen. A survey recent advances in Ilm-based multi-turn dialogue systems. arXiv preprini arXiv:2402. 18013, 2024.
[2] Philippe Laban, Wojeiech Kryšciński, Divyansh Agarwal, Alexander Richard Fabbri, Caiming Xiong, Shafiq Joty, and Chien-Sheng Wu. Summedits: measuring Ilm ability at factual reasoning through the lens of summarization. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 9662-9676, 2023.
[3] Qiuhan Gu. Llm-based code generation method for golang compiler testing. In Proceedings of the 31 ACM Joint European Software Engineering Conference and Symposium on the Foundations of Software Engineering, pages 2201-2203, 2023.
[4] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, 1lge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidr, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
[5] Anthropic. The claude 3 model family: Opus, sonnet, haiku, March 2024. Accessed: 2024-07-
09.
[6] Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean- baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, arXiv preprint 2:XXi:24003.055000 2024.
[7] Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for lims. arXiv preprint arX0::2310.01801, 2023.
[8] 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 generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024.
[9] Dongjie Yang. Xiaodong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. PyramidInfer: Pyramid KV cache compression for thg-uhhruughput LLM inference. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar, editors, Findings of the Association for Computational Linguistics ACL 2024, pages 3258-3270, Bangkok, Thailand and virtual meeting, August 2024. Association for Computational Linguistics.
[10] Yichi Zhang, Bofei Gao, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, Wen Xiao, et al. Pyramidkv: Dynamic ky cache compression based on pyramidal information funneling. arXiv preprint arXiv:2406.0220 2024.
[11) 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. In The Tiihth- Annual Conference on Neural Information Processing Systems, 2024.
[12] Maximilian Jeblick Simon Jegou. Llm ky cache compression made casy, 2024.
[13] Kefei Wang and Feng Chen. Catalyst: Optimizing cache management for large in-memory key-value systems. Proceedings of the VLDB Endowment, (6(13):4339-4355 2023.

---

[14] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long- document transformer. arXiv preprint #xXi::2004.05750, 2020.
[15] Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Zero-shot extreme length generalization for large language models, 2024.
[16] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint (xXiv2339.1745.3 2023.
[17] Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidiss and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of impor- tance hypothesis for Ilm kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024.
[18] Siyu Ren and Kenny Q Zhu. On the efficacy of eviction policy for key-value constrained generative language model inference. arXiv preprint arXiv:2002.06662, 2024.
[19] Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin- Lin, et al. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. Advances in Neural Information Processing Systems, 37:52481-525111 2024.
[20] Yucheng Li. Huiqiang Jiang, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Amir H Abdi, Dongsheng Li, Jianfeng Gao, Yuqing Yang, et al. Mminference: Accelerating pre-filling for long-context vlms via modality-aware permutation sparse attention. arXiv preprint arXiv:2504. 60833, 2025.
[21] Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chengruidong Zhang, Bailu Ding. Kai Zhang. al. Retrievalattention: Accelerating long- context IIm inference via vector retrieval. arXiv preprint arXiv:2409. 2005:::2099.0116,,, 2024.
[22] Renze Chen, Zhuofeng Wang, Beiquan Cao, Tong Wu, Size Zheng, Xiuhong Li, Xuechao Wei, Shengen Yan, Meng Li, and Yun Liang. Arkvale: Efficient generative Ilm inference with recallable key-value eviction. Advances in Neural Information Processing Systems, 37:113134
113155, 2024.
[23] Hailin Zhang, Xiaodong Ji, Yilin Chen, Fangcheng Fu, Xupeng Miao, Xiaonan Nie, Weipeng Chen, and Bin Cui. Pqcache: Product quantization- -based kvcache for long context Ilm inference. Proceedings of the ACM on Management of Data, 3(3):1 2025.
[24] He Sun, Li Li, Mingjun Xiao, and Chengzhong Xu. Breaking the boundaries of long- context IIm inference: Adaptive kv management on single commodity gpu. arXiv preprint 2506.22500.2188 2025.
[25] Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivas- tava, Ce Zhang, Yuandong Tian, Christopher Re, and Beidi Chen. Deja vu: Contextual sparsity for efficient ilms at inference time, 2023.
[26] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 35:16344-16359) 2022.
[27] Tri Dao. Flashattention-22: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.00861, 2023.
[28] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 611-62 2023.
[29] Aaron Grattafiori, Abhimanyu Dubey, and Abhinav Jauhri eet al. The lama 3 herd of models, 2024.
12

---

[30] 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 arX:v:2310.06825, 2023.
[31] Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023.
[32] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.06654 2024.
[33] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, al. Longbench: A bilingual, multitask benchmark for long context understanding, arXiv preprint arXiv:2308. 14508, 2023.
[34] In Gim, Guojun Chen, Seung- Lee, Nikhil Sarda, Anurag Khandelwal, and Lin Zhong. Prompt cache: Modular attention reuse for low-latency inference. Proceedings of Machine Learning and Systems, 6:325-338, 2024.
[35] Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E Gonzalez, et al. Sglang: Efficient execution of structured language model programs. arXiv preprint arXiv:2312.07104, 2024.
[36] Hyng-Hyun Kim, Jinuk Kim, Sangwoo Kwon, Jae W Lee, Sangdoo Yun, and Hyun Oh Song. Kvzip: Query- kv cache compression with context reconstruction. arXiv preprint arXiv:2505.23416, 2025.
[37] Alessio Devoto, Maximilian Jeblick, and Simon Jégou. Expected attention: Ky cache compres- sion by estimating attention from future queries distribution. arXiv preprint 1rXiv:2510.00636 2025.
[38] Kevin Galim, Ethan Ewer, Wonjun Kang, Minjae Lee, Hyung Koo, and Kangwook Lee. Draft-based approximate inference for llms. arXiv preprint arXiv:2506.088: 2025.
[39] Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, and S Kevin Zhou. Identify critical kv cache in Ilm inference from an output perturbation perspective. arXiv preprint wrXiv:2502.03805 2025.
[40] Yuan Feng. Haoyu Guo, JunLin Lv, S. Kevin Zhou, and Xike Xie. Taming the fragility of kv cache eviction in Ilm inference, 2025.
[41] Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. Duoattention: Efficient long-context Ilm inference with retrieval and streaming heads. arXiv preprint arXiv:2410.1083 2024.
[42] Yu Fu, Zefan Cai, Abedelkadir Asi, Wayne Xiong, Yue Dong, and Wen Xiao. Not all heads matter: A head-level kv cache compression method with integrated retrieval and reasoning. arXiv preprint arXiv:2410. 19258, 2024.
[43] Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. Advances in Neural Information Processing Systems, 15:27168-27188 2022.
[44] 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 kv cache. arXiv preprint arXiv:2402.022
2024.
[45] Shichen Dong, Wen Cheng. Jiayu Qin, and Wei Wang. Qaq: Quality adaptive quantization for Ilm kv cache. arXiv preprint arXiv:2403. 04643, 2024.
[46] Hanshi Sun, Zhuoming Chen, Xinyu Yang. Yuandong Tian, and Beidi Chen. Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding. arXiv preprint arXiv:2404.1199 11912, 2024.
13

---

[47] Penghui Yang, Cunxiao Du, Fengzhuo Zhang, Haonan Wang, Tianyu Pang, Chao Du, and Bo An. Longspec: Long-contexs lossless speculative decoding with efficient drafting and verification. In ES-FoMo III: 3rd Workshop on Efficient Systems for Foundation Models.
[48] Yicheng Ji, Jun Zhang, Heming Xia, Jinpeng Chen, Lidan Shou, Gang Chen, and Huan Li. Specvlm: Enhancing speculative decoding of video llms via verifier-guided token pruning. arXiv preprint arXiv:2508. 16201, 2025.
[49] Gregory Kamradt. Needle In A Haystack pressure testing LLMs. Github, 2023.
[50] Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. In JCLR, 2024.
[51] Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The narrativeqa reading comprchension challenge. Transac- tions of the Association for Computational Linguistics, 6:317-328, 2018.
[52] Pradeep Dasigi. Kyle Lo, Iz Beltagy, Arman Cohan, Noah A Smith, and Matt Gardner, A dataset of informattion-sccking questions and answers anchored in research papers. arXiv preprint arXiv:2105.03001, 2021.
[53] Zhilin Yang, Peng Qi. Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhut- dinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. arXiv preprint arXiv:/809.0900 2018.
[54] Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing multi-hop QA dataset for comprehensive evaluation of reasoning steps. In Donia Scott, Nuria Bel, and Chengqing Zong, editors, Proceedings of the 28th International Conference on Computational Linguistics, pages 6609-6625, Barcelona, Spain (Online), December 2020. International Committee on Computational Linguistics.
[55] Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. Musique: Multihop questions via single-hop question composition. Transactions of the Association for Computational Linguistics, 10:539-554, 2022.
[56] Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. Efficient attentions for long document summarization. arXiv preprint arXhv:2104.02112 2021.
[57] Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan Awadallah, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, et Qmsum: new benchmark for query-based multi multi-domain meeting summarization. arXiv preprint arXiv:2104.05938 2021.
[58] Alexander R Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir R Radev. Multi-news: large-scale multi-document summarization dataset and abstractive hierarchical model. arXiv preprint arXiv:1906.00 01749, 2019.
[59] Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension, 2017.
[60] Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. Samsum cor- pus: A human-annotated dialogue dataset for abstractive summarization. arXiv preprint arXiv: 1. 12237,, 2019.
[61] Xin Li and Dan Roth. Learning question classifiers. In COLING 2002: The 19th International Conference Computational Linguistics, 2002.
[62] Daya Guo, Canwen Xu, Nan Duan, Jian Yin, and Julian McAuley. Longcoder: long-range pre-trained language model for code completion, 2023.
[63] Tianyang Liu, Canwen Xu, and Julian McAuley. Repobench: Benchmarking repository-level code auto-completion systems, 2023.
14

---

[64] Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. arXiv preprint arXiv:2307.0317,, 2023.
[65] Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing molti-hop qa dataset for comprehensive evaluation of reasoning steps. In Proceedings of the 28th International Conference Computational Linguistics, pages 6609-6666 6625, 2020.
[66] Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Pupers), pages 1601-1611 1, 2017.
[67] Daya Guo, Canwen Xu, Nan Duan, Jian Yin, and Julian McAuley. Longcoder: A long-range pre-trained language model for code completion. arXiv preprint arXiv:2306. 14899, 2023.
[68] Tianyang Liu, Canwen Xu, and Julian McAuley. Repobench: Benchmarking repository-level code auto-completion systems. arXiv preprint arXiv:2306.03000 2023.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The claims made in the abstract and introduction are fully aligned with the paper's s main contributions, as detailed in Section All stated results are supported by Section
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations of our work are explicitly discussed in Appendix A.6
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
16

---

Answer: [Yes]
Justification: All theoretical results in our paper are accompanied by a full set of clearly stated assumptions and complete proofs. Theorems and formulas properly numbered and cross-reeerenced within the main text. Detailed proofs are provided in Section 3.2 and Appendix A.8 A.9 A.10
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: All experimental settings and implementation details are described in Sec- tion 3 and Section 4 This ensures that the main claims and conclusions in the paper are reprodueible.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
17

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Ye]]
Justification: We have already release our code repository, which includes detailed instruc- tions for reproducing all main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https //nips.com ptapps.ccppublic/uuiieeeee wttp://guuiieeeeecoo wtwiiddeeeeooeeeemeeeeeeee
ssiiiiolicyi for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Our method does not involve any training process; instead, perform KV cache compression during the inference phase large language models. This a plug-and- play approach. All experimental parameters and settings provided in Section 4
Guidelines:
The answer NA means th2at the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: We use greedy decoding for LLM inference, which is deterministic and does not introduce aandmmess--consiee
with related work [IT J. Therefore, unnecessary to report error bars. Additional, we conduct comprehensive evaluations across a wide range of Ruler and LongBench benchmarks under various settings, especially including comprchensive assessments in both question-aware and question-agnostic compression
18

---

setting. This far exceeds the standard evaluation scope in the field and provides strong evidence of our method's effectiveness.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer:
[Yes]
Justification: We specify the computational resources and efficiency in our experiments in Section 4.4 which enables reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: : httt://nnssiss/e
Answer: [Yes]
Justification:
Our research complies with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
19

---

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Our proposed methods advances the state of the art in KV cache compression, enabling more efficient model inference and lowering computational resource requirements. We do not anticipate any significant negative societal consequences associated with this work, as primarily offers foundational enhancements to current techniques.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not release any models or datasets that have high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
20

---

Justification: All third-pan -arty assets (including code, datasets, and pretrained models) used in our work are properly credited in the paper with appropriate citations.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, pappesswitttteee thcode danasatss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's ss creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our work does not introduce any new datasets, code, models.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve any crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
21

---

Answer:
[NA]
Justification:
Our work does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
The development of the core methods in this research does not involve the use of large language models important, original, or non-standard components. LLMs were not used in any way that affects the scientific validity or originality of our results.
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy http:///eurripps.ee/eeeeeeeeeeee for what should or should not be described.
22

---

A Appendix
A.1 Additional Related Works
Additional works also mitigate the challenges posed by massive KV Caches during long-sequence inference while not reducing the number of cache elements. These works are orthogonal to our work. For instance, in our implementation, we have integrated the Flash Attention [26] technique to enhance efficient computation. Similar efforts, such as Paged Attention [28], employ efficient memory management strategies to reduce I/O latency without altering the size of the KV Cache. KV cache quantization methods [4 45), reduce the size of cache by lowering the precision of individual elements. The cache eviction techniques focused in this paper also be further combined and complemented with quantization in the future. Some studies also employ speculative decoding to accelerate long-sequence generation [46 47 48), typically using model with reduced KV cache to generate drafts. Future work could integrate more advanced cache eviction methods to further improve the efficiency of such approaches.
A.2 Detail results of Ruler Evaluation
(ecureecomplmmentt the missing result for the multi-hop QA subtask in Figure Jue to the space limitation. Figures 8 I0 further provide a comprehensive overview of the performance of Ada- SnapKV and Ada-Pyramid across all 16K Ruler subtasks on two LLMs in both question-agnosti and question-aware scenarios. The results demonstrate that regardless of the scenario (question- annesiio-aanootic agnostic or question-awarc model (LLama or Mistral), or cache size budgets, Ada-SnapKV and Ada-Pyramid outperform the original SnapKV and Pyramid approaches in most tasks. This highlights the broad applicability and effectiveness of the adaptive allocation strategy.
Avc agee ScoR 23
Cacte Sudget
Figure 7: M-QA Subtask Analysis on Ruler for Question-agnosti Scenario (Llama3. 1-8B (S--nsstuuct)
A.3 Robustness Analysis of Safeguard
Rather than fine-tuning the safeguard parameter, a, for each model or budget-a process that would introduce considerable complexity-we opted for fixed value. As shown in Table [Table Number], we conducted a robustness analysis of a using Mistral-7B on the LongBench benchmark. The results indicate that smaller allows for more aggressive budget allocation, improving performance under limited budgets. Conversely, a larger a performs slightly better with higher budgets. Based on this trade-off, we selected balanced value for 0.2 in this work.
Table 4: Robustness Analysis of a (Question-aware
Longbench, Mistral-7B
Budget 128
Budget 256
Budget 512
Budget 1024
Ada-SnapKV 0
35.69
38.9
40.53
41.44
Ada-SnapKV 0.2
35.48
38.61
40.58
41.61
Ada-SnapKV a 0.5
35.2
38.46
40.66
41.62
A.4 Detail results of LongBench Evaluation
Tables 6 g and present the detailed evaluation results of different methods across datasets on Bench for two models. As shown, the Ada-SnapK and Ada-Pyramid methods, with our adaptive allocation enhancement, achieve consistent improvements in overall generation quality across different budgets and models.
23

---

Ada- SnapkV
Ada-Pyramid
SnapkV
Pyramid
StreamingLMM
(a) CWB
FWE
(c) S-NIAH-1
(d) S-NIAH-2
S-NIAH-3
(1) MK-NIAH-I
1973
(g) MK-NIAH-2 (h) MK-NIAH-3
(i) MV-NIAH
ii MQ-NIAH
(k) VT
(l) S-QA
(m) M-QA
Figure 8: Subtask Analysis on Ruler (Question-aware, 1-8--nstrct)))).
Ada-SnapKV
Ada-Pyramid
+ SnapkV
Pyramid
StreamingLLM
(a) CWE
(b) FWE
(c) S-NIAH-1
(d) S-NIAH-2
(e) S-NIAH-3
(i) MK-NIAH-1
MK-NIAH-2
(h) MK-NIAH-3
ii MV-NIAH
(i) MQ-NIAHHH
(k)
()) S-QA
(m) M-QA
Figure 9: Subtask Analysis on Ruler (Question-agnostiic, Mistral-7B-Insttruttv0...
A.5
Detailed Information of Ruler Benchmark
Below is a brief overview of the different tasks in Ruler. For detailed examples, please refer to Tables and We adhere to the input prompt format from KVPress [12), dividing the input into context and question segments. The question segment is highlighted in green, while other colors represent the context segment. In question-aware compression, both the context and question segments are input into the model and compressed. For question-agnostic compression, where future questions are unpredictable, only the context segment is input for compression. After compression, the question segment is input for answer generation.
Single NIAH (S-NIAH): basic information retrieval task. In this scenario, keyword sentence, referred to as the "needle," is embedded within a lengthy text, called the "haystack." The goal retrieve the "needle" from the context. The "haystack" may consist of repetitive, noisy sentences or essays by Paul Graham [49].
Multi-keys NIAH (MK-NIAH): Multiple "needles" are inserted into the "haystack," " but the task requires retrieving only one of them.

---

Ada- Sha-SnakKV
Ada-Pyramid
SnapkV
Pyramid
StreamingLLM
(a) CWB
FWE
(c) S-NIAH-1
(d) S-NIAH-2
(c) S-NIAH-3
(1) MK-NIAH-I
275
1ar
( MK-NIAH-2 (h) MK-NIAH-3
(i) MV-NIAH
ii MQ-NIAH
(k) VT
()) S-QA
(m) M-QA
Figure 10: Subtask Analysis on Ruler (Question-aware, Distral-7B-Instruutt-v..
Multiple "needles" in the form of key-value pairs are placed within the "haystack."" " The task is retrieve all values associated with given key. Multi-queries NIAH (MQ-NIAH): Multiple "needle" in the form of key-value pairs are inserted into the "haystack". All "needles" corresponding to the keys specified in the queries need to be retrieved. (author?) (50)
Multi-values NIAH (MV-NIAH):
Variable Tracking (VT): A series of variable binding statements (e.g., X2 X1, = X2,...) are inserted at various positions within long text to simulate the recognition of entities and relationships. The task objective is to return all variable names that point to the same value V.
Common words extraction task (CWE): The input sequence is constructed by sampling words from a discrete uniform distribution within pre-defined (synthetic) word list. Words that appear most frequently are termed "common words". The model is required to identify all common words. The number of common words remains constant while the number of uncommon words increases with the length of the sequence.
Frequent words extraction task (FWE): The input sequenee is constructed by sampling words from a Zeta distribution within a pre-defined (synthetic) word list. The model need to return the top-K frequent words in the context.
Single-Hop QA (S-QA) Answer questions based on passages. The context extended to simulate long-context input by adding distracting information(golden paragraphs). And answers come from a single piece of evidence.
Multi-Hop QA (M-QA) Answer questions based on passages. For Multi-hop QA, which requires integrating information from multiple sources.
A.6 Limitations
This paper, for the first time, introduces head deie-wise allocation strategy for KV cache compression, supported by both theoretical analysis and empirical validation. However, this head-wise allocation is still limited to within single layer. In future work, plan to extend the head-wise allocation mechanism across layers and further develop corresponding theoretical analysis to demonstrate advantages.
A.7 Detailed Information of LongBench Benchmark
LongBench
33 spans multiple task domains with an average length of 6711, including single-
document QA 51 52 multi-document QA 54 555, summarization [56 57 few-shot
learning
59 60 61 synthetic tasks 33 and code generation (2, 63
25

---

Single-Doc QA: Single- document QA requires models to obtain the answer from a single piece of source. The test samples are derived from diverse datasets including Narrative
[51], qasper 122 and MultifieldQA 664, encompassing a range of documents such as legal files, governmental reports, encyclopedia, and academic papers.
Multi-Doc QA: Multi-document QA requires models to extract and combine information from several documents to obtaín the answer. The test samples are built from three Wikipedia- based multi-hop QA datasets: HotpotQA [53], 2WikiMulthhoOAA
665 and MuSiQue
[55].
Summarization task requires comprehensive understanding of the con- text. The test samples are drawn from the GovReport dataset [56], and QMSum [57]. Additionally, the MultiNews dataset originates from multi-document summarization corpus detailed in [58]
Summarization:
Few-shot Learning: Few-shot Learning task have integrated classification, summarization, and reading comprehension tasks to maintain task diversity. For classification, the TREC dataset [61] is included. The SAMSum dataset [60], which comprises messenger-style conversations with human-annotated summaries, has been incorporated for summarization tasks. Additionally, TriviaQA [66], a dataset of question-answer pairs, utilized for reading comprehension tasks.
Synthetic Tasks are designed to assess the model's ability in handling specific scenarios and patterns. Two synthetic tasks, PussageRetrievaleen and PassageCount, are included LongBench. PassageRetrieval-cot passages is randomly selected and use GPT-3.5 3-rbbbboo to summarize one of them. The task challenges the model to identify the original paragraph that matches the generated summary. PassageCount is designed to be more challenging. this task requires the model to leverage the entire context to solve it. Several passages are randomly pickedfrom English Wikipedia, randomly duplicate each paragraph multiple times, and then shuffle the paragraphs. The task is to determine the number of unique passages within the provided set.
Synthetic Task:
derived from English Wikipedia. 30
Code Completion: Code Completion task is to assist users by completing code based on previous code input and context. The LCC dataset, derived from the original Long Code Completion dataset [67]. The RepoBench-P dataset [68] is designed to aggregating information from code across files. Model is required to predict the next line of code.
Table 12 provides a comprehensive description of information pertaining to 16 datasets in LongBench. In the question-agmostic scenarios, we separate the questions from the context, with the question portions highlighted in Tables 13
15 and [8
A.8 Proof of Theorem
Theorem A.1. The L1 eviction loss can be bounded by e:
La Eviction Loss S e 2hC - 2C Σ Σ
I/A
99)
JE[L.)
where C = Max (IV, ] is a constant number, representing the max row norm among all matrices.
Proof. Consider the softmax function as: softmar(x)) Σj expp(a))
exp((x)
(10)
Thus, the attention weight after eviction procedure is:
A cofttax(-oo0 (1 II, 8;) where 8; K
(11)
- -0000
A E, exp((s - 00 I(-)
(12)
Ieeap(s)) Rerp(s?) exp((()
(13)
E,Kerp(s) E, exp(s!) E, I exp(s))
26

---

Table 5: Detailed results of Llama-3.1-8B Instrrcct Bonggeenc.(Quussion--aveaa
Single-Doc
Multi-Doc.
Summarization
Few-shotl euuming
Synthetic
Code
Qusper
MIF-nn
HatpoQA
Musique
GovReport
QMSum
MuhhiNws
TREC
TrivisQA
SAMSum
PCount
PRe
Ave.
Score
Full Cache
30.22
45,37
5088
55.97
45.00
31,26
35.12
25.38
27.20
72.50
91.64
43.57
9.41
99.50
62.88
56,43
49.20
B-
SLM
22.24
20.87
31.22
44.02
55
24.54
18.76
21.00
48
40.50
84.41
38.82
8.00
99
57.02 47 29
38.43
Pyramid
25.70
24.69
47.74
52.87
50
27.23
20.02
22.38
19.74
44.50
88 81
40.30
7.22
99.50
57.25 49 90
41.78
SnapKV
25.54
24.45
48.03
53.31
40.75
28. 19
20.13
22.36
19.55
45.50
89 20
40.62
6.97
99,50
58.45 4945549.90 90
42.03
Ada- Para-Pyramid
27.07
25.61
49.30
53.02
41 29 27.83
20.70
23.18
20.38
51.50
90.76
40.62
6.92
99.00
59.30
50.85 88
42.96
Ada- Ada-SnupKV
24.90
24.41
49.95
53.15
41.73
28.55
20.54
23.21
20.28
50.50
89
49.49
40.71
7.45
99.00
58.
58.7452.40
42.81
B=256
SLM
22.71
23.79
31.80
43.43
55
21.
20.68
20.67
46.00
87
40.8
7.20
99.
59.89499 19
39.76
Pyramid
25.53
35.15
51 44
55.03
42
28.62
22.57
23.37
22.33
56.50
91.
19
41.28
6.97
99.50
60.3651.18
44.47
SnapKV
26.02
32.49
51.62
$4.40
42.77
28.94
22.83
23.54
22.55
$3,50
91 10
40.95
7.48
99.50
60.67 39
44.48
Ada-Pyraminn
25.12
35,06
52.28
$4.66
41 89
28.76
23.14
23. 36
22.67
63,00
72
41.21
7.75
99.50
61.47 53.0793.09
45.23
Ada-SnupKV
26.11
33.39
51.44
54.94
42.15
29.54
23.01
23.85
22.88
63.50
91.57
40.94
8,00
99.50
61,9554.33
45.44
SLM
25.75
45
01
91
24.93
23.61
23.57
57. $7.50
86
6.98
60,8551.02
41.37
Pyramid
28.71
39.89
82.86
54.00
44,20
31.22
24.74
23.73
24.28
66.00
91.07
41.42
8,39
99.50
51,9953.44
46.59
SnapKY
29.22
40.01
53.15
$4.47
43.63
31.32
25.04
23.77
24.19
64.00
92.05
41.57
8.01
99.50
63.2155.05
46.76
Ada Pyramid 28.04
40.63
$1.03
54.71
43.39
M0.26
25.35
24.12
24.61
69.00
91 79
42.55
7.95
99.50
62.2855449
46.98
Ada SmapkV
29.07
40.16
52.44
53.90
43.05
31.10
25.75
24.39
24.85
69.00
92.34
42.05
7.98
99.50
63,4355.32
47.15
1024
SLM
24.97
30.22
37.06
46.57
25.24
25.72
42. 42.28
Pyramid
29.62
43.66
$4.10
55.06
44 22
31.30
27.27
24.30
25.68
68,50
97.27
41,96
7.73
99.50
63.1355.85
47.70
SnapKV
29.28
43.64
54.34
54.24
44.34
31 52
27 80
24, 39
25.95
69.00
91.72
42.50
7,80
99.50
62 29 99.956.45
47.84
Ada Pyramid
28.76
44.57
51.73
$4.89
15
31.97
27.75
25.26
25.84
70.50
91
42.37
7.67
99.50
62.9656.52
48.00
Ada-SnupKV
29.23
44.09
53.82
54.80
44.01
31.40
28.86
24.73
26.04
72.50
91.72
42.56
7.82
99.50
63,2256.33
48.16
2048
SLM
26.25
35.81
38.78
48
55
67. 50
Pyramid
29.81
45.13
54.88
55
44
31.71
30,28
24,83
26.64
71.00
91
42.42
10.43
99.50
64.6658.0 64
48.84
SnapKY
29.75
45.
55.06
56.09
42 31
31.13
24.92
26.93
72.00
91.65
42.81
9.99 99.50
64.8259.30
49.09
Ada- Pyramid
30.88
44.64
54.96
S5
45.00
31.04
30.30
25.14
26.93
72.00
91.64
43.25
10.09
99.50
64, 64.60 20
49.05
Ada Adn-SiupKV 30.64
45.23
55.46
55.55
45.26
31.14
31,28
24.89
26.72 72
72.50
91.64
42.75
9.68
99.50
64.88 59.25
49.15
Table 6: Detailed results of Mistral-7B --add-77--nntt---il on LongBench. (Question-award
Single-Due Duc. QA
Multi Duc. QA
Summarization
Syuthetic
Code
NitvOn
Qusper
MF-4n
HutpoiQA
#WIkMO4
Musique
GovRepon
QMSam
MuliiNows
TREC
TrivisQA
SAMSum
PCaum
PRe
Ave.
Score
Cache
26.74
50.00
43.45
27.77
18.49
32.91
24.64
26.99
43.32
86.31
57.39
04.3954..3
42,83
B=128
SLM
18.02
13.32
27
30.49
21 81
11.85
15.49
19.42
17.5 84
44.00
80.74
37.3 35.35
3.57
22 93
49.07
43.7743.74
28.57
Pyramid
20.78
19.76
42.92
36.31
37
13.91
18.84
21 84
20.42
46.50
84.55
40.23
2.79
65.46
51.4946.83
34.69
SnapKV
20.98
19:05 44.63
35.31
22 91
13.95
18.76
21.36
20.31
45.00
83.83
39.52
3.23
64.53
52. 1947 30
34.55
Ada-Pyramid
21.18
20.23
44.54
37.97
22.85
85
14.54
21.91
20.84
52.00
39.49
2.87
71.81
52,3447 50
35.83
Ada-SnupKV
20.10
20.14
44.20
37.32
23,90
15.29
19.15
22.27
20.71
50.00
39.6
3.09
67.11
52.26
48.25
35.48
B=256
SLM
19:08
15.30
28.27
31 87
22 26
11:37
18.08
19.37
19.99
$1 $1.00
80
39.62
3.57
15,90
51 21.44-2522
29.60
Pyramid
20.39
22.63
46.67
38.64
23.73
15.82
21.34
22.30
22.01
$7.50
83.63
40.49
2.99
75.84
53.5650.03
37.35
SHapKV
20.90
21.95
47.49
39.62
24,89
89
15:04
21.46
22.80
22.62
$7.50
84 86
40.51
3.76
77.95
$4.08 50.98
37.90
Ada Pyramid
22.21
23.26
46.92
39 04
25.01
16.99
21.43
23.07
28.28
62.50
85 70
41.11
2.71
78.86
54,26 50.49
38.49
SnapKy
20.97
23.71
47.52
38.48
15.76
22.
22.70
40.75
80.24
54.4251.46
BES12
SLM
21.12
15.99
30.82
30.39
22
10.92
21.43
19.98
22.88 88
61.50
82
11
41. 899
3.21
17.
53.43447.0 05
31,40
Pyramid
21.81
24.33
48.67
39.31
25 14
17.51
23.22
23.16
23.87
66.00
85
20
41.96
3.08
85. 71
55.10 98
39.69
SnapKV
24.29
27.83
49.00
39.6
25
26
17.63
23.53
23 36
24.44
65
00
86
18
42.01
3.27
85, 29
56.
10 73
40.35
Ada- Ada-Pyramid
22.92
26.37
47.66
39.49
25 52
18.50
23.43
23.53
23.87
66.50
85 42
41.
2.59
85
55,1452.29
40.04
Ada-SnupKV
23.62
27.34
48.70
39.81
26.42
17.36
23.42
23.26
24.50
67.50
42.04
3.04
86.64
53.01
40.58
B= 1024
SLM
22.15
18.64
31.03
34.94
45
11.93
23.89
20.60
48
64.00
84.71
41.
3.49
22, 15
53.7249.11 19
33,00
Pyramid
24.12
29.44
48.83
40.30
25
94
19.42
25 29
23.52
25.77
68.00
86.30
41.62
2.84
86.07
56.0652.57
41.01
SnapKV
24.10
30.11
48.97
40.97
26
89
18.1
25.93
23,
26.02
67.50
86 25
42.33
2.94
87.23
57.07 53.43
41.35
Ada- Ada-Pyrrmii
24.82
28.94
48.47
41.44
26
58
19.75
25. 25.08
23.84
25.54
68
85 80
42.94
3.51 85.68
52.75
41.22
Ada-SmupKV
25.11
29.98
49.27
40.62
27.05
18.54
25.85
23,46
26. 08
68.50
86.30
42.77
2.92
85.27
56.88
41.61
B=2048
SLM
22.63
22.68
33.66
92
13.76
26.96
20,80
71
66.00
00
42
2.40
25.33
$2,88
Pyramid
25.68 31 39
49.25
41.01
27.06
19.35
27.52
23.46
26.52
70.00
86
30 42.65
2.72 85.93
53.9633.33
41.84
SnapKV
25.79
32.54
49.18
42.09
21.31
18.91
28.50
24.01
26.83
70.00
86.46
42.86
2.95
86.56
57.2453.89
42.20
Ada-Pyramid
26.41
31.74 49.08
40.95
27.79
18.89
27.54
23.88
26. 78
70.50
86 30
43.44
2.61
85.93
57.2253.56
42.04
Ada- Sda-S
26.60
32.68
49.10
42.65
27.27
18.97
28.60
24.17
26.86
70.00
86
30
43.12
2.45
86.81
57.3354.11
42.31
Given 1, = softmax(s:) where 8, = 96 KT we can get Al = expls?)
Σ, exple!))
E, expp(()))
I/A4 | I,111
- E exp(s))
(14)
27

---

Table 7: Detailed results of -n-ssruttttt LongBench. (Question-agnostic)
Single-Doc.
Multi-Doc.
Summarization
Few-shotl eurming
Synthetic
Code
NATQA
Qusper
MIF-nn
HatpoQA
Musique
GovReport
QMSum
MuhtiNews
TREC
TrivisOA
SAMSum
PCount
PRe
Ave.
Score
Full Cache
30.22
45,37
55.80
55.97
45.00
31,26
35.12
25.38
27.20
72.50
91.64
43.57
9.41
99.50
62.88
56,43
49.20
B-8
SLM
13.57
11.48
24.47
34.25
25.35
11.31
20.06
17.5
80
30.50
50.14
35.33
3.00
5.50
15.50 60.0060.00 25
23.49
Pyramid
13:91
13:05
18.71
27.43
14.36
6.35
17.93
17.25
19.52
21.00
89 09
38.03
1.11
5.00
61.9357.0 00
26.35
SnapKV
12.23
11.88
17.93
25.32
12:41
8.20
17.88
16.84
19.44
22.00
90.97
37.83
2.50
5.50
61.6257.54
26.26
Ada- Para-Pyramid
13.30
13.78
17.92
31.35
15.90
7.65
19.20
17.48
19.83
24.00
90.67
38.23
4.10
6.00
63.40 56 38
27.45
A
13.52
12.79
18.30
28.55
14.13
9.35
19.04
17.26
19.94
25.00
91.44
37.96
1.54
4,00
63,34
56. 39
27.06
SLM
14.25
12.92
28.25
31
19
11.39
22.86
17.26
22.73
44.00
52.80 80
38. 38.77
4.00
9.00
16.11 59.79
25.35
Pyramid
13.57
17.36
19:29
30.95
18. 10
8.21
20.30
17.65
20.90
26.00
90.94
40.02
4.73
4.50
64.71 55.06
28.27
SnapKV
14,82
16.32
18.44
29.53
14.25
9.26
20.32
17. 17.56
21.70
28.50
91 49
40.15
4.12
5.00
64. 605478
28. 18
Ada- Pyramid
14.39
18.66
20.79
32.07
20.08
6.52
20.84
18.
21.51
31.00
91 39
41. 26
7.10
7.00
65.87 56,06
29.58
Ada Ada-Si
15.19
15.97
20.05
35.08
16.31
8.32
21.31
17.99
22.11
33,00
91.68
41.28
5.93
5.50
66.1454476
29.41
SLM
14.47
16.86
11.00
26.01
18.20
10
53.00
04
4.00
9.50
19.7457.87
27. 66
Pyramid
15.67
21.59
23.56 56
35.63
24 36
9.77
22,18
19.05
23.25
34,00
91.11
42.44
4.61
17.50
65.87 54.41
31.56
SnapKY
15.22
22.90
23.41
33.00
22.
8,91
22.63
18.31
23.50
35.00
91 39
41.85
4.50
15.50
56.79
31,20
Ada Pyramid 16.05
25.46
25.75
37.04
26.72
10,65
22.80
19.75
23.28
19,50
91.65
42.64
5.75
15.00
65.5555577
32.71
Ada SmapKV
17.72
24.22
25.38
38.00
23.47
9.25
23.55
18,91
23,82
43.00
92.14
42.81
6.34
15.00
66,1255.32
32,82
SLM
13.36
21.55
41.70
32
19.36
16.76
2.00
22.91 05
30.68
Pyramid
16.67
29.11 28.74
40.13
27 89
13.43
24.63
20.36
25.01
43,00
91.86
43.48
6.78
$2.50
65.34
36.49
SnapKV
19:57
31.67
31.52
42.04
27 89
12.94
25.31
19.66
25. 12
40
91 37
42.83
5.99
56.00
66.5653
37.56
Ada- Pyramid
20.04
31.37
32.77
40.91
30.40
14.35
24.97
21.36
25.37
46.50
91.61
43.75
9.13 54,50
65.35 55
37.98
Ada-SnupKV
19.74
30.47
31.42
40.82
28.98
16.07
25,35
20.57
25,57
$3,50
91.76
43.81
53,50 4.63
65.75
55,90
37,99
B=2048
SLM
16.86
31.83
48.22
33.59
29
15:74
30.11
68
65. 50
44.1
39.75 $6.66 66
36.13
Pyramid
21.16
26.66
37.56
43.28
19.72
27.90
21.73
26.35
53,00
92.36
44.61
6.94
89.50
64.2353.11
42.15
SnapKY
21.24
39.86
38.21
45.96
36
21.20
28.68
21.49
26.50
58.00
92.36
44.06
6.31
88.50
64.26
42,86
Ada- Pyramid
20.75
39. 12 40.18
44.49
41.17
18.94
27.90
22.06
26.71
$7.50
91 86
44.35
10.30
84.50
63.7254811
43.02
Adn Adn-SnupKV
22.42
40.26
40.13
49.60
38.23
19.70
28.49
21.76
26.77
61,50
92
35
44.04
8.27
85.00
64.11 54.09
43.55
Table 8: Detailed results of Mistral-7B--nsstuutt000222 on LongBench. (Question-agnostic)
Single. Single-Doc.C QA
Multi Doc.. QA
Summarization
Few Few-shnil Leamning
Synthetic
Code
Qusper
MF-en
2WikiMOA
Musiquue
GovRepon
QMSam
MultiNews
TREC
TriviiQA
SAMSum
PCount
PRo
Ave. Score
Pull Cache
26,74
32.84
50.00
43.45
27.77
18.49
32.91
24.64
26.99
71.00
86.23
43.32
2.94
86,31
57.39
17.3955..3 54.32
42.83
B 128
SLM
7.62
7.49
26.12
18.97
14.4 44
7.26
19.46
19.46 46
26 50
73
27
36
1.50
4.50
14.58 59
22.14
Pyramid
10.69
8.37
20.65
18. 14
13.7
5.99
17.32
17.71
19,20
69
37.48
3.25
3.00
53.63
24.55
SnapKY
10.51
9.52
21.28
18.76
13.85
5.57
17.68
18.05
18
21.50
82
36.61
3.74
2.50
52.8459.04
24.51
Ada- edd-Pyramid
12.53
9.61
20.94
18.54
14.38
4.90
17.54
18.15
19.02
22.00
84.44
37.72
3.67
1.33
55.49 59 28
74.97
Ada SmapKV
11.13
9.33
21.78
19.41
14.96
5.95
17.82
17.39
19.15
24.50
K3.71
37.33
3.73
1.50
54.91
59.92
25.16
B=256
SLM
8:37
7.47
28.82
21.42
15 30
6.98
21.74
18.23
22.92
38, 50
74 46
37.77
1,25 3.75
17 58 95
23.97
Pyramid
11.24
16.86
20.71
6.58
18.21
21.04
39.
8.68
26.21
SnapKV
10.99
10.94
23.27
21.67
13.87
6.42
26
18. 18,26
21.54
28.50
85 72
39.21
2.56
6.50
57.74
26.59
Ada- Ada-Pyramid
11.01
9.98
22.
22 22
15.42
7.17
19.78
18.74
21.27
31.00
87.30
39 32
3.42
8.43
58. 18 58 08
27.09
Ada-SnupKV
10.45
10.47
23.67
22.45
13.48
7.29
20.58
18.21
21.60
32.00
87.05
39.74
3.66
8.13
58.99
27.26
B=512
SLM
9.54
9.05
32.27
21 64
15.17
8.76
24.89
18.33
25. 39
49.00 00
75 51
39. 40
1,50
5.75
21.43 64
25 95
Pyramid
11.72
12.41
23.61
23.77
14.36
6.61
19.09
87.18
40.61
59.68
28.54
SnapKV
13.37
25.33
6.37
19.00
40.05
60.1857.5
29.04
Ada Pyramid
11.51
12.59
24.90
23.29
15.62
7.10
21.78
19.41
23.13
45.00
87
40.96
3.37
19.10
60.65 58.46
79.64
Adn-SnupKV
13.07
12.48
26.69
26.33
14.12
7.74
23.02
19.13
23.43
47.00
87.81
41.05
3.77
17.83
60.9557.20
30.10
1024
SLM
11:05
14.15
38.
23.79
18 25
9.71
27.20
20 25
26 20
$2 $2.50
76
96
41 79
0.00
8.50
24
24
56, 58
28.
Pyramid
12.22
16.32
28:01
26.55
15.58
7.59
23.91
20.04
25 38
46.00
87
41
41.37
3.74
34,29
63,
63.44
31.88
SuapKV
13.43
16.30
30.57
27.04
88
8.48
25.64
2.87
28.21
63.55 S7.38 38
Ada-Pyramid 12.34
17.09
30.24
27.13
14.01
8.26
23.60
20.03
25.01
$2.50
87
42.75
3.96
39,33
62.02
32.76
Ada SimpKV
13.86
18.39
31,07
29.03
14.92
8.24
24.82
19.86
25.59
55,50
87
41.98
3.04
35,96
64.32557.55
33.27
B-2048
SLM
11.80
20.40
43.35
25.01
18 21
9,80
29.40
40
26.73 73
$9.00
83
06
43 43.26
0.38
15
30.11 56.05
30.86
Pyramid
14.04
21.19
36.29
27.07
15.74
10.21
15.62
20.71
26 25
$2.50
87
31
43.53
4.05
65.42
64.97 58.04
35.81
SnapKV
15.91
21.64
29.
3.
11.49
27.29
20.67
26.78
56.00
87 46
42. 19
4.93
61.83
65.3456.5 88
36.26
Ada- Adr-Pyrmmid
15.44
36.86
16.88
10.75
25.52
20.74
26.60
43.05
2.89
70.96
65, 58.66
Ada-SnupKV
15.84
22.24
38.67
32
16.54
12.68
26.35
20.93
26.65
57.50
87.49
43.55
4,66
68.67
64.2458.25
37,27
Then we can obtain:
A, I.
Å = 11A OZ
(15)
28

---

Thus:
カ= Σ [E[L.]]
O
V.W?
(16)
By calculating the LI distance between their outputs, we can obtain
I
1l - 0111= 11 Σ (1
O A,V,, 11
(17)
11A,
Iill1
xE[1,N]
I
II(1
@ A.V.W |
(18)
|| の 111
(E[1.]]
I
Σ
0 A:ll V.W 100
(19)
|| Q Zi|l
VE[LAA]
I
S C
Σ II(1 [E[1,M]
0 A:ll
(20)
|4 O 11
where C' = Max [11
By expanding A,, we can further simplify the expression.
Let ||A, I,ll1 as F, E (0,1)
(21)
I
11 SC
Σ
11(1
@
@A1l
(22)
|A, + I
<E[LAA]
I IAA
F. WP-II1
c Σ Σ
(23)
F.
AE JELL.]]
Considering
I if Ki and V are retained
I
11
10 otherwise, evict K. and v?² and
Σ ]E[1,n]
A =1
(24)
V 0
<sT =1
(1- F.)
(1
1ly ا|1 SC Σ Σ
A CC
Σ
A²
(25)
F
iE(I,A) jE[I.N]
(E[1,A]
jE(I.M)
if
Due to F. = 0 = Σ
IA-=
Σ
A
JE(1,M]
jE[L,x]
=
<1
ll - u11 SC Σ
Σ
1 C Σ (1
Σ A
(26)
(E(T,A]
11
jE(I.m)
=/
2C Σ
Σ A
(27)
(E[L,]] [E[L,n]
2C Σ
Σ (1-I/)A
(28)
KE]L,AJE[LL..
= 2hC - 2C
Σ
Σ
(29)
(ELAAEEELLL
Finally,
L1 Eviction Loss <== 2hC - 2C Σ
Σ
(30)
sE[LA]
29

---

A.9 Proof of Theorem 3.2
Theorem A.2. The Top-k cache eviction {I?} minimizes the upper bound of L1 eviction loss:
Top-k eviction decision {Iī}
mine,
(31)
(2,)
mine 24C - 20 (7.)
Σ VE[I,N]
Σ
A
(32)
jE[1,M]
Proof. Given the budget allocation results {B:}. the objective of minimizing can be decomposed into i E [1, h1] independent subproblems, each expressed as follows:
arg max
I/A jE(L.T]
Σ s.t.
ΣI
=
(33)
JE(I.n]
The Top-k cache eviction maximizes the h independent subproblems, thereby minimizing the upper bound Formally:
I arg max I jE(1.m]
Σ TAA
(34)
The optimal solution satisfics:
A ETop-kk (TTp-k(A.,BB) Ac)
min Σ I I
Σ
A
(35)
]E[L,x] jE(1,n
jElLaa
A.10 Proof of Theorem 3.3
Theorem A.3. The adaptive budget allocation {Bi} derived from Algorithum Jachieves the minimal upper bound e** for loss associated with Top-k eviction methods:
min [B.}
(36)
Proof. From line of Algorithm which adaptively allocates budgets {Bi} based on the Top-k indices T = Top-k(A, indices, it follows that the resulting Top-k eviction leads to an upper bound et"
2hC - 2C Σ
Σ
A
(37)
(E[LA)
JE[L,n]
K(TT
- Σ
Σ
A
(38)
SE[LA]]
SE[1.M]
Given a fixed overall budget, i.e., Σie(I,)) B, = EiE(L,) B we can derive that ex
for any budget allocation results {B.). This is because, in the upper bound, the term A/ETTp-(AA...)) AC. This result can also be understood as global optimal solution always outperforming local optimal solution.
EieL, Σ ]j (TT Af 2 EiE(L,)) Ejen..)
A.11 Detailed Visualization of Head Concentration
Figure supplements Figure Ia in the main paper by presenting the visualization results across all layers. can be observed that in all layers, different heads exhibit significant variations in attention concentration. This indicates that the adaptive allocation strategy has great potential to reduce the eviction loss in practice.
30

---

Layer
Layer 2
Layer
Layer 4
Heed Indee:
Head HeedIndox Index
Wead fndes
40
20
20
40
20
20
40
20
Aggregared Thp Ratio
Aggregated The Rabio Layer
Aggresated Top Ratio
Apgre regieed Thp Ratio Layer 0
Layer 5
Layer 7
Treed Inboo
Head Inbox
20
2
10
20
20
40
20
20
40
20
Appropaced Tap Rato
Agpr caated Top Rabi Layer
Agaragated fatio
popropoooo Top Racio
9
Layer
Layer 12
Hewd Intex)
Head Index
Head Intex
20
20
20
20
20
46
25
Aggreguties Ratio
Agtregated The Rabio
Apgregxted Ratio
Radi
Layer
Teed Indec
Hoxt Indox
Nexd Indexc
20
2
20
Apgregated Ratio Layer
Aggregated The Robio Layer
Apgregated Radio
Layer
Layer 20
Head Indee?
Hea Index!
Head Inters
20
Layer
Layer
Layer
Heed Endeet
Head Index
Intec:
Asprogaced Rato Layer
Aggr agated Tea Layer
Aggragated Ratio Layer
Approgozod Top Layer 28
Fweed Indea
Hard Index!
2
40
40
Apgregatied Ratio Layer 29
Aggregated The Layer 30
Appregated Top Retio Layer
Appi geggeeeeeedd Tapp
Layer 32
Fraed (ndee
Head intex
2
30
Heed
20
40
20
40
20 Agg rogrego Top Rotio
Aggrogated Tho Rotio
Agpr od Top Robio
Apgrogated Top Robio
0.5
5.8
0.9
1.0
Aggregated Altontion Weight
Figure 11: Visualization of Heads' Concentrations
31

---

Table 9: S-NIAH and MK-NIAH templates in Ruler Benchmark. 32
Task Template:
Some special magic numbers hidden within the following Make sure memorize will quiz you about the numbers afterwards.
Single NIAH Subtask-1 (S-NIAH-I)
The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.
One of the special magic numbers for /word] is: [number).
What is the special magic number for [word] mentioned in the provided text?
special magic number for [word] mentioned in the provided text
Task Template:
Some special magic numbers are hidden within the following text. Make sure to memorize it. will quiz you about the numbers afterwards.
Single NIAII Subtask-2 (S-NTAH-2)
Paul Graham Essays.
One of special magic numbers for [word] is: [nmmber]. som.
What the special magic number for [word] mentioned in the provided text?
The special magic number for [word] mentioned in the provided text is
Task Template:
Some special magic words are hidden within the following text. Make sure to memorize it. will quiz you about the words afterwards.
Single NIAH Subtask-3 (S-NIAII-3)
Paul Graham Essays.
One special magic words for (word) [word].
What the special magic word for [word) mentioned in the provided text?
The special magic word for [word) mentioned in the provided text is
Task Template:
Some special magic numbers hidden within the following Make sure memorize will
quiz you about the numbers afterwards.
Paul Graham Essays.
Multi-keys NIAH
One of the special magic numbers for [word-1] 1] is: [number-I I).
Subtask-1 (MK-NIAH-I)
One of the special magic numbers for [word-2] 22 is: (number-2)
One special magic numbers for [word-3) is: (number-3)
One special magic numbers [word-4] is: (number-4) )) cer.
What the special magic number for (word-4) mentioned the provided text?
The special magic number for [word-4) mentioned in the provided text is
Task Template:
Some special magic numbers hidden within the following text. Make memorize it. will
quiz you about the numbers afterwards.
One of the special magic numbers for [word-1] 1] is: (numher-I }]-
Multi-keys NIAH
One of the special magic numbers for [word-2] ] is: [number-2]
Subtask-2 (MK-NIAH-2)
One of the special magic numbers for [word-x] is: [number-x] m...
One the special magic numbers for [word-n-1] is: (number-n-1],
One the special magic numbers for [word-n] [numher-n].
What the special magic number for (word-x) mentioned in the provided text?
The special magic number for [word-x) mentioned in the provided text is
Task Template:
Some special magic uuids are hidden within the following text. Make sure to memorize will quiz about the uuids afterwards.
One of the special magic uuids for [uuid-1] is: [uuid-1]. 1).
Multi-keys NIAII Subtask-3 (MK-NIAH-3)
One of the special magic uuids for [uuid-2] is: [uuid-2].
One
special magic uuids for [uuid-x) is: [uuid-x].
special magic uuids [uuid-n-1] is: [uuid-n-1).
One the special magic uuids for [uuid-n] [uuid-n].
What the special magic number for (uuid-x) mentioned in the provided text?
The special magic number for [uuid-x) mentioned in the provided text is
32

---

Table 10: MV-NIAH, MQ-NIAH, VT, CWE, and FWE templates in Ruler Benchmark. 32
Task Template:
Some special magic numbers are hidden within the following text. Make sure to memorize it. wwll
quiz you ahout the numbers afterwards.
Paul Graham Essays.
One of the special magic numbers for [word] is: [number-] ]
Multi-values NIAH (MV-NIAH)
One of the special magic numbers for [word] [number-2) ]. .....
One of the special magic numbers for [word] is: [number-3] ].
One of the special magic numbers for [word] is: |numher-4 4]. -.-.
What
special magic numbers [word] mentioned in the provided text?
The special magic numbers for [word) mentioned provided text are
Task Template:
Some special magic numbers are hidden within the following text. Make memorize it. will quiz about the numbers afterwards.
Paul Graham Essays.
One of the special magic numbers for [word-1] || [number-] ] s00
One of the special magic numbers for [word 2| is: [numher-2 ].
Multi-queries NIAH (MQ-NIAH)
One of the special magic numbers for [word-3) is: [number-3)
One of the special magic numbers for [word-4] [number-4] co..
What are all the special magic numbers for [word-1), (word-2), (word-3). and (word-4) mentioned
in the provided text?
The special magic numbers for (word-1), 1), (word-2), (word-3). 3) and (word-4) mentioned in the provided text are
Template:
(one task example)
Memorize and track the chain(s) of variable assignment hidden in the following text.
The grass green. The sky blue. The sun is yellow. Here go. There and back again. VAR
[X]] = [number]
Variahle Tracking
[X2]
(VT)
VAR X3
X2
VAR [X4 [X3
VAR (XS) (X4)
Question: Find all variables that are assigned the value [number] in the text above.
Answer: According the chain(s) variable assignment above, variables assigned the value [number). they are:
Task Template:
(one task example)
Below numbered list words. In these words, some appear often than others. Memorize the ones that appear most often.
Common Words Extraction
(CWE)
I. word-a word-b word-c word-a word-d word-a word-e 8. word-f.... .....
Question: What the 10 most common words in the above list?
Answer: The top 10 words that appear most often in the are:
Task Template:
Read the following coded text and track the frequency each coded word. Find the three frequently appeared coded words... word-a word word-a. word-f
word e word a ... word dl word word-g ... word-h word-a ... word-i ..... Question: Do not provide any explanation. Please ignore the dots '..... What the three frequently appeared words in the above coded text?
Frequent Words Extraction
(FWE)
Answer: According the coded text above, the three most frequently appeared words
33

---

Table 11: QA templates in Ruler Benchmark. 32
Task Template: Answer the question based on the given documents, Only give me answer not oulput any other words.
The following are given documents.
Document 1:
[document-1] 1
Single Hop QA
Document
[document-x]
(S-QA)
-.-.
Document
[document-n]
Answer the question hased on the given documents. Only give me the answer and do not output any other words.
Question:
(question)
Answer:
Task Template:
Answer the question based on the given documents. Only give me the answer and do not output other words.
The following given ducuments.
Document
[document- }
Document
[document-x]
Multi Hop
Document
(M-QA)
(document-y]
Document
[document-n]
Answer the question based on given documents. Only give me answer output any other words.
Question: (question)
Answer:
34

---

Table 12: Details 16 Datasets LongBench
Label
Task
Task Type
Eval metric
Avg len
Language
Sample Num
NrtvQA
NarrativeQA
Single-Doe. QA
F1
18,409
200
Qasper
Qasper
Single-Doc. QA
F1
3,619
EN
200
MF-en
MultiFieldQA-en
Single-Doc. QA
FI
4,559
EN
150
HotporQA
HotpotQA
Multi-Doc. QA
FI
9,151 EN
200
2WikiMQA
WWiiiMulthopQA
Multi-Doc. QA
FI
4,887 EN
200
Musique
MuSiQue
Multi-Doc. QA
Fl
11,214
EN
200
Gov GovReport
GovReport
Summarization
Rouge-I,
8,734
EN
200
QMSum
QMSum
Summarization
Rouge-L
10,614
EN
200
MultiNews
MultiNews
Summarization
Rouge-L
2,113
EN
200
TREC
TREC
Few-shotLearning
Accuracy
5,177
EN
200
TriviaQA
TriviaQA
Few-shotl earrning
F1
8,209
EN
200
SAMSum
SAMSum
Few-shotl
Rouge-L.
6,258
EN
200
PCount
PassageCount
Synthetic
Accuracy
11,141
EN
200
PRe
PassageRetrieval-en
Synthetic
Accuracy
9,289
EN
200
Loc
LCC
Code
Edit Sim
1,235
Python/C#/lava
500
RB-P
RepoBench-P
Code
Edit Sim
4,206
Python/Java
500
Table 13: LongBench templates. Single-Doc. QA Tasks.
Task Template:
are given story, which can be either novel or movie script, and question. Answer the question asconciscly you using single phrase possible. Do provide any explanation.
Story: (context)
NarrativeQAA
Now, answer the question based on the story asconcisely as you can, using a single phrase if possible. Do not provide any explanation.
Question: [question)
Task Template:
You are given scientific article and a question. Answer the question as concisely as you can, using a single phrase or sentence if possible. question cannol be answered based the information the article, write "unanswerable". If the question is yes/no question, "yes", "no". "unanswerable". Do provide any explanation.
Article: [context]
Qasper
Answer the question hased on ahove article as concisely as you can, using single phrase sentence possible. question cannot answered based on the information the article, write "unanswcrable" If the question is yes/no question, "yes", "no". "unanswerable", Do provide any explanation.
Question: (question)
Task Template:
Read the following and answer briefly.
(context)
MultiticldQA EN
Now, answer following question based on the above text, only give me the answer do not output any other words.
Question: (question)
35

---

Table 14: LongBench templates. Multi-Doc. QA Tasks.
Task Template: Answer the question based given passages. Only give me the and output any other words.
The following are given passages. [context]
HotpotQA
Answer the question based given passages. Only give me the and do output any other words.
Question: {question)
Task Template:
Answer the question based on the given passages. Only give me the answer and do not output any other words.
The following are given passages. [context)
2WikimQA
Answer the question based on the given passages. Only give me the answer and do not output any other words.
Question: (question)
Task Template:
Answer the question based on given passages. Only give me the answer and do not output any other words.
The following are given passages. [context]
Musique
Answer the question based on given passages. Only give me the answer and do not output any other words.
Question: {question)
Table 15: LongBench templates. Summarization Tasks.
Task Template:
You are given report by government ageney. Write a one-page summary of the report.
Gov Report
Report:
[context]
Now, write one-page summary of the report.
Task Template:
You given meeting transcript and query containing question instruction. Answer query one more sentences.
Transcript:
QMSum
[context]
Now, query based on above meeting transcript one more sentences.
Query: (question)
Task Template:
You are given several news passages. Write one-page summary of all news.
Multi News
News: [context]
Now, write one-page summary of all the news.
36

---

Table 16: LongBench templates. Few-shot Learning Tasks.
Task Template:
Please determine the type of the question below. Here are some examples of questions.
TREC
[context]
[question]
Task Template:
Answer the question based on the given passage. Only give me the answer and do not output any other words. The following are some examples.
TriviaQA
[context]
[question]
Task Template:
Summarize dialogue into few short sentences. The following are some examples.
SAMSum
[context]
[question]
Table 17: LongBench templates. Synthetic Tasks.
Task Template:
There some paragraphs below sourced from Wikipedia. of them may duplicates. curefully read these paragraphs and deternine how many unique paragraphs there are after removing duplicates. In other words, how many non renrreeeting puragraphs are there in total?
Passage Count
[context]
Please enter the final count of unique paragraphs after removing duplicates. The output format should only contain the number, such as 1, 2. 3. and so on.
Task Template:
Here are 30 paragraphs from Wikipedia. along with an abstract. Please determine which paragraph the abstract from.
[context)
Passage Retrieval EN
following is an abstract.
[question]
Please enter the number the paragraph that the abstract is from. The answer format must be like "Paragraph 1", "Paragraph ete. etc.
Table 18: LongBench templates. Code Tasks.
Task Template:
Please complete the code given below.
Lcc
[context]
Next line of code:
Task Template:
Please complete the code given below. (context)
Repobench-P
[question)
Next line code:
37


Reviews:
[{"review_id": "6pjFGeonqQ", "reviewer": "Reviewer_8Yky", "summary": "This paper makes the observation that different attention heads within a layer have highly differing concentrations of attention scores across tokens in the KV cache. The authors first formulate the L1 Eviction Loss, prove that it is upper-bounded w.r.t. attention scores and eviction decisions, and show that KV cache eviction methods, and prove that the Top-$k$ eviction decision achieves the tightest upper-bound.\n\nThe authors then propose Ada-KV, a budget allocation algorithm adaptive to the token concentrations of attention heads, and show that this method is optimal w.r.t. the L1 Eviction Loss. The method, implemented with a custom CUDA kernel, is integrated as a subroutine and improves upon the strong baselines of SnapKV and PyramidKV for a number of long-context benchmarks.", "questions": "Can the authors please clarify the interaction of Ada-KV's \"budget decision\" with the actual eviction policy of the respective baselines SnapKV and PyramidKV in Algorithm 2?\n\nI would appreciate it if the authors could address the questions raised in the \"Weaknesses\" section above.", "limitations": "The authors discuss the limitations (e.g. that the budget allocation is within-layer and not cross-layer) in Appendix A.5.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n\n- The per-head allocation of the total token budget is formulated intuitively, and the simple Top-$k$ solution is shown to be a minimizer of the upper bound within the L1 Eviction Loss formulation. The approach is also strongly supported by observations of the behavior and visualization of attention heads.\n\n- The CUDA kernel for Ada-KV retains the benefits of FlashAttention even when different heads have to attend to a variable number of tokens.\n\n### Weaknesses\n\n- While the L1 loss formulation appears to simplify the derivation of Ada-KV as minimizer for cache eviction loss, I do not see the intuition behind this theoretical formulation, i.e. that minimizing the Manhattan distance between attention scores is more useful than Euclidean distance. Additionally, the statement \"[t]he degradation in generation quality after cache eviction stems from changes in the attention output\" is trivially true but does not yield much insight into this theoretical formulation — especially since it's possible to construct larger attention score perturbations (w.r.t. L1 or L2 distance) which would actually result in better quality.\n\n- It is not entirely clear to me to what extent the budget allocation of Ada-KV, which is directly responsible for which tokens get evicted, overrides the PyramidKV and SnapKV cache eviction strategies in Algorithm 2. While the role of Algorithm 1 is evident inside Algorithm 2, it is hard to understand the role of either of the two strategies. For example, PyramidKV allocates a smaller and smaller token budget at higher and higher layers. What if Ada-KV produces an \"inverted pyramid\" budget, and choose to \"safeguard\" more and more tokens at higher and higher layers? What would be the precise interaction of this method with PyramidKV in such a case?\n\n- My previous point of confusion about the role of PyramidKV and SnapKV seems to be reflected in the experimental results, since the red and blue curves are nearly identical on the majority of datasets. This suggests to me that Ada-KV might be more of a KV cache eviction method in its own right than a \"plug-and-play\" integration with other methods.", "was_revised": true, "final_justification": "My confusion surrounding the novel within-layer allocation of Ada-KV versus the cross-layer allocation of Pyramid-KV has been addressed, and has led to me appreciating the impact of Ada-KV over simply Snap-KV or simply Pyramid-KV. The paper appears strong to me, and contains a theoretical motivation, substantial empirical gains, and an efficient implementation.", "rebuttal": "# Reviewer 8Yky\nThank you for your recognition of our work and your insightful comments. We address your questions as follows.\n\n## Reply to W1:\n\n> While the L1 loss formulation appears to simplify the derivation ... Additionally, the statement ... is trivially true but does not yield much insight into this theoretical formulation — especially since it's possible to construct larger attention score perturbations (w.r.t. L1 or L2 distance) which would actually result in better quality.\n\nWe chose L1 distance primarily for its simplicity. That said, we agree that other metrics—such as L2 or more task-specific distances—may better capture attention perturbations in certain contexts. Exploring these alternatives is a valuable direction for future work.\n\nWhile it is true that larger perturbations can occasionally result in better generation quality, our theoretical motivation relies on the commonly assumed Lipschitz continuity of neural networks. Under this assumption, smaller input perturbations generally lead to smaller output variations. Therefore, reducing attention output perturbations after cache eviction decreases variation in the subsequent FFN block and deeper layers, ultimately narrowing the gap between the token distributions of the original and post-eviction outputs. This keeps generation after cache eviction as close as possible to the original, which is the primary goal of compression quality. We will clarify this intuition in the theoretical section.\n\n\n## Reply to W2 & Q1:\n\n> ...What if Ada-KV produces an \"inverted pyramid\" budget, and choose to \"safeguard\" more and more tokens at higher and higher layers? What would be the precise interaction of this method with PyramidKV in such a case?...\n\nSorry for the confusion. In fact, Ada-KV performs dynamic allocation among all attention heads **within a layer**, and thus does not conflict with previous methods that handle allocation **across layers**.\n\nSnapKV uses a uniform token budget for all layers, while PyramidKV sets different layer-wise budgets. However, both methods distribute each layer’s budget evenly across all heads within that layer, without further head-wise differentiation, mostly due to the complexity of dynamic allocation and parallelization at that granularity.\nAda-KV addresses this by dynamically allocating each layer's budget among the attention heads within this layer, enabled by our theoretical formulation and an efficient CUDA kernel that facilitates the implementation.\n\n\nThus, Ada-KV and SnapKV/PyramidKV serve orthogonal roles:\n- SnapKV/PyramidKV define how many tokens to retain per layer.\n- Ada-KV decides how to split that per-layer budget among the heads.\n\nThese mechanisms do not conflict or override one another. In future work, we plan to extend Ada-KV's adaptive allocation strategy to layer-wise budget distribution, which may bring further improvements.\n\n## Reply to W3:\n\n> ... since the red and blue curves are nearly identical on the majority of datasets. This suggests to me that Ada-KV might be more of a KV cache eviction method in its own right than a \"plug-and-play\" integration .\n\nThank you for your question! The red and blue curves are very close because the previous two SOTA methods SnapKV and PyramidKV have similar performance, which leads to a small performance gap between Ada-SnapKV and Ada-PyramidKV. This reflects the similarity of the underlying baselines, rather than a limitation of Ada-KV.\n\nImportantly, Ada-KV remains a plug-and-play enhancement: it dynamically allocates budget within layers and can be applied to base cache eviction method. \nSeveral new methods were released publicly after our submission (to maintain double-blind review, we refer to them as Anonymous Methods; links will be provided after publication). We found that Ada-KV could substantially improve their performance. We will include these new results as case studies in the Appendix of the revised version after publication.\n\n\n\n| Question-Agnostic LongBench Ave. Score | 20% Cache | 40%  Cache |\n|---|---|---|\n| SnapKV  | 40.21 | 44.92 |\n| SnapKV + Ada-KV (Ada-SnapKV) | 41.39 | 45.74 |\n| Anonymous Method | 43.78 | 47.76 |\n| **Anonymous Method + Ada-KV** | **46.68** | **49.21** |"}, {"review_id": "lNqRreC1PD", "reviewer": "Reviewer_QpJ1", "summary": "The paper introduces Ada-KV, a novel method for optimizing Key-Value (KV) cache eviction in Large Language Models (LLMs) by employing adaptive head-wise budget allocation. Unlike prior cache compression methods that assign a uniform memory budget to each attention head, Ada-KV dynamically allocates budget based on attention concentration patterns. The paper establishes a theoretical framework showing that minimizing eviction loss can be guided by allocating more budget to heads with dispersed attention. This strategy is implemented in a plug-and-play fashion, improving existing methods like SnapKV and PyramidKV. Extensive experiments across datasets from Ruler and LongBench benchmarks, and multiple models (Llama-3.1-8B, Mistral-7B, Llama-3.1-70B), demonstrate significant improvements in generation quality, especially under low memory budgets and in the challenging question-agnostic setting.", "questions": "1. Have you explored applying Ada-KV during the decoding phase? What challenges or considerations arise there?\n2. Regarding the safeguard, what would happen if you don't have safeguard, or have a high or low alpha?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. This is the first paper to propose head-level KV cache budget allocation, addressing a clear limitation in prior work where all heads share the same fixed budget. The authors provide both theoretical motivation and empirical justification, showing that heads exhibit varied attention patterns, making uniform allocation suboptimal.\n2. Ada-KV can be seamlessly integrated into existing eviction frameworks like SnapKV and PyramidKV. This modular design makes the method broadly applicable and easy to adopt without retraining or significant architectural changes.\n3. The experiments is comprehensive, they do the experiments on multiple datasets with multiple model families. The method consistently outperforms baseline techniques across diverse datasets and models. \n4. Ada-KV extends to models that use GQA, avoiding unnecessary cache duplication.\n\nWeakness:\n1. The method is only applied during the prefilling stage, meaning it does not affect the autoregressive decoding steps where cache size also grows with generation. While this is not a fundamental flaw, a discussion on extending adaptive budgeting to the generation phase would be a valuable addition.\n2. While the safeguard mechanism for avoiding under-allocation to sparse heads is intuitively reasonable, its impact is not fully studied. Some ablation or sensitivity analysis on α would help understand the trade-offs.", "was_revised": false, "final_justification": "", "rebuttal": "# Reviewer QpJ1\nThank you for your valuable review and for acknowledging our work. We address your questions as follows.\n\n## Reply to W1 & Q1:\n> ... a discussion on extending adaptive budgeting to the generation phase would be a valuable addition.\n\nThank you for your suggestion. Like prior SOTA methods such as SnapKV and PyramidKV, Ada-KV currently applies cache compression only during the prefilling stage. That said, extending adaptive budgeting to the decoding phase is an interesting direction and certainly feasible.\n\nOne possible approach is to segment the generated context into chunks and apply adaptive budgeting on these chunks (e.g., after decoding each chunk). We agree that it is a valuable direction for further work and plan to investigate compression strategies for the decoding phase.\n\n\n## Reply to W2 & Q2:\n> While the safeguard mechanism for avoiding under-allocation to sparse heads is intuitively reasonable, its impact is not fully studied. Some ablation or sensitivity analysis on $\\alpha$ would help understand the trade-offs.\n\n\nThanks for your insightful question. Due to the high cost of evaluating LLMs on long-text benchmarks, we conducted only small-scale analyses for $\\alpha$ setting in the initial stage and simply chose a fixed value rather than seeking the optimal setting for different models or budgets, as this would introduce significant complexity. \n\nOur results show that Ada-KV is generally robust to the choice of $\\alpha$. A smaller $\\alpha$ allows for more aggressive budget allocation, improving performance under limited budgets, while a larger $\\alpha$ performs slightly better with higher budgets. Based on this trade-off, we set $\\alpha = 0.2$ as a balanced result.\n\n\n| Longbench, Mistral-7B | Budget 128 | Budget 256 | Budget 512 | Budget 1024 |\n|---|---|---|---|---|\n|Ada-SnapKV $\\alpha=$0 | **35.69** | **38.9** | 40.53 | 41.44 |\n|Ada-SnapKV $\\alpha=$0.2 | 35.48 | 38.61 | 40.58 | 41.61 |\n|Ada-SnapKV $\\alpha=$0.5 | 35.2 | 38.46 | **40.66** | **41.62** |"}, {"review_id": "2xKg76CGmX", "reviewer": "Reviewer_zQ8b", "summary": "The paper addresses the efficiency challenges of long-sequence inference in large language models (LLMs) by focusing on Key-Value (KV) cache eviction to reduce memory and I/O latency without sacrificing generation quality. It identifies a critical limitation in existing Top-k eviction methods: uniform allocation of a fixed cache budget across attention heads ignores the diverse concentration patterns exhibited by different heads. To overcome this, the authors propose Ada-KV, a head-wise adaptive budget allocation strategy guided by a theoretical upper bound on L1 eviction loss, which dynamically reallocates budget from sparse to dispersed heads. Ada-KV is designed to be plug-and-play, seamlessly integrating with prior SOTA methods such as SnapKV and Pyramid to form Ada-SnapKV and Ada-Pyramid variants. Extensive experiments on 13 Ruler datasets and 16 LongBench datasets, under both question-aware and question-agnostic scenarios, demonstrate that Ada-KV substantially improves post-eviction generation quality across various budget settings and base models (Llama-3.1-8B-Instruct and Mistral-7B-instruct-v0.2). Additionally, the paper presents an efficient CUDA implementation with a flattened cache layout and GQA compatibility, maintaining computational efficiency comparable to uniform eviction methods.", "questions": "1. Since Ada‑KV allocates a different number of tokens to each head, CUDA kernels will need to handle heterogeneous cache sizes per head. This variability may incur additional scheduling or memory fragmentation overhead, especially when batch‑processing multi-head operations. Could the authors clarify how this challenge is addressed in practice?\n\n2. The Top‑k eviction step requires gathering token entries from non-contiguous memory locations, which breaks coalesced memory access patterns. Such gather operations are known to reduce memory bandwidth efficiency and increase latency. Could the authors clarify how this challenge is addressed in practice?", "limitations": "See weakness and questions.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "# Strengths\n1. Introduces the head-wise adaptive budget allocation framework for KV cache eviction, moving beyond uniform budget schemes.\n\n2. Demonstrates seamless integration of Ada-KV with existing Top-k eviction methods (SnapKV and Pyramid), creating Ada-SnapKV and Ada-Pyramid with minimal changes to upstream frameworks.\n\n3. Evaluates on two major benchmarks (Ruler and LongBench), across two base LLMs and both question-aware and question-agnostic settings, underscoring the broad applicability of the method.\n\n4. Reports significant gains even under constrained budgets (e.g., 20% cache size), illustrating robustness in low-resource scenarios.\n\n5. Proposes a flattened cache storage layout and custom CUDA kernels for variable-length FlashAttention, ensuring that adaptive allocation does not incur additional runtime or memory overhead compared to uniform methods.\n\n6. Addresses Group Query Attention (GQA) compatibility, reducing redundancy in grouped KV caches and enabling a 4× cache size reduction in practice.\n\n7. Offers detailed methodological descriptions, pseudocode (Algorithms 1 and 2), facilitating understanding and reproducibility.\n\n# Weaknesses \n\n1. The main idea of this work remains relatively straightforward. As such, the contribution may be perceived as a modest advancement over existing techniques\n\n2. The theoretical analysis, while clearly structured, lacks sufficient depth and insight to represent a primary contribution of the work. Given its relatively simplicity, it would be more appropriate to condense this section into one or two succinct paragraphs, ensuring the paper remains focused on more impactful elements.\n\n3. The safeguard parameter $\\alpha$ (defaulted to 0.2) is introduced to prevent excessively small budgets for sparse heads, but there is limited analysis of its sensitivity or optimal selection across benchmarks.", "was_revised": true, "final_justification": "The score I have provided accurately reflects my current positive evaluation of the paper.\n\nEspecially, there are two \"5\" scores. I don't think this paper deserves three \"5\" scores considering that he main idea of this work remains relatively straightforward and incremental.", "rebuttal": "# Reviewer zQ8b\n\nThank you for your thorough review. We address your questions as follows.\n\n\n## Reply to W1:\n\n> The main idea of this work remains relatively straightforward. As such, the contribution may be perceived as a modest advancement over existing techniques\n\n\nThe core strengths of Ada-KV lie in its simplicity, generality, and plug-and-play design. It has the potential to enhance a wide range of existing and future cache eviction methods without requiring additional training or model changes. Moreover, achieving effective head-wise budget allocation demands both algorithmic insight and efficient implementation—a non-trivial task when aiming to maintain high performance and scalability.\n\n\n## Reply to W2:\n\n> The theoretical analysis, while clearly structured, .. Given its relatively simplicity, it would be more appropriate to condense this section into one or two succinct paragraphs...\n\n\nWe appreciate your suggestion and will revise Section 3.2 to be more concise, so as to reallocate space to present more experimental results in the main text. While the theoretical analysis is simple by design, it provides a formal justification for head-wise allocation via a loss upper bound and improves interpretability. We will refine the section accordingly. \n\n\n\n## Reply to W3:\n\n> The safeguard parameter (defaulted to 0.2) is introduced to prevent excessively small budgets for sparse heads, but there is limited analysis of its sensitivity or optimal selection across benchmarks.\n\nDue to the high cost of evaluating LLMs on long-text benchmarks, we conducted only small-scale sensitivity analysis for $\\alpha$ setting in the initial stage and simply chose a fixed value rather than seeking the optimal setting for different models or budgets, as this would introduce significant complexity.\n\nOur results show that Ada-KV is generally robust to the choice of $\\alpha$. A smaller $\\alpha$ allows for more aggressive budget allocation, improving performance under limited budgets, while a larger $\\alpha$ performs slightly better with higher budgets. Based on this trade-off, we set $\\alpha = 0.2$ as a balanced result.\n\n\n| Longbench, Mistral-7B | Budget 128 | Budget 256 | Budget 512 | Budget 1024 |\n|---|---|---|---|---|\n|Ada-SnapKV $\\alpha=$0 | **35.69** | **38.9** | 40.53 | 41.44 |\n|Ada-SnapKV $\\alpha=$0.2 | 35.48 | 38.61 | 40.58 | 41.61 |\n|Ada-SnapKV $\\alpha=$0.5 | 35.2 | 38.46 | **40.66** | **41.62** |\n\n\n\n## Reply to Q1:\n\n> Since Ada‑KV allocates a different number of tokens to each head, CUDA kernels will need to handle heterogeneous cache sizes per head. This variability may incur additional scheduling or memory fragmentation overhead, especially when batch‑processing multi-head operations. Could the authors clarify how this challenge is addressed in practice?\n\n\nFirstly, the Top-K cache eviction method selects the most important subset from the original KV cache, copies them to new contiguous memory, and releases the original space. However, head-wise budget allocation does introduce management challenges.\n\nWe address this challenge by implementing a CUDA kernel that flattens the cache of each head into a contiguous layout, along with metadata indicating the start position and length for each head. For batch processing, we use a variable-length flash attention technique that uses the metadata to extract relevant KV entries per head. Computation proceeds in a block-wise fashion along the sequence dimension, accomodating varying KV lengths across heads. Once computation for a shorter head is completed, GPU resources are dynamically reallocated to process blocks from other heads. \n\nRegarding memory fragmentation, the flattened layout ensures each head’s cache is contiguous in memory, so Ada-KV is no different from other cache eviction methods in this respect. Additionally, Ada-KV can further reduce memory fragmentation by supporting vLLM paged-attention. In accordance with the double-blind review policy, the code link will be provided after publication.\n\n## Reply to Q2:\n\n> The Top‑k eviction step requires gathering token entries from non-contiguous memory locations, which breaks coalesced memory access patterns. Such gather operations are known to reduce memory bandwidth efficiency and increase latency. Could the authors clarify how this challenge is addressed in practice?\n\nThe Top-K cache eviction method selects the most important cache entries from a large set, copying them into a new contiguous memory space, and releasing the original space. This ensures that the retained entries remain contiguous in memory after eviction. \nIn practice, eviction occurs only once per layer during the prefill phase, meaning the operation is infrequent.\nSo, the overhead of a single gather operation is negligible compared to the inference cost and the benefits it provides."}, {"review_id": "boXg6xY3Bm", "reviewer": "Reviewer_1zqk", "summary": "The paper addresses efficiency bottlenecks in LLM inference caused by the size of key-value caches, particularly under long context scenarios. Existing cache eviction strategies typically allocate compression budgets uniformly across attention heads, which fails to exploit varying head-wise attention sparsity. The authors propose Ada-KV, a novel head-wise adaptive budget allocation strategy that minimizes a derived upper bound on attention output loss. Ada-KV is shown to be plug-and-play with prior Top-k based eviction schemes (e.g., SnapKV, Pyramid). Extensive evaluations on Ruler and LongBench benchmarks in both question-aware and question-agnostic settings show consistent improvements in generation quality while preserving computational efficiency.", "questions": "1. Some very related works [1,2,3] are overlooked. The authors also need to explain the difference between the proposed method and these works in detail.\n2. After carefully reviewing Section 3.2 and Appendix A.8, I believe the theoretical component is not essential to the core contributions of the paper. Prior works [1,2,3] have already established that attention heads exhibit varying degrees of sparsity. Given this, it is rather intuitive that assigning the same budget to heads with different sparsity levels is suboptimal. In my view, the introduction of Section 3.2 does not significantly strengthen the paper.\n3. It is unclear whether the head-wise budget sizes are dynamically adjusted during the decoding phase. Based on the current description, they appear to be fixed. However, it is possible that the optimal budget sizes could evolve as new tokens are generated. This raises the question of whether the budgets are determined by the model itself or influenced by the input text. If the latter is true, further analysis would be necessary to understand this dynamic behaviour. Compared with the current theoretical component, I also believe this question warrants a more formal theoretical analysis.\n4. Although the idea of allocating different budget sizes for different heads is interesting and has been proved by lots of previous papers, I do not observe substantial improvements over the baseline methods to which Ada-KV is applied (around 1% improvement on average). Could the authors explain why?\n5. When the budget size is small (10% of the original cache size) the performance degradation is huge compared with Full Cache (37.45% vs. 49.20%, as presented in Table 1). Given that the current KV-cache optimization approaches [4,5,6] can achieve much better performance with even smaller budget sizes (5% or less), I suggest applying Ada-KV to other base models.\n\n\nReference:\n1. Minference: https://arxiv.org/pdf/2407.02490\n2. HeadKV: https://arxiv.org/pdf/2410.19258\n3. DuoAttention: https://arxiv.org/pdf/2410.10819\n4. ArkVale: https://openreview.net/pdf?id=4oAt5L4lYe\n5. SparQ: https://arxiv.org/pdf/2312.04985\n6. PQCache: https://arxiv.org/pdf/2407.12820", "limitations": "The paper does not cover the limitations.", "rating": 2, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. The paper addresses a critical and timely challenge in long-context LLM inference;\n2. Two benchmarks (LongBench and Ruler) are tested in the experiments.\n\nWeaknesses:\n1. The paper overlooks several relevant prior works that also exploit per-head or per-group cache eviction;\n2. The performance of the proposed method is not convincing.\n\nMore details can be found in [Questions].", "was_revised": true, "final_justification": "I appreciate the authors’ reply. In their response, they highlight the distinction between eviction-based approaches and dynamic sparsity methods. My main concern is that eviction-based methods were widely used in prior work primarily because, at the time, dynamic approaches were considered costly due to retrieval overhead and GPU–CPU offloading. However, recent research has shown that many dynamic approaches are both efficient and capable of achieving much better approximations, as I noted in Question 5. Ada-KV could be seen as an inspiring contribution in 2024, but given the current state of the field, I believe that for the work to be accepted, it should demonstrate its efficacy when integrated with modern dynamic sparsity methods such as Arkvale and PQCache. Unfortunately, even though I already suggested it in my review, the authors did not provide such evidence in either the main paper or the rebuttal. Therefore, I will maintain my current score.", "rebuttal": "# Reviewer 1zqk\n\nThank you for your time and effort in reviewing our work. We address your concerns below and are happy to discuss any further questions.\n\n## Reply to W1&Q1:\n\n> Some very related works [1,2,3] are overlooked. The authors also need to explain the difference between the proposed method and these works in detail.\n\n**About MInference [1] and Dynamic Sparse Attention Methods** \n\nMInference focuses on offline analysis of head patterns for prefilling acceleration. Along with methods such as [4,5,6], it belongs to the **category of dynamic sparse attention**. These methods differ fundamentally from the **cache eviction** methods we study. Specifically, dynamic sparse attention methods retain the full KV cache but dynamically select subsets for computation at each step, typically to reduce computation. In contrast, cache eviction removes less important KV entries to reduce memory usage directly. Please see our response to Q5 for more detailed discussion. We will incorporate these distinctions in the revised Related Work section.\n\n**About HeadKV [2] and DuoAttention [3]**\n\nBoth HeadKV and DuoAttention rely on **offline**, **training-based** analysis of attention heads to evaluate cache importance. As noted in the DuoAttention paper, **their method requires significant computational resources** (e.g., 8×NVIDIA A100 GPUs) for training with 7–8B parameter models, and no results are reported at the 70B scale—likely due to the prohibitive cost. In contrast, **our method is plug-and-play**, requiring no additional training cost. \n\nFunctionally, DuoAttention does not support adaptive allocation—it only converts certain heads into sliding window style. Notably, HeadKV, on the other hand, uses the Ada-KV's CUDA kernel and code to implement its adaptive head-wise budget allocation after its offline training.\n\n**On General Applicability**\n\nAda-KV is designed as a plug-and-play enhancement that can be easily integrated into existing cache eviction methods. In our paper, we demonstrate this by applying it to the widely recognized SOTA methods such as SnapKV and PyramidKV. Additionally, several new methods were released publicly after our submission (to maintain double-blind review, we refer to them as Anonymous Methods; links will be provided after publication). We found that Ada-KV could deliver substantial improvements on these concurrent methods, outperforming even training-based approaches such as DuoAttention and HeadKV. Below, we present an example result achieved by enhancing a method with Ada-KV. We will include these new results as case studies of Ada-KV's effectiveness in the Appendix of the revised version.\n\n| Question-Agnostic LongBench Ave. Score | 20% Cache | 40% Cache |\n|---|---|---|\n| DuoAttention (Training-based) | 39.52 | 48.17 |\n| HeadKV (Training-based) | 42.64 | 47.23 |\n| SnapKV  | 40.21 | 44.92 |\n| SnapKV + Ada-KV (Ada-SnapKV) | 41.39 | 45.74 |\n| Anonymous Method | 43.78 | 47.76 |\n| **Anonymous Method + Ada-KV** | **46.68** | **49.21** |\n\n\n## Reply to W2 & Q4:\n\n> ... I do not observe substantial improvements over the baseline methods to which Ada-KV is applied (around 1% improvement on average). ...\n\nFirst, the improvements brought by Ada-KV on the Ruler benchmark—specifically designed for long-context memory capacity of LLMs—are substantial and should not be overlooked. For example, **in the challenging question-agnostic setting, Ada-KV improves SnapKV from 78.08 to 88.17 at 60% cache size.** On average, this results in a 9-point improvement across different cache sizes.\n\nRegarding LongBench, it covers six task types across sixteen datasets. As the task-level analysis in Section 4.3 (line 303) of our paper, certain tasks, especially the Code task, are largely insensitive to cache compression. As shown in Table 2, the Code task incurs no loss even at 20% cache budget, indicating minimal room for further improvement. In contrast, Ada-KV increases performance from 32.54 to 36.05 on Single-Doc QA task. This suggests that performance gains are correlated with the difficulty of long-context task. This also explains why improvements are much more significant on challenging long-context benchmarks Ruler.\n\nFinally, it is worth emphasizing that Ada-KV is a plug-and-play enhancement with no additional training cost. Its consistent improvements across existing and emerging methods (see Reply Q1: On General Applicability) demonstrate broad applicability that goes beyond average percentage gains.\n\n\n## Reply to Q2:\n\n> ... it is rather intuitive that assigning the same budget to heads with different sparsity levels is suboptimal. In my view, the introduction of Section 3.2 does not significantly strengthen the paper.\n\nWe appreciate your suggestion and will revise Section 3.2 to be more concise, so as to reallocate space to present more experimental results in the main text.\n\nThat said, while the intuition behind head-wise budget allocation is straightforward, we believe that a formal theoretical analysis still adds significant value. Our theoretical analysis for the first time formally provides an upper bound on the loss under head-wise budget allocation, which offers a sound justification for the design of Ada-KV. This not only enhances the interpretability of our method but also lays a foundation for future research. \n\n\n\n\n## Reply to Q3:\n\n> It is unclear whether the head-wise budget sizes are dynamically adjusted during the decoding phase. Based on the current description, they appear to be fixed. However, it is possible that the optimal budget sizes could evolve as new tokens are generated. This raises the question of whether the budgets are determined by the model itself or influenced by the input text. If the latter is true, further analysis would be necessary to understand this dynamic behaviour. Compared with the current theoretical component, I also believe this question warrants a more formal theoretical analysis.\n\nWe clarify that Ada-KV follows the standard practice of previous methods (SnapKV and PyramidKV): cache eviction is applied once after the prefilling phase at each layer. No further cache eviction is performed during decoding, so the budget remains fixed. \n\nWe agree that optimal budgets may depend on both the model and the input, and could shift during generation. However, in cache eviction settings, dynamically adjusting budgets is difficult—once entries are evicted, they cannot be recovered. In contrast, dynamic sparse attention retains the full KV cache and allows for such adaptability. While this makes dynamic budget adjustment feasible in that context, it lies outside the scope of our work (see response to Q1). We see this as an interesting direction for future research and discuss further the connection to dynamic sparse attention (see response to Q5).\n\n\n\n## Reply to Q5:\n\n> ...Given that the current KV-cache optimization approaches [4,5,6] can achieve much better performance with even smaller budget sizes (5% or less), I suggest applying Ada-KV to other base models.\n\n\nFirst, we would like to clarify that the three referenced papers [4, 5, 6] fall under the category of dynamic sparse attention, which differs fundamentally from the cache eviction studied in our work. **So, a direct comparison is not appropriate due to the different definitions of \"budget\".**\n\n\nDynamic sparse attention [4, 5, 6] retains the **entire** KV cache in memory, selecting only a portion (the \"budget\") per step for computation. While this approach reduces computation, it does **not** reduce memory usage—regardless of whether the full KV cache is stored on the CPU [4, 6] or GPU [5]. Moreover, if stored on the CPU, retrieval indexing overhead and PCIe bandwidth quickly become bottlenecks. In contrast, cache eviction permanently removes less important KV entries, reducing both memory consumption and attention computation.\n\n\nBecause of their complementary strengths and differing goals, dynamic sparsity and cache eviction are viewed as **orthogonal** approaches, which can be combined in practice. For example, one can first apply cache eviction to reduce the KV cache to a moderate size (e.g., 50%), then apply dynamic sparsity (e.g., using only 5% of the retained cache per step) to further reduce computation cost. We believe exploring such hybrid strategies is a promising direction for future work.\n\n\n[1] Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention\n\n[2] HeadKV: Not all heads matter: A head-level kv cache compression method with integrated retrieval and reasoning\n\n[3] DuoAttention: Efficient long-context llm inference with retrieval and streaming heads\n\n[4] ArkVale: Efficient generative llm inference with recallable key-value eviction\n\n[5] SparQ attention: Bandwidth-efficient llm inference\n\n[6] PQCache: Product quantization-based kvcache for long context llm inference"}]

Author rebuttal:

Meta-review: Ada-KV introduces the first head-wise adaptive KV cache allocation strategy, optimizing LLM inference. Its strengths include plug-and-play integration without retraining, significantly improving generation quality under low memory budgets. Reviewers initially questioned the L1 loss intuition, interaction with baselines, and safeguard parameter sensitivity. Authors clarified L1 was for simplicity and explained Ada-KV's orthogonal within-layer/cross-layer interaction with existing methods. Crucially, the core disagreement regarding the conflation of KV cache eviction with dynamic sparsity was refuted, emphasizing these are orthogonal and active research fields. The paper was accepted due to its novelty, effectiveness, and clarified contribution to the field.
Decision: Accept (poster)