💬 NLP

KVzap: Fast, Adaptive, and Faithful KV Cache Pruning

KVzap is a fast, input-adaptive KV cache pruning method that achieves state-of-the-art compression with negligible accuracy loss across various large language models and tasks, addressing the critical inference bottleneck caused by growing context lengths.

Original authors: Simon Jegou, Maximilian Jeblick

Published 2026-02-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Simon Jegou, Maximilian Jeblick

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are reading a very long book, and every time you turn a page, you have to keep a mental note of every single word you've read so far to understand the story. As the book gets longer, your brain (or in this case, the computer's memory) gets overwhelmed trying to hold onto every single word. This is the problem with modern AI models: as they read longer and longer texts, they run out of "mental space" (memory) to keep track of the important parts.

This paper introduces a new tool called KVzap that acts like a super-efficient librarian for these AI models. Here is how it works, broken down into simple concepts:

The Problem: The "Mental Clutter"

When an AI reads a sentence, it creates a "Key-Value" (KV) pair for every word. Think of these as sticky notes with the word written on them, stored in a giant pile.

The Bottleneck: If the AI reads a 100,000-word document, it has 100,000 sticky notes. Keeping all of them takes up a massive amount of memory, slowing the AI down and making it expensive to run.
The Old Way: Previous methods tried to throw away sticky notes based on simple rules (like "keep only the last 100 words" or "keep the words that appeared most often"). But these rules were often too blunt, accidentally throwing away important information and making the AI make mistakes.

The Solution: KVzap (The Smart Librarian)

KVzap is a new method that decides which sticky notes to keep and which to throw away, but it does it fast, smartly, and faithfully (without losing the story).

Here is the analogy of how KVzap works:

1. The "Double-Check" Problem (The Old Method)
A previous top-tier method called KVzip was very accurate. It worked like this: To decide if a word was important, the AI would pretend to read the whole book again and see which words it paid attention to.

The Flaw: This is like asking a student to read a chapter, then read it a second time just to decide which sentences were important. It's too slow and takes too much energy. It also couldn't work while the AI was writing a story (decoding), only while reading.

2. The KVzap Shortcut
KVzap solves this by training a tiny, super-fast "assistant" (a small neural network) to guess which words are important without re-reading the text.

The Analogy: Imagine a seasoned librarian who has read thousands of books. Instead of re-reading a new book to find the important parts, the librarian glances at the page and instantly knows, "This paragraph is crucial, keep it. That one is just filler, throw it away."
How it learns: The team trained this "assistant" by showing it the AI's internal thoughts (hidden states) and asking it to predict what the slow, double-check method would have said. The assistant learned to mimic the expert perfectly but in a fraction of a second.

3. The "Sliding Window" Safety Net
To make sure the AI doesn't forget the immediate context (like the last few sentences it just wrote), KVzap keeps a "sliding window" of the most recent 128 words.

The Analogy: Even if you are summarizing a whole novel, you always keep the last few pages in your hand so you don't lose your place. KVzap keeps these recent words safe and never throws them away.

Why is this a big deal?

The paper claims KVzap is a game-changer for three main reasons:

It's Fast: It adds almost no extra time to the process. It's like having a librarian who sorts the books while you are still walking into the library.
It's Adaptable: It doesn't use a fixed rule (like "keep 50% of the notes"). Instead, it looks at the text. If the text is complex and dense, it keeps more notes. If the text is repetitive, it throws more away. It adjusts automatically.
It's Accurate: In tests on long reading tasks (like answering questions from a 4,000-word document) and reasoning tasks (like solving math problems), KVzap managed to shrink the memory usage by 2 to 4 times while keeping the AI's accuracy almost exactly the same as if it had kept all the notes.

The Results

The researchers tested this on popular AI models (like Qwen and Llama). They found that:

KVzap beat 15 other existing methods.
It achieved the best scores on a leaderboard for long-context tasks.
It works for both reading long documents (prefilling) and writing long stories or solving problems step-by-step (decoding).

In Summary

KVzap is like giving an AI a "smart filter" that instantly knows which parts of a long conversation or document are essential and which are just noise. It allows the AI to handle much longer texts without running out of memory or getting confused, all without slowing down the process. The authors believe this makes it ready for real-world use in major AI systems.

Technical Summary: KVzap - Fast, Adaptive, and Faithful KV Cache Pruning

1. Problem Statement

As context lengths in transformer-based language models expand to tens or hundreds of thousands of tokens, the Key-Value (KV) cache has become a dominant bottleneck for efficient Large Language Model (LLM) inference. The KV cache, which stores intermediate attention states, consumes significant GPU memory (e.g., 335 GB for a 65B model at 128k sequence length) and increases time-to-first-token while reducing decoding throughput.

While architectural modifications (e.g., Grouped Query Attention, Multi-head Latent Attention) and sparse attention mechanisms have addressed compression along the head ( $H$ ), dimension ( $D$ ), and layer ( $L$ ) axes, the sequence length axis ( $T$ ) remains largely unaddressed by widely adopted methods. Existing ad-hoc pruning methods often fail to meet four critical criteria for production deployment:

