Skip to content

README.md

Latest commit

 

History

History
159 lines (128 loc) · 6.46 KB

README.md

File metadata and controls

159 lines (128 loc) · 6.46 KB
Raw

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.