KV cache quantization: fp8 (KV8) + 4-bit TurboQuant (KV_TQ) on CPU, CUDA, and Metal

[NeuralNotwerk]

Summary

Quantized KV cache for the MLA latent across all three backends, as opt-in tiers:

• KV8=1 — fp8 e4m3 + per-row scale (~3.9× less KV RAM than f32), CPU + CUDA + Metal.
• KV_TQ=4 — 4-bit KV (~7.8× less): rotated-int4 + Lloyd-Max codebook by default (full reconstruction, so it corrects the MLA key and value), or PolarQuant (arXiv 2502.02617, KV_TQ_POLAR=1, variable 2-6 bit). CPU + Metal; under CUDA the quantized paths self-disable (guarded, falls back to CPU attention).
• Metal backend: fp8/TQ encode+decode run inside the fused GPU attention (producer and consumer both on-GPU), and Metal becomes opt-out (default-on when a device is present, COLI_METAL=0 to force CPU) — Apple unified memory has no separate VRAM pool to overflow, so the CUDA-style opt-in gate doesn't apply.

Supersedes #299/#300, which I withdrew for fleet soak-testing: the VRAM-budget interaction I flagged there (fp8 shadow vs CUDA_EXPERT_GB=auto) is fixed in this series ("KV8 shadow-aware auto VRAM budget" + "KV_SHADOW auto keyed on ragged decode"), and the stack has since gathered the review fixes, vectorized fp8 kernels (15.8× prefill, 4-6× decode vs scalar), the 4-bit tier, and a second architecture (Metal) worth of validation.

Commits (bisectable, each builds)

1. KV8=1 — fp8 e4m3 latent KV (CPU + CUDA), .coli_kv v2 — encoder bit-exact with the engine's consumers (RNE, per-row amax/448 scale); disk format v2 with quantize-on-upgrade from v1, reject-on-mismatch, self-heal.
2. KV8 long-context CUDA — fp8 device shadow (each row uploads once instead of re-streaming the whole history per layer per block: ~170 GB of PCIe saved on a 124k prompt), split-T kernels, DSA-aware gather.
3. KV8 shadow-aware auto VRAM budget + KV_SHADOW auto policy — the KV8=1: fp8 e4m3 latent KV cache (CPU + CUDA) with .coli_kv v2 #299 withdrawal issue.
4. KV_SHADOW auto keyed on ragged decode; VRAM floor at shadow alloc.
5. KV8 review fixes + vectorized fp8 kernels (15.8× prefill, 4-6× decode).
6. Metal fp8/int4 KV attention + CPU TurboQuant/int4 KV tiers — kv_tq.h codecs (randomized-Hadamard rotation + recursive polar / rotated-int4 + Lloyd codebook), MSL kernels a_fp8enc/a_score8/a_clat8/a_tq_enc/a_tq_deq, 3-way qmode in the fused layer decode, .coli_kv v3 (COLIKV3, h[7]=codec<<8|bits), Metal opt-out default, tests.

Why rotated-int4 for the 4-bit tier

MLA's latent is both key and value, so score-only fixes (QJL-style) leave the value error uncorrected. Full reconstruction through a randomized-Hadamard rotation + fixed 16-level Lloyd-Max codebook (N(0,1)-optimal, scaled by radius/√n) measures ~0.096 rel-L2 at 4 bits vs ~0.147 for recursive-polar — and the rotation makes the radius invariant, so one per-row scale survives. 4-bit KV is a quality trade (tiny-oracle: f32 32/32, KV8 30/32, TQ4-int4 23/32) — that's the physics of 4 bits against MLA, not codec slack; KV8 is the lossless-feeling tier, TQ4 the memory-constrained one.

Validation

