Skip to content

ersiner/kvqbench

Repository files navigation

kvqbench — tensor-level KV-cache quantization benchmark

A reproducible, tensor/kernel-level benchmark that compares KV-cache quantization methods as first-class peers on the metrics that can be isolated from a full model: reconstruction error, decode-time error accumulation, outlier structure, and GPU throughput/memory.

Scope. Tensor-level KV-cache quantization metrics only — K/V reconstruction, decode-time error accumulation, outlier structure, and memory/throughput. No end-to-end model accuracy (MATH500 / AIME24 / HumanEval / …) is measured; that needs full autoregressive generation with a vLLM backend and ~50 GPU-days, and is out of scope (see What is NOT covered).

Every method runs through the same Quantizer ABC and the same experiments, so no method is privileged. Results are reported as per-metric per-method leaderboards — not as pass/fail against any single paper — alongside a per-method fidelity tag describing how faithfully each implementation tracks its published form.

Origin: kvqbench began as an implementation of the tensor-level benchmark suite described in the KVarN study ("Variance-Normalized KV-Cache Quantization Mitigates Error Accumulation in Reasoning Tasks", Müller et al., Huawei CSL, arXiv:2606.03458, code huawei-csl/KVarN) — its experiments, metrics and method set were derived from that paper. It has since moved to be an independent, method-neutral benchmark where KVarN is one method among eleven, with no privileged status (see Acknowledgments).

Everything is explicitly seeded, runs on a single modest GPU (or CPU for a smoke test), and emits human-readable tables plus machine-readable JSON. A local run writes the full results/summary.json + results/leaderboard.json (gitignored); the committed, provenance-clean leaderboards from one reference H100 run live in results/published/ (synthetic + real, kept separate) and are the in-repo source of truth behind RESULTS.md.


The 11 methods

Registry id → class (kvqbench/methods/__init__.py). Fidelity is each implementation's faithfulness to its published form (faithful = numerically matches the paper / reference; approximate = deliberately simplified — deltas below). Sources + simplifications for every method are in docs/PAPERS.md; to add your own, see docs/ADDING_METHODS.md.

id class K granularity V granularity rotation fidelity bits/elem (k2v2)
fp16 FP16Ref identity (lossless ref) identity (lossless ref) no faithful 16.0
rtn_tensor RTNPerTensor per-tensor asym RTN per-tensor asym RTN no faithful ~2.0
rtn_token RTNPerToken per-token asym RTN per-token asym RTN no faithful ~2.x
rtn_channel RTNPerChannel per-channel asym RTN per-channel asym RTN no faithful ~2.x
kivi KIVI per-channel grouped per-token grouped no faithful 2.1875
quarot QuaRot Hadamard(D) → RTN/head_dim Hadamard(D) → RTN/head_dim yes approximate ~2.19
kvquant KVQuant per-channel NUQ + 1% fp16 per-token NUQ + 1% fp16 no approximate ~2.6 (NUQ+sparse)
kvarn KVarN Hadamard → VarN → per-row RTN Hadamard → VarN → per-row RTN yes faithful 2.25
polarquant PolarQuant per-pair polar (radius+angle) per-token RTN no approximate ~k bits + V
turboquant TurboQuant Hadamard → Lloyd-Max codebook Hadamard → Lloyd-Max codebook yes approximate ~2.x
kitty Kitty magnitude-ranked mixed precision per-token RTN no approximate ~2.5 (mixed)

