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.

[JustVugg]

Rebase needed: #298 just merged into dev — it reworked the dense/attention CUDA kernels' scale handling (per-group fmt=4 scales, new gs/ng kernel params), which is exactly where KV8/KV_TQ dequant sits. The conflict is semantic, so you're the right person to resolve it: your quantized-KV dequant should now compose with absorb_scale in the attention kernels rather than the old per-row assumption. This PR is on the v1.0.1 shortlist — happy to help if you want a second pair of eyes.

[JustVugg]

Status: this conflicts with dev. Several PRs landed today across colibri.c and openai_server.py (#525, #526, #532), so the overlap is partly self-inflicted on my side — apologies for the moving target.

On the substance: KV quantization is a real lever here and I'd like it, but the shape of the PR is what's slowing it down. It's ~3,400 lines across 20 files covering two schemes (fp8 KV8 and 4-bit TurboQuant) on three backends (CPU, CUDA, Metal). I can validate the CPU path; I have neither NVIDIA nor Apple hardware, so the other two thirds rest entirely on your measurements. That's a lot to accept in one merge.

What would unblock it: land the CPU path first, on its own. It's the one I can verify end to end against the oracle, and once the format and the KV plumbing are settled and merged, the CUDA and Metal kernels become much smaller, more focused follow-ups — each with a hardware owner who can gate it (this project's rule for GPU changes is a token-identity check from someone with the silicon).

Also worth stating explicitly in the PR: what happens to KV persistence (.coli_kv, restored across restarts) when the quantization mode changes between sessions. A cache written as fp16 and read back as fp8 — or the reverse — must be refused rather than misread, the same way the loader refuses containers whose byte counts don't add up.

Tell me if splitting is workable and I'll prioritise the CPU half.

[NeuralNotwerk]

Splitting is workable — I'll carve the CPU path out as its own PR on current dev, with CUDA and Metal as hardware-gated follow-ups. Meanwhile, here's the quality data you asked for, on the tiny oracle you can reproduce yourself.

Teacher-forcing token-exact vs the transformers oracle

tools/make_glm_oracle.py (transformers 5.14.1, seed 1234 — the regenerated ref_glm.json is byte-identical to the committed one), then SNAP=./glm_tiny TF=1 [KV8=1|KV_TQ=6|KV_TQ=4] ./colibri 64 16 16 per backend:

KV mode	CPU-x86 (AVX2)	CPU-ARM (NEON)	Metal	CUDA	cross-backend
f32	32/32	32/32	32/32	32/32	exact
KV8 (fp8)	30/32	30/32	30/32	30/32	identical positions + tokens
TQ6	30/32	30/32	30/32	30/32	identical
TQ4	23/32	23/32	23/32	23/32	identical

Greedy generation (decode path — this is where the native CUDA/Metal quantized-KV kernels engage):

KV mode	CPU-x86	CPU-ARM	Metal	CUDA
f32	20/20	20/20	20/20	20/20
KV8	14/20	14/20	14/20	14/20
TQ6	17/20	17/20	17/20	17/20
TQ4	10/20	10/20	10/20	10/20

The TQ4 generated sequence is byte-identical CUDA vs Metal vs CPU.

Reading this honestly:

• f32 stays 32/32 and 20/20 everywhere — the PR does not perturb the full-precision path (the earlier Metal fast-math regression is fixed and stays fixed).
• Every quantized flip happens at the same position with the same token on all four backends. The divergence vs the oracle is the quantization itself — deterministic and implementation-independent — not a backend bug. This is the specific failure mode you called out from cuda+engine: full fmt=4 (grouped int4 gs=64) support + diagnostic harness #298 ("looked right, produced garbage"), and it's the thing the cross-backend identity rules out.
• The tiny random-weight model is a worst case: logit margins are near zero, so tiny KV perturbations flip argmaxes that a trained model's margins absorb (the 744B stays coherent under TQ4 end-to-end). KV8 and TQ6 lose 2/32 even on this stress test; TQ4 is the aggressive tier — TQ6 is the quality-parity choice when the 4-bit loss matters.

