Repository navigation
fix(qwen4-exp): register the QSA FP8 scale planes for KV offload - #2655
Conversation
Decode-sized QSA index scoring (the target's MTP verify and the draft, at most 256 rows) ran on the CUDA-core kernel, which reads every request's compressed keys once per query row. The new Triton MFMA kernel stacks 4 rows x 4 index heads into one M = 16 operand, so a request's keys are read once per verify and meet every head in a single dot; the per-head ReLU is applied before the heads are summed back out of M. On Qwen3.8-Flash-Next (MI355X, TP1, MTP 3) the target's decode scoring drops from 1130 to 320 us a step and the draft's from 147 to 85 us, -6.2% of decode kernel time. ATOM_QSA_STACKED_DECODE_SCORER=0 keeps the old kernel. Co-authored-by: Cursor <cursoragent@cursor.com>
Qwen QSA refused anything but a BF16 KV cache and ignored
--index-cache-dtype. Serve --kv_cache_dtype fp8 and --index-cache-dtype
{bf16,fp8,fp4} (unset follows the KV dtype, as Config resolves it):
- K/V: E4M3 with one FP32 scale per (token, KV head) = row amax / FP8 max,
written by a new Triton kernel (K and V in one launch). Dynamic scales
need no calibration and never saturate. The sparse GQA kernel reads
E4M3 tiles, multiplies the key scale into its score column and folds the
value scale into the probability before PV; FP8 tiles have their own
tuned block config.
- Compressed index keys: FP8 with a per-row FP32 scale, or MXFP4 (packed
E2M1 pairs, one E8M0 exponent per 32 dims, rounded up so a group never
clips). Every scorer -- the stacked decode one, the CUDA-core one and the
prefill MFMA one -- reads them; FP4 decodes with bit operations. The raw
index keys stay BF16: they are mean-pooled before normalization and a
group can straddle a prefill chunk.
- The MTP draft layer keeps BF16 caches in a `qsa_draft` row space of its
own (ATOM_QSA_DRAFT_KV_BF16, default 1): an FP8 draft KV lowers the real
acceptance length by ~3.5%, a BF16 one not at all.
- `_qsa_fields` declares the scale planes, so the arena sizes the pool by
bytes; the memory profile runs on caches of the served format. Plugin
hosts keep binding their own BF16 caches.
On Qwen3.8-Flash-Next (MI355X, TP1, MTP 3) FP8 KV fits 1.64-1.67x the
tokens, cuts prefill sparse GQA by 17%, and leaves GSM8K and the MTP
acceptance length unchanged. With --block-size 64 and a BF16 index cache
the prefill scorer stays on FlyDSL.
Co-authored-by: Cursor <cursoragent@cursor.com>
… into lingzha/qwen38-lmcache-fp8-kv Co-authored-by: Cursor <cursoragent@cursor.com>
build_kv_cache_tensor bound the FP8 K/V scale planes (and the quantized compressed-index scale) to the attention module but registered k_scale / v_scale as None, so lmcache_offload's DenseKVByteCodec stored and restored the FP8 rows without their scales. A prefix fetched back from CPU then decoded under whatever scales the reused blocks held: GSM8K pass 2 dropped from 0.98 to 0.50, every CPU-restored request returning empty or running to max_tokens. MTP draft rows also registered the target's index fields at the draft row. Register every paged field the arena declares, with the row space's own suffix, and move index_compressed_scale as a new codec tail. BF16 caches register the same tensors as before. Co-authored-by: Cursor <cursoragent@cursor.com>
|
No conflict with the updated base ( Verified (MI355X,
Note that the tree tested here does not include #2635. With #2635 merged, offload retrieval regresses; details are in #2636. |
HaonanWang98
left a comment
HaonanWang98
left a comment
There was a problem hiding this comment.
LGTM. Reviewed the scale-plane registration and suffix-aware target/draft bindings at a1a3abf. The registered fields match the arena bindings, and the codec includes the compressed-index scale plane. No blocking issues found in this fix.
|
@HaonanWang98 @valarLip A heads-up on how this PR relates to #2636. valarLip's review of #2636 today flagged two issues in
Both are what this PR fixes: it registers every paged field the arena declares, including the FP8 K/V scale planes and Plan: merge #2636 first, then retarget this PR to |
…#2655 The in-process codec copies only what build_kv_cache_tensor hands it: no k_scale / v_scale / index_compressed_scale, and the target's index rows for a BF16 draft layer. With FP8 KV (or an FP8/FP4 index cache, including one that follows --kv-cache-dtype) a CPU hit restores quantized rows under the scales of the blocks' previous owners; GSM8K pass 2 fell to 0.54 with no error. check_qwen38_kv_transfer now refuses lmcache_offload with any quantized QSA cache, pointing at bf16 caches or lmcache_mp, which publishes every arena field and stays allowed. #2655 carries the scale planes and the draft rows and removes this check. Addresses review 5698904199 (must change 1). Co-authored-by: Cursor <cursoragent@cursor.com>
*_draft rows exist only under quantized target caches with ATOM_QSA_DRAFT_KV_BF16 (QSAAttention.qsa_row_space). lmcache_offload now refuses that combination until #2655, which fixes the row lookup in build_kv_cache_tensor. Pin the case that stays reachable: with BF16 caches the MTP layer is one more row of the target's fields, and the KVCacheTensor the in-process codec copies names exactly the rows the layer binds. Addresses review 5698904199 (must change 2). Co-authored-by: Cursor <cursoragent@cursor.com>
Add k_scale/v_scale, index_compressed_scale and the *_draft copies, and note that lmcache_offload refuses quantized QSA caches until #2655. Co-authored-by: Cursor <cursoragent@cursor.com>
|
@valar-cop review this pr |
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
LGTM, but layout version and the duplicated field list should be fixed..
any e2e with fp8 index cache or mtp? old tier upgrade run?
ci green.
| getattr(kvt, "index_raw", None), | ||
| getattr(kvt, "index_compressed", None), | ||
| getattr(kvt, "rope_positions", None), | ||
| getattr(kvt, "index_compressed_scale", None), |
There was a problem hiding this comment.
bytes_per_block changed but namespace didn't, old tier chunks will hit 'too small'.. bump PAGE_LAYOUT_VERSION or fold segment sizes into build_page_namespace?
There was a problem hiding this comment.
@valar-cop Not bumping it, and here is why:
- BF16 QSA chunks keep their bytes; the scale planes exist only for quantized caches.
- No tree has written FP8 QSA chunks with
lmcache_offload. [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636 refuses that combination at startup, so this PR is the first to store them. main's Qwen chunks could not be read here anyway. Their digest differs because the page namespace keyskv_cache_dtype/index_cache_dtypeand [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636'sqsa_page_segments.lmcache_mphas its own::lmcache-mp-vsuffix.
A PAGE_LAYOUT_VERSION bump would invalidate every model's persisted PAGE objects. Details are in the PR comment (de10df694).
There was a problem hiding this comment.
ok, checked all four. _qsa_fields only adds the scale planes when qsa_cache_row_layout returns a scale tail, so bf16 chunks keep their bytes; the quantized cases were refused at startup in the base, and main's qwen digest differs through qsa_page_segments (atom/kv_transfer/offload/config.py:386). keeping v3 is fine, resolved on my side..
| arena.view(f.name).numel() * arena.view(f.name).element_size() // blocks | ||
| for f in fields | ||
| ) | ||
| assert codec.bytes_per_block == per_block |
There was a problem hiding this comment.
only sum checked, a swapped k_scale/v_scale would still pass.. assert each field data_ptr and per-segment sizes?
There was a problem hiding this comment.
@valar-cop Agreed: a swapped or wrong-row scale passed the sum. This file is back to its base. The coverage now lives in tests/test_qwen4_exp_offload.py::test_every_qsa_layer_hands_the_codec_the_rows_it_binds, which covers five cache formats and two target layers plus the MTP draft layer, with rope positions. For each layer, every registered tensor must be that layer's own row of the right field, compared by data_ptr in stored order, and the set must equal what the layer binds. Registering k_scale from v_scale fails 4 of the 5 cases, and every quantized case fails on the base (de10df694).
There was a problem hiding this comment.
checked, resolved on my side: test_qsa_quantized_cache.py is back to base, and test_every_qsa_layer_hands_the_codec_the_rows_it_binds does compare each field's own row by data_ptr in stored order, plus the bound set at tests/test_qwen4_exp_offload.py:501. the four quantized cases fail on c90317566.
| rope_positions=( | ||
| arena.view("rope_positions")[row] if self.cache_rope_positions else None | ||
| ), | ||
| k_scale=optional("k_scale"), |
There was a problem hiding this comment.
same field list twice, that drift is what caused this.. build the views once and use for both?
There was a problem hiding this comment.
@valar-cop Done in de10df694. build_kv_cache_tensor now looks each field up once, and bind_caches and KVCacheTensor get the same views.
There was a problem hiding this comment.
yes, looks right. one lookup at atom/model_ops/attentions/qwen4_exp_attn.py:572 now feeds both bind_caches and KVCacheTensor, and the rope gate follows the arena field (_qsa_fields only declares it when cache_rope_positions). resolved, thanks.
build_kv_cache_tensor bound the FP8 K/V scale planes and the quantized compressed-index scale to the attention module but did not register them, so the in-process codec stored and restored FP8 rows without their scales: GSM8K pass 2 fell from 0.98 to 0.54. Register k_scale / v_scale and append index_compressed_scale to `extra_planes`, all through the row space's own suffix (MTP draft rows registered the target's index rows), and drop #2636's refusal of lmcache_offload with quantized QSA caches. Co-authored-by: Cursor <cursoragent@cursor.com>
… lookup Look each field up once and hand the same views to bind_caches and KVCacheTensor. Fold the codec coverage into the existing draft-row test: every QSA layer, five cache-format cases, each registered tensor matched by data_ptr against its own row space and row in the stored order. Co-authored-by: Cursor <cursoragent@cursor.com>
…mcache-fp8-kv #2636 moved the QSA rules into qwen4_exp/qsa_cache_rules.py; this branch's changes to them (no lmcache_offload refusal for quantized caches, qsa_caches_quantized folded into qsa_draft_rows) move with them. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Whole-PR cleanup, pushed as Size. The diff GitHub showed before was 12 files, +1619/-111. Most of that was a pre-merge copy of #2652, which is now on
What was merged or removed.
@valar-cop, your three inline comments:
Validation (
@HaonanWang98 your approval predates this rework; could you take another look? cc @valarLip |
|
ci triage for de10df6:
|
…mcache-fp8-kv #2636 folded qwen4_exp/qsa_cache_rules.py back into qwen4_exp_attn.py; this branch's rule changes (no lmcache_offload refusal for quantized caches, qsa_caches_quantized folded into qsa_draft_rows) follow them. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Merged the new #2636 head ( |
Drop the refusal of quantized QSA caches under in-process offload and hand the codec the k/v scale rows and the index_compressed scale plane through extra_planes, with a test that every quantized layer gives the codec the rows it binds. Co-authored-by: Cursor <cursoragent@cursor.com>
Bind the layer and describe its KVCacheTensor from one set of views, as the earlier review asked, keep the recipe's P/D line as it was, and drop _check_kv_transfer from the readers of the QSA field rules. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Update Size vs #2636. Before: 3 files, +87/−111. After: 3 files, +94/−31. The change now adds lines instead of rewriting #2636's. Per file.
Not changed. Validation.
Pending. CI is running. #2636 lands first. The description is updated to match this head. |
resident_grid_refusals takes the KV connectors a model has run beside its resident kernels (default none, so Kimi-K3 and the DSA indexer are unchanged). Qwen Mono passes ``lmcache_offload``: its stores and loads are bounded gather / scatter kernels and host copies on the connector's own pack / copy streams, ordered by events (the pack stream waits on the forward's producer event), with no persistent kernel and nothing that waits on the model, so they can share or delay CUs but not hold one. Other connectors stay refused. GDNPools also checks that both pools are dense, as the kernels address them. Checked on one MI355X, on this branch merged with #2655 (#2636's Qwen offload), Zhang Ling's A line with OFFLOAD=1, a 3,613-block GPU pool and 16 sequences: 32 12K-token prompts pushed the early prefixes to the CPU tier, then three 4000-token Mono decodes ran while 16 of them came back (195,584 KV tokens and 16 GDN state images loaded from the CPU tier, stores throughout). There were 0 errors and 0 timeouts; the decode event gap was 7.5 / 84 / 89 ms (median / p99 / max), against 9.4 / 84 / 87 ms on Native. GSM8K 200 with real MTP and offload: 0.985, acceptance 0.72. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Motivation
With FP8 QSA caches and in-process
lmcache_offload, a prefix fetched back from the CPU tier decodes under the wrong scales. The dense codec moves exactly whatbuild_kv_cache_tensorregisters onKVCacheTensor, and the FP8 K/V scale planes and the compressed-index scale plane were not registered. The quantized rows came back, but the scales stayed whatever the reused block held. Nothing is logged: on the GSM8K pass served from CPU, accuracy fell to 0.5432 and MTP acceptance to 0.06. #2636 therefore refuses this combination at startup. This PR registers the planes and lifts that refusal.Base: #2636 (
lingzha/qwen38-lmcache). The FP8 QSA caches themselves (#2652) are onmain.Changes
Qwen4ExpMetadataBuilder.build_kv_cache_tensorlooks each arena field up once, with the row space's suffix, and hands the same views tobind_cachesand toKVCacheTensor. It now registersk_scale/v_scale, and appendsindex_compressed_scaletoextra_planes. The dense codec readsextra_planesafter its existing tails, so BF16 byte layouts do not change.qwen4_exp_attn._check_kv_transferno longer refuseslmcache_offloadwith quantized QSA caches.recipes/Qwen3.8-Flash-Next.md: drops the note asking for BF16 caches withlmcache_offload.tests/test_qwen4_exp_offload.py:test_qwen38_runs_in_process_offload_of_quantized_caches).test_every_quantized_qsa_layer_hands_the_codec_the_rows_it_bindscovers the four quantized cache formats, using two target layers plus an MTP draft layer with rope positions. Each registered tensor must be its own row space's row of the right field, in stored order, and the set must equal what the layer binds. All four cases fail whenk_scaleis left unregistered.test_a_bf16_target_gives_the_mtp_draft_layer_its_own_rows.Stored layout and older tiers
PAGE_LAYOUT_VERSIONis unchanged, for these reasons:lmcache_offload: [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636 refuses that combination at startup.main's Qwen chunks digest differently anyway, because the namespace keyskv_cache_dtype/index_cache_dtypeand [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636'sqsa_page_segments.lmcache_mpalready copied whole arenas and has its own::lmcache-mp-vsuffix.A version bump would invalidate every model's persisted PAGE objects.
Not changed
KVCacheTensorandDenseKVByteCodecare unchanged;PAGE_LAYOUT_VERSIONtoo (above).main'sbuild_kv_cache_tensordocstring still says "three caches"; it is out of date but left alone.qsa_caches_quantized(now one caller) and theleaveslocal stay as in [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636.tests/test_qwen4_exp_offload.pyis gated on AITER, so the CI CPU runner skips it. The module under test imports AITER at module level, so the new test has to live there; it runs in the GPU unit tests.Validation (
f0bea2199; GSM8K onde10df694and4aff317a5, see below)mainfail too. Two engram host-path tests, which this PR does not touch, failed once and pass on rerun.test_qwen4_exp_offload.py,model_ops/test_qsa_quantized_cache.py,test_aiter_flydsl_gdn.py,test_qwen4_exp_mtp.py,test_lmcache_offload_config.py,test_lmcache_offload_connector.py, MI355X): 558 passed.amd/Qwen3.8-Flash-Next-Quark-MXFP4, TP1, MTP 3, block 64, in-processlmcache_offload(320 GiB CPU pool), questions 0-439, state interval 256. Pass 1 stores. Pass 2 resends the same prompts after HBM was flushed, so it is served from CPU.de10df694registers the same views the same way as this head.4aff317a5differs only in looking each field up twice, and in [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636's comment, line-order and default-argument changes.de10df694)de10df694)4aff317a5)