AttnVQ submission

.gitignore

@@ -0,0 +1,2 @@
+.DS_Store
+.DS_Store

README.md

@@ -1,3 +1,91 @@
 ---
 license: apache-2.0
 ---
+
+# AttnVQ — Attention-Aware KV Cache Quantization
+
+Training-free **product vector quantization** of the KV cache for long-context LLMs. AttnVQ fits small per-subspace codebooks with LBG, but scores distortion by **attention-output error** (and key cosine / inner-product bias), not cache MSE. Calibration is light: **10–15 agent traces, ~15 s on GPU** — enough to capture the **model's** K/V geometry (data-aware, not corpus-dependent).
+
+Primary target: **Laguna-XS.2** (model-agnostic). Only the **10 full-attention layers** are compressed; 30 sliding-window layers stay fp16.
+
+## What this repo offers
+
+| Component | Description |
+|---|---|
+| **`generate.py`** | Minimal inference: `VQQuantizedCache` → `model.generate()` (uint8 indices, real memory savings) |
+| **`vqkv/`** | Quantizers (ProductVQ, RoPESplit, scalar/KIVI baselines), attention-aware metrics, compressed cache |
+| **`benchmark.py`** | Fit codebooks + **cheap metrics** (key cosine, attn-output error, ip-bias) on real cache dumps |
+| **`turbo_benchmark.py`** | Faithful **TurboQuant** baseline (Haar rotation + Lloyd-Max + QJL) |
+| **`longbench_eval.py`** | LongBench v1 proxy metrics + optional end-to-end task scoring |
+| **`artifacts/`** | Pre-fit codebooks and LongBench results |
+
+**Variants:** `productvq-*` (AttnVQ), `ropesplit-1b` (RoPE-half split for Laguna), scalar/KIVI/sign/ternary baselines, TurboQuant MSE/Prod.
+
+## Headline results (Laguna-XS.2)
+
+**Memory @ 131K context** (full-attention layers only):
+
+| Config | KV cache |
+|---|---|
+| fp16 | 5.4 GB |
+| AttnVQ 2-bit (`productvq-32x256-2b`) | 0.73 GB (7.4×) |
+| AttnVQ 1-bit (`productvq-16x256-1b`) | 0.40 GB (14×) |
+
+**LongBench v1** (mean F1 over qasper, 2wikimqa, hotpotqa, repobench-p; single 15-trace codebook):
+
+- **2-bit:** ~96% of fp16 — TurboQuant ~83%, INT2 ~75%
+- **1-bit:** AttnVQ and RoPESplit beat every iso-budget baseline on every task
+- **0.5-bit:** only VQ reaches this regime at all
+
+Full numbers: `artifacts/longbench_results.json`, `artifacts/longbench_cheap_metrics.json`.
+
+**Note:**  Wall-clock speedup requires a fused dequant kernel.
+
+
+## Quick start
+Tested on CUDA 12.4 / NVIDIA A100.
+
+```bash
+pip install "git+https://github.com/huggingface/transformers.git" \
+    accelerate datasets torch==2.9.1 torchvision tqdm
+python generate.py
+```
+
+Use fitted codebooks for memory efficient long context generation:
+```py
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from vqkv.compressed_cache import VQQuantizedCache
+
+
+# load model
+tok = AutoTokenizer.from_pretrained("poolside/Laguna-XS.2", trust_remote_code=True, fix_mistral_regex=True)
+model = AutoModelForCausalLM.from_pretrained("poolside/Laguna-XS.2", torch_dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True).eval()
+
+# load codebooks or fit and use your own
+CODEBOOKS_PATH = "artifacts/codebooks.pt"
+codebooks = torch.load(CODEBOOKS_PATH, map_location="cuda", weights_only=False)
+
+# build cache
+quantizers, layers = codebooks["fitted"]["productvq-32x256-2b"], blob["meta"]["full_layers"]
+cache = VQQuantizedCache(quantizers, layers)  # persists uint8 codebook indices
+
+# generate
+ids = tok("Hello", return_tensors="pt").to(model.device)
+out = model.generate(**ids, max_new_tokens=32, past_key_values=cache, use_cache=True) 
+print(tok.decode(out[0, ids["input_ids"].shape[1]:], skip_special_tokens=True))
+
+# print memory footprint
+print(cache.memory_footprint())
+```
+
+## Reproduce
+
+Precomputed results are under `artifacts/`. To re-fit and evaluate:
+
+```bash
+python benchmark.py --stage fit
+python turbo_benchmark.py --stage fit 
+python longbench_eval.py --stage cheap --n_eval 50  # cheap metrics
+python longbench_eval.py --stage generate --n_eval 50   # slow: full generation & task metrics
+```

artifacts/codebooks.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:30f0ac721169cfc655f1d6ca5ac07bc817ec434b6149437fb247e5052be11028
+size 17489175

artifacts/longbench_cheap_metrics.json

@@ -0,0 +1,1192 @@
+[
+  {
+    "config": "scalar-int4",
+    "bits_per_elt": 4.25,
+    "n_traces": 50,
+    "key_cos": 0.99257,
+    "val_cos": 0.99503,
+    "key_mse": 0.02338,
+    "val_mse": 0.0001,
+    "attn_cos": 0.99966,
+    "attn_output_error": 0.02973,
+    "ip_rel": 1.34455,
+    "ip_bias": 0.0001,
+    "task": "qasper"
+  },
+  {
+    "config": "scalar-int2",
+    "bits_per_elt": 2.25,
+    "n_traces": 50,
+    "key_cos": 0.86766,
+    "val_cos": 0.89301,
+    "key_mse": 0.46759,
+    "val_mse": 0.00257,
+    "attn_cos": 0.99046,
+    "attn_output_error": 0.15708,
+    "ip_rel": 4.9622,
+    "ip_bias": 0.00096,
+    "task": "qasper"
+  },
+  {
+    "config": "kivi-int2",
+    "bits_per_elt": 2.25,
+    "n_traces": 50,
+    "key_cos": 0.86152,
+    "val_cos": 0.89301,
+    "key_mse": 0.54308,
+    "val_mse": 0.00257,
+    "attn_cos": 0.99043,
+    "attn_output_error": 0.15729,
+    "ip_rel": 5.98639,
+    "ip_bias": -5e-05,
+    "task": "qasper"
+  },
+  {
+    "config": "turboquant~-int2",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.64016,
+    "val_cos": 0.64168,
+    "key_mse": 2.17849,
+    "val_mse": 0.01394,
+    "attn_cos": 0.36733,
+    "attn_output_error": 2.47328,
+    "ip_rel": 9.66254,
+    "ip_bias": -0.00053,
+    "task": "qasper"
+  },
+  {
+    "config": "productvq-16x256-1b",
+    "bits_per_elt": 1.0,
+    "n_traces": 50,
+    "key_cos": 0.8606,
+    "val_cos": 0.82426,
+    "key_mse": 0.43329,
+    "val_mse": 0.00342,
+    "attn_cos": 0.97447,
+    "attn_output_error": 0.33079,
+    "ip_rel": 4.05013,
+    "ip_bias": 0.00015,
+    "task": "qasper"
+  },
+  {
+    "config": "productvq-8x256-0.5b",
+    "bits_per_elt": 0.5,
+    "n_traces": 50,
+    "key_cos": 0.73386,
+    "val_cos": 0.67074,
+    "key_mse": 0.76039,
+    "val_mse": 0.00583,
+    "attn_cos": 0.91293,
+    "attn_output_error": 0.50347,
+    "ip_rel": 4.9189,
+    "ip_bias": 0.00143,
+    "task": "qasper"
+  },
+  {
+    "config": "ropesplit-1b",
+    "bits_per_elt": 1.0,
+    "n_traces": 50,
+    "key_cos": 0.84851,
+    "val_cos": 0.82369,
+    "key_mse": 0.46147,
+    "val_mse": 0.00342,
+    "attn_cos": 0.97455,
+    "attn_output_error": 0.33236,
+    "ip_rel": 3.94593,
+    "ip_bias": 0.00051,
+    "task": "qasper"
+  },
+  {
+    "config": "sign-1bit",
+    "bits_per_elt": 1.125,
+    "n_traces": 50,
+    "key_cos": 0.82715,
+    "val_cos": 0.79949,
+    "key_mse": 0.51141,
+    "val_mse": 0.00367,
+    "attn_cos": 0.97079,
+    "attn_output_error": 0.40264,
+    "ip_rel": 3.99374,
+    "ip_bias": 0.00027,
+    "task": "qasper"
+  },
+  {
+    "config": "ternary-bitnet",
+    "bits_per_elt": 1.71,
+    "n_traces": 50,
+    "key_cos": 0.91306,
+    "val_cos": 0.90047,
+    "key_mse": 0.26966,
+    "val_mse": 0.00192,
+    "attn_cos": 0.98765,
+    "attn_output_error": 0.2377,
+    "ip_rel": 3.23327,
+    "ip_bias": -6e-05,
+    "task": "qasper"
+  },
+  {
+    "config": "productvq-64x256-4b",
+    "bits_per_elt": 4.0,
+    "n_traces": 50,
+    "key_cos": 0.99599,
+    "val_cos": 0.99485,
+    "key_mse": 0.01346,
+    "val_mse": 0.00012,
+    "attn_cos": 0.99959,
+    "attn_output_error": 0.03627,
+    "ip_rel": 0.76086,
+    "ip_bias": -2e-05,
+    "task": "qasper"
+  },
+  {
+    "config": "productvq-32x256-2b",
+    "bits_per_elt": 2.0,
+    "n_traces": 50,
+    "key_cos": 0.95936,
+    "val_cos": 0.94781,
+    "key_mse": 0.13408,
+    "val_mse": 0.00108,
+    "attn_cos": 0.99537,
+    "attn_output_error": 0.13766,
+    "ip_rel": 2.72405,
+    "ip_bias": 0.00031,
+    "task": "qasper"
+  },
+  {
+    "config": "turbo-mse-4b",
+    "bits_per_elt": 4.125,
+    "n_traces": 50,
+    "key_cos": 0.99504,
+    "val_cos": 0.99504,
+    "key_mse": 0.01612,
+    "val_mse": 0.0001,
+    "attn_cos": 0.9996,
+    "attn_output_error": 0.03475,
+    "ip_rel": 0.9187,
+    "ip_bias": -3e-05,
+    "task": "qasper"
+  },
+  {
+    "config": "turbo-prod-4b (K:+qjl)",
+    "bits_per_elt": 5.125,
+    "n_traces": 50,
+    "key_cos": 0.99504,
+    "val_cos": 0.99504,
+    "key_mse": 0.01613,
+    "val_mse": 0.0001,
+    "attn_cos": 0.9996,
+    "attn_output_error": 0.03478,
+    "ip_rel": 1.04771,
+    "ip_bias": 4e-05,
+    "task": "qasper"
+  },
+  {
+    "config": "turbo-mse-2b",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.94071,
+    "val_cos": 0.9407,
+    "key_mse": 0.18395,
+    "val_mse": 0.00118,
+    "attn_cos": 0.99373,
+    "attn_output_error": 0.16035,
+    "ip_rel": 3.38128,
+    "ip_bias": 0.00046,
+    "task": "qasper"
+  },
+  {
+    "config": "turbo-prod-2b (K:+qjl)",
+    "bits_per_elt": 3.125,
+    "n_traces": 50,
+    "key_cos": 0.94071,
+    "val_cos": 0.94071,
+    "key_mse": 0.18395,
+    "val_mse": 0.00118,
+    "attn_cos": 0.99373,
+    "attn_output_error": 0.1604,
+    "ip_rel": 3.49197,
+    "ip_bias": -0.00021,
+    "task": "qasper"
+  },
+  {
+    "config": "turbo-mse-1b",
+    "bits_per_elt": 1.125,
+    "n_traces": 50,
+    "key_cos": 0.7997,
+    "val_cos": 0.79979,
+    "key_mse": 0.57262,
+    "val_mse": 0.00367,
+    "attn_cos": 0.97083,
+    "attn_output_error": 0.40155,
+    "ip_rel": 5.04924,
+    "ip_bias": -0.00028,
+    "task": "qasper"
+  },
+  {
+    "config": "turbo-prod-1b (K:+qjl)",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.7997,
+    "val_cos": 0.79979,
+    "key_mse": 0.57262,
+    "val_mse": 0.00367,
+    "attn_cos": 0.97084,
+    "attn_output_error": 0.40162,
+    "ip_rel": 7.50776,
+    "ip_bias": -6e-05,
+    "task": "qasper"
+  },
+  {
+    "config": "scalar-int4",
+    "bits_per_elt": 4.25,
+    "n_traces": 47,
+    "key_cos": 0.99254,
+    "val_cos": 0.99503,
+    "key_mse": 0.02381,
+    "val_mse": 0.0001,
+    "attn_cos": 0.99968,
+    "attn_output_error": 0.02836,
+    "ip_rel": 0.95936,
+    "ip_bias": 0.00013,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "scalar-int2",
+    "bits_per_elt": 2.25,
+    "n_traces": 47,
+    "key_cos": 0.86637,
+    "val_cos": 0.89305,
+    "key_mse": 0.48028,
+    "val_mse": 0.00262,
+    "attn_cos": 0.99088,
+    "attn_output_error": 0.15136,
+    "ip_rel": 4.48127,
+    "ip_bias": 0.00021,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "kivi-int2",
+    "bits_per_elt": 2.25,
+    "n_traces": 47,
+    "key_cos": 0.85589,
+    "val_cos": 0.89305,
+    "key_mse": 0.57652,
+    "val_mse": 0.00262,
+    "attn_cos": 0.99085,
+    "attn_output_error": 0.1516,
+    "ip_rel": 5.84995,
+    "ip_bias": -0.0007,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "turboquant~-int2",
+    "bits_per_elt": 2.125,
+    "n_traces": 47,
+    "key_cos": 0.64105,
+    "val_cos": 0.64241,
+    "key_mse": 2.21716,
+    "val_mse": 0.0142,
+    "attn_cos": 0.39052,
+    "attn_output_error": 2.37009,
+    "ip_rel": 13.81316,
+    "ip_bias": 0.0022,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "productvq-16x256-1b",
+    "bits_per_elt": 1.0,
+    "n_traces": 47,
+    "key_cos": 0.86048,
+    "val_cos": 0.82425,
+    "key_mse": 0.44288,
+    "val_mse": 0.00349,
+    "attn_cos": 0.97475,
+    "attn_output_error": 0.34054,
+    "ip_rel": 4.03823,
+    "ip_bias": 0.00044,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "productvq-8x256-0.5b",
+    "bits_per_elt": 0.5,
+    "n_traces": 47,
+    "key_cos": 0.7317,
+    "val_cos": 0.66903,
+    "key_mse": 0.78213,
+    "val_mse": 0.00598,
+    "attn_cos": 0.91045,
+    "attn_output_error": 0.52276,
+    "ip_rel": 4.68996,
+    "ip_bias": 0.00175,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "ropesplit-1b",
+    "bits_per_elt": 1.0,
+    "n_traces": 47,
+    "key_cos": 0.84739,
+    "val_cos": 0.8237,
+    "key_mse": 0.4739,
+    "val_mse": 0.0035,
+    "attn_cos": 0.97487,
+    "attn_output_error": 0.34152,
+    "ip_rel": 4.43052,
+    "ip_bias": 0.00079,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "sign-1bit",
+    "bits_per_elt": 1.125,
+    "n_traces": 47,
+    "key_cos": 0.82765,
+    "val_cos": 0.79952,
+    "key_mse": 0.51975,
+    "val_mse": 0.00374,
+    "attn_cos": 0.97069,
+    "attn_output_error": 0.40168,
+    "ip_rel": 4.59382,
+    "ip_bias": -0.00051,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "ternary-bitnet",
+    "bits_per_elt": 1.71,
+    "n_traces": 47,
+    "key_cos": 0.91339,
+    "val_cos": 0.90049,
+    "key_mse": 0.27376,
+    "val_mse": 0.00196,
+    "attn_cos": 0.98759,
+    "attn_output_error": 0.23654,
+    "ip_rel": 3.65978,
+    "ip_bias": 0.00062,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "productvq-64x256-4b",
+    "bits_per_elt": 4.0,
+    "n_traces": 47,
+    "key_cos": 0.99596,
+    "val_cos": 0.99487,
+    "key_mse": 0.01392,
+    "val_mse": 0.00012,
+    "attn_cos": 0.9996,
+    "attn_output_error": 0.03607,
+    "ip_rel": 0.86073,
+    "ip_bias": -5e-05,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "productvq-32x256-2b",
+    "bits_per_elt": 2.0,
+    "n_traces": 47,
+    "key_cos": 0.95941,
+    "val_cos": 0.94804,
+    "key_mse": 0.13688,
+    "val_mse": 0.0011,
+    "attn_cos": 0.99554,
+    "attn_output_error": 0.13911,
+    "ip_rel": 2.10964,
+    "ip_bias": 5e-05,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "turbo-mse-4b",
+    "bits_per_elt": 4.125,
+    "n_traces": 47,
+    "key_cos": 0.99505,
+    "val_cos": 0.99504,
+    "key_mse": 0.01637,
+    "val_mse": 0.00011,
+    "attn_cos": 0.99961,
+    "attn_output_error": 0.03411,
+    "ip_rel": 0.86796,
+    "ip_bias": 8e-05,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "turbo-prod-4b (K:+qjl)",
+    "bits_per_elt": 5.125,
+    "n_traces": 47,
+    "key_cos": 0.99505,
+    "val_cos": 0.99504,
+    "key_mse": 0.01637,
+    "val_mse": 0.00011,
+    "attn_cos": 0.99961,
+    "attn_output_error": 0.03415,
+    "ip_rel": 1.15955,
+    "ip_bias": 1e-05,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "turbo-mse-2b",
+    "bits_per_elt": 2.125,
+    "n_traces": 47,
+    "key_cos": 0.94077,
+    "val_cos": 0.94072,
+    "key_mse": 0.18695,
+    "val_mse": 0.0012,
+    "attn_cos": 0.99375,
+    "attn_output_error": 0.15968,
+    "ip_rel": 2.78259,
+    "ip_bias": 0.00031,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "turbo-prod-2b (K:+qjl)",
+    "bits_per_elt": 3.125,
+    "n_traces": 47,
+    "key_cos": 0.94077,
+    "val_cos": 0.94072,
+    "key_mse": 0.18696,
+    "val_mse": 0.0012,
+    "attn_cos": 0.99375,
+    "attn_output_error": 0.15975,
+    "ip_rel": 3.76935,
+    "ip_bias": -0.00026,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "turbo-mse-1b",
+    "bits_per_elt": 1.125,
+    "n_traces": 47,
+    "key_cos": 0.79981,
+    "val_cos": 0.79987,
+    "key_mse": 0.58236,
+    "val_mse": 0.00374,
+    "attn_cos": 0.97061,
+    "attn_output_error": 0.40143,
+    "ip_rel": 4.56288,
+    "ip_bias": -0.00021,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "turbo-prod-1b (K:+qjl)",
+    "bits_per_elt": 2.125,
+    "n_traces": 47,
+    "key_cos": 0.79981,
+    "val_cos": 0.79986,
+    "key_mse": 0.58237,
+    "val_mse": 0.00374,
+    "attn_cos": 0.97061,
+    "attn_output_error": 0.40152,
+    "ip_rel": 6.74448,
+    "ip_bias": -0.00105,
+    "task": "2wikimqa"
+  },
+  {
+    "config": "scalar-int4",
+    "bits_per_elt": 4.25,
+    "n_traces": 50,
+    "key_cos": 0.99257,
+    "val_cos": 0.99502,
+    "key_mse": 0.02373,
+    "val_mse": 0.0001,
+    "attn_cos": 0.99969,
+    "attn_output_error": 0.02754,
+    "ip_rel": 1.21463,
+    "ip_bias": -0.0001,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "scalar-int2",
+    "bits_per_elt": 2.25,
+    "n_traces": 50,
+    "key_cos": 0.86642,
+    "val_cos": 0.893,
+    "key_mse": 0.48208,
+    "val_mse": 0.0026,
+    "attn_cos": 0.99106,
+    "attn_output_error": 0.14817,
+    "ip_rel": 4.89558,
+    "ip_bias": 0.00069,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "kivi-int2",
+    "bits_per_elt": 2.25,
+    "n_traces": 50,
+    "key_cos": 0.84639,
+    "val_cos": 0.893,
+    "key_mse": 0.61765,
+    "val_mse": 0.0026,
+    "attn_cos": 0.99102,
+    "attn_output_error": 0.14845,
+    "ip_rel": 6.11069,
+    "ip_bias": 0.0004,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "turboquant~-int2",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.64118,
+    "val_cos": 0.64228,
+    "key_mse": 2.22279,
+    "val_mse": 0.01408,
+    "attn_cos": 0.3972,
+    "attn_output_error": 2.2878,
+    "ip_rel": 10.22977,
+    "ip_bias": 0.00486,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "productvq-16x256-1b",
+    "bits_per_elt": 1.0,
+    "n_traces": 50,
+    "key_cos": 0.85939,
+    "val_cos": 0.82403,
+    "key_mse": 0.44685,
+    "val_mse": 0.00347,
+    "attn_cos": 0.97507,
+    "attn_output_error": 0.34394,
+    "ip_rel": 3.87267,
+    "ip_bias": -0.00045,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "productvq-8x256-0.5b",
+    "bits_per_elt": 0.5,
+    "n_traces": 50,
+    "key_cos": 0.72974,
+    "val_cos": 0.66901,
+    "key_mse": 0.78885,
+    "val_mse": 0.00593,
+    "attn_cos": 0.9105,
+    "attn_output_error": 0.52521,
+    "ip_rel": 4.54048,
+    "ip_bias": 0.00067,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "ropesplit-1b",
+    "bits_per_elt": 1.0,
+    "n_traces": 50,
+    "key_cos": 0.8452,
+    "val_cos": 0.82348,
+    "key_mse": 0.48031,
+    "val_mse": 0.00347,
+    "attn_cos": 0.97524,
+    "attn_output_error": 0.34501,
+    "ip_rel": 4.81682,
+    "ip_bias": -0.00074,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "sign-1bit",
+    "bits_per_elt": 1.125,
+    "n_traces": 50,
+    "key_cos": 0.82812,
+    "val_cos": 0.79949,
+    "key_mse": 0.51988,
+    "val_mse": 0.00371,
+    "attn_cos": 0.97041,
+    "attn_output_error": 0.40237,
+    "ip_rel": 4.76672,
+    "ip_bias": -0.00047,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "ternary-bitnet",
+    "bits_per_elt": 1.71,
+    "n_traces": 50,
+    "key_cos": 0.91369,
+    "val_cos": 0.90047,
+    "key_mse": 0.27374,
+    "val_mse": 0.00194,
+    "attn_cos": 0.98747,
+    "attn_output_error": 0.23722,
+    "ip_rel": 3.34637,
+    "ip_bias": 0.00039,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "productvq-64x256-4b",
+    "bits_per_elt": 4.0,
+    "n_traces": 50,
+    "key_cos": 0.99591,
+    "val_cos": 0.99486,
+    "key_mse": 0.01412,
+    "val_mse": 0.00012,
+    "attn_cos": 0.9996,
+    "attn_output_error": 0.03707,
+    "ip_rel": 0.79562,
+    "ip_bias": -7e-05,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "productvq-32x256-2b",
+    "bits_per_elt": 2.0,
+    "n_traces": 50,
+    "key_cos": 0.9591,
+    "val_cos": 0.94788,
+    "key_mse": 0.13817,
+    "val_mse": 0.00109,
+    "attn_cos": 0.99563,
+    "attn_output_error": 0.14155,
+    "ip_rel": 3.0345,
+    "ip_bias": 0.00029,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "turbo-mse-4b",
+    "bits_per_elt": 4.125,
+    "n_traces": 50,
+    "key_cos": 0.99504,
+    "val_cos": 0.99504,
+    "key_mse": 0.01644,
+    "val_mse": 0.0001,
+    "attn_cos": 0.99962,
+    "attn_output_error": 0.0335,
+    "ip_rel": 0.95323,
+    "ip_bias": -1e-05,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "turbo-prod-4b (K:+qjl)",
+    "bits_per_elt": 5.125,
+    "n_traces": 50,
+    "key_cos": 0.99504,
+    "val_cos": 0.99504,
+    "key_mse": 0.01644,
+    "val_mse": 0.0001,
+    "attn_cos": 0.99962,
+    "attn_output_error": 0.03354,
+    "ip_rel": 1.13378,
+    "ip_bias": 0.00023,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "turbo-mse-2b",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.94075,
+    "val_cos": 0.94072,
+    "key_mse": 0.18757,
+    "val_mse": 0.00119,
+    "attn_cos": 0.99379,
+    "attn_output_error": 0.15863,
+    "ip_rel": 3.15094,
+    "ip_bias": -0.00013,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "turbo-prod-2b (K:+qjl)",
+    "bits_per_elt": 3.125,
+    "n_traces": 50,
+    "key_cos": 0.94075,
+    "val_cos": 0.94072,
+    "key_mse": 0.18757,
+    "val_mse": 0.00119,
+    "attn_cos": 0.99379,
+    "attn_output_error": 0.1587,
+    "ip_rel": 3.73498,
+    "ip_bias": 9e-05,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "turbo-mse-1b",
+    "bits_per_elt": 1.125,
+    "n_traces": 50,
+    "key_cos": 0.79977,
+    "val_cos": 0.79987,
+    "key_mse": 0.58414,
+    "val_mse": 0.00371,
+    "attn_cos": 0.97038,
+    "attn_output_error": 0.40117,
+    "ip_rel": 3.84738,
+    "ip_bias": 7e-05,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "turbo-prod-1b (K:+qjl)",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.79976,
+    "val_cos": 0.79986,
+    "key_mse": 0.58415,
+    "val_mse": 0.00371,
+    "attn_cos": 0.97038,
+    "attn_output_error": 0.40125,
+    "ip_rel": 6.33608,
+    "ip_bias": 0.0003,
+    "task": "hotpotqa"
+  },
+  {
+    "config": "scalar-int4",
+    "bits_per_elt": 4.25,
+    "n_traces": 50,
+    "key_cos": 0.99253,
+    "val_cos": 0.99503,
+    "key_mse": 0.02406,
+    "val_mse": 0.00011,
+    "attn_cos": 0.99968,
+    "attn_output_error": 0.02813,
+    "ip_rel": 1.03821,
+    "ip_bias": -0.00024,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "scalar-int2",
+    "bits_per_elt": 2.25,
+    "n_traces": 50,
+    "key_cos": 0.86483,
+    "val_cos": 0.89305,
+    "key_mse": 0.493,
+    "val_mse": 0.00263,
+    "attn_cos": 0.99115,
+    "attn_output_error": 0.14928,
+    "ip_rel": 5.79172,
+    "ip_bias": 0.00156,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "kivi-int2",
+    "bits_per_elt": 2.25,
+    "n_traces": 50,
+    "key_cos": 0.84588,
+    "val_cos": 0.89305,
+    "key_mse": 0.62044,
+    "val_mse": 0.00263,
+    "attn_cos": 0.99111,
+    "attn_output_error": 0.14955,
+    "ip_rel": 4.84881,
+    "ip_bias": 0.00017,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "turboquant~-int2",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.64121,
+    "val_cos": 0.64206,
+    "key_mse": 2.23468,
+    "val_mse": 0.01425,
+    "attn_cos": 0.386,
+    "attn_output_error": 2.34058,
+    "ip_rel": 11.80499,
+    "ip_bias": 0.00126,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "productvq-16x256-1b",
+    "bits_per_elt": 1.0,
+    "n_traces": 50,
+    "key_cos": 0.86038,
+    "val_cos": 0.82379,
+    "key_mse": 0.44569,
+    "val_mse": 0.00351,
+    "attn_cos": 0.97498,
+    "attn_output_error": 0.34639,
+    "ip_rel": 3.6548,
+    "ip_bias": 0.00048,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "productvq-8x256-0.5b",
+    "bits_per_elt": 0.5,
+    "n_traces": 50,
+    "key_cos": 0.73171,
+    "val_cos": 0.66755,
+    "key_mse": 0.78719,
+    "val_mse": 0.00602,
+    "attn_cos": 0.9103,
+    "attn_output_error": 0.52946,
+    "ip_rel": 4.94806,
+    "ip_bias": -0.00091,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "ropesplit-1b",
+    "bits_per_elt": 1.0,
+    "n_traces": 50,
+    "key_cos": 0.84593,
+    "val_cos": 0.82324,
+    "key_mse": 0.48009,
+    "val_mse": 0.00352,
+    "attn_cos": 0.97512,
+    "attn_output_error": 0.34746,
+    "ip_rel": 4.24382,
+    "ip_bias": 0.00107,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "sign-1bit",
+    "bits_per_elt": 1.125,
+    "n_traces": 50,
+    "key_cos": 0.82941,
+    "val_cos": 0.7995,
+    "key_mse": 0.51788,
+    "val_mse": 0.00375,
+    "attn_cos": 0.9717,
+    "attn_output_error": 0.40183,
+    "ip_rel": 4.47009,
+    "ip_bias": -0.00038,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "ternary-bitnet",
+    "bits_per_elt": 1.71,
+    "n_traces": 50,
+    "key_cos": 0.91408,
+    "val_cos": 0.90049,
+    "key_mse": 0.27345,
+    "val_mse": 0.00197,
+    "attn_cos": 0.98803,
+    "attn_output_error": 0.23646,
+    "ip_rel": 3.31466,
+    "ip_bias": 0.00039,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "productvq-64x256-4b",
+    "bits_per_elt": 4.0,
+    "n_traces": 50,
+    "key_cos": 0.99592,
+    "val_cos": 0.99482,
+    "key_mse": 0.01414,
+    "val_mse": 0.00012,
+    "attn_cos": 0.9996,
+    "attn_output_error": 0.03737,
+    "ip_rel": 0.97058,
+    "ip_bias": 0.00012,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "productvq-32x256-2b",
+    "bits_per_elt": 2.0,
+    "n_traces": 50,
+    "key_cos": 0.9593,
+    "val_cos": 0.94794,
+    "key_mse": 0.13801,
+    "val_mse": 0.0011,
+    "attn_cos": 0.99558,
+    "attn_output_error": 0.1428,
+    "ip_rel": 2.48128,
+    "ip_bias": -0.00033,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "turbo-mse-4b",
+    "bits_per_elt": 4.125,
+    "n_traces": 50,
+    "key_cos": 0.99504,
+    "val_cos": 0.99504,
+    "key_mse": 0.01652,
+    "val_mse": 0.00011,
+    "attn_cos": 0.99962,
+    "attn_output_error": 0.03385,
+    "ip_rel": 0.90062,
+    "ip_bias": -0.00015,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "turbo-prod-4b (K:+qjl)",
+    "bits_per_elt": 5.125,
+    "n_traces": 50,
+    "key_cos": 0.99504,
+    "val_cos": 0.99504,
+    "key_mse": 0.01652,
+    "val_mse": 0.00011,
+    "attn_cos": 0.99962,
+    "attn_output_error": 0.0339,
+    "ip_rel": 1.23727,
+    "ip_bias": -0.00021,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "turbo-mse-2b",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.94075,
+    "val_cos": 0.94073,
+    "key_mse": 0.18858,
+    "val_mse": 0.00121,
+    "attn_cos": 0.99395,
+    "attn_output_error": 0.1586,
+    "ip_rel": 3.42166,
+    "ip_bias": 0.00037,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "turbo-prod-2b (K:+qjl)",
+    "bits_per_elt": 3.125,
+    "n_traces": 50,
+    "key_cos": 0.94075,
+    "val_cos": 0.94073,
+    "key_mse": 0.18858,
+    "val_mse": 0.00121,
+    "attn_cos": 0.99395,
+    "attn_output_error": 0.15867,
+    "ip_rel": 3.99742,
+    "ip_bias": 0.00031,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "turbo-mse-1b",
+    "bits_per_elt": 1.125,
+    "n_traces": 50,
+    "key_cos": 0.79977,
+    "val_cos": 0.79989,
+    "key_mse": 0.58729,
+    "val_mse": 0.00375,
+    "attn_cos": 0.97165,
+    "attn_output_error": 0.4003,
+    "ip_rel": 4.29184,
+    "ip_bias": 0.00013,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "turbo-prod-1b (K:+qjl)",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.79976,
+    "val_cos": 0.79989,
+    "key_mse": 0.5873,
+    "val_mse": 0.00375,
+    "attn_cos": 0.97165,
+    "attn_output_error": 0.4004,
+    "ip_rel": 6.17042,
+    "ip_bias": -0.00022,
+    "task": "passage_retrieval_en"
+  },
+  {
+    "config": "scalar-int4",
+    "bits_per_elt": 4.25,
+    "n_traces": 50,
+    "key_cos": 0.99235,
+    "val_cos": 0.99503,
+    "key_mse": 0.02448,
+    "val_mse": 0.0001,
+    "attn_cos": 0.99966,
+    "attn_output_error": 0.02895,
+    "ip_rel": 1.16444,
+    "ip_bias": -8e-05,
+    "task": "repobench-p"
+  },
+  {
+    "config": "scalar-int2",
+    "bits_per_elt": 2.25,
+    "n_traces": 50,
+    "key_cos": 0.86482,
+    "val_cos": 0.89308,
+    "key_mse": 0.48425,
+    "val_mse": 0.00243,
+    "attn_cos": 0.99033,
+    "attn_output_error": 0.15413,
+    "ip_rel": 5.37172,
+    "ip_bias": -0.00046,
+    "task": "repobench-p"
+  },
+  {
+    "config": "kivi-int2",
+    "bits_per_elt": 2.25,
+    "n_traces": 50,
+    "key_cos": 0.85308,
+    "val_cos": 0.89308,
+    "key_mse": 0.59033,
+    "val_mse": 0.00243,
+    "attn_cos": 0.9903,
+    "attn_output_error": 0.1544,
+    "ip_rel": 5.35,
+    "ip_bias": -0.0002,
+    "task": "repobench-p"
+  },
+  {
+    "config": "turboquant~-int2",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.64046,
+    "val_cos": 0.64283,
+    "key_mse": 2.21611,
+    "val_mse": 0.01318,
+    "attn_cos": 0.38991,
+    "attn_output_error": 2.35685,
+    "ip_rel": 11.24984,
+    "ip_bias": -0.00501,
+    "task": "repobench-p"
+  },
+  {
+    "config": "productvq-16x256-1b",
+    "bits_per_elt": 1.0,
+    "n_traces": 50,
+    "key_cos": 0.86444,
+    "val_cos": 0.8248,
+    "key_mse": 0.43171,
+    "val_mse": 0.0032,
+    "attn_cos": 0.97492,
+    "attn_output_error": 0.32172,
+    "ip_rel": 4.49992,
+    "ip_bias": -0.0003,
+    "task": "repobench-p"
+  },
+  {
+    "config": "productvq-8x256-0.5b",
+    "bits_per_elt": 0.5,
+    "n_traces": 50,
+    "key_cos": 0.74333,
+    "val_cos": 0.67596,
+    "key_mse": 0.75351,
+    "val_mse": 0.00546,
+    "attn_cos": 0.91412,
+    "attn_output_error": 0.49052,
+    "ip_rel": 6.03475,
+    "ip_bias": -0.00016,
+    "task": "repobench-p"
+  },
+  {
+    "config": "ropesplit-1b",
+    "bits_per_elt": 1.0,
+    "n_traces": 50,
+    "key_cos": 0.8528,
+    "val_cos": 0.82423,
+    "key_mse": 0.45971,
+    "val_mse": 0.0032,
+    "attn_cos": 0.97494,
+    "attn_output_error": 0.32322,
+    "ip_rel": 4.06041,
+    "ip_bias": -0.0002,
+    "task": "repobench-p"
+  },
+  {
+    "config": "sign-1bit",
+    "bits_per_elt": 1.125,
+    "n_traces": 50,
+    "key_cos": 0.82897,
+    "val_cos": 0.79969,
+    "key_mse": 0.5174,
+    "val_mse": 0.00347,
+    "attn_cos": 0.96847,
+    "attn_output_error": 0.40384,
+    "ip_rel": 4.84825,
+    "ip_bias": -0.00062,
+    "task": "repobench-p"
+  },
+  {
+    "config": "ternary-bitnet",
+    "bits_per_elt": 1.71,
+    "n_traces": 50,
+    "key_cos": 0.91402,
+    "val_cos": 0.90054,
+    "key_mse": 0.27299,
+    "val_mse": 0.00182,
+    "attn_cos": 0.9867,
+    "attn_output_error": 0.23861,
+    "ip_rel": 2.86088,
+    "ip_bias": -0.0002,
+    "task": "repobench-p"
+  },
+  {
+    "config": "productvq-64x256-4b",
+    "bits_per_elt": 4.0,
+    "n_traces": 50,
+    "key_cos": 0.99606,
+    "val_cos": 0.99508,
+    "key_mse": 0.01358,
+    "val_mse": 0.00011,
+    "attn_cos": 0.99962,
+    "attn_output_error": 0.0346,
+    "ip_rel": 0.71153,
+    "ip_bias": -6e-05,
+    "task": "repobench-p"
+  },
+  {
+    "config": "productvq-32x256-2b",
+    "bits_per_elt": 2.0,
+    "n_traces": 50,
+    "key_cos": 0.9603,
+    "val_cos": 0.94747,
+    "key_mse": 0.13409,
+    "val_mse": 0.001,
+    "attn_cos": 0.99551,
+    "attn_output_error": 0.13328,
+    "ip_rel": 2.39037,
+    "ip_bias": -5e-05,
+    "task": "repobench-p"
+  },
+  {
+    "config": "turbo-mse-4b",
+    "bits_per_elt": 4.125,
+    "n_traces": 50,
+    "key_cos": 0.99504,
+    "val_cos": 0.99503,
+    "key_mse": 0.01641,
+    "val_mse": 0.0001,
+    "attn_cos": 0.99958,
+    "attn_output_error": 0.03482,
+    "ip_rel": 0.9464,
+    "ip_bias": 6e-05,
+    "task": "repobench-p"
+  },
+  {
+    "config": "turbo-prod-4b (K:+qjl)",
+    "bits_per_elt": 5.125,
+    "n_traces": 50,
+    "key_cos": 0.99504,
+    "val_cos": 0.99503,
+    "key_mse": 0.01641,
+    "val_mse": 0.0001,
+    "attn_cos": 0.99958,
+    "attn_output_error": 0.03486,
+    "ip_rel": 1.26089,
+    "ip_bias": 7e-05,
+    "task": "repobench-p"
+  },
+  {
+    "config": "turbo-mse-2b",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.94071,
+    "val_cos": 0.94064,
+    "key_mse": 0.18716,
+    "val_mse": 0.00112,
+    "attn_cos": 0.99335,
+    "attn_output_error": 0.16086,
+    "ip_rel": 3.0084,
+    "ip_bias": 0.00063,
+    "task": "repobench-p"
+  },
+  {
+    "config": "turbo-prod-2b (K:+qjl)",
+    "bits_per_elt": 3.125,
+    "n_traces": 50,
+    "key_cos": 0.94071,
+    "val_cos": 0.94064,
+    "key_mse": 0.18716,
+    "val_mse": 0.00112,
+    "attn_cos": 0.99336,
+    "attn_output_error": 0.1609,
+    "ip_rel": 4.07001,
+    "ip_bias": 8e-05,
+    "task": "repobench-p"
+  },
+  {
+    "config": "turbo-mse-1b",
+    "bits_per_elt": 1.125,
+    "n_traces": 50,
+    "key_cos": 0.79966,
+    "val_cos": 0.79959,
+    "key_mse": 0.58273,
+    "val_mse": 0.00347,
+    "attn_cos": 0.96834,
+    "attn_output_error": 0.40316,
+    "ip_rel": 4.66947,
+    "ip_bias": -0.00025,
+    "task": "repobench-p"
+  },
+  {
+    "config": "turbo-prod-1b (K:+qjl)",
+    "bits_per_elt": 2.125,
+    "n_traces": 50,
+    "key_cos": 0.79965,
+    "val_cos": 0.79959,
+    "key_mse": 0.58274,
+    "val_mse": 0.00347,
+    "attn_cos": 0.96834,
+    "attn_output_error": 0.40325,
+    "ip_rel": 7.08782,
+    "ip_bias": -0.00021,
+    "task": "repobench-p"
+  }
+]