The fidelity tag and its deltas are first-class metadata on the Quantizer ABC (methods/base.py) and appear in every run's summary.json["method_fidelity"]. The approximate methods carry concrete deltas (full text in docs/PAPERS.md):

  • QuaRot — deterministic Sylvester–Walsh Hadamard instead of QuaRot's randomized fast-Hadamard kernel; cache-only, so the offline V-rotation fusion (here applied directly to the V tile) and the Q-side rotation that cancels K's inside Q·Kᵀ are not modeled. (Axes now match the paper: head-wise Hadamard on K and V + asymmetric groupwise RTN over head_dim, group=128, 0.95 clip ratio.)
  • KVQuant — data-driven non-uniform NUQ codebook fit by unweighted (MSE) Lloyd on the tensor at hand, in place of KVQuant's Fisher-sensitivity-weighted k-means (Fisher needs gradients on a calibration set); ~1% fp16 dense-and-sparse outliers; pre-RoPE quantization omitted.
  • PolarQuant — canonical method is K-only (V fp16); here V is quantized (per-token RTN) to join the bit-width sweep; omits the decode-time LUT.
  • TurboQuant — a deterministic Walsh–Hadamard rotation stands in for the QR-of-Gaussian random rotation; the 1-bit QJL residual is omitted (the rotation + Lloyd-Max codebook, the defining pieces, are kept).
  • Kitty — omits the page-centric two-tensor layout + Triton kernels (system-only; numerics-neutral).

Conventions that affect accounting only:

  • Aux precision — scales/zeros are stored fp16 for exact round-trip numerics, while bits_per_element uses an FP8-scale + FP16-zp budget. Affects Exp4 memory accounting only.
  • bits >= 16 is an fp16 passthrough on every method → exact round-trip (MSE == 0), used as a data-quality sanity gate.

Method mechanisms in brief

  • RTN (round-to-nearest, three granularities) — asymmetric affine quantization with one scale/zero per tensor, per token, or per channel; the granularity ladder from coarsest to finest.
  • KIVI — grouped asymmetric RTN: per-channel for K (group along tokens), per-token for V (group along channels).
  • QuaRot — a Walsh–Hadamard rotation along head_dim (on both K and V) spreads per-channel outliers, then asymmetric groupwise RTN over head_dim in the rotated frame (inverse applied on dequant).
  • KVQuant — a data-driven non-uniform (NUQ) codebook over the per-vector normalized body plus a sparse fp16 outlier side-store (see deltas).
  • KVarN — Hadamard rotation → dual-axis variance normalization (VarN, a Sinkhorn-like log-domain std balancing that equalizes per-token and per-channel variance) → asymmetric per-row RTN → scale absorption (folds the RTN scale/zero into the matching VarN axis so dequant costs ~one extra multiply/element).
  • PolarQuant — re-encodes consecutive key-dimension pairs in polar form (per-channel radius + uniform angle), spreading channel outliers across a smooth representation.
  • TurboQuant — rotates (Hadamard) so coordinates are near-Gaussian, then a per-coordinate Lloyd-Max (Gaussian-optimal) codebook beats uniform RTN.
  • Kitty — keeps the key cache at a low bit-width but boosts the highest-magnitude channels (top fraction) to a higher bit-width.

The 4 experiments

Each lives in kvqbench/experiments/exp{N}_*.py and exposes run(args, device) -> dict. The harness ranks methods per metric.

# module what it measures
1 exp1_reconstruction round-trip MSE vs bit-width {2,3,4,8,16}; per-token magnitude/direction (E_M/E_D) decomposition; E_M/E_T at error quantiles; K-vs-V error split; reconstruction bias (signed magnitude shift); codebook-generalization gap + group-size sensitivity.
2 exp2_accumulation attention-output MAE/MSE/cosine-drift vs context length, static vs pseudo-decode regimes; signed gap series + accumulation slope; attention-distribution KL (fwd + rev); inner-product bias/variance; attention-weighted reconstruction error.
3 exp3_outliers per-token magnitude-diff density; ‖K_dq‖ vs ‖K‖ joint-density off-diagonal mass; E_M/E_T at top quantiles; per-channel incoherence pre/post Hadamard.
4 exp4_throughput quantize / dequant latency; bits/element (nominal + effective) + peak memory; compression vs fp16.

