KV-Surgeon: Compensation-Aware Offline KV-Cache Compaction

Code for KV-Surgeon.

Offline compaction builds a smaller KV cache for a fixed context once and reuses it for every later query. KV-Surgeon decides which cache entries to delete with a quadratic surrogate of the attention-output error. It deletes entries greedily, and after each deletion it updates the remaining entries in closed form to compensate. A final refit sets the retained keys' logit biases and values.

This repository is built on the Attention Matching code base (Zweiger et al., 2026). It also contains the baselines and ablations from the paper.

Installation

Python 3.12 and a CUDA GPU are required.

pip install -r requirements.txt

Quick start

Compact the KV cache of a short article with KV-Surgeon and compare answers before and after:

python -m examples.qa_demo --model Qwen/Qwen3-4B --target-size 0.02

Evaluate on QuALITY without wandb:

python -m evaluation.run_qa_evaluation \
  --model-name Qwen/Qwen3-4B \
  --algorithm-config kv_surgeon --methods original kv_surgeon \
  --query-config ss-plus-repeat \
  --precomputed-budget-path head_budget_optimization/head_budgets/Qwen3-4B/optimized_agnostic.json \
  --dataset-name quality --n-articles 1 --target-size 0.02

For KV-Surgeon-fast use --methods kv_surgeon_fast --query-config repeat. --target-size 0.02 keeps 2% of the context (50x compaction).

Reproducing the paper

Every experiment is a wandb grid sweep in sweeps/. See sweeps/README.md for the layout and how to launch them.

Experiment	Sweeps
Main results, all models and ratios	sweeps/frontier/<model>/{quality,longhealth}/
Long Code Arena	sweeps/frontier/qwen3-4b/lca/
Component ablations	sweeps/ablation_*.yaml
Head budgets and final refit	sweeps/frontier/qwen3-4b/quality/uniform/
Ported baselines, published vs. our pipeline	sweeps/frontier/qwen3-4b/quality/{kvzip.yaml,uniform/,optimized/}

Each run writes a result JSON to the log_dir of its sweep. The ablation tables are aggregated with scripts/aggregate_refit_budget_ablation.py and scripts/aggregate_selection_ablations.py. The loss-level analyses of compensation use scripts/surrogate_check.py, scripts/kv_surgeon_oracle.py and scripts/refresh_diag.py.

Data.

• QuALITY is included.
• LongHealth is downloaded from its repository on first use.
• Long Code Arena must be built once: python scripts/build_lca_repo_completion.py.

Precomputed files.

• Per-head budgets: head_budget_optimization/head_budgets/<Model>/optimized_agnostic.json. Recompute them with scripts/head-budget-optimization.sh, then run the solve step described in head_budget_optimization/run.py.
• Q-Filters vectors: qfilters/pile/<Model>.pt. Recompute them with scripts/compute_qfilters.py.

The reference answers used for perplexity are generated and cached on first use. They are sampled, so perplexities can differ slightly from ours.

Methods

Name in the paper	Algorithm	Reference
KV-Surgeon, KV-Surgeon-fast	kv_surgeon	this work
LOO	loo_evict	this work; generalizes CAOTE and ReST-KV
AM-HAK, AM-OMP, AM-OMP-CS	highest_attention_keys, omp	Attention Matching
OBCache	loo_evict (score: obcache_k)	Gu et al., 2025
H2O	highest_attention_keys (score_method: mean)	Zhang et al., 2023
KVzip	highest_attention_keys (score_method: max), global_highest_attention_keys	Kim et al., 2025
KVMerger	kvmerger	Wang et al., 2024
Q-Filters	qfilters	Godey et al., 2025
StreamingLLM	magnitude_evict (score: recency)	Xiao et al., 2024
L2 key-norm	magnitude_evict (score: norm)	Devoto et al., 2024
Summarization	summarize

Calibration queries come from self-study (Cartridges) or a repeat pass (KVzip).

Benchmarks: QuALITY, LongHealth, Long Code Arena. Models: Qwen3, Llama 3.1, Gemma 3. The Q-Filters vectors are calibrated on the Pile (NeelNanda/pile-10k). Generation uses vLLM and Transformers.

Repository layout

Path	Contents
compaction/algorithms/	Per-head compaction algorithms
compaction/compaction_methods/	Whole-cache wrappers: per-head budgets, calibration queries, chunking
compaction/query_generation/	Self-study, repeat-pass and prefill calibration queries
evaluation/	Datasets, QA evaluation, algorithm and query configs
head_budget_optimization/	Per-head budget computation and the budget files
models/	Qwen3, Qwen3-MoE, Llama and Gemma 3 with compacted-cache support
sweeps/	Experiment definitions
scripts/	Data building, Q-Filters, ablation aggregation and analyses
main.py, controller.py, runner.sh	wandb sweep entry point and SLURM launcher

Citation

If you use this code, please cite KV-Surgeon and Attention Matching:

@misc{wagner2026kvsurgeon,
      title={The KV-Surgeon: Compensation-Aware Offline KV-Cache Compaction},
      author={Moritz Wagner and Christophe Roux and Max Zimmer and Sebastian Pokutta},
      year={2026},
      eprint={2610.25933},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2610.25933},
}

@misc{zweiger2026fastkvcompaction,
      title={Fast {KV} Compaction via {Attention Matching}},
      author={Adam Zweiger and Xinghong Fu and Han Guo and Yoon Kim},
      year={2026},
      eprint={2602.16284},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2602.16284},
}

License

MIT, see LICENSE. Some files are under other licenses; see NOTICE.