ForesightKV: Optimizing KV Cache Eviction for Reasoning Models by Learning Long-Term Contribution

Recently, large language models (LLMs) have shown remarkable reasoning abilities by producing long reasoning traces. However, as the sequence length grows, the key-value (KV) cache expands linearly, incurring significant memory and computation costs. Existing KV cache eviction methods mitigate this issue by discarding less important KV pairs, but often fail to capture complex KV dependencies, resulting in performance degradation. To better balance efficiency and performance, we introduce ForesightKV, a training-based KV cache eviction framework that learns to predict which KV pairs to evict during long-text generations. We first design the Golden Eviction algorithm, which identifies the optimal eviction KV pairs at each step using future attention scores. These traces and the scores at each step are then distilled via supervised training with a Pairwise Ranking Loss. Furthermore, we formulate cache eviction as a Markov Decision Process and apply the GRPO algorithm to mitigate the significant language modeling loss increase on low-entropy tokens. Experiments on AIME2024 and AIME2025 benchmarks of three reasoning models demonstrate that ForesightKV consistently outperforms prior methods under only half the cache budget, while benefiting synergistically from both supervised and reinforcement learning approaches. Code is available at \url{https://github.com/RUCAIBox/ForesightKV}.

LLMs can solve harder problems by working through many steps, but this process requires them to keep a growing record of past information, which uses a lot of computer memory and slows them down. This paper introduces ForesightKV, a method that teaches these systems how to keep the most useful parts of that record and safely discard less useful parts while producing long answers. Rather than relying on simple hand-written rules, our method learns from examples where we look ahead to see which stored information will actually matter later, and it is further trained to avoid forgetting details that are especially important for accurate predictions. On challenging math reasoning benchmarks, ForesightKV helps several reasoning models perform better than previous memory-saving methods while using only about half of the usual storage budget.