KV8 long-context CUDA: fp8 device shadow, split-T kernels, DSA-aware gather#300
Closed
NeuralNotwerk wants to merge 2 commits into
Closed
KV8 long-context CUDA: fp8 device shadow, split-T kernels, DSA-aware gather#300NeuralNotwerk wants to merge 2 commits into
NeuralNotwerk wants to merge 2 commits into
Conversation
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>
Contributor
Author
|
Withdrawing along with #299 (stacked) — same reason: more soak testing on our fleet first. |
Contributor
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #299 (first commit is that PR; review the top commit).
The absorb kernels keep all T scores in shared memory, so CUDA attention rejects T>4096 (decode) / 8192 (batch) and long-context attention falls back to CPU — measured 22–500 ms per layer-token at 16k–262k. Under KV8 this PR removes the caps:
kv_dev_sync8— the fp8 latent KV lives in a device-resident shadow on each layer's home GPU (~149 MB/layer at 256k, vs 604 MB f32 that never fit in VRAM); decode uploads only the new rows. Lazy allocation, fail-soft to the host-upload paths.dim3(H,P)partial online-softmax blocks plus a log-sum-exp merge. A one-block-per-head tiled kernel measured 103 ms/layer at 262k (64 blocks can't fill the GPU); split-T brings it to 24.5 ms — 11–17× faster than the CPU fallback at 16k–262k. Batch shapes keep a tiled streaming kernel (S×H blocks already saturate) with the cap lifted.attention_absorb_sel_kernel8attends the selected rows straight 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 exactly like the CPU reference. The f32 GPU flow is untouched and keeps its historical dense-over-selection behavior — flagging that divergence separately in case you want it addressed for f32 too.Validation: kernel tests vs double-precision references on sm_120 and sm_89, capped-vs-shadow parity at T=3000; engine teacher-forcing at T=10,000 matches the CPU KV8 engine on 9999/10000 positions (one fp-summation-order near-tie); forced-selection generation is token-identical to dense with the gather kernel confirmed launching via nsys; a 9936-step REPLAY through the split path runs clean.
🤖 Generated with Claude Code