KV Admission: Learning What to Write for Efficient Long-Context Inference

TL;DR

This paper introduces WG-KV, a learned cache management method for long-context LLM inference that significantly reduces memory and speeds up processing by predicting token utility before writing to the cache.

Contribution

It formalizes KV cache management as a causal system and proposes WG-KV, a lightweight mechanism that learns to predict token utility, improving efficiency in long-context inference.

Findings

Reduces memory usage by 46-68%

Achieves 3.03-3.70x prefill speedup

Achieves 1.85-2.56x decode speedup

Abstract

Long-context LLM inference is bottlenecked by the quadratic attention complexity and linear KV cache growth. Prior approaches mitigate this via post-hoc selection or eviction but overlook the root inefficiency: indiscriminate writing to memory. In this paper, we formalize KV cache management as a causal system of three primitives: KV Admission, Selection, and Eviction. We instantiate KV Admission via Write-Gated KV (WG-KV), a lightweight mechanism that learns to predict token utility before cache entry. By filtering out low-utility states early to maintain a compact global cache alongside a sliding local cache, WG-KV reduces memory usage by 46-68% and delivers 3.03-3.70x prefill and 1.85-2.56x decode speedups on Llama and Qwen models, while maintaining compatibility with FlashAttention and Paged-KV systems. These results demonstrate that learning what to write is a principled and practical recipe for efficient long-context inference. Code is available at https://github.com/EMCLab-Sinica/WG-KV.