Skip to content

TurboQuant KV Compression

TurboQuant KV compression — a cathedral memory vault where token-cards collapse into hexagonal cells, each labeled with a fractional inscription, Apple-Silicon silhouettes embossed on the back wall.

The KV cache is where long conversations go to eat memory. Every token a model generates leaves behind a pair of tensors — the keys and values from every attention head. On a Mac Mini that also runs sanctum-mlx-devstral, the Home Assistant gateway, Sanctum TTS, and a dozen other services, that cache decides whether you get to have a conversation at all.

The production model, Qwen3.6-35B-A3B-4bit, is mercifully hybrid: only 10 of its 40 layers run full attention (the other 30 are linear-attention and hold no standard KV cache), but each of those 10 hoards keys and values at 256-dim heads. TurboQuant is how we shrink that. The idea isn’t ours — Google’s ICLR 2026 paper (arXiv 2504.19874), layered on KVQuant and QJL. Our job: implement it cleanly in sanctum-mlx, wire it into Apple’s MLX runtime, and prove quality holds when the bits come down.

incoming K,V tensor [1, H_kv, L, D]
↓ ┌──── keys ────┐
↓ ▼ ▼
concat-grow plain on-device Array (bf16, bit-exact, no compression)
↓ │
↓ ┌── values ─────┘
↓ ▼
group-affine quantize on device → store (indices u8, scale f16, zero f16)
↓
attention: sdpa_dequant_v fused Metal kernel reads compressed V state
and dequantizes it inline in registers, no full V tensor
ever materializes

Both halves live in services/sanctum-mlx/src/turboquant/ and hand off to a CompressedKVCache implementing mlx_lm::cache::KeyValueCache — a drop-in for stock ConcatKeyValueCache, routed via the KeyValueCache::fused_attention trait method.

The production com.sanctum.mlx LaunchAgent passes --turboquant against Qwen3.6-35B-A3B-4bit, and the server log emits a turboquant ... ratio=1.32x line for each of the 10 full-attention layers (3, 7, 11, … 39) on every request.

The paper’s recipe — historical context

Section titled “The paper’s recipe — historical context”

Slice 1 followed the TurboQuant paper literally:

keys: normalize → Hadamard rotate → Lloyd-Max quantize → QJL 1-bit sign correction → pack
values: group-affine quantize (KVQuant style) → store scale + zero per group
on every read: dequant all stored tokens → Array → attention

This is the algorithm the protocol below validated; the bit-budget math behind the “~3.5 bits/channel quality-neutral” claim is real. The receipts (Slice 1 PPL sweep, Slice 1.5 negative result, Slice 4a kernel correctness, Slice 4a-final pivot) are preserved in agent memory.

The first end-to-end run on a Qwen3.5-27B-4bit model with default settings (3-bit keys + QJL, 4-bit values, group_size=128) delivered this:

ppl (fp16) = 9.3055
ppl (turboquant) = 28.6451
Δppl = +19.34 absolute, +207.83% relative

Catastrophic. The plumbing works — no crashes, the bit stream roundtrips — but the model has forgotten what it was looking at. This is where “tune a few knobs and hope” fails. We need a protocol.

Five principles, in priority order.

Do not touch bit widths until each stage is proven correct in isolation. Loss could hide in the Array↔f32 bridge, Hadamard rotation, Lloyd-Max, QJL, or the value group-affine path — and perplexity collapse can’t tell them apart.

Run Config What it isolates Expected
A1 8-bit keys + 8-bit values, QJL off Bridge + minimal-loss quant ≈ fp16
A4 8-bit keys + 8-bit values, QJL on QJL math correctness at high precision ≈ fp16
A5 8-bit keys, 4-bit values, QJL off Value stage only measures value loss
A6 3-bit keys + QJL, 8-bit values Key stage only measures key loss

If A1 fails, the bug is in the bridge or rotation — bisect, don’t sweep. If A5 is clean but A6 broken, it’s in the key pipeline (Lloyd-Max, QJL, rotation). And so on.

