Search papers, labs, and topics across Lattice.
This paper introduces DistillCache, a reinforcement learning-based framework for adaptive Key-Value (KV) cache eviction in large language models (LLMs), addressing the memory bottleneck associated with long-context inference. By formulating KV-cache eviction as a sequential decision-making problem, DistillCache employs a lightweight policy network trained with KL-divergence rewards, achieving 94.2% accuracy on LongBench with only 25% of the cache budget. The method outperforms existing heuristic baselines and concurrent RL-based approaches, demonstrating significant improvements in both accuracy and throughput for memory-efficient LLM inference.
DistillCache retains over 94% accuracy on long-context tasks while slashing memory usage by 75%, outperforming traditional methods and setting a new standard for efficient LLM inference.
Transformer-based large language models (LLMs) achieve strong performance across many tasks, but their Key-Value (KV) cache grows linearly with sequence length, creating a severe memory bottleneck for long-context inference. Existing heuristic eviction methods (e.g., H$_2$O and SnapKV) rely on static attention or positional signals that often fail to capture a token's future predictive influence. We propose DistillCache, a reinforcement learning framework that formulates KV-cache eviction as a sequential decision problem. DistillCache learns a lightweight policy network using rich internal model signals (attention statistics, value norms, entropy, and position) and trains it with REINFORCE via a per-step KL-divergence reward to preserve the full-cache output distribution. On a 7B-parameter instruction-tuned Transformer (Mistral-7B-Instruct-v0.3), DistillCache retains 94.2% of full-cache accuracy on LongBench at a 25% cache budget, outperforming both strong heuristic baselines (H$_2$O, SnapKV) by up to 2.7 absolute points and, under our re-implementations, concurrent RL-based methods (ForesightKV, RLKV) by up to 1.4 points on long-context tasks. On reasoning benchmarks, DistillCache is competitive with the best concurrent method and surpasses it under aggressive compression. It also delivers up to 2.1x full-cache throughput while maintaining competitive practical efficiency. These results highlight the effectiveness of learned, distribution-aware policies for memory-efficient long-context LLM inference.
