Conversation
…3 NaN The SM70 grouped-page4 QSA prefill route can emit NaN for every query row and head on a fixed set of head dims when the KV cache is E4M3 and a hybrid (GDN/Mamba) layout co-locates non-attention state in the paged KV pool. The same batch is correct through the XQA route. Root cause: the grouped planner pads each category to a multiple of eight with the null block -- (physical microblock 0, mask 0) -- and counts the padding in seq_len. The forward loads page 0's K/V for those padded rows and masks them multiplicatively (P = 0). But 0 * NaN = NaN survives the P@V tensor-core reduction, and under E4M3 the null block's bytes decode to NaN. Every group's every tile reads the same page-0 microblock, so the whole tile's rows and heads go NaN on identical dims. XQA never reads the null block, which is why it stays clean on the same inputs. Fix (defense in depth, both default-on): - Kernel: in flash_attention_grouped_verify_e5m2_partial_kernel under SPARSE_PAGE4, after loading each K/V panel, zero every row no query attends (unattended iff (mask & (0x11111111u << (row & 3))) == 0) so the reduction sees 0 * 0 = 0. A leading __syncthreads() before the V-panel zeroing keeps the panel load from racing past the zero and resurrecting the NaN; the trailing one publishes the zeros before P@V. - Python (VLLM_SM70_QSA_GROUPED_PAD_FIX, default 1): after the planner, repoint every mask==0 padding entry at the group's first real microblock. torch.where + copy_ has no host sync, so it stays CUDA-graph capturable. This backstops binaries built before the kernel fix. The grouped route stays default-on (VLLM_SM70_QSA_GROUPED_PAGE4=1); only what the padded rows contribute changes. Adds tests/models/qwen4_exp/test_qsa_e4m3.py::test_qsa_grouped_page4_null_block_padding_no_nan, an SM70-only regression that poisons the null block's value plane with E4M3 NaN and asserts the grouped route stays finite and matches XQA, parametrized over the Python pad-fix off (kernel-only) and on. It reproduces the NaN on an unfixed kernel and passes on either fix. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> (cherry picked from commit cab8358)
…adding test The regression added in the previous commit used the XQA page4 route as its clean reference. On Flash-V100 builds whose XQA E4M3 page4 kernel requires an fp32 partition buffer, that reference raises "XQA decode tmp_out must be fp32 for E4M3 KV" before the grouped route is exercised, so the test cannot run against such a build. Use a kernel-independent einsum ground truth instead: decode the E4M3 K/V to fp32 and attend only the selected tokens (which map to physical pages >= 1, so the null block is excluded). This is .so-agnostic and directly checks that the grouped route stays finite and numerically correct. The grouped route and the parametrized pad-fix coverage are unchanged. Verified on V100 (SM70): the grouped null-block NaN reproduces with the stock kernel and pad-fix off (test fails), and passes with either the kernel fix or the Python pad-fix. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> (cherry picked from commit 31d137c)
Adapt vllm-project/vllm#51865 (0a94d85a66499cad8297ead86a470967de5c0212) to the current target and Eagle/MTP draft paths. A K+1-token prompt tail must not replay a decode graph merely because its shape matches. Preserve dummy capture and genuine decode classification. Co-authored-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Janelle Cai <janelle.cai@modal.com> Assisted-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> (cherry picked from commit 1892fb2)
Staging branch for the separate E4M3 + MTP draft-scale enablement; NOT for merge and no PR. Ports two work-in-progress patch sets on top of the null-block-padding fix branch: - phase2: opt-in gate (VLLM_QWEN4EXP_QSA_E4M3_MTP) that relaxes the QSA E4M3 MTP0 requirement, an extracted _verify_e4m3_kv_requirements, the envs entry, an MTP weight-remap hook, calibration-overlay tooling, and CPU/GPU tests. - phase2b: a CUDA-graph-capture-safe XQA workspace (torch.full on-device instead of a host torch.tensor copy) with capture/envelope tests. CPU check (CUDA hidden): 27 passed, 19 skipped (GPU-only), 2 failed. The two failures (test_validate_scale_overlay_lists_missing_names, test_finalize_qsa_scale_load_success_and_missing) exercise draft-side q-scale visibility whose model.py scale-load changes are not part of these two patches; that piece is out of scope for this staging branch. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> (cherry picked from commit 464f538)
Finish the staged draft-scale path without changing the existing main-model fallback contract. Add a repository-owned, revision-checked target/draft scale overlay materializer and make its CPU tests portable. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> (cherry picked from commit bb4bc78)
The 1CatAI#647 overlay kept the stricter native E4M3 ABI but lost the matching Python workspace dtype selection. Key the cached workspace by KV dtype and allocate FP32 temporary output for E4M3 while retaining FP16 for the existing path. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> (cherry picked from commit b86e641)
ruff-format and clang-format only; no behaviour change. Keeps the native tree byte-identical to the already-qualified SM70 build so its compiled binaries can be reused. Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
test_qsa_e4m3_loader_requires_all_24_scales asserted that an incomplete overlay raises, but left VLLM_QWEN4EXP_QSA_E4M3_STRICT_SCALES unset, so _validate_qsa_e4m3_scale_load took the non-strict warning path and the test failed on main. The gate it claimed to cover was never exercised. Set the env explicitly and assert both contracts: non-strict returns the missing scale names so the caller can fall back to unit scales, strict raises with the loaded/required count. Also assert the non-FP8 cache short circuit returns an empty set rather than only that it does not raise. Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
GPU acceptance: PASS at C1, C2 and C4Image Startup gate on every run: calibrated main QSA GatesSixteen generations per run: four real prompts (Chinese scene, Chinese story,
E4M3 buys +74% to +78% KV capacity at the same memory fraction.
The corruption screen is deliberately narrow: doubled CJK punctuation, U+FFFD, Same-seed A/BAll 64 generations (4 runs x 16) are corruption-free. The candidate is MTP
Prefill and decodeWarm-up request discarded, then three measured repeats; medians. Output pinned
Prefill is 23% to 46% faster. Decode is 22% to 30% slower, and the That is the trade this lane buys: substantially more context and faster prefill Warming matters when reading these numbers. The first request at a new shape Not covered
Evidence outside the containers on llm252: |
2 participants
Why this replaces the previous E4M3 lane
PRs #16, #17 and #18 grew into a single stack that mixed the E4M3 work with
the CPU KV offload lane (1CatAI#598, 26 commits), 1CatAI#637, 1CatAI#647 and an unmerged
experimental branch. That stack qualified, was promoted, and was rolled back
the same day. Untangling it turned out to be cheaper than continuing it,
because the E4M3 delta is far smaller than the stack suggested.
Against 1CatAI
main(b711d53045), the whole QSA E4M3 main KV path isalready present. What was actually missing is four things, and this branch is
exactly those four plus their two prerequisites:
28697df2c7,c0c6a3fd879e8ea99b61e1aa0103ffVLLM_QWEN4EXP_QSA_E4M3_MTP.fdebb3e0a687af57de30004db582d6The native tree is byte-identical to the already-qualified SM70 build, so the
compiled binaries are reused and only the Python layer is repacked.
This branch deliberately does not carry 1CatAI#598. CPU KV offload stays its own
lane; the separate H2D restore-ordering fix went into 1CatAI#598 directly.
Relationship to upstream vllm-project
vllm-project/vllm has merged its own QSA FP8 work: #54890 (FP8 indexer cache,
2026-09-07) and #55557 (fp8_e4m3 main KV on the QSA path, 2026-09-16). Neither
is in 1CatAI
mainyet.They are not duplicates of this branch. Upstream's
qsa_sparse_paged_attentionFP8 support targets the Triton split-K kernel, asserts
q.dtype == torch.bfloat16, and adds an sm_120 tuning table. The SM70 deployment here runsFP16 activations through the native page4/XQA routes. The two live in the same
files but on different code paths.
The contract is the same on both sides, which is what matters for the eventual
rebase: per-tensor host-side
k_scale/v_scalefloats, the K scale foldedinto the attention scale, the V scale applied as an output scale, and an
E4M3 cache stored as uint8 and viewed as
float8_e4m3fn. No device-side scalebuffers on either path.
Known constraint:
draft_sample_methoddraft_sample_method=probabilisticwith MTP corrupts output atmax_num_seqs=4and any temperature above 0. This is not introduced here andis not E4M3-related: it reproduces on the current production image
(
1.5.0+ct252.20260914.gf21df5b279) with FP16 KV by changing onlymax_num_seqsfrom 2 to 4. It is temperature-gated because temperature 0 usesthe greedy argmax verify path, while temperature > 0 goes through the
probabilistic rejection-sampling ratio test.
greedyis upstream's default for MTP (arg_utils.py) and is unaffected atevery concurrency tested. This lane uses
greedy; the probabilistic defect istracked separately.
Tests
CPU, this branch:
tests/models/qwen4_exp/303 passed, 103 skipped, 0 failed.The 168 errors in that run are teardown-only (
cleanup_dist_env_and_memoryneeds an accelerator) and reproduce identically on
main.pre-commiton all 19 changed files: all hooks pass, including ruff, mypy,clang-format, SPDX, forbidden imports and the CUDA API policy.
GPU acceptance is in progress and will be posted here: quality gates
(real-text prompts, temperature sweep, >=256-token generations, same-seed A/B
against the known-good baseline), C1/C2/C4, MTP, and 1k/8k prefill and decode
throughput.
AI assistance
Claude produced the re-plan, the rebuilt stack, the test repair and this
description. Commits keep their original authors:
28697df2c7,c0c6a3fd87and
e1aa0103ffare areslp's. The human submitter must review every changedline before merge.