Title: Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference
Venue: NeurIPS 2025
Accepted: yes
Keywords: Efficient AI

Large Language Model

LLM Inference

KV Cache
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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].\n\nQUESTIONS:\n1. 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\n\nLIMITATIONS:\nThe paper does not cover the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n# 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nSee weakness and questions.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\nCan 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.\n\nLIMITATIONS:\nThe authors discuss the limitations (e.g. that the budget allocation is within-layer and not cross-layer) in Appendix A.5.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}