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 4 commits 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) |
|
Good news — I think this one is solved. The GPU freed up here, so I put What found itThe other two tools were informative by staying quiet: The bug
_kv_quant_scatter_kernel[(tokens, heads)](k, v, ..., k.stride(0), ...) # K's pitch only# kernel, line 108
src = t * stride_xs + h * D + d
xk = tl.load(k_src + src, ...) # K -- correct, this is K's pitch
xv = tl.load(v_src + src, ...) # V -- reads V with K's pitchThe guard at line 153 checks only the inner stride: assert k.stride(1) == 1 and v.stride(1) == 1, "K/V rows must be contiguous"It never asserts The test breaks that contract at k = k_rows.view(tokens, heads, dim) # view -> stride(0) = 1152
v = v_rows.view(tokens, heads, dim).clamp(-FP8_MAX, FP8_MAX) # clamp -> stride(0) = 384
That is why it is V only and K byte-perfect: K's pitch is 1152. Why both of our boxes were rightThis is the part I think resolves the ordering question. Splitting the mismatches by region, same
The deterministic half is exactly 764 every time — tokens 0–2, where the wrong address still So @Kaempferia's allocator hypothesis was half right, and the better half: the values really do FixPass V's pitch as its own argument and load each tensor with its own: idx_ptr,
- stride_xs, # source row pitch, in elements (the qkv slice is wider than one row)
+ stride_xs, # K source row pitch, in elements (the qkv slice is wider than one row)
+ stride_vx, # V source row pitch: K and V need not share one. A .clamp()/.contiguous()
+ # on one side alone leaves it densely packed while the other keeps the
+ # qkv pitch, and reusing K's pitch then reads V off the end of its rows.
stride_kd, # K cache row pitch, in elements (== HEADS * D)
@@
- src = t * stride_xs + h * D + d
- xk = tl.load(k_src + src, mask=mask, other=0.0).to(tl.float32)
- xv = tl.load(v_src + src, mask=mask, other=0.0).to(tl.float32)
+ off = h * D + d
+ xk = tl.load(k_src + t * stride_xs + off, mask=mask, other=0.0).to(tl.float32)
+ xv = tl.load(v_src + t * stride_vx + off, mask=mask, other=0.0).to(tl.float32)
@@
k.stride(0),
+ v.stride(0),
k_cache.stride(0),Verification on sm_89, each file in its own process, across the seven
The remaining 11 are the other items from my first comment (the How bad is itLatent, not a live serving bug — which I think is the honest reading, and it also explains why The models that use fp8 KV go But the contract is invisible, and the failure mode is silent. Any future backend that materialises And to reinforce @Kaempferia's point, because this is the sharp edge: the test's second Happy to open this as a PR against your branch, or leave it here for you to take — whichever you Written with AI assistance; every number above was measured on my hardware (RTX 4090, sm_89) |
|
Excellent work — compute-sanitizer naming the line on the first run, and the explanation that one shared pitch reads both tensors fits every data point we exchanged. The "memcheck is blind because the caching allocator rounds up and the bad read stays inside the pooled segment" note is worth keeping for anyone else debugging a Triton kernel — that is a genuinely non-obvious failure mode. Independent confirmation from the Blackwell end, with your patch applied to our tree (the fp8 PR head cherry-picked onto the PLE-disk line, same environment as before — sm_120, RTX 5090 D, torch 2.11.0+cu130, triton 3.6.0):
On your allocator-hypothesis verdict — agreed, "half right, and the better half": the addressing is fixed and wrong, only the contents of the uninitialised region were layout-dependent. That also cleanly explains why our standalone script came out 10/10 clean on sm_120: luck, with different allocator histories buying different luck. It was never two bugs. And I fully agree the sharper edge is the reconstruction assertion passing at 0.035 — worth re-stating in the PR that codes-level checking is the only thing that catches this class. One practical point on where the fix should land: our tree is a local branch of a clone — there is no GitHub fork or branch behind it that a PR could target. The right home for this fix is the PR itself: if you push it to Thank you for chasing this to the actual line — that is the kind of root-cause work worth recording. Ada and Blackwell both stand ready to test anything else you want run. |
1 similar comment
|
Excellent work — compute-sanitizer naming the line on the first run, and the explanation that one shared pitch reads both tensors fits every data point we exchanged. The "memcheck is blind because the caching allocator rounds up and the bad read stays inside the pooled segment" note is worth keeping for anyone else debugging a Triton kernel — that is a genuinely non-obvious failure mode. Independent confirmation from the Blackwell end, with your patch applied to our tree (the fp8 PR head cherry-picked onto the PLE-disk line, same environment as before — sm_120, RTX 5090 D, torch 2.11.0+cu130, triton 3.6.0):
On your allocator-hypothesis verdict — agreed, "half right, and the better half": the addressing is fixed and wrong, only the contents of the uninitialised region were layout-dependent. That also cleanly explains why our standalone script came out 10/10 clean on sm_120: luck, with different allocator histories buying different luck. It was never two bugs. And I fully agree the sharper edge is the reconstruction assertion passing at 0.035 — worth re-stating in the PR that codes-level checking is the only thing that catches this class. One practical point on where the fix should land: our tree is a local branch of a clone — there is no GitHub fork or branch behind it that a PR could target. The right home for this fix is the PR itself: if you push it to Thank you for chasing this to the actual line — that is the kind of root-cause work worth recording. Ada and Blackwell both stand ready to test anything else you want run. |
|
Thanks for running it on the Blackwell end — same single flip, same multiset property, from a On where the fix should live: you are right that the PR is the natural home. Since git fetch https://github.com/MT-z/FreeToken.git fix/kv-fp8-vstore-pitch
git cherry-pick FETCH_HEAD
@ArqAlice — this is yours to take or leave. It is a fix to your PR, not a competing one, and I The commit message carries the full derivation rather than just the diff — the read offsets per I also put your point about the reconstruction assertion in it, since it is the part most likely to
For completeness, the numbers on this side against the real PR head The other items from that first comment — the Written with AI assistance; every number above was measured on my hardware (RTX 4090, sm_89) |
|
@MT-z san I'd be very happy to accept your fix. If you don't mind, could you open a PR against my feat/fp8-quantization branch? I think that would be the cleanest way to preserve your contribution and the investigation behind it. Once it's up, I'll review and merge it into #354. Thanks again for all the testing and debugging — especially for verifying this on sm_89. It's been extremely helpful. |
`quantize_kv_to_cache` passed `k.stride(0)` as the only source pitch and
`_kv_quant_scatter_kernel` used it for both tensors:
src = t * stride_xs + h * D + d
xk = tl.load(k_src + src, ...)
xv = tl.load(v_src + src, ...)
The guard above it checks only the inner stride (`k.stride(1) == 1 and
v.stride(1) == 1`), never `k.stride(0) == v.stride(0)`, so the kernel carries an
undocumented contract: K and V must share one row pitch.
When they do not, V is read at K's pitch. In the failing test K is a view of the
qkv slice (pitch 1152) while V is materialised by `.clamp()` (pitch 384), so with
8 tokens of 3072 elements:
token 0 reads 0 correct by coincidence
token 1-2 reads 1152, 2304 in range, WRONG rows
token 3-7 reads 3456 .. 8064 past the initialised data
2684 of 3072 codes wrong, all in V, K byte-perfect. Deterministic addressing;
only the contents of the uninitialised tail vary with allocator history, which is
why the mismatch count drifts (2684 / 2663 / 2676 across runs) while the
mismatching positions do not -- the in-range half is exactly 764 every time.
Found with `compute-sanitizer --tool initcheck` (TRITON_DISABLE_LINE_INFO=0),
which named `kv_quant.py:110`. `memcheck` reports 0 errors because PyTorch's
caching allocator rounds allocations up and the bad read stays inside the pooled
segment; `racecheck` reports 0 hazards because it is not a race.
Fix: pass `v.stride(0)` as its own kernel argument and load each tensor with its
own pitch.
Verified on RTX 4090 (sm_89): tests/kernels/test_kv_fp8.py 2 failed -> 1 failed,
the flip being test_codes_match_the_reference_quantizer_and_reconstruction_is_close;
five consecutive standalone runs give K 0/3072 and V 0/3072 with got.sort() ==
exp.sort(); compute-sanitizer initcheck reports 0 errors on the patched build.
Independently confirmed on RTX 5090 D (sm_120) by @Kaempferia: same single flip,
same multiset property, 5 runs clean.
Note for reviewers: the test's SECOND assertion (dequantised error <= 0.08)
passes at 0.035 while V is wrong, so a reconstruction-level check does not catch
this class. Only the exact-code assertion does.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
It is up: ArqAlice#1 -- one commit, one file, +7/-4. Before opening it I re-ran the counts on your branch head (
I also shortened the new kernel comment from three lines to two, having noticed the "one or two One thing I noticed while testing and have not mentioned yet. On this box the KV/expert budget The Ada box stays available for anything you want run on sm_89, any time. Assisted-by: Claude Opus 5 |
fix(kernels): give the V tensor its own row pitch in the fp8 KV store
|
Ran this on 2 x RTX 6000 Ada (sm_89, 48 GB each, PCIe 4 x16, no NVLink) with Qwen3.8-Flash-Next (RadixArk NVFP4, modelopt), TP=2 over both cards, offload backend, pinned PLE, vision tower loaded, load-time per-tensor FP8 dense (#389). Tree: main af71ba4 + TP (#385) + #389 + #392 + this PR at 3e5bbdd (which already carries MT-z's V-pitch fix), merged with one docstring conflict. torch 2.11.0+cu130, triton 3.6.0 (the pins). Tests, one process per file on one card:
The same list MT-z and Kaempferia reported, plus the emulation one. Nothing in the failing set is on the path this model serves through (QSA pool + qsa_sparse), and the serving numbers below agree with that. Serving. Same harness as my other reports: single-stream = median of three 256-token generations, aggregate = eight concurrent 256-token requests, TTFT on a ~1k-token prompt, residency = expert slots / 24,576, probe = an 8-question greedy smoke test (thinking off). The cache status route reports 13,056 B per token per rank in bf16 and 7,009 B in fp8 for this model (12 QSA layers, 2 local KV heads, the bf16 index tier included).
(*) Client-side numbers of the first fp8 process: your Triton kernels compiled during that benchmark (the Triton cache shows the entries landing in the bench window), while the server's own per-batch decode rates in the same run were 94-101 tok/s at one running request and ~448 at eight, the same as every 262k-pool run. The same flags with the kernels cached give the production row. So on this box the fp8 attend kernel costs nothing measurable at an equal pool (the 262k fp8 row also had two compiles inside its bench; its server-side rates were 92-101 vs 99-101), and the halved KV is what makes eight full contexts free: 70.9% residency decodes like 95.8%, where bf16's 33.7% loses 15% single-stream and 42% at 8 concurrent. Long-context checks on the fp8 8-context run, all greedy: a planted-fact transcript recalled at 105,683 and 214,673 prompt tokens (same 3/4 as bf16; the fourth is my scorer); 8 concurrent users with 177k-token prompts each all got their planted id back (1.42M prompt tokens in flight, 389 s wall for the eight prefills), and a second round on the same eight conversations hit the full cached prefix for every user (cached 177,472 per user in the log) and decoded at 324 tok/s aggregate; image prompts keep their prefix-cache hits (320 / 26,304 cached tokens, no false hit on a different image). Greedy text against the bf16 production run: the fp8 runs diverge after 37 / 29 / 38 words on the three prompts (the bf16 run at the 8x pool is word-identical to production on all three, so that is the cache quantization, not run-to-run noise). This is now the production configuration on that box (fp8, 8 x 262,144 tokens, 8 running). |
…ike text 037f102 narrowed the rule from "the whole prompt in one chunk" to "the image span in one chunk", which is what a 196-token sprite in a 166k-token turn needs. The span is [first image token, last+1) because ``mm_embeds`` is one concatenated tensor scattered in one forward -- so it grows with the TEXT between two screenshots, not just with the pictures. An agent conversation reaches the limit by talking: 400 prompt with images needs 10392 contiguous tokens in one prefill chunk (the image tokens span [160334, 170726) and cannot be split) Nothing configurable moves that. Cheaper images (~490 tokens each after the clamp) only buy more turns before the gap between the first and last one exceeds a chunk, and raising --max-prefill-length OOMs long before it helps: a 32k chunk's activations do not fit beside a 5 GiB KV pool on a 24 GiB card (measured -- it took the worker down twice today). So the concatenated tensor stops being scattered whole. ``_merge_multimodal`` takes the rows belonging to the placeholders inside ITS OWN forward -- the ones an earlier chunk or a prefix-cache hit already consumed sit in front of the window -- and the adder chunks an image prompt exactly like a text one. ``Req.mm_scatter`` and the whole pull-back / reject path go away with it, ~90 lines. Both families that carry a tower here are converted; the approach is gdevenyi's, from FlashML-org#386 (28fd56d). The span cap 09ea814 put in ``match_req`` goes too. It existed because a hit landing inside a placeholder run left half the run cached and half to forward, which the all-in-one-forward scatter could not represent; the window skips the cached half instead. Without the cap a prompt that ends with its image keeps its prefix -- 20,800 of 20,840 tokens on the repeat here, 6.0 s -> 1.2 s, and a different image at the same position still misses (answered "Green" where the cached one answers "Blue"). Measured on Ornith-1.5-35B-A3B-NVFP4, one 4090, --max-prefill-length left at its 8192 default: 2 images with 9k of text between span ~19k 10,186 tokens, 3.2 s (was a 400) 6 images with 9k between each span ~50k 55,360 tokens, 18.1 s (was a 400) A(blue) 9k B(green), and reversed "Blue, Green" / "green blue" -- read across the boundary, in order tests/tokenizer 58, tests/scheduler 90, tests/kvcache/radix 142: all passed. Twelve tests pinning the removed rule are gone and three cover the window (a span wider than a chunk now admits; a chunk scatters only its own rows; a chunk holding no placeholder scatters nothing). The ``_NoSwa`` stub gained the ``page_size`` the reservation math has been reading, which is what had six of these failing on this branch already. A cold system-test run is character-identical to the same branch without this commit, all seven cases. Assisted-by: Claude Opus 5 Re-verified on this branch (no FlashML-org#337/FlashML-org#354/FlashML-org#287 under it): tests/tokenizer 58, tests/scheduler 88, tests/kvcache/radix 142 all passed; a cold system-test run is character-identical to the same change on the daily branch, all seven cases; the two shapes that used to 400 (spans of ~19k and ~50k tokens) answer at the 8192 default.
|
Thank you for the table -- it sent me back to my own report, which was not wrong but was Repeating the serving numbers on one 24 GiB cardOrnith-1.5-35B-A3B-NVFP4, offload backend, single RTX 4090 (sm_89), i9-14900KF, torch
Same conclusion as yours, from the other end of the hardware range. The attend kernel costs Following the four diagnoses
Each of the four you named is one line. I measured them one at a time, so each accounts for a
One line in the report moved that I had not prescribed: Two details that were not in your table, in case they are useful. The write at row 256 is
They sit on small outputs, and the 1.33 relative figure is an element whose reference is The four came straight out of your table; I am in your debt for it, and for the patience Assisted-by: Claude Opus 5 |
4 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).