PaperlayerView exact revision on arXiv

Source-linked AI summary

KVzap: Fast, Adaptive, and Faithful KV Cache Pruning

Simon Jegou, Maximilian Jeblick

arXiv:2601.07891v2cs.LGcs.AIcs.CL

TL;DR

Growing contexts make KV-cache memory a major inference bottleneck, while existing pruning methods face speed and deployment limitations. KVzap approximates KVzip with lightweight hidden-state surrogate models and adaptive thresholding, achieving 2–4× KV-cache compression with negligible accuracy loss across multiple models and tasks. Its post-hoc design remains a limitation compared with possible end-to-end pruning.

  • Problem

    Long contexts make the KV cache a dominant inference bottleneck, while existing pruning methods have not been integrated into major inference engines because of practical speed and phase limitations.

  • Method

    KVzap uses lightweight linear or MLP models over hidden states to approximate KVzip+ scores and prunes below-threshold KV pairs adaptively during prefilling and decoding.

  • Results

    2–4× KV cache compression with negligible accuracy loss is achieved across Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B on long-context and reasoning benchmarks.

  • Takeaways & Limitations

    KVzap’s simplicity, compression, and robust performance make it a candidate for production deployment and for bridging academic pruning research with inference engines.

  • Takeaways & Limitations

    KVzap is not training-free and remains a post-hoc addition, while end-to-end pruning may eventually yield better performance.

Abstract

from arXiv · show

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.

1. Introduction

Longer contexts make the KV cache a major inference bottleneck, while existing architectural and pruning approaches do not yet provide broadly adopted, efficient compression along the sequence-length axis. KVzap is introduced as a fast, phase-agnostic, faithful approximation of KVzip.

  • 335 GB of bfloat16 KV-cache memory is required by Llama1-65B at 128k tokens, making cache size a dominant inference bottleneck.The cache also increases peak GPU memory and time to first token while reducing decoding throughput.
  • 4×, 12×, and up to 16× KV-cache compression are achieved by GQA in Llama3, GLM 4.5, and Qwen3-235B-A22B, respectively.These architectural changes target the head axis rather than the sequence-length axis.
  • No widely adopted architectural change compresses the sequence-length axis, while sparse attention improves throughput without reducing resident KV-cache memory.This leaves memory-size reduction along the T-axis as an open practical need.
  • More than 20 KV-cache pruning methods are implemented in NVIDIA/kvpress, yet none are integrated into major engines such as vLLM, SGLang, or TRT-LLM.The paper identifies speed, phase coverage, optimization compatibility, and faithfulness as requirements for practical pruning.
  • KVzap is introduced as a fast approximation of KVzip that uses surrogate models over hidden states for KV-cache pruning.The contribution includes KVzip+ scoring, surrogate approximation, and DMS-style per-head threshold pruning with delayed eviction.

2. Related work

Related work reduces attention computation, evicts less useful KV pairs, or lowers KV precision, but important methods remain constrained by memory, speed, phase, or training requirements. KVzap is positioned against these approaches as a practical pruning method.

  • Block-Sparse Attention and Offloading: Sparse attention methods reduce computation by selecting relevant KV blocks or segments, but they focus on throughput rather than permanently reducing resident-cache memory.MoBA can reduce prefilling latency but often requires fine-tuning to recover performance.
  • Training-Free Eviction Methods: Training-free eviction methods remove KV pairs judged unlikely to be used, generally leaving the forward pass unchanged.These methods aim to reduce memory while preserving information needed for faithful generation.
  • Training-Free Eviction Methods: KVzip requires a computationally expensive double forward pass on an extended prompt and cannot be applied during decoding.These limitations make it unsuitable for online inference and decoding-heavy reasoning tasks.
  • End-to-End Learned Eviction: DMS learns per-token eviction decisions from hidden states using trainable weights, differentiable Gumbel-sigmoid reparametrization, and an output-preservation loss.Its training objective combines KL divergence with a sparsity penalty.
  • Quantization: Quantization methods such as KIVI and ZipCache reduce cache memory by storing KV pairs in INT4 or INT2 formats with specialized kernels.Quantization is presented as orthogonal to token pruning and combinable with it.

3. Method