KV persistence across mode changes

Already handled the way you describe, and covered by tests/test_kv_disk (in this PR): the .coli_kv magic encodes the tier (COLIKV/COLIKV2/COLIKV3) and h[7] carries the TQ codec+bits. On any mismatch the file is refused and restarted, never misread, with an explicit message, e.g.:

[KV] .coli_kv is fp8 (saved under KV8=1): starting over (set KV8=1 to resume it)
[KV] .coli_kv KV_TQ (codec 1, 4-bit) != current (codec 1, 3-bit): starting over

The test covers v1→v2/v3 upgrade (old file untouched until the first save rewrites it), v2/v3-under-f32 reject, TQ codec/bits mismatch reject, and the append-side self-heal.

Perf reference (already posted above, environment for the record)

Resident 744B, driver 610.43.02, CUDA 12.8.1 runtime, 4×RTX 5090 + RTX 4090 (sm_120+sm_89): decode FP32 6.96 / FP8 7.01 / TQ4 6.95 tok/s — parity.

Next from me: the CPU-only PR (KV8 + KV_TQ consumers, kv_fp8.h/kv_tq.h, kv_persist.h v2/v3, the flags/plumbing) rebased on current dev, with the table above's CPU column as its validation. CUDA and Metal follow once that lands.

[NeuralNotwerk]

Split done — the CPU path is now its own PR: #553 (1,282 lines / 12 files, no backend code), rebased on today's dev, with the tiny-oracle TF/generation tables and the .coli_kv mode-mismatch behavior documented in the description. It carries !g_kv8&&!g_tq guards on the existing CUDA/Metal fast paths so quantized KV is safe (CPU-consumer-bound) on any build today.

Converting this PR to draft: it stays as the reference for the full three-backend implementation, and once #553 lands I'll rebase the CUDA and Metal kernels out of it as the two focused follow-ups (each with its hardware validation + token-identity data — the CUDA half is already validated on 4×5090+4090 at f32-parity decode on the 744B).

[NeuralNotwerk]

Rebased onto current dev (26629e7) — head is 23b7628, and the PR is MERGEABLE again after sitting DIRTY. Still a draft, still parked behind #553 by design; this is to keep it from rotting while the CPU half lands, and to surface the semantic merges early rather than at carve-off time.

Four files conflicted. Two were mechanical, three were semantic and worth review:

1. DSA gather — dev's COLI_DSA_GATHER vs this branch's fp8-shadow gather. dev grew an f32 DSA path (#510) that stages the selected top-k rows into a compact pair and hands them to the dense absorb kernel. This branch had replaced that whole region with "DSA selection ⇒ only the fp8 shadow has a gather kernel, else CPU", which would have silently deleted the new f32 feature. Both are kept and routed by ownership: KV8 with a live shadow for this layer takes the shadow's absolute-index gather; everything else falls to dev's staging path — except under KV8-without-shadow or any KV_TQ tier, where the f32 Lc/Rc rows that path reads are not allocated and there is no TQ gather kernel, so it drops to the CPU consumer. That last guard is new code, not a merge artifact.

2. CUDA_EXPERT_GB=auto headroom. dev replaced the hardcoded 2 GB reserve with the configurable g_cuda_reserve_gb; this branch subtracts the KV8 shadow projection. Merged so the shadow composes with the configurable reserve rather than reinstating the 2 GB constant.

3. HIP portability — this one is a real behavior change. dev added backend_gpu_compat.h and a HIP syntax-check CI job after this branch was written. The KV8 CUDA path includes cuda_fp8.h and decodes through __nv_cvt_fp8_to_halfraw, neither of which exists under hipcc — as written, the rebase would have broken the HIP build. Following that header's own stated design ("every platform difference lives in this header"), the fp8 include moved into its CUDA branch behind a new COLI_GPU_HAS_FP8, and fp8_e4m3() gained a portable bit-math fallback for HIP: normals (1+M/8)*2^(E-7), denormals M*2^-9, S.1111.111 reading back as 0. Same value grid as the CUDA cvt and as the host LUT in kv_fp8.h, just not the hardware instruction. The HIP path is untested on real hardware — it compiles, and the decode matches the CPU LUT by construction, but nobody has run KV8 on ROCm.

4. tests/test_backend_cuda.cu — took dev's coli_cuda_tensor_upload_g signature (#298 grouped scales) at all three sites. CPU attention hunks — same resolution as #553: quantized branches first, dev's #442 SIMD f32 dots preserved verbatim in the else.

Validation on the rebased tree

Run here on macOS 26.5.1 / Apple M4 / 32 GB:

• make check green — 167 Python tests, every C suite including the four KV ones (test_kv_fp8, test_kv_tq, test_kv_disk, test_kv_alloc).
• CPU and METAL=1 both build and link clean.
• make metal-test passes end to end on Apple silicon — the Metal third of this PR is hardware-validated on this exact rebased tree:
  • fp8 encode vs the CPU twin: byte_diffs=0, scale_err=0 at n=1 and n=256
  • fused fp8 attention: nerr 3.6e-06 – 5.3e-06 across S=1/3/4 and MTP
  • TQ round-trip GPU-vs-CPU: 0 – 1.1e-07 rel (polar b3/b4 and int4, n=64 and 512)
  • TQ int4 native-vs-staging equivalence ≤ 8.2e-07, with 1.33× / 1.43× / 1.35× / 1.89× speedups at T = 65 / 257 / 1025 / 2049
  • full-layer decode residual: f32, KV8, TQ-int4, TQ-polar4 all ≤ 4.7e-06
• CUDA host paths syntax-check clean under -DCOLI_CUDA. nvcc/hipcc aren't available on this machine, so the compile check is CI's — the CUDA and HIP syntax jobs both pass on 23b7628.

What this does not re-validate

The fleet-box CUDA numbers I posted on the 20th were measured on the pre-split tree, and the DSA-gather routing and the reserve accounting both changed in this rebase. Those need a re-run on 23b7628 before any CUDA carve-off is offered for merge — I'm not carrying the old numbers forward as if they still applied. Same for the ROCm path, which has never been exercised. Metal is the one third that is currently validated on the rebased code.

[NeuralNotwerk]

Rebased onto current dev (fb0a726), head 85852e2, CI 10/10 green. Still a draft, still parked behind #553 — this is to keep it from rotting and, more usefully, because the rebase produced the AMD data point this PR never had.

The rebase

One conflict, in backend_gpu_compat.h, against the rocWMMA work from #599: dev now #errors without rocWMMA headers, defines __CUDA_ARCH__ 700 under HIP, and remaps __half to rocwmma::float16_t. Both sides kept. The KV8 COLI_GPU_HAS_FP8 fallback deliberately does not route through that __half mapping — rocwmma::float16_t is a WMMA fragment type and the fp8 fallback is plain float arithmetic, so the two never meet; noted in the header so nobody "simplifies" them together later.

Also folded in from the #553 review pass (cherry-picked, so it collapses cleanly when #553 lands):

• KV_TQ now refuses non-power-of-two latent row widths. Both codecs rotate through a radix-2 FWHT and return an inert radius 0 otherwise, so on such a model every row quantized to zero and the engine produced confident garbage with no diagnostic. test_kv_tq pins that behavior for all four entry points so the guard can't quietly become a no-op.
• KV_TQ=1 lands on the recommended 4-bit tier, not the most aggressive one it used to clamp up to.
• The GPU test binary's summary said q8/q4/q2/f32 correctness ok while also covering the fp8/KV8 and rotated-int4 cases inline. It now names them, so a green run stops understating itself.

Token identity across four backends, on this tree

The claim this PR has been making is that the quantized flips are deterministic quantization loss rather than backend divergence. That now holds across two ISAs and two GPU vendors, all re-measured on 85852e2 — not carried over from the pre-split measurements:

KV mode	CPU ARM/NEON	CPU x86/AVX2	CUDA (RTX 4090, sm_89)	HIP (Radeon AI PRO R9700, gfx1201)
f32	32/32	32/32	32/32	32/32
KV8=1	30/32	30/32	30/32	30/32
KV_TQ=4	23/32	23/32	23/32	23/32

Identical positions, identical tokens, four backends.

The AMD column is new — these kernels had never been run on ROCm. make hip-test passes (exit 0) on ROCm 6.4.2 / gfx1201, covering the fp8/KV8 absorb kernels and the rotated-int4 attention_absorb_tq kernel against host-dequantized references. The single-source claim in backend_gpu_compat.h holds through the KV-quant additions with no AMD-specific kernel code.

Other backends on the rebased tree:

• Metal (macOS 26.5.1, M4): make metal-test green end to end — fp8 encode byte-identical to the CPU twin (byte_diffs=0), fused fp8 attention nerr ≤ 5.3e-06, TQ round-trip GPU-vs-CPU ≤ 1.1e-07, TQ int4 native-vs-staging equivalent to ≤ 8.2e-07 with 1.46–1.93× speedups at T = 65…2049, full-layer decode ≤ 4.7e-06 for f32/KV8/TQ-int4/TQ-polar4.
• CPU: make check green on macOS and Linux (197 tests).

What this still does not claim

No perf numbers are carried forward. The fleet-box CUDA throughput figures I posted on the 20th were measured pre-split, and this rebase changed both the DSA-gather routing and the CUDA_EXPERT_GB=auto reserve accounting, so they need re-running on 85852e2 before any CUDA carve-off goes up for merge. Everything above is correctness, not speed.

One observation for whoever owns the AMD path: on this box the tiny-model teacher-forcing rate went from ~3230 pos/s to ~159 pos/s after #599 enabled rocWMMA. It's a 5-layer random-weight model where launch overhead dominates, so it may well be meaningless — but it's a 20× delta on dev's own change and probably deserves a look with a real model before anyone quotes AMD throughput.

[JustVugg]

Status check, and a question rather than a rebase request.

dev has moved substantially and this now conflicts in c/colibri.c, c/backend_metal.mm and the Metal test. Before asking you to rebase 3,483 lines across those, I want to check the sequencing with you, because #553 is the CPU half of this PR, split out at our request, and it is the one that has been reviewed and independently reproduced.

As I read it:

• CPU: KV cache quantization — KV8 (fp8 e4m3) + KV_TQ (rotated-int4 / PolarQuant) #553 — CPU KV8 + KV_TQ, the format and the KV plumbing. Reviewed end to end, numbers reproduced here, and you rebased it and fixed two no-silent-error gaps after that review.
• KV cache quantization: fp8 (KV8) + 4-bit TurboQuant (KV_TQ) on CPU, CUDA, and Metal #399 — the same thing plus the CUDA and Metal kernels, i.e. this PR is a superset that includes the CPU work already sitting in CPU: KV cache quantization — KV8 (fp8 e4m3) + KV_TQ (rotated-int4 / PolarQuant) #553.

If that is right, rebasing both is wasted effort for you: they will keep conflicting with each other as well as with dev. The path I would suggest is to land #553 first — it is the one with the oracle-verifiable half — and then reduce this PR to just the CUDA and Metal kernels on top, which is a much smaller and much more reviewable diff, and the shape you described yourself when you proposed the split.

If I have the relationship backwards, tell me and I will treat #399 as the live one instead.

Either way, one thing changed on our side that makes the next rebase cheaper: c/Makefile used to be a single hand-written TEST_BINS line that every test-adding PR appended to, so any two of them conflicted by construction. Gates are derived from the build rules now (#733) — drop your TEST_BINS entry and your test is picked up by its own rule. Merge order is also oldest-clean-first from here on.

[ThefloorMiner]

Ampere data point: this branch builds and runs on 4× RTX A6000 (sm_86) with CUDA 13, on GLM-5.2 744B. Posting it because @JustVugg's sequencing note says the CUDA follow-up would be gated on a hardware owner's check, and this is the hardware. It is not the token-identity check — I explain below exactly why I cannot claim that one, so nobody treats this as more than it is.

Host

GPU	4× RTX A6000 48 GB, sm_86, driver 580.173.02, all PCIe Gen4 x16
CPU	EPYC 7402P, Zen 2, 24c/48t — AVX2, no AVX-512
RAM	264 GB DDR4-2400, 8/8 channels
Model	mastouri/GLM-5.2-colibri-int4-g64-with-int8-mtp, 142 shards, 429.3 GB
Toolchain	CUDA 13 via micromamba (nvcc 12.6 + gcc 13 fails with _Float32 undefined)

Build

make colibri CUDA=1 CUDA_ARCH=sm_86 CUDA_HOME=~/cuda13

Exit 0. Two warnings, one of them yours-adjacent and harmless:

backend_cuda.cu(1315): warning #177-D: variable "ng" was declared but never referenced
      int ng = tensor->ng > 0 ? tensor->ng : 1;

Binary links libcudart.so.13 and carries the [CUDA] markers, i.e. a real CUDA build rather than a silent CPU fallback.

Runs

Three conditions, COLI_MODEL on NVMe, COLI_CUDA=1 COLI_GPUS=0,1,2,3 CUDA_EXPERT_GB=176 CUDA_DENSE=1 RAM_GB=235 CTX=8192 COLI_TEMP=0 NGEN=48, engine stopped between runs, .coli_usage restored byte-for-byte before each:

	loads	throughput	output
baseline	23.01 s	1.14 tok/s	coherent
KV8=1	23.35 s	1.14 tok/s	coherent
KV_TQ=4	23.88 s	1.14 tok/s	coherent

No crash, no out of memory, no fallback message. Throughput identical across all three, which is what your PR description predicts for CUDA — the quantized paths self-disable there.

Generated text, prompt "Explique en une phrase ce qu'est un mixture of experts":

baseline : Un "mixture of experts" (ou mélange d'experts) est une architecture
           d'intelligence artificielle qui divise une tâche complexe...
KV8=1    : Un « mixture of experts » (mélange d'experts) est une architecture
           d'intelligence artificielle qui divise une tâche complexe...
KV_TQ=4  : Un "mixture of experts" est une technique d'intelligence artificielle
           qui divise une tâche complexe entre plusieurs modèles spécialisés...

All three fluent, correct, and semantically equivalent.

What I did not establish, and why it matters here

The three outputs are not token-identical despite COLI_TEMP=0. The tempting reading is "KV8 changes the output". I do not think that reading is supported, and I am not making it, for two reasons:

1. Your PR states the quantized paths self-disable under CUDA. If that holds, KV8=1 should have been a no-op — so a difference between baseline and KV8=1 is more likely to be run-to-run variation than an effect of the flag.
2. I never ran the baseline twice. Without that, I have no idea whether this build is token-reproducible at all on 4 GPUs. Cross-device FP reduction order under #pragma omp parallel for if(g_cuda_ndev>1) (colibri.c:4722) is an obvious candidate for run-to-run drift, and KV_TQ=4 — which certainly self-disabled — also differs from baseline, which points the same way.

So the honest statement is: on this host I cannot currently perform a token-identity check, because I have not shown the baseline is token-identical to itself. That seems worth knowing before the CUDA follow-up is gated on exactly that test, since if multi-GPU is non-deterministic the gate needs to be a distributional check rather than an equality one.

The host is powered down for now. When it is back the first thing I will run is baseline ×3 with the same seed and config, and report whether identity holds at all. If it does, the KV8 comparison becomes meaningful and I will run it properly.

Unrelated context that argues for landing this

Per-phase profile from the baseline run above, same host, decode:

expert-disk    2.345 s   ( 4.8 %)
expert-matmul  9.074 s   (18.6 %)
attention     31.264 s   (64.2 %)   ← of which score-softmax-value 20.398 s
other          5.987 s   (12.3 %)

Attention dominates decode by a wide margin on this class of host, so KV cache width is not a memory footnote here — it is on the critical path. And on the capacity side, cap_for_ram() reserves 47.7 GB of fp32 KV at --ctx 262144 on a 264 GB machine, which is what drives the expert cache to its cap=1 floor (#768). Halving that reservation is the difference between long context being usable and being theoretical on 4× A6000.

Happy to run whatever specific check you want when the machine is back up — including on #553's CPU path, where determinism should be much easier to establish.

[ThefloorMiner]

Follow-up to my previous comment: the prerequisite is answered, and it came out the way that helps you.

I said I could not perform the token-identity check because I had never shown the baseline was token-identical to itself on 4 GPUs — and that if multi-GPU reduction order drifted, the gate you proposed would need to be distributional rather than an equality.

Ran it. Six identical requests, GLM-5.2 744B, 4× RTX A6000, temperature 0, production config:

run 1 : 338 bytes · e86d393c506f
run 2 : 338 bytes · e86d393c506f
run 3 : 338 bytes · e86d393c506f
run 4 : 338 bytes · e86d393c506f
run 5 : 338 bytes · e86d393c506f
run 6 : 338 bytes · e86d393c506f

6 runs -> 1 distinct output

Deterministic to the token on 4 GPUs. My worry about cross-device FP reduction order at colibri.c:4722 was unfounded — whatever that path does, it is order-stable here.

So the check you wanted to gate the CUDA follow-up on is achievable on this hardware, and it can stay an equality test rather than becoming a distributional one. That is the simpler gate and the stronger one.

What this means for my earlier numbers

It also removes the ambiguity I flagged. In my first comment the three conditions produced different text:

baseline : Un "mixture of experts" (ou mélange d'experts) est une architecture…
KV8=1    : Un « mixture of experts » (mélange d'experts) est une architecture…
KV_TQ=4  : Un "mixture of experts" est une technique d'intelligence artificielle…

I refused to read that as "KV8 changes the output" because I could not rule out run-to-run noise. I now can: there is no run-to-run noise. Those three outputs differ because something in the configuration differed — not because the engine wanders.

I am still not claiming it is KV8, for one concrete reason: KV_TQ=4 also differs from baseline, and by your own PR description KV_TQ self-disables under CUDA and falls back to CPU attention. A flag that is supposed to be a no-op on this path is not behaving as one. Either the self-disable is not firing, or falling back to CPU attention is itself changing the arithmetic, or both flags are perturbing something unrelated to the KV format.

That is a real question and it now has a clean way to be answered, which it did not have before. The proper run is: baseline ×2 (confirm identity within the PR build, not just the production build), then KV8=1, then KV_TQ=4, all one-shot, same seed-free greedy config, diffed byte-for-byte. If baseline ×2 matches and KV8=1 does not, that is a genuine signal about the CUDA path and worth your attention before the follow-up lands.

The host is running other measurements right now. I will post that diff.

One correction to my earlier comment, since it affects how you read the profile I posted

I attached a per-phase profile claiming attention is 64 % of decode. Fresh runs at the same CTX but 64 generated tokens instead of ~750 put expert-matmul at 46 % and attention at 26 %. The share is generation-length dependent — attention grows with the KV as the answer gets longer — so "attention is 64 %" is true of a long generation and misleading as a general statement. I am re-checking how the two PROFILE blocks compose before I restate it anywhere.

The number in that profile that survives and is worth your eye is different, and it is from PROF=1:

P0-EXEC: routed CPU 10.967s / 19.67 GB/s (10158 row) | routed GPU critical 1.666s
         | router 1.809s | residual P2P 0.000s / 0 hop | orchestration 2.449s

19.67 GB/s for the CPU-side routed expert path, on 8-channel DDR4-2400 that sustains ~85. That is 23 % of the host's memory bandwidth, and it caps this machine at ~3 tok/s on its own regardless of anything else. I am sweeping OMP_NUM_THREADS now to separate "the path is serialised somewhere" from "4–5 cores is all it can use". If it is the former I will file it separately with the sweep attached.