KVpop — Key-Value Cache Compression with Predictive Online Pruning

Lukas Hauzenberger, Niklas Schmidinger, Anamaria-Roberta Hartl, David Stap, Thomas Schmied, Sebastian Böck, Günter Klambauer, Sepp Hochreiter

KVpop learns to prune KV cache tokens by supervising eviction decisions with future-attention mass.

How can we compress the KV cache in LLMs by predicting which tokens will be useful for future attention, rather than relying on static heuristics?

Large language models suffer from linearly growing memory costs as context length increases, but existing eviction methods rely on static heuristics or proxy scores that fail to predict which tokens will be useful for future reasoning. KVpop introduces a learned eviction policy that supervises token retention by regressing against the actual attention mass a token receives after it leaves a protected recent window. This target is computed efficiently during training using a transposed-attention pass that avoids materializing dense attention maps. On mathematical reasoning benchmarks, KVpop retains 98% of full-attention performance on Qwen3-4B at 75% cache compression, consistently outperforming heuristic and learned baselines.

Paper Primer

KVpop treats KV cache reduction as a supervised prediction problem at the eviction boundary. It uses lightweight head-wise scorers to rank tokens, training them with a boundary-aware loss that pushes the student's retention decisions to match a teacher's future-attention signal.

The core mechanism is a stateful scorer that defers the keep-or-drop decision until a token exits the protected window. By delaying the score, the model integrates near-future context into the eviction decision, effectively acting like a mail sorter that waits for the full address to appear before deciding which bin to route an envelope into.

KVpop maintains high reasoning performance under aggressive cache compression.

Pass@1 on AIME and HMMT benchmarks for Qwen3-4B and Qwen3-8B. 98% of dense-attention performance at 75% compression; 97% at 88% compression.

Delayed stateful scoring improves token retention accuracy.

Ablation comparing mLSTM scorers with and without delayed readout. 0.2-point increase in token accuracy over immediate scoring.

Why is this approach better than simply using attention scores to evict tokens?

Attention scores are local proxies that often fail to capture long-range utility, especially in reasoning tasks where early tokens are reused after a delay. KVpop directly supervises the eviction decision using the actual future-attention mass, ensuring the policy is optimized for long-term relevance.

Does this method require changing the base model architecture?

No, KVpop is a sparse-attention retrofit. It adds lightweight, head-wise scoring modules that are trained during a distillation phase, leaving the base model architecture and inference-time execution flow unchanged.

The KV Cache Bottleneck

We expose why KV cache growth limits long‑context decoding and preview KVpop’s predictive eviction.

Key‑Value (KV) cache growth scales linearly with the number of generated tokens, turning memory and bandwidth into a hard limit for long‑context inference. Existing eviction heuristics are static or rely on proxy scores, which cannot reliably predict a token’s future usefulness and therefore cause brittle pruning as relevance shifts.

The KV cache stores a key and a value for every token that has already been generated, so its size grows one‑to‑one with context length, making memory and data movement increase proportionally.

KVpop addresses this bottleneck by learning a predictive scoring module that decides which tokens to keep or drop at the eviction boundary, supervised directly by a future‑attention target. The target measures the attention mass a token receives after it leaves the protected window, and it is obtained with a transposed‑attention pass that avoids building a dense $S\times S$ attention matrix. Additionally, KVpop can employ a delayed memory‑based scorer that postpones the scoring decision for a fixed number of steps, allowing the module to incorporate near‑future context before eviction.

Linear scaling of memory with context length is the primary obstacle to efficient long‑context decoding.

KVpop System Overview

Related work categorizes prior KV‑cache reduction approaches and positions KVpop among them.

Prior work on reducing KV‑cache size falls into three families: sparse retrieval that keeps the full cache but limits which keys are attended, heuristic eviction that enforces a hard cache bound with fixed rules, and learned eviction that trains a predictor to decide which tokens to drop.

**Figure 1.** **Overview of KVpop.** (1) **Training and inference:** In each attention layer, keys and values are passed to a lightweight per-KV-head scoring module. The cache always retains sink tokens and a protected recent window, and uses predicted scores to select the remaining long-range top-k entries. Attention is computed over the retained set. (2) **Training only:** Future-attention targets are computed on the fly and used to supervise the eviction decision induced by the fixed budget. **Sparse Attention Pattern:** KVPOP yields a sparse attention pattern for efficient decoding, maintaining constant memory per head during inference.