Once each stage is proven, search the Pareto frontier. Don’t grid-search 120 cells — use a staircase:

  1. Fix value_bits=8, sweep key_bits ∈ {3, 4, 5, 6, 8} → find minimum that meets budget
  2. Fix key_bits at step-1 winner, sweep value_bits ∈ {2, 4, 6, 8}
  3. Fix both, check group_size ∈ {32, 64, 128} sensitivity
  4. At the winning (key_bits, value_bits), compare use_qjl: on vs off — does QJL actually help?

About 20 runs total; ~10 seconds each on an M4 Max once the model is cached.

Every run appends one line to services/sanctum-mlx/bench/turboquant_sweep.jsonl:

{
"ts": "2026-04-18T11:40:00Z",
"host": "mbp",
"config": {"key_bits": 4, "value_bits": 6, "group_size": 128, "use_qjl": false, "head_dim": 256},
"eval": {"text_sha": "3f4a…", "n_tokens": 54},
"ppl_fp16": 9.3055,
"ppl_tq": 10.12,
"abs_delta": 0.8145,
"rel_delta": 0.0876,
"runtime_sec": {"fp16": 4.2, "tq": 6.8},
"verdict": "warn"
}

The text SHA is the comparability key — different text means a different row, not a different config. The verdict classifier:

Verdict Budget
pass abs_delta ≤ 0.3 AND rel_delta ≤ 5%
warn abs_delta ≤ 1.0
fail beyond warn
invalid non-finite — means a bug

Never change two variables between runs. The harness reads the five SANCTUM_TQ_* knobs (see reference below) into a OnceLock before the model loads — no recompile between runs. The shell driver bench/tune.sh orchestrates the staircase; individual stages (ablation, key, value, group, qjl) run in isolation.

Don’t over-measure. The protocol has two exit conditions:

  • Confirmed bug (any ablation fails): stop, fix, restart — sweeping over a bug produces garbage data.
  • Confirmed winner (a config meets both budgets): stop; the knee is located. Tighten later if needed.
Terminal window
cd ~/Projects/sanctum-rs
SANCTUM_MLX_TEST_MODEL=/path/to/Qwen3.5-27B-4bit \
SANCTUM_TQ_KEY_BITS=4 \
SANCTUM_TQ_VALUE_BITS=6 \
SANCTUM_TQ_USE_QJL=false \
SANCTUM_TQ_JSONL=services/sanctum-mlx/bench/turboquant_sweep.jsonl \
SANCTUM_TQ_SOFT=1 \
cargo test -p sanctum-mlx --test turboquant_ppl --release \
-- --ignored --nocapture

SANCTUM_TQ_SOFT=1 disables the assertion failure on regression — during sweeps, bad numbers are data. Omit it when gating.

Knob Env Range Meaning
key_bits SANCTUM_TQ_KEY_BITS 3–8 Codebook precision per key dim. Rotation runs regardless.
value_bits SANCTUM_TQ_VALUE_BITS 2–8 Affine quant precision per value dim, per group.
value_group_size SANCTUM_TQ_GROUP_SIZE 32, 64, 128 Granularity of per-group scale/zero for values.
use_qjl SANCTUM_TQ_USE_QJL true / false Whether to add the QJL 1-bit sign correction after Lloyd-Max. Only meaningful at key_bits ≤ 4.
head_dim SANCTUM_TQ_HEAD_DIM model-specific Must match the model. Qwen3.5 = 256, Qwen2.5 = 128.

Metal memory ceilings (orthogonal to TurboQuant, but relevant to the why):

Flag Env Effect
--metal-cache-limit-mb SANCTUM_MLX_CACHE_LIMIT_MB Cap MLX’s reusable buffer cache. Stops sanctum-mlx from fighting sanctum-mlx-devstral for the Metal heap.
--metal-memory-limit-mb SANCTUM_MLX_MEMORY_LIMIT_MB Hard cap total MLX device memory.
--metal-wired-limit-mb SANCTUM_MLX_WIRED_LIMIT_MB Cap non-swappable memory — keep hot tensors resident without starving the rest of the system.

