Learning to Evict from Key-Value Cache

TL;DR

This paper introduces KV Policy, a reinforcement learning framework that learns to efficiently manage Key-Value cache eviction in large language models, significantly improving performance and generalization over heuristic methods.

Contribution

The paper proposes a novel RL-based approach for cache eviction in LLMs, training lightweight per-head agents to predict token utility without modifying the model or adding inference overhead.

Findings

KVP outperforms heuristic baselines on RULER and OASST2-4k benchmarks.

KVP generalizes well to downstream tasks and longer contexts.

The RL approach effectively predicts future token utility for cache management.

Abstract

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 future utility, which evaluates the quality of the ranking across all cache budgets, requiring no modifications to the underlying LLM or additional inference. Evaluated across two different model families on the long-context benchmark RULER and the multi-turn dialogue benchmark OASST2-4k, KVP significantly outperforms baselines. Furthermore, zero-shot tests on standard downstream tasks (e.g., LongBench, BOOLQ, ARC) indicate that KVP generalizes well beyond its training distribution and to longer context lengths. These results demonstrate that learning to predict future token utility is a powerful and scalable paradigm for adaptive KV cache management.