KVzap: Fast, Adaptive, and Faithful KV Cache Pruning

TL;DR

KVzap addresses the KV cache memory bottleneck in long-context transformer inference by introducing a fast, input-adaptive pruning approach that uses lightweight per-layer surrogates to predict KV importance from hidden states and applies threshold-based pruning with a sliding window. By combining a normalization-inspired scoring (KVzip+) with efficient surrogates (KVzap-Linear and KVzap-MLP), it enables $2\times$ to $4\times$ KV cache compression while maintaining accuracy across prefilling and decoding tasks and achieving state-of-the-art results on KVpress benchmarks. The method is trained on a large, diverse dataset to learn per-layer mappings from hidden states to KV importance and includes thorough ablations and overhead analyses showing negligible compute and memory penalties. KVzap demonstrates strong generalization across models (Qwen3-8B, Qwen3-32B, Llama-3.1-8B-Instruct) and tasks (RULER, LongBench, AIME25), offering a practical path toward production deployment in inference engines with fast, adaptive, and faithful KV cache pruning.

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.

Figures (15)

• Figure 1: https://huggingface.co/spaces/nvidia/kvpress-leaderboard for Qwen3-8B (left) and Llama-3.1-8B-Instruct (right) comparing different KV cache pruning methods. The plots compare the accuracy on the RULER 4k dataset ruler (y-axis) against the KV cache compression ratio (x-axis). KVzap achieves state-of-the-art performance on both models, matching KVzip kvzip — which it approximates — while outperforming 15 other methods, including Expected Attention expected_attention, Duo Attention duoattention, and Compactor compactor.
• Figure 2: RULER 4k results for Qwen3-8B (left), Llama-3.1-8B-Instruct (middle), and Qwen3-32B (right). Zoomed-in view (y-axis range [90, 100]) of the high-performance region from Figure \ref{['fig:leaderboard']}. KVzap surrogates perform comparably to---and sometimes exceed---the KVzip+ oracle they approximate.
• Figure 3: LongBench results for Qwen3-8B (left), Llama-3.1-8B-Instruct (middle), and Qwen3-32B (right). KVzap models again maintain accuracy close to the full KV cache baseline. The elevated scores for Expected Attention are primarily driven by outliers in TREC, one of the 21 subsets of LongBench; see Figure \ref{['fig:longbench_no_trec']} for results excluding TREC.
• Figure 4: AIME25 Reasoning Performance. Comparison of pass@1 (solid lines) and pass@4 (dashed lines) accuracy for Qwen3-8B (left) and Qwen3-32B (right). KVzap-MLP maintains robust performance even when discarding over 50% of the KV cache.
• Figure 5: Distribution of compression ratios for Qwen3-8B and KVzap-MLP on RULER 4k, LongBench, and AIME25 (left), and comparison to an alternative pruning method (right).