The gate is abs_delta ≤ 0.3 AND rel_delta ≤ 5% on a short eval — permissive next to the paper’s sub-0.1 Δppl by design: our implementation omits some refinements (learned rotations, full TurboQuant value pipeline) to ship Slice 1. Passing it was only good enough to ship behind a flag. Promotion to default required a higher bar:

  • Clean ablation pass across all stages
  • Wider eval — ≥1024 tokens of real prose, not the calibration string
  • Sample-quality check: 200-token continuation from the same prefix, both caches, compared
  • No regression across prompt families (code, prose, reasoning)

The Bold config cleared it in Slice 1b — the evidence follows, starting with the day the protocol paid for itself.

The first end-to-end run on Qwen3.5-27B-4bit with scaffold defaults (3-bit keys + QJL, 4-bit values) produced a catastrophe:

ppl (fp16) = 9.3055
ppl (turboquant) = 28.6451 Δppl = +19.34 (+208%)

The instinct, staring at +19.34, is to turn knobs. The protocol said ablate first — and the ablation revealed what knob-turning never would: all four bit configurations gave the same +20 PPL collapse, 3-bit + QJL and 8-bit/8-bit with QJL off alike. That’s impossible if the quant math were the issue; the bug was upstream.

We added BypassMode::{IdentityPassthrough, BridgeOnly} — short-circuits that skip quantization (Slice-1 branch only; the pivot later deleted the CPU bridge they probed):

Run Mode Δppl
A0 IdentityPassthrough (raw concat) 0.0000 PASS
A0.5 BridgeOnly (f32 round-trip, no quant) +20.62 FAIL

Bisected in two runs: the f32 Array↔CPU bridge, not quantization. MLX Arrays from transpose_axes([0, 2, 1, 3]) are strided views, and as_slice::<T>() reads the raw buffer, ignoring logical layout. Without flatten() to force materialization, a logical [B, H, L, D] view yielded pre-transpose [B, L, H, D] bytes — heads and tokens permuted before any quantization ran. Fix: three lines in turboquant/cache.rs, flatten before as_slice.

Post-fix sweep — the actual Pareto frontier

Section titled “Post-fix sweep — the actual Pareto frontier”

Full staircase on the same 27B model, 54-token eval, protocol compliant — one knob at a time, pre-fix rows archived as invalid.

Run Config Δppl
A1 8-bit / 8-bit, no QJL 0.023
A4 8-bit / 8-bit, QJL on 0.014
A5 8-bit keys / 4-bit values 0.024
A6 3-bit + QJL / 8-bit values 0.137
3-bit Δ=0.254 4-bit Δ=0.069 5-bit Δ=0.090
6-bit Δ=0.031 8-bit Δ=0.023

Keys quantize gracefully. All pass.

2-bit Δ=1.22 FAIL breaks 4-bit Δ=0.12 PASS
6-bit Δ=0.05 8-bit Δ=0.07

Hard floor: 4-bit values. The value pathway (group-affine, no rotation) is less tolerant than keys.

key_bits values QJL off QJL on
4 6 0.050 0.184 QJL hurts
3 8 0.254 0.137 QJL helps

QJL costs one bit from the codebook. At 4-bit keys that trade is net-negative; at 3-bit keys the sign-correction wins. Rule: enable QJL only when key_bits ≤ 3.

At (k=4, v=6) the group sizes 32, 64, 128 all land within 0.04 absolute — noise. Use 128 for minimum metadata overhead.

Config Δppl Rel Est. Compression Risk
Safe: k=4, v=4, g=128, QJL=off 0.123 1.32% ~3.2× solid pass
Bold: k=3, v=4, g=128, QJL=on 0.338 3.63% ~4.0× 0.04 over abs budget — noise-level margin

On a 54-token eval we can’t responsibly discriminate between these at the 0.3-boundary. Principle 5: do not flip a default without a confirmed winner. Wider eval (wikitext-2, ≥1024 tokens) is the gate.

Ran both candidates on a 1256-token corpus of real sanctum-docs architecture prose. The baseline shifted up — real prose is harder than the 54-token calibration string — but it’s fairer.

Candidate Δppl abs Rel Est. Compression
Safe: k=4, v=4, g=128, QJL=off 0.157 1.18% ~3.2×
Bold: k=3, v=4, g=128, QJL=on 0.196 1.47% ~4.0×

Both decisively inside the 0.3 abs / 5% rel budgets. The wider sample firmed up the short eval’s hint: Bold’s borderline 0.34 tightened to 0.20, Safe’s 0.12 drifted to 0.16.

TurboQuantConfig::default() on main is still key_bits=3, value_bits=4, use_qjl=true — the scaffold had the right intent; the protocol made it evidence. Slice 4 then made the fused sdpa_dequant_v kernel the production path, retiring the O(T²) CPU round-trip that had justified the opt-in caution. The flag flipped; only value_bits and value_group_size steer storage now.

What stays narrow is model coverage: only the 35B and the original Qwen3.5-27B are validated. Qwen2.5-Coder, Gemma4-31B, and the rest still need their own protocol run — see Slice 1c for the wall we hit trying.

Slice 1c — cross-architecture de-risk (ecosystem wall + a finding)

Section titled “Slice 1c — cross-architecture de-risk (ecosystem wall + a finding)”

We tried to validate the Bold default on a non-Qwen3.5 architecture before committing to Gemma-4 Rust work. It hit an ecosystem wall — and produced a finding anyway.

Python mlx-lm 0.29.1 (latest on Python 3.9) supports qwen2, qwen3, gemma3, llama4, and more — but not qwen3_5 (our validated model) and not gemma4 (the fine-tune target). Our Rust mlx-lm fork has the inverse: qwen3_5 and qwen3, but not qwen2 or any gemma.

Python ∩ Rust = { }

No model type loads in both stacks. Cross-validation means either porting TurboQuant to Python as a custom KVCache subclass (~2–3 hr) or adding an architecture to the Rust fork (~4–6 hr Qwen2, ~8–16 hr Gemma-4). Neither is a quick de-risk.

We ran Python mlx-lm’s built-in QuantizedKVCache (plain affine per-group quant, not TurboQuant’s rotation path) on Qwen2.5-Coder-7B (head_dim=128, 7:1 GQA), same 1238-token eval:

KV bits PPL Δppl abs Verdict
fp16 baseline 25.99 — —
8 25.94 0.05 PASS
4 339,210 339,184 FAIL — catastrophic
3 712,519 712,493 FAIL
2 648,756 648,730 FAIL

A PPL of 339,210 on a ~152k-vocab model is worse than uniform random — not a quality gradient, a cliff. Stock 4-bit affine quant isn’t losing information, it’s poisoning state.

That sets the forward slices:

Slice What Gate
1d Port TurboQuant to Python mlx-lm as a custom KVCache (~3-hr session; validates qwen2 + qwen3 + gemma3) should precede Slice G
G Gemma-4 Rust loader upstream Python gemma4 support AND Slice 1d passing
Q2 Qwen2 Rust loader — unlocks Qwen2.5-Coder in production adjacent practice for Slice G

Details in bench/ANALYSIS.md and the raw JSON at bench/qwen25coder_kvq_validation.json.

Two things carry forward:

  1. Any MLX Array read via as_slice after a transpose must flatten first. Documented in cache.rs; future backends follow suit or justify skipping.
  2. The value pathway is the quality floor, not the key pathway. Any compression plan should prioritize better value quantization (rotation? learned group scales?) over squeezing key bits.

Neither was visible from the perplexity number alone; both fell out of ablating before sweeping. The compression is invisible now — a ratio=1.32x line in a log nobody reads. The best infrastructure is the kind you get to forget.