artifacts/longbench_results.json

@@ -0,0 +1,632 @@
+[
+  {
+    "task": "qasper",
+    "config": "fp16",
+    "bpe": 16.0,
+    "f1": 0.3584,
+    "delta_f1": 0.0
+  },
+  {
+    "task": "qasper",
+    "config": "scalar-int4",
+    "bpe": 4.25,
+    "f1": 0.3565,
+    "delta_f1": 0.0031
+  },
+  {
+    "task": "qasper",
+    "config": "scalar-int2",
+    "bpe": 2.25,
+    "f1": 0.2531,
+    "delta_f1": -0.1003
+  },
+  {
+    "task": "qasper",
+    "config": "kivi-int2",
+    "bpe": 2.25,
+    "f1": 0.27,
+    "delta_f1": -0.0834
+  },
+  {
+    "task": "qasper",
+    "config": "turboquant~-int2",
+    "bpe": 2.125,
+    "f1": 0.0927,
+    "delta_f1": -0.2607
+  },
+  {
+    "task": "qasper",
+    "config": "productvq-16x256-1b",
+    "bpe": 1.0,
+    "f1": 0.2049,
+    "delta_f1": -0.1485
+  },
+  {
+    "task": "qasper",
+    "config": "productvq-8x256-0.5b",
+    "bpe": 0.5,
+    "f1": 0.146,
+    "delta_f1": -0.2074
+  },
+  {
+    "task": "qasper",
+    "config": "ropesplit-1b",
+    "bpe": 1.0,
+    "f1": 0.2108,
+    "delta_f1": -0.1426
+  },
+  {
+    "task": "qasper",
+    "config": "sign-1bit",
+    "bpe": 1.125,
+    "f1": 0.1731,
+    "delta_f1": -0.1803
+  },
+  {
+    "task": "qasper",
+    "config": "ternary-bitnet",
+    "bpe": 1.71,
+    "f1": 0.221,
+    "delta_f1": -0.1325
+  },
+  {
+    "task": "qasper",
+    "config": "productvq-64x256-4b",
+    "bpe": 4.0,
+    "f1": 0.349,
+    "delta_f1": -0.0044
+  },
+  {
+    "task": "qasper",
+    "config": "productvq-32x256-2b",
+    "bpe": 2.0,
+    "f1": 0.3092,
+    "delta_f1": -0.0442
+  },
+  {
+    "task": "qasper",
+    "config": "turbo-mse-4b",
+    "bpe": 4.125,
+    "f1": 0.3582,
+    "delta_f1": 0.0048
+  },
+  {
+    "task": "qasper",
+    "config": "turbo-prod-4b (K:+qjl)",
+    "bpe": 5.125,
+    "f1": 0.3363,
+    "delta_f1": -0.0171
+  },
+  {
+    "task": "qasper",
+    "config": "turbo-mse-2b",
+    "bpe": 2.125,
+    "f1": 0.2605,
+    "delta_f1": -0.0929
+  },
+  {
+    "task": "qasper",
+    "config": "turbo-prod-2b (K:+qjl)",
+    "bpe": 3.125,
+    "f1": 0.2684,
+    "delta_f1": -0.085
+  },
+  {
+    "task": "qasper",
+    "config": "turbo-mse-1b",
+    "bpe": 1.125,
+    "f1": 0.1645,
+    "delta_f1": -0.1889
+  },
+  {
+    "task": "qasper",
+    "config": "turbo-prod-1b (K:+qjl)",
+    "bpe": 2.125,
+    "f1": 0.188,
+    "delta_f1": -0.1654
+  },
+  {
+    "task": "2wikimqa",
+    "config": "fp16",
+    "bpe": 16.0,
+    "f1": 0.3163,
+    "delta_f1": 0.0
+  },
+  {
+    "task": "2wikimqa",
+    "config": "scalar-int4",
+    "bpe": 4.25,
+    "f1": 0.3041,
+    "delta_f1": -0.0121
+  },
+  {
+    "task": "2wikimqa",
+    "config": "scalar-int2",
+    "bpe": 2.25,
+    "f1": 0.2227,
+    "delta_f1": -0.0936
+  },
+  {
+    "task": "2wikimqa",
+    "config": "kivi-int2",
+    "bpe": 2.25,
+    "f1": 0.1952,
+    "delta_f1": -0.1211
+  },
+  {
+    "task": "2wikimqa",
+    "config": "turboquant~-int2",
+    "bpe": 2.125,
+    "f1": 0.0392,
+    "delta_f1": -0.2771
+  },
+  {
+    "task": "2wikimqa",
+    "config": "productvq-16x256-1b",
+    "bpe": 1.0,
+    "f1": 0.2314,
+    "delta_f1": -0.0849
+  },
+  {
+    "task": "2wikimqa",
+    "config": "productvq-8x256-0.5b",
+    "bpe": 0.5,
+    "f1": 0.0979,
+    "delta_f1": -0.2184
+  },
+  {
+    "task": "2wikimqa",
+    "config": "ropesplit-1b",
+    "bpe": 1.0,
+    "f1": 0.2485,
+    "delta_f1": -0.0678
+  },
+  {
+    "task": "2wikimqa",
+    "config": "sign-1bit",
+    "bpe": 1.125,
+    "f1": 0.1601,
+    "delta_f1": -0.1562
+  },
+  {
+    "task": "2wikimqa",
+    "config": "ternary-bitnet",
+    "bpe": 1.71,
+    "f1": 0.244,
+    "delta_f1": -0.0723
+  },
+  {
+    "task": "2wikimqa",
+    "config": "productvq-64x256-4b",
+    "bpe": 4.0,
+    "f1": 0.2542,
+    "delta_f1": -0.0621
+  },
+  {
+    "task": "2wikimqa",
+    "config": "productvq-32x256-2b",
+    "bpe": 2.0,
+    "f1": 0.3099,
+    "delta_f1": -0.0064
+  },
+  {
+    "task": "2wikimqa",
+    "config": "turbo-mse-4b",
+    "bpe": 4.125,
+    "f1": 0.301,
+    "delta_f1": -0.0152
+  },
+  {
+    "task": "2wikimqa",
+    "config": "turbo-prod-4b (K:+qjl)",
+    "bpe": 5.125,
+    "f1": 0.3204,
+    "delta_f1": 0.0041
+  },
+  {
+    "task": "2wikimqa",
+    "config": "turbo-mse-2b",
+    "bpe": 2.125,
+    "f1": 0.2739,
+    "delta_f1": -0.0424
+  },
+  {
+    "task": "2wikimqa",
+    "config": "turbo-prod-2b (K:+qjl)",
+    "bpe": 3.125,
+    "f1": 0.2502,
+    "delta_f1": -0.0661
+  },
+  {
+    "task": "2wikimqa",
+    "config": "turbo-mse-1b",
+    "bpe": 1.125,
+    "f1": 0.2113,
+    "delta_f1": -0.105
+  },
+  {
+    "task": "2wikimqa",
+    "config": "turbo-prod-1b (K:+qjl)",
+    "bpe": 2.125,
+    "f1": 0.1891,
+    "delta_f1": -0.1272
+  },
+  {
+    "task": "hotpotqa",
+    "config": "fp16",
+    "bpe": 16.0,
+    "f1": 0.4018,
+    "delta_f1": 0.0
+  },
+  {
+    "task": "hotpotqa",
+    "config": "scalar-int4",
+    "bpe": 4.25,
+    "f1": 0.3873,
+    "delta_f1": 0.0155
+  },
+  {
+    "task": "hotpotqa",
+    "config": "scalar-int2",
+    "bpe": 2.25,
+    "f1": 0.2828,
+    "delta_f1": -0.089
+  },
+  {
+    "task": "hotpotqa",
+    "config": "kivi-int2",
+    "bpe": 2.25,
+    "f1": 0.2158,
+    "delta_f1": -0.156
+  },
+  {
+    "task": "hotpotqa",
+    "config": "turboquant~-int2",
+    "bpe": 2.125,
+    "f1": 0.0118,
+    "delta_f1": -0.36
+  },
+  {
+    "task": "hotpotqa",
+    "config": "productvq-16x256-1b",
+    "bpe": 1.0,
+    "f1": 0.1762,
+    "delta_f1": -0.1956
+  },
+  {
+    "task": "hotpotqa",
+    "config": "productvq-8x256-0.5b",
+    "bpe": 0.5,
+    "f1": 0.1524,
+    "delta_f1": -0.2194
+  },
+  {
+    "task": "hotpotqa",
+    "config": "ropesplit-1b",
+    "bpe": 1.0,
+    "f1": 0.2318,
+    "delta_f1": -0.14
+  },
+  {
+    "task": "hotpotqa",
+    "config": "sign-1bit",
+    "bpe": 1.125,
+    "f1": 0.1167,
+    "delta_f1": -0.2551
+  },
+  {
+    "task": "hotpotqa",
+    "config": "ternary-bitnet",
+    "bpe": 1.71,
+    "f1": 0.3327,
+    "delta_f1": -0.0391
+  },
+  {
+    "task": "hotpotqa",
+    "config": "productvq-64x256-4b",
+    "bpe": 4.0,
+    "f1": 0.4212,
+    "delta_f1": 0.0494
+  },
+  {
+    "task": "hotpotqa",
+    "config": "productvq-32x256-2b",
+    "bpe": 2.0,
+    "f1": 0.3679,
+    "delta_f1": -0.0039
+  },
+  {
+    "task": "hotpotqa",
+    "config": "turbo-mse-4b",
+    "bpe": 4.125,
+    "f1": 0.4002,
+    "delta_f1": 0.0284
+  },
+  {
+    "task": "hotpotqa",
+    "config": "turbo-prod-4b (K:+qjl)",
+    "bpe": 5.125,
+    "f1": 0.3697,
+    "delta_f1": -0.0021
+  },
+  {
+    "task": "hotpotqa",
+    "config": "turbo-mse-2b",
+    "bpe": 2.125,
+    "f1": 0.2987,
+    "delta_f1": -0.0731
+  },
+  {
+    "task": "hotpotqa",
+    "config": "turbo-prod-2b (K:+qjl)",
+    "bpe": 3.125,
+    "f1": 0.3223,
+    "delta_f1": -0.0495
+  },
+  {
+    "task": "hotpotqa",
+    "config": "turbo-mse-1b",
+    "bpe": 1.125,
+    "f1": 0.1526,
+    "delta_f1": -0.2192
+  },
+  {
+    "task": "hotpotqa",
+    "config": "turbo-prod-1b (K:+qjl)",
+    "bpe": 2.125,
+    "f1": 0.1677,
+    "delta_f1": -0.2041
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "fp16",
+    "bpe": 16.0,
+    "f1": 1.0,
+    "delta_f1": 0.0
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "scalar-int4",
+    "bpe": 4.25,
+    "f1": 0.98,
+    "delta_f1": -0.02
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "scalar-int2",
+    "bpe": 2.25,
+    "f1": 0.74,
+    "delta_f1": -0.26
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "kivi-int2",
+    "bpe": 2.25,
+    "f1": 0.62,
+    "delta_f1": -0.38
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "turboquant~-int2",
+    "bpe": 2.125,
+    "f1": 0.02,
+    "delta_f1": -0.98
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "productvq-16x256-1b",
+    "bpe": 1.0,
+    "f1": 0.26,
+    "delta_f1": -0.74
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "productvq-8x256-0.5b",
+    "bpe": 0.5,
+    "f1": 0.06,
+    "delta_f1": -0.94
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "ropesplit-1b",
+    "bpe": 1.0,
+    "f1": 0.1,
+    "delta_f1": -0.9
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "sign-1bit",
+    "bpe": 1.125,
+    "f1": 0.08,
+    "delta_f1": -0.92
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "ternary-bitnet",
+    "bpe": 1.71,
+    "f1": 0.26,
+    "delta_f1": -0.74
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "productvq-64x256-4b",
+    "bpe": 4.0,
+    "f1": 1.0,
+    "delta_f1": 0.0
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "productvq-32x256-2b",
+    "bpe": 2.0,
+    "f1": 0.72,
+    "delta_f1": -0.28
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "turbo-mse-4b",
+    "bpe": 4.125,
+    "f1": 1.0,
+    "delta_f1": 0.0
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "turbo-prod-4b (K:+qjl)",
+    "bpe": 5.125,
+    "f1": 1.0,
+    "delta_f1": 0.0
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "turbo-mse-2b",
+    "bpe": 2.125,
+    "f1": 0.84,
+    "delta_f1": -0.16
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "turbo-prod-2b (K:+qjl)",
+    "bpe": 3.125,
+    "f1": 0.9,
+    "delta_f1": -0.1
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "turbo-mse-1b",
+    "bpe": 1.125,
+    "f1": 0.06,
+    "delta_f1": -0.94
+  },
+  {
+    "task": "passage_retrieval_en",
+    "config": "turbo-prod-1b (K:+qjl)",
+    "bpe": 2.125,
+    "f1": 0.02,
+    "delta_f1": -0.98
+  },
+  {
+    "task": "repobench-p",
+    "config": "fp16",
+    "bpe": 16.0,
+    "f1": 0.196,
+    "delta_f1": 0.0
+  },
+  {
+    "task": "repobench-p",
+    "config": "scalar-int4",
+    "bpe": 4.25,
+    "f1": 0.1949,
+    "delta_f1": 0.0089
+  },
+  {
+    "task": "repobench-p",
+    "config": "scalar-int2",
+    "bpe": 2.25,
+    "f1": 0.1671,
+    "delta_f1": -0.0189
+  },
+  {
+    "task": "repobench-p",
+    "config": "kivi-int2",
+    "bpe": 2.25,
+    "f1": 0.1744,
+    "delta_f1": -0.0116
+  },
+  {
+    "task": "repobench-p",
+    "config": "turboquant~-int2",
+    "bpe": 2.125,
+    "f1": 0.0757,
+    "delta_f1": -0.1103
+  },
+  {
+    "task": "repobench-p",
+    "config": "productvq-16x256-1b",
+    "bpe": 1.0,
+    "f1": 0.1693,
+    "delta_f1": -0.0167
+  },
+  {
+    "task": "repobench-p",
+    "config": "productvq-8x256-0.5b",
+    "bpe": 0.5,
+    "f1": 0.1483,
+    "delta_f1": -0.0376
+  },
+  {
+    "task": "repobench-p",
+    "config": "ropesplit-1b",
+    "bpe": 1.0,
+    "f1": 0.1618,
+    "delta_f1": -0.0241
+  },
+  {
+    "task": "repobench-p",
+    "config": "sign-1bit",
+    "bpe": 1.125,
+    "f1": 0.1309,
+    "delta_f1": -0.0551
+  },
+  {
+    "task": "repobench-p",
+    "config": "ternary-bitnet",
+    "bpe": 1.71,
+    "f1": 0.1675,
+    "delta_f1": -0.0185
+  },
+  {
+    "task": "repobench-p",
+    "config": "productvq-64x256-4b",
+    "bpe": 4.0,
+    "f1": 0.1891,
+    "delta_f1": 0.0032
+  },
+  {
+    "task": "repobench-p",
+    "config": "productvq-32x256-2b",
+    "bpe": 2.0,
+    "f1": 0.192,
+    "delta_f1": 0.0061
+  },
+  {
+    "task": "repobench-p",
+    "config": "turbo-mse-4b",
+    "bpe": 4.125,
+    "f1": 0.1854,
+    "delta_f1": -0.0005
+  },
+  {
+    "task": "repobench-p",
+    "config": "turbo-prod-4b (K:+qjl)",
+    "bpe": 5.125,
+    "f1": 0.1832,
+    "delta_f1": -0.0028
+  },
+  {
+    "task": "repobench-p",
+    "config": "turbo-mse-2b",
+    "bpe": 2.125,
+    "f1": 0.1874,
+    "delta_f1": 0.0015
+  },
+  {
+    "task": "repobench-p",
+    "config": "turbo-prod-2b (K:+qjl)",
+    "bpe": 3.125,
+    "f1": 0.193,
+    "delta_f1": 0.007
+  },
+  {
+    "task": "repobench-p",
+    "config": "turbo-mse-1b",
+    "bpe": 1.125,
+    "f1": 0.1579,
+    "delta_f1": -0.0281
+  },
+  {
+    "task": "repobench-p",
+    "config": "turbo-prod-1b (K:+qjl)",
+    "bpe": 2.125,
+    "f1": 0.1571,
+    "delta_f1": -0.0289
+  }
+]

