Binqin Shi

Author directory

2026

The prohibitive memory footprint of the Key-Value (KV) cache imposes a critical bottleneck for efficient long-context LLM serving. Current eviction methods, ranging from attention-based heuristics to recent learnable approaches, rely on scoring criteria decoupled from the downstream task performance that ultimately matters. To address this, we introduce KV-Retriever, a layer-wise learnable importance scoring module directly optimized by the end-to-end downstream task loss. We demonstrate that aligning the scoring objective with the actual task signal—rather than intermediate proxies such as attention statistics or oracle labels—yields a fundamentally more accurate notion of token importance. Functioning as a lightweight, plug-and-play module with only ∼0.01% additional parameters on a frozen 8B backbone, KV-Retriever is compatible with diverse host eviction frameworks. Extensive experiments show that, for each backbone, one trained retriever can be reused unchanged across multiple KV cache compression architectures and delivers consistent improvements on LongBench, LongBench-v2, and RULER. KV-Retriever requires one-time supervised training with golden answers; “plug-and-play” refers to deployment into otherwise unchanged host eviction frameworks, not to training-free adaptation. Our code is available at https://github.com/Taylor-Gavel/kv_retriever.