PaperlayerView exact revision on arXiv

Source-linked AI summary

KVzip: Query-Agnostic KV Cache Compression with Context Reconstruction

Jang-Hyun Kim, Jinuk Kim, Sangwoo Kwon, Jae W. Lee, Sangdoo Yun, Hyun Oh Song

arXiv:2505.23416v2cs.DBcs.LG

TL;DR

Long contexts make KV caching costly in memory and attention computation, while query-aware eviction either repeats prefilling or fails to generalize across queries. KVzip instead scores KV pairs through LLM-based context reconstruction and evicts less important pairs, achieving strong multi-query performance with substantially smaller caches and lower decoding latency.

  • Problem

    Longer contexts substantially increase KV-cache memory consumption and attention computation, while existing query-aware eviction methods do not support effective reuse across different queries.

  • Method

    KVzip simulates context reconstruction with teacher-forced LLM decoding and scores each KV pair by its maximum received attention before eviction.

  • Results

    KVzip reduces KV-cache size by up to 70% with negligible performance loss and improves FlashAttention decoding latency by approximately 2× across diverse multi-query tasks and models.

  • Takeaways & Limitations

    KVzip provides a reusable compressed cache that remains robust across diverse queries, long-context benchmarks, quantized models, and KV-cache structures.

  • Takeaways & Limitations

    KVzip’s reconstruction-based compression can cause an instruction-finetuned LLaMA3.1-8B model to answer private-context queries that it refuses with the full KV cache.

Abstract

from arXiv · show

Transformer-based large language models (LLMs) cache context as key-value (KV) pairs during inference. As context length grows, KV cache sizes expand, leading to substantial memory overhead and increased attention latency. This paper introduces KVzip, a query-agnostic KV cache eviction method enabling effective reuse of compressed KV caches across diverse queries. KVzip quantifies the importance of a KV pair using the underlying LLM to reconstruct original contexts from cached KV pairs, subsequently evicting pairs with lower importance. Extensive empirical evaluations demonstrate that KVzip reduces KV cache size by $3$-$4\times$ and FlashAttention decoding latency by approximately $2\times$, with negligible performance loss in question-answering, retrieval, reasoning, and code comprehension tasks. Evaluations include various models such as LLaMA3.1, Qwen2.5, and Gemma3, with context lengths reaching up to 170K tokens. KVzip significantly outperforms existing query-aware KV eviction methods, which suffer from performance degradation even at a 90% cache budget ratio under multi-query scenarios.

1 Introduction

KVzip addresses the memory and attention costs of long-context inference by compressing KV caches once for reuse across diverse queries. Its reconstruction-based, query-agnostic eviction maintains accuracy across multi-query benchmarks while reducing cache size and decoding latency.

  • 33 GB of KV-cache memory is required to cache 120K tokens in Qwen2.5-14B with FP16 precision, exceeding the model’s 28 GB parameter storage.
  • KVzip optimizes a reusable compressed KV cache for each context, targeting offline-prepared caches in conversational agents and enterprise retrieval systems.
  • KVzip scores KV pairs by maximum attention received during LLM-based context reconstruction, then evicts pairs with lower importance.
  • 2× FlashAttention decoding-latency reduction accompanies a 394× reduction in KV-cache size with negligible performance loss on diverse queries.
  • KVzip maintains inference accuracy while evicting up to 70% of the KV cache across question-answering, reasoning, retrieval, and code-comprehension benchmarks with contexts up to 170K tokens.

2 Preliminary

KV caching stores context representations for efficient conditioned generation, but existing eviction methods score KV pairs using query information available during prefill. These query-aware methods require repeated prefills for new queries, while reusing a cache compressed for one query can substantially reduce performance on others.

  • An autoregressive Transformer with L layers and H KV heads caches hidden representations as KV pairs during inference.
  • L × H × n_c KV pairs are generated during prefill for a context tokenized into n_c tokens.
  • Existing methods such as SnapKV and PyramidKV score KV-pair importance from queries in a trailing context window and retain query-relevant pairs.
  • Query-aware eviction requires repetitive cache prefills for each new query, despite effectiveness on single-query benchmarks.
  • Reusing a cache compressed for an initial query significantly reduces performance on different queries, motivating query-agnostic eviction.