artifacts/turbo_codebooks.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6167b29f4be685cb01d41aab6443f7b0e491512a05f729bfeb14d46de721fde
+size 9917315

benchmark.py

@@ -0,0 +1,782 @@
+"""
+benchmark.py — fit AttnVQ codebooks and evaluate on real Laguna-XS.2 caches.
+
+Stages:
+    dump      capture post-RoPE K/V from full-attention layers
+    fit       per-layer LBG codebooks → artifacts/codebooks.pt
+    cheap     proxy metrics (key cosine, attn-output error, ip-bias)
+    swebench  optional resolve-rate delta on SWE-bench Verified (needs Docker)
+
+Usage:
+    python benchmark.py --stage fit
+    python benchmark.py --stage cheap --n_eval 64
+    python benchmark.py --stage dump
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, asdict
+
+import torch
+
+from vqkv.quantizers import (ScalarKV, KIVIScalarKV, ProductVQKV, RoPESplitVQKV,
+                            SignScalarKV, TernaryScalarKV)
+from vqkv.metrics import (key_cosine, cache_mse, inner_product_distortion,
+                          attention_output, attn_output_cosine, attn_output_error)
+
+MODEL_ID = os.environ.get("LAGUNA_ID", "poolside/Laguna-XS.2")
+ARTIFACT_DIR = os.environ.get("VQKV_ARTIFACTS", "./artifacts")
+os.makedirs(ARTIFACT_DIR, exist_ok=True)
+
+CALIB_DATASET = os.environ.get("CALIB_DATASET", "SWE-bench/SWE-smith-trajectories")
+CALIB_SPLIT = os.environ.get("CALIB_SPLIT", "tool")
+CALIB_SOURCE = os.environ.get("CALIB_SOURCE", "swesmith")  # swesmith | longbench-hotpotqa
+
+_HOTPOTQA_PROMPT = (
+    "Answer the question based on the given passages. "
+    "Only give me the answer and do not output any other words.\n\n"
+    "The following are given passages.\n{context}\n\n"
+    "Answer the question based on the given passages. "
+    "Only give me the answer and do not output any other words.\n\n"
+    "Question: {input}\nAnswer:"
+)
+
+# ============================================================================
+# Cache configurations under test. Each is a (name, factory) where factory()
+# returns an unfitted quantizer. `None` means the fp16 baseline (no quant).
+# ============================================================================
+def cache_configs():
+    return [
+        ("fp16 (baseline)",      None),
+        ("scalar-int4",          lambda: ScalarKV(nbits=4)),
+        ("scalar-int2",          lambda: ScalarKV(nbits=2)),
+        ("kivi-int2",            lambda: KIVIScalarKV(nbits=2)),
+        ("productvq-64x256-4b",   lambda: ProductVQKV(n_sub=64, n_codes=256, iters=15)), 
+        ("productvq-32x256-2b",   lambda: ProductVQKV(n_sub=32, n_codes=256, iters=15)), 
+        ("productvq-16x256-1b",  lambda: ProductVQKV(n_sub=16, n_codes=256, iters=15)),
+        ("productvq-8x256-0.5b", lambda: ProductVQKV(n_sub=8,  n_codes=256, iters=15)),
+        ("ropesplit-1b",         lambda: RoPESplitVQKV(n_sub_half=8, n_codes=256, iters=15)),
+        ("sign-1bit",            lambda: SignScalarKV(per_channel_key=True)),
+        ("ternary-bitnet",       lambda: TernaryScalarKV(alpha=0.7, per_channel_key=True)),
+    ]
+
+
+# ============================================================================
+# Model + layer-structure loading
+# ============================================================================
+def load_model_and_meta():
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID, torch_dtype=torch.bfloat16, device_map="cuda",
+        trust_remote_code=True)
+    model.eval()
+    cfg = model.config
+    full_layers = [i for i, t in enumerate(cfg.layer_types) if t == "full_attention"]
+    meta = {
+        "full_layers": full_layers,
+        "n_kv_heads": cfg.num_key_value_heads,
+        "n_q_heads":  cfg.num_attention_heads,
+        "head_dim": cfg.head_dim,
+        "n_layers": cfg.num_hidden_layers,
+    }
+    print(f"[meta] full-attention layers ({len(full_layers)}): {full_layers}")
+    print(f"[meta] kv_heads={meta['n_kv_heads']} head_dim={meta['head_dim']}")
+    assert len(full_layers) > 0, "no full-attention layers found; check config"
+    return model, tok, meta
+
+
+# Trace flattening (SWE-smith messages JSON, nebius trajectory/tool roles).
+def flatten_trace(example, tok) -> str:
+    """Flatten one trajectory to a string via the model's chat template."""
+    raw = (example.get("messages") or example.get("trajectory")
+           or example.get("conversations"))
+    if raw is None:
+        return json.dumps(example)[:200_000]
+
+    # SWE-smith stores 'messages' as a JSON string, not a native list.
+    if isinstance(raw, str):
+        try:
+            raw = json.loads(raw)
+        except json.JSONDecodeError:
+            return raw[:200_000]
+
+    norm = []
+    for m in raw:
+        role = m.get("role") or m.get("from") or "user"
+        content = m.get("content") or m.get("value") or ""
+
+        # SWE-smith: content can be a list of {"type":"text","text":"..."} blocks
+        if isinstance(content, list):
+            content = "\n".join(
+                item.get("text", str(item)) if isinstance(item, dict) else str(item)
+                for item in content
+            )
+
+        # nebius: assistant turns carry tool_calls alongside content
+        tool_calls = m.get("tool_calls")
+        if tool_calls:
+            tc_text = json.dumps(tool_calls, ensure_ascii=False)
+            content = (content + "\n" + tc_text).strip() if content else tc_text
+
+        # 'tool' role (observation) has no equivalent in most chat templates;
+        # map it to 'user' so the template accepts it.
+        role = {"human": "user", "gpt": "assistant", "tool": "user"}.get(role, role)
+
+        if not content.strip():
+            continue
+        norm.append({"role": role, "content": content})
+
+    # Merge consecutive same-role messages produced by tool->user collapsing.
+    merged: list[dict] = []
+    for m in norm:
+        if merged and merged[-1]["role"] == m["role"]:
+            merged[-1]["content"] += "\n\n" + m["content"]
+        else:
+            merged.append(dict(m))
+
+    try:
+        return tok.apply_chat_template(merged, tokenize=False,
+                                       add_generation_prompt=False)
+    except Exception:
+        return "\n\n".join(f"{m['role']}: {m['content']}" for m in merged)
+
+
+def flatten_longbench(example) -> str:
+    """Format a LongBench hotpotqa example (context + input) as a plain string."""
+    return _HOTPOTQA_PROMPT.format(
+        context=example["context"], input=example["input"])
+
+
+def _load_longbench_hotpotqa():
+    """Load THUDM/LongBench hotpotqa, bypassing the deprecated dataset script."""
+    from datasets import load_dataset as _ld
+    for fname in ("hotpotqa_e.jsonl", "hotpotqa.jsonl"):
+        try:
+            return _ld(
+                "json",
+                data_files=f"hf://datasets/THUDM/LongBench/data/{fname}",
+                split="train",
+            )
+        except Exception:
+            continue
+    return _ld("THUDM/LongBench", name="hotpotqa", split="test")
+
+
+def _load_calib_dump_dataset(n_calib: int, calib_source: str, tok):
+    """Return (dataset, text_fn, label) for stage_dump."""
+    from datasets import load_dataset
+
+    if calib_source == "longbench-hotpotqa":
+        ds = _load_longbench_hotpotqa()
+        n_total = len(ds)
+        # Last n_calib rows avoid overlap with longbench_eval.py (range(0, n_eval)).
+        start = max(0, n_total - n_calib)
+        ds = ds.select(range(start, n_total))
+        label = f"LongBench hotpotqa (rows {start}–{n_total - 1})"
+        return ds, flatten_longbench, label
+
+    ds = load_dataset(CALIB_DATASET, split=CALIB_SPLIT)
+    label = f"{CALIB_DATASET} split={CALIB_SPLIT}"
+    return ds, lambda ex: flatten_trace(ex, tok), label
+
+
+# STAGE: dump -- run real traces, capture post-RoPE K/V from full-attn layers
+def stage_dump(n_calib=16, max_len=32768, calib_source: str | None = None,
+               min_len=2048):
+    from transformers.cache_utils import DynamicCache
+
+    calib_source = calib_source or CALIB_SOURCE
+    model, tok, meta = load_model_and_meta()
+    full = set(meta["full_layers"])
+
+    # device_map="auto" can leave model.device as 'meta'; always use cuda:0.
+    input_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    class DumpingCache(DynamicCache):
+        def __init__(self, *a, **k):
+            super().__init__(*a, **k)
+            self.dump = {i: {"k": [], "v": []} for i in full}
+        def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
+            if layer_idx in full:
+                self.dump[layer_idx]["k"].append(
+                    key_states.detach()[0].permute(1, 0, 2).float().cpu())
+                self.dump[layer_idx]["v"].append(
+                    value_states.detach()[0].permute(1, 0, 2).float().cpu())
+            return super().update(key_states, value_states, layer_idx, cache_kwargs)
+
+    ds, get_text, source_label = _load_calib_dump_dataset(n_calib, calib_source, tok)
+    print(f"[dump] source={calib_source}  {source_label}  rows={len(ds)} "
+          f"schema={list(ds[0].keys())}")
+
+    agg = {i: {"k": [], "v": []} for i in full}
+    used = 0
+    for ex in ds:
+        if used >= n_calib:
+            break
+        text = get_text(ex)
+        ids = tok(text, return_tensors="pt", truncation=True,
+                  max_length=max_len).to(input_device)
+        if ids["input_ids"].shape[1] < min_len:
+            continue
+        cache = DumpingCache(config=model.config)  # fresh per trace
+        with torch.no_grad():
+            model.model(**ids, past_key_values=cache, use_cache=True)  # skip lm_head
+        for i in full:
+            agg[i]["k"].append(torch.cat(cache.dump[i]["k"]))
+            agg[i]["v"].append(torch.cat(cache.dump[i]["v"]))
+        used += 1
+        print(f"[dump] trace {used}/{n_calib} len={ids['input_ids'].shape[1]}")
+
+    calib = {i: {"k": torch.cat(agg[i]["k"]), "v": torch.cat(agg[i]["v"])}
+             for i in full}
+    for i in list(full)[:3]:
+        k = calib[i]["k"].reshape(-1, meta["head_dim"])
+        cut = meta["head_dim"] // 2
+        print(f"[dump] layer {i}: {tuple(calib[i]['k'].shape)} | "
+              f"rope-half std {k[:, :cut].std():.3f} "
+              f"pass-half std {k[:, cut:].std():.3f} "
+              f"max|ch| {k.abs().amax(0).max():.2f}")
+    path = os.path.join(ARTIFACT_DIR, "calib_caches.pt")
+    torch.save({"calib": calib, "meta": meta}, path)
+    print(f"[dump] saved -> {path}")
+
+
+# STAGE: fit -- per-layer codebooks for every (data-dependent) config
+def stage_fit(only: list[str] | None = None):
+    """Fit quantizers and write artifacts/codebooks.pt.
+
+    By default fits every config in cache_configs(). Pass ``only=["sign-1bit", ...]``
+    to fit just those names and merge into an existing codebooks.pt (skips the rest).
+    Tuning-free quantizers (Sign, Ternary, Scalar, KIVI) finish in seconds; only
+    ProductVQ / RoPE-split need calib_caches.pt.
+    """
+    codebooks_path = os.path.join(ARTIFACT_DIR, "codebooks.pt")
+    calib_path = os.path.join(ARTIFACT_DIR, "calib_caches.pt")
+
+    if only and os.path.exists(codebooks_path):
+        existing = torch.load(codebooks_path, weights_only=False)
+        fitted = existing["fitted"]
+        meta = existing["meta"]
+        print(f"[fit] merging into existing {codebooks_path} ({len(fitted)} configs)")
+    else:
+        if not os.path.exists(calib_path):
+            raise FileNotFoundError(
+                f"{calib_path} not found; run --stage dump first, or use "
+                f"--only with an existing codebooks.pt for tuning-free configs")
+        blob = torch.load(calib_path)
+        calib, meta = blob["calib"], blob["meta"]
+        fitted = {}
+
+    calib = None
+    if os.path.exists(calib_path):
+        calib = torch.load(calib_path)["calib"]
+
+    hd = meta["head_dim"]
+    layer_ids = meta["full_layers"]
+
+    configs = [(n, f) for n, f in cache_configs() if f is not None]
+    if only is not None:
+        only_set = set(only)
+        configs = [(n, f) for n, f in configs if n in only_set]
+        unknown = only_set - {n for n, _ in configs}
+        if unknown:
+            raise ValueError(f"unknown --only config(s): {sorted(unknown)}")
+
+    # GPU fitting: LBG is pure torch; moving calib to CUDA makes bmm/argmin ~20x
+    # faster. Serial over layers on GPU (CUDA is already async; threading adds no
+    # benefit). Parallel over layers on CPU (BLAS releases GIL; real concurrency).
+    fit_device = "cuda" if torch.cuda.is_available() else "cpu"
+    if fit_device == "cuda":
+        print(f"[fit] GPU available ({torch.cuda.get_device_name()}) -- fitting LBG on CUDA")
+    n_workers = 1 if fit_device == "cuda" else min(len(layer_ids), os.cpu_count() or 1)
+
+    for name, factory in configs:
+        t0 = time.time()
+
+        def _fit_layer(i, _factory=factory, _device=fit_device):
+            if calib is not None and i in calib:
+                kf = calib[i]["k"].reshape(-1, hd)[:200_000].to(_device)
+                vf = calib[i]["v"].reshape(-1, hd)[:200_000].to(_device)
+            else:
+                kf = torch.zeros(1, hd, device=_device)
+                vf = torch.zeros(1, hd, device=_device)
+            q = _factory().fit(kf, vf)
+            # Codebooks are saved to disk as CPU tensors; move back before returning.
+            if _device != "cpu" and hasattr(q, "to"):
+                q.to("cpu")
+            return i, q
+
+        with ThreadPoolExecutor(max_workers=n_workers) as ex:
+            per_layer = dict(ex.map(_fit_layer, layer_ids))
+
+        fitted[name] = per_layer
+        print(f"[fit] {name}: {len(per_layer)} layer-codebooks in {time.time()-t0:.1f}s")
+
+    torch.save({"fitted": fitted, "meta": meta}, codebooks_path)
+    print(f"[fit] saved -> {codebooks_path} ({len(fitted)} configs total)")
+
+
+# A drop-in cache that applies a per-layer quantizer to the target layers only.
+def make_vq_cache_class(per_layer_quantizers, target_layers, model_config, device=None):
+    """Build a VQCache class with codebooks pre-moved to `device`.
+
+    Pass device="cuda" (or the model's device) so the roundtrip runs entirely
+    on GPU with no CPU<->GPU transfers. Without this, every update() call
+    implicitly transfers key_states to CPU and back, making decode very slow.
+    """
+    from transformers.cache_utils import DynamicCache
+
+    if device is not None:
+        for q in per_layer_quantizers.values():
+            if hasattr(q, "to"):
+                q.to(device)
+
+    class VQCache(DynamicCache):
+        def __init__(self, *a, **k):
+            super().__init__(*a, **k)
+            self.q = per_layer_quantizers
+            self.target = set(target_layers)
+        def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
+            if layer_idx in self.target and layer_idx in self.q:
+                q = self.q[layer_idx]
+                b, h, s, d = key_states.shape
+                kf = key_states[0].transpose(0, 1).reshape(-1, d).float()
+                vf = value_states[0].transpose(0, 1).reshape(-1, d).float()
+                # See stage_cheap: per-channel-key quantizers reduce the key
+                # along the token axis and need a per-head block, not a flatten.
+                per_channel_key = (
+                    isinstance(q, KIVIScalarKV)
+                    or getattr(q, "per_channel_key", False)
+                )
+                if per_channel_key:
+                    kk = key_states[0].permute(1, 0, 2).float()   # (s, h, d)
+                    k_hat = torch.stack([q.roundtrip_k(kk[:, hh, :]) for hh in range(h)], 1)
+                    vv = value_states[0].permute(1, 0, 2).float()
+                    v_hat = torch.stack([q.roundtrip_v(vv[:, hh, :]) for hh in range(h)], 1)
+                    k_hat = k_hat.permute(1, 0, 2)[None]
+                    v_hat = v_hat.permute(1, 0, 2)[None]
+                else:
+                    k_hat = q.roundtrip_k(kf).reshape(s, h, d).permute(1, 0, 2)[None]
+                    v_hat = q.roundtrip_v(vf).reshape(s, h, d).permute(1, 0, 2)[None]
+                key_states = k_hat.to(key_states.dtype).to(key_states.device)
+                value_states = v_hat.to(value_states.dtype).to(value_states.device)
+            return super().update(key_states, value_states, layer_idx, cache_kwargs)
+    return VQCache
+
+
+# STAGE: cheap -- tier-1/2 metrics on held-out trace windows (no test suites)
+def stage_cheap(n_eval=64, max_len=16384):
+    import collections
+    from datasets import load_dataset
+    blob = torch.load(os.path.join(ARTIFACT_DIR, "codebooks.pt"), weights_only=False)
+    fitted, meta = blob["fitted"], blob["meta"]
+    hd, full = meta["head_dim"], meta["full_layers"]
+    n_q = meta.get("n_q_heads", 48)   # fallback for codebooks.pt written before this field
+
+    # Window for O(T²) attention metrics. 512 tokens keeps peak mem <200MB on GPU.
+    ATTN_WIN = 512
+
+    model, tok, _ = load_model_and_meta()
+    from tqdm import tqdm
+    from transformers.cache_utils import DynamicCache
+
+    # Move every fitted codebook to the model device ONCE. The roundtrip then
+    # runs entirely on-GPU against the on-GPU eval caches (see EvalDump below).
+    # ScalarKV/KIVI are tuning-free and have no .to(); they follow their input.
+    for per_layer in fitted.values():
+        for q in per_layer.values():
+            if hasattr(q, "to"):
+                q.to(model.device)
+
+    class EvalDump(DynamicCache):
+        def __init__(self):
+            super().__init__(); self.d = {i: {} for i in full}
+        def update(self, ks, vs, li, ck=None):
+            if li in set(full):
+                # Keep on-GPU: the quantizer roundtrip is a nearest-neighbour
+                # search -- a GPU job. Moving to CPU here dominated runtime.
+                self.d[li]["k"] = ks.detach()[0].permute(1, 0, 2).float()
+                self.d[li]["v"] = vs.detach()[0].permute(1, 0, 2).float()
+            return super().update(ks, vs, li, ck)
+
+    ds = load_dataset(CALIB_DATASET, split=CALIB_SPLIT)
+
+    # Per-trace rows (not saved to disk; aggregated below)
+    trace_rows = []
+
+    # Held-out slice. Start at 500 to safely clear any rows the dump stage
+    # consumed (dump skips short traces, so actual rows used >> n_calib=16).
+    for ex in tqdm(ds.select(range(500, 500 + n_eval))):
+        text = flatten_trace(ex, tok)
+        ids = tok(text, return_tensors="pt", truncation=True, max_length=max_len).to(model.device)
+        cache = EvalDump()
+        with torch.no_grad():
+            # model.model() skips lm_head: avoids a ~3 GB allocation per trace
+            # (max_len * vocab_size * 2 bytes) that is not needed for KV metrics.
+            model.model(**ids, past_key_values=cache, use_cache=True)
+
+        # Synthetic Q for attention/IP metrics: generated once per (trace, layer)
+        # and reused across all configs so the comparison is apples-to-apples.
+        # Unit-normalised so inner-product scale doesn't swamp the bias signal.
+        synth_q = {}
+        for i in full:
+            s = cache.d[i]["k"].shape[0]
+            win = min(s, ATTN_WIN)
+            q_rand = torch.randn(win, n_q, hd, device=cache.d[i]["k"].device)
+            synth_q[i] = q_rand / q_rand.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+
+        for name, _ in cache_configs():
+            if name == "fp16 (baseline)":
+                continue
+            per_layer = fitted[name]
+            acc = collections.defaultdict(float)
+            nL = 0
+            for i in full:
+                k = cache.d[i]["k"]      # (s, h, d)
+                v = cache.d[i]["v"]
+                q = per_layer[i]
+                s, h, d = k.shape
+
+                # Per-channel-key quantizers (KIVI, Sign, Ternary with
+                # per_channel_key=True) reduce the KEY along the TOKEN axis
+                # (dim=0). They must see one (s, d) block PER HEAD; flattening
+                # (s,h,d)->(s*h,d) would mix tokens across heads into one scale
+                # and corrupt the key metric. VQ/scalar-per-token quantizers
+                # reduce along dim=-1 and are safe to flatten.
+                per_channel_key = (
+                    isinstance(q, KIVIScalarKV)
+                    or getattr(q, "per_channel_key", False)
+                )
+                if per_channel_key:
+                    k_hat = torch.stack([q.roundtrip_k(k[:, hh, :]) for hh in range(h)], 1)
+                    v_hat = torch.stack([q.roundtrip_v(v[:, hh, :]) for hh in range(h)], 1)
+                else:
+                    k_hat = q.roundtrip_k(k.reshape(-1, d)).reshape(s, h, d)
+                    v_hat = q.roundtrip_v(v.reshape(-1, d)).reshape(s, h, d)
+
+                acc["key_cos"] += key_cosine(k, k_hat)
+                acc["val_cos"] += key_cosine(v, v_hat)
+                acc["key_mse"] += cache_mse(k, k_hat)
+                acc["val_mse"] += cache_mse(v, v_hat)
+
+                # Windowed attention/IP metrics on last ATTN_WIN tokens
+                win = min(s, ATTN_WIN)
+                kw, kw_hat = k[-win:], k_hat[-win:]
+                vw, vw_hat = v[-win:], v_hat[-win:]
+                q_syn = synth_q[i]                         # (win, n_q, d)
+
+                out_ref, _ = attention_output(q_syn, kw, vw, n_q)
+                out_hat, _ = attention_output(q_syn, kw_hat, vw_hat, n_q)
+                acc["attn_cos"]          += attn_output_cosine(out_ref, out_hat)
+                acc["attn_output_error"] += attn_output_error(out_ref, out_hat)
+
+                ip = inner_product_distortion(q_syn, kw, kw_hat)
+                acc["ip_rel"]  += ip["ip_rel_err"]
+                acc["ip_bias"] += ip["ip_bias"]
+
+                nL += 1
+
+            trace_rows.append({
+                "trace_len":          ids["input_ids"].shape[1],
+                "config":             name,
+                "key_cos":            acc["key_cos"]            / nL,
+                "val_cos":            acc["val_cos"]            / nL,
+                "key_mse":            acc["key_mse"]            / nL,
+                "val_mse":            acc["val_mse"]            / nL,
+                "attn_cos":           acc["attn_cos"]           / nL,
+                "attn_output_error":  acc["attn_output_error"]  / nL,
+                "ip_rel":             acc["ip_rel"]             / nL,
+                "ip_bias":            acc["ip_bias"]            / nL,
+            })
+
+    # Aggregate across traces: one summary row per config (what gets saved)
+    agg = collections.defaultdict(lambda: collections.defaultdict(list))
+    for r in trace_rows:
+        for col in ("key_cos", "val_cos", "key_mse", "val_mse",
+                    "attn_cos", "attn_output_error", "ip_rel", "ip_bias"):
+            agg[r["config"]][col].append(r[col])
+
+    COLS = ("key_cos", "val_cos", "key_mse", "val_mse",
+            "attn_cos", "attn_output_error", "ip_rel", "ip_bias")
+
+    summary = []
+    for name, _ in cache_configs():
+        if name not in agg:
+            continue
+        cols = agg[name]
+        n = len(cols["key_cos"])
+        q0 = next(iter(fitted[name].values()))
+        row = {"config": name, "bits_per_elt": round(q0.bits_per_element(hd), 4), "n_traces": n}
+        for col in COLS:
+            row[col] = round(sum(cols[col]) / n, 5)
+        summary.append(row)
+
+    print(f"\n[cheap] mean metrics over {n_eval} held-out traces:")
+    print(f"  {'config':24s}  {'bpe':>5}  {'key_cos':>8}  {'val_cos':>8}  "
+          f"{'key_mse':>9}  {'val_mse':>9}  {'attn_cos':>9}  {'attn_err':>9}  "
+          f"{'ip_rel':>8}  {'ip_bias':>9}")
+    for row in summary:
+        print(f"  {row['config']:24s}  {row['bits_per_elt']:5.2f}  "
+              f"{row['key_cos']:8.4f}  {row['val_cos']:8.4f}  "
+              f"{row['key_mse']:9.5f}  {row['val_mse']:9.5f}  "
+              f"{row['attn_cos']:9.4f}  {row['attn_output_error']:9.4f}  "
+              f"{row['ip_rel']:8.5f}  {row['ip_bias']:9.6f}")
+
+    out_path = os.path.join(ARTIFACT_DIR, "cheap_metrics.json")
+    json.dump(summary, open(out_path, "w"), indent=2)
+    print(f"[cheap] saved -> {out_path}")
+
+
+# Optional SWE-bench Verified eval (requires Docker + swebench).
+_AGENT_SYSTEM = (
+    "You are an expert software engineer fixing a GitHub issue. "
+    "You have a bash shell inside the repository checked out at the failing commit. "
+    "Use <bash>command</bash> tags to run shell commands. "
+    "Explore the code, implement the fix, then output <submit> when done."
+)
+
+
+def _agent_loop(model, tok, task: dict, cache_factory, max_turns: int,
+                max_new: int = 1024) -> str:
+    """Run a minimal ReAct-bash loop on one SWE-bench task.
+
+    Clones the repo at base_commit into a temp dir, runs the model in a
+    generate→bash→observe loop, and returns the final `git diff HEAD` patch.
+    Each generate call rebuilds the full context from scratch so the VQCache
+    sees the compounding long-context pressure that the project targets.
+    """
+    import re
+    import shutil
+    import subprocess
+    import tempfile
+
+    repo_dir = tempfile.mkdtemp(prefix="sweagent_")
+    try:
+        subprocess.run(
+            ["git", "clone", f"https://github.com/{task['repo']}.git", repo_dir],
+            check=True, capture_output=True, timeout=120,
+        )
+        subprocess.run(
+            ["git", "checkout", task["base_commit"]],
+            check=True, capture_output=True, cwd=repo_dir, timeout=30,
+        )
+    except Exception as exc:
+        print(f"  [agent] clone/checkout failed for {task['instance_id']}: {exc}")
+        shutil.rmtree(repo_dir, ignore_errors=True)
+        return ""
+
+    try:
+        messages = [
+            {"role": "system",  "content": _AGENT_SYSTEM},
+            {"role": "user",    "content": (
+                f"Repository: {task['repo']}\n\n"
+                f"Issue:\n{task['problem_statement']}"
+            )},
+        ]
+
+        for _ in range(max_turns):
+            prompt = tok.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True)
+            ids = tok(prompt, return_tensors="pt", truncation=True,
+                      max_length=32768).to(model.device)
+
+            cache = cache_factory()
+            with torch.no_grad():
+                out = model.generate(
+                    **ids, max_new_tokens=max_new, do_sample=False,
+                    past_key_values=cache, use_cache=True,
+                )
+            gen = tok.decode(out[0, ids["input_ids"].shape[1]:],
+                             skip_special_tokens=True)
+            messages.append({"role": "assistant", "content": gen})
+
+            if re.search(r"<submit\s*/?>", gen, re.I):
+                break
+
+            cmds = re.findall(r"<bash>(.*?)</bash>", gen, re.DOTALL)
+            if not cmds:
+                break  # model stopped issuing commands; take whatever diff we have
+
+            obs_parts = []
+            for cmd in cmds:
+                try:
+                    r = subprocess.run(
+                        cmd, shell=True, capture_output=True, text=True,
+                        timeout=30, cwd=repo_dir,
+                    )
+                    obs_parts.append(
+                        f"$ {cmd.strip()}\n{(r.stdout + r.stderr)[:2000]}")
+                except subprocess.TimeoutExpired:
+                    obs_parts.append(f"$ {cmd.strip()}\n[timeout after 30s]")
+            messages.append({"role": "user", "content": "\n\n".join(obs_parts)})
+
+        diff = subprocess.run(
+            ["git", "diff", "HEAD"],
+            capture_output=True, text=True, cwd=repo_dir,
+        )
+        return diff.stdout
+    finally:
+        shutil.rmtree(repo_dir, ignore_errors=True)
+
+
+def run_swebench_subset(model, tok, cache_factory_per_layer, target_layers,
+                        task_ids, max_turns=100):
+    """Run the official SWE-bench Verified harness over `task_ids`.
+
+    Builds per-config VQCache (or None for fp16), runs the mini-bash agent on
+    each task, writes predictions.jsonl, evaluates with the swebench harness,
+    and returns resolve rate in [0, 1].  Keep `task_ids` FIXED across configs.
+    """
+    import glob
+    from datasets import load_dataset
+
+    # Build cache factory: VQCache for quantized configs, None for fp16.
+    if cache_factory_per_layer is not None:
+        VQCls = make_vq_cache_class(cache_factory_per_layer, target_layers,
+                                    model.config, device=model.device)
+        def cache_factory(): return VQCls()
+    else:
+        def cache_factory(): return None  # model allocates DynamicCache internally
+
+    # Load task metadata indexed by instance_id.
+    verified = load_dataset("princeton-nlp/SWE-bench_Verified", split="test")
+    tasks = {r["instance_id"]: r for r in verified
+             if r["instance_id"] in set(task_ids)}
+
+    # Generate patches with the agent.
+    run_id = f"vqkv_{int(time.time())}"
+    predictions = []
+    for task_id in task_ids:
+        task = tasks.get(task_id)
+        if task is None:
+            print(f"[swebench] {task_id}: not in Verified dataset, skipping")
+            continue
+        print(f"[swebench] {task_id} ({len(predictions)+1}/{len(task_ids)})")
+        patch = _agent_loop(model, tok, task, cache_factory, max_turns=max_turns)
+        predictions.append({
+            "instance_id": task_id,
+            "model_patch": patch,
+            "model_name_or_path": "laguna-vqkv",
+        })
+
+    preds_path = os.path.join(ARTIFACT_DIR, f"{run_id}.jsonl")
+    with open(preds_path, "w") as f:
+        for p in predictions:
+            f.write(json.dumps(p) + "\n")
+    print(f"[swebench] wrote {len(predictions)} predictions -> {preds_path}")
+
+    # Run the official harness (needs Docker daemon).
+    # pip install swebench
+    from swebench.harness.run_evaluation import main as run_evaluation
+    run_evaluation(
+        dataset_name_or_path="princeton-nlp/SWE-bench_Verified",
+        split="test",
+        instance_ids=task_ids,
+        predictions_path=preds_path,
+        max_workers=4,
+        force_rebuild=False,
+        cache_level="env",
+        clean=False,
+        open_file_limit=4096,
+        run_id=run_id,
+        timeout=1800,
+    )
+
+    # Parse results. swebench writes a JSON summary; location varies by version.
+    result_files = (
+        glob.glob(os.path.join(ARTIFACT_DIR, f"{run_id}*.json"))
+        + glob.glob(f"{run_id}*.json")  # also check cwd
+    )
+    if not result_files:
+        print(f"[swebench] WARNING: no results file found for run_id={run_id}")
+        return 0.0
+
+    with open(result_files[0]) as f:
+        results = json.load(f)
+
+    if isinstance(results, list):
+        n_resolved = sum(1 for r in results if r.get("resolved", False))
+    elif isinstance(results, dict):
+        # some harness versions use {instance_id: {resolved: bool, ...}}
+        n_resolved = sum(1 for v in results.values()
+                         if (v.get("resolved") if isinstance(v, dict) else v))
+    else:
+        n_resolved = 0
+
+    return n_resolved / len(task_ids)
+
+
+def stage_swebench(n_tasks=50, seed=0):
+    import random
+    from datasets import load_dataset
+    blob = torch.load(os.path.join(ARTIFACT_DIR, "codebooks.pt"))
+    fitted, meta = blob["fitted"], blob["meta"]
+    model, tok, _ = load_model_and_meta()
+
+    verified = load_dataset("princeton-nlp/SWE-bench_Verified", split="test")
+    rng = random.Random(seed)
+    task_ids = [verified[i]["instance_id"]
+                for i in rng.sample(range(len(verified)), n_tasks)]
+    json.dump(task_ids, open(os.path.join(ARTIFACT_DIR, "task_subset.json"), "w"))
+    print(f"[swebench] fixed subset of {n_tasks} tasks (seed={seed}) saved.")
+
+    results = []
+    for name, _ in cache_configs():
+        quantizers = None if name == "fp16 (baseline)" else fitted[name]
+        try:
+            rate = run_swebench_subset(model, tok, quantizers,
+                                       meta["full_layers"], task_ids)
+        except NotImplementedError as e:
+            print(f"[swebench] {name}: STUB -- {e}")
+            rate = None
+        results.append({"config": name, "resolve_rate": rate})
+        print(f"[swebench] {name}: resolve_rate={rate}")
+
+    # report DELTAS vs fp16 (robust to absolute-score contamination)
+    base = next((r["resolve_rate"] for r in results
+                 if r["config"] == "fp16 (baseline)"), None)
+    print("\n[swebench] resolve rate on fixed subset (delta vs fp16):")
+    for r in results:
+        d = (None if (r["resolve_rate"] is None or base is None)
+             else round(r["resolve_rate"] - base, 4))
+        print(f"  {r['config']:24s} {r['resolve_rate']}  (Δ {d})")
+    json.dump(results, open(os.path.join(ARTIFACT_DIR, "swebench_results.json"), "w"),
+              indent=2)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--stage", required=True,
+                    choices=["dump", "fit", "cheap", "swebench"])
+    ap.add_argument("--n_calib", type=int, default=16)
+    ap.add_argument(
+        "--calib_source", type=str, default=None,
+        choices=["swesmith", "longbench-hotpotqa"],
+        help="dump stage: calibration corpus (default: CALIB_SOURCE env or swesmith)")
+    ap.add_argument("--n_eval", type=int, default=64)
+    ap.add_argument("--n_tasks", type=int, default=50)
+    ap.add_argument(
+        "--only", type=str, default=None,
+        help="fit stage: comma-separated config names to fit/merge (e.g. "
+             "sign-1bit,ternary-bitnet). Skips refitting other configs.")
+    args = ap.parse_args()
+
+    if args.stage == "dump":
+        stage_dump(n_calib=args.n_calib, calib_source=args.calib_source)
+    elif args.stage == "fit":
+        only = [s.strip() for s in args.only.split(",")] if args.only else None
+        stage_fit(only=only)
+    elif args.stage == "cheap":
+        stage_cheap(n_eval=args.n_eval)
+    elif args.stage == "swebench":
+        stage_swebench(n_tasks=args.n_tasks)
+
+
+if __name__ == "__main__":
+    main()