• Unit: make test-c green (new suites: test_kv_fp8 exhaustive e4m3 roundtrip + RNE; test_kv_tq rotation self-inverse, roundtrip, distortion, inert rows; test_kv_disk v1/v2/v3 round-trip + upgrade/reject/self-heal; test_kv_alloc).
• Metal: make metal-test green — fp8 GPU encoder byte-exact vs coli_fp8_enc; TQ GPU↔CPU roundtrip ~3e-8; fused attention parity vs CPU ~5e-6 for fp8/int4/polar; full-layer residual parity for f32/KV8/TQ (~4-5e-6).
• Live (GLM-5.2 744B int4, macOS ARM, this exact branch): f32/KV8/KV_TQ=4 all run the fused Metal attention with 0 CPU fallbacks, greedy output Metal == CPU; KV_TQ=4 + MTP DRAFT=2 decodes at 0.30 tok/s with 100% draft acceptance on a counting prompt (MTP head repaired per tools: repair int4-converted MTP heads in place; warn on --mtp --ebits <8 #397 — the two features compose).
• CUDA: the KV8 CUDA stack is the same code soak-tested on our fleet since KV8=1: fp8 e4m3 latent KV cache (CPU + CUDA) with .coli_kv v2 #299/KV8 long-context CUDA: fp8 device shadow, split-T kernels, DSA-aware gather #300, now with the VRAM-budget fix; the new TQ tier is explicitly excluded from every CUDA path (!g_tq guards on the absorb gates — no TQ CUDA kernels exist yet, phase 2). Not re-compiled on this macOS machine; last CUDA build/test was the fleet run.

Notes for review

• The f32 path is untouched when both tiers are off (KV8/KV_TQ unset → byte-identical behavior; the quantized byte-caches are NULL and every consumer branches on g_kv8||g_tq).
• KV8 and KV_TQ are mutually exclusive (KV_TQ wins with a warning).
• Disk cache: v1 (f32) loads under KV8 via quantize-on-load without rewriting the file until the first save; v2/v3 under a mismatched mode are rejected (start clean) rather than misread.
• Metal-off (COLI_METAL=0) and non-Metal builds behave exactly as the CPU tier.

🤖 Generated with Claude Code

[JustVugg]

Reviewed against post-v1.0.0 dev. The design is right for this codebase — opt-in tiers, bit-exact encoder/consumer contract, disk format v2 with upgrade+self-heal, per-backend tests — and 18 of 19 files merge clean. The one conflict is the important one: c/glm.c in the CUDA attention dispatch. dev recently grew the ragged batched-attention branch (COLI_CUDA_ATTN → coli_cuda_attention_project_ragged, #365/#372 line), which your series predates — and your sel_any DSA-gating for KV8 rewrites exactly that region.

How the ragged branch interacts with the fp8 device shadow (does a ragged gather read the shadow rows, or does it re-upload?) is a semantic call I'm not comfortable resolving for you, and I have no CUDA hardware here to validate it. Could you rebase the series onto current dev and make that call explicitly? Everything else looks mergeable as-is, and with your soak-testing infra behind it this lands with confidence. The OLMoE testbed (#362, just merged) may also be handy for cheap CPU-side KV-quant A/Bs.

[NeuralNotwerk]

Rebased onto current main (post-#411, 72874f3) — all six commits re-validated (CPU build + test suite, Metal build + parity tests green; CUDA via CI).

The ragged × fp8-shadow call, made explicit in the code (comment at the gate in attention_rows): neither — the ragged path doesn't read the shadow and doesn't re-upload. The shadow is primary-sequence-scoped by design: it mirrors only m->kv incrementally, and per-slot ragged KVStates are never shadowed. More decisively, under KV8 the f32 rows the ragged wrapper passes (kvs[s]->Lc/Rc) aren't even allocated (the quantized byte-caches replace them), so letting the ragged branch run would deref NULL. So the gate is kvs && !g_kv8 && !g_tq: quantized multi-slot rows fall through to the CPU ragged path, which decodes fp8/TQ natively — the same fallback contract the rest of the series uses for paths without quantized kernels. A KV8-aware ragged gather (per-slot shadow, or per-call fp8 upload + an fp8 ragged kernel) is marked as future work in the comment; happy to do it as a follow-up if the multi-slot + KV8 combination matters to anyone's deployment.

Also folded in during the rebase: the batch-absorb gate keeps its sel_any DSA-skip semantics alongside upstream's new ragged branch (they're disjoint by construction — ragged requires kvs, absorb requires !kvs), and ragged_kv_append/project_ragged are untouched. Thanks for the OLMoE pointer — that testbed should make CPU-side KV-quant A/Bs much cheaper than the 744B.

[JustVugg]

Heads-up: c/glm.c was just split into c/colibri.c + header modules (#391, now merged into dev). This PR touches the old glm.c, so it will need a rebase onto current dev with its changes moved into colibri.c (or the relevant extracted header: quant.h, sample.h, kv_persist.h, telemetry.h, grammar.h). The make glm target still works (alias of make colibri), so no build scripts break. Apologies for the churn — ping me if the rebase gets thorny and I can help place the hunks.

[NeuralNotwerk]

CUDA validation on fleet hardware (consolidated)

Ran the rebased series on our fleet box — 4× RTX 5090 (sm_120) + RTX 4090 (sm_89), EPYC 7532, 503 GB RAM, CUDA 12.8.1. (Context for the sparse CUDA numbers lately: this box pulls ~4.5 kW and North Carolina is mid-heat-wave, so it only comes up in cool windows.)

Build & tests: real nvcc compile for both arches via the same two-stage image build we deploy — zero errors; full make test-c green on Linux (incl. the new test_kv_fp8 / test_kv_tq / test_kv_disk v3 suites).

Correctness, live on GLM-5.2 744B (COLI_CUDA_ATTN=1, DRAFT=2, greedy): f32 and KV8 both generate correctly with the fp8 device shadow engaged and identical 98% MTP acceptance (39/40); KV_TQ=4 exercises the !g_tq gates exactly as documented — attention falls back to CPU, no faults (TQ's CUDA kernels are explicitly future work).

Performance — measured at the production operating point (serve mode over HTTP, exact prod env: CUDA_EXPERT_GB=auto CUDA_DENSE=1 KV8=1 COLI_KV_SLOTS=4 CTX=262144, tuned OMP, IPC_LOCK): the branch built into our production image pinned 1501 experts / 28.4 GB on device 0 (marginally more than the deployed image's 1315 / 24.9 from the same profile) and served a 256-token completion in ≤41 s wall including first-request prefill — ≥6 tok/s, matching what the currently-deployed e03bd86 image delivers on this box. A controlled binary A/B (deployed e03bd86 vs this branch, identical harness and env) is indistinguishable across every metric — pins, hit rate, acceptance, tokens/forward, throughput. Performance-neutral vs the soak-tested KV8 stack.

---

On the #391 heads-up (glm.c → colibri.c + extracted headers): thanks for the warning and the offer — happy to do that rebase as soon as the split lands in main (or retarget the PR to dev if you'd rather take it there; whichever you prefer). Most of this series' surface (KV producer/consumer, disk format, env setup) should map cleanly onto quant.h/kv_persist.h; I'll ping if any hunk placement is ambiguous.

[michael-denyer]

Tested this branch on an M5 Pro, 64 GB, macOS 26.5.2, GLM-5.2 REAP-504B int4 container, expert-pruned to 168 per layer, streamed from disk. Two findings, one of which you probably want before merge.

1. With KV8 and KV_TQ both unset, greedy output diverges from this branch's own merge-base at the first generated token. Merge-base 3cd2674 opens "We need to explain how", this branch opens " Could you explain it based on the paper". Same env for both: COLI_METAL=1 DIRECT=1 MTP=0 PIPE=1 PIPE_WORKERS=8 RAM_GB=50 TOPP=0.7 TEMP=0, same .coli_usage, .coli_kv wiped before each run. Both outputs are deterministic across reruns, so this is a real logits change, not noise. That contradicts the "f32 path untouched when both tiers are off" note, at least on Metal with this container. First place I'd look is the 3-way qmode dispatch in the fused layer decode.

2. Short-context throughput, 13-token prompt, 128 generated, two rounds per arm: f32 2.95 and 3.65 tok/s, KV8 2.70 and 2.90, KV_TQ=4 1.89 and 2.16. The TQ4 attention wall roughly doubles, 8.9s baseline to 16.2s and 14.2s. Context this short has no KV-size upside, so treat these only as a floor for the tier costs on Metal.

Caveat on both: the CLI prompt is untemplated and the model degenerates into a token loop after about 30 tokens in every arm, merge-base included. The degeneration is deterministic, so the first-token divergence and the relative throughput stand, but the absolute tok/s is not a chat workload.

I can rerun with different knobs or capture fuller logs if that helps narrow down the qmode question.

[NeuralNotwerk]

Ready for review — validated on Metal, CPU, and CUDA

Both review findings are root-caused and fixed, and the KV-quant stack is now hardware-validated across all three backends (incl. an end-to-end run on GLM-5.2 744B).

@michael-denyer's findings — both fixed

1. f32 first-token divergence — root cause was not the qmode dispatch. The series compiled the entire Metal shader library with MTLMathModeSafe (added for the fp8 RNE encoder), silently recompiling every f32 kernel fast-math-off. A byte-identical A/B proved mm_gemv alone drifts ~4e-6, compounding to flip greedy argmax at token 1. Fix: two libraries — fast-math for all f32/MoE/consumer kernels (restores merge-base parity), safe-math only for the encoders. fp8 stays byte-exact.
2. TQ4 slower than f32 — TQ4 was the only tier with no native quantized consumer (Metal re-dequantized the whole cache to f32 staging every step; CUDA silently fell to CPU; CPU re-dequantized each shared latent row 2H=128×). Fix: the rotated-int4 codec's randomized-Hadamard rotation is orthogonal, so q·x̂ == rotate(q)·c and Σ wₜx̂ₜ == unrotate(Σ wₜcₜ) — rotate the query once per head, dot packed nibbles through the codebook, never reconstruct the cache. Native kernels on Metal + CUDA + the CPU consumer; codec-0 (PolarQuant) keeps the staging path.

Validation

• Metal — make metal-test green; native TQ4 matches the old staging path to ~5e-7 and is 1.4–1.9× faster, gap widening with context; f32 regression fixed (fp8 byte-exact).
• CPU — score/context orthogonality identities + full MLA-consumer equivalence unit-tested (test_kv_tq.c).
• CUDA — built on 4×RTX 5090 + RTX 4090 (sm_89+sm_120), compiles clean; make cuda-test passes incl. a new attention_absorb_tq case; end-to-end on 744B the native kernel engages and TQ4 output is coherent with MTP acceptance identical to f32.
• Fastest CUDA config (TQ4 + MTP=2, 5 GPUs, fully pinned): prefill ~4.95 tok/s, decode 3.61 tok/s, 100% resident (nothing streams), coherent output.

Also in this push

• openai_server.py: the stdout dispatcher no longer tears down on an unrecognized telemetry frame (which failed every in-flight request) — single-slot and multi-slot now present the identical HTTP interface, resilient to engine/server version skew.

Remaining

One open item: the rebase onto the colibri.c split (#391) — the glm.c hunks move into colibri.c / quant.h / kv_persist.h. @JustVugg — happy to do that now, or take you up on the offer to help place the ambiguous hunks; your call on landing it here vs. retargeting to dev.

[NeuralNotwerk]

Rebased onto current dev (colibri.c + extracted headers, post-#391) — the series is additive so this was mostly relocation: KV8/TQ globals + kv_alloc + the attention consumers → colibri.c; the .coli_kv v2/v3 disk format → kv_persist.h; kv_fp8.h/kv_tq.h are new headers. The KV-tier flags moved above the kv_persist.h include so the disk format can see them, and the pin_load KV8-shadow VRAM projection is merged with #445's additive prefix budget (both present). Also folds in the openai_server.py unified single/multi-slot fix.

Rebased tree is green locally: make colibri, make test-c (incl. kv_fp8/kv_tq identities + kv_disk v1→v2/v3 upgrade/reject/self-heal), and make metal-test (f32 byte-exact, native TQ4 1.3–1.5× over staging). CUDA is unchanged from the hardware-validated version (compiles sm_89+sm_120, cuda-test incl. attention_absorb_tq, coherent end-to-end on GLM-5.2 744B). No conflicts vs dev.