feat(kvcache): store the KV cache as fp8 e4m3 codes (--kv-cache-dtype… - #354
feat(kvcache): store the KV cache as fp8 e4m3 codes (--kv-cache-dtype…#354ArqAlice wants to merge 1 commit into
Conversation
… fp8)
One (token, kv head) row of K and of V becomes head_dim e4m3 codes plus ONE
fp32 symmetric scale, in a code buffer with exactly the geometry of the 16-bit
KV buffer -- only the element type changes. That halves the bytes per cached
token (the scale sidecar costs 4/head_dim of it back, ~3% at head_dim 128), and
it is what lets Qwen3.8-Flash-Next serve a 1M-token context on this card.
Codes are kept in a plain uint8 buffer on EVERY architecture, and the fp8e4nv
type never appears in a kernel signature. Both ways of choosing that per target
failed on real hardware and are recorded here so nobody reopens them: the
compile-time fp8-native probe (e4m3_compat.e4m3_native_cx) answers the question
independently from the host that allocated the buffer and disagreed with it on
sm_100, and branching on a pointer's element type is NOT statically pruned --
triton still type-checked the dead arm, whose int mask fill is illegal against
an fp8 pointer ("cannot cast int32 to fp8e4nv", raised at CUDA graph capture).
What remains is the software encode/decode that already runs wherever the fp8
type is unavailable and is bit-exact per e4m3_compat's header, so the cache
holds the same bytes and produces the same numbers on every card (docs/cli.md).
- server/args.py, engine/config.py: --kv-cache-dtype {auto,bf16,fp8}, refused at
startup for the pools and backends that cannot apply the row scales
(attention/__init__.py: BackendInfo.supports_fp8_kv) rather than ignored.
- kernel/triton/kv_quant.py: fused quantize+scatter -- one launch under CUDA
graph capture, where the slot ids arrive as a device tensor.
- kvcache: unit_bytes() counts codes plus the scale sidecar, so ft ctl stats and
cache --kv N follow the smaller footprint, and rebuild reallocates the scale
buffers alongside the codes (mha, hybrid-SWA and QSA pools).
- kvcache/base.py: pool.dtype is the COMPUTE dtype -- what store_kv receives and
what a backend sizes its scratch with -- while pool.store_dtype is what the
buffer holds. Reporting codes as dtype handed e4m3 to QSA's 16-bit indexer and
died compiling qsa_mqa_paged; the contract is now asserted at backend init and
in the kernel wrapper. QSA's block-selection keys stay 16-bit: only the
selected K/V rows are read back as codes.
Tested on: sm_100, 148 SMs, Linux; 524,480 fp8 KV tokens = 6.47 GiB,
Qwen3.8-Flash-Next with: ft serve --kv-cache-dtype fp8 -> 1M-token context.
Covered by tests/kernels/test_kv_fp8.py, tests/kernels/test_qsa_fp8.py,
tests/kernels/test_triton_attention.py, tests/kernels/test_e4m3_compat.py,
tests/kvcache/test_mha_pool_fp8.py, tests/kvcache/test_qsa_pool_fp8.py and
tests/engine/test_kv_quant_config.py (CUDA-gated; not run on the Windows
development box, which has neither triton nor pytest installed).
Not included here, on purpose: unifying the two fp8-native probes (triton's
cache-key walk rejects a constexpr function that defers to a host one, so
warn_if_probes_disagree() reports the disagreement instead), and a hardware
decode fast path on sm_89+ (that needs a constexpr flag threaded from the host
plus the matching AOT variants, since testing the dtype does not prune).
|
Thanks for building this -- a smaller KV cache is the single thing that would help this box most, I could not tell from the description whether the suite has been run anywhere yet -- the note says
and the Setup. RTX 4090 24 GB (sm_89, 128 SMs), i9-14900KF, driver 595.84 (CUDA 13.2), CUDA toolkit Group A -- fails without touching the GPU (4)Both are upstream tests from Reads like the remap indexes the scale buffers with the unmapped layer id. The scale-one path reaches the kernel with a CPU tensor. Group B -- the store kernel mis-reads a strided qkv slice (1, the interesting one)
Making the input contiguous fixes it, and swapping which slice carries the small values fixes it Worth knowing: the second assertion in that same test passes -- dequantised error is 0.035 This is the one I would guess could be arch-specific (block shape or vector width making the Group C -- a device-side assert, and why the raw failure count misleads (1)
The extra 7 are collateral. A host-side bounds check on that index would turn this into a readable Group D -- remaining (3)
And the good news: Not your bug
I have not measured throughput or context length -- with V codes wrong there is nothing worth Written with AI assistance; every number above was measured on my hardware and I can |
Environment
Result on PR's own tests42 passed, 5 failed (the 5 that actually execute the new Triton kernels): Clearing Narrowed it down (not a precision problem — stores are lost)Driving
Key observations:
The minimal driver (matrix above) is ~30 lines around |
|
Independent confirmation from the other end of the hardware range: I see the same thing on Where our two runs agree, with the test's own parameters (
The third row is the one I would add to your 2x2: swapping which slice carries the small values Two more data points from my side, one of which does not obviously fit the dropped-writes model:
One thing worth flagging for whoever picks this up: the second assertion in Happy to run your minimal driver here for a second architecture, or to test a fix -- an Ada box is Written with AI assistance; every number above was measured on my hardware and I can |
|
Thank you for the detailed follow-up, and for taking the time to verify this independently — between the two machines this now spans sm_89 and sm_120, which is exactly what whoever fixes it needs. Two things in your data move this forward. The magnitude-swap row (K small / V large on the same views → clean) is a sharper probe than my original 2x2: it helps rule out a magnitude-dependent path and narrows things toward the third slice's addressing — since with the swap K becomes the "V-like" tensor and stays correct. And the five-run table — the zero-slot count drifting 15 → 9 → 0 while junk codes grow, with the input held fixed — is what rules out a data-dependent cause and points at undefined ordering, which fits the fingerprint we saw on our side too. One point I would ask to be carried into any fix, because it plausibly explains how the serving run in the PR description could look healthy: the reconstruction assertion only checks dequantised error (0.035 < 0.08 here) and passes while the codes are wrong. Validation of a fix needs to happen at the codes level, not by perplexity-style error. A corroborating observation from our side, offered only in case it shortens the search — a hypothesis rather than a verified mechanism: we see the same split in environment, not just run-to-run. Inside the pytest process the test fails consistently (2684/3072, K clean), while an identical standalone script — same build, same seed, same tensors, codes compared against a torch-side RNE reference — ran clean in 10/10 repeats (0/3072 each). With the kernel and inputs byte-identical, the trigger appears to depend on process memory layout (allocator state when Both machines remain available to test a fix — ours covers the Blackwell end (sm_120, RTX 5090 D), yours the Ada end. We would be glad to run any candidate patch through the PR's own tests plus a service-level smoke check on our side. About this reply: like your own note, drafted with AI assistance. Every number above was measured on our hardware (RTX 5090 D, sm_120) — please take the measurements over the phrasing, as AI-assisted wording can misrepresent the intended meaning across languages. |
|
Thanks — and likewise for the careful write-up. It is genuinely useful having a Blackwell box on Your allocator-state hypothesis is worth pinning down, because on sm_89 it comes out the other way My five-run table was not from pytest. It was a standalone script:
(I just re-ran the pytest side three times to be sure of that row: 2684 every time, no drift.) So the stable-vs-drifting halves are swapped between our two machines. That does not contradict If it helps narrow it: my standalone run allocates the code buffers through the test module's own Fully agreed on validating at the codes level. For whoever picks this up, the concrete check that Ada box stays available for any candidate patch, on the PR's own tests plus the standalone matrix. Written with AI assistance; every number above was measured on my hardware (RTX 4090, sm_89) |
3 participants
… fp8)
One (token, kv head) row of K and of V becomes head_dim e4m3 codes plus ONE fp32 symmetric scale, in a code buffer with exactly the geometry of the 16-bit KV buffer -- only the element type changes. That halves the bytes per cached token (the scale sidecar costs 4/head_dim of it back, ~3% at head_dim 128), and it is what lets Qwen3.8-Flash-Next serve a 1M-token context on this card.
Codes are kept in a plain uint8 buffer on EVERY architecture, and the fp8e4nv type never appears in a kernel signature. Both ways of choosing that per target failed on real hardware and are recorded here so nobody reopens them: the compile-time fp8-native probe (e4m3_compat.e4m3_native_cx) answers the question independently from the host that allocated the buffer and disagreed with it on sm_100, and branching on a pointer's element type is NOT statically pruned -- triton still type-checked the dead arm, whose int mask fill is illegal against an fp8 pointer ("cannot cast int32 to fp8e4nv", raised at CUDA graph capture). What remains is the software encode/decode that already runs wherever the fp8 type is unavailable and is bit-exact per e4m3_compat's header, so the cache holds the same bytes and produces the same numbers on every card (docs/cli.md).
Tested on: sm_100, 148 SMs, Linux; 524,480 fp8 KV tokens = 6.47 GiB,
Qwen3.8-Flash-Next with: ft serve --kv-cache-dtype fp8 -> 1M-token context.
Covered by tests/kernels/test_kv_fp8.py, tests/kernels/test_qsa_fp8.py,
tests/kernels/test_triton_attention.py, tests/kernels/test_e4m3_compat.py,
tests/kvcache/test_mha_pool_fp8.py, tests/kvcache/test_qsa_pool_fp8.py and
tests/engine/test_kv_quant_config.py (CUDA-gated; not run on the Windows
development box, which has neither triton nor pytest installed).
Not included here, on purpose: unifying the two fp8-native probes (triton's cache-key walk rejects a constexpr function that defers to a host one, so warn_if_probes_disagree() reports the disagreement instead), and a hardware decode fast path on sm_89+ (that needs a constexpr flag threaded from the host plus the matching AOT variants, since testing the dtype does not prune).