CPU: KV cache quantization — KV8 (fp8 e4m3) + KV_TQ (rotated-int4 / PolarQuant)

[NeuralNotwerk]

The CPU half of #399, split out as requested — the part you can verify end-to-end against the oracle. 1,282 lines / 12 files (vs #399's ~3,400 / 20): no backend files, no new kernels. Once the format and KV plumbing land here, the CUDA and Metal kernels become small focused follow-ups, each gated by a hardware owner's token-identity check.

What's in it

• KV8=1 — latent KV in fp8 e4m3 + per-row f32 scale (~3.9× less KV RAM). The consumer LUT-decodes inline in the score/context dots; the per-row scale factors out of the sums.
• KV_TQ=4 — rotated-int4 codec (randomized-Hadamard rotation + Lloyd codebook; radius rides the per-row scale). The consumer uses the rotation's orthogonality (q·x̂ == rotate(q)·c): rotate the query once per head, dot the packed nibbles through the 16-level codebook, unrotate the accumulated context once — the cache is never reconstructed. KV_TQ=2|3|5|6 / KV_TQ_POLAR=1 use the PolarQuant codec.
• .coli_kv across mode changes — the file magic encodes the tier (COLIKV/COLIKV2/COLIKV3) and h[7] the TQ codec+bits; any mismatch is refused and restarted with an explicit message, never misread (e.g. [KV] .coli_kv is fp8 (saved under KV8=1): starting over). v1 files upgrade in place, old file untouched until the first save rewrites it. Covered by tests/test_kv_disk (upgrade, both reject directions, TQ-params reject, append-side self-heal).
• Memory budgeting — kv_pool_bytes accounts the quantized row widths, so cap_for_ram/PIN recover the KV savings for the expert tiers.
• Guards on the existing GPU fast paths — the CUDA/Metal attention branches read f32 rows that aren't allocated under quantized KV, so they gain !g_kv8&&!g_tq and quantized runs fall to the CPU consumer (one-time stderr notice under CUDA). A CUDA or Metal build with KV8/KV_TQ set is therefore safe today — just CPU-bound in attention until the kernel PRs.
• Second commit (small, optional to keep): PROF=1 disk-load throughput lines (pin-load GB/s + experts/s, live streaming rate). Byte-identical output with PROF unset.

Validation (reproducible on your machine)

pip install "transformers>=5.11" && python tools/make_glm_oracle.py   # byte-identical to committed ref_glm.json
SNAP=./glm_tiny TF=1 [KV8=1|KV_TQ=6|KV_TQ=4] ./colibri 64 16 16

KV mode	TF (prefill)	greedy gen (decode)
f32	32/32	20/20
KV8	30/32	14/20
TQ6	30/32	17/20
TQ4	23/32	10/20

• f32 stays token-exact — the quantized paths are strictly additive.
• The quantized flips are identical positions + identical tokens on CPU-x86 (AVX2) and CPU-ARM (NEON) — and on the CUDA/Metal implementations waiting in the follow-ups — i.e. deterministic quantization loss, not implementation divergence.
• The tiny random-weight model is a worst case (near-zero logit margins); on the real 744B, TQ4 decodes coherently end-to-end. KV8/TQ6 lose 2/32 even here; TQ4 is the aggressive tier, TQ6 the quality-parity choice.
• make test-c green (incl. the four new kv suites); builds clean with CUDA=1 and METAL=1 (the guards compile against the existing backends — no new entry points referenced).

Follow-ups queued behind this: CUDA kernels (fp8 shadow + absorb8/TQ native, hardware-validated on 4×5090+4090, sm_89+sm_120, resident-744B decode at f32 parity) and Metal kernels (fused-decode fp8/TQ + the two-library fast-math split), each with its perf + token-identity data attached.

🤖 Generated with Claude Code

[NeuralNotwerk]

CPU performance validation for this branch, measured on its exact tree (744B, fully resident — PIN=auto PIN_GB=all RAM_GB=490, 32-thread AVX2, greedy templated prompt, NGEN=80, 100% expert hit so KV precision is the only variable):

decode tok/s	MTP off	DRAFT=1	DRAFT=2
f32	3.02	2.71	2.25
KV8	2.96	2.72	2.18
TQ4	2.90	2.73	2.11