generate.py

@@ -0,0 +1,25 @@
+"""AttnVQ: VQQuantizedCache wired into model.generate()."""
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from vqkv.compressed_cache import VQQuantizedCache
+
+
+# load model
+tok = AutoTokenizer.from_pretrained("poolside/Laguna-XS.2", trust_remote_code=True, fix_mistral_regex=True)
+model = AutoModelForCausalLM.from_pretrained("poolside/Laguna-XS.2", torch_dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True).eval()
+
+# load codebooks or fit and use your own
+CODEBOOKS_PATH = "artifacts/codebooks.pt"
+codebooks = torch.load(CODEBOOKS_PATH, map_location="cuda", weights_only=False)
+
+# build cache
+quantizers, layers = codebooks["fitted"]["productvq-32x256-2b"], blob["meta"]["full_layers"]
+cache = VQQuantizedCache(quantizers, layers)  # persists uint8 codebook indices
+
+# generate
+ids = tok("Hello", return_tensors="pt").to(model.device)
+out = model.generate(**ids, max_new_tokens=32, past_key_values=cache, use_cache=True) 
+print(tok.decode(out[0, ids["input_ids"].shape[1]:], skip_special_tokens=True))
+
+# print memory footprint
+print(cache.memory_footprint())

longbench_eval.py

