KVpop — Key-Value Cache Compression with Predictive Online Pruning

Summary (Overview)

• KVpop is a novel learned KV cache eviction method that enforces a fixed per-head memory budget in pretrained transformer models by training lightweight importance scorers to predict which tokens will be useful for future attention.
• The method introduces a future-attention supervision target computed via a transposed-attention pass that avoids materializing dense attention maps, and a boundary-aware pairwise loss that trains scorers on the exact keep/drop decision at the cache retention boundary.
• KVpop supports both stateless (MLP-based) and stateful (mLSTM-based) scorers, with a novel delayed scoring mechanism that leverages the protected recent window to incorporate near-future context before making eviction decisions.
• On mathematical reasoning benchmarks (AIME, HMMT), KVpop retains 97–100% of full-attention performance at 88% compression and 98–100% at 75% compression on Qwen3-4B and Qwen3-8B, outperforming prior methods including StreamLLM, TOVA, and DMS.
• KVpop achieves constant memory during inference (19GB vs. 36GB for dense attention at 131k tokens) and faster decoding than DMS due to its homogeneous per-head cache budget.

Introduction and Theoretical Foundation

Background

Transformer-based large language models (LLMs) rely on a key-value (KV) cache during autoregressive decoding, storing key and value representations of all past tokens to avoid recomputing the full sequence history. However, this cache grows linearly with context length, creating a significant memory bottleneck for long-context inference.

The Eviction Problem

KV cache reduction via eviction requires deciding which tokens to permanently remove from the cache. This is fundamentally a prediction problem: which tokens will be useful for future queries? The challenge is that locally salient tokens may become irrelevant later, while tokens receiving little immediate attention may become critical many steps ahead—particularly in reasoning traces where earlier statements are reused after a delay.

Prior Approaches and Their Limitations

Key Insight

KVpop addresses two critical gaps in prior work:

1. Direct supervision of the eviction decision: Instead of differentiable relaxations, train an explicit token-level predictor of future utility, supervised by the attention mass a token receives after it exits the protected window.
2. Delayed scoring: Since importance is only needed at the eviction boundary (when a token leaves the protected window), the scorer can incorporate near-future context before making the keep/drop decision.

Methodology

Fixed-Budget Sparse Attention

Each KV head retains:

• Sink tokens: first $s$ tokens
• Protected window: $w$ most recent tokens
• Long-range top-k: selected from eligible tokens

The per-head KV budget during decoding is:

$$B = s + w + k \tag{1}$$

Future-Attention Supervision Target

For a token $t$ in KV head $h$ with query head group $g$, the future-attention mass is defined as the mean attention probability it receives after leaving the protected window:

$$m_t^{(h,g)} = \frac{1}{N_t} \sum_{d=t+w}^{S-1} p_{d \to t}^{(h,g)}, \quad N_t = \max(1, S - (t+w)) \tag{2}$$

The target score aggregates log-masses across query heads sharing KV head $h$:

$$r_{h,t}^{\mathrm{tgt}} = \mathrm{Agg}_g \left[ \log \left(\epsilon + m_t^{(h,g)}\right) \right] \tag{3}$$

Recency Decay

To prevent high-scoring tokens from occupying cache slots indefinitely, both teacher and student apply a recency decay:

$$r_{h,t}^{\mathrm{tgt}}(q) = r_{h,t}^{\mathrm{tgt}} + \left\lfloor \frac{q - t}{n} \right\rfloor \log \gamma_h \tag{5}$$

where $\gamma_h \in (0,1)$ is a per-head decay factor and $n$ is the decay step size.

Boundary-Aware Retention Loss

The loss focuses on the single comparison that determines cache membership: the newly eligible token $t_{\mathrm{new}} = q - w$ versus the teacher's cutoff token $t_{\mathrm{bnd}}$ (the lowest-scoring token still in the top-k):

$$\mathcal{L}_{\mathrm{score}} = \mathbb{E}_{q,h} \left[ \omega_{q,h} \, \mathrm{softplus} \left( -y_{q,h} \frac{\hat{r}_{h,t_{\mathrm{new}}}(q) - \hat{r}_{h,t_{\mathrm{bnd}}}(q)}{\tau} \right) \right] \tag{7}$$

where $y_{q,h} \in \{+1, -1\}$ indicates whether the teacher retains $t_{\mathrm{new}}$, $\tau$ is a temperature, and $\omega_{q,h}$ weights ambiguous decisions.

Efficient Target Computation via Transposed Attention

The key insight is that the future-attention mass can be computed without materializing the dense $S \times S$ attention map. By swapping queries and keys in a second attention-like call:

$$\log m_t^{(h,g)} = \log \sum_{d=t+w}^{S-1} \exp\left(\ell^{(h,g)}(d,t) - \mathrm{LSE}_d^{(h,g)}\right) - \log N_t \tag{8}$$

The transposed pass reuses the sparse LSE values from the student's forward pass, making target computation efficient and adding no inference overhead.

Running Top-k with Fenwick Tree