Unlike sparse retrieval, KVpop permanently evicts low‑utility tokens, yielding a bounded KV cache while preserving the ability to predict future attention utility via the learned scorer.

The KVpop Scoring Module

KVpop learns to rank tokens so a fixed per‑head budget can be enforced without hurting future attention.

The KV cache grows linearly with sequence length, turning autoregressive decoding into a memory bottleneck. KVpop solves this by learning a lightweight scorer that decides which tokens to keep so that each head respects a fixed budget $B$.

The scorer predicts each token’s future‑attention utility and ranks tokens so that only the most useful ones survive the budget constraint.

Sink tokens $t_1,t_2$ are kept automatically (2 tokens).

Recent window keeps the latest $w=3$ tokens: $t_8,t_9,t_{10}$ (3 tokens).

Apply decay (factor $0.9^{\text{age}}$) to remaining scores, yielding decayed scores $[0.18,0.45,0.09,0.72,0.27,0.54,0.36]$ for $t_3\ldots t_7$.

Select top‑$k=4$ decayed scores: $t_4$ (0.72), $t_6$ (0.54), $t_5$ (0.27), $t_7$ (0.36) → keep $t_4,t_6,t_7,t_5$.

Mask $M$ now retains $t_1,t_2,t_4,t_5,t_6,t_7,t_8,t_9,t_{10}$ — exactly $9$ tokens, satisfying $B$.

Decay ensures that older high‑scoring tokens lose priority, so the budget always favors recent relevance while still allowing a few long‑range tokens to survive.

Training step of KVpop (Algorithm 1)

How does the KVpop scorer differ from a naïve magnitude‑based pruning that simply keeps the largest absolute values?

Magnitude pruning ignores temporal relevance; KVpop first decays scores so older tokens lose weight, then selects top‑$k$ after decay. This prevents stale high‑magnitude tokens from occupying the budget, whereas naïve pruning would keep them regardless of when they appear.

Defining Future Attention Utility

Define a supervision target that measures a token’s future attention utility and use it to train the scorer.

The KV cache grows linearly during autoregressive decoding, so we need a target that tells the scorer which tokens will be useful later, not just which ones look important now.

Future attention mass is like the interest a savings account earns: a token that will be consulted many times in the future accrues higher “interest”, and the model learns to predict that interest before the token leaves the protected window.

How does this future‑attention target differ from the usual attention‑score supervision that looks at current attention?

Standard supervision treats the attention weight at the current step as the signal, which only reflects immediate relevance. The future‑attention target instead measures how much a token will be attended to after it leaves the protected window, providing a forward‑looking utility estimate that guides eviction decisions.

Compute $N_t = \max(1,\,6-(2+2)) = 2$ future positions.

Group $g=0$: sum $0.30+0.20=0.50$, average $m(0,0)_2 = 0.50/2 = 0.25$.

Group $g=1$: sum $0.10+0.10=0.20$, average $m(0,1)_2 = 0.20/2 = 0.10$.

Apply $\log(\epsilon+m)$ (with $\epsilon\approx0$): $\log(0.25) \approx -1.386$, $\log(0.10) \approx -2.303$.

Aggregate with max: $r^{\text{tgt}}_{0,2}= \max(-1.386,\,-2.303) = -1.386$.

The target assigns a higher score to tokens that receive more future attention, even if their immediate attention is low, which directly informs which tokens the cache should retain.

**Figure 2.** Transposed Attention. (1) Blue tiles show the future attention mass for token $t$. (2) Transposing $q$ and $k$ turns the per-key column-sum into a per-query row-sum, which attention kernels return as their auxiliary LSE.

Delayed Memory Scoring

We implement delayed scoring and efficient top‑k pruning to keep KV cache size low.

The KV cache grows linearly with sequence length, turning it into a decoding bottleneck. To avoid materializing full $S\times S$ attention matrices, we compute a cheap proxy target and prune tokens on the fly.