3 Method

KVzip scores KV pairs by how strongly they support reconstructing the original context, then evicts low-scoring pairs. Chunked scoring makes this feasible for long contexts while preserving attention efficiency and cross-task usefulness.

  • 3.2 KV Importance Scoring: KVzip uses an LLM forward pass with a repeat prompt and original context to score each KV pair by its maximum attention during reconstruction.The resulting scores determine eviction priorities across KV pairs and can also support head-level eviction.
  • 3.4 Technical Challenge and Solution: KVzip computes chunk-wise attention scores using subsampled keys, softmax normalization, slicing, and a maximum over grouped queries.For later chunks, the repeat prompt includes the last 8 tokens of the preceding chunk.
  • 3.3 Observation: Reconstruction-critical KV features substantially overlap with features used by question-answering, summarization, and reasoning tasks.Cross-task attention distributions concentrate where reconstruction attention is high, unlike the more query-specific overlap between distinct QA tasks.
  • 3.4 Technical Challenge and Solution: Chunked scoring reduces importance-scoring complexity from O(n_c^2) to O(mn_c) while keeping peak memory overhead constant in context length.The context is partitioned into fixed-size chunks, each scored independently; the method uses m = 2K across experiments because chunk size negligibly affects performance.
  • 3.4 Technical Challenge and Solution: Compression introduces approximately twice standard prefill’s computational overhead with under 2% additional memory, while reducing inference latency and cache size.Compression is performed once per context or model, and experiments report negligible performance degradation at compression ratios as low as 30%.

4 Experiment

KVzip is evaluated as a query-agnostic compression method across diverse queries, tasks, models, architectures, and cache settings. Results show strong task generalization, robustness under quantization, and the importance of reconstructing the full context when scoring KV importance.

  • Task Generalization: 30% cache ratio preserves performance on most retrieval-intensive tasks, while baseline methods degrade notably at 90% retention under multi-query evaluation.The evaluation uses Qwen2.5-7B-1M across 12 benchmark datasets grouped by retrieval, contextual-understanding, and redundancy categories.
  • Model Scale and Architecture: KVzip generalizes across larger models, distinct model families, and hybrid attention architectures, outperforming baseline methods in averaged 12-task comparisons.Gemma3 uses global and sliding-window attention, with eviction applied to global layers because they dominate cache size at long contexts.
  • KV Quantization: 4-bit KV quantization leaves KVzip robust, while the base LLaMA3-8B model shows greater contextual sparsity than LLaMA3.1-8B.The same quantization scheme is applied during prefill, importance scoring, and decoding.
  • Context-Independent Eviction: One-time static head-level scoring enables context-independent eviction without post-deployment compression overhead.KVzip derives head scores by aggregating pair-level importance scores from an 88K-token English book sample.
  • Context-Independent Eviction: KVzip’s reconstruction-derived head scores outperform DuoAttention’s passkey-optimized scores while requiring substantially less importance-scoring effort.DuoAttention requires several hours of optimization on an 8-GPU node, whereas KVzip uses context reconstruction on a natural-language textbook.
  • Necessity of Context Reconstruction: Full context reconstruction is essential for preventing degraded importance scoring compared with using only the prompt or partial context.The comparison evaluates the repeat prompt alone, the first or last 10% of context, and the complete context.

5 Related Work

Prior work reduces inference cost through sparse attention, cache compression, or attention-aware eviction. These approaches motivate KVzip’s focus on compressing reusable KV caches without requiring task-specific query information.

  • KV Cache Compression: Sparse Transformer and Compressive Transformer methods reduce KV-cache or attention costs through architectural sparsity or learned cache compression.These approaches generally modify model training or merge cached representations during training.
  • Efficient LLM Inference: Inference-time efficiency methods reduce attention computation through sparse attention, cache offloading, or retrieval rather than directly applying KV-cache eviction.Examples include BigBird, MInference, and Quest.

6 Conclusion