For efficient decoding, KVpop sorts tokens once by static priority and uses a Fenwick tree to maintain the running top-k as new tokens become eligible:

• Time complexity: $O(S \log S)$ per head
• Space complexity: $O(S)$ per head
• Implemented as a custom Triton kernel parallelized over batch and KV heads

Scorer Architectures

Stateless (KVpop_mlp): Two-layer MLP with SiLU activation, input is concatenated key and value: $\mathbf{x}_{h,t} = [\mathbf{k}_{h,t}; \mathbf{v}_{h,t}]$

Stateful (KVpop): mLSTM-based scorer with delayed readout:

$$\mathbf{h}_{h,t_{\mathrm{new}}} = \frac{\tilde{\mathbf{q}}_{h,t_{\mathrm{new}}}^\top \mathbf{C}_{h,q}}{\tilde{\mathbf{q}}_{h,t_{\mathrm{new}}}^\top \mathbf{z}_{h,q}} \tag{10}$$ $$\hat{r}_{h,t_{\mathrm{new}}} = \mathbf{a}_h^\top \mathrm{SiLU}(\mathbf{h}_{h,t_{\mathrm{new}}}) + b_h \tag{11}$$

The delayed readout uses the mLSTM state after processing up to position $q$ to score token $t_{\mathrm{new}} = q - w$, incorporating near-future context.

Empirical Validation / Results

Experimental Setup

• Models: Qwen3-4B-Instruct-2507 and Qwen3-8B (thinking mode enabled)
• Training data: Nemotron-Math v2 (high reasoning effort subset), sequence length $S = 16384$
• Budgets: $B \in \{2048, 4096\}$ corresponding to 88% and 75% compression ratios
• Baselines: StreamLLM, TOVA, StreamLLM+ (trained variant), DMS (trained with matched parameter count)

Main Results on Mathematical Reasoning

Table 1: Pass@1 on AIME and HMMT

Out-of-Domain Generalization

Table 2: GPQA-D and LCB results for Qwen3-4B

KVpop remains competitive with the dense teacher despite training only on mathematical reasoning data.

Inference Efficiency (Qwen3-8B)

• Peak VRAM at 131k tokens: Dense = 36GB, DMS = 19GB, KVpop = 19GB (constant memory)
• Latency: KVpop outperforms both dense attention and DMS at long generation lengths due to homogeneous per-head cache budgets that compile more efficiently than DMS's ragged dynamic gates

Ablation: Delayed Scoring

Delayed mLSTM scoring achieves a 0.2-point increase in token accuracy over immediate scoring after 2,000 steps, confirming that aligning eviction decisions with available near-future context is beneficial.

Eviction Pattern Analysis

• KVpop preferentially retains reasoning-structural tokens (discourse markers like "Thus", operation words like "multiplies", symbolic tokens like "=") while dropping purely numeric tokens
• The first layer retains nearly all tokens, consistent with observations that early layers perform broader processing
• Global mean recall of oracle top-k decisions: 81%

Theoretical and Practical Implications

Theoretical Contributions

1. Direct supervision of eviction decisions: KVpop is the first method to train KV cache eviction policies by directly supervising the exact keep/drop decision at the retention boundary, rather than using differentiable relaxations or proxy objectives.

2. Future-attention as supervision signal: The target based on future attention mass provides a theoretically grounded measure of token utility that captures delayed relevance, which proxy scores (recency, cumulative attention) cannot.

3. Efficient target computation: The transposed-attention trick shows that column-wise reductions over attention matrices can be computed as row-wise reductions in a transposed pass, reusing kernel LSE outputs—a potentially general technique for efficient attention-based supervision.

4. Delayed scoring principle: The protected window enables a principled form of delayed decision-making where the scorer can integrate near-future context before committing to an eviction decision.

Practical Implications

1. Memory efficiency: KVpop enables aggressive KV cache compression (75–88%) with minimal quality loss, directly addressing the memory bottleneck in long-context inference.

2. Inference speed: The homogeneous per-head budget enables efficient GPU execution, outperforming methods with dynamic per-head allocations (like DMS) in latency.

3. Domain transfer: The learned eviction policy generalizes beyond training data distribution to code generation and STEM reasoning, suggesting the learned importance signals capture generalizable properties of reasoning traces.

4. Architecture flexibility: The scorer design is modular—any module assigning scalar importance scores can be supervised with the future-attention target, enabling future exploration of alternative memory-based scorers.

Conclusion

KVpop introduces a fixed-budget KV cache compression method that learns predictive online eviction 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.

Key results: KVpop retains 95–100% of dense-attention performance at 75–88% compression on mathematical reasoning benchmarks, outperforming both training-free and trained baselines, while maintaining constant inference memory and competitive latency.

Future directions:

• Exploring alternative memory-based scorers beyond mLSTM
• Hybrid dense-sparse layer configurations for improved quality-efficiency trade-offs
• Evaluation under paged KV-cache managers (vLLM, SGLang)
• Application to models trained from scratch with sparse attention patterns