PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast KVzip: Efficient and Accurate LLM Inference with Gated KV Eviction

Jang-Hyun Kim, Dongyoon Han, Sangdoo Yun

arXiv:2601.17668v2cs.LGcs.CL

TL;DR

Long-context inference makes KV-cache compression necessary, but existing methods trade accuracy against compression overhead. Fast KVzip adds lightweight sink-attention gating to frozen-weight LLMs and trains the gates without backpropagation. Across diverse models and tasks, it maintains near-lossless performance while reducing KV-cache size to 30%.

  • Problem

    Existing KV-cache eviction methods often trade performance degradation against prohibitive compression overhead, limiting practical long-context deployment.

  • Method

    Fast KVzip uses lightweight sink-attention gates over input hidden states to score and evict KV pairs during prefill and decoding while keeping LLM weights frozen.

  • Results

    Across prefill- and decoding-intensive tasks, Fast KVzip achieves near-lossless performance while reducing KV-cache size to 30% of its original size.

  • Takeaways & Limitations

    The method generalizes across retrieval, code comprehension, mathematical reasoning, diverse open-source LLMs, quantized weights, and hybrid attention mechanisms.

  • Takeaways & Limitations

    Applying the method during pretraining to induce hardware-efficient structured computation patterns remains future work.

Abstract

from arXiv · show

Efficient key-value (KV) cache management is crucial for the practical deployment of large language models (LLMs), yet existing compression techniques often incur a trade-off between performance degradation and computational overhead. We propose a novel gating-based KV cache eviction method for frozen-weight LLMs that achieves high compression ratios with negligible computational cost. Our approach introduces lightweight sink-attention gating modules to identify and retain critical KV pairs, and integrates seamlessly into both the prefill and decoding stages. The proposed gate training algorithm relies on forward passes of an LLM, avoiding expensive backpropagation, while achieving strong task generalization through a task-agnostic reconstruction objective. Extensive experiments across the Qwen2.5-1M, Qwen3, and Gemma3 families show that our method maintains near-lossless performance while evicting up to 70% of the KV cache. The results are consistent across a wide range of tasks, including long-context understanding, code comprehension, and mathematical reasoning, demonstrating the generality of our approach.

1. Introduction

Long-context inference makes KV-cache memory a bottleneck, while existing eviction methods trade compression efficiency against accuracy or runtime. Fast KVzip uses lightweight gating to assess KV importance during prefill and decoding, with experiments showing strong compression and broad task performance.

  • KV-cache memory scales linearly with sequence length, creating a dominant inference bottleneck for long-context LLMs.
  • Fast KVzip decodes future KV-pair utility from input hidden states and integrates lightweight gating into both prefill and decoding.The gating mechanism assesses KV importance without reconstructing the entire context.
  • Gates work best on input hidden states, while an attention-sink-inspired architecture provides the strongest performance.
  • 25% KV budget preserves model accuracy on RULER-4K while Fast KVzip matches KVzip’s compression performance and outperforms existing KVPress baselines.
  • Fast KVzip generalizes across Qwen and Gemma models, long-context and reasoning benchmarks, quantized weights, and sliding-window hybrid attention.

2. Related Work

Prior KV-cache methods exploit attention sparsity, query-agnostic importance, or predefined patterns, while Transformer gating methods target selective computation. Fast KVzip instead retains attention execution while using gating specifically for KV eviction.

  • KV Cache Compression: Efficient attention architectures often require training from scratch or substantial fine-tuning to accommodate structural changes.
  • KV Cache Compression: Inference-time sparsity methods can overfit current inputs, whereas query-agnostic methods estimate KV importance independently of specific queries for multi-query robustness.
  • Gating in Transformers: MoD gates attention skipping and MoBA gates block selection, while Fast KVzip performs attention for every token and improves efficiency through KV eviction.
  • KV Cache Compression: Locret is limited to prefill compression, while TrimKV degrades in prefill-intensive scenarios and specializes in mathematical reasoning.
  • KV Caching with Eviction: Fast KVzip computes importance scores and evicts low-importance KV features during chunked prefill, while decoding uses buffered parallel gating.

3. Method

