KV cache quantization: fp8 (KV8) + 4-bit TurboQuant (KV_TQ) on CPU, CUDA, and Metal#399
KV cache quantization: fp8 (KV8) + 4-bit TurboQuant (KV_TQ) on CPU, CUDA, and Metal#399NeuralNotwerk wants to merge 8 commits into
Conversation
dc2bf2c to
613236f
Compare
Release v1.0.0
release: fix the Windows job shell in release.yml (v1.0.0 tag build)
|
Reviewed against post-v1.0.0 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 |
Store the MLA latent rows (Lc kv_lora + Rc qk_rope) as fp8 e4m3 bytes with a per-row f32 amax/448 scale instead of f32: ~3.9x less KV RAM (a 256k slot drops 47.7 -> ~13 GB), and kv_pool_bytes reports the smaller pool so cap_for_ram/expert_avail stop demoting experts to disk under multi-slot KV. Same regime as DeepSeek-V3's fp8 latent-KV practice: the latent is both key and value in the absorbed attention; per-row scaling keeps e4m3 safe. - c/kv_fp8.h: 256-entry decode LUT + RNE bit-math encoder (saturating at +/-448, signed zero preserved, NaN stored as inert 0), per-row quantizer with a subnormal-amax guard (448/amax would overflow to +inf). - attention_rows: producer quantizes the fresh normed/roped row in place of the f32 memcpy; absorb score/context loops LUT-dequant inline with the scale hoisted out of the dot; the non-absorb prefill dequants Lc into an f32 staging row for the kv_b matmul and reads Rc via the LUT. - CUDA: attention_absorb_kernel8/_batch_kernel8 (hw fp8 cvt, sm_89+) behind coli_cuda_attention_absorb8/_absorb_batch8/_project_batch8 — 1/4 the PCIe traffic of the f32 uploads; wired at all three COLI_CUDA_ATTN call sites and into the Windows DLL loader. COLI_CUDA_PIPE and Metal still read f32 rows, so KV8 forces the pipe off and auto-disables under Metal. - .coli_kv v2 (COLIKV2 magic, dtype in the header) through the persistent- handle + staging-buffer append path: fp8 rows + scales at ~46 KB/token; v1 files quantize on resume and stay intact on disk until the first append rewrites them as v2 (magic self-heal in kv_disk_open), so a crash before the first save loses nothing. - Tests: exhaustive e4m3 roundtrip + RNE ties incl. the binade-boundary mantissa carry, KV8 kv_alloc transitions, v1/v2 disk round-trip/upgrade/ reject/self-heal, CUDA fp8 kernel checks vs dequantized-host reference. Validation: tiny-oracle KV8=0 bit-identical (TF 32/32, gen 20/20, CPU and CUDA); KV8=1 flips only two positions that transformers' own f32 forward flips across versions (near-tie margins on a degenerate random-weight model); CPU and CUDA fp8 paths agree token-for-token; ASan/UBSan clean. Real-model A/B (GLM-5.2 744B int4, 3x ~450-token log-lik requests): ppl deltas +2.4% / +0.6% / -3.7% — net ~0, inside the int4 noise floor. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The absorb kernels kept all T scores in shared memory, so CUDA attention rejected T>4096 (decode) / 8192 (batch) and long-context attention fell back to CPU (22-500 ms per layer-token at 16k-262k, measured). Three legs remove the cap for KV8: - kv_dev_sync8: the fp8 latent KV lives in a device-resident shadow on the layer's home GPU (~149 MB/layer at 256k vs 604 f32 that never fit) and decode uploads only the new rows; lazy, fail-soft to the host paths. - split-T (flash-decoding style): dim3(H,P) partial online-softmax blocks (cl[K], m, z) + a log-sum-exp merge kernel. A one-block-per-head tiled kernel measured 103 ms/layer at 262k (64 blocks cannot fill the GPU); split-T brings it to 24.5 ms — 11-17x faster than the CPU fallback at 16k-262k. The batch/prefill path keeps the tiled streaming kernel (S*H blocks already saturate) with the cap lifted to 1M. - DSA gather: attention_absorb_sel_kernel8 attends the selected rows from the shadow (absolute indices, NS<=4096). Under KV8 the batch branches are skipped when a selection is active, so GPU decode RESPECTS the DSA selection like the CPU reference does (the f32 GPU flow is untouched and keeps its historical dense-over-selection behavior). Validation: cuda-test streaming/shadow/gather vs double references on sm_120 + sm_89, capped-vs-shadow parity at T=3000; engine TF at T=10000 matches the CPU KV8 engine on 9999/10000 positions (one fp-order near-tie); DSA_FORCE generation is token-identical to dense with 95 gather launches; 9936-step REPLAY through the split path clean at 271.8 tok/s. Stacked on the KV8 base commit (fp8 KV cache + .coli_kv v2). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
First prod boot with KV8=1 + CUDA_EXPERT_GB=auto OOM-disabled dense tensors to CPU mid-prefill: the fp8 device shadow allocates lazily on each layer's kv_b home device (~CTX*584 B/layer, ~2.4 GB/device at 256k) AFTER the auto budget measured free VRAM, blowing the fixed 2 GB/device reserve. Three fixes so auto is correct by construction: - pin_load's auto budget projects the shadow per device exactly like g_cuda_dense_projected (est_ctx from CTX, kv_b home from qt_load's round-robin; sharded/ineligible kv_b layers never shadow and are skipped). - KV_SHADOW policy: the shadow only pays where decode uses the contiguous CUDA attention path (kvs==NULL). Multi-slot mux decode is ragged -> always CPU, and prefill uploads the same bytes either way, so the shadow would cost VRAM for nothing there. AUTO: on iff KV_SLOTS==1; KV_SHADOW=0/1 forces; the [KV8] banner states the decision. - kv_dev_sync8 partial-alloc hazard: if the 2nd-4th of the four shadow buffers failed to allocate, the next call saw L8 non-NULL and uploaded to NULL scale buffers. Failure now frees and NULLs all four. Validated: tiny-oracle 30/32 in all three modes (shadow ON, auto-off under KV_SLOTS=4, forced off with host-upload fallback); 10k-token TF matches the CPU engine 9999/10000 through BOTH the shadow and host-upload streaming paths; make test-c + cuda-test green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Review findings on 2143450, both confirmed and fixed: - The auto-off predicate was kv_slot_count()==1, but plain run_serve with KV_SLOTS>1 decodes CONTIGUOUSLY (step/spec_decode, kvs==NULL) — the CUDA attention path is live there and killing the shadow would re-upload the whole fp8 history per layer per token. Only the multi-slot MUX (SERVE_BATCH) has ragged decode; the policy now keys on that. run_serve multi-slot and single-slot mux (where MTP verify forwards are contiguous) keep the shadow. - pin_load's projection reserves per CTX, but PROMPT mode sizes max_t from the prompt and never reads CTX, so a 200k-token PROMPT could still blow the reserve. kv_dev_sync8 now checks free VRAM at the point of allocation and declines the shadow unless >=1.5 GB remains after it — the host-upload fallback takes over. This guards every mode, present and future, not just the ones the projection can predict. Policy matrix validated on the tiny model (plain / serve x slots x mux x forced); test-c and cuda-test green; oracle unchanged (30/32 fp8). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Correctness (from the fp8 KV review + claims-ledger audit): - shadow uploads: coli_cuda_pipe_upload_async (pinned ring, cudaMemcpyAsync on the kernel stream) replaces 4 blocking NULL-stream memcpys per layer per token in kv_dev_sync/sync8 — stream-ordered with the consuming kernels, so the documented pageable-memcpy race on the newest rows is gone - fmt=4 grouped-int4 tensors are no longer cuda_eligible (row_bytes()=0 made every upload fail, silently killing the whole KV8 GPU path and leaking projected VRAM budget); layer_cuda_shard_kvb skips fmt=4 - spec-verify batches (S=5..64, g_spec_live) now compute and consume DSA selection like decode instead of attending dense — drafts and verify use the same attention function (the JustVugg#163 acceptance-collapse class) - absorb_forced in expert_avail/cap_for_ram mirrors the runtime gate (CUDA build + g_cuda_enabled + kv_lora<=512), not the raw env var: a CPU-only run with COLI_CUDA_ATTN=1 no longer drops the ~30 GB kvb_all reserve it will later allocate - g_scratch_failed is atomic and cleared at coli_cuda_matmul entry: stale scratch-OOMs from attention paths can't reclassify a later hard matmul error as transient - fused o_proj + expert-MLP quant_matmul launches chunk S (grid-Y max 65535): PREFILL_CHUNK=0 prompts >65k rows no longer fail after the absorb ran - .coli_kv: v2 load sanitizes e4m3 NaN codes (CPU lut reads 0, GPU cvt read NaN) and non-finite scales; kv_disk_truncate clamps nrec to the records physically written (append-skip + truncate could declare a gap of garbage); appends fsync data before publishing the counter (power-loss ordering) - chunked prefill zeroes the straddled MTP KV row per chunk boundary: mtp_draft read malloc-uninitialized memory there under DRAFT>0 - kv_dev_sync8 caches the 1.5 GB floor rejection (retry every ~1024 calls) instead of a synchronizing cudaMemGetInfo per layer per token Performance (ncu: stream kernel was L1-pipe-bound 97.5%, SM 19%, DRAM ~0 — one load instruction per fp8 byte): - fp8_row_dot: uint4 vector loads for every score dot; context accumulate reworked to coalesced uchar4 k-quads per thread (was per-byte column loads at stride K); per-row softmax weight*scale hoisted out of the k loop - single-row and sel kernels: softmax on thread 0 -> tree reduction (red[]); sel path with NS>=1024 splits the list dim3(H,P) into online-softmax partials + existing merge (was 64 blocks on a 170-SM GPU, 16.7% occupancy) - reserve()/reserve_bytes() round growth up 25%: no more per-token cudaFree+cudaMalloc churn as lb/rb grow monotonically each decode token - prefill KV producer loop (rmsnorm+rope+quant) parallelized with OMP (kv_dev_valid shrink hoisted out as a serial min-pass) Measured (RTX 5090, GLM dims H=64 K=512, int4 kv_b): project_batch_kvdev8 S=2048 T=32k: 25336 -> 1602 ms/call (15.8x) absorb_kvdev8 T=32k/128k: 4.04 -> 0.96 / 14.23 -> 3.25 ms (4.2-4.4x) absorb_kvdev8_sel nsel=2048/4096: 0.78 -> 0.21 / 1.47 -> 0.26 ms (3.7-5.7x) identical fixture outputs; cuda-test parity green (incl. new small-NS and split-NS sel cases); test-c green Docs: KV_SHADOW row + code comment now state the real always-on-auto policy (57204d0) and the multi-slot rebuild cost; README KV8 numbers put both endpoints on the same basis (47.7 -> 12.1 GB latent pool, ~13 with the f32 DSA indexer), notes the 1M-token fp8 GPU cap and the lazy v1->v2 rewrite; kv_fp8.h documents the NaN-mixed-row behavior. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
CPU (kv_tq.h + glm.c): - PolarQuant codec (randomized-Hadamard rotation + recursive polar) and a rotated-int4 + Lloyd-codebook codec (the competitive 4-bit tier, full reconstruction so it fixes the MLA key AND value). coli_kvq_* dispatch; KV_TQ=4 defaults to int4, KV_TQ_POLAR=1 or other bit widths use polar. - KV_TQ env, producer/consumer, alloc, budget, .coli_kv v3 (COLIKV3, h[7]=codec<<8|bits). - tests: test_kv_tq.c (host), test_kv_disk.c v3 round-trip. Metal (backend_metal.mm/.h): - fp8 kernels a_fp8enc (RNE, byte-exact vs coli_fp8_enc) / a_score8 / a_clat8. - TQ kernels a_tq_enc / a_tq_deq with a codec arg (polar or rotated-int4). - encode_attention is 3-way qmode (f32 / KV8 / KV_TQ); entry points coli_metal_attn_decode8 / _tq, and coli_metal_layer_decode gains qmode. - Metal is now opt-OUT (default-on when a device is present; COLI_METAL=0 forces CPU). CUDA stays opt-in (discrete VRAM can OOM). - metal-test: fp8/TQ/int4 attention parity + run_layer full-layer harness. Validated: make test-c green; metal-test all-green (fp8 byte-exact, TQ ~3e-8 roundtrip, ~1e-6 attention parity). Live on GLM-5.2 744B int4 (Metal, greedy): f32, KV8, and KV_TQ=4 int4 all run the fused attention on the GPU (0 CPU fallbacks) and produce coherent output; Metal output == CPU output. Tiny-oracle: f32 32/32, KV8 30/32, KV_TQ=4 (int4) 23/32. KV_TQ=4 + MTP (DRAFT=2) on the real model: coherent output, no crash (MTP acceptance is low under int4 KV, but speculative decode always emits the verified token, so output correctness is unaffected). Hardening (found by auditing every g_kv8 site for a missing g_tq sibling — the same "quant arrays are NULL under f32" NULL-deref class): - layer_decode must NOT index m->Lc8/Rc8/Lsc/Rsc under f32 (those arrays are NULL unless KV8/KV_TQ is on) — pass NULL there. (Regression found by diffing against vanilla on the real model; metal-test's f32 harness passed because it hard-codes nullptr for those args.) - step_decode_batch ragged-mux validity check gated on g_kv8 only -> rejected every KV_TQ row in multi-slot serve; now (g_kv8||g_tq). - kv_mtp_row_zero had no g_tq branch -> the MTP straddle row was left uninitialized under KV_TQ; now zeros the packed row + sets radius 0 (coli_kvq_dequant_row then yields the zero vector). - CUDA per-row absorb branch (COLI_CUDA_ATTN) was gated on g_cuda_enabled but not !g_tq, so KV_TQ on a CUDA box would fall to the f32 arm and read NULL Lc; now !g_tq (matches cuda_absorb, which already excluded TQ). TQ stays CPU/Metal-only in phase 1. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Rebased onto current The ragged × fp8-shadow call, made explicit in the code (comment at the gate in Also folded in during the rebase: the batch-absorb gate keeps its |
b02bc92 to
606fe19
Compare
|
Heads-up: |
Upstream PR JustVugg#399 (KV8 fp8 / TQ4 quantized latent KV) leaves Lc/Rc NULL when a quantized tier is active — the mirror sync would deref NULL. Guard on the arrays themselves (env-independent), falling back to the CPU attention path; an fp8-aware VK mirror (upload Lc8+scale, decode e4m3 in the shader, 4x less mirror traffic) is the follow-up once JustVugg#399 merges.
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 Correctness, live on GLM-5.2 744B ( Performance — measured at the production operating point (serve mode over HTTP, exact prod env: On the #391 heads-up ( |
|
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.
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. |
3 participants
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).COLI_METAL=0to 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)
.coli_kvv2 — 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.kv_tq.hcodecs (randomized-Hadamard rotation + recursive polar / rotated-int4 + Lloyd codebook), MSL kernelsa_fp8enc/a_score8/a_clat8/a_tq_enc/a_tq_deq, 3-wayqmodein the fused layer decode,.coli_kvv3 (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
make test-cgreen (new suites:test_kv_fp8exhaustive e4m3 roundtrip + RNE;test_kv_tqrotation self-inverse, roundtrip, distortion, inert rows;test_kv_diskv1/v2/v3 round-trip + upgrade/reject/self-heal;test_kv_alloc).make metal-testgreen — fp8 GPU encoder byte-exact vscoli_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).DRAFT=2decodes 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).!g_tqguards 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
KV8/KV_TQunset → byte-identical behavior; the quantized byte-caches are NULL and every consumer branches ong_kv8||g_tq).COLI_METAL=0) and non-Metal builds behave exactly as the CPU tier.🤖 Generated with Claude Code