Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference

TL;DR

Ada-KV introduces a head-wise adaptive cache eviction strategy for large language models, significantly improving inference efficiency by tailoring cache management to attention head patterns, outperforming existing uniform approaches.

Contribution

This paper presents Ada-KV, the first adaptive, head-wise cache eviction method guided by a theoretical loss bound, enhancing LLM inference efficiency and quality.

Findings

Significant quality improvements over existing methods.

Effective integration with prior cache eviction techniques.

Validated on multiple datasets with diverse scenarios.

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.