KVzap: Fast, Adaptive, and Faithful KV Cache Pruning

Abstract: Growing context lengths in transformer-based LLMs have made the key-value (KV) cache a critical inference bottleneck. While many KV cache pruning methods have been proposed, they have not yet been adopted in major inference engines due to speed--accuracy trade-offs. We introduce KVzap, a fast, input-adaptive approximation of KVzip that works in both prefilling and decoding. On Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B across long-context and reasoning tasks, KVzap achieves $2$--$4\times$ KV cache compression with negligible accuracy loss and achieves state-of-the-art performance on the KVpress leaderboard. Code and models are available at https://github.com/NVIDIA/kvpress.

• The paper presents KVzap, a fast, adaptive method for pruning transformer key-value caches with negligible accuracy loss at 2–4× compression rates.
• It employs a lightweight, trainable surrogate model that approximates KVzip+ scores from hidden states, reducing compute overhead while maintaining performance.
• Experimental results across multiple LLM benchmarks demonstrate KVzap’s robust performance and state-of-the-art accuracy–compression trade-offs.

KVzap: Fast, Adaptive, and Faithful KV Cache Pruning

Introduction and Motivation

Transformer-based LLMs require the storage of all key-value (KV) pairs in a memory cache throughout autoregressive generation. As context lengths scale to tens or hundreds of thousands of tokens, the KV cache emerges as a prohibitive bottleneck impacting inference throughput and memory usage, particularly for large-scale deployment on GPUs. Prior architectural interventions (GQA, MLA, hybrid stacking) optimize cache size along head, feature, or layer axes. However, efficient, practical, and faithful compression along the sequence (time, $T$) axis remains unsolved, with ad-hoc pruning methods exhibiting either non-negligible compute overheads or fragile accuracy trade-offs, thereby limiting adoption in real-world inference engines.

The paper introduces KVzap (2601.07891), a fast, input-adaptive, and faithful cache pruning method that approximates the KVzip scoring function with a lightweight, trainable surrogate operating purely on the hidden states. Crucially, KVzap achieves negligibly small accuracy loss at $2$–$4\times$ compression ratios on benchmark tasks across Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B—surpassing state-of-the-art (SOTA) baselines—while being compatible with both prefilling and decoding modes. The method is designed for minimal compute and memory overhead and exhibits strong generalization across synthetic, real-world, and reasoning datasets.

Figure 1: Accuracy vs. compression ratio on the RULER 4k dataset for Qwen3-8B (left) and Llama-3.1-8B-Instruct (right): KVzap matches or exceeds competing methods including KVzip, Expected Attention, and Compactor.

Method: Surrogate Scoring for Efficient KV Eviction

KVzip and KVzip+

KVzip operates by scoring each cached KV pair via a context-reconstruction mechanism and leveraging attention statistics to measure informativeness. For a given prompt, the copy-and-paste pretext task scores each KV as the maximum attention weight received during a repeat phase. Next, the lowest-scoring pairs are pruned. This method, while effective, is computationally prohibitive, requiring both a lengthy extended context and dual passes (prefilling, repeated context), rendering it impractical for both prefill and decoding.

To regularize and improve the informativeness measure, the authors propose KVzip+, a variant augmenting KVzip's score with a normalization, reflecting each token's true impact on the post-attention residual stream. Explicitly, the importance score for position $i$ is

$$s_i^+ = \max_{j \in \text{repeat prompt}} a_{ji} \frac{||W_O v_i||}{||h_j||}$$

where $a_{ji}$ is the attention weight, $W_O$ is the output projection, $v_i$ is the value vector, and $h_j$ is the hidden state at position $j$.

KVzap: Surrogate Model for Fast, Adaptive Pruning

KVzap learns to approximate the logarithm of the KVzip+ score using a lightweight per-layer surrogate (either linear or a small MLP) mapping hidden states $h_t$ directly to importance scores across all heads. For each position, the surrogate predicts $H$ scores in log-space, which are then thresholded by a user-selected $\tau$ to determine if each KV should be pruned. KVzap is trained post-hoc on a large, heterogeneous corpus of prompts, with explicit supervision from the oracle KVzip+ statistics.

Adaptive thresholding enables prompt-dependent compression, dynamically adjusting the ratio in accordance with information redundancy or query density. To protect recent (local) context, a sliding window of $w=128$ tokens is always retained, following StreamingLLM conventions.

Experimental Results

Accuracy–Compression Trade-off

Experiments are conducted on Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B using a suite of synthetic and real-world benchmarks (RULER, LongBench, AIME25). Both MLP and linear surrogates are evaluated. The surrogate's predictive accuracy, measured by Pearson $R^2$ on validation sets, ranges $0.63$–$0.77$, highlighting strong capacity to infer cache saliency from hidden states.

KVzap achieves $2$–$4\times$ cache reduction with negligible or no measurable accuracy degradation. On RULER 4k, KVzap not only matches but sometimes exceeds the KVzip+ oracle. Across LongBench, near-baseline performance is sustained for up to $3\times$ compression, with observed accuracy loss almost entirely explained by increased information density or real-world data heterogeneity. On AIME25 (Figure 2), even at $50\%$ reduction, pass@$k$ metrics remain stable.

Figure 3: RULER 4k results for Qwen3-8B, Llama-3.1-8B-Instruct, Qwen3-32B: High-accuracy region ([90, 100]) demonstrating negligible performance loss at high compression.

Figure 4: LongBench results for Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B: KVzap approximates full-KV baseline even at moderate-to-high compression ratios.

Figure 2: AIME25, a reasoning-intensive benchmark: pass@1 and pass@4 accuracy confirm robust performance, even when discarding over 50% of the KV cache.

Ablations and Surrogate Robustness

Distributional analysis (Figure 5) validates KVzap's ability to produce input-adaptive compression ratios, showing up to $20\%$ variation across prompt types. Surrogates trained with per-head MLPs always outperform linear variants, particularly at high compression, except in Llama-3.1-8B-Instruct where linear surrogates suffice, potentially due to lower overall saliency variance.

Figure 5: Distribution of achieved compression ratios across tasks (left) and superiority over fixed-ratio top-$k$ or AdaKV policies (right).

Detailed head-wise analysis (Figures 6–8) confirms that MLP surrogates achieve higher $R^2$ and better alignment with the oracle scores across all but the first transformer layers, and that their predictions are robust both in distribution and in correlation structure.

Implementation and Overhead

KVzap is implemented as a simple set of additional linear projections per transformer layer, introducing no more than $1.1\%$ relative compute overhead (MLP case) and $0.02\%$ (linear case), a cost dwarfed by quadratic attention and entirely hidden during memory-bound cache accesses (as shown in Table 1 of the paper). Memory overhead is similarly negligible. Integration with variable-length efficient attention kernels (PagedAttention, FlashAttention2) is feasible but not yet addressed in open-source inference engines—a practical engineering challenge for next steps.

Implications and Future Outlook

KVzap provides conclusive empirical support that LLMs underutilize their context memory and that informative KVs can be reliably identified and pruned via post-hoc analysis of hidden state geometry. This finding implicates substantial theoretical redundancy in standard transformer architectures, a gap that could be exploited for both faster inference and reduced deployment cost. Unlike methods enforcing architecture changes at pretraining (GQA, MLA), KVzap is a data-driven, backward-compatible “drop-in” solution.

Nevertheless, KVzap remains a post-hoc optimization—future work will likely focus on integrating pruning objectives into pretraining, enabling joint optimization of attention, memory, and cache usage end-to-end (as in DMS or fully-sparse attention models). Another line of work is large-scale engineering to fuse non-uniform pruning with production-quality kernel implementations, closing the gap between academic methods and practical LLM deployment.

Conclusion

KVzap offers a data-driven, highly efficient, and accurate framework for KV cache pruning in LLM inference, achieving SOTA accuracy–compression trade-offs with minimal engineering burden. Its design as a lightweight surrogate over hidden states makes it readily integrable and broadly applicable, with theoretical and practical implications for both LLM deployment and future transformer architecture research. The adaptability and empirical fidelity demonstrated by KVzap motivate further research toward cache-efficient, context-adaptive generative modeling.