Instead of scoring a token the moment it enters the cache, we wait until it exits the protected sliding window and then assign its importance based on how much future attention it would have received.

For query $d=2$, the allowed keys are $t'=1$ and $t'=2$ (mask $M_{2,1}=M_{2,2}=1$).

Compute logits $\ell(2,1)$ and $\ell(2,2)$, exponentiate, and sum: $\exp(\ell(2,1))+\exp(\ell(2,2))$.

Take $\log$ to obtain $\text{LSE}(2)$.

When token $t=1$ exits the window at $q=2$, use $\text{LSE}(2)$ in Eq. (8) to compute its delayed score.

The resulting score ranks token 1 higher than token 0, so token 1 is retained while token 0 may be evicted.

Delayed scoring lets us reuse the already‑computed $\text{LSE}$ values, turning a potentially $O(S^2)$ operation into $O(S)$ per head.

How does delayed scoring differ from the usual immediate scoring of KV entries?

Immediate scoring evaluates importance as soon as a token is written to the cache, which cannot foresee how future queries will use it. Delayed scoring postpones evaluation until the token has left the protected window, allowing the model to incorporate actual future‑attention contributions (the Future Attention Mass) into the importance score.

Reuse the sparse log‑normalizers $\text{LSE}_{h,g}(d)$ from the student forward pass.

Swap queries and keys: treat original key positions $t$ as queries and original query positions $d$ as keys.

Run a second attention call (implemented with FlexAttention) with a block mask enforcing $d \ge t + w$.

Obtain the column‑wise LSE term of Eq. (8) for all keys in parallel.

Subtract $\log N_t$ and add $\varepsilon$ to form $\log m_{h,g}(t)$.

Aggregate across heads sharing KV head $h$ to produce the final target $r_{tgt}^{h,t}$.

Pre‑sort all tokens by a static priority score and store their ranks.

Maintain a Fenwick tree over the ranks of tokens that have become eligible (i.e., have exited the window).

At each query $q$, insert the newly eligible token $t = q - w$ into the Fenwick tree.

Query the tree for the rank $\tau_q$ of the $k$‑th best eligible token using binary lifting.

Form the predicate “rank $\le \tau_q$” and feed it to FlexAttention, which builds the sparse mask on the fly.

The attention kernel then attends to the union of sink tokens, the protected window, and the top‑$k$ eligible tokens.

Ranks after sorting: token 3 (rank 1), token 0 (rank 2), token 2 (rank 3), token 5 (rank 4), token 4 (rank 5), token 1 (rank 6).

At query $q=3$, token $t=1$ becomes eligible and is inserted (rank 6) into the Fenwick tree.

Binary lifting finds $\tau_3$ = rank 3 (the $k$‑th best eligible token).

Tokens with rank $\le 3$ (tokens 3, 0, 2) are kept; the mask includes them plus the protected window tokens $2$ and $3$.

FlexAttention builds the sparse mask and computes attention only over this reduced set.

Because ranks are static, the Fenwick tree lets us update the top‑$k$ set in $O(\log S)$ time instead of recomputing a full sort at every query.

Fenwick‑tree insertion and $k$‑th rank query used in the top‑$k$ pruning step.

**Figure 3.** KVpop stateful eviction policy. (1) The mLSTM-memory is updated with the most recent KV pair. (2) For delayed scoring, the memory is read with the token at position $q - w$ that has just exited the protected sliding window. (3) The mLSTM emits an importance score, and ranking determines whether the token is kept or evicted.

Efficient Online Pruning

We introduce Fenwick‑tree pruning with stateless and stateful scorers to score KV tokens online.

KV‑cache growth scales linearly with generated length, quickly exhausting GPU memory during autoregressive decoding. A naïve eviction policy discards tokens arbitrarily, often removing useful context. Our solution is an online importance‑scoring pipeline that decides which tokens to keep as they become eligible.

The Fenwick tree lets us maintain a running count of tokens whose static rank is high enough, so for each query we can instantly retrieve the rank threshold that yields the top‑$k$ eligible tokens.

Initialize a Fenwick tree of length $5$ with all zeros.

Insert token 0 (rank 3) → tree counts become $[0,0,1,0,0]$.