@@ -0,0 +1,763 @@
+"""
+longbench_eval.py — LongBench v1 evaluation for AttnVQ and baselines.
+
+Stages:
+    cheap     proxy KV metrics on task contexts (fast, default)
+    generate  end-to-end task scoring via model.generate (slow)
+
+Tasks: qasper, 2wikimqa, hotpotqa, passage_retrieval_en, repobench-p
+
+Usage:
+    python longbench_eval.py --stage cheap --n_eval 50
+    python longbench_eval.py --stage generate --tasks hotpotqa qasper
+"""
+
+from __future__ import annotations
+
+import argparse
+import collections
+import json
+import os
+import re
+import string
+from difflib import SequenceMatcher
+
+import torch
+from tqdm import tqdm
+
+from turbo_benchmark import TurboQuantMSE, QJLResidualIP, turbo_configs  # noqa: F401 — unpickle
+from vqkv.quantizers import KIVIScalarKV  # noqa: F401
+from vqkv.metrics import (key_cosine, cache_mse, inner_product_distortion,
+                          attention_output, attn_output_cosine, attn_output_error)
+
+MODEL_ID     = os.environ.get("LAGUNA_ID", "poolside/Laguna-XS.2")
+ARTIFACT_DIR = os.environ.get("VQKV_ARTIFACTS", "./artifacts")
+
+ATTN_WIN = 512
+CHEAP_COLS = ("key_cos", "val_cos", "key_mse", "val_mse",
+              "attn_cos", "attn_output_error", "ip_rel", "ip_bias")
+
+# Scorers (referenced from TASK_CONFIGS)
+def _normalize(s: str) -> str:
+    s = s.lower()
+    s = s.translate(str.maketrans("", "", string.punctuation))
+    return " ".join(s.split())
+
+
+def qa_f1_score(prediction: str, ground_truths: list[str], **_) -> float:
+    """Max token-level F1 over all reference answers (QA tasks)."""
+    pred_toks = _normalize(prediction).split()
+    best = 0.0
+    for ref in ground_truths:
+        ref_toks = _normalize(ref).split()
+        if not pred_toks or not ref_toks:
+            best = max(best, float(pred_toks == ref_toks))
+            continue
+        common = collections.Counter(pred_toks) & collections.Counter(ref_toks)
+        n_common = sum(common.values())
+        if n_common == 0:
+            continue
+        p = n_common / len(pred_toks)
+        r = n_common / len(ref_toks)
+        best = max(best, 2 * p * r / (p + r))
+    return best
+
+
+def retrieval_score(prediction: str, ground_truths: list[str], **_) -> float:
+    """Passage-retrieval accuracy: extract digit from prediction, match gold.
+
+    Ground truths are strings like 'Paragraph 3'; we extract the number and
+    check whether it appears among the first 10 digits in the prediction.
+    Mirrors the LongBench retrieval_score implementation.
+    """
+    gt_ids = set()
+    for ref in ground_truths:
+        m = re.findall(r"\d+", ref)
+        gt_ids.add(m[0] if m else _normalize(ref))
+    pred_nums = re.findall(r"\d+", prediction)
+    if not pred_nums:
+        return 0.0
+    right = sum(1 for n in pred_nums[:10] if n in gt_ids)
+    return right / len(pred_nums[:10])
+
+
+def classification_score(prediction: str, ground_truths: list[str],
+                         all_classes: list[str] | None = None, **_) -> float:
+    """Classification accuracy (TREC).
+
+    If all_classes is provided (from the dataset example), restrict matches
+    to valid class labels to avoid spurious substring hits — matches the
+    LongBench classification_score behaviour.
+    """
+    if all_classes:
+        matched = [c for c in all_classes if c.lower() in prediction.lower()]
+        return float(any(ref.lower() in [m.lower() for m in matched]
+                         for ref in ground_truths))
+    # Fallback: normalised substring match
+    pred_norm = _normalize(prediction)
+    return float(any(_normalize(ref) in pred_norm for ref in ground_truths))
+
+
+def edit_similarity_score(prediction: str, ground_truths: list[str], **_) -> float:
+    """Character-level edit similarity for code completion (RepoBench-P).
+
+    Uses difflib.SequenceMatcher.ratio() — equivalent to
+    fuzz.ratio() (not partial_ratio) and zero-dep.
+    """
+    pred = prediction.strip()
+    best = 0.0
+    for ref in ground_truths:
+        ref_s = ref.strip()
+        if not pred and not ref_s:
+            best = 1.0
+        elif pred and ref_s:
+            best = max(best, SequenceMatcher(None, pred, ref_s).ratio())
+    return best
+
+
+# ============================================================================
+# Task configs
+# ============================================================================
+# QA tasks: increasing context length (~3.6K → ~18K tokens) directly probes
+# the compounding-error-vs-length axis from RUNBOOK Step 6.
+# Three extra tasks cover retrieval accuracy, classification, and code completion.
+# ============================================================================
+TASK_CONFIGS: dict[str, dict] = {
+    "qasper": dict(
+        score_fn=qa_f1_score,
+        max_new_tokens=20,
+        max_len=16384,
+        prompt_template=(
+            "Answer the question based on the given passages. "
+            "Only give me the answer and do not output any other words.\n\n"
+            "The following are given passages.\n{context}\n\n"
+            "Answer the question based on the given passages. "
+            "Only give me the answer and do not output any other words.\n\n"
+            "Question: {input}\nAnswer:"
+        ),
+    ),
+    "2wikimqa": dict(
+        score_fn=qa_f1_score,
+        max_new_tokens=10,
+        max_len=16384,
+        prompt_template=(
+            "Answer the question based on the given passages. "
+            "Only give me the answer and do not output any other words.\n\n"
+            "The following are given passages.\n{context}\n\n"
+            "Answer the question based on the given passages. "
+            "Only give me the answer and do not output any other words.\n\n"
+            "Question: {input}\nAnswer:"
+        ),
+    ),
+    "hotpotqa": dict(
+        score_fn=qa_f1_score,
+        max_new_tokens=32,
+        max_len=32768,
+        prompt_template=(
+            "Answer the question based on the given passages. "
+            "Only give me the answer and do not output any other words.\n\n"
+            "The following are given passages.\n{context}\n\n"
+            "Answer the question based on the given passages. "
+            "Only give me the answer and do not output any other words.\n\n"
+            "Question: {input}\nAnswer:"
+        ),
+    ),
+    "passage_retrieval_en": dict(
+        score_fn=retrieval_score,
+        max_new_tokens=10,
+        max_len=32768,
+        prompt_template=(
+            "Below is a record of a series of paragraphs, each from different "
+            "documents. Tell me which paragraph the given passage is from.\n\n"
+            "{context}\n\n"
+            "Based on the above paragraphs, which paragraph does the following "
+            "passage come from? Only output the paragraph number. "
+            "Do not output any other characters.\n\n"
+            "{input}\nAnswer:"
+        ),
+    ),
+    "repobench-p": dict(
+        score_fn=edit_similarity_score,
+        max_new_tokens=64,
+        max_len=16384,
+        prompt_template=(
+            "Please complete the code given below.\n{context}\n{input}\n"
+        ),
+    ),
+}
+
+
+# ============================================================================
+# Model loading
+# ============================================================================
+def load_model_and_meta():
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, fix_mistral_regex=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID, torch_dtype=torch.bfloat16, device_map="cuda",
+        trust_remote_code=True)
+    model.eval()
+    cfg = model.config
+    full_layers = [i for i, t in enumerate(cfg.layer_types) if t == "full_attention"]
+    meta = {
+        "full_layers": full_layers,
+        "n_kv_heads":  cfg.num_key_value_heads,
+        "n_q_heads":   cfg.num_attention_heads,
+        "head_dim":    cfg.head_dim,
+    }
+    print(f"[meta] full-attention layers ({len(full_layers)}): {full_layers}")
+    print(f"[meta] kv_heads={meta['n_kv_heads']} head_dim={meta['head_dim']}")
+    return model, tok, meta
+
+
+# ============================================================================
+# Generic VQCache
+# ============================================================================
+def make_cache_class(per_layer_fns: dict, target_layers: list):
+    """Build a DynamicCache subclass that round-trips K/V through the given fns.
+
+    per_layer_fns[layer_idx] = {
+        "k_fn": callable (N, d) -> (N, d),
+        "v_fn": callable (N, d) -> (N, d),
+        "per_channel": bool,   # True for KIVI/Sign/Ternary (key dim=0 reduction)
+    }
+
+    For non-per-channel quantizers, K and V are flattened to (s*h, d) before
+    calling k_fn/v_fn. For per-channel, each head's (s, d) block is passed
+    separately so the quantizer's token-axis statistics are per-head (not
+    cross-head), matching the KIVI / benchmark.py convention.
+    """
+    from transformers.cache_utils import DynamicCache
+    _target = set(target_layers)
+
+    class VQCache(DynamicCache):
+        def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
+            if layer_idx in _target and layer_idx in per_layer_fns:
+                entry = per_layer_fns[layer_idx]
+                k_fn        = entry["k_fn"]
+                v_fn        = entry["v_fn"]
+                per_channel = entry.get("per_channel", False)
+
+                _b, h, s, d = key_states.shape
+
+                if per_channel:
+                    # key_states[0]: (h, s, d) -> (s, h, d)
+                    kk    = key_states[0].permute(1, 0, 2).float()
+                    k_hat = torch.stack([k_fn(kk[:, hh, :]) for hh in range(h)], dim=1)
+                    vv    = value_states[0].permute(1, 0, 2).float()
+                    v_hat = torch.stack([v_fn(vv[:, hh, :]) for hh in range(h)], dim=1)
+                    # (s, h, d) -> (h, s, d) -> (1, h, s, d)
+                    k_hat = k_hat.permute(1, 0, 2).unsqueeze(0)
+                    v_hat = v_hat.permute(1, 0, 2).unsqueeze(0)
+                else:
+                    # key_states[0]: (h, s, d) -> (s, h, d) -> (s*h, d)
+                    kf    = key_states[0].permute(1, 0, 2).reshape(-1, d).float()
+                    vf    = value_states[0].permute(1, 0, 2).reshape(-1, d).float()
+                    # (s*h, d) -> (s, h, d) -> (h, s, d) -> (1, h, s, d)
+                    k_hat = k_fn(kf).reshape(s, h, d).permute(1, 0, 2).unsqueeze(0)
+                    v_hat = v_fn(vf).reshape(s, h, d).permute(1, 0, 2).unsqueeze(0)
+
+                key_states   = k_hat.to(key_states.dtype).to(key_states.device)
+                value_states = v_hat.to(value_states.dtype).to(value_states.device)
+
+            return super().update(key_states, value_states, layer_idx, cache_kwargs)
+
+    return VQCache
+
+
+# ============================================================================
+# Build unified config list from both codebook files
+# ============================================================================
+def build_all_configs(meta: dict, device, only: list[str] | None = None):
+    """Return list of (name, bpe, cache_cls_or_None).
+
+    cache_cls_or_None: class (not instance) to call as cls() each generation,
+    or None for the fp16 baseline (DynamicCache allocated internally by model).
+    """
+    hd     = meta["head_dim"]
+    layers = meta["full_layers"]
+    configs: list[tuple[str, float, type | None]] = []
+
+    configs.append(("fp16", 16.0, None))
+
+    # -- Regular codebooks (ProductVQ, RoPESplit, Scalar, KIVI, Sign, Ternary) --
+    cb_path = os.path.join(ARTIFACT_DIR, "codebooks.pt")
+    if os.path.exists(cb_path):
+        blob   = torch.load(cb_path, weights_only=False)
+        fitted = blob["fitted"]
+        for name, per_layer in fitted.items():
+            for q in per_layer.values():
+                if hasattr(q, "to"):
+                    q.to(device)
+            q0          = next(iter(per_layer.values()))
+            bpe         = round(q0.bits_per_element(hd), 4)
+            per_channel = (isinstance(q0, KIVIScalarKV)
+                           or getattr(q0, "per_channel_key", False))
+            per_layer_fns = {
+                i: {"k_fn": per_layer[i].roundtrip_k,
+                    "v_fn": per_layer[i].roundtrip_v,
+                    "per_channel": per_channel}
+                for i in layers if i in per_layer
+            }
+            configs.append((name, bpe, make_cache_class(per_layer_fns, layers)))
+    else:
+        print(f"[warn] {cb_path} not found — skipping ProductVQ / scalar configs")
+
+    # -- TurboQuant codebooks --
+    tc_path = os.path.join(ARTIFACT_DIR, "turbo_codebooks.pt")
+    if os.path.exists(tc_path):
+        blob       = torch.load(tc_path, weights_only=False)
+        fitted     = blob["fitted"]
+        bits_list  = blob["bits"]
+        cfg_lookup = {n: (b, use_qjl) for n, b, use_qjl in turbo_configs(bits_list)}
+        for name, per_layer in fitted.items():
+            b, use_qjl = cfg_lookup[name]
+            bpe        = round(b + (1 if use_qjl else 0) + 16.0 / hd, 4)
+            for entry in per_layer.values():
+                entry["kq"].to(device)
+                entry["vq"].to(device)
+                if entry["qjl"] is not None:
+                    entry["qjl"].to(device)
+            per_layer_fns = {
+                i: {"k_fn": per_layer[i]["kq"].roundtrip,
+                    "v_fn": per_layer[i]["vq"].roundtrip,
+                    "per_channel": False}
+                for i in layers if i in per_layer
+            }
+            configs.append((name, bpe, make_cache_class(per_layer_fns, layers)))
+    else:
+        print(f"[warn] {tc_path} not found — skipping TurboQuant configs")
+
+    if only is not None:
+        only_set = set(only)
+        configs  = [(n, b, c) for n, b, c in configs if n in only_set]
+        unknown  = only_set - {n for n, _, _ in configs}
+        if unknown:
+            print(f"[warn] unknown --configs names: {sorted(unknown)}")
+
+    print(f"[build] {len(configs)} configs loaded")
+    for n, b, _ in configs:
+        print(f"  {n:<32}  {b:.3f} bpe")
+    return configs
+
+
+# ============================================================================
+# Per-task generation + scoring
+# ============================================================================
+def _load_longbench(task_name: str):
+    """Load a LongBench task, compatible with datasets >= 3.0.
+
+    datasets >= 3.0 dropped custom dataset-script support and raises
+    'Dataset scripts are no longer supported' for THUDM/LongBench.
+    We load the underlying JSONL files from the HF Hub directly instead.
+    Tries the English-suffixed file first (_e.jsonl), then the plain name.
+    """
+    from datasets import load_dataset as _ld
+    # Candidates in priority order:
+    #   {task}_e.jsonl   – English-suffixed bilingual tasks
+    #   {task}.jsonl     – single-language tasks
+    #   {task_underscored}.jsonl – hyphenated names (repobench-p → repobench_p)
+    slug = task_name.replace("-", "_")
+    candidates = [f"{task_name}_e.jsonl", f"{task_name}.jsonl",
+                  f"{slug}_e.jsonl",      f"{slug}.jsonl"]
+    # deduplicate while preserving order
+    seen: set[str] = set()
+    for fname in [c for c in candidates if not (c in seen or seen.add(c))]:
+        try:
+            return _ld(
+                "json",
+                data_files=f"hf://datasets/THUDM/LongBench/data/{fname}",
+                split="train",
+            )
+        except Exception:
+            continue
+    # Fallback for older datasets versions that still support scripts
+    return _ld("THUDM/LongBench", name=task_name, split="test")
+
+
+def prompt_text(task_cfg: dict, example: dict) -> str:
+    return task_cfg["prompt_template"].format(
+        context=example["context"], input=example["input"])
+
+
+def load_vqkv_fitted(device):
+    path = os.path.join(ARTIFACT_DIR, "codebooks.pt")
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"{path} not found; run benchmark.py --stage fit")
+    blob = torch.load(path, weights_only=False)
+    fitted, meta = blob["fitted"], blob["meta"]
+    for per_layer in fitted.values():
+        for q in per_layer.values():
+            if hasattr(q, "to"):
+                q.to(device)
+    return fitted, meta
+
+
+def load_turbo_fitted(device):
+    path = os.path.join(ARTIFACT_DIR, "turbo_codebooks.pt")
+    if not os.path.exists(path):
+        print(f"[warn] {path} not found — skipping TurboQuant configs")
+        return {}, {}
+    blob = torch.load(path, weights_only=False)
+    fitted = blob["fitted"]
+    cfg_lookup = {n: (b, use_qjl) for n, b, use_qjl in turbo_configs(blob["bits"])}
+    for per_layer in fitted.values():
+        for entry in per_layer.values():
+            entry["kq"].to(device)
+            entry["vq"].to(device)
+            if entry["qjl"] is not None:
+                entry["qjl"].to(device)
+    return fitted, cfg_lookup
+
+
+def _aggregate_cheap_rows(trace_rows: list[dict], bpe_lookup: dict) -> list[dict]:
+    agg = collections.defaultdict(lambda: collections.defaultdict(list))
+    for r in trace_rows:
+        for col in CHEAP_COLS:
+            agg[r["config"]][col].append(r[col])
+    summary = []
+    for name in bpe_lookup:
+        if name not in agg:
+            continue
+        cols = agg[name]
+        n = len(cols["key_cos"])
+        row = {"config": name, "bits_per_elt": bpe_lookup[name], "n_traces": n}
+        for col in CHEAP_COLS:
+            row[col] = round(sum(cols[col]) / n, 5)
+        summary.append(row)
+    return summary
+
+
+def _print_cheap_table(task_name: str, summary: list[dict]):
+    print(f"\n[cheap/{task_name}] mean metrics:")
+    print(f"  {'config':32s}  {'bpe':>5}  {'key_cos':>8}  {'val_cos':>8}  "
+          f"{'key_mse':>9}  {'val_mse':>9}  {'attn_cos':>9}  {'attn_err':>9}  "
+          f"{'ip_rel':>8}  {'ip_bias':>9}")
+    for row in summary:
+        print(f"  {row['config']:32s}  {row['bits_per_elt']:5.2f}  "
+              f"{row['key_cos']:8.4f}  {row['val_cos']:8.4f}  "
+              f"{row['key_mse']:9.5f}  {row['val_mse']:9.5f}  "
+              f"{row['attn_cos']:9.4f}  {row['attn_output_error']:9.4f}  "
+              f"{row['ip_rel']:8.5f}  {row['ip_bias']:9.6f}")
+
+
+def run_cheap_task(model, tok, meta, task_name: str, task_cfg: dict,
+                   vqkv_fitted: dict, turbo_fitted: dict,
+                   turbo_cfg_lookup: dict, bpe_lookup: dict,
+                   config_filter: set[str] | None, n_eval: int,
+                   min_len: int = 2048) -> list[dict]:
+    """Dump fp16 caches on n_eval LongBench prompts; score all quantizer configs."""
+    from transformers.cache_utils import DynamicCache
+
+    ds = _load_longbench(task_name)
+    subset = ds.select(range(min(n_eval, len(ds))))
+    max_len = task_cfg["max_len"]
+    full = meta["full_layers"]
+    hd = meta["head_dim"]
+    n_q = meta.get("n_q_heads", 48)
+    dev = model.device
+    if str(dev) == "meta":
+        dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    print(f"\n[cheap/{task_name}] {len(subset)} examples, max_len={max_len}")
+
+    class EvalDump(DynamicCache):
+        def __init__(self):
+            super().__init__()
+            self.d = {i: {} for i in full}
+
+        def update(self, ks, vs, li, ck=None):
+            if li in set(full):
+                self.d[li]["k"] = ks.detach()[0].permute(1, 0, 2).float()
+                self.d[li]["v"] = vs.detach()[0].permute(1, 0, 2).float()
+            return super().update(ks, vs, li, ck)
+
+    trace_rows: list[dict] = []
+    n_used = 0
+
+    for ex in tqdm(subset, desc=f"cheap/{task_name}"):
+        text = prompt_text(task_cfg, ex)
+        ids = tok(text, return_tensors="pt", truncation=True,
+                  max_length=max_len).to(dev)
+        if ids["input_ids"].shape[1] < min_len:
+            continue
+
+        cache = EvalDump()
+        with torch.no_grad():
+            model.model(**ids, past_key_values=cache, use_cache=True)
+        n_used += 1
+
+        synth_q = {}
+        for i in full:
+            s = cache.d[i]["k"].shape[0]
+            win = min(s, ATTN_WIN)
+            q_rand = torch.randn(win, n_q, hd, device=dev)
+            synth_q[i] = q_rand / q_rand.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+
+        # -- vqkv configs (ProductVQ, scalar, KIVI, …) --
+        for name, per_layer in vqkv_fitted.items():
+            if config_filter and name not in config_filter:
+                continue
+            acc = collections.defaultdict(float)
+            nL = 0
+            for i in full:
+                k = cache.d[i]["k"]
+                v = cache.d[i]["v"]
+                q = per_layer[i]
+                s, h, d = k.shape
+                per_channel = (
+                    isinstance(q, KIVIScalarKV)
+                    or getattr(q, "per_channel_key", False)
+                )
+                if per_channel:
+                    k_hat = torch.stack([q.roundtrip_k(k[:, hh, :]) for hh in range(h)], 1)
+                    v_hat = torch.stack([q.roundtrip_v(v[:, hh, :]) for hh in range(h)], 1)
+                else:
+                    k_hat = q.roundtrip_k(k.reshape(-1, d)).reshape(s, h, d)
+                    v_hat = q.roundtrip_v(v.reshape(-1, d)).reshape(s, h, d)
+
+                acc["key_cos"] += key_cosine(k, k_hat)
+                acc["val_cos"] += key_cosine(v, v_hat)
+                acc["key_mse"] += cache_mse(k, k_hat)
+                acc["val_mse"] += cache_mse(v, v_hat)
+
+                win = min(s, ATTN_WIN)
+                kw, kw_hat = k[-win:], k_hat[-win:]
+                vw, vw_hat = v[-win:], v_hat[-win:]
+                q_syn = synth_q[i]
+
+                out_ref, _ = attention_output(q_syn, kw, vw, n_q)
+                out_hat, _ = attention_output(q_syn, kw_hat, vw_hat, n_q)
+                acc["attn_cos"] += attn_output_cosine(out_ref, out_hat)
+                acc["attn_output_error"] += attn_output_error(out_ref, out_hat)
+
+                ip = inner_product_distortion(q_syn, kw, kw_hat)
+                acc["ip_rel"] += ip["ip_rel_err"]
+                acc["ip_bias"] += ip["ip_bias"]
+                nL += 1
+
+            trace_rows.append({
+                "task": task_name,
+                "trace_len": ids["input_ids"].shape[1],
+                "config": name,
+                **{col: acc[col] / nL for col in CHEAP_COLS},
+            })
+
+        # -- TurboQuant configs --
+        for name, per_layer in turbo_fitted.items():
+            if config_filter and name not in config_filter:
+                continue
+            acc = collections.defaultdict(float)
+            nL = 0
+            for i in full:
+                k = cache.d[i]["k"]
+                v = cache.d[i]["v"]
+                e = per_layer[i]
+                s, h, d = k.shape
+                k_hat = e["kq"].roundtrip(k.reshape(-1, d)).reshape(s, h, d)
+                v_hat = e["vq"].roundtrip(v.reshape(-1, d)).reshape(s, h, d)
+
+                acc["key_cos"] += key_cosine(k, k_hat)
+                acc["val_cos"] += key_cosine(v, v_hat)
+                acc["key_mse"] += cache_mse(k, k_hat)
+                acc["val_mse"] += cache_mse(v, v_hat)
+
+                win = min(s, ATTN_WIN)
+                kw, kw_hat = k[-win:], k_hat[-win:]
+                vw, vw_hat = v[-win:], v_hat[-win:]
+                q_syn = synth_q[i]
+
+                out_ref, _ = attention_output(q_syn, kw, vw, n_q)
+                out_hat, _ = attention_output(q_syn, kw_hat, vw_hat, n_q)
+                acc["attn_cos"] += attn_output_cosine(out_ref, out_hat)
+                acc["attn_output_error"] += attn_output_error(out_ref, out_hat)
+
+                q0 = q_syn[:, 0, :]
+                kr0 = kw[:, 0, :]
+                kh0 = kw_hat[:, 0, :]
+                qi = torch.randint(0, win, (4096,), device=dev)
+                ki = torch.randint(0, win, (4096,), device=dev)
+                ip_ref = (q0[qi] * kr0[ki]).sum(-1)
+                if e["qjl"] is not None:
+                    ip_hat = e["qjl"].estimate_ip(q0[qi], kr0[ki])
+                else:
+                    ip_hat = (q0[qi] * kh0[ki]).sum(-1)
+                acc["ip_bias"] += (ip_hat - ip_ref).mean().item()
+                acc["ip_rel"] += ((ip_hat - ip_ref).abs() /
+                                  ip_ref.abs().clamp_min(1e-6)).mean().item()
+                nL += 1
+
+            trace_rows.append({
+                "task": task_name,
+                "trace_len": ids["input_ids"].shape[1],
+                "config": name,
+                **{col: acc[col] / nL for col in CHEAP_COLS},
+            })
+
+    print(f"[cheap/{task_name}] used {n_used}/{len(subset)} traces (min_len={min_len})")
+    return trace_rows
+
+
+def stage_cheap(tasks: list[str], n_eval: int, config_filter: set[str] | None):
+    model, tok, model_meta = load_model_and_meta()
+    dev = model.device
+    if str(dev) == "meta":
+        dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    vqkv_fitted, cb_meta = load_vqkv_fitted(dev)
+    turbo_fitted, turbo_cfg_lookup = load_turbo_fitted(dev)
+    meta = cb_meta if cb_meta else model_meta
+    hd = meta["head_dim"]
+
+    bpe_lookup: dict[str, float] = {}
+    for name, per_layer in vqkv_fitted.items():
+        q0 = next(iter(per_layer.values()))
+        bpe_lookup[name] = round(q0.bits_per_element(hd), 4)
+    for name in turbo_cfg_lookup:
+        if name in turbo_fitted:
+            b, use_qjl = turbo_cfg_lookup[name]
+            bpe_lookup[name] = round(b + (1 if use_qjl else 0) + 16.0 / hd, 4)
+
+    if config_filter:
+        bpe_lookup = {k: v for k, v in bpe_lookup.items() if k in config_filter}
+        print(f"[cheap] config filter: {len(bpe_lookup)} configs")
+
+    all_trace_rows: list[dict] = []
+    all_summary: list[dict] = []
+
+    for task_name in tasks:
+        task_cfg = TASK_CONFIGS[task_name]
+        # trec few-shot prompts can be shorter than agentic traces
+        min_len = 512 if task_name == "trec" else 2048
+        rows = run_cheap_task(
+            model, tok, meta, task_name, task_cfg,
+            vqkv_fitted, turbo_fitted, turbo_cfg_lookup, bpe_lookup,
+            config_filter, n_eval, min_len=min_len,
+        )
+        all_trace_rows.extend(rows)
+        summary = _aggregate_cheap_rows(rows, bpe_lookup)
+        for row in summary:
+            row["task"] = task_name
+        all_summary.extend(summary)
+        _print_cheap_table(task_name, summary)
+
+    out_path = os.path.join(ARTIFACT_DIR, "longbench_cheap_metrics.json")
+    with open(out_path, "w") as fh:
+        json.dump(all_summary, fh, indent=2)
+    print(f"\n[cheap] saved -> {out_path}  ({len(all_summary)} task×config rows)")
+
+
+# ============================================================================
+# Per-task generation + scoring (--stage generate)
+# ============================================================================
+def run_task(model, tok, task_name: str, task_cfg: dict,
+             configs: list, n_eval: int) -> dict[str, float]:
+    """Run every config on n_eval examples. Returns {config_name: mean score}."""
+    ds     = _load_longbench(task_name)
+    subset = ds.select(range(min(n_eval, len(ds))))
+    print(f"\n[{task_name}] {len(subset)} examples, "
+          f"max_len={task_cfg['max_len']}, max_new={task_cfg['max_new_tokens']}")
+
+    score_fn = task_cfg.get("score_fn", qa_f1_score)
+    max_new  = task_cfg["max_new_tokens"]
+    max_len  = task_cfg["max_len"]
+    scores: dict[str, list[float]] = {name: [] for name, _, _ in configs}
+
+    for ex in tqdm(subset, desc=task_name):
+        prompt = prompt_text(task_cfg, ex)
+        refs   = (ex["answers"] if isinstance(ex["answers"], list)
+                  else [ex["answers"]])
+        # all_classes is present in TREC examples; ignored by other scorers
+        all_classes = ex.get("all_classes") or None
+        ids    = tok(prompt, return_tensors="pt", truncation=True,
+                     max_length=max_len).to(model.device)
+
+        for name, _bpe, cache_cls in configs:
+            cache = cache_cls() if cache_cls is not None else None
+            with torch.no_grad():
+                out = model.generate(
+                    **ids,
+                    max_new_tokens=max_new,
+                    do_sample=False,
+                    past_key_values=cache,
+                    use_cache=True,
+                )
+            pred = tok.decode(out[0, ids["input_ids"].shape[1]:],
+                              skip_special_tokens=True).strip()
+            scores[name].append(score_fn(pred, refs, all_classes=all_classes))
+
+    return {name: sum(vs) / len(vs) for name, vs in scores.items() if vs}
+
+
+# ============================================================================
+# Main
+# ============================================================================
+def stage_generate(tasks: list[str], n_eval: int, config_filter: list[str] | None):
+    model, tok, meta = load_model_and_meta()
+    device = model.device
+    if str(device) == "meta":
+        device = torch.device("cuda:0")
+
+    configs = build_all_configs(meta, device, only=config_filter)
+    print(f"\n[generate] {len(configs)} configs × {len(tasks)} tasks × "
+          f"{n_eval} examples each")
+
+    all_results: dict[str, dict[str, float]] = {}
+    for task_name in tasks:
+        all_results[task_name] = run_task(
+            model, tok, task_name, TASK_CONFIGS[task_name], configs, n_eval)
+
+    records = []
+    print(f"\n{'task':<22} {'config':<32} {'bpe':>5}  {'score':>6}  {'Δ':>8}")
+    print("-" * 78)
+    for task_name in tasks:
+        task_scores = all_results[task_name]
+        fp16_score = task_scores.get("fp16")
+        for name, bpe, _ in configs:
+            score = task_scores.get(name)
+            delta = (None if (score is None or fp16_score is None)
+                     else round(score - fp16_score, 4))
+            delta_str = f"{delta:+.4f}" if delta is not None else "       —"
+            print(f"{task_name:<22} {name:<32} {bpe:5.2f}  "
+                  f"{score:6.4f}  {delta_str}")
+            records.append({
+                "task": task_name,
+                "config": name,
+                "bpe": bpe,
+                "score": round(score, 4) if score is not None else None,
+                "delta": delta,
+            })
+        print()
+
+    out_path = os.path.join(ARTIFACT_DIR, "longbench_results.json")
+    with open(out_path, "w") as fh:
+        json.dump(records, fh, indent=2)
+    print(f"[generate] saved -> {out_path}")
+
+
+def main():
+    ap = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--stage", choices=["cheap", "generate"], default="cheap",
+                    help="cheap: KV proxy metrics (default); generate: task scores")
+    ap.add_argument("--tasks", nargs="+", default=list(TASK_CONFIGS),
+                    choices=list(TASK_CONFIGS),
+                    help="LongBench tasks to run (default: all 6)")
+    ap.add_argument("--n_eval", type=int, default=50,
+                    help="Examples per task (default: 50)")
+    ap.add_argument("--configs", nargs="+", default=None,
+                    help="Config names to include (default: all). "
+                         "E.g. --configs productvq-32x256-2b turbo-mse-2b")
+    args = ap.parse_args()
+
+    config_filter = set(args.configs) if args.configs else None
+
+    if args.stage == "cheap":
+        stage_cheap(args.tasks, args.n_eval, config_filter)
+    else:
+        stage_generate(args.tasks, args.n_eval, args.configs)
+
+
+if __name__ == "__main__":
+    main()