Prefill ~5.5–5.6 tok/s across all cells. Two takeaways:

• The quantized consumers hold decode parity with f32 in every MTP mode (≤4% spread, within run noise; TQ4 actually edges f32 under DRAFT=1). Resident CPU decode is expert-matmul-bound, so the KV consumers' cost is hidden — the tiers buy their memory savings for free.
• The f32 column is byte-identical behavior to dev (this branch doesn't touch the f32 path — the token-exact oracle CI confirms), and its numbers match our pre-rebase baseline within 1%, so [Performance]: hot-path float reductions are scalar and don't auto-vectorize (router matmul + MLA-absorb attention) #442's SIMD absorb dots carry over intact.

Environment: EPYC-class 32C host, ZFS storage, dockerized gcc 13 build, colibri 64 4 4.

[JustVugg]

This is on-priority — a smaller KV cache is exactly what buys longer context on a low-RAM box — and the review is positive: opt-in (KV8/KV_TQ default off), format-versioned (.coli_kv stays byte-compatible for default users, resume-safe), heavily unit-tested, and the token-exact oracle is green. I want it in. But it now conflicts with dev — several PRs landed today in the KV/attention area. Please rebase onto current dev.

Two things to confirm on the rebase, given the 'no silent errors' bar here: (1) the producer-loop parallelization in attention_rows (#pragma omp parallel for … if(S>8)) is on ALL paths including default f32 — the oracle passing says it's bit-identical, but please keep that schedule(static) deterministic through the rebase; (2) it's lossy quant with round-trip tests but no automated accuracy gate — that's inherent and fine as opt-in, just noting it's why it stays behind a flag. Rebase and I'll merge.

[NeuralNotwerk]

Rebased onto current dev (5724dae). One trivial conflict (TEST_BINS union with the new test_router_nan); every code hunk applied clean.

On your two points:

1. Producer-loop determinism — preserved verbatim: #pragma omp parallel for schedule(static) if(S>8) on the KV producer, all paths including default f32. Each s writes only its own row (no shared state; the CUDA shadow-shorten is hoisted out of the loop), and schedule(static) keeps the row→thread mapping fixed, so f32 stays bit-identical — the token-exact oracle in CI re-verifies this on the rebased tree.
2. Lossy-by-design, opt-in — agreed; that's the contract. The quality bounds we ship are the tiny-oracle tables above (KV8/TQ6 30/32, TQ4 23/32, deterministic and backend-identical) plus the round-trip/distortion unit tests. A model-level perplexity gate would need a real-weights fixture in CI — happy to discuss as a follow-up if you want one.

Also verified through the rebase: your new COLI_DSA_GATHER path composes correctly with the quant guards — it stages f32 rows, and under KV8/KV_TQ the enclosing branch is skipped (byte caches, no f32 rows), so quantized runs keep falling to the CPU consumer.

Post-rebase validation: CPU + METAL=1 builds clean, make test-c green (incl. the four KV suites), teacher-forcing/generation tables unchanged (f32 32/32 & 20/20; KV8 30/32; TQ6 30/32; TQ4 23/32 — identical positions on both builds). The CPU perf table posted above was measured on the pre-rebase tree; the rebase changes none of the measured CPU code paths (the dev-side deltas are the CUDA gather + allocator hardening), so the parity numbers stand.

[NeuralNotwerk]

Rebased onto current dev (26629e7) — conflict-free this time, head is b98cf47. CI is 10/10 green on that exact SHA.

Flagging why this needed a second pass: the rebase I posted on the 23rd targeted 5724dae, and dev moved 33 commits past it within hours, so the PR has been sitting in a conflicting state ever since — my fault for not re-checking. It's current now.

On your point (1) — producer-loop determinism: preserved verbatim. #pragma omp parallel for schedule(static) if(S>8) is at c/colibri.c:2464 on the rebased head, unchanged, on all paths including default f32. The Inkling token-exact oracle passing on this SHA re-verifies the f32 path is bit-identical through the rebase, which is the empirical check rather than my say-so.

On your point (2) — no automated accuracy gate: agreed and unchanged; both tiers stay opt-in and default off.

One correction I made while rebasing, worth calling out because it changes what this PR claims. The docs/ENVIRONMENT.md hunk was describing the other PR's work, not this one: the KV8 row advertised CUDA coverage, a device-resident fp8 shadow, split-T kernels and "11–17× faster than the CPU fallback" numbers, and there was a whole row for KV_SHADOW — an env var that appears nowhere in this PR's code. All of that lives in #399 and none of it is in these 12 files. Meanwhile KV_TQ / KV_TQ_POLAR, which this PR does add, weren't documented at all.

That was a leftover from the split, and shipping it would have meant the docs promising GPU acceleration that the merged code doesn't have. The rebased commit now documents the CPU-only reality:

• KV8 — fp8 e4m3 + per-row scale, CPU consumer only; the CUDA/Metal fused paths read f32 rows so they fall back (one-time notice under COLI_CUDA_ATTN=1); forces COLI_CUDA_PIPE=0; native kernels are follow-ups.
• KV_TQ — new row: KV_TQ=4 rotated-int4 as the recommended tier, 2|3|5|6 and KV_TQ_POLAR=1 for PolarQuant, power-of-two row-width requirement, .coli_kv v3 + the refuse-and-restart behavior on any mode/codec/bit-width mismatch, same CPU-only status.

Local re-verification on the rebased tree (macOS 26.5.1, Apple M4, arm64/NEON — the other half of the "identical flips across ISAs" claim): make check green, 167 Python tests + every C suite including the four new KV ones; METAL=1 builds and links clean with the guards compiled in.

Nothing else changed — no code differences from what you reviewed, same two commits.

[JustVugg]

Reviewed this end to end and reproduced your numbers independently — thank you for the clean split from #399, and for gating GPU behind the CPU consumer so a KV8/KV_TQ build is safe today.

What I verified locally (all good):

• The four codec unit tests pass (test_kv_fp8 exhaustive roundtrip, test_kv_tq distortion — int4 ~0.096 rel-L2 / polar4 ~0.147, test_kv_alloc, test_kv_disk upgrade/reject/self-heal). Format safety is genuinely careful — the .coli_kv tier magic + refuse-and-restart is exactly right.
• The default path is unchanged: with no KV env, TF is 32/32, byte-identical to baseline. Opt-in, default off — nobody who doesn't set the flag is affected. That's the right shape.
• Reproduced your oracle table exactly on glm_tiny: f32 32/32, KV8 30/32, TQ6 30/32, TQ4 ~23/32.

The one thing blocking a merge — real-model quality. glm_tiny is random-weight, so its logit margins are near-tied and any fp8 rounding flips argmax; 30/32 there tells us the codec perturbs, not whether it costs quality on a model that's actually confident. I fully expect KV8 to flip far fewer tokens on the real 744B — but "expect" isn't "measured", and this is a decode-quality feature.

Could you post a real-model quality datapoint — the #108 protocol (hellaswag n=200, or MMLU), f32 vs KV8 vs KV_TQ on the real GLM-5.2 gs64 container? That's the number that decides it:

• If KV8 ≈ f32 at n=200 (within CI), it merges as a recommended opt-in for RAM-constrained long-context serving, and I'll say so.
• If KV8 costs measurable quality, it still merges as experimental/opt-in with an honest "trades quality for KV RAM" note, and KV_TQ (which is clearly lossier — 9.6% codec L2) stays behind a louder caveat.

Either way the merge is close — the mechanics are solid and it's off by default. I just won't call a decode-quality feature safe on the strength of a random tiny model. Ping me with the n=200 row and we land it.

[NeuralNotwerk]

Thanks for reproducing it end to end — and for naming a decision rule that works either way. The real-model number is being worked; one question about the container at the bottom.

Meanwhile, rebased onto current dev (fb0a726 — it moved 20 commits) and fixed two no-silent-error gaps that came out of a review pass on my side. New commit 3678834, kept separate so it's clear what changed since you read this:

1. KV_TQ on a model with non-power-of-two latent rows was silent garbage. Both codecs rotate through a radix-2 FWHT, and coli_kvq_quant_row returns an inert radius 0 for any other width. On such a model every latent row would quantize to zero and the engine would generate confident nonsense with no diagnostic — the same silent-misread class the .coli_kv tier magic exists to prevent, just reached through model shape instead of file format. It now checks once after model_init and refuses, naming the shapes and pointing at KV8 (no width constraint):

[KV_TQ] this model's latent rows are kv_lora=N qk_rope=M, but the rotation needs
power-of-two widths (>=2): every row would quantize to zero. Refusing to run
quantized -- unset KV_TQ (or use KV8=1, which has no width constraint).

GLM-5.2 is 512/64, so nothing that works today changes. test_kv_tq now pins the inert behavior for all four entry points (both codecs, both dispatch paths) so the guard can't become quietly harmless if someone deletes it.

2. KV_TQ=1 clamped up to 2 — handing anyone who meant "just turn it on" the most aggressive, lowest-quality tier. It now lands on the recommended 4-bit tier and says so; >6 still clamps down to the grid.

Post-rebase verification

Re-ran everything after the rebase rather than assuming it carried:

	macOS 26.5.1 / M4 (NEON)	Linux / EPYC 7532 (AVX2)
make check	green	green, 197 tests
f32 (no KV env)	32/32	32/32
KV8=1	30/32	30/32
KV_TQ=6	30/32	30/32
KV_TQ=4	23/32	23/32
KV_TQ=1	23/32 (== TQ4)	23/32 (== TQ4)

Same numbers you reproduced, unchanged by the fixes, and still bit-identical across the two ISAs. CI 10/10 green on 3678834.

On the real-model quality datapoint

Agreed on the reasoning — glm_tiny is random-weight, its logit margins are near-tied, and 30/32 there measures perturbation, not quality cost. It shouldn't decide a decode-quality feature.

One thing before I burn the GPU-hours, so I measure the right thing: which container do you mean by gs64? The one I have locally reports quant_method: fp8, weight_block_size: [128,128] in its config.json, which reads as the upstream FP8 release rather than a colibrì group-scaled int4 conversion. Since #579 only recently made gs64 the converter default, I want to be sure I'm not reporting hellaswag on a per-row container and calling it the gs64 number — a wrong-container result is worse than no result here.

If you point me at the exact container (or confirm "convert with current dev defaults"), I'll run the #108 harness at --limit 200 for f32 / KV8 / KV_TQ=4 and post the row with Wilson intervals, same as #108 did.

[JustVugg]

Rebase request — and an apology for how many of these you have had.

dev moved a lot in the last day: 19 PRs landed, including the Vulkan backend (#418), the Kimi K3 GPU tier (#705), Metal grouped-int4 (#457), the shared routing telemetry (#716/#719), and several CUDA and launcher fixes. This PR now conflicts.

Before asking, three things changed on our side so that this is the last one of these you should need for a while.

1. The main source of these conflicts is closed. c/Makefile had a single hand-written TEST_BINS line listing every test gate. Every PR that added a test appended to that same line, so any two such PRs conflicted by construction, even when they touched entirely unrelated code — c/Makefile appeared in 26 of 40 open PRs. #386 hit it twice while being rebased, and said so, which is what sent us looking.

Gates are now derived from the build rules (#733). Adding a test means adding your .c and its own rule, which land in different places in the file. There is no shared list left to conflict on.

2. We resolved what we could ourselves instead of asking. Eight PRs were unblocked by a maintainer pushing the merge to the contributor's branch rather than requesting a rebase — including three where both sides carried a real change and had to be merged rather than picked. No commits were rewritten. Yours is here because its conflict is in engine code (c/colibri.c or a backend), where guessing your intent would be worse than asking.

3. Merge order is oldest-clean-first from now on. A PR that is green and unconflicted merges ahead of anything opened after it. The reason some of you rebased many times is that newer, smaller PRs kept jumping the queue and resetting you — a starvation loop we built, not bad luck on your side.

What we need: one rebase onto current dev. If your only conflict was the TEST_BINS line, just drop your entry — your test is picked up by its own rule now.

If you would rather not, say so and we will close it with thanks and the branch stays yours to reopen. No pressure either way, and no hard feelings — several of these have been open a while through no fault of the author.

[NeuralNotwerk]

All good, I can have an agent rebase and try again if this kind of feature hasn't made it in yet. I just moved on to slightly different projects with smaller models and this hasn't been in front of me much.

[JustVugg]

Reviving this: KV quantization just moved near the top of our priority list, and this PR is months ahead of where we thought we were.

Context: #768 measured our fp32 KV at 3.2x llama.cpp's cost on the same checkpoint (0.316 vs 0.098 MB/token on GLM-5.2), forcing a choice between long context and expert residency — 5.46 tok/s at ctx 32k vs 0.86 at 262k. Your KV8 addresses exactly that, and the KV_TQ rotated-int4 approach is the same outlier-spreading idea llama.cpp recently added for its quantized KV (a Walsh–Hadamard rotation before storage) — you were there first on our side.

Would you rebase onto current dev? It has drifted a lot since this was opened. If the conflicts are heavy, say so and we can help carve it into smaller pieces — e.g. KV8 first (the simpler win), KV_TQ second. A perplexity/token-exactness note per mode in the PR description would also help review: our bar is that anything non-bit-identical ships behind an explicit flag with a measured quality delta.

[ZacharyZcR]

Offering to carry the rebase, unless either of you would rather keep it. Context for why now: in #1140 the maintainer ranked this PR as the single highest-value unblock for the attention-bandwidth wall on resident hosts (fp32 KV is bandwidth the S=1 MLA kernel has to move, every token), and described it as stalled on time rather than on any technical doubt — and @NeuralNotwerk, you mentioned having moved on to other projects, so this looks like exactly the kind of hand-off that keeps good work from rotting.

What I'd do, following the carve-up already suggested in this thread:

1. KV8 (fp8-e4m3) first as its own PR rebased onto current dev — the simpler win, and the one KV cache is stored in fp32 and kvb_all materialises the decompressed K/V (x3.2 vs llama.cpp) #768's arithmetic prices directly (~half of the 57 GB avoidable at ctx 262144 on GLM-5.2).
2. KV_TQ (rotated-int4 / PolarQuant) second, once KV8 has survived review.
3. Your commits keep your authorship through the rebase; anything I have to write fresh to resolve drift lands as separate commits on top, clearly labelled.

One reason I'm a reasonable person to do this specifically: the maintainer flagged that three of the eleven conflicts sit inside attention_rows and are semantic rather than mechanical — and attention_rows just changed again in #1142 (the tiled kvb reconstruction), which is my change, so I already hold both sides of that merge in my head. I'd post the three semantic resolutions as a review checklist for @NeuralNotwerk (or @JustVugg) to approve before anything merges, per the plan already stated in this thread.

Say go and I'll start with KV8; say no and I'll stay out of the way.

[JustVugg]

@ZacharyZcR for me ok is important to merge this in colibri

[ZacharyZcR]

Rebase delivered: #1147. All three of your commits with authorship preserved, @NeuralNotwerk; the PROF disk-metrics commit is dropped (dev's PROF accounting was rebuilt since — it can return separately). The four semantic resolutions are itemized in the PR body for your veto, including the one new overlap that appeared since the maintainer's count: dev's #1142 tiled kvb reconstruction now shares attention_rows with your dequant staging — resolved by keeping your one-shot design under quantized KV and gating the tiled arm to f32 KV, so neither side's semantics moved. Default path verified byte-identical (tiny TF oracle 32/32 with the flags unset); KV8/KV_TQ opt-in behavior and all four kv suites pass as in your original.

[JustVugg]

Closing this because the work is in dev — and it went in as yours.

@ZacharyZcR carried the rebase in #1151 with all three of your commits intact and authorship preserved through it, listing every conflict resolution for your veto rather than quietly deciding them. Two follow-ups have since built on that base: #1155 taught the flash and gather kvb arms about KV8/KV_TQ, and #1207 added grouped latent scales (KV8_GS) following FlashMLA's fp8 cache geometry.

That is also why this PR now shows as conflicting: it collides with its own contents.

Thank you for it. KV quantization on the CPU path was the single highest-value unblock on the list — fp32 KV is bandwidth the S=1 MLA kernel has to move for every token, and it is now the base other people are extending rather than a branch waiting its turn. Five weeks between opening and landing is longer than it should have been, and that is on our sequencing, not on the work.