Insert token 1 (rank 1) → counts $[1,0,1,0,0]$.

Prefix sum reaches $k=2$ at rank 3, so $\tau_q=3$; tokens with rank ≤ 3 (tokens 0,1,3) are eligible.

When token 2 (rank 4) arrives, prefix sum at rank 4 becomes $3$, still $\tau_q=3$ because the top‑$k$ threshold does not change.

The tree instantly adapts the cutoff as new tokens become eligible, without scanning the entire cache.

How does this differ from a static “evict oldest” heuristic?

Oldest‑first eviction ignores the pre‑computed relevance rank and can discard a token that would later receive high attention. The Fenwick‑tree approach keeps tokens based on their predicted future utility, guaranteeing that the $k$ most promising tokens survive regardless of insertion order.

A tiny head‑wise MLP predicts a token’s importance using only its own key and value vectors, providing a fast, memory‑free score at insertion time.

Why detach the scorer inputs from the main graph?

Detachment isolates the scorer’s learning signal, so the retention loss updates only $f_h$ while the backbone continues to be trained solely by the distillation loss. This avoids unwanted gradient coupling that could destabilize the primary task.

The scorer maintains a recurrent memory that aggregates context up to the moment a token leaves the protected window, allowing it to evaluate importance with richer, near‑future information.

Compute bilinear interaction $h_{h,2}= \frac{[0.2,0.4]\cdot[0.5,0.3]}{1.0}=0.2\cdot0.5+0.4\cdot0.3=0.10+0.12=0.22$.

Apply SiLU: $\text{SiLU}(0.22)=0.22/(1+e^{-0.22})\approx0.20$.

Project: $\hat{r}_{h,2}= a^{\top}0.20 + b_h$; with $a=0.8$, $b_h=0.05$ gives $\hat{r}=0.8\cdot0.20+0.05=0.21$.

Since $k=1$, the token with the highest $\hat{r}$ among eligible tokens is kept; here $\hat{r}=0.21$ is the sole candidate, so it survives.

The stateful scorer can raise a token’s score after seeing subsequent context, whereas a stateless scorer would have fixed its score at insertion.

In what way does this differ from Delayed Memory Scoring (DMS)?

DMS postpones the eviction decision but still scores tokens based on their *initial* representation. Our stateful scorer recomputes the score after the token has observed $w$ future tokens, using the mLSTM’s accumulated context, thus capturing information that DMS cannot.

**Algorithm 2:** Running top-$k$ cutoff computation for one KV head. A Fenwick tree maintains eligible long-range tokens ordered by their static rank. For each query position, the cutoff $\tau_q$ is the rank threshold corresponding to the current top-$k$ eligible tokens. ``` Require: static ranks rank(0),...,rank(S - 1) budget (s,w,k) 1: initialize Fenwick tree T with zeros 2: τ_q ← -1 for all q 3: for q = 0,...,S - 1 do 4: t ← q - w 5: if t ≥ s then 6: insert rank(t) into T 7: end if 8: if |E(q)| ≥ k then 9: τ_q ← FindByPrefixCount(T, k) 10: end if 11: end for 12: return τ_0,...,τ_{S-1} ```

Experimental Setup

KVpop sets new state-of-the-art Pass@1 scores on Qwen models under compression.

KVpop attains the highest Pass@1 scores on both Qwen3‑4B and Qwen3‑8B at the 75 % compression ratio.

Table 1 shows KVpop reaching 0.62 absolute and 0.44 teacher‑relative scores, surpassing all other variants.

**Table 1.** Pass@1 on AIME and HMMT shown as absolute scores (Abs) and w.r.t. teachers (Rel)

These are instruction‑tuned language models with 4 B and 8 B parameters that serve as the backbone for evaluating KVpop.

Generalization and Robustness

KVpop retains accuracy on out‑of‑domain tasks despite heavy compression.

KVpop matches the dense teacher within 1% on both GPQA‑D and LCB benchmarks despite being trained only on mathematical reasoning data.

Table 2

**Table 2.** Results for QWEN3-4B on GPQA-D and LCB.

