Skip to content

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

Open
NeuralNotwerk wants to merge 4 commits into
JustVugg:devfrom
NeuralNotwerk:cpu-kv-quant
Open

CPU: KV cache quantization — KV8 (fp8 e4m3) + KV_TQ (rotated-int4 / PolarQuant)#553
NeuralNotwerk wants to merge 4 commits into
JustVugg:devfrom
NeuralNotwerk:cpu-kv-quant

Conversation

@NeuralNotwerk

Copy link
Copy Markdown
Contributor

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 budgetingkv_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

Copy link
Copy Markdown
Contributor Author

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

Copy link
Copy Markdown
Owner

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

Copy link
Copy Markdown
Contributor Author

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
NeuralNotwerk force-pushed the cpu-kv-quant branch 2 times, most recently from c4b27ed to b98cf47 Compare July 24, 2026 16:24
@NeuralNotwerk

Copy link
Copy Markdown
Contributor Author

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

Copy link
Copy Markdown
Owner

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

Copy link
Copy Markdown
Contributor Author

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.

@NeuralNotwerk
NeuralNotwerk force-pushed the cpu-kv-quant branch 2 times, most recently from 5059f5d to e6252ca Compare July 28, 2026 11:41
@JustVugg JustVugg added enhancement New feature or request performance Velocità / tok-s / ottimizzazioni labels Jul 28, 2026
NeuralNotwerk and others added 4 commits July 28, 2026 18:03
…olarQuant)

Latent-KV cache quantization for the MLA attention, CPU path only — the CUDA
and Metal kernels are follow-up PRs, each gated by a hardware owner.

- KV8=1: fp8 e4m3 byte cache + per-row f32 scale (~3.9x less KV RAM). The CPU
  consumer LUT-decodes inline in the score/context dot products; 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_hat == 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 and
  KV_TQ_POLAR=1 use the PolarQuant codec (recursive-polar, variable bits).
- .coli_kv persistence: v2 (KV8) / v3 (KV_TQ) formats; the file magic encodes
  the tier and h[7] the TQ codec+bits. Any mode mismatch is refused and the
  cache restarted with an explicit message — never misread. v1 files upgrade
  in place (old file untouched until the first save rewrites it).
- Existing CUDA/Metal attention fast paths gain !g_kv8&&!g_tq guards: they
  read f32 rows that are not allocated under quantized KV, so quantized runs
  fall to the CPU consumer (with a one-time stderr notice under CUDA).
- Memory budgeting (kv_pool_bytes) accounts the quantized row widths, so
  cap_for_ram/PIN recover the KV savings for the expert tiers.
- Tests: test_kv_fp8 (exhaustive e4m3 roundtrip + RNE), test_kv_tq (rotation
  orthogonality, roundtrip, distortion, inert rows), test_kv_disk (v1->v2/v3
  upgrade + reject + self-heal), test_kv_alloc (f32<->KV8 transitions).

Tiny-oracle validation (tools/make_glm_oracle.py, seed 1234): teacher-forcing
f32 32/32 (exact), KV8 30/32, TQ6 30/32, TQ4 23/32 — flips are identical
across CPU-x86/CPU-ARM (and the CUDA/Metal implementations of the follow-ups),
i.e. deterministic quantization loss, not backend divergence.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Opt-in (PROF=1) instrumentation for expert-weight disk I/O rates, previously
reported only as seconds:

- pin_load: '[PROF] pin load: X GB read in Ys = Z GB/s aggregate | N experts
  @ M experts/s' — the initial hot-expert load rate off disk at startup.
- prof_report: '[PROF] disk stream: E experts/s | A GB/s aggregate over the
  phase (Cx avg read concurrency, P GB/s per loader thread)' — the live
  streaming rate during prefill/decode.

Additive only; with PROF unset every mode's output stays byte-identical.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Two no-silent-error gaps in the KV_TQ env handling, both found in review.

1. Power-of-two guard. Both codecs rotate through a radix-2 FWHT, and
   coli_kvq_quant_row returns an inert radius 0 for any other width. On a model
   whose kv_lora/qk_rope are not powers of two, that meant EVERY latent row
   quantized to zero and the engine generated confident garbage 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. Now checked
   once after model_init and refused with the shapes named and KV8 (no width
   constraint) suggested. GLM-5.2 is 512/64, so nothing that works today
   changes; this only fires where the codec cannot represent the model.

   test_kv_tq pins the underlying behavior for all four entry points (both
   codecs, both dispatch paths) so the guard cannot become quietly harmless.

2. KV_TQ=1 clamped UP to 2, handing "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.

Unchanged: default (no KV env) is 32/32 token-exact, KV8 30/32, KV_TQ=4 23/32
on the tiny oracle -- identical to the pre-change numbers.
…e_row: refuse negative scale

A power-of-two width >512 (e.g. qk_rope=1024) passed the pow2 gate but
overflows the 512-wide qtl/qtr and Lf/Rf stack staging buffers in the
attention consumers -- same refuse-don't-corrupt rule, checked once at
startup. And kv8_sanitize_row accepted a corrupt *negative* row scale,
which sign-flips the whole row on load; the encoder only emits scales
> 0, so anything else is now an inert row, matching the TQ twin.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request performance Velocità / tok-s / ottimizzazioni

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants