CPU: KV cache quantization — KV8 (fp8 e4m3) + KV_TQ (rotated-int4 / PolarQuant)#553
CPU: KV cache quantization — KV8 (fp8 e4m3) + KV_TQ (rotated-int4 / PolarQuant)#553NeuralNotwerk wants to merge 2 commits into
Conversation
|
CPU performance validation for this branch, measured on its exact tree (744B, fully resident —
Prefill ~5.5–5.6 tok/s across all cells. Two takeaways:
Environment: EPYC-class 32C host, ZFS storage, dockerized gcc 13 build, |
|
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 Two things to confirm on the rebase, given the 'no silent errors' bar here: (1) the producer-loop parallelization in |
…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>
|
Rebased onto current On your two points:
Also verified through the rebase: your new Post-rebase validation: CPU + |
bf57e2d to
c4b27ed
Compare
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=1use the PolarQuant codec..coli_kvacross mode changes — the file magic encodes the tier (COLIKV/COLIKV2/COLIKV3) andh[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 bytests/test_kv_disk(upgrade, both reject directions, TQ-params reject, append-side self-heal).kv_pool_bytesaccounts the quantized row widths, socap_for_ram/PIN recover the KV savings for the expert tiers.!g_kv8&&!g_tqand 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.PROF=1disk-load throughput lines (pin-load GB/s + experts/s, live streaming rate). Byte-identical output withPROFunset.Validation (reproducible on your machine)
make test-cgreen (incl. the four new kv suites); builds clean withCUDA=1andMETAL=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