KVzap predicts KV-pair importance from hidden states using lightweight per-layer surrogate models trained to approximate KVzip+ scores. It prunes below-threshold pairs adaptively while preserving a recent-token window.

  • KVzap applies a lightweight hidden-state model in each transformer layer to predict importance scores and discard KV pairs below threshold τ.The surrogate approximates the scoring policy of an improved KVzip variant.
  • KVzip reaches up to 4× compression with minimal accuracy loss but requires two extended-prompt prefills and cannot operate during decoding.These constraints motivate the KVzap approximation.
  • KVzip scores KV pairs using attention weights from a copy-and-paste repetition prompt, then removes the lowest-scoring pairs across heads and layers.The method treats low attention during repetition as evidence that a KV pair carries little information.
  • KVzip+ adds a normalization term based on each value vector’s contribution to the residual stream before surrogate prediction.The contribution term is described as a_j_i W_O v_i for token i at decoding step j.
  • Per-layer linear or two-layer MLP surrogates map each hidden state h_t to H log-space scores using one or two matrix multiplications.The models are trained on 1.2M hidden-state and log-score pairs per KV head from diverse text.
  • Thresholding makes KVzap input-adaptive: higher thresholds increase compression, while information-dense prompts retain more tokens than redundant prompts.A sliding window of the most recent w=128 tokens is always retained to preserve local context.

4. Experiments

Experiments evaluate KVzap on three language models using surrogate training and validation procedures designed for diverse text. The surrogates approximate KVzip+ scores well, while KVzap adds little computational and memory overhead.

  • Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B are evaluated, with reproducible code, trained models, and evaluation logs provided.The experiments use the NVIDIA/kvpress and NVIDIA/KVzap resources.
  • 1.2M training pairs per KV head are produced from roughly 2.4k selected prompts and 500 sampled tokens per prompt, with 23k pairs held out for validation.The source covers common crawl, multilingual, math, and code subsets, among others.
  • KVzap-Linear and KVzap-MLP predict log(s+) with outputs for 8 KV heads, using hidden dimensions of 4096 or 5120.The MLP uses one hidden layer with width D_h/8 and GELU activation.
  • R^2 values of 0.60–0.80 show that KVzip+ scores can be approximated from hidden states, with KVzap-MLP consistently outperforming KVzap-Linear.The reported metric is average Squared Pearson correlation over all KV heads and layers.
  • 1.1% and 0.02% are the respective upper bounds on relative compute cost for KVzap-MLP and KVzap-Linear when counting linear projections.Relative memory overhead matches these bounds, and long-context attention cost dominates overall computation.

4.3. Prefilling and decoding tasks

KVzap is evaluated across long-context prefilling benchmarks and reasoning-oriented decoding, spanning RULER, LongBench, and AIME25. On RULER, it achieves state-of-the-art results and preserves accuracy at substantial compression.

  • Evaluation scope: KVzap is evaluated on RULER, LongBench, and AIME25 to cover long-context prefilling and reasoning-oriented decoding regimes.The benchmarks contain 6500, 4750, and 30 samples, respectively.
  • LongBench: LongBench evaluates six long-context task categories across 21 English and Chinese subsets, with average performance reported across subsets.The categories include question answering, summarization, few-shot learning, synthetic tasks, and code completion.
  • RULER: KVzap achieves state-of-the-art RULER 4k results for Qwen3-8B and Llama-3.1-8B-Instruct, outperforming 15 concurrent pruning methods.The result is reported for both evaluated models.
  • RULER: KVzap maintains perfect RULER accuracy up to 3–4× compression, while KVzip+ consistently matches or exceeds KVzip.KVzap-MLP outperforms KVzap-Linear on Qwen models, whereas KVzap-Linear performs especially well on Llama-3.1-8B-Instruct.

4.5. LongBench

LongBench results show that KVzap retains near-perfect accuracy at 2–3× compression, while AIME25 results indicate that KVzap-MLP preserves reasoning accuracy beyond 2× compression. LongBench averages require caution because small, low-accuracy subsets create high variance.

  • LongBench: KVzap maintains near-perfect LongBench accuracy up to 2–3× compression, with the same thresholds producing lower ratios than on RULER.The difference is attributed to LongBench’s denser, mostly real-world data compared with synthetic and repetitive RULER samples.
  • LongBench: Expected Attention’s elevated LongBench scores are largely driven by TREC outliers, while KVzap remains close to the full-cache baseline across several subsets.The reported interpretation points to over-prompting as a possible explanation for higher TREC accuracy at high compression.
  • LongBench: LongBench’s generally low subset accuracy and small subset size, typically n=200, produce high variance that complicates conclusive interpretation.This caveat applies to the benchmark’s per-subset results and their aggregate interpretation.
  • AIME25: KVzap-MLP preserves AIME25 reasoning accuracy at compression ratios exceeding 2× using average pass@1 and pass@4 over four rollouts per question.AIME25 contains 30 Olympiad-level integer-answer problems, with a 32k-token generation limit.

4.7. Adaptive compression