KVzip introduces query-agnostic KV-cache eviction based on reconstructing original context from cached KV pairs. Across multi-query evaluations, it achieves strong compression and latency improvements with negligible performance loss.

  • Conclusion: KVzip reduces KV-cache sizes by up to 70% with negligible performance loss across diverse tasks, models, and long-context benchmarks.The conclusion emphasizes robust compression in multi-query settings.
  • Conclusion: Approximately 2× lower decoding attention latency accompanies KVzip’s compressed-cache performance.The latency improvement is reported for decoding attention.
  • Conclusion: KVzip estimates KV-pair importance by reconstructing the original context, enabling compressed caches to be reused across queries.This query-agnostic design distinguishes the method’s compression objective from query-specific eviction.

A Implementation Details

The implementation precomputes KV importance through chunked context reconstruction, retains the highest-scoring pairs with non-uniform head budgets, and evaluates the method against several eviction baselines.

  • KV Importance Scoring: The scoring algorithm prefills the context KV cache, partitions the context into fixed-size chunks, and reconstructs each chunk using a trailing context span.The chunk size is fixed to 2K tokens, and the reconstruction prompt begins by asking the model to repeat previous context.
  • KV Importance Scoring: Attention received by each chunk’s cached keys is summarized across grouped-query groups and reconstruction positions to produce layer- and head-level importance scores.The algorithm uses maximum attention values before aggregating across sequence positions.
  • Baseline Methods: SnapKV and PyramidKV use official implementations with observation-window, pooling, and layer-budget settings adapted for the evaluated context lengths.SnapKV uses uniform layer budgets, while PyramidKV uses linearly decreasing layer budgets.
  • Baseline Methods: H2O uses maximum prefilling self-attention scores, whereas KVzip uses self-attention scores from context reconstruction.This establishes the principal scoring distinction between the two methods.

B Broader Impacts and Limitations

KVzip improves computational efficiency by compressing KV caches, while the paper notes unresolved information-loss, privacy, and deployment-overhead concerns.

  • KVzip compresses KV caches to reduce computational resources and infrastructure requirements.
  • The study provides no theoretical guarantees concerning compression-induced information loss.
  • KV eviction may raise privacy-leakage concerns through changes in model behavior, although practical implications appear limited under presumed user consent.
  • Context-dependent eviction incurs compression overhead, while context-independent head-level eviction removes deployment overhead but generally achieves lower compression efficiency.

C.1 Reconstruction Chunk Size

The ablations examine reconstruction chunk size, repeat-prompt robustness, scoring overhead, and head-budget allocation, identifying practical trade-offs in KVzip’s design.

  • Reconstruction Chunk Size: Average performance differences remain below 2% at a 0.3 KV cache ratio across scoring chunk sizes.The experiments adopt a 2K chunk size for computational efficiency.
  • Repeat Prompts: KVzip remains robust across original, paraphrased, and absent repeat prompts.The repeat prompt contains 7 tokens, compared with contexts of at least several hundred tokens.
  • Repeat Prompts: 98.1% of KV pairs in a 2K-token NIAH context receive maximum reconstruction attention from the repeated context, rising to 99.4% among retained features after 30% compression.
  • Scoring Variants: Softmax-free scoring removes approximately 10% of scoring overhead but causes approximately a 10% degradation in compression ratios.
  • Head-Budget Allocation: Uniform head-budget allocation outperforms the baseline, while non-uniform allocation achieves superior compression performance.

D Individual Dataset Performance

KVzip is evaluated across models, benchmarks, compression ratios, and multi-task settings, with results generally favoring the method but several model-specific context limits requiring adjusted evaluations.

  • Model Scale and Architecture: Individual-dataset results cover Qwen2.5-14B-1M, LLaMA3.1-8B, Gemma3-12B, and LLaMA3-8B-W8A8KV4 across compression ratios from 0.1 to 1.0.
  • Model Scale and Architecture: Gemma evaluations shorten Retr.KV and Retr.Prefix-Suffix contexts because their approximately 170K-token inputs exceed the model’s 128K-token limit.
  • Model Scale and Architecture: LLaMA3-8B-W8A8KV4 evaluations shorten several datasets because the base model lacks capability on those tasks, producing near-zero accuracy.
  • Multi-Task Datasets: KVzip consistently outperforms baselines on SCBench multi-task datasets and on the smaller LLaMA3.1-3B model.
  • RULER Benchmark: KVzip maintains performance at a 25% compression rate on RULER, while other state-of-the-art eviction methods degrade significantly.
Loading 2505.23416v2…