Fast and lightweight: Negligible overhead.
Phase-agnostic: Applicable to both prefilling (long context) and decoding (reasoning).
Optimization-friendly: Compatible with kernels like FlashAttention2 and PagedAttention.
Faithful: Minimal accuracy degradation.

Current state-of-the-art methods like KVzip achieve high compression but require a computationally expensive double forward pass on extended prompts, rendering them unsuitable for online inference and decoding tasks.

2. Methodology

The authors introduce KVzap, a fast, input-adaptive approximation of KVzip that operates during both prefilling and decoding. The methodology consists of three core components:

2.1. Enhanced Scoring Policy (KVzip+)

The authors first refine the scoring mechanism of KVzip. KVzip scores tokens based on "context reconstruction," measuring the maximum attention weight a token receives when the prompt is repeated. The authors enhance this by incorporating a normalization term inspired by Devoto et al. (2025), which accounts for the contribution of a token's value vector to the residual stream.
The KVzip+ score ( $s^+_i$ ) is defined as:
$s^+_i = \max_{j \in \text{<prompt>}} \frac{a_{ji}}{\|W_O v_i\| / \|h_j\|}$
where $a_{ji}$ is the attention weight, $W_O$ is the output projection matrix, $v_i$ is the value vector, and $h_j$ is the hidden state.

2.2. Surrogate Modeling

To avoid the prohibitive cost of the KVzip+ double forward pass, KVzap trains lightweight surrogate models (either a linear layer or a two-layer MLP) to predict the log-scores ( $\log(s^+)$ ) directly from the model's input hidden states ( $h$ ).

Input: Hidden states $h_t \in \mathbb{R}^{D_h}$ .
Output: Scores for each KV head in $\mathbb{R}^H$ .
Training: The models are trained on 1.2 million pairs of (hidden state, log-score) sampled from the Nemotron-Pretraining-Dataset. The MLP variant uses a hidden width of $D_h/8$ with GELU activation.
Performance: The surrogates achieve an average $R^2$ correlation of 0.60–0.80 with the KVzip+ oracle, demonstrating that expensive scores can be accurately approximated from hidden states.

2.3. Adaptive Pruning and Eviction

Unlike KVzip, which enforces a fixed budget (e.g., keeping exactly 50% of tokens), KVzap employs threshold-based pruning:

Thresholding: Tokens with predicted scores below a fixed threshold $\tau$ are discarded. This allows the compression ratio to adapt dynamically to the information density of the input prompt.
Sliding Window: To preserve local context, the most recent $w=128$ tokens are always retained, following the StreamingLLM approach.
Implementation: The pruning occurs after the attention operation. The process is computationally efficient, adding negligible overhead (bounded by 1.1% for MLP and 0.02% for Linear models).

3. Key Contributions

KVzip+: An improved scoring metric for KV cache importance that incorporates hidden-state normalization.
Surrogate Approximation: A demonstration that KVzip+ scores can be effectively approximated by lightweight models trained on hidden states, eliminating the need for double forward passes.
KVzap: A new pruning technique that applies these surrogates with threshold-based eviction and sliding windows, achieving state-of-the-art performance while meeting all four practical criteria (speed, phase-agnosticism, optimization-friendliness, and faithfulness).

4. Experimental Results

Experiments were conducted on Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32B across three benchmarks: RULER (long-context retrieval/aggregation), LongBench (multi-task long-context), and AIME25 (mathematical reasoning).

Compression vs. Accuracy: KVzap achieves 2–4× KV cache compression with negligible accuracy loss.
- On RULER 4k, KVzap matches or exceeds the performance of the KVzip+ oracle and outperforms 15 other pruning methods (including Expected Attention and Duo Attention).
- On LongBench, KVzap maintains near-perfect accuracy up to 2–3× compression.
- On AIME25, KVzap-MLP preserves reasoning accuracy even when discarding over 50% of the KV cache.
Adaptivity: The threshold-based approach automatically adjusts compression ratios based on task difficulty. For instance, it achieves higher compression on synthetic RULER data compared to the denser, real-world LongBench data.
Efficiency: The method adds negligible compute and memory overhead. During decoding, the additional FLOPs utilize GPU cycles that would otherwise be stalled by memory bandwidth limitations.

5. Significance and Claims

The paper claims that KVzap bridges the gap between academic pruning research and real-world inference engines. Its significance lies in:

Practicality: It is the first method to achieve state-of-the-art compression without sacrificing speed or requiring architectural changes, making it compatible with major inference engines (vLLM, SGLang, TRT-LLM).
Faithfulness: It demonstrates that LLMs do not fully exploit the KV cache and that unused pairs can be identified via hidden states without retraining the base model end-to-end.
Scalability: The approach is validated across models ranging from 8B to 32B parameters and is designed to be compatible with future sparse attention architectures.

The authors acknowledge limitations, noting that KVzap is a post-hoc addition rather than an end-to-end trained objective, and that further engineering is required to translate compression ratios into wall-clock speedups via optimized kernels (e.g., PagedAttention). However, they posit that KVzap's combination of simplicity, high compression, and robust performance makes it a prime candidate for production deployment.