exp3 is presently a KVarN-component ablation (KIVI vs Hadamard-only vs VarN-only vs full KVarN) and is reported as diagnostics only; it joins the symmetric per-method leaderboard once it sweeps all methods.

Output-impact-aware metrics

Layered on top of the raw experiments (and ranked the same neutral way) is a set of calibration-free, single-attention-layer metrics that move from raw reconstruction toward does the error matter — still tensor-level (no full-model inference), still computed identically for every method through the Quantizer round-trip (a method that doesn't "do" the thing lands at the trivial value). See docs/ROADMAP.md for the full item→key map.

  • Attention-distribution KLKL(softmax(q·Kᵀ/√d) ‖ softmax(q·K̂ᵀ/√d)) and its reverse: how much quantization reshapes which tokens attention lands on. Up-weights the tail errors average MAE/MSE flatten.
  • Inner-product bias + variance — splits the logit error q·k̂ − q·k into bias² + variance, separating unbiased-but-noisy (QJL/TurboQuant) from biased-but-tight (asymmetric RTN, KVarN scale-absorption) — identical under a plain MSE.
  • Attention-weighted reconstruction error — weights each token's K error by the attention mass it actually receives (a calibration-free stand-in for Fisher sensitivity); rewards methods that protect high-attention tokens.
  • Reconstruction bias, accumulation slope d(error)/d(context), codebook-generalization gap (in-sample vs cross-batch transfer; ~0 for purely local-per-group methods, positive for data-fit codebooks / token-global statistics), and a group-size sensitivity sweep.

How results are reported

run_all.py ranks every method on each experiment's metrics and writes per-metric per-method leaderboards to results/summary.json["rankings"] and a standalone results/leaderboard.json. Each metric carries an explicit better-direction (e.g. reconstruction MSE: lower is better; compression: higher is better) and ranks all methods that ran — no method is privileged and no paper threshold is applied.

Alongside the rankings, the harness reports method-agnostic data-quality flags (summary.json["data_quality"]): fp16 round-trips to MSE == 0, every method is exact at 16-bit, reconstruction MSE is monotonic in bit-width, the fp16 control accumulates ~0 error (and ~0 attention-distribution KL), and a purely local-per-group method shows a byte-zero codebook-generalization gap (a neutrality sanity check on the new transfer-test metric). These are sanity checks on the measurement, never a verdict on a method.

Budget-aware comparison + multi-seed CI

summary.json["pareto"] (and leaderboard.json) carry a quality-vs-bits/element Pareto built from exp1 (each (method, bit-width) cell carries both its bits_per_element and reconstruction MSE) plus a matched-budget selector — the lowest-MSE method at each target bits/elem budget. This compares methods at an equal effective budget rather than at a nominal bit-width, which is the fair way to rank mixed-precision methods (e.g. KVQuant's fp16 outliers cost extra bits). The quality_vs_bits_pareto.png figure plots it.

For run-to-run variance, multi_seed_runner.py re-runs the quality experiments over N seeds and reports mean ± 95% CI per metric (ranked by mean):

python multi_seed_runner.py --n-seeds 5 --experiments 1,2 --out results/
# -> results/multi_seed_leaderboard.json

What is NOT covered (and why)

End-to-end accuracy on MATH500 / AIME24 / HumanEval / IFEval / Line-Retrieval / NiaH, plus the end-to-end output-distribution KL (over the model's vocabulary, across all layers and the full generation), is not isolatable at the tensor level. It needs full autoregressive generation with the quantizer wired into a model's attention KV path (vLLM backend), multiple models, and ~50 GPU-days. It is intentionally omitted; the tensor-level suite here measures the mechanism (reconstruction, outliers, accumulation, cost) that those tables depend on.

What is now isolated — a single-attention-layer attention-distribution KL (exp2's attn_kl), the tensor-level bridge toward that end-to-end KL — measures how quantization reshapes one layer's attention weights without running the model.


Data sources

kvqbench/data.py:

  • Synthetic (default, always available)make_synthetic_kv builds [B, Hkv, T, D] tensors with heavy-tailed per-token norms (lognormal σ=1.0) × random unit directions, plus persistent K channels scaled by a channel_gain (the per-channel outlier structure). V omits the channel outliers. Fully seeded.
  • Real (optional, --real)capture_real_kv does one HF forward pass of Qwen/Qwen3-4B (fallback meta-llama/Llama-3.2-1B), hooks k_proj/v_proj per layer, and returns stacked [num_layers, Hkv, T, D]. Requires transformers; degrades gracefully with a clear RuntimeError so the synthetic path always works.

Running locally

# 0) install (torch + numpy; add optional extras as desired)
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
# optional: pip install transformers accelerate matplotlib

# 1) fast self-check (tiny shapes, seconds; CPU is fine)
python run_all.py --smoke

# 2) full synthetic run, all four experiments
python run_all.py --experiments 1,2,3,4 --out results/

# 3) restrict to a subset of methods
python run_all.py --experiments 1,2,4 --methods kivi,quarot,kvarn

# 4) also capture real KV (needs transformers + a model + VRAM)
python run_all.py --experiments 1,2,3,4 --real --out results/

# 5) render figures (needs matplotlib; never blocks the run)
python run_all.py --experiments 1,2,3,4 --plots
#   or standalone:
python results/plots.py results/

# 6) re-render tables + leaderboard from saved JSON, no torch needed
python run_all.py --report-only results/

CLI flags (run_all.py): --experiments 1,2,3,4 · --methods a,b,c · --smoke · --real · --device auto|cpu|cuda|cuda:N · --out results/ · --seed N · --plots · --report-only DIR · --no-tables. Outputs: results/exp{N}_{synthetic|real}.json, results/summary.json (rankings + data-quality + method fidelity), and results/leaderboard.json.

A modest GPU suffices for everything (an 8-head × 32768 × 128 fp16 K tensor is ~67 MB; 24 GB is comfortable). CUDA bf16/fp16 tensor cores matter only for the Exp4 kernel timings; on CPU those numbers are recorded but flagged unreliable.


Running on RunPod

One command does everything (create pod → sync → run → pull results → terminate):

# API key in $RUNPOD_API_KEY or ~/.runpod_key first, then:
scripts/run_on_runpod.sh                                  # full run, auto-terminate
scripts/run_on_runpod.sh --keep                           # leave the pod up afterwards
scripts/run_on_runpod.sh --gpu "NVIDIA H100 80GB HBM3" --cloud SECURE

Results + figures land in results/. The pieces it orchestrates, if you want to run them by hand:

  • scripts/runpod_ctl.py — self-contained (stdlib-only) RunPod GraphQL CLI: create / status <id> / ssh <id> / terminate <id>. Reads the API key from RUNPOD_API_KEY or ~/.runpod_key (never hardcoded), injects your public key from remote/id_ed25519.pub, and provisions an on-demand GPU pod with SSH exposed.
  • scripts/remote_run.sh — runs on the pod: prints nvidia-smi + torch/CUDA versions, pip-installs requirements.txt (+ optional transformers/matplotlib), runs python run_all.py --smoke, then on success the full --experiments 1,2,3,4 and --real if a model is reachable.

The pinned image targets torch 2.4.0 / CUDA 12.4.1; any Hopper/Ada GPU works. (Blackwell / sm_120 GPUs need a newer image with torch ≥ 2.7 / CUDA ≥ 12.8.)


Repo layout

<repo root>/
  run_all.py                       # harness CLI (entry point)
  multi_seed_runner.py             # multi-seed mean±CI leaderboard
  requirements.txt
  README.md   RESULTS.md           # overview + results digest
  LICENSE     NOTICE               # Apache-2.0 text + attribution
  docs/
    PAPERS.md                      # per-method sources + simplifications
    ADDING_METHODS.md              # how to add a quantizer
  kvqbench/
    common.py                      # hadamard, packing, RTN, Lloyd-Max, VarN, metrics
    data.py                        # synthetic + real KV
    SPEC.txt                       # benchmark methodology (paper distillation)
    methods/                       # base ABC + 11 quantizers + registry
      kvarn/reference/             # vendored huawei-csl KVarN source (Apache-2.0 provenance)
    experiments/                   # exp1..exp4 + kvarn_deepdive, each run(args, device)->dict
  results/                         # generated output (gitignored, except below)
    plots.py                       # matplotlib figures (tracked source)
    exp{N}_{synthetic|real}.json   # per-experiment output (gitignored)
    summary.json / summary_real.json       # rankings + pareto + dq + fidelity (gitignored)
    leaderboard.json / leaderboard_real.json  # standalone deliverable (gitignored)
    multi_seed_leaderboard.json    # mean±CI from multi_seed_runner.py (gitignored)
    published/                     # TRACKED, provenance-clean leaderboards (see its README)
      leaderboard_{synthetic,real}.json
      multi_seed_leaderboard.json
  scripts/
    run_on_runpod.sh               # one command: create pod → sync → run → pull → terminate
    remote_run.sh                  # runs ON the pod
    runpod_ctl.py                  # RunPod GraphQL control CLI (local)
    publish_results.py             # derive results/published/ from saved JSON (torch-free)
  remote/                          # SSH keypair + API key (gitignored, never committed)

Reproducibility & fidelity

All randomness derives from --seed (forwarded into every experiment); there is no wall-clock or unseeded RNG. Data-quality gates hold across runs: at bits >= 16 every method round-trips to MSE == 0, and reconstruction MSE is monotonic in bit-width. Confirmed bit budgets: kvarn k2v2 = 2.25, kivi k2v2 = 2.1875, fp16 = 16.0.

Fidelity is verified, not asserted: KVarN matches the vendored huawei-csl reference (kvqbench/methods/kvarn/reference/) exactly (round-trip error 0.0, VarN byte-exact against the reference Sinkhorn), and the faithful/approximate tags above record where each implementation stands. See kvqbench/SPEC.txt for the full methodology and per-method notes.


License

kvqbench is released under the Apache License 2.0 (see LICENSE).

It bundles, verbatim, Apache-2.0 source from the vLLM project — obtained via the KVarN repo (huawei-csl/KVarN, a vLLM fork) — under kvqbench/methods/kvarn/reference/, kept as provenance for the kvarn method (the implementation in kvqbench/methods/kvarn.py is an independent reimplementation). Those files retain their original SPDX Apache-2.0 headers; attribution is in NOTICE. Per-method sources are in docs/PAPERS.md.


Acknowledgments

kvqbench started as an implementation of the tensor-level benchmark suite described in the KVarN study ("Variance-Normalized KV-Cache Quantization Mitigates Error Accumulation in Reasoning Tasks", Müller et al., Huawei CSL, arXiv:2606.03458) — its experiments, metrics, and method set were derived from that paper and its open-source reference implementation, huawei-csl/KVarN. It has since moved to be an independent, method-neutral benchmark: KVarN is now one method among eleven, with no privileged status and no paper-reproduction pass/fail goal.

Grateful acknowledgment to the KVarN authors — their paper defined the experimental protocol this benchmark generalizes, and their Apache-2.0 release is vendored under kvqbench/methods/kvarn/reference/ as the kvarn provenance reference.

About

Neutral tensor-level KV-cache quantization benchmark: 11 methods as first-class peers, per-metric leaderboards, matched-budget Pareto, multi-seed CI, and per-method fidelity tags.

Topics

Resources

License

Stars

0 stars

Watchers

0 watching

Forks

Packages

 
 
 

Contributors