KVzap: Fast, Adaptive, and Faithful KV Cache Pruning

AI summary

Overview

Research area: Efficient large language model inference, specifically key-value (KV) cache compression for transformer attention.

Technical level: Advanced. The paper assumes familiarity with transformer attention, KV caching, grouped query attention (GQA), and inference serving internals, though its core idea can be understood without that background.

Scope: The paper introduces KVzap, a lightweight learned surrogate that predicts which KV cache entries are safe to discard, evaluated on Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B across RULER, LongBench, and AIME25.

What This Paper Is About

As context lengths grow into the tens or hundreds of thousands of tokens, the KV cache — the stored key and value vectors that let a model reuse past tokens during generation — becomes a dominant memory and speed bottleneck. Many pruning methods have been proposed, but none have been adopted in major inference engines such as vLLM, SGLang, or TRT-LLM because they trade too much accuracy for speed, or are too slow to run. The paper's goal is a pruning method that is fast, works during both prefilling and decoding, fits existing optimized attention kernels, and loses almost no accuracy.

Key Contributions

1. KVzip+: An enhanced version of the KVzip scoring rule that adds a normalization term (inspired by Expected Attention) based on each token's contribution to the residual stream relative to the norm of the hidden state.
2. A surrogate approximation: The paper shows that the expensive KVzip+ scores can be approximated by a lightweight per-layer surrogate model — either a single linear layer (KVzap-Linear) or a two-layer MLP (KVzap-MLP) — trained on the model's own hidden states.
3. KVzap: A new pruning technique that applies these surrogates at inference using a DMS-style recipe: per-head compression, score-based pruning via thresholding, and delayed eviction with a sliding window.
4. State-of-the-art benchmark results: KVzap achieves 2x–4x KV cache compression with negligible accuracy loss and reaches state-of-the-art performance on the KVpress Leaderboard, outperforming 15 other methods while matching the KVzip method it approximates.

Main Findings

• State-of-the-art on the KVpress Leaderboard: On RULER 4k, KVzap achieves top results for both Qwen3-8B and Llama-3.1-8B-Instruct, matching KVzip while outperforming 15 other methods, including Expected Attention, Duo Attention, and Compactor.
• Surrogates approximate the oracle well: Average squared Pearson correlation (R²) between KVzip+ scores and KVzap predictions on validation data ranged from 0.671 (Qwen3-8B, Linear) and 0.711 (Qwen3-8B, MLP), to 0.743 / 0.772 for Llama-3.1-8B-Instruct, and 0.629 / 0.668 for Qwen3-32B. KVzap-MLP consistently outperformed KVzap-Linear.
• KVzip+ matches or beats KVzip: Across RULER and LongBench, the added normalization term consistently matched or exceeded the original KVzip scoring.
• Negligible overhead: Relative compute cost is bounded by 1.1% for KVzap-MLP and 0.02% for KVzap-Linear when considering linear projections only. Per-model figures: Qwen3-8B 1.09% / 0.02%, Llama-3.1-8B-Instruct 0.96% / 0.02%, Qwen3-32B 0.67% / 0.01%.
• Accuracy preserved under compression (Table 2):
  • Qwen3-8B (KVzap-MLP, τ = -4): RULER 4k 95.32 → 95.09; RULER 16k 92.99 → 92.78; LongBench 46.74 → 46.49; AIME25 pass@4 0.77 → 0.77; average compression ratio 0.72 (3.5x).
  • Llama-3.1-8B (KVzap-Linear, τ = -7): RULER 4k 95.69 → 95.55; RULER 16k 93.42 → 93.29; LongBench 45.25 → 44.65; average compression ratio 0.67 (3.0x).
  • Qwen3-32B (KVzap-MLP, τ = -4): RULER 4k 95.65 → 95.95; RULER 16k 95.19 → 94.96; LongBench 50.56 → 50.40; AIME25 pass@4 0.83 → 0.87; average compression ratio 0.63 (2.7x).
• Reasoning is preserved: KVzap-MLP preserves AIME25 reasoning accuracy even at compression ratios exceeding 2x, whereas KVzap-Linear with τ = -3 hit 96% compression on Qwen3-8B and 93% on Qwen3-32B, producing zero correct answers.
• Compression adapts to input: Thresholding produces up to 20% variation in compression ratio across prompts, and outperforms fixed-ratio top-k selection, whether applied per-head or per-layer (AdaKV).
• The sliding window matters: On LongBench-LCC with Qwen3-8B, KVzap-MLP, and τ = -4, accuracy without a local window (w = 0) drops to 28.37%. Enforcing w = 128 restores it to 62.51%, while w = 512 yields no additional gain (62.37%).
• Task-dependent compression: The same thresholds yield lower compression ratios on LongBench than on RULER, likely because RULER samples are synthetic and repetitive while LongBench is real-world data with higher information density.
• A caveat on LongBench: Expected Attention's apparently above-baseline scores are largely driven by outliers on the TREC subset, and the small size of LongBench subsets (typically n = 200) leads to high variance.

Methodology in Plain English

The starting point is KVzip, the strongest existing pruning method. KVzip works by making the model do a "copy-and-paste" exercise: it appends a repeated copy of the prompt and asks the model to reproduce it, then scores each KV pair by how much attention the repeated copy pays back to the original. High-scoring pairs are kept; low-scoring pairs are dropped. This is accurate but expensive — it requires a double forward pass on a prompt twice as long as the input — and it cannot run during decoding.

KVzap sidesteps this. Rather than running the expensive scoring procedure at inference, the authors train small models (one linear layer or a two-layer MLP per transformer layer) to predict the KVzip+ scores directly from the hidden states the model already computes. Because these surrogates need only one or two matrix multiplications and depend only on hidden states, they add almost no cost and can run during both prefilling and decoding.

Training data came from Nemotron-Pretraining-Dataset-sample: 27k prompts across 9 subsets (common crawl, multilingual, math, code, and others), filtered to 750–1,250 tokens, with up to 500 prompts per subset selected for training and 5 for validation, giving roughly 2.4k prompts. From these, 500 tokens were randomly sampled per prompt to yield 1.2M training pairs per KV head, with 23k held out for validation. The surrogates map a hidden state to one score per KV head (H = 8), with input dimension equal to the model hidden size (4,096 or 5,120) and a GELU-activated hidden layer of width D_h/8 (512 or 640) for the MLP.

At inference, KVzap discards any KV pair whose predicted score falls below a threshold τ. Higher thresholds mean more aggressive compression. A sliding window of the most recent w = 128 tokens is always retained. Compression is applied after the attention operation.

Why This Matters

KV cache pruning has been an active academic area for years, but the paper's central claim is a practical one: none of these methods have made it into vLLM, SGLang, or TRT-LLM, because they fail at least one of four criteria — being fast and lightweight, working in both prefilling and decoding, being compatible with FlashAttention2 or PagedAttention kernels, and being faithful to the uncompressed model's accuracy. KVzap is framed explicitly as an attempt to satisfy all four.

Real-world applications:

• Long-document processing: summarizing, searching, or question-answering over inputs of tens to hundreds of thousands of tokens, where the KV cache otherwise dominates GPU memory.
• Reasoning and chain-of-thought systems: models that generate thousands of output tokens, a regime where KVzip cannot be applied at all but KVzap can.
• Cost reduction in serving: 2.7x–3.5x average compression, as measured here, directly reduces the memory footprint per concurrent request.
• Deployment on memory-constrained hardware: holding more concurrent sessions or longer contexts within fixed GPU memory.

Industry relevance: The work comes from NVIDIA and is released alongside code in NVIDIA/kvpress and trained models in the NVIDIA/KVzap collection, with the stated goal of bridging academic pruning research and production inference engines. The authors note that turning compression into wall-clock speedups requires PagedAttention kernels that handle variable-length blocks across heads, and that prior systems (DMS, Compactor, AdaKV) have shown this is feasible.

Future Directions

• Scale and generalization: Validate beyond the 32B model studied here, on larger open-source models such as GLM 4.7 and Qwen3-235B-A22B, and on architectures with sparse attention such as DeepSeek V3.2. The authors also suggest extending evaluation to more reasoning benchmarks, agentic tasks, and short-context knowledge tasks.
• End-to-end training: KVzap is a post-hoc addition, not training-free. The authors argue that end-to-end pruning objectives such as DMS may eventually yield better performance, drawing an analogy to how Multi-Token Prediction is superseding ad-hoc speculative decoding techniques such as Medusa.
• Kernel and systems work: Converting compression into actual wall-clock speedups and GPU memory savings was not explored. KVzap creates non-uniform cache lengths across heads, requiring PagedAttention kernels that handle variable-length blocks.
• Pre-attention pruning: Since KVzap relies only on hidden states, the authors note pruning could be applied before attention to directly accelerate prefilling.
• Training improvements: The paper acknowledges a potential train-test distribution shift because KVzap was trained on prompts limited to 1,250 tokens, and suggests better data selection and hyperparameter tuning as paths to further accuracy gains.

Target Audience

This paper is most valuable to inference and ML systems engineers working on LLM serving infrastructure, researchers working on KV cache compression and efficient attention, and practitioners who need long-context or long-generation inference under tight memory budgets. Readers evaluating whether to adopt pruning in production will find the four deployment criteria and the overhead analysis (Appendix B) directly relevant. Readers new to transformer inference will find the framing accessible but will need background on attention and KV caching to follow the scoring derivations.

Authors’ abstract

Growing context lengths in transformer-based language models 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.