**Figure 7.** Alignment with the future-attention oracle. We plot the top-k retention recall distribution over heads for all layers of QWEN3-4B at 75% KV-cache compression. The metric measures what fraction of the full-attention teacher's top-k retained tokens are also kept by the learned KVPOPscorer under the same cache budget. Higher values indicate better agreement with the teacher retention policy.

Inference Performance Analysis

We revisit the KV cache bottleneck and show how KVpop reshapes inference efficiency.

The KV cache growth limits autoregressive decoding speed; KVpop replaces static eviction heuristics with a learned scorer that predicts future attention utility.

KVpop limits peak VRAM to 19 GB at 131 k tokens, a 47 % reduction versus dense attention.

Dense attention peaks at 36 GB, while both DMS and KVpop grow only 19 % to 19 GB under the same generation length.

**Figure 4.** Inference efficiency on Qwen3-8B. We measure end-to-end decoding latency and peak allocated VRAM at batch size 1, 75% KV compression, and varying generation lengths. Dense attention incurs growing memory cost as the KV cache expands, whereas DMS and KVPOP grow more slowly. KVPOP further reduces latency, maintaining stable throughput across long generations and outperforming both dense attention and DMS at the longest sequence lengths.

KVpop decodes faster than DMS because it enforces a fixed budget per KV head, whereas DMS’s dynamic gates create ragged per‑head caches that are harder for hardware to compile and execute.

**Figure 5.** Effect of delayed scoring. Token accuracy with and without delayed scoring.

Eviction‑policy analysis reveals KVpop preferentially evicts purely numeric tokens while preserving reasoning‑structural tokens such as discourse markers and operators; the first layer is an exception, retaining nearly all tokens.

**Figure 6.** Eviction patterns in a mathematical reasoning trace. We visualize the last 112 tokens of a randomly sampled sequence. Rows denote layers, columns denote token strings, and color indicates how many attention heads in a layer retained the corresponding token. KVPOP tends to evict purely numeric tokens more often, while retaining reasoning-structural tokens.

Alternative Aggregation Methods

Mathematical details, implementation tricks, and experimental settings for KVpop and DMS.

We first describe an alternative probability‑space mean aggregation for the future‑attention target.

Count normalization removes the bias that early tokens see more future queries.

Temporal decay introduces a per‑head exponential decay factor $\gamma_{h}$ that biases older tokens toward lower scores.

Margin weighting down‑weights boundary examples whose target scores are nearly tied.

Keep/Drop balancing equalizes the aggregate weight of positive and negative boundary decisions.

We sample only a small set of query positions after the long‑range budget is saturated.

Running top‑k sparse attention during training and pre‑fill requires a fixed‑budget mask.

Fenwick‑tree cut‑off computation (Algorithm 2)

Future‑attention targets can be obtained by a transposed attention pass.

Run the main attention pass and record per‑query log‑normalizers (dense or sparse).

Swap queries and keys and run a second attention pass.

Subtract the original log‑normalizer from each logit (Eq. 33) and mask out positions with $d < t+w$.

Read the auxiliary log‑sum‑exp output to obtain $\log \bar{m}_{h,g}(t)$ for all keys.

Aggregate across groups (Eq. 15) and apply count normalization (Eq. 16) to get the final target.

Scorer modules consume cached keys and values to predict eviction scores.

Training uses mixed‑precision on 8 H100 GPUs, gradient clipping, and a consistent decoding configuration.

**Table 3.** Training and Inference Hyperparameters.

**Figure 8.** DMS evict ratio. We compute the median eviction ratio over 100 randomly sampled sequences for QWEN3-4B with DMS at a 75% compression ratio. A value of 0 indicates dense/full attention, whereas a value of 1 indicates that attention is restricted to the local sliding window. DMS exhibits highly heterogeneous eviction patterns: some heads use the full long-range token budget and behave close to dense-attention heads, while many heads—especially in early and late layers—collapse to sliding-window-only attention.

**Figure 26.** Visualization of attention weights across five layers of a transformer model, with tokens on the x-axis and attention heads on the y-axis.

KVpop’s learned sparsity retains contiguous semantic blocks rather than a uniform pattern, especially in later layers where symbolic tokens are preserved.

