Predicting Future Utility: Global Combinatorial Optimization for Task-Agnostic KV Cache Eviction
Abstract
Given the quadratic complexity of attention, KV cache eviction is vital to accelerate model inference. Current KV cache eviction methods typically rely on instantaneous heuristic metrics, implicitly assuming that score magnitudes are consistent proxies for importance across all heads. However, this overlooks the heterogeneity in predictive fidelity across attention heads. While certain heads prioritize the instantaneous contribution of tokens, others are dedicated to capturing long-horizon utility. In this paper, we propose that optimal budget allocation should be governed by the marginal utility in preserving long-term semantic information. Building on this insight, we propose LU-KV, a novel framework that formulates head-level budget allocation as a global combinatorial optimization problem to maximize the long-horizon marginal contribution of reserved tokens. To solve this non-convex problem, we employ a convex-hull relaxation and a marginal-utility-based greedy solver, achieving near-optimal solutions. Furthermore, we implement a data-driven offline profiling protocol to facilitate the practical deployment of LU-KV. Evaluations on LongBench and RULER benchmarks demonstrate that LU-KV reduces KV cache size by 80\% with minimal performance degradation, while also decreasing inference latency and GPU memory footprint.
Lay Summary
Modern AI assistants, such as chatbots and document-analysis tools, need to process and remember long conversations or documents. To do this efficiently, they store intermediate information called a "KV cache" — essentially a memory bank that helps the model recall what it has already read. However, this memory grows rapidly with text length, consuming enormous amounts of computing resources and slowing down responses. To address this, researchers have developed methods that selectively discard less important memories. Current approaches treat all parts of the AI model equally, using simple rules of thumb — such as "keep whatever seemed most relevant just now" — to decide what to forget. This works reasonably well, but ignores the fact that different components of the model play very different roles: some focus on what is immediately useful, while others track information that only becomes important much later in a conversation. In this paper, we introduce LU-KV, a smarter memory management system that respects these differences. Rather than applying the same deletion rule everywhere, LU-KV asks: which memories, if kept, will provide the greatest benefit over the long run? It allocates more storage budget to the model components that handle long-range information, and less to those that only need short-term context — much like prioritizing a long-term planner over a short-term to-do list. We show that LU-KV can reduce KV cache by 80% while preserving the quality of the model's responses, making AI assistants faster, cheaper to run, and more accessible — without sacrificing their ability to understand long documents.