KVzap uses thresholding to adapt compression to prompt complexity rather than enforcing a fixed budget. Across models and datasets, it averages 2.7–3.5× compression while preserving accuracy, with sliding-window size affecting LongBench performance.

  • Adaptive compression: The maximum accuracy-preserving compression is task-dependent, and KVzap’s thresholding automatically yields different compression ratios across benchmarks.The same threshold can produce different ratios because benchmark characteristics differ.
  • Adaptive compression: KVzap achieves 2.7–3.5× average KV cache compression while maintaining accuracy across model scales and tasks.Table 2 reports the best Linear/MLP and threshold configuration for each model.
  • Adaptive compression: Thresholding adapts compression to prompt complexity, producing up to 20% variation across prompts and outperforming fixed-ratio top-k selection.The comparison includes both per-head and per-layer AdaKV baselines.
  • Sliding window: On LongBench-LCC, increasing the sliding window from w=0 to w=128 raises accuracy from 28.37% to 62.51%, while w=512 yields 62.37%.The experiment uses Qwen3-8B, KVzap-MLP, and τ=−4; the window restores position information absent from input hidden states.

5. Discussion

KVzap achieves 2–4× KV cache compression with negligible accuracy loss across multiple models and benchmarks, but broader validation and systems engineering remain necessary for deployment.

  • Discussion: 2–4× KV cache compression with negligible accuracy loss is reported across Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B on RULER, LongBench, and AIME25.The design uses lightweight linear or MLP models applied to hidden states.
  • Scope and Generalization: Further validation is needed on larger open-source models and architectures with sparse attention.The authors specifically identify GLM 4.7, Qwen3-235B-A22B, and DeepSeek V3.2 as future validation targets.
  • Scope and Generalization: Evaluation could be extended to more reasoning benchmarks, agentic tasks, and short-context knowledge tasks.
  • Ad-hoc vs. End-to-End Training: KVzap is a post-hoc addition rather than a training-free method, while end-to-end pruning objectives may eventually yield better performance.The authors contrast KVzap with end-to-end approaches such as DMS.
  • Implementation Challenges: Turning compression into wall-clock speedups and GPU memory savings requires engineering for variable-length blocks, and this optimization was not explored.KVzap’s non-uniform per-head cache lengths require kernels such as PagedAttention.
  • Implementation Challenges: Because KVzap relies only on hidden states, pruning could also be applied before attention to accelerate prefilling.
  • Conclusion: The authors present KVzap as a candidate for production deployment based on its simplicity, compression ratios, and performance across tasks and models.

A. KVzap model training

KVzap trains lightweight surrogate models to infer KVzip+ scores from hidden states, with MLPs generally approximating the scores better than linear models.

  • Training analysis: KVzap reports detailed KVzip+ and R^2 score distributions for Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B.
  • Surrogate comparison: KVzap-MLP consistently achieves higher R^2 than KVzap-Linear across models.
  • Threshold selection: Llama-3.1-8B-Instruct has a lower score distribution than Qwen3-8B and Qwen3-32B, motivating lower pruning thresholds.
  • Layer analysis: Both surrogates perform worse in the first transformer layer, where KVzip+ scores are harder to infer from token embeddings alone.
  • Input representation: Predicting scores from keys and values instead of hidden states produces strictly lower R^2.
  • Implementation: KVzap-Linear is trained with scikit-learn, while KVzap-MLP is trained with skorch.

B. KVzap compute and memory overhead

KVzap adds very little compute or memory overhead relative to a transformer layer, although aggressive compression can eliminate benchmark accuracy in some settings.

  • Compute overhead: 1.1% maximum relative compute overhead is reported for KVzap-MLP, compared with 0.02% for KVzap-Linear.The comparison considers linear projections only.
  • Compute overhead: In long-context regimes, quadratic attention cost dominates overall complexity, making KVzap’s overhead effectively negligible.
  • Accuracy boundary: 96% compression for Qwen3-8B and 93% for Qwen3-32B with KVzap-Linear at τ = −3 explains the zero AIME25 scores.
  • Memory overhead: KVzap’s relative memory overhead matches its compute-overhead bounds because the FLOPs-counting factor cancels in the ratio.
  • Overall overhead: KVzap’s additional parameters introduce no meaningful memory or compute overhead.

C. Detailed benchmark results

The detailed evaluation reports per-subset and aggregate results across RULER, LongBench, and AIME25 for the three evaluated models.

  • RULER: RULER Figures 9–11 provide per-subset results for Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B.
  • LongBench: LongBench Figures 13–15 provide per-subset results for all three models, with Figure 12 reporting average scores excluding TREC.
  • AIME25: AIME25 results are reported for each of four rollouts in Table 4.
  • RULER: Figure 9 reports RULER 4k results for Qwen3-8B across 13 subsets.
  • RULER: Figure 10 reports RULER 4k results for Llama-3.1-8B-Instruct across 13 subsets.
  • RULER: Figure 11 reports RULER 4k results for Qwen3-32B across 13 subsets.
  • LongBench: Figure 12 reports LongBench average scores across 20 or 21 subsets after excluding TREC for the three models.
  • LongBench: Figures 13–15 report LongBench results for Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B across 21 subsets.
Loading 2601.07891v2…