Source-linked AI summary
Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference
Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, S. Kevin Zhou
TL;DR
Long-sequence LLM inference is constrained by rapidly growing KV caches, while existing eviction methods allocate budgets uniformly despite differing attention patterns across heads. Ada-KV derives an eviction-loss upper bound to guide head-wise adaptive allocation and integrates with prior Top-k methods. Across Ruler and LongBench under question-aware and question-agnostic settings, it demonstrates improved performance over existing methods.
Problem
Existing KV-cache eviction methods uniformly allocate compression budgets across attention heads, overlooking their distinct attention patterns during long-sequence inference.
Method
Ada-KV derives a theoretical upper bound on eviction loss and uses it to adaptively allocate budgets across heads while integrating with existing Top-k eviction methods.
Results
Ada-KV improves performance across 13 Ruler and 16 LongBench datasets under both question-aware and question-agnostic compression scenarios.
Takeaways & Limitations
Ada-KV provides a plug-and-play strategy for optimizing existing KV-cache eviction methods through adaptive head-wise budget allocation.
Takeaways & Limitations
The theoretical guarantee concerns an upper bound whose reduction does not strictly guarantee reduced practical eviction loss.
Abstract
from arXiv · showhide
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.
1 Introduction
Long-sequence inference makes KV-cache growth a major memory and runtime challenge. Ada-KV addresses inefficient uniform head-wise budgets with theory-guided adaptive allocation and improves cache eviction across benchmarks and compression settings.
- Motivation: An 8B LLM processing 2M tokens can require up to 256GB of KV cache, burdening GPU memory and inference runtime.
- Motivation: Existing eviction methods retain top-k cache elements under predefined budgets, reducing memory usage and accelerating decoding.
- Problem: Uniform budgets overlook head-specific attention concentration, wasting capacity on sparse heads and increasing eviction loss for dispersed heads.
- Ada-KV: Ada-KV reallocates budgets from sparse to dispersed heads, adapting cache allocation to distinct attention patterns.
- Evaluation: Evaluations across 13 Ruler and 16 LongBench datasets show improvements under both question-aware and question-agnostic compression scenarios.
- Ada-KV: Ada-KV derives a theoretical eviction-loss upper bound and uses it to guide adaptive allocation while integrating plug-and-play with existing Top-k methods.
2 Related Works
Related work reduces long-sequence inference costs through KV-cache eviction or selective attention computation. These approaches differ in whether they remove cache entries or retain all entries while narrowing computation.
- KV-Cache Eviction: KV-cache eviction methods reduce memory and decoding costs by removing non-critical cache elements, including through sliding-window and Top-k strategies.
- Top-k Methods: Top-k methods prioritize critical cache elements using attention weights, with prior systems combining special, punctuation, recent, and selected elements.
- Sparse Attention: Sparse attention keeps the full KV cache but selectively uses critical entries during computation, so it does not reduce KV-cache memory footprint.
3 Methodology
Ada-KV formalizes eviction loss as an upper-bound minimization problem and adaptively allocates per-head budgets according to attention concentration. The resulting allocation integrates with Top-k eviction methods and is theoretically optimal for the bound while reducing practical eviction loss.
- 3.1 Multi-Head Self-Attention: The methodology models multi-head self-attention, KV-cache eviction decisions, and post-eviction outputs for a fixed overall budget.Each head retains only its allocated number of cache elements, while the post-eviction output is compared with the original attention output.
- 3.2 Theoretical Foundation: L1 eviction loss is defined as the distance between pre- and post-eviction self-attention outputs, and a row-norm analysis derives an upper bound.The bound provides the optimization target for the subsequent eviction and allocation strategies.
- 3.2 Theoretical Foundation: Under a fixed per-head budget, Top-k eviction retains the cache elements with the highest attention weights and achieves the tightest upper bound.The decision retains only the top B_i elements in head i and evicts the rest.
- 3.3 Adaptive Budget Allocation: Ada-KV selects the B largest attention weights across heads and assigns each head a budget according to its frequency among those selections.Sparse heads receive smaller budgets, while saved capacity is reallocated to dispersed heads.
- 3.3 Adaptive Budget Allocation: Theorem 3.3 states that Ada-KV minimizes the upper bound among Top-k allocation strategies, including uniform allocation.The paper notes that a lower bound does not strictly guarantee lower practical loss, but treats bound minimization as the design objective.
4 Experiments
Experiments evaluate Ada-KV integrations across Ruler and LongBench, using multiple models, cache budgets, and question-aware or question-agnostic compression. Across these settings, adaptive allocation improves quality while retaining comparable computational efficiency.
- Experimental Setup: Evaluations use Llama-3.1-8B-Instruct and Mistral-7B-instruct-v0.2 on Ruler’s 13 tasks and LongBench’s 16 datasets.SnapKV and Pyramid are the primary baselines; StreamingLLM provides a sliding-window reference.
- Experimental Setup: Experiments cover cache budgets of 20%, 40%, 60%, and 80% of the original cache size in question-aware and question-agnostic settings.Question-agnostic compression reveals questions only after compression and is intended to reflect prompt-caching and multi-turn scenarios.
- Ruler Benchmark: Ada-SnapKV improves SnapKV’s question-agnostic Llama-3.1-8B scores from 87.59 to 92.67 at 80% cache and from 44.02 to 53.29 at 20% cache.The comparison is reported for the Ruler benchmark’s average score.
- Ruler Benchmark: At 80% cache on difficult Ruler tasks, Ada-SnapKV raises SnapKV’s scores from 62.4 to 97.6 on S-NIAH-3 and from 85.2 to 99.6 on MK-NIAN-2.The results are reported for Llama-3.1-8B-Instruct in the question-agnostic setting.
- LongBench Benchmark: On LongBench, Ada-SnapKV and Ada-Pyramid consistently improve their base methods across fixed budgets, with performance approaching lossless quality at budget 2048.This finding is reported under question-aware compression; question-agnostic compression produces a performance drop across all methods.
- Computation Efficiency: With a fixed budget of 1024, Ada-SnapKV has peak memory usage and decoding latency comparable to SnapKV, while both outperform full-cache inference.The efficiency comparison uses FlashAttention-2 and is supported by flattened cache layout and custom CUDA kernels.
5 Broad Benefits of the Adaptive Budget Allocation Strategy
Ada-KV’s plug-and-play design supports applications beyond Ada-SnapKV and Ada-Pyramid. Follow-up work has adopted the strategy, while concurrent approaches also investigate training-based head profiling.
- Broad Applicability: Ada-KV is broadly applicable beyond the two presented integrations because of its plug-and-play design.The paper reports follow-up applications including CriticalKV and DefensiveKV.
- Related Extensions: Concurrent methods explore budget allocation through training-based profiling, including DuoAttention’s categorization of heads into “full attention” and “streaming” types.This represents an alternative direction to direct integration of Ada-KV.
6 Conclusion
The paper identifies adaptive allocation across attention heads as an overlooked factor in KV-cache eviction and proposes Ada-KV to optimize it. Evaluations show benefits across question-aware and question-agnostic compression on Ruler and LongBench.
- Conclusion: Adaptive budget allocation across attention heads is identified as an overlooked factor in existing KV-cache eviction strategies.The paper links this issue to differing attention patterns across heads.
- Conclusion: Ada-KV uses a theoretical loss upper bound to guide head-wise allocation and integrates with existing SOTA methods as Ada-SnapKV and Ada-Pyramid.The integrations demonstrate the strategy’s plug-and-play design.
- Conclusion: Across Ruler and LongBench, adaptive allocation improves cache eviction in both question-aware and question-agnostic settings.The conclusion presents these results as evidence of adaptive allocation’s effectiveness.
A.1 Additional Related Works
The paper distinguishes cache eviction from orthogonal techniques that improve computation or memory management without reducing the number of KV-cache elements.
- Orthogonal Approaches: FlashAttention is integrated to improve efficient computation while remaining orthogonal to methods that do not reduce KV-cache elements.Paged Attention reduces I/O latency through memory management without changing cache size.
- Orthogonal Approaches: Paged Attention mitigates long-sequence inference challenges by reducing I/O latency without altering the size of the KV cache.This contrasts with cache eviction, which selectively removes cache elements.
A.2 Detail results of Ruler Evaluation
Across Ruler subtasks, Ada-SnapKV and Ada-Pyramid outperform their original methods in most tasks across scenarios, models, and cache budgets.
- Ruler Evaluation: Ada-SnapKV and Ada-Pyramid outperform original SnapKV and Pyramid approaches in most Ruler subtasks.This pattern holds across question-agnostic and question-aware settings, two models, and different cache-size budgets.
A.3 Robustness Analysis of Safeguard α
The safeguard parameter α exhibits a budget-dependent trade-off: smaller values support more aggressive allocation under limited budgets, whereas larger values perform slightly better at higher budgets.
- Robustness Analysis: Smaller α enables more aggressive budget allocation and improves performance under limited budgets.The robustness analysis uses Mistral-7B on LongBench.
- Robustness Analysis: Larger α performs slightly better with higher cache budgets.The authors use a fixed value rather than tuning α separately for each model or budget.
A.4 Detail results of LongBench Evaluation
On LongBench, adaptive allocation consistently improves overall generation quality across datasets, cache budgets, models, and both question-aware and question-agnostic settings.
- LongBench Evaluation: Ada-SnapKV and Ada-Pyramid consistently improve overall generation quality across 16 LongBench datasets.Detailed results cover two models and multiple budgets.
- LongBench Evaluation: The LongBench evaluations include question-aware and question-agnostic scenarios for Llama and Mistral models.The reported comparisons cover both adaptive integrations and their original baselines.
A.5 Detailed Information of Ruler Benchmark
Ruler covers retrieval, tracking, extraction, and question-answering tasks, while the paper formalizes eviction loss and derives adaptive allocation for Top-k cache eviction. The benchmark also distinguishes question-aware from question-agnostic compression, and the analysis reports broad empirical gains with a single-layer scope limitation.
- Compression Scenarios: Question-aware compression uses context and question segments, whereas question-agnostic compression compresses only the context before the question is supplied.The question segment is separated from the context in the question-agnostic setting.
- Ruler Benchmark: Ruler includes needle-in-a-haystack retrieval, variable tracking, common-word extraction, frequent-word extraction, and question-answering tasks.The benchmark descriptions specify synthetic task constructions and retrieval objectives.
- Theoretical Analysis: Theorem 3.1 bounds L1 eviction loss using the row norm, with the constant representing the maximum row norm among matrices.The bound connects pre- and post-eviction attention outputs through an upper-bound objective.
- Theoretical Analysis: Top-k cache eviction minimizes the theoretical upper bound for fixed budget allocations by retaining cache elements with the highest attention weights.The proof decomposes the objective into independent head-level subproblems.
- Adaptive Allocation: Algorithm 1’s adaptive allocation achieves the minimal upper bound associated with Top-k eviction under a fixed overall budget.The paper characterizes this as a global optimum outperforming a local optimum.
- Scope and Motivation: Attention concentration varies substantially across heads in all layers, motivating adaptive allocation, but the proposed strategy remains limited to within a single layer.Future work aims to extend allocation across layers and develop corresponding theory.