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.

[NeuralNotwerk]

Rebased onto dev with #298 — the quantized-KV dequant now composes with absorb_scale instead of the old per-row scale assumption. The semantic merge, concretely:

• KV8 + TQ attention kernels (attention_absorb_kernel8 / batch / stream / split / merge / sel / sel_split / kernel_tq) now take int gs, int ng and route their weight scale through absorb_scale(wscale, fmt, gs, ng, row, k) — grouped-int4 (fmt==4) kv_b projections get per-group scales, per-row otherwise. The 12 launch sites pass w->gs, w->ng.
• quant_matmul_rows (the o_proj helper in the KV host entries) threads gs/ng into its inner quant_matmul launch.
• attention_absorb_ragged_kernel stays on your per-row path — cuda+engine: full fmt=4 (grouped int4 gs=64) support + diagnostic harness #298 left it per-row and it has no gs/ng, so I explicitly did not force absorb_scale there (that was a spot I had to be careful not to over-apply during the sweep).
• Windows DLL loader stub for coli_cuda_attention_absorb_tq so the MSVC build links.

Validated on hardware (5×RTX 5090 + RTX 4090, sm_89+sm_120):

• make glm CUDA=1 clean; make cuda-test passes incl. the grouped-int4 attention_absorb_tq case (q8/q4/q2/f32 correctness ok).
• End-to-end on GLM-5.2 744B — both f32 and TQ4 coherent, the native CUDA TQ4 kernel engages (codec-1 single-sequence), MTP acceptance matches f32.

And this closes @michael-denyer's original findings. The "TQ4 slower than f32" one is fixed — fully-resident 744B, CUDA decode (MTP off):

KV	decode tok/s
FP32	6.96
FP8	7.01
TQ4	6.95

TQ4 is now at parity (was ~2× slower) at ~28% less KV memory — the native rotated-int4 consumer landed, so it no longer re-dequantizes the cache to f32 every step. I ran the full FP32/FP8/TQ4 × MTP matrix across CUDA / CPU-x86 / Metal / CPU-ARM (resident and disk-streaming); happy to drop the whole table here if it's useful for the v1.0.1 notes.

One additive extra on the branch: PROF=1 now reports expert-load disk throughput (GB/s + experts/s) for both the initial pin and live streaming — handy for diagnosing streaming-bound boxes. Output stays byte-identical with PROF unset.

[JustVugg]

To merge this I need validation data I can't produce myself (no such GPU here). The bar every backend holds: (1) correctness first — a teacher-forcing token-exact check on a tiny model (SNAP=<tiny> TF=1 ./coli ... → N/32 vs the CPU oracle), or at minimum a clearly-coherent chat transcript; (2) perf — before/after tok/s on the same prompt + your GPU/driver/runtime version. Rebase on current dev (CI is green), get CI passing, post that data, and I merge. Correctness matters more than speed — #298 was a GPU path that looked right and produced garbage.

For KV8 / KV_TQ quantization: this changes numerics, so I need the quality picture, not just speed — perplexity or a token-exact-within-tolerance check vs full-precision KV on a tiny model, per backend (CPU/CUDA/Metal) you touched. That''s the gate for a precision-changing feature. Thanks @NeuralNotwerk.

[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.