Fix Q UE8M0 quant and require fp32 LN params in fused DSv3.2 indexer kernel - #3451
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
94eab23 to
0300efa
Compare
…kernel The fused `indexer_qk_rope_quant_and_cache_kernel` (added in ROCm#3185) diverged numerically from the unfused vLLM path it replaces, causing a catastrophic GSM8K accuracy regression on DeepSeek-V3.2 (0.95 -> 0.49 strict-match, 20-shot) when fusion was enabled. Two independent bugs in this kernel: 1. Q quant skipped UE8M0 rounding and used direct division by a constexpr fp8_max. The unfused path is `per_token_group_quant_fp8` in vllm/.../fp8_utils.py which (a) multiplies by the reciprocal of fp8_max and (b) applies UE8M0 (`scale = 2^ceil(log2(scale))`) when the caller passes `scale_fmt = "ue8m0"`. DSv3.2's indexer always passes ue8m0, so the fused kernel disagreed with the unfused one on ~94% of Q values per call. The K-quant block in the same kernel already applied UE8M0 correctly (lines ~1429); only the Q block was missing it. Direct division also triggers GPU fast-division which adds 1-ULP noise that flips the rounded fp8 value at representable boundaries. 2. `norm_weight` and `norm_bias` were typed `scalar_t*`, which in production is bf16. vLLM stores `LayerNorm.weight`/`bias` as fp32 and `F.layer_norm` consumes them as fp32 in the unfused path. The implicit fp32 -> bf16 cast at the call site lost ~16 mantissa bits per element, causing K to drift from the unfused reference by 1+ fp8 ULP on ~3% of values per call. Both drifts compounded across the 58 indexer layers in DSv3.2 and broke long-context accuracy. With both fixes applied (and a separate cache-layout fix on the vLLM call site), GSM8K strict-match recovered to 0.95. Changes: - Apply UE8M0 + multiply-by-reciprocal to the Q scale, mirroring the unfused Triton `per_token_group_quant_fp8`. - Hard-type `norm_weight` and `norm_bias` as `float*` in the kernel signature and the launch macro, drop the `static_cast<float>` on each load, and assert `dtype == fp32` on the host so the bf16 footgun cannot be reintroduced. ABI note: this tightens the dtype contract for `norm_weight`/`norm_bias` from "match q.dtype" to "must be fp32". vLLM is the only known caller of this kernel today and already has fp32 LN params natively, so the change is benign for it. Any future caller passing bf16 will get a loud host check failure rather than silent ~3% K drift.
0300efa to
423f22a
Compare
There was a problem hiding this comment.
Pull request overview
This PR fixes two numerical-correctness issues in the fused indexer_qk_rope_quant_and_cache_kernel to match the unfused vLLM reference path for DeepSeek-V3.2, restoring accuracy when fusion is enabled.
Changes:
- Fix Q-side FP8 quant scale computation to match vLLM (reciprocal-multiply + optional UE8M0 power-of-two scaling).
- Require LayerNorm
norm_weight/norm_biasto be FP32 end-to-end (kernel signature, launcher casts, and host-side dtype checks), avoiding silent bf16 precision loss.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…ix lands The DSv3.2 indexer eval hangs and hits the 3600s timeout (HIP backtrace) on current AITER main. Verified the ROCm#3451 fix (cache_kernels.cu) cherry-picked does not resolve it yet, so keep DeepSeek-V3.2 Accuracy disabled (run_on_pr/run_on_schedule = False -> reported in the 'disabled' bucket, not deleted) to keep the downstream signal clean. Re-enable once the DSv3.2 indexer kernel fix merges (ROCm#3451 / dsv32-indexer-fused-kernel-fixes).
|
Data point from a downstream verification run: I cherry-picked this PR's fix ( |
…+ SGLang) on MI350X (#3441) * ci(atom-downstream): add Kimi-K2.5-MXFP4 TP8 accuracy gate on MI350X Add Kimi-K2.5-MXFP4 to the ATOM downstream accuracy matrix so AITER changes are regression-checked against the Kimi e2e workload. Verified on MI355X (gfx950) 2026-05-30 with rocm/atom-dev:vllm-latest (aiter ef114b0), vllm serve TP8: gsm8k 3-shot flexible-extract 0.9409, matching the amd/Kimi-K2.5-MXFP4 reference. Runs on the requested linux-aiter-do-mi350x-8 runner; threshold 0.92 leaves margin. Triggered with the existing ci:atom / ci:all labels (and on push/ schedule/workflow_dispatch). * ci(kimi-downstream): add vLLM (OOT) + SGLang Kimi-K2.5 accuracy gates on MI350X Complements the ATOM in-tree gate in atom-test.yaml so AITER is regression-checked against Kimi-K2.5-MXFP4 across all three downstream serving stacks. New workflow kimi-downstream.yaml rebuilds the PR's AITER (gfx950) into the backend image and runs gsm8k via the in-image launch scripts: - vllm: atom-dev:vllm-latest + atom_oot_test.sh (vllm serve), blocking, threshold 0.92 (verified 0.9409 on MI355X, aiter ef114b0) - sglang: atom-dev:sglang-latest + atom_sglang_test.sh, continue-on-error (non-blocking) -- SGLang currently can't load AMD Quark MXFP4 per-expert weights; kept as a visible signal until that loader bug is fixed. Runs on linux-aiter-do-mi350x-8 under ci:kimi / ci:vllm / ci:sglang / ci:all labels (and push/schedule/workflow_dispatch). * ci(kimi-downstream): use official ROCm nightly images for vLLM + SGLang Move the downstream Kimi-K2.5 gates off the AMD-internal atom-dev:*-latest images onto the official nightly downstream containers, so the gate validates AITER against what actually ships: - vLLM: rocm/vllm-dev:nightly (upstream-native KimiK25ForConditionalGeneration) - SGLang: lmsysorg/sglang-rocm:v0.5.12.post1-rocm720-mi35x-20260531 (no floating nightly tag exists; pin latest MI35X/ROCm7.2 build) These images lack the ATOM launch scripts + lm_eval, so add self-contained launchers (.github/scripts/kimi_{vllm,sglang}_accuracy.sh) that vllm serve / sglang.launch_server + lm_eval gsm8k and print KIMI_FLEX_EXTRACT. The PR's AITER is built from source into the image at runtime (gfx950), avoiding any torch ABI mismatch. SGLang stays non-blocking (loader bug). * ci(kimi-downstream): fix vLLM + SGLang Kimi-K2.5 gates on official nightly images Reproduced both gate failures on MI350X (gfx950) TP8 and fixed them. vLLM (rocm/vllm-dev:nightly): the worker crash was --load-format fastsafetensors -- the official image doesn't ship the fastsafetensors package, so every TP worker died at weight load with an ImportError (surfaced only as 'WorkerProc initialization failed'). Drop the flag; the default safetensors loader handles the 521GB checkpoint in ~27s/worker. Reference accuracy 0.9409. SGLang (lmsysorg/sglang-rocm:...mi35x-20260531): two upstream-SGLang bugs on this Kimi-K2.5 path (not AITER): 1) MoE loader crash with shared-experts fusion vs unfused per-expert Quark MXFP4 shard shape -> --disable-shared-experts-fusion. 2) Default triton/fused-MLA path crashes on Kimi-K2.5 head dims (and fp8 KV cache dies on the triton fp8 dot) -> --attention-backend aiter + bf16 KV (drop --kv-cache-dtype fp8_e4m3). This routes MLA through the PR's AITER kernels end to end. Measured gsm8k flexible-extract 0.9272. SGLang lane is now blocking (continue_on_error: false). * ci(atom-downstream): expand coverage to InferenceX MI355X frontier set Add DeepSeek-V4-Pro, Qwen3.5-397B-A17B-FP8, MiniMax-M2.7, GLM-5.1-FP8 to the ATOM accuracy generator's default set, matching the model coverage of SemiAnalysisAI/InferenceX on MI355X. All configs come from ATOM models_accuracy.json; runners pinned to the AITER MI350X pool by TP size (V4-Pro/GLM-5.1 -> do-mi350x-8, Qwen3.5-397B -> -4, MiniMax-M2.7 -> -2). GLM-5.1-FP8 chosen over GLM-5-FP8 (GLM-5 has a known ATOM sparse-attn- indexer crash near the gsm8k request count). * ci(kimi-perf): add Kimi-K2.5 perf gates (vLLM + SGLang) on official nightly New kimi-perf-downstream.yaml runs a Kimi-K2.5 throughput sweep (ISL/OSL 1024/1024, concurrency 4..64) on linux-aiter-do-mi350x-8 and gates on c=64 output tok/s. Triggered by ci:performance / ci:all label + nightly schedule (20:43 UTC) + manual dispatch. Builds PR AITER from source into the official nightly images, reusing the validated accuracy-gate launch flags. Validated on MI350X gfx950 TP8 (PR AITER built from source): vLLM (rocm/vllm-dev:nightly) c=64 3126.4 tok/s -> floor 2250 SGLang (lmsysorg/sglang-rocm:...531) c=64 3284.7 tok/s -> floor 2400 * ci(atom-downstream): keep PR-default set small; frontier models via ci:atom_full Move the InferenceX frontier models (DeepSeek-V4-Pro, Qwen3.5-397B-A17B-FP8, GLM-5.1-FP8, MiniMax-M2.7) out of the always-on PR default set and back to on-demand ci:atom_full coverage. Running all of them on every PR overloaded the do-mi350x pool and flaked the Kimi accuracy + perf gates. The runner pins (do-mi350x by TP size) are retained, so ci:atom_full still exercises them on the AITER MI350X cluster. MiniMax-M2.7 (needs an ATOM-side HSA_NO_SCRATCH_RECLAIM env) is therefore no longer red on every PR. * ci(sglang-downstream): disable DeepSeek-V3.2 accuracy until indexer fix lands The DSv3.2 indexer eval hangs and hits the 3600s timeout (HIP backtrace) on current AITER main. Verified the #3451 fix (cache_kernels.cu) cherry-picked does not resolve it yet, so keep DeepSeek-V3.2 Accuracy disabled (run_on_pr/run_on_schedule = False -> reported in the 'disabled' bucket, not deleted) to keep the downstream signal clean. Re-enable once the DSv3.2 indexer kernel fix merges (#3451 / dsv32-indexer-fused-kernel-fixes). --------- Co-authored-by: Peng Sun <sunway513@users.noreply.github.com>
|
Thanks @sunway513 - worth flagging: SGLang's DSv3.2 indexer ( This PR only touches |
|
The core fix is correct: aligning the Q scale with the unfused per_token_group_quant_fp8 (reciprocal-multiply) and making the LN params fp32 end-to-end. Two non-CI items to address before merge:
|
The kernel indexes norm_weight[dim] / norm_bias[dim] as a flat array; a non-1D or non-contiguous view would silently produce wrong results. Add explicit dim() == 1 and is_contiguous() checks before launch so a bad input is caught immediately rather than corrupting K values.
|
Thanks for the review @ChuanLi1101! I've adressed both points:
|
ChuanLi1101
left a comment
There was a problem hiding this comment.
LGTM - the Q UE8M0 + fp32 LN fixes are correct, and my earlier review points are now addressed (contiguity/1D guards added, breaking dtype change documented in the ABI note + caller audit). One thing before merge (non-blocking on the code): please confirm the red CI gates (Kimi-K2.5 vllm, Kimi perf vllm, gpt-oss-120b fp8_kvcache benchmark) are pre-existing/unrelated - none of them exercise the DSv3.2 indexer path this PR touches, and DeepSeek-R1 / SGLang / Standard tests are all green.
Thanks for the review. The Kimi errors here are due to this bug in upstream vllm that is hopefully going to be fixed in the next nightlies: vllm-project/vllm#45596 The GPT-OSS error is related to refactoring triton version update. The current version in vllm is 3.6 but installing aiter automatically updated that to 3.7 (got released 2 weeks ago) this version moved |
|
let me know once ci passed |
|
Thanks @valarLip! The CI failures are unrelated to this PR, see @amd-mghanimi's comment above. I rebased the PR but need approval to re-start the CI |
|
New CI issues:
In the past, some of these errors could get resolved by rerun. Because part of the slow loading time for these tests were due to model was getting downloaded on the runner machine, so on second attempts often they could pass but I think for Kimi tests the source script should add some waiting time to be able to pass the tests. |
|
PR #3451 updates LayerNorm weight and bias to FP32, which requires matching changes in ATOM’s LayerNorm module. Merging this PR without corresponding ATOM adjustments will break existing production models. Background Required ATOM LayerNorm model-op updates
Merge Requirement |
|
Thanks @zufayu for the comment, agreed this needs an ATOM-side companion. This PR only changes the dtype contract on indexer_qk_rope_quant_and_cache (norm dtype == q.dtype → == fp32), it doesn't touch ATOM's LayerNorm or layernorm2d_fwd, so the CK kernel shouldn't need any change. Also, are you seeing an actual accuracy regression in ATOM, or is this from review? On vLLM this fix recovers GSM8K 0.49 → 0.95 (with the vLLM companion #43907) when running with |
|
Checked vLLM #43907 (deepseek_v2.py): it calls .to(torch.float32).contiguous() on k_norm.weight/bias explicitly, commenting "casting to bf16 here loses precision and drifts K from the unfused path." — so the companion ATOM change is required. |
|
Rebased cleanly on latest main (no conflicts) - needs approval from a maintainer to re-run the CI. Code is unchanged since @ChuanLi1101's LGTM: contiguity/1D guards and the breaking fp32 dtype-contract doc are in. @junhaha666 (already requested) / @valarLip - could one of you take a look and review this PR? The ATOM dependency is covered by ROCm/ATOM#1359 (optional |
|
btw looks like no accuracy diff for ATOM side? |
No accuracy diff on ATOM, @cpersson-amd's numbers in ROCm/ATOM#1359 show DSv3.2 gsm8k 5-shot at 0.9560 vs 0.9591 strict-match on main. ATOM already calls the indexer with preshuffle=True, so it never hit the K-cache layout bug that tanked vLLM (fixed separately in vllm#43907) — here the fp32 LN change is just a precision alignment and accuracy stays flat. #1359 keeps dtype optional (defaults None) so all other models are unchanged, and only pins DSv3.2 k_norm to fp32. Suggest landing the two together. CI's green now |
Summary
The fused
indexer_qk_rope_quant_and_cache_kernel(added in #3185) diverged from the unfused vLLM path it replaces, dropping DeepSeek-V3.2 GSM8K strict-match (lm-eval, 20-shot) from 0.95 → 0.49 when fusion was enabled. Two independent kernel bugs incsrc/kernels/cache_kernels.cu:Bug A — Q quant. The unfused reference
per_token_group_quant_fp8multiplies by1 / fp8_max(avoids GPU fast-division noise) and applies UE8M0 rounding whenscale_fmt = "ue8m0", which DSv3.2's indexer always passes. The fused kernel did neither, while the K-side of the same kernel applied UE8M0 correctly. Net: ~94% Q disagreement per call. Fix mirrors the K-side.Bug B —
norm_weight/norm_biastypedscalar_t*(= bf16 at runtime). vLLM storesLayerNorm.weight/biasas fp32 andF.layer_normconsumes them as fp32; the call site was casting fp32 → bf16 before launch, drifting K on ~3% of values per call. Fix hard-types both asfloat*, drops the now-redundantstatic_cast<float>in the kernel body, and asserts fp32 on the host. Also adds explicitdim() == 1and contiguity guards so non-1D or strided views are caught at launch rather than silently producing wrong results.Breaking change
norm_weightandnorm_biasnow must be fp32. The dtype contract tightened from "must matchq.dtype" to "must be fp32"; any caller passing bf16 LN params will now get a loudAITER_CHECKfailure at launch instead of silent ~3% K drift. Caller audit: vLLM DSv3.2 indexer and ROCm/ATOMattention_mla_sparse.pyboth store LN params as fp32 natively and are unaffected. SGLang does not call this kernel.End-to-end result
DSv3.2, MI355X, TP=4, lm-eval GSM8K (
num_fewshot=20,num_concurrent=256, full 1319-question suite):preshuffle=True, fp8preshuffle=True, bf16The fix is independent of
--kv-cache-dtype(the indexer's internal cache is always fp8). Drift harness: fused vs unfused Q fp8 + Q scale and K (withpreshuffle=Trueon both) bit-exact post-fix.Companion vLLM PR
vllm-project/vllm#43907 — ships
preshuffle=Trueand switches the call site to pass fp32 LN params.