A Probabilistic Interpretation of KV Cache Eviction
Renato Geh, Alex Chen, Daniel Israel, Aditya Grover, Guy Van den Broeck
cs.CL, cs.AI
2026-08-28
UCLA proves exact KV eviction is NP-complete and recasts top-k as unbounded-bias estimators; sampling plus SNIS correction is more robust across LongBench and RULER.
KV cache eviction sells a simple trade: drop some cached keys and values, gain memory and throughput, and hope generation quality barely moves. H2O, SnapKV, TOVA, and StreamingLLM already do this in practice. They score tokens and keep a top-k set.
What they do not do is state the problem. The quantity that matters is attention after eviction, on later decode steps, when dropped entries are gone. Most papers score at eviction time and then decode as if the missing mass never existed. This UCLA paper writes the decision problem down, proves it is hard, and shows that today's heuristics are estimators whose bias is unbounded.
KVEviction asks: given K, Q, V, an error vector ε, and a keep ratio r, does there exist a subset of size ⌊n·r⌋ whose softmax attention stays within ε of the full cache? A reduction from Partition shows NP-completeness. That version is already easier than deployment, because it only has to match attention at eviction time. Decode-time matching is harder: the dropped keys are no longer there.
Attention is an expectation. softmax(qKᵀ) is a categorical distribution p over n positions; the head output is Ep[V]. Eviction restricts the support and renormalizes. Whatever probability mass was dropped rescales the survivors, and the distortion compounds across layers and future tokens.
Top-k scoring is a zero-variance estimator. The same scores always keep the same entries. Variance is zero; bias can be arbitrarily large. Probabilistic eviction reads those scores as an unnormalized proposal π, draws m samples with replacement, and drops entries with count zero. At decode time, self-normalized importance sampling (SNIS) splits the expectation into the pre-eviction span, which must be estimated, and the post-eviction span, which is exact. The two pieces are mixed by an estimate of the missing normalizing constant. Bias and mean squared error of SNIS are both O(1/m).
A temperature τ on the importance weights trades the two errors: τ=1 is ordinary correction; τ→∞ turns correction off and cuts variance. For grouped-query attention, one mixture proposal covers a whole KV group so every head in the group stays consistent.
The models are Llama3.2-3B and Qwen3-4B. Baselines are StreamingLLM, SnapKV, TOVA, H2O, and K-norm from KVPress. LongBench is restricted to HotpotQA, QASPER, and TriviaQA prompts under 3,000 tokens; RULER uses 130 examples. In these plots the compression ratio r is the eviction fraction: r near 1 means almost the whole prompt is gone.
On average scores, corrected probabilistic eviction matches or slightly beats the heuristics in the low-to-mid compression band. On pairwise win score, the minimum-variance proposal with a harmonic prior, πmin-h, ranks first across both models and the task mix. Deterministic heuristics still crater on individual splits: StreamingLLM leads Llama3 question answering and finishes last on MultiKey-NIAH; K-norm is often top-three on Qwen3 MultiKey-NIAH and last on CommonWords. Lowering bias reduces that kind of task-specific collapse.
The budget is a global soft constraint, so heads spend it according to their own distributions. Lower layers keep more, higher layers evict more, in line with PyramidKV. Temperature only helps when m is tiny and variance explodes; when samples are plentiful, raising τ adds bias and hurts. The paper does not publish a numeric per-split F1 table, and it reports no end-to-end tokens per second.
KV eviction now has a statistical vocabulary. An H2O score can stay in the pipeline as a proposal, then be corrected at decode. A production stack that keeps top-k with no correction has chosen unbounded bias in exchange for determinism.
This is closer to a change of viewpoint than a drop-in kernel. Correction stores sample counts, and per-head variable budgets want a sparse attention implementation.
The authors are explicit: SNIS is still biased and still has variance. Zero bias and zero variance would mean reconstructing the dropped cache from less memory than was removed. Extra state is O(k·h·b) floats, small next to the O(d·k·h·b) cache, but the attention kernel still has to change. Variable-length compression also wants an efficient sparse tensor path; they point at Ada-KV and do not ship a production kernel.
The eval is narrow: 3B/4B models, LongBench cut at 3,000 tokens, a RULER subset. The hardness result protects eviction-time attention, not future decode steps. Wall-clock speed after correction is not reported.