DMS exhibits heterogeneous eviction: a few heads keep most of the long‑range budget while many heads fall back to pure sliding‑window attention, contrasting with KVpop’s uniform per‑head allocation.

Training and Inference Details

We evaluate KVpop’s inference efficiency against baselines across two cache budgets.

We compare KVpop to both training‑free eviction methods (StreamingLLM and TOVA) and learned sparsification (DMS), as well as a trained variant of StreamingLLM (StreamingLLM+) that matches KVpop’s budget. All baselines share the same cache budget $B$ ∈ {2048, 4096}, corresponding to compression ratios of 88 % and 75 %. For the hybrid reasoning model Qwen3‑8B we enable thinking mode to reflect its intended use.

At 88 % compression KVpop retains 97 % of Qwen3‑4B’s performance and 100 % of Qwen3‑8B’s, while the stateless KVpopmlp reaches 93 % and 98 % respectively. Training‑free StreamingLLM and TOVA collapse under this budget, and StreamingLLM+ recovers only partially. DMS, which uses a Gumbel‑sigmoid policy, trails KVpop by 8–16 points, and the ordering of methods remains consistent at 75 % compression, where all improve but KVpop stays ahead.

Supplementary Mathematical Analysis

Appendix summarizing conclusions, limitations, acknowledgments, and references.

While mLSTM scoring provides a strong stateful KV pruning mechanism, we leave a broader exploration of alternative memory‑based scorers for future work. KVpop is designed as an efficient post‑training retrofit for Transformers with dense attention, rather than a compressed‑cache architecture trained from scratch. Finally, our homogeneous per‑head cache budget enables efficient GPU execution, but hybrid dense‑sparse layers may further improve the quality‑efficiency trade‑off.

We introduced KVpop, a fixed‑budget KV cache compression method that learns predictive online decisions from future‑attention supervision. By training scorers at the eviction boundary and optionally delaying stateful scoring until near‑future context is available, KVpop directly targets the token‑retention decision that determines cache membership. The resulting sparse retrofit bounds inference memory and improves the quality‑efficiency trade‑off over prior eviction methods.

Across AIME and HMMT, KVpop retains 95 % of dense‑attention performance on Qwen3‑4B at 75 % KV cache compression and 94 % at 88 % compression. On Qwen3‑8B, retention reaches 95 % and 99 % at the same compression ratios. Moreover, despite being distilled on mathematical reasoning data, the learned eviction policy remains competitive on code generation and STEM reasoning benchmarks.

We acknowledge EuroHPC Joint Undertaking for awarding us access to Leonardo at CINECA, Italy, Deucalion at MACC, Portugal, Discoverer at SofiaTech, Bulgaria. This work was supported by the European Union’s Horizon Europe research and innovation programme under grant agreement number 101214398 (ELLIOT).

References are listed below, covering foundational works on LSTM, attention mechanisms, linear attention, and recent advances in KV cache compression.

Background: Linear Attention

Background on linear attention and the mLSTM scorer.

Linear attention replaces the softmax similarity $\kappa_{\exp}(q, k)=\frac{\exp(q^{\top}k)}{\sqrt{d_{qk}}}$ with a kernel $\kappa_{\phi}(q, k)=\phi(q)^{\top}\phi(k)$ that admits an explicit feature map $\phi:\mathbb{R}^{d_{qk}}\to\mathbb{R}^{d_{qk}}$ (Katharopoulos et al., 2020). Via associativity this yields two equivalent causal attention implementations—a parallel chunkwise computation for training and prefill, and an online recurrent update where each head maintains a running key–value summary $C_t\in\mathbb{R}^{d_{qk}\times d_v}$ (and optional normalizer $z_t\in\mathbb{R}^{d_{qk}}$), updated by rank‑one outer products $C_t=C_{t-1}+\phi(k_t)\otimes v_t$, $z_t=z_{t-1}+\phi(k_t)$, and producing output $h_t=\frac{\phi(q_t) C_t}{\phi(q_t) z_t}$. The mLSTM, introduced as a sublayer of xLSTM, serves as a stateful importance scorer by augmenting the outer‑product update with expressive gates, and has demonstrated strong performance across language modelling, vision, and biology.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers