Learned KV-cache eviction

Context. To decode without recomputing the past, a transformer keeps the key and value of every earlier token, the KV cachethe stored keys and values of all earlier tokens, per layer and per KV head; it grows by one entry per head with every generated token, and reads all of it at every step, so memory and the time per token grow with the length of the output. Reasoning models now write traces of tens of thousands of tokens, which makes the cache the main cost of decoding. Evictionkeeping a fixed number of cached tokens per head and permanently deleting the rest, which bounds memory; sparse retrieval methods instead keep the full cache and read only part of it keeps a fixed number of entries and deletes the rest, and the question is which to keep. Fixed rules keep the first few tokens and a recent window (StreamingLLMan eviction rule that keeps the first few tokens (attention sinks) and a window of the most recent tokens, and drops everything in between (Xiao et al., 2023)) or the tokens with the most attention right now (TOVAan eviction rule that drops, once the cache is full, the token with the lowest attention from the current query (Oren et al., 2024)), and these are guesses about the future from the past. The learned method DMSdynamic memory sparsification: binary keep-or-evict gates per token and head, trained through a continuous relaxation, with eviction delayed by a short window (Łańcucki et al., 2025) trains eviction gates through a continuous relaxation but is given no target for which tokens to keep. KVpop (Hauzenberger et al., 2026) supervises that choice directly.

Which cached tokens a model will need later is largely predictable, and there is a label for it in the model's own attention: the attention each token receives from the queries that come after it leaves the recent window. A small scorer per KV head, trained on that label, keeps Qwen3-4B at 97% and Qwen3-8B at 100% of their accuracy on competition math with a cache of 2,048 tokens per head, 12.5% of the 16,384-token training length, where the fixed rules keep between 19% and 58%. The scorer can also wait to score a token until it leaves the window, and so see what followed it. The model itself is fine-tuned along with the scorer, and that training alone, with a fixed sparse pattern, already recovers part of the gap; the learned choice of tokens accounts for the rest.

Setting. Each KV headone set of keys and values; under grouped-query attention several query heads share it, and the cache is stored per KV head $h$ keeps $s=4$ sink tokensthe first tokens of the sequence, kept permanently because attention heads put much of their weight on them, a protected window of the $w=256$ most recent tokens and $k$ long-range slots, a budget $B=s+w+k$ of 2,048 or 4,096 entries. Positions are numbered from 0, so the sinks are positions 0 to $s-1$. At query position $q$ the tokens $E(q)=\{t:\ s\le t\le q-w\}$ compete for the $k$ slots, and exactly one token, $t_{\rm new}=q-w$, joins them. The label is the mean attention a token receives once it has left the window. With $p_{u\to t}$the dense causal attention probability from query position $u$ to key $t$, for one query head $g$ of the group sharing KV head $h$ and $S$ the training sequence length,$$\begin{aligned}m_t&=\frac{1}{N_t}\sum_{u=t+w}^{S-1}p_{u\to t},\\N_t&=\max(1,S-t-w),\\r^{\rm tgt}_{h,t}&=\max_{g}\,\log\big(\epsilon+m^{(h,g)}_t\big),\end{aligned}$$where the maximum over the query heads sharing the KV entry keeps a token that any one of them relies on, and $\epsilon$ is a small constant. Both the target ranking and the scorer's ranking subtract the token's age times a learned per-head rate, adding $(q-t)\log\gamma_h$ with $\gamma_h$a learned decay per KV head, constrained to $[0.999,\,0.999999]$; at 0.999 a token's score falls by 1 per 1,000 tokens of age, so a token with a high score cannot hold its slot forever. The teacher (the original model with its full cache, whose attention gives the target) keeps the $k$ best tokens of $E(q)$ under this ranking; let $t_{\rm bnd}$ be the lowest-ranked token it keeps and $y_{q,h}=+1$ if it keeps $t_{\rm new}$, $-1$ if not. The scorer's predicted scores $\hat r$, age-adjusted the same way to $\hat r_{h,t}(q)=\hat r_{h,t}+(q-t)\log\gamma_h$, are trained with a logistic loss on that single pair:$$\begin{aligned}L_{\rm score}=\mathbb E_{q,h}\Big[&\,\omega_{q,h}\,\mathrm{softplus}\Big(\\&-y_{q,h}\,\Delta_{q,h}/\tau\Big)\Big],\\\Delta_{q,h}&=\hat r_{h,t_{\rm new}}(q)-\hat r_{h,t_{\rm bnd}}(q).\end{aligned}$$Here $\tau$ is a temperature, $\omega_{q,h}$ a weight that discounts near-ties and balances keep and drop labels, and softplus$\mathrm{softplus}(x)=\log(1+e^x)$, a smooth version of $\max(x,0)$ makes the loss small once the decision is right by a margin. Comparing with the cutoff token trains exactly the decision that changes what is in the cache.

Computing the target, and running the cache. The target needs, for each key, a sum over all later queries, a column sum of the $S\times S$ attention matrix, which an attention kernel never forms. Write $\ell(u,t)$ for the attention logit$\ell(u,t)=q_u^\top k_t/\sqrt{d_k}$, the score before the softmax, with $d_k$ the key dimension and $\mathrm{LSE}_u=\log\sum_{t'\le u}e^{\ell(u,t')}$ for the softmax normalizer of query $u$, which FlexAttentionPyTorch's compiled path for attention variants written as a score modifier and a mask, which also returns each row's log normalizer returns with its output. Then$$\begin{aligned}&\log\sum_{u\ge t+w}p_{u\to t}\\&\quad=\log\sum_{u\ge t+w}e^{\ell(u,t)-\mathrm{LSE}_u},\end{aligned}$$a log-sum-exp over $u$ for fixed $t$, which is what an attention call returns for each query once queries and keys swap roles: key $t$ becomes the query, the score modifier subtracts $\mathrm{LSE}_u$ and the mask keeps $u\ge t+w$. One extra attention-like pass during training gives every target, and the paper uses the normalizers of the sparse attention it already computes in place of the dense ones, which it finds works as well. Selection is cheap too: with a fixed decay rate the ranking at every query equals a ranking by one static priority per token, so each new eligible token is inserted into a Fenwick treea binary indexed tree: an array that supports adding a count and finding where a running count reaches a value, each in logarithmic time of ranks and the $k$-th best is read off, $O(S\log S)$ per head, and the resulting mask is evaluated inside the attention kernel. The scorer is either a two-layer MLP on the token's key and value (KVpop$_{\rm mlp}$) or an mLSTMa gated linear-attention cell from xLSTM: a fixed-size matrix memory that each token updates with a gated outer product of its key and value features, read with a query vector per KV head that has read the keys and values up to $q$ when it scores $t_{\rm new}=q-w$, so the decision uses the 256 tokens that followed the token. Scorer inputs are detached: the scorer learns only from $L_{\rm score}$ and the model only from a KL divergence to the dense teacher's top 256 logits.

In practice KVpop is a retrofit: about 2,000 steps of full fine-tuning (2B tokens on 8 H100s in the paper) of an existing model on long reasoning traces, after which every head carries the same fixed-size cache. For Qwen3 models on competition math, a 4,096-entry cache per head lost at most 2% of accuracy and a 2,048-entry cache cost 0 to 3% of accuracy, with peak memory flat at about 19 GB while dense decoding reached 36 GB at 131K tokens. The uniform per-head budget matters for speed: DMS lets a few heads keep most of the long-range budget, which gives ragged caches that are slow to execute, and KVpop decoded 131K tokens in about 40% of DMS's time (read off a plot). The authors have not yet tested it in a paged serving system such as vLLM, so the speed figures come from a research implementation at batch size 1.