Papers
Topics
Authors
Recent
Search
2000 character limit reached

KVpop -- Key-Value Cache Compression with Predictive Online Pruning

Published 6 Jul 2026 in cs.LG and cs.AI | (2607.05061v2)

Abstract: Key-value (KV) cache growth is a major bottleneck in autoregressive decoding, as memory and bandwidth scale linearly with context length. Existing KV eviction methods often rely on static heuristics or proxy scores, which poorly track future token utility and cause brittle eviction as relevance shifts. To address this, we introduce KVpop, which learns a fixed-budget KV eviction policy by directly supervising the keep-or-drop decision. The scorer is trained against a novel future-attention target, computed efficiently without materializing dense attention maps. We further introduce a delayed memory-based scorer that, uniquely among learned eviction methods, defers scoring for a fixed number of steps to exploit near-future context. On AIME and HMMT mathematical reasoning, KVpop retains 98% of full-attention performance on Qwen3-4B at 75% KV cache compression and 97% at 88% compression, consistently outperforming established eviction baselines. Qwen3-8B shows even stronger results, reaching near-full teacher performance. These results show that supervising eviction with future-attention signals cuts memory costs while maintaining quality.

Summary

  • The paper introduces KVpop, which uses future-attention supervision and delayed mLSTM scoring to accurately prune transformer KV caches.
  • It enforces a fixed per-head cache budget, achieving up to 88% compression while retaining near-teacher performance on benchmarks.
  • KVpop outperforms heuristic methods by prioritizing essential structural tokens and delivering efficient, predictable GPU throughput.

Predictive Online Pruning of Key-Value Caches in Long-Context LLM Decoding

Motivation and Problem Statement

Autoregressive decoding in transformer-based LLMs depends crucially on maintaining a key-value (KV) cache, enabling efficient retrieval of past context. However, as context length increases, the KV cache grows linearly, becoming a primary bottleneck for both memory usage and bandwidth during inference. Traditional approaches to cache reduction, such as sliding-window mechanisms or static sink-token retention, remove tokens based on fixed heuristics, frequently failing to identify the future utility of tokens as relevance shifts dynamically. Score-based and learned policies have emerged, with recent methods like Dynamic Memory Compression (DMC) and Dynamic Memory Sparsification (DMS) retrofitting sparse attention patterns via auxiliary scoring and gating. Nevertheless, none directly supervise the eviction decisions with respect to future token utility, nor intelligently exploit context that accumulates after a token enters the cache.

KVpop: Predictive Online Pruning Framework

The paper presents KVpop, a sparse-attention retrofit that enforces a fixed per-head KV cache budget without altering the base model architecture. KVpop augments each KV head with:

  • Sink token retention (typically the initial context),
  • A protected window of recent tokens immune to eviction,
  • A long-range top-kk retention policy, governed by lightweight importance scoring modules.

These scorers predict token utility at the eviction boundary, trained via a boundary-aware loss against a novel future-attention supervision target. This target efficiently estimates the future attention mass each token receives, computed by a transposed attention pass using log-sum-exp normalizers, obviating the need for full dense attention computation (Figure 1). Figure 1

Figure 1: Overview of KVpop, illustrating token scoring, cache management, and sparse attention patterns for efficient decoding.

Furthermore, KVpop supports stateful importance scoring via mLSTM modules, uniquely enabling delayed scoring at the eviction boundary. This allows the scorer to integrate near-future evidence observed while the token remains in the protected window, improving eviction policy fidelity compared to prior approaches that commit at insertion time (Figure 2). Figure 2

Figure 2: mLSTM-based stateful eviction policy with delayed scoring, exploiting accumulated context for retention decisions.

Supervision Targets and Scoring Mechanisms

Attention supervision is formulated as the mean future-attention mass received by a token post-protection window, aggregated using max or mean across grouped-query heads. Temporal decay is applied to scores, creating a recency bias so that cache slots are not monopolized by historically high-utility tokens. At each eviction boundary, the scoring module is trained to distinguish between the newly eligible token and the teacher’s cutoff token, with a pairwise logistic loss. For efficient mask generation during training and prefill, KVpop leverages a Fenwick tree for running top-kk threshold computation and integrates with sparse attention kernels such as FlexAttention.

Experimental Results: Quality-Efficiency Trade-offs

KVpop was distilled onto Qwen3-4B and Qwen3-8B models, evaluated on mathematical reasoning benchmarks (AIME24/25, HMMT) at context lengths up to 16,38416{,}384 tokens. Strong quantitative results were obtained:

  • At 75% KV cache compression, KVpop retained 98% of teacher performance for Qwen3-4B and reached near-teacher accuracy for Qwen3-8B.
  • At the more aggressive 88% compression, retention remained at 97% and 100%, respectively.
  • KVpop consistently outperformed heuristic baselines (Sliding Window, TOVA) and DMS under matched memory budgets.
  • Retention policy generalized robustly to GPQA-D and LiveCodeBench (code generation and STEM reasoning), with negligible drop in performance, despite the sparsification policy being trained solely on math reasoning data.

Eviction pattern analysis revealed that KVpop does not merely enact recency or uniform retention; reasoning-structural tokens (discourse markers, operation words, symbolic pivots) are preferentially retained, while purely numerical tokens are evicted at higher rates, especially in upper layers (Figure 3, Figure 4). Figure 3

Figure 3: Layer-wise retention visualization for mathematical reasoning traces, indicating structural token retention.

Figure 4

Figure 4: Cross-layer retention patterns on multiple sequences under 75% KV compression.

Alignment with a full-attention “oracle” showed KVpop achieves 81% mean recall of teacher top-kk retention per layer, demonstrating strong agreement (Figure 5). Figure 5

Figure 5: Top-kk retention recall distribution, measuring KVpop’s agreement with the future-attention oracle.

Contrast with DMS revealed KVpop’s fixed per-head budget yields homogeneous and predictable cache allocation, enabling efficient GPU execution and stable throughput for long generations. DMS, in contrast, exhibits winner-take-all allocation, with cache capacity concentrated in few heads (Figure 6). Figure 6

Figure 6: DMS eviction ratio heatmap, showing heterogeneous allocation across heads and layers.

KVpop also improved end-to-end latency due to predictable cache size per head, outperforming both dense attention and DMS at extreme sequence lengths (Figure 7). Figure 7

Figure 7

Figure 7: End-to-end decoding latency for Qwen3-8B, showing KVpop’s throughput advantage across generation lengths.

Delayed scoring via mLSTM modules provided additional benefit, increasing token accuracy by 0.2 points versus immediate scoring ablations.

Theoretical and Practical Implications

KVpop’s formulation—direct supervision of eviction at the boundary using a future-attention target—addresses the core challenge of predicting token utility in the evolving context. By decoupling scorer optimization from the backbone and exploiting delayed stateful scoring, the framework achieves aggressive cache reduction with minimal performance loss. Practically, it enables LLM deployment in memory-constrained environments and facilitates long-context processing (e.g., legal, scientific, and code generation domains). Theoretically, it demonstrates that content-dependent eviction can be learned and generalized efficiently, outperforming strictly recency-based or local scoring heuristics.

The transferability of the retention policy, even across out-of-domain benchmarks, suggests that learned cache sparsification can generalize, potentially catalyzing work on architecture-agnostic sparse retrofitting. Future directions may encompass paged KV cache management, hybrid dense-sparse allocation across layers, and exploration of alternative memory scorer architectures beyond mLSTM.

Conclusion

KVpop advances the state of the art in KV cache compression for transformer LLMs by formulating eviction as a boundary-supervised prediction task, anchored to future-attention signals. The method achieves near-full teacher performance under substantial memory reduction, with homogeneous cache allocation, efficient inference throughput, and robust generalization. The framework lays the foundation for post-training sparse retrofitting, unlocking long-context language modeling with bounded memory and high quality.

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Collections

Sign up for free to add this paper to one or more collections.