Learning to Evict from Key-Value Cache

The growing size of Large Language Models (LLMs) makes efficient inference challenging, primarily due to the memory demands of the autoregressive Key-Value (KV) cache. Existing eviction or compression methods reduce cost but rely on heuristics, such as recency or past attention scores, which serve only as indirect proxies for a token’s future utility and introduce computational overhead. We reframe KV cache eviction as a reinforcement learning (RL) problem: learning to rank tokens by their predicted usefulness for future decoding. To this end, we introduce KV Policy (KVP), a framework of lightweight per-head RL agents trained on pre-computed generation traces using only key and value vectors. Each agent learns a specialized eviction policy guided by a holistic reward, derived from future utility, that evaluates the quality of the ranking across all cache budgets, requiring no modifications to the underlying LLM or additional inference. Evaluated across two model families on the long-context benchmark RULER (up to 128K tokens) and the multi-turn dialogue benchmark OASST2-4k, KVP significantly outperforms strong baselines. Zero-shot tests on standard downstream tasks (BoolQ, LongBench passage retrieval, GovReport) further show that KVP generalizes beyond its training distribution and to considerably longer sequence lengths. These results demonstrate that learning to predict future token utility is a powerful and scalable paradigm for adaptive KV cache management.

Large language models generate text one token at a time and store a "key-value (KV) cache" of intermediate results for every prior token to avoid redundant recomputation. This cache grows linearly with sequence length and can reach tens of gigabytes for long inputs, capping context length and inflating cost. Existing methods shrink it by evicting tokens, but rely on heuristics like recency or past attention scores, indirect proxies for whether a token will actually be needed later. We reframe eviction as a learning problem: train a lightweight model to rank cached tokens by predicted future usefulness, so the lowest-ranked can be dropped to fit any memory budget. We use reinforcement learning, with a reward that measures how much future attention the eviction destroys, summed across all possible budgets. A separate small ranker is trained per attention head, offline on pre-recorded model traces, leaving the underlying model unchanged. Across two model families and benchmarks up to 128000 tokens, our method consistently beats heuristic and attention-based baselines and generalizes zero-shot to tasks and context lengths it was never trained on. Thus, we show that a small policy, learned once and offline, can replace hand-crafted eviction rules and substantially reduce the cost of long-context inference.