Fast KVzip formulates KV compression as learned gating: layer-wise gates predict KV importance, evict low-importance features during prefill and decoding, and are trained from reconstruction-derived targets with frozen LLM weights. The method uses sink-attention gates and achieves efficient, accurate inference without runtime reconstruction overhead.

  • 3.1. Gating Mechanism: Fast KVzip adds an independent gating branch that processes each layer’s hidden features and outputs importance scores for that layer’s KV features.The branch does not alter the original attention outputs.
  • 3.1. Gating Mechanism: Low-importance KV features are evicted after chunked prefill, maintaining a compressed cache at a fixed retention ratio or within a predefined memory budget.The same gated eviction strategy is applied during decoding.
  • 3.2. Interpretation of Gating Formulation: Existing compression methods can be viewed as constrained instances of a generalized gating framework, whereas Fast KVzip permits differentiable, data-driven gating without empirical or heuristic constraints.This reframes KV cache compression as optimization of a gating function.
  • 3.3. Gate Training: Reconstruction-derived KV importance targets generalize significantly better than targets from next-token prediction or instruction-based question answering.The comparison evaluates Qwen2.5-7B-1M on synthetic key retrieval and SQuAD QA tasks.
  • 3.3. Gate Training: Hidden states outperform key states and pre-RoPE key states as gate inputs, while position-encoded features degrade performance.The method therefore separates positional information from gating and uses a local-window strategy for recency.
  • 3.4. Architecture: The gate uses a sink-attention architecture with fixed low-rank key features per layer, avoiding additional KV caching and simplifying integration into existing architectures.This architecture is competitive on contextual QA and avoids the degraded retrieval performance observed with alternative architectures.
  • 3.5. Efficiency: Fast KVzip reduces latency and memory usage versus the base model and improves prefilling efficiency over KVzip by removing runtime KV cache compression overhead.Inference uses a non-uniform KV cache eviction structure.

4. Experiment

Fast KVzip is evaluated across prefill-intensive and decoding-intensive settings, multiple open-source LLM families, and diverse long-context and reasoning benchmarks. It matches or exceeds baselines while preserving performance at substantially reduced KV-cache budgets.

  • Evaluation Setup: Fast KVzip is trained once per model on FineWeb-Edu and evaluated across downstream tasks using a single NVIDIA H100 80GB GPU.The evaluation uses PyTorch and FlashAttention-2.
  • Evaluation Setup: Fast KVzip is evaluated on Qwen2.5-7B/14B-1M, Qwen3-8B/14B, Qwen3-8B-FP8, and Gemma3-12B, including quantized and hybrid-attention models.For Gemma3-12B, only the global-attention KV cache is compressed.
  • Prefill-Intensive Tasks: Across 12 prefill-intensive datasets, Fast KVzip outperforms all baselines while matching KVzip, which requires twice the prefilling computational cost for compression.The evaluation covers retrieval, contextual understanding, and high-context-redundancy tasks.
  • Cross-Model Generality: Fast KVzip maintains effectiveness across Qwen3-8B, Qwen2.5-14B-1M, Gemma3-12B, and Qwen3-8B-FP8.The comparison averages relative performance across 12 datasets after normalizing each dataset to full-cache performance.
  • Decoding-Intensive Tasks: 4K KV budget size yields near-lossless performance on AIME24 and MATH, outperforming decoding-scenario compression baselines on Qwen3-8B and Qwen3-14B.The early-stopping-of-thinking method drastically reduces reasoning performance in this comparison.

5. Qualitative Analysis

The gating analysis shows that KV retention varies substantially across attention heads. Sparse heads, especially in early layers, favor structural tokens, while other heads retain data-frequent tokens more broadly.

  • Retention Patterns: 36% KV cache budget produces highly sparse retention in half of the attention heads, primarily in early layers.Qwen2.5-7B-1M has 28 layers with 4 KV heads per layer.
  • Interpretation: The observed head behavior suggests that gates prioritize high-level features from middle-to-late layers over low-level features from early layers.This interpretation is based on the distribution of retention rates across attention heads.
  • Token Preferences: Sparse heads predominantly retain punctuation or line-break tokens, whereas medium heads both retain and evict frequent source-data tokens.Head types are defined by retention rate: sparse r < 0.05, medium 0.05 ≤ r < 0.9, and dense 0.9 ≤ r.

6. Conclusion

Fast KVzip is presented as a gating-based KV-cache compression method that preserves task performance with low cache budgets and negligible runtime compression overhead. The authors identify structured gating during pretraining and unified multi-choice gating as future directions.

  • Conclusion: Fast KVzip matches state-of-the-art performance while eliminating runtime compression overhead.The method generalizes across retrieval, code comprehension, and reasoning tasks while reducing KV cache size to 30%.
  • Future Work: Applying the method during pretraining to induce hardware-efficient structured computation patterns remains future work.The authors also propose integrating decisions to skip computation, discard KV features, or store them for selective retrieval.

A. Experimental Details

This section describes the training hyperparameters and gate architecture used in the experiments.

  • Experimental Details: The experimental details cover training hyperparameters and the gate architecture.These details correspond to Sections 3.3 and 3.4.

A.1. Dataset

Fast KVzip trains its gates on 1M tokens from FineWeb-Edu, combining shorter sequences with an additional long-context dataset.

  • 1M tokens from FineWeb-Edu train the gates, comprising 500K tokens from 10K–30K-token sequences and another 500K tokens from 100K-token sequences.The long-context sequences are constructed by concatenating randomly subsampled samples.

A.2. Training Hyperparameters

Gate optimization uses precomputed hidden-state and KV-importance tuples, parallel stochastic-gradient updates, and inference-time eviction policies for prefill and decoding. Training converges stably, while prediction quality varies across layers and token-level comparisons avoid over-smoothing.

  • Training Hyperparameters: Hidden states and KV-importance scores are precomputed, then gates are optimized in parallel with binary cross-entropy for 5K updates at batch size 1K.The learning rate is 0.2, and the schedule corresponds to five epochs over 1M training tokens.
  • Training Hyperparameters: Training loss converges stably, with lower loss in earlier layers and higher loss in middle and later layers.The authors conjecture that higher layers are harder to predict from a single hidden state because their attention mechanisms are more complex.
  • Training Hyperparameters: Token-level prediction comparisons show that the trained gating mechanism captures token-level dynamics and avoids over-smoothing.
  • Inference Configuration: During prefill, eviction follows each 16K-token chunk while retaining the most recent 4K-token window.For contexts shorter than 16K tokens, the method retains the last 2% of tokens.
  • Inference Configuration: During decoding, eviction starts at the cache budget, uses a hidden cache of 128, and preserves a 128-token local attention window.Equivalent numbers of KV entries are evicted after parallel gating decisions to maintain the budget.

B.1. Training Data

Gate performance improves with more training data and long-context examples, requires sufficient sink keys, and benefits from retaining a local window whose size is not critical between 1K and 8K tokens.

  • Training Data: Increasing FineWeb-Edu training data improves gate performance, particularly on SCBench.KV, and adding 100K-token samples is effective.The comparison removes long-context samples to form a 500K-short dataset.
  • Gate Configuration: Gate performance is robust to projection dimension but depends on using a sufficient number of sink keys.
  • Inference Configuration: Retaining a local window improves performance, while results remain comparable for local window sizes from 1K to 8K tokens.All configurations have identical total KV cache size after eviction.

C. Extended Comparison with Prior Gating Approaches

Fast KVzip is compared with prior gating approaches under general-task settings, where it outperforms TrimKV particularly on retrieval tasks while using one dataset across models and evaluations.

  • Comparison with TrimKV: Fast KVzip outperforms TrimKV in general task settings on Qwen3-4B-Instruct-2507, particularly on retrieval tasks.The comparison uses the official TrimKV model available for general language tasks.
  • Training Setup: DMS and TrimKV train separate gates for reasoning and instruction-following tasks, whereas Fast KVzip uses one dataset across models and evaluations.
  • Prediction Comparison: Figure 20 compares Fast KVzip and KVzip token-level predictions on SCBench.En.QA at a 36% KV cache budget ratio.The plotted values are averaged over heads for each layer, using a dataset not used during training.
Loading 2601.17668v2…