turbo_benchmark.py

@@ -0,0 +1,531 @@
+"""
+turbo_benchmark.py — faithful TurboQuant baseline (Zandieh et al., ICLR 2026).
+
+Haar rotation + per-coordinate Lloyd-Max codebooks + optional QJL residual for
+unbiased inner-product estimation. Separate from benchmark.py; writes
+artifacts/turbo_codebooks.pt.
+
+Usage:
+    python benchmark.py --stage dump   # or use existing calib_caches.pt
+    python turbo_benchmark.py --stage fit
+    python turbo_benchmark.py --stage cheap
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+
+import torch
+
+from vqkv.metrics import (key_cosine, cache_mse, inner_product_distortion,
+                          attention_output, attn_output_cosine, attn_output_error)
+
+ARTIFACT_DIR = os.environ.get("VQKV_ARTIFACTS", "./artifacts")
+MODEL_ID = os.environ.get("LAGUNA_ID", "poolside/Laguna-XS.2")
+CALIB_DATASET = os.environ.get("CALIB_DATASET", "SWE-bench/SWE-smith-trajectories")
+CALIB_SPLIT = os.environ.get("CALIB_SPLIT", "tool")
+EVAL_SOURCE = os.environ.get("EVAL_SOURCE", "swesmith")
+
+_HOTPOTQA_PROMPT = (
+    "Answer the question based on the given passages. "
+    "Only give me the answer and do not output any other words.\n\n"
+    "The following are given passages.\n{context}\n\n"
+    "Answer the question based on the given passages. "
+    "Only give me the answer and do not output any other words.\n\n"
+    "Question: {input}\nAnswer:"
+)
+
+
+# ============================================================================
+# Lloyd-Max 1-D scalar quantizer (the MSE-optimal levels for a given sample)
+# ============================================================================
+def lloyd_max_1d(samples: torch.Tensor, n_levels: int, iters: int = 30,
+                 seed: int = 0) -> torch.Tensor:
+    """Optimal scalar (Lloyd-Max) codebook for a 1-D distribution given samples.
+
+    samples: (M,) 1-D values drawn from the (concentrated) coordinate dist.
+    returns: (n_levels,) sorted reconstruction levels.
+
+    This is 1-D k-means; for the rotated unit-norm coordinates the distribution
+    is the same across coordinates in high dim, so one shared codebook per
+    bit-width suffices (matching the paper's single precomputed Beta codebook).
+    """
+    s = samples.flatten()
+    s = s[torch.isfinite(s)]
+    # init levels at quantiles so empty bins are rare
+    qs = torch.linspace(0.0, 1.0, n_levels + 2, device=s.device)[1:-1]
+    levels = torch.quantile(s, qs)
+    for _ in range(iters):
+        # assign each sample to nearest level
+        idx = torch.bucketize(s, (levels[1:] + levels[:-1]) / 2)
+        new = levels.clone()
+        for j in range(n_levels):
+            m = idx == j
+            if m.any():
+                new[j] = s[m].mean()
+        shift = (new - levels).abs().max()
+        levels = new
+        if shift < 1e-6:
+            break
+    return torch.sort(levels).values
+
+
+def _haar_rotation(d: int, seed: int = 0, device=None) -> torch.Tensor:
+    g = torch.Generator().manual_seed(seed)
+    a = torch.randn(d, d, generator=g)
+    q, r = torch.linalg.qr(a)
+    # sign-correct so Q is Haar-distributed (QR sign ambiguity)
+    q = q * torch.sign(torch.diag(r)).unsqueeze(0)
+    if device is not None:
+        q = q.to(device)
+    return q
+
+
+# ============================================================================
+# TurboQuant-MSE : rotation -> per-coordinate Lloyd-Max -> norm rescale
+# ============================================================================
+class TurboQuantMSE:
+    def __init__(self, nbits: int = 3, seed: int = 0):
+        self.nbits = nbits
+        self.seed = seed
+        self._rot = None       # (d, d)
+        self._levels = None    # (K,) shared Lloyd-Max levels
+
+    def fit(self, calib: torch.Tensor):
+        """calib: (N, d). Fit rotation + shared per-coordinate level set."""
+        d = calib.shape[-1]
+        self._rot = _haar_rotation(d, self.seed, calib.device)
+        xn = calib / calib.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+        r = xn @ self._rot                       # (N, d) rotated unit-norm coords
+        # all coordinates share the same concentrated dist -> pool them
+        pool = r.flatten()
+        if pool.numel() > 2_000_000:             # cap for speed
+            pool = pool[torch.randperm(pool.numel(), device=pool.device)[:2_000_000]]
+        self._levels = lloyd_max_1d(pool, 1 << self.nbits)
+        return self
+
+    def _quantize(self, x):
+        norms = x.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+        xn = x / norms
+        y = xn @ self._rot
+        edges = (self._levels[1:] + self._levels[:-1]) / 2
+        idx = torch.bucketize(y, edges).clamp(0, self._levels.numel() - 1)
+        return idx, norms
+
+    def roundtrip(self, x):
+        idx, norms = self._quantize(x)
+        y_hat = self._levels[idx]
+        x_hat = y_hat @ self._rot.T
+        return x_hat * norms
+
+    def to(self, device):
+        self._rot = self._rot.to(device)
+        self._levels = self._levels.to(device)
+        return self
+
+
+# ============================================================================
+# QJL inner-product estimator (faithful, ASYMMETRIC) -- NOT a reconstruction.
+#
+# Important: QJL does not reconstruct a vector. It estimates <q, k> directly,
+# with the query JL-transformed but UNQUANTIZED and only the key residual
+# sign-quantized (the asymmetric estimator of Zandieh et al.). It therefore
+# cannot be expressed as roundtrip(x)->x_hat without reintroducing the very
+# bias it removes. We expose it as a standalone estimator used only in the
+# inner-product-bias evaluation, not in the cosine/MSE reconstruction metrics.
+#
+#   <q, k> ~= <q, k_hat_mse>  +  (sqrt(pi/2) * ||r|| / m) * sum_i (g_i . q) * sign(g_i . r)
+#
+# where r = k_n - k_hat_mse_n is the residual in unit-norm space, g_i are the
+# rows of a gaussian JL matrix, and ||r|| is the stored residual norm.
+# ============================================================================
+class QJLResidualIP:
+    def __init__(self, mse: "TurboQuantMSE", qjl_rows: int = None, seed: int = 0):
+        self.mse = mse
+        self.seed = seed
+        self.qjl_rows = qjl_rows
+        self._G = None
+
+    def fit(self, calib: torch.Tensor):
+        d = calib.shape[-1]
+        m = self.qjl_rows or d
+        g = torch.Generator().manual_seed(self.seed + 7)
+        self._G = torch.randn(m, d, generator=g).to(calib.device)
+        return self
+
+    def estimate_ip(self, q, k):
+        """Unbiased estimate of rowwise <q, k> using MSE stage + QJL residual.
+        q, k: (N, d). Returns (N,) inner-product estimates."""
+        m = self._G.shape[0]
+        # stage-1 reconstruction in unit-norm space
+        knorm = k.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+        kn = k / knorm
+        idx, _ = self.mse._quantize(k)
+        kn_hat = (self.mse._levels[idx]) @ self.mse._rot.T
+        score1 = (q * (kn_hat * knorm)).sum(-1)           # <q, k_hat_mse>
+        # stage-2 residual correction (asymmetric: q unquantized, residual signed)
+        r = (kn - kn_hat)
+        rnorm = r.norm(dim=-1, keepdim=True).clamp_min(1e-12)
+        rn = r / rnorm
+        qg = q @ self._G.T                                # (N, m)  query JL, unquantized
+        rs = torch.sign(rn @ self._G.T)                   # (N, m)  residual sign bits
+        score2 = ((torch.pi / 2) ** 0.5 / m) * (qg * rs).sum(-1) * rnorm.squeeze(-1) * knorm.squeeze(-1)
+        return score1 + score2
+
+    def to(self, device):
+        self.mse.to(device)
+        self._G = self._G.to(device)
+        return self
+
+
+# ============================================================================
+# Config table. Each TurboQuant variant uses TurboQuant-MSE for RECONSTRUCTION
+# (the metric that feeds cosine/MSE). For keys we ALSO build a QJL inner-product
+# estimator (the faithful Prod second stage), reported only in the ip_bias
+# column since QJL is an IP estimator, not a reconstruction.
+#
+# `qjl=True` marks variants that add the QJL key estimator. 1-bit is included
+# as the aggressive floor; at 1 bit the MSE stage is effectively sign, so the
+# QJL residual carries most of the IP fidelity there.
+# ============================================================================
+def turbo_configs(bits_list):
+    cfgs = []
+    for b in bits_list:
+        cfgs.append((f"turbo-mse-{b}b", b, False))
+        cfgs.append((f"turbo-prod-{b}b (K:+qjl)", b, True))
+    return cfgs
+
+
+# ============================================================================
+# STAGE: fit -- per-layer codebooks from calib_caches.pt
+# ============================================================================
+def stage_fit(bits_list):
+    path = os.path.join(ARTIFACT_DIR, "calib_caches.pt")
+    if not os.path.exists(path):
+        raise FileNotFoundError(
+            f"{path} not found. Run `python benchmark.py --stage dump` first "
+            f"to produce the calibration caches this script reuses.")
+    blob = torch.load(path, weights_only=False)
+    calib, meta = blob["calib"], blob["meta"]
+    hd = meta["head_dim"]
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+
+    fitted = {}
+    for name, b, use_qjl in turbo_configs(bits_list):
+        per_layer = {}
+        t0 = time.time()
+        for i, c in calib.items():
+            kf = c["k"].reshape(-1, hd)[:200_000].to(dev)
+            vf = c["v"].reshape(-1, hd)[:200_000].to(dev)
+            kq = TurboQuantMSE(nbits=b).fit(kf)
+            vq = TurboQuantMSE(nbits=b).fit(vf)
+            qjl = QJLResidualIP(kq).fit(kf) if use_qjl else None
+            per_layer[i] = {"kq": kq, "vq": vq, "qjl": qjl}
+        fitted[name] = per_layer
+        print(f"[fit] {name}: {len(per_layer)} layer-codebooks in {time.time()-t0:.1f}s")
+    torch.save({"fitted": fitted, "meta": meta, "bits": bits_list},
+               os.path.join(ARTIFACT_DIR, "turbo_codebooks.pt"))
+    print(f"[fit] saved -> {ARTIFACT_DIR}/turbo_codebooks.pt")
+
+
+# ============================================================================
+# Trace flattening -- minimal copy (kept independent of benchmark.py)
+# ============================================================================
+def flatten_trace(example, tok) -> str:
+    raw = (example.get("messages") or example.get("trajectory")
+           or example.get("conversations"))
+    if raw is None:
+        return json.dumps(example)[:200_000]
+    if isinstance(raw, str):
+        try:
+            raw = json.loads(raw)
+        except json.JSONDecodeError:
+            return raw[:200_000]
+    norm = []
+    for m in raw:
+        role = m.get("role") or m.get("from") or "user"
+        content = m.get("content") or m.get("value") or ""
+        if isinstance(content, list):
+            content = "\n".join(
+                it.get("text", str(it)) if isinstance(it, dict) else str(it)
+                for it in content)
+        tc = m.get("tool_calls")
+        if tc:
+            tct = json.dumps(tc, ensure_ascii=False)
+            content = (content + "\n" + tct).strip() if content else tct
+        role = {"human": "user", "gpt": "assistant", "tool": "user"}.get(role, role)
+        if not content.strip():
+            continue
+        norm.append({"role": role, "content": content})
+    merged = []
+    for m in norm:
+        if merged and merged[-1]["role"] == m["role"]:
+            merged[-1]["content"] += "\n\n" + m["content"]
+        else:
+            merged.append(dict(m))
+    try:
+        return tok.apply_chat_template(merged, tokenize=False,
+                                       add_generation_prompt=False)
+    except Exception:
+        return "\n\n".join(f"{m['role']}: {m['content']}" for m in merged)
+
+
+def flatten_longbench(example) -> str:
+    """Format a LongBench hotpotqa example (context + input) as a plain string."""
+    return _HOTPOTQA_PROMPT.format(
+        context=example["context"], input=example["input"])
+
+
+def _load_longbench_hotpotqa():
+    """Load THUDM/LongBench hotpotqa, bypassing the deprecated dataset script."""
+    from datasets import load_dataset as _ld
+    for fname in ("hotpotqa_e.jsonl", "hotpotqa.jsonl"):
+        try:
+            return _ld(
+                "json",
+                data_files=f"hf://datasets/THUDM/LongBench/data/{fname}",
+                split="train",
+            )
+        except Exception:
+            continue
+    return _ld("THUDM/LongBench", name="hotpotqa", split="test")
+
+
+def _load_cheap_eval_dataset(n_eval: int, eval_source: str, tok):
+    """Return (dataset_slice, get_text, max_len, label) for stage_cheap."""
+    from datasets import load_dataset
+
+    if eval_source == "longbench-hotpotqa":
+        ds = _load_longbench_hotpotqa()
+        # First n_eval rows — disjoint from dump calib (last n_calib rows).
+        end = min(n_eval, len(ds))
+        ds = ds.select(range(0, end))
+        label = f"LongBench hotpotqa (rows 0–{end - 1})"
+        return ds, flatten_longbench, 32768, label
+
+    ds = load_dataset(CALIB_DATASET, split=CALIB_SPLIT)
+    start = 500  # held-out offset (matches benchmark.py --stage cheap)
+    end = min(start + n_eval, len(ds))
+    ds = ds.select(range(start, end))
+    label = f"{CALIB_DATASET} split={CALIB_SPLIT} (rows {start}–{end - 1})"
+    return ds, lambda ex: flatten_trace(ex, tok), 16384, label
+
+
+def load_model_and_meta():
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID, dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True)
+    model.eval()
+    cfg = model.config
+    full = [i for i, t in enumerate(cfg.layer_types) if t == "full_attention"]
+    meta = {"full_layers": full, "n_kv_heads": cfg.num_key_value_heads,
+            "n_q_heads": cfg.num_attention_heads,
+            "head_dim": cfg.head_dim, "n_layers": cfg.num_hidden_layers}
+    print(f"[meta] full-attention layers ({len(full)}): {full}")
+    return model, tok, meta
+
+
+# ============================================================================
+# STAGE: cheap -- full metric set on held-out traces (mirrors benchmark.py)
+# ============================================================================
+def stage_cheap(n_eval=64, max_len=None, eval_source: str | None = None,
+                min_len=2048):
+    import collections
+    from tqdm import tqdm
+    from transformers.cache_utils import DynamicCache
+
+    eval_source = eval_source or EVAL_SOURCE
+    blob = torch.load(os.path.join(ARTIFACT_DIR, "turbo_codebooks.pt"),
+                      weights_only=False)
+    fitted, meta = blob["fitted"], blob["meta"]
+    full = meta["full_layers"]
+    hd   = meta["head_dim"]
+    n_q  = meta.get("n_q_heads", 48)   # fallback for turbo_codebooks.pt without this field
+
+    # Window for O(T²) attention metrics (matches benchmark.py)
+    ATTN_WIN = 512
+
+    # Lookup: config name -> (nbits, use_qjl) for bits_per_elt reporting
+    cfg_lookup = {name: (b, use_qjl) for name, b, use_qjl in turbo_configs(blob["bits"])}
+
+    model, tok, _ = load_model_and_meta()
+    dev = model.device
+    if str(dev) == "meta":
+        dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    ds, get_text, default_max_len, source_label = _load_cheap_eval_dataset(
+        n_eval, eval_source, tok)
+    if max_len is None:
+        max_len = default_max_len
+    print(f"[turbo-cheap] eval_source={eval_source}  {source_label}  "
+          f"rows={len(ds)}  max_len={max_len}")
+
+    for per_layer in fitted.values():
+        for entry in per_layer.values():
+            entry["kq"].to(dev); entry["vq"].to(dev)
+            if entry["qjl"] is not None:
+                entry["qjl"].to(dev)
+
+    class EvalDump(DynamicCache):
+        def __init__(self):
+            super().__init__(); self.d = {i: {} for i in full}
+        def update(self, ks, vs, li, ck=None):
+            if li in set(full):
+                self.d[li]["k"] = ks.detach()[0].permute(1, 0, 2).float()
+                self.d[li]["v"] = vs.detach()[0].permute(1, 0, 2).float()
+            return super().update(ks, vs, li, ck)
+
+    trace_rows = []
+    n_used = 0
+
+    for ex in tqdm(ds, desc=f"turbo-cheap/{eval_source}"):
+        text = get_text(ex)
+        ids = tok(text, return_tensors="pt", truncation=True,
+                  max_length=max_len).to(dev)
+        if ids["input_ids"].shape[1] < min_len:
+            continue
+        cache = EvalDump()
+        with torch.no_grad():
+            model.model(**ids, past_key_values=cache, use_cache=True)
+        n_used += 1
+
+        # Synthetic Q: generated once per (trace, layer), reused across configs.
+        # Using synthetic Q for the QJL estimator is more faithful than the
+        # previous keys-as-proxy approach: q0[qi] are proper query vectors,
+        # kr0[ki] are the keys being estimated.
+        synth_q = {}
+        for i in full:
+            s = cache.d[i]["k"].shape[0]
+            win = min(s, ATTN_WIN)
+            q_rand = torch.randn(win, n_q, hd, device=dev)
+            synth_q[i] = q_rand / q_rand.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+
+        for name, per_layer in fitted.items():
+            acc = collections.defaultdict(float)
+            nL = 0
+            for i in full:
+                k = cache.d[i]["k"]; v = cache.d[i]["v"]   # (s, h, d)
+                s, h, d = k.shape
+                e = per_layer[i]
+                k_hat = e["kq"].roundtrip(k.reshape(-1, d)).reshape(s, h, d)
+                v_hat = e["vq"].roundtrip(v.reshape(-1, d)).reshape(s, h, d)
+
+                acc["key_cos"] += key_cosine(k, k_hat)
+                acc["val_cos"] += key_cosine(v, v_hat)
+                acc["key_mse"] += cache_mse(k, k_hat)
+                acc["val_mse"] += cache_mse(v, v_hat)
+
+                # Windowed attention/IP metrics on last ATTN_WIN tokens
+                win = min(s, ATTN_WIN)
+                kw, kw_hat = k[-win:], k_hat[-win:]
+                vw, vw_hat = v[-win:], v_hat[-win:]
+                q_syn = synth_q[i]                          # (win, n_q, d)
+
+                out_ref, _ = attention_output(q_syn, kw, vw, n_q)
+                out_hat, _ = attention_output(q_syn, kw_hat, vw_hat, n_q)
+                acc["attn_cos"]          += attn_output_cosine(out_ref, out_hat)
+                acc["attn_output_error"] += attn_output_error(out_ref, out_hat)
+
+                # IP metrics. For Prod configs, use the faithful QJL asymmetric
+                # estimator (designed for unbiased <q,k> estimation). For MSE
+                # configs, use the plain reconstructed inner product.
+                q0  = q_syn[:, 0, :]    # (win, d) — head 0 of synthetic Q
+                kr0 = kw[:, 0, :]       # (win, d) — head 0 of reference K
+                kh0 = kw_hat[:, 0, :]   # (win, d) — head 0 of reconstructed K
+                qi  = torch.randint(0, win, (4096,), device=dev)
+                ki  = torch.randint(0, win, (4096,), device=dev)
+                ip_ref = (q0[qi] * kr0[ki]).sum(-1)
+                if e["qjl"] is not None:
+                    ip_hat = e["qjl"].estimate_ip(q0[qi], kr0[ki])
+                else:
+                    ip_hat = (q0[qi] * kh0[ki]).sum(-1)
+                acc["ip_bias"] += (ip_hat - ip_ref).mean().item()
+                acc["ip_rel"]  += ((ip_hat - ip_ref).abs() /
+                                   ip_ref.abs().clamp_min(1e-6)).mean().item()
+                nL += 1
+
+            trace_rows.append({
+                "trace_len":         ids["input_ids"].shape[1],
+                "config":            name,
+                "key_cos":           acc["key_cos"]           / nL,
+                "val_cos":           acc["val_cos"]           / nL,
+                "key_mse":           acc["key_mse"]           / nL,
+                "val_mse":           acc["val_mse"]           / nL,
+                "attn_cos":          acc["attn_cos"]          / nL,
+                "attn_output_error": acc["attn_output_error"] / nL,
+                "ip_rel":            acc["ip_rel"]            / nL,
+                "ip_bias":           acc["ip_bias"]           / nL,
+            })
+
+    # Aggregate across traces: one summary row per config
+    agg = collections.defaultdict(lambda: collections.defaultdict(list))
+    for r in trace_rows:
+        for col in ("key_cos", "val_cos", "key_mse", "val_mse",
+                    "attn_cos", "attn_output_error", "ip_rel", "ip_bias"):
+            agg[r["config"]][col].append(r[col])
+
+    COLS = ("key_cos", "val_cos", "key_mse", "val_mse",
+            "attn_cos", "attn_output_error", "ip_rel", "ip_bias")
+
+    summary = []
+    for name in cfg_lookup:
+        if name not in agg:
+            continue
+        b, use_qjl = cfg_lookup[name]
+        # bpe: nbits/coord + 1 bit/coord for QJL signs (m=d rows) + fp16 norm
+        bpe = b + (1 if use_qjl else 0) + 16.0 / hd
+        cols = agg[name]
+        n = len(cols["key_cos"])
+        row = {"config": name, "bits_per_elt": round(bpe, 4),
+               "n_traces": n, "eval_source": eval_source}
+        for col in COLS:
+            row[col] = round(sum(cols[col]) / n, 5)
+        summary.append(row)
+
+    print(f"\n[turbo-cheap] mean metrics over {n_used} traces ({eval_source}):")
+    print(f"  {'config':30s}  {'bpe':>5}  {'key_cos':>8}  {'val_cos':>8}  "
+          f"{'key_mse':>9}  {'val_mse':>9}  {'attn_cos':>9}  {'attn_err':>9}  "
+          f"{'ip_rel':>8}  {'ip_bias':>9}")
+    for row in summary:
+        print(f"  {row['config']:30s}  {row['bits_per_elt']:5.2f}  "
+              f"{row['key_cos']:8.4f}  {row['val_cos']:8.4f}  "
+              f"{row['key_mse']:9.5f}  {row['val_mse']:9.5f}  "
+              f"{row['attn_cos']:9.4f}  {row['attn_output_error']:9.4f}  "
+              f"{row['ip_rel']:8.5f}  {row['ip_bias']:9.6f}")
+
+    out_name = ("turbo_cheap_metrics_hotpotqa.json"
+                if eval_source == "longbench-hotpotqa"
+                else "turbo_cheap_metrics.json")
+    out_path = os.path.join(ARTIFACT_DIR, out_name)
+    json.dump(summary, open(out_path, "w"), indent=2)
+    print(f"[turbo-cheap] saved -> {out_path}")
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__,
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--stage", required=True, choices=["fit", "cheap"])
+    ap.add_argument("--bits", type=int, nargs="+", default=[4, 2, 1],
+                    help="bit-widths to sweep (default: 4 2 1)")
+    ap.add_argument("--n_eval", type=int, default=64)
+    ap.add_argument(
+        "--eval_source", type=str, default=None,
+        choices=["swesmith", "longbench-hotpotqa"],
+        help="cheap stage: eval corpus (default: EVAL_SOURCE env or swesmith)")
+    ap.add_argument("--max_len", type=int, default=None,
+                    help="cheap stage: max prompt tokens (default: 16384 swesmith, "
+                         "32768 longbench-hotpotqa)")
+    args = ap.parse_args()
+    if args.stage == "fit":
+        stage_fit(args.bits)
+    elif args.stage == "cheap":
+        stage_cheap(n_eval=args.n_eval, max_len=args.max_len,
+                    eval_source=args.eval_source)
+
+
+if __name__ == "__main__":
+    main()

vqkv/compressed_cache.py

@@ -0,0 +1,245 @@
+"""
+vqkv.compressed_cache — inference-ready KV cache with real memory savings.
+
+VQQuantizedCache persists uint8 codebook indices for target layers (not bf16
+reconstructions). Dequantizes transiently per layer on each attention read.
+Without a fused kernel this saves memory but not wall-clock; see memory_footprint().
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+try:
+    from vqkv.quantizers import ProductVQKV, ScalarKV, KIVIScalarKV
+except ModuleNotFoundError:  # flat layout (files beside this one)
+    from quantizers import ProductVQKV, ScalarKV, KIVIScalarKV
+
+
+# ----------------------------------------------------------------------------
+# encode / decode for a *fitted* ProductVQKV  (no change to quantizers.py)
+# ----------------------------------------------------------------------------
+# These reach into the already-built batched stacks (cb, mu, sd) that
+# ProductVQKV._ensure_stacked() prepares, so encode(x) then decode(idx) is
+# bit-identical to the existing q.roundtrip_k(x). Promote them to methods on
+# ProductVQKV if you prefer; kept as free functions to leave your file untouched.
+
+def _stacked(q: ProductVQKV, which: str):
+    q._ensure_stacked()
+    return q._k_stacked if which == "k" else q._v_stacked
+
+
+def pvq_encode(q: ProductVQKV, x: torch.Tensor, which: str = "k") -> torch.Tensor:
+    """x: (N, head_dim) -> idx: (N, n_sub) uint8 (nearest codeword per sub-block)."""
+    cb, mu, sd = _stacked(q, which)            # cb: (n_sub, K, sub_dim)
+    n_sub, K, sub_dim = cb.shape
+    if K > 256:
+        raise ValueError(f"n_codes={K} > 256 needs >1 byte/index; this path is uint8. "
+                         f"Use K<=256 (all headline configs do) or add bit-packing.")
+    N = x.shape[0]
+    c_sq = (cb * cb).sum(-1).unsqueeze(1)      # (n_sub, 1, K)
+    mu_b = mu.permute(1, 0, 2) if mu is not None else None
+    sd_b = sd.permute(1, 0, 2) if sd is not None else None
+    chunk = max(1, (256 * 1024 * 1024) // (n_sub * K * 4))
+    out = []
+    for s0 in range(0, N, chunk):
+        xc = x[s0:s0 + chunk]
+        c = xc.shape[0]
+        xb = xc.reshape(c, n_sub, sub_dim).permute(1, 0, 2).contiguous()
+        if mu_b is not None:
+            xb = (xb - mu_b) / sd_b
+        x_sq = (xb * xb).sum(-1, keepdim=True)            # (n_sub, c, 1)
+        cross = torch.bmm(xb, cb.transpose(1, 2))         # (n_sub, c, K)
+        idx = (x_sq - 2 * cross + c_sq).argmin(-1)        # (n_sub, c)
+        out.append(idx.to(torch.uint8))
+    return torch.cat(out, dim=1).T.contiguous()           # (N, n_sub) uint8
+
+
+def pvq_decode(q: ProductVQKV, idx: torch.Tensor, which: str = "k") -> torch.Tensor:
+    """idx: (N, n_sub) uint8 -> x_hat: (N, head_dim) (codebook dtype, e.g. fp32)."""
+    cb, mu, sd = _stacked(q, which)
+    n_sub, K, sub_dim = cb.shape
+    mu_b = mu.permute(1, 0, 2) if mu is not None else None
+    sd_b = sd.permute(1, 0, 2) if sd is not None else None
+    idxT = idx.T.long()                                   # (n_sub, N)
+    rec = torch.gather(cb, 1, idxT.unsqueeze(-1).expand(-1, -1, sub_dim))  # (n_sub,N,sub_dim)
+    if mu_b is not None:
+        rec = rec * sd_b + mu_b
+    N = idxT.shape[1]
+    return rec.permute(1, 0, 2).reshape(N, n_sub * sub_dim)
+
+
+def pvq_codebook_bytes(q: ProductVQKV) -> int:
+    """Fixed per-layer codebook overhead (bytes), amortized over all tokens."""
+    cb_k, _, _ = _stacked(q, "k")
+    cb_v, _, _ = _stacked(q, "v")
+    return cb_k.numel() * cb_k.element_size() + cb_v.numel() * cb_v.element_size()
+
+
+# ----------------------------------------------------------------------------
+# Inference-ready cache: compressed store for target layers, native for rest
+# ----------------------------------------------------------------------------
+try:
+    from transformers.cache_utils import DynamicCache
+    _HAVE_TF = True
+except Exception:  # let the file import for the standalone memory harness
+    DynamicCache = object
+    _HAVE_TF = False
+
+
+class VQQuantizedCache(DynamicCache):
+    """Drop-in cache that persists ProductVQ indices for ``target_layers`` and
+    leaves all other layers in native precision (Laguna's sliding-window layers
+    are bounded at 512 tokens and don't dominate, so we don't touch them).
+
+    Memory model: ``key_cache``/``value_cache`` for target layers are never
+    populated with full tensors. We keep ``k_codes[layer]`` / ``v_codes[layer]``
+    as (seq, n_kv_heads, n_sub) uint8 and dequantize the whole buffer transiently
+    each time the layer runs attention.
+
+    NOTE ON transformers VERSIONS: signatures around DynamicCache shift between
+    releases. This targets the modern
+    ``update(key_states, value_states, layer_idx, cache_kwargs=None) -> (k, v)``
+    interface and overrides ``get_seq_length``. If your pinned version calls
+    additional hooks (``reorder_cache`` for beam search, ``crop`` for assisted
+    decoding), forward them to the code buffers the same way ``update`` does.
+    """
+
+    def __init__(self, per_layer_quantizers: dict, target_layers, *a, **k):
+        super().__init__(*a, **k)
+        self.q = per_layer_quantizers
+        self.target = set(int(i) for i in target_layers)
+        self.k_codes: dict[int, torch.Tensor] = {}
+        self.v_codes: dict[int, torch.Tensor] = {}
+        self._dtype = None
+        self._device = None
+
+    # -- the hot path --------------------------------------------------------
+    def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
+        if layer_idx not in self.target or layer_idx not in self.q:
+            # native path for sliding-window / non-targeted layers
+            return super().update(key_states, value_states, layer_idx, cache_kwargs)
+
+        self._dtype = key_states.dtype
+        self._device = key_states.device
+        q = self.q[layer_idx]
+        b, h, s, d = key_states.shape  # b == 1 in this harness
+
+        # encode the NEW tokens (token-major, head-minor rows -> (s, h, n_sub))
+        kf = key_states[0].transpose(0, 1).reshape(-1, d).float()
+        vf = value_states[0].transpose(0, 1).reshape(-1, d).float()
+        kc = pvq_encode(q, kf, "k").reshape(s, h, -1)
+        vc = pvq_encode(q, vf, "v").reshape(s, h, -1)
+
+        # append to the persistent compressed buffer
+        if layer_idx in self.k_codes:
+            self.k_codes[layer_idx] = torch.cat([self.k_codes[layer_idx], kc], dim=0)
+            self.v_codes[layer_idx] = torch.cat([self.v_codes[layer_idx], vc], dim=0)
+        else:
+            self.k_codes[layer_idx] = kc
+            self.v_codes[layer_idx] = vc
+
+        # transiently dequantize the FULL buffer for this layer's attention
+        allk, allv = self.k_codes[layer_idx], self.v_codes[layer_idx]
+        S = allk.shape[0]
+        kfull = pvq_decode(q, allk.reshape(-1, allk.shape[-1]), "k").reshape(S, h, d)
+        vfull = pvq_decode(q, allv.reshape(-1, allv.shape[-1]), "v").reshape(S, h, d)
+        kfull = kfull.permute(1, 0, 2)[None].to(self._dtype).to(self._device)
+        vfull = vfull.permute(1, 0, 2)[None].to(self._dtype).to(self._device)
+        return kfull, vfull
+
+    def get_seq_length(self, layer_idx: int = 0) -> int:
+        for li in sorted(self.target):
+            if li in self.k_codes:
+                return self.k_codes[li].shape[0]
+        return super().get_seq_length(layer_idx) if _HAVE_TF else 0
+
+    # -- live memory readout for the demo ------------------------------------
+    def memory_footprint(self) -> dict:
+        """Persistent bytes actually held on device, split by source."""
+        code_bytes = sum(t.numel() for t in self.k_codes.values()) \
+            + sum(t.numel() for t in self.v_codes.values())          # uint8 = 1 B
+        cb_bytes = sum(pvq_codebook_bytes(self.q[li]) for li in self.k_codes)
+        native_bytes = 0
+        if _HAVE_TF:
+            for kc in getattr(self, "key_cache", []):
+                if isinstance(kc, torch.Tensor):
+                    native_bytes += kc.numel() * kc.element_size()
+            for vc in getattr(self, "value_cache", []):
+                if isinstance(vc, torch.Tensor):
+                    native_bytes += vc.numel() * vc.element_size()
+        return {"compressed_indices_B": code_bytes,
+                "codebooks_B": cb_bytes,
+                "native_layers_B": native_bytes,
+                "total_B": code_bytes + cb_bytes + native_bytes}
+
+
+# ----------------------------------------------------------------------------
+# Honest byte-accounting (model-free) -- drives the memory demo & the harness
+# ----------------------------------------------------------------------------
+class LagunaGeom:
+    """Laguna-XS.2 cache geometry (from the config / proposal)."""
+    n_layers = 40
+    full_layers = 10          # full_attention layers that hold the growing cache
+    sliding_layers = 30
+    sliding_window = 512
+    n_kv_heads = 8
+    head_dim = 128
+
+
+def kv_cache_bytes(context_len: int, bits_per_elt_full: float,
+                   geom: LagunaGeom = LagunaGeom(),
+                   bits_per_elt_sliding: float = 16.0) -> dict:
+    """Total KV-cache bytes at a context length.
+
+    Only the full-attention layers carry the growing cache; sliding layers are
+    capped at ``sliding_window`` tokens. ``bits_per_elt_full`` is the rate the
+    quantizer reports for the compressed layers (use 16.0 for the fp16 baseline).
+    The K and V tensors are both counted.
+    """
+    elts_per_token = geom.n_kv_heads * geom.head_dim * 2          # K and V
+    full_tokens = context_len * geom.full_layers
+    slide_tokens = min(context_len, geom.sliding_window) * geom.sliding_layers
+    full_B = full_tokens * elts_per_token * bits_per_elt_full / 8
+    slide_B = slide_tokens * elts_per_token * bits_per_elt_sliding / 8
+    return {"full_B": full_B, "sliding_B": slide_B, "total_B": full_B + slide_B}
+
+
+def compression_vs_fp16(context_len: int, bits_per_elt_full: float,
+                        geom: LagunaGeom = LagunaGeom()) -> float:
+    base = kv_cache_bytes(context_len, 16.0, geom)["total_B"]
+    comp = kv_cache_bytes(context_len, bits_per_elt_full, geom)["total_B"]
+    return base / comp
+
+
+# ----------------------------------------------------------------------------
+# self-test: verify encode/decode == roundtrip and show the real byte win
+# ----------------------------------------------------------------------------
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    N, hd = 4096, 128
+    k = torch.randn(N, hd)
+    v = torch.randn(N, hd)
+
+    q = ProductVQKV(n_sub=32, n_codes=256, iters=10).fit(k, v)   # 2 bits/elt
+    idx = pvq_encode(q, k, "k")
+    k_hat = pvq_decode(q, idx, "k")
+    rt = q.roundtrip_k(k)
+
+    print(f"idx dtype/shape : {idx.dtype} {tuple(idx.shape)}")
+    print(f"encode->decode == roundtrip : "
+          f"{torch.allclose(k_hat, rt, atol=1e-4)}  "
+          f"(max abs diff {(k_hat - rt).abs().max():.2e})")
+
+    raw_B = k.numel() * 2                       # fp16
+    comp_B = idx.numel() * 1                     # uint8 indices
+    print(f"stored bytes  fp16={raw_B}  vq-2b={comp_B}  ratio={raw_B / comp_B:.1f}x")
+    print(f"reported bits/elt: {q.bits_per_element(hd):.3f}")
+
+    for L in (4096, 32768, 131072):
+        r = compression_vs_fp16(L, q.bits_per_element(hd))
+        gb = kv_cache_bytes(L, q.bits_per_element(hd))["total_B"] / 1e9
+        base = kv_cache_bytes(L, 16.0)["total_B"] / 1e9
+        print(f"context {L:>7}: fp16={base:6.2f} GB  vq-2b={gb:6.2f} GB  ({r:.1f}x)")

vqkv/metrics.py

@@ -0,0 +1,138 @@
+"""
+vqkv.metrics — attention-aware distortion metrics for KV-cache quantization.
+
+Reports key cosine, inner-product bias, and attention-output error (not cache
+MSE alone). attention_output() matches Laguna GQA geometry (8 KV / 48 Q heads).
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+def attention_output(q, k, v, n_q_heads, gate=None, scaling=None):
+    """Single forward of GQA attention for one layer.
+
+    q: (T, n_q_heads, head_dim)
+    k: (T, n_kv_heads, head_dim)
+    v: (T, n_kv_heads, head_dim)
+    Returns: (T, n_q_heads, head_dim) attention output (pre o_proj).
+    """
+    T, n_kv_heads, head_dim = k.shape
+    groups = n_q_heads // n_kv_heads
+    if scaling is None:
+        scaling = head_dim ** -0.5
+
+    # expand kv heads to q heads (GQA)
+    k_exp = k.repeat_interleave(groups, dim=1)  # (T, n_q_heads, d)
+    v_exp = v.repeat_interleave(groups, dim=1)
+
+    # (n_q_heads, T, d)
+    qh = q.transpose(0, 1)
+    kh = k_exp.transpose(0, 1)
+    vh = v_exp.transpose(0, 1)
+
+    scores = torch.matmul(qh, kh.transpose(-1, -2)) * scaling  # (H, T, T)
+    causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=q.device))
+    scores = scores.masked_fill(~causal, float("-inf"))
+    attn = F.softmax(scores, dim=-1)
+    out = torch.matmul(attn, vh)  # (H, T, d)
+    out = out.transpose(0, 1)     # (T, H, d)
+
+    if gate is not None:  # Laguna per-head softplus gate
+        out = out * F.softplus(gate).unsqueeze(-1)
+    return out, attn
+
+
+def cache_mse(x, x_hat):
+    return F.mse_loss(x_hat, x).item()
+
+
+def attn_output_error(out_ref, out_hat):
+    """Relative L2 error in attention output (magnitude + direction combined)."""
+    num = (out_hat - out_ref).norm()
+    den = out_ref.norm().clamp_min(1e-8)
+    return (num / den).item()
+
+
+def attn_output_cosine(out_ref, out_hat, per_head=True):
+    """Cosine similarity between reference and quantized attention outputs.
+
+    This isolates DIRECTION, which is what propagates through o_proj into the
+    residual stream. A quantizer can preserve direction while distorting
+    magnitude (or vice versa); relative-L2 conflates the two. On Laguna the
+    per-head softplus gate (g_proj) rescales each head AFTER attention, so a
+    magnitude error is partly re-absorbed by the gate -- which makes the
+    directional (cosine) error the more faithful quality signal here.
+
+    out_*: (T, H, d). Returns mean cosine over (valid token, head) pairs.
+    If per_head=False, flattens heads and measures the full per-token vector.
+    """
+    if per_head:
+        a = out_ref.reshape(-1, out_ref.shape[-1])
+        b = out_hat.reshape(-1, out_hat.shape[-1])
+    else:
+        a = out_ref.reshape(out_ref.shape[0], -1)
+        b = out_hat.reshape(out_hat.shape[0], -1)
+    cos = torch.nn.functional.cosine_similarity(a, b, dim=-1, eps=1e-8)
+    return cos.mean().item()
+
+
+def magnitude_direction_split(out_ref, out_hat):
+    """Decompose attention-output error into a directional part (1 - cosine)
+    and a magnitude part (relative norm error per vector). Reporting both shows
+    WHICH kind of error a quantizer makes -- the discriminating diagnostic when
+    comparing a rotation-based scalar method (TurboQuant) against product VQ.
+    """
+    a = out_ref.reshape(-1, out_ref.shape[-1])
+    b = out_hat.reshape(-1, out_hat.shape[-1])
+    cos = torch.nn.functional.cosine_similarity(a, b, dim=-1, eps=1e-8)
+    na, nb = a.norm(dim=-1).clamp_min(1e-8), b.norm(dim=-1)
+    mag_rel = ((nb - na).abs() / na).mean().item()
+    return {"dir_err (1-cos)": round((1 - cos.mean()).item(), 5),
+            "mag_rel_err": round(mag_rel, 5)}
+
+
+def attn_kl(attn_ref, attn_hat):
+    """Mean KL(attn_ref || attn_hat) over query positions and heads."""
+    a = attn_ref.clamp_min(1e-9)
+    b = attn_hat.clamp_min(1e-9)
+    kl = (a * (a.log() - b.log())).sum(dim=-1)  # (H, T)
+    # only count valid (non-masked) rows: those with > 0 mass already normalized
+    return kl.mean().item()
+
+
+def key_cosine(k_ref, k_hat):
+    """Mean cosine similarity between original and reconstructed KEY vectors.
+
+    The attention logit is q.k, which depends on the DIRECTION of k. Key cosine
+    is a cheap, query-free proxy for logit fidelity: if directions are preserved,
+    logits are preserved up to per-key magnitude. k_*: (T, H, d)."""
+    a = k_ref.reshape(-1, k_ref.shape[-1])
+    b = k_hat.reshape(-1, k_ref.shape[-1])
+    return torch.nn.functional.cosine_similarity(a, b, dim=-1, eps=1e-8).mean().item()
+
+
+def inner_product_distortion(q, k_ref, k_hat, n_pairs=4096):
+    """Relative error in the q.k inner products that drive attention logits.
+
+    This is the objective TurboQuant (ICLR 2026) argues for over cache MSE: it
+    shows MSE-optimal quantizers are BIASED for inner-product estimation. We
+    measure |<q,k_hat> - <q,k_ref>| / |<q,k_ref>| over sampled query-key pairs.
+
+    Also returns the signed mean (bias): a method can be unbiased (mean ~0) yet
+    high-variance, or biased -- the distinction TurboQuant's QJL stage targets.
+    q: (T, Hq, d); k_*: (T, Hk, d). Uses head 0 of each for a cheap estimate.
+    """
+    qd = q[:, 0, :]
+    kr = k_ref[:, 0, :]
+    kh = k_hat[:, 0, :]
+    T = qd.shape[0]
+    qi = torch.randint(0, T, (n_pairs,), device=qd.device)
+    ki = torch.randint(0, T, (n_pairs,), device=qd.device)
+    ip_ref = (qd[qi] * kr[ki]).sum(-1)
+    ip_hat = (qd[qi] * kh[ki]).sum(-1)
+    rel = ((ip_hat - ip_ref).abs() / ip_ref.abs().clamp_min(1e-6)).mean().item()
+    bias = (ip_hat - ip_ref).mean().item()
+    return {"ip_rel_err": round(rel, 5), "ip_bias": round(bias, 6)}

vqkv/quantizers.py

@@ -0,0 +1,618 @@
+"""
+vqkv.quantizers — KV-cache quantizers for AttnVQ.
+
+ScalarKV, KIVIScalarKV, ProductVQKV (LBG product VQ), RoPESplitVQKV,
+SignScalarKV, TernaryScalarKV. Each exposes fit() + roundtrip_k/v() and
+bits_per_element(). Distortion is evaluated via vqkv.metrics (attention-output
+error, not cache MSE alone).
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+
+import torch
+
+
+# ----------------------------------------------------------------------------
+# Utilities
+# ----------------------------------------------------------------------------
+def _affine_quantize(x: torch.Tensor, nbits: int, dim: int):
+    """Symmetric-range affine quantization along `dim`. Returns (q, scale, zero)."""
+    qmax = (1 << nbits) - 1
+    xmin = x.amin(dim=dim, keepdim=True)
+    xmax = x.amax(dim=dim, keepdim=True)
+    scale = (xmax - xmin).clamp_min(1e-8) / qmax
+    zero = torch.round(-xmin / scale)
+    q = torch.clamp(torch.round(x / scale) + zero, 0, qmax)
+    return q, scale, zero
+
+
+def _affine_dequantize(q, scale, zero):
+    return (q - zero) * scale
+
+
+# ----------------------------------------------------------------------------
+# Scalar baselines
+# ----------------------------------------------------------------------------
+@dataclass
+class ScalarKV:
+    """Per-token affine quantization of both K and V (the transformers default)."""
+    nbits: int = 4
+
+    def fit(self, k_calib, v_calib):  # scalar quant is tuning-free
+        return self
+
+    def roundtrip_k(self, k):
+        # k: (..., head_dim) -> quantize along head_dim (per-token, per-head)
+        q, s, z = _affine_quantize(k, self.nbits, dim=-1)
+        return _affine_dequantize(q, s, z)
+
+    def roundtrip_v(self, v):
+        q, s, z = _affine_quantize(v, self.nbits, dim=-1)
+        return _affine_dequantize(q, s, z)
+
+    def bits_per_element(self, head_dim):
+        # nbits per element + scale/zero (fp16 each) amortized over head_dim
+        return self.nbits + (2 * 16) / head_dim
+
+
+@dataclass
+class KIVIScalarKV:
+    """KIVI-style: keys quantized per-channel, values per-token.
+
+    Keys are quantized along the TOKEN axis (per channel); values along the
+    head_dim axis (per token). Requires a token-axis view, so fit/roundtrip
+    operate on a full (N_tokens, head_dim) block per head.
+    """
+    nbits: int = 4
+
+    def fit(self, k_calib, v_calib):
+        return self
+
+    def roundtrip_k(self, k):
+        # k: (N_tokens, head_dim). per-channel == quantize along token axis (dim=0)
+        q, s, z = _affine_quantize(k, self.nbits, dim=0)
+        return _affine_dequantize(q, s, z)
+
+    def roundtrip_v(self, v):
+        q, s, z = _affine_quantize(v, self.nbits, dim=-1)
+        return _affine_dequantize(q, s, z)
+
+    def bits_per_element(self, head_dim):
+        return self.nbits + (2 * 16) / head_dim
+
+
+# ----------------------------------------------------------------------------
+# LBG / k-means codebook (the 1980 algorithm, the engine of the project)
+# ----------------------------------------------------------------------------
+def lbg_codebook(data: torch.Tensor, n_codes: int, iters: int = 25,
+                 seed: int = 0) -> torch.Tensor:
+    """Linde-Buzo-Gray vector quantizer design (== Lloyd / k-means for MSE).
+
+    data:    (N, d) training vectors
+    returns: (n_codes, d) codebook
+    """
+    g = torch.Generator().manual_seed(seed)
+    N, d = data.shape
+    dev = data.device
+    idx = torch.randperm(N, generator=g)[:n_codes].to(dev)
+    cb = data[idx].clone()
+    ones = torch.ones(N, dtype=data.dtype, device=dev)
+    for _ in range(iters):
+        # Assignment: BLAS GEMM path avoids the O(N*K*d) intermediate that
+        # torch.cdist allocates for small d (e.g. d=8 for n_sub=16).
+        assign = _sq_l2(data, cb).argmin(dim=1)
+
+        # Update: vectorized scatter_add replaces the Python loop over n_codes.
+        # Was: 256 Python iterations; now: 3 torch ops.
+        new_cb = torch.zeros_like(cb)
+        counts = torch.zeros(n_codes, dtype=data.dtype, device=dev)
+        new_cb.scatter_add_(0, assign.unsqueeze(1).expand(-1, d), data)
+        counts.scatter_add_(0, assign, ones)
+        live = counts > 0
+        new_cb[live] /= counts[live].unsqueeze(1)
+        # Re-seed dead centroids (LBG splitting heuristic)
+        dead = (~live).nonzero(as_tuple=True)[0]
+        if dead.numel() > 0:
+            ri = torch.randint(0, N, (dead.numel(),), generator=g).to(dev)
+            new_cb[dead] = data[ri]
+
+        shift = (new_cb - cb).norm()
+        cb = new_cb
+        if shift < 1e-5:
+            break
+    return cb
+
+
+def lbg_codebook_batched(xb: torch.Tensor, n_codes: int, iters: int = 25,
+                          seed: int = 0) -> torch.Tensor:
+    """Fit n_sub independent LBG codebooks in one batched pass.
+
+    Mirrors the structure of ProductVQKV._roundtrip: replaces the Python loop
+    over sub-blocks with a single bmm-based assignment and a scatter_add-based
+    update, so n_sub sub-block fits become one operation at each Lloyd step.
+
+    xb:      (n_sub, N, sub_dim)  training vectors, pre-normalized if needed
+    returns: (n_sub, n_codes, sub_dim) codebooks, one per sub-block
+    """
+    g = torch.Generator().manual_seed(seed)
+    n_sub, N, sub_dim = xb.shape
+    dev = xb.device
+
+    # Initialization: random subset per sub-block (same RNG sequence as serial)
+    idx = torch.stack([torch.randperm(N, generator=g)[:n_codes]
+                       for _ in range(n_sub)]).to(dev)          # (n_sub, K)
+    cb = xb[torch.arange(n_sub, device=dev).unsqueeze(1), idx].clone()  # (n_sub, K, sub_dim)
+
+    ones = torch.ones(N, dtype=xb.dtype, device=dev)
+
+    for _ in range(iters):
+        # Assignment — one bmm replaces n_sub independent GEMMs
+        x_sq  = (xb * xb).sum(-1, keepdim=True)            # (n_sub, N, 1)
+        c_sq  = (cb * cb).sum(-1).unsqueeze(1)              # (n_sub, 1, K)
+        cross = torch.bmm(xb, cb.transpose(1, 2))           # (n_sub, N, K)
+        assign = (x_sq - 2 * cross + c_sq).argmin(dim=-1)  # (n_sub, N)
+
+        # Update — batched scatter_add
+        new_cb = torch.zeros_like(cb)
+        counts = torch.zeros(n_sub, n_codes, dtype=xb.dtype, device=dev)
+        assign_exp = assign.unsqueeze(-1).expand(-1, -1, sub_dim)
+        new_cb.scatter_add_(1, assign_exp, xb)
+        counts.scatter_add_(1, assign, ones.unsqueeze(0).expand(n_sub, -1))
+
+        live = counts > 0
+        new_cb[live] /= counts[live].unsqueeze(-1)
+
+        # Re-seed dead centroids per sub-block (rare; loop is fine)
+        dead_any = ~live
+        if dead_any.any():
+            for s in range(n_sub):
+                dead = dead_any[s].nonzero(as_tuple=True)[0]
+                if dead.numel():
+                    ri = torch.randint(0, N, (dead.numel(),), generator=g).to(dev)
+                    new_cb[s, dead] = xb[s, ri]
+
+        shift = (new_cb - cb).norm()
+        cb = new_cb
+        if shift < 1e-5:
+            break
+
+    return cb
+
+
+def _sq_l2(x: torch.Tensor, cb: torch.Tensor) -> torch.Tensor:
+    """Squared L2 distance matrix (N, K) via BLAS GEMM.
+
+    torch.cdist for small d (e.g. d=8 for n_sub=16) falls back to a naive
+    expand-and-subtract path that allocates an (N, K, d) intermediate.
+    For N=131072, K=256, d=8 that is 1.07 GB per call, causing massive
+    CPU allocation pressure and 100x slowdowns versus the BLAS path.
+
+    ||x-c||^2 = ||x||^2 - 2*(x @ c.T) + ||c||^2 uses only (N,K) memory.
+    """
+    return ((x * x).sum(1, keepdim=True)
+            + (cb * cb).sum(1)
+            - 2 * (x @ cb.T))
+
+
+def vq_encode(x: torch.Tensor, cb: torch.Tensor) -> torch.Tensor:
+    """Nearest-codeword indices. x:(N,d) cb:(K,d) -> (N,) long."""
+    return _sq_l2(x, cb).argmin(dim=1)
+
+
+# ----------------------------------------------------------------------------
+# Product VQ
+# ----------------------------------------------------------------------------
+@dataclass
+class ProductVQKV:
+    """Product vector quantization of head-dim sub-blocks.
+
+    Each head's `head_dim` vector is split into `n_sub` contiguous sub-vectors
+    of length `sub_dim = head_dim / n_sub`; each sub-vector is quantized
+    against its own codebook of size `n_codes`.
+
+    If `normalize=True`, each sub-vector is standardized by its per-dimension
+    calibration mean/std before codebook design and restored after decode.
+    This decorrelates the scale variation RoPE introduces in the rotated half
+    of the key, and is the mechanism that lets RoPE-split specialize.
+
+    Rate (bits/element) = n_sub * log2(n_codes) / head_dim.
+    e.g. head_dim=128, n_sub=8, n_codes=256  -> 8*8/128 = 0.5 bits/element.
+    """
+    n_sub: int = 8
+    n_codes: int = 256
+    iters: int = 25
+    normalize: bool = False
+    k_codebooks: list = field(default_factory=list)
+    v_codebooks: list = field(default_factory=list)
+    _k_stats: list = field(default_factory=list)
+    _v_stats: list = field(default_factory=list)
+    _k_stacked: tuple = None   # lazily-built (cb, mu, sd) batched tensors
+    _v_stacked: tuple = None
+
+    def _split(self, x):
+        # x: (N, head_dim) -> list of (N, sub_dim)
+        return list(torch.chunk(x, self.n_sub, dim=-1))
+
+    def _fit_one(self, x):
+        N, head_dim = x.shape
+        sub_dim = head_dim // self.n_sub
+        # (n_sub, N, sub_dim) -- same layout _roundtrip uses, so batching mirrors inference
+        xb = x.reshape(N, self.n_sub, sub_dim).permute(1, 0, 2).contiguous()
+
+        if self.normalize:
+            mu = xb.mean(dim=1, keepdim=True)          # (n_sub, 1, sub_dim)
+            sd = xb.std(dim=1, keepdim=True).clamp_min(1e-6)
+            xb = (xb - mu) / sd
+            # unstack into (1, sub_dim) tuples so _stack / to() are unchanged
+            stats = [(mu[s], sd[s]) for s in range(self.n_sub)]
+        else:
+            stats = [None] * self.n_sub
+
+        cb_batched = lbg_codebook_batched(xb, self.n_codes, self.iters)
+        cbs = list(cb_batched.unbind(dim=0))            # n_sub x (K, sub_dim)
+        return cbs, stats
+
+    def fit(self, k_calib, v_calib):
+        self.k_codebooks, self._k_stats = self._fit_one(k_calib)
+        self.v_codebooks, self._v_stats = self._fit_one(v_calib)
+        return self
+
+    def _stack(self, codebooks, stats):
+        """Lazily stack per-sub-block codebooks/stats into batched tensors so the
+        whole product-VQ encode is a few large ops instead of n_sub small ones.
+
+        Returns:
+          cb_stacked: (n_sub, K, sub_dim)
+          mu_stacked: (1, n_sub, sub_dim) or None   (None => no normalization)
+          sd_stacked: (1, n_sub, sub_dim) or None
+        Requires uniform sub_dim and n_codes across sub-blocks, which ProductVQ
+        guarantees (torch.chunk into equal pieces, single n_codes).
+        """
+        cb_stacked = torch.stack(codebooks, dim=0)  # (n_sub, K, sub_dim)
+        if any(st is not None for st in stats):
+            mu = torch.stack([st[0].reshape(-1) for st in stats], dim=0)  # (n_sub, sub_dim)
+            sd = torch.stack([st[1].reshape(-1) for st in stats], dim=0)
+            mu_stacked = mu.unsqueeze(0)  # (1, n_sub, sub_dim)
+            sd_stacked = sd.unsqueeze(0)
+        else:
+            mu_stacked = sd_stacked = None
+        return cb_stacked, mu_stacked, sd_stacked
+
+    def _ensure_stacked(self):
+        if getattr(self, "_k_stacked", None) is None:
+            self._k_stacked = self._stack(self.k_codebooks, self._k_stats)
+        if getattr(self, "_v_stacked", None) is None:
+            self._v_stacked = self._stack(self.v_codebooks, self._v_stats)
+
+    def _roundtrip(self, x, stacked):
+        """Batched product-VQ round-trip.
+
+        x: (N, head_dim). Splits into (N, n_sub, sub_dim), then does ONE batched
+        squared-L2 (n_sub, N, K), one argmin, one gather -- replacing the Python
+        loop over sub-blocks and its 3*n_sub small kernels.
+
+        Chunks over N so peak memory (the (n_sub, chunk, K) distance tensor)
+        stays bounded: at N=131072, n_sub=16, K=256 the un-chunked tensor is
+        ~2 GB in fp32. The batched path is a GPU optimization -- on CPU it is
+        roughly on par with the per-sub-block loop (no launch overhead to hide).
+        """
+        cb, mu, sd = stacked                      # cb: (n_sub, K, sub_dim)
+        n_sub, K, sub_dim = cb.shape
+        N = x.shape[0]
+        c_sq = (cb * cb).sum(-1).unsqueeze(1)           # (n_sub, 1, K)
+        if mu is not None:
+            mu_b = mu.permute(1, 0, 2)                  # (n_sub, 1, sub_dim)
+            sd_b = sd.permute(1, 0, 2)
+
+        # cap the distance tensor at ~256 MB fp32: n_sub * chunk * K * 4 bytes
+        chunk = max(1, (256 * 1024 * 1024) // (n_sub * K * 4))
+        out_chunks = []
+        for start in range(0, N, chunk):
+            xc = x[start:start + chunk]                 # (c, head_dim)
+            c = xc.shape[0]
+            xb = xc.reshape(c, n_sub, sub_dim).permute(1, 0, 2).contiguous()
+            if mu is not None:
+                xb = (xb - mu_b) / sd_b
+            x_sq = (xb * xb).sum(-1, keepdim=True)      # (n_sub, c, 1)
+            cross = torch.bmm(xb, cb.transpose(1, 2))   # (n_sub, c, K)
+            d2 = x_sq - 2 * cross + c_sq
+            idx = d2.argmin(dim=-1)                     # (n_sub, c)
+            idx_exp = idx.unsqueeze(-1).expand(-1, -1, sub_dim)
+            rec = torch.gather(cb, 1, idx_exp)          # (n_sub, c, sub_dim)
+            if mu is not None:
+                rec = rec * sd_b + mu_b
+            out_chunks.append(rec.permute(1, 0, 2).reshape(c, n_sub * sub_dim))
+        return torch.cat(out_chunks, dim=0)
+
+    def roundtrip_k(self, k):
+        self._ensure_stacked()
+        return self._roundtrip(k, self._k_stacked)
+
+    def roundtrip_v(self, v):
+        self._ensure_stacked()
+        return self._roundtrip(v, self._v_stacked)
+
+    def to(self, device):
+        self.k_codebooks = [cb.to(device) for cb in self.k_codebooks]
+        self.v_codebooks = [cb.to(device) for cb in self.v_codebooks]
+        self._k_stats = [(st[0].to(device), st[1].to(device)) if st is not None else None
+                         for st in self._k_stats]
+        self._v_stats = [(st[0].to(device), st[1].to(device)) if st is not None else None
+                         for st in self._v_stats]
+        # invalidate cached stacks; they rebuild on next roundtrip on the new device
+        self._k_stacked = None
+        self._v_stacked = None
+        return self
+
+    def bits_per_element(self, head_dim):
+        sub_dim = head_dim / self.n_sub
+        return math.log2(self.n_codes) / sub_dim
+
+
+@dataclass
+class TurboQuantKV:
+    """Data-oblivious rotation + per-coordinate scalar quantization, in the
+    spirit of TurboQuant (Zandieh et al., ICLR 2026).
+
+    Pipeline: random rotation Pi (so coordinates become near-iid in high dim),
+    then a per-coordinate scalar codebook. We use a uniform Lloyd-Max-style
+    codebook on the rotated coordinates as a faithful stand-in for their
+    precomputed Beta-optimal codebook (the exact codebook is an implementation
+    detail; the rotation is the load-bearing idea). We store the per-vector L2
+    norm in fp16 and rescale on dequant, as the paper specifies.
+
+    This is included so we can compare a rotation-based SCALAR method against
+    product VQ on the attention-output COSINE metric -- the comparison the
+    TurboQuant paper does not make (it optimizes cache-vector MSE / inner prod).
+    """
+    nbits: int = 4
+    seed: int = 0
+    _rot: torch.Tensor = None
+    _levels: torch.Tensor = None
+
+    def _make_rotation(self, d):
+        g = torch.Generator().manual_seed(self.seed)
+        a = torch.randn(d, d, generator=g)
+        q, _ = torch.linalg.qr(a)
+        return q
+
+    def fit(self, k_calib, v_calib):
+        d = k_calib.shape[-1]
+        self._rot = self._make_rotation(d)
+        # Fit levels on UNIT-NORMALIZED, rotated coordinates -- the same domain
+        # the round-trip quantizes in. (Earlier bug: levels were fit on
+        # un-normalized rotated data, so the per-vector-normalized values fell
+        # outside the level range and collapsed to one bin.)
+        kn = k_calib / k_calib.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+        rk = kn @ self._rot
+        lo = torch.quantile(rk.flatten()[:200000], 0.001)
+        hi = torch.quantile(rk.flatten()[:200000], 0.999)
+        self._levels = torch.linspace(lo.item(), hi.item(), (1 << self.nbits))
+        return self
+
+    def _roundtrip(self, x):
+        norms = x.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+        xn = x / norms
+        y = xn @ self._rot                              # rotate
+        idx = torch.bucketize(y, self._levels)
+        idx = idx.clamp(0, self._levels.numel() - 1)
+        y_hat = self._levels[idx]                       # per-coord scalar quant
+        x_hat = y_hat @ self._rot.T                     # un-rotate
+        return x_hat * norms                            # rescale by stored norm
+
+    def roundtrip_k(self, k):
+        return self._roundtrip(k)
+
+    def roundtrip_v(self, v):
+        return self._roundtrip(v)
+
+    def to(self, device):
+        self._rot    = self._rot.to(device)
+        self._levels = self._levels.to(device)
+        return self
+
+    def bits_per_element(self, head_dim):
+        # nbits/coord + fp16 norm amortized over head_dim
+        return self.nbits + 16 / head_dim
+
+
+@dataclass
+class RoPESplitVQKV:
+    """ProductVQ with separate codebooks for the RoPE'd vs pass-through halves
+    of each KEY head.
+
+    Laguna full-attention layers use partial_rotary_factor=0.5: the first half
+    of head_dim is rotated (position-dependent, broad distribution), the second
+    half is identity (position-independent). A single codebook must straddle
+    two regimes; splitting lets each codebook specialize.
+
+    Values receive no RoPE, so V uses a plain ProductVQ.
+    """
+    n_sub_half: int = 4          # sub-vectors PER HALF for keys
+    n_codes: int = 256
+    iters: int = 25
+    rotary_fraction: float = 0.5
+    _rope_vq: ProductVQKV = None
+    _pass_vq: ProductVQKV = None
+    _v_vq: ProductVQKV = None
+
+    def fit(self, k_calib, v_calib):
+        d = k_calib.shape[-1]
+        cut = int(d * self.rotary_fraction)
+        k_rope, k_pass = k_calib[..., :cut], k_calib[..., cut:]
+        self._cut = cut
+        self._rope_vq = ProductVQKV(self.n_sub_half, self.n_codes, self.iters,
+                                     normalize=True).fit(k_rope, k_rope)
+        self._pass_vq = ProductVQKV(self.n_sub_half, self.n_codes, self.iters,
+                                    normalize=False).fit(k_pass, k_pass)
+        self._v_vq = ProductVQKV(2 * self.n_sub_half, self.n_codes, self.iters,
+                                 normalize=True).fit(v_calib, v_calib)
+        return self
+
+    def roundtrip_k(self, k):
+        kr = self._rope_vq.roundtrip_k(k[..., :self._cut])
+        kp = self._pass_vq.roundtrip_k(k[..., self._cut:])
+        return torch.cat([kr, kp], dim=-1)
+
+    def roundtrip_v(self, v):
+        return self._v_vq.roundtrip_v(v)
+
+    def to(self, device):
+        self._rope_vq.to(device)
+        self._pass_vq.to(device)
+        self._v_vq.to(device)
+        return self
+
+    def bits_per_element(self, head_dim):
+        # both halves use n_sub_half codes over head_dim/2 elements
+        half = head_dim / 2
+        sub_dim = half / self.n_sub_half
+        return math.log2(self.n_codes) / sub_dim
+
+
+# ----------------------------------------------------------------------------
+# Data-oblivious baseline: a simplified TurboQuant-style quantizer
+# ----------------------------------------------------------------------------
+@dataclass
+class RandomRotationScalarKV:
+    """Simplified, data-OBLIVIOUS rotation-then-scalar quantizer in the spirit
+    of TurboQuant / PolarQuant (Zandieh et al., ICLR 2026).
+
+    NOT the full method: TurboQuant adds PolarQuant's normalization-free polar
+    transform and a 1-bit QJL residual correction for UNBIASED inner-product
+    estimation. This stand-in captures only the core data-oblivious idea --
+    apply a fixed random orthogonal rotation so coordinates concentrate (a
+    Beta/Gaussian-like distribution), then scalar-quantize each coordinate with
+    a fixed range. It exists so the harness can run the central scientific
+    comparison of this project:
+
+        data-OBLIVIOUS rotation+scalar   vs.   data-DEPENDENT product VQ
+
+    on real Laguna cache statistics. If you want the real thing, drop in the
+    unofficial impl (github.com/0xSero/turboquant or hackimov/turboquant-kv)
+    behind this same fit/roundtrip interface. The OpenReview discussion of the
+    paper is contested precisely on the oblivious-vs-data-dependent claim, so a
+    clean head-to-head on a NEW model is a genuine contribution either way.
+    """
+    nbits: int = 3
+    seed: int = 0
+    _R: torch.Tensor = None        # rotation
+    _Rk_range: tuple = None
+    _Rv_range: tuple = None
+
+    def _rotation(self, d):
+        g = torch.Generator().manual_seed(self.seed)
+        a = torch.randn(d, d, generator=g)
+        q, _ = torch.linalg.qr(a)   # random orthogonal matrix
+        return q
+
+    def fit(self, k_calib, v_calib):
+        d = k_calib.shape[-1]
+        self._R = self._rotation(d)
+        # fixed (data-oblivious-ish) ranges from a robust quantile of rotated calib
+        rk = k_calib @ self._R
+        rv = v_calib @ self._R
+        self._Rk_range = (rk.quantile(0.001), rk.quantile(0.999))
+        self._Rv_range = (rv.quantile(0.001), rv.quantile(0.999))
+        return self
+
+    def _rt(self, x, rng):
+        xr = x @ self._R
+        lo, hi = rng
+        qmax = (1 << self.nbits) - 1
+        scale = (hi - lo).clamp_min(1e-8) / qmax
+        q = torch.clamp(torch.round((xr - lo) / scale), 0, qmax)
+        xr_hat = q * scale + lo
+        return xr_hat @ self._R.T   # rotation is orthogonal, inverse == transpose
+
+    def roundtrip_k(self, k):
+        return self._rt(k, self._Rk_range)
+
+    def roundtrip_v(self, v):
+        return self._rt(v, self._Rv_range)
+
+    def bits_per_element(self, head_dim):
+        # rotation is a fixed matrix (no per-token overhead); ranges are global.
+        return float(self.nbits)
+
+
+# ----------------------------------------------------------------------------
+# 1-bit baselines (the floor of the scalar family; rate-neighbors to VQ)
+# ----------------------------------------------------------------------------
+@dataclass
+class SignScalarKV:
+    """Symmetric 1-bit (sign) quantization with a per-group scale.
+
+        x_hat = scale * sign(x),   scale = mean(|x|) over the group.
+
+    This is the honest 1-bit floor of the KIVI/quanto scalar family: unlike the
+    affine `ScalarKV(nbits=1)`, it stores NO zero-point (cache K/V are ~zero-mean
+    after norm), so it is both cheaper and better-centered. Group axis matches
+    KIVI conventions: keys per-channel (token axis), values per-token.
+    `per_channel_key` toggles the key axis.
+    """
+    per_channel_key: bool = True
+    group_dim_k: int = 0     # 0 = per-channel (token axis); -1 = per-token
+    bits: float = 1.0
+
+    def fit(self, k_calib, v_calib):
+        return self
+
+    def _sign_q(self, x, dim):
+        scale = x.abs().mean(dim=dim, keepdim=True).clamp_min(1e-8)
+        return torch.sign(x) * scale
+
+    def roundtrip_k(self, k):
+        dim = 0 if self.per_channel_key else -1
+        return self._sign_q(k, dim)
+
+    def roundtrip_v(self, v):
+        return self._sign_q(v, dim=-1)
+
+    def bits_per_element(self, head_dim):
+        # 1 bit/element + one fp16 scale per group amortized over the group.
+        # per-token value group = head_dim elements; per-channel key group is the
+        # token axis (large), so its scale overhead is negligible. Report ~1 +
+        # 16/head_dim as a conservative upper bound for the per-token case.
+        return 1.0 + 16.0 / head_dim
+
+
+@dataclass
+class TernaryScalarKV:
+    """1.58-bit ternary quantization {-1, 0, +1} with a per-group scale, in the
+    style of BitNet b1.58 -- the "just go (almost) 1-bit" school that the
+    Bonsai/PrismML line popularized. Zeros are assigned by a threshold at a
+    fraction of the mean-abs, letting small-magnitude coordinates drop out.
+
+        thr = alpha * mean(|x|);  x_hat = scale * {sign(x) if |x|>thr else 0}
+
+    Rate ~ log2(3) ~ 1.58 bits/element. Included so the table spans the full
+    aggressive regime and so a reviewer who knows BitNet sees the comparison.
+    """
+    alpha: float = 0.7
+    per_channel_key: bool = True
+
+    def fit(self, k_calib, v_calib):
+        return self
+
+    def _tern_q(self, x, dim):
+        m = x.abs().mean(dim=dim, keepdim=True).clamp_min(1e-8)
+        thr = self.alpha * m
+        mask = (x.abs() > thr).to(x.dtype)
+        scale = (x.abs() * mask).sum(dim=dim, keepdim=True) / \
+                mask.sum(dim=dim, keepdim=True).clamp_min(1.0)
+        return torch.sign(x) * mask * scale
+
+    def roundtrip_k(self, k):
+        dim = 0 if self.per_channel_key else -1
+        return self._tern_q(k, dim)
+
+    def roundtrip_v(self, v):
+        return self._tern_q(v, dim=-1)
+
+    def bits_per_element(self, head_dim):
+        import math as _m
+        return _m.log2(3) + 16.0 / head_dim