fix(ck): fix gfx950 bpreshuffle determinism and illegal-type codegen - #8639
Conversation
|
Reviewed and validated against the sgl-project/sglang#28685 repro on MI350X / ROCm 7.2 (gfx950). I rebuilt the aiter Correctness — the per-M-tile sparse wrong-rows are gone:
Determinism — 30 repeated launches on the worst prior cases (M = 96, M = 128) are bitwise-identical ( This matches the root cause analysis in sglang#28685: the packed 2-lane post-scale FMA ( LGTM — thanks @hdt98. Once this lands in CK/aiter it supersedes the sglang-side |
…50 warning Point the GLM-5.2 gfx950 block-FP8 accuracy warning at the validated upstream CK kernel fix (scalar FMA + accumulator anchor) instead of the interim plain-kernel workaround. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
|
works for long context if set SGLANG_FP8_PAGED_MQA_LOGITS_TORCH=0 |
|
@hdt98, can we merge it? |
|
@andyluo7 I am still testing/validating this with the integration with your PR sgl-project/sglang#28471 for full SGLang model path |
|
@andyluo7 please take a look at this PR andyluo7/sglang#2 SGLang upstream PR tracking here: sgl-project/sglang#29275 |
There was a problem hiding this comment.
Pull request overview
This PR updates Composable Kernel’s blockscale weight-preshuffle FP8 GEMM path to eliminate a gfx950/ROCm 7.2 nondeterminism/wrong-row failure mode by avoiding a packed accumulator update pattern that depended on an LLVM illegal-type coercion workaround. It also extends the CK profiler/test harness to detect bytewise nondeterminism and sparse per-row correctness failures on representative regression shapes.
Changes:
- Switch bpreshuffle v1/v3 post-scale accumulator updates from packed 2-lane vector FMAs to scalar
AccDataTypeFMAs, and add a VGPR “anchor” after each update to stabilize repeated-launch determinism. - Add profiler support for repeated-run bytewise determinism checks plus a per-row relative error diagnostic to catch sparse bad-row failures.
- Add long-running FP8 regression tests (gated by an env var) for the reported WKV determinism repro shapes and GLM-5.2 out-of-allowlist shapes.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/test/gemm_blockscale_wp/test_gemm_common.hpp | Extends the shared test runner to pass through determinism-check count and optional verification control. |
| projects/composablekernel/test/gemm_blockscale_wp/test_gemm_blockscale_wp_fp8.cpp | Adds env-gated long regression tests covering reported determinism and GLM-5.2 out-of-allowlist accuracy/determinism shapes. |
| projects/composablekernel/profiler/include/profiler/profile_gemm_blockscale_wp_impl.hpp | Implements determinism repeat-run memcmp checks and row-relative error diagnostics in the profiler path. |
| projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle_v3.hpp | Reworks bpreshuffle v3 post-scale accumulation to scalar FMAs and adds VGPR anchoring. |
| projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle_v1.hpp | Reworks bpreshuffle v1 post-scale accumulation to scalar FMAs and adds VGPR anchoring. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
aosewski
left a comment
There was a problem hiding this comment.
Thanks for tackling this — the scalar-FMA rewrite reads as a clean, mathematically-equivalent fix for the gfx950 illegal-type miscompile (both packed lanes always held the same scale, so dropping to a scalar is a pure codegen-shape change), and the determinism + row-relative tooling is a welcome addition. I left a handful of inline notes, mostly suggestions. Two things I'd most like your take on:
-
Is the
asm volatileanchor actually load-bearing, or does the scalar rewrite alone already give you determinism? Right now it reads as two mechanisms for one fix, and it isn't clear which is doing the work — an A/B on the failing shapes (or a comment explaining why the anchor is essential) would settle it, and would stop a future cleanup from deleting it as a harmless no-op. -
The four sibling MoE pipelines still carry the identical packed
GetRegSizePerXdlops() / 2pattern. Since the root cause is compiler-level (illegal-type legalization, not specific to this path), they look exposed to the same gfx950 issue — is leaving them untouched deliberate (not reachable by the reported models) or just out of scope here? Worth a note + tracking issue either way.
No blockers from me; details inline.
Thanks @aosewski for careful reviews.
|
…fle_stabilizer_prready
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
|
@andriy-ca and reviewers, could you please help approve/trigger the remaining CI runs for the current head? I updated PR title to be compliant with policy checks and are now passing, but the required workflows appear to be in action_required state. |
|
@andriy-ca I pushed a new commit to address @aosewski remaining in-code comments/clarifications: added the AccDataType guard, named/rationalized the row-relative thresholds, hoisted the row-relative check out of the FP8 tolerance branch, documented the bit-exact repeat-run check, and added M/N/K SCOPED_TRACE coverage. Validation on MI350:
The current therock-pr-bot failure appears to be because required workflows in action_required state, not because of a test failure. Could you or a maintainer please approve/trigger the remaining CI runs? Thank you. |
fix(ck): fix gfx950 bpreshuffle determinism and illegal-type codegen (#8639) ## Motivation Fix the CK `gemm_a8w8_blockscale_bpreshuffle` path behind: - ROCm/aiter#3261: nondeterministic `gemm_a8w8_blockscale_bpreshuffle` results on gfx950/MI350/MI355 for fixed inputs and launch shapes. - sgl-project/sglang#28685: GLM-5.2-FP8 can silently produce wrong rows on gfx950/ROCm 7.2 when out-of-allowlist block-FP8 linear shapes route to AITER's CK bpreshuffle path. The issue is isolated to the bpreshuffle CK kernel path on gfx950/ROCm 7.2, not general FP8 numerics. This PR updates the CK source so the kernel no longer depends on the removed LLVM illegal-type coercion workaround. ## Technical Details - Replace the packed 2-lane post-scale accumulator FMA update in bpreshuffle v1/v3 with scalar `AccDataType` updates. - Add a VGPR read/write accumulator anchor after each post-scale scalar update to enable deterministic repeated launches. - Add profiler/test support for: - repeated-run bytewise determinism checks - row-relative error checks to catch sparse bad-row failures near tile boundaries - GLM-5.2 representative/out-of-allowlist shapes Relevant issues: - ROCm/aiter#3261 - sgl-project/sglang#28685 ## Test Plan Validated on gfx950 with ROCm 7.2 using the CK profiler/test path. Coverage includes: - ROCm/aiter#3261 WKV determinism sweep: - `M=8192`, `N=512`, `K={384,640,3968,4096,4224}` - GLM-5.2-FP8 representative shapes: - `q_a_proj`, `q_b_proj`, `mlp.down_proj`, and `kv_b_proj` - row-boundary sweep including `M={32,48,56,64,72,96,128}`, `N=28672`, `K=512` - CPU reference comparison for correctness cases - row-relative error reporting for sparse bad-row failures - repeated-run determinism checks - before/after perf regression sweep - manual pre-commit run on changed CK files ## Test Result - Kernel passed the tests: - `ReportedWkvDeterminism` - `Glm52OutOfAllowlistAccuracyAndDeterminism` - Build/test compile command validation showed `COERCE_COUNT 0`, confirming the tests ran without `-mllvm -amdgpu-coerce-illegal-types=1`. - Perf A/B showed no material regression; worst measured slowdown was about `+2.47%`, with other representative WKV/GLM shapes flat or faster. - Manual pre-commit on changed CK files passed ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
fix(ck): fix gfx950 bpreshuffle determinism and illegal-type codegen ## Motivation Fix the CK `gemm_a8w8_blockscale_bpreshuffle` path behind: - ROCm/aiter#3261: nondeterministic `gemm_a8w8_blockscale_bpreshuffle` results on gfx950/MI350/MI355 for fixed inputs and launch shapes. - sgl-project/sglang#28685: GLM-5.2-FP8 can silently produce wrong rows on gfx950/ROCm 7.2 when out-of-allowlist block-FP8 linear shapes route to AITER's CK bpreshuffle path. The issue is isolated to the bpreshuffle CK kernel path on gfx950/ROCm 7.2, not general FP8 numerics. This PR updates the CK source so the kernel no longer depends on the removed LLVM illegal-type coercion workaround. ## Technical Details - Replace the packed 2-lane post-scale accumulator FMA update in bpreshuffle v1/v3 with scalar `AccDataType` updates. - Add a VGPR read/write accumulator anchor after each post-scale scalar update to enable deterministic repeated launches. - Add profiler/test support for: - repeated-run bytewise determinism checks - row-relative error checks to catch sparse bad-row failures near tile boundaries - GLM-5.2 representative/out-of-allowlist shapes Relevant issues: - ROCm/aiter#3261 - sgl-project/sglang#28685 ## Test Plan Validated on gfx950 with ROCm 7.2 using the CK profiler/test path. Coverage includes: - ROCm/aiter#3261 WKV determinism sweep: - `M=8192`, `N=512`, `K={384,640,3968,4096,4224}` - GLM-5.2-FP8 representative shapes: - `q_a_proj`, `q_b_proj`, `mlp.down_proj`, and `kv_b_proj` - row-boundary sweep including `M={32,48,56,64,72,96,128}`, `N=28672`, `K=512` - CPU reference comparison for correctness cases - row-relative error reporting for sparse bad-row failures - repeated-run determinism checks - before/after perf regression sweep - manual pre-commit run on changed CK files ## Test Result - Kernel passed the tests: - `ReportedWkvDeterminism` - `Glm52OutOfAllowlistAccuracyAndDeterminism` - Build/test compile command validation showed `COERCE_COUNT 0`, confirming the tests ran without `-mllvm -amdgpu-coerce-illegal-types=1`. - Perf A/B showed no material regression; worst measured slowdown was about `+2.47%`, with other representative WKV/GLM shapes flat or faster. - Manual pre-commit on changed CK files passed ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: SONLE5 <sonle5@OM-SONLE5-M5.local> Co-authored-by: Andriy Roshchenko <107577548+andriy-ca@users.noreply.github.com>
…P works, fp8 KV legal Replaces the single gfx950 cell with three and closes both gaps it declared. gfx942 is untouched. Rows generated by gen_glm52_mi355x_rows.py; each cell's published launch command is built from the argv.txt of the server that produced its numbers, so a recipe cannot drift from its own table. 1. The two mandatory bpreshuffle patches are retired on ROCm >= 7.2. The cell stated its own exit condition -- 'the CK rewrite ROCm/rocm-libraries#8639 supersedes the disable workaround; not in aiter 7d604afe5'. This image ships aiter d9e5ef7ce and #8639 is in its CK submodule, and sglang 0.5.17 gates rather than disables. A/B on identical argv: GSM8K 0.980 both arms, and the patched arm is slower at every concurrency on both request shapes (-8.8/-8.5/-9.0% at ISL 8192, -7.2/-7.0/-7.0% at 76k), with output-token counts identical to 0.00% in all six pairs. The gotcha is rewritten, not deleted: below ROCm 7.2 the patches are still required. 2. gaps[1] 'MTP not enabled on AMD' -- it is. GlmMoeDsaForCausalLMNextN, draft in the same checkpoint, one flag. Isolated against a control differing by exactly that flag: 2.58-2.76x decode from 8k to 256k for 1.03-1.06x prefill. 3. 'KV-cache dtype must be bfloat16 with DSA tilelang' is a CUDA rule. On ROCm fp8_e4m3 takes the pool 1,654,976 -> 3,194,368 tokens at +0.0pp GSM8K and -0.38% accept length. 4. gaps[0] 'balanced / high-throughput' -- both cells filled. Two things the cells say out loud rather than bury. The balanced recipe ships --mem-fraction-static 0.85, not 0.88: at 0.88 it aborts the process on a 131k prefill with HSA_STATUS_ERROR_OUT_OF_RESOURCES at a KV token usage of 0.04, and 0.85 clears 262k for 5.5% of pool and no measurable speed. And the benchmark rows are --dataset-name random, which this repo already established overstates accept length; degeneracy_probe.py run against GLM-5.2 returns 9 unique tokens of 512 with a top 8-gram repeated 27 times, and accept length is 3.994 on random against 2.969 on ShareGPT and 3.555 on GSM8K. The rows stay for comparability with the cell they replace; the cell says to size from real text. Three honest gaps replace the two closed ones: long-context accuracy under fp8 KV (unverified, with the diff harness to run), AIME25 not re-run on 0.5.17, and DP-attention + DeepEP still untried.
…y cells (#3) * feat(glm-5.2-fp8): add the MI355X 0.5.17 playbook and its row generator Re-measures the gfx950 cell on rocm/sgl-dev:v0.5.17-rocm724-mi35x-20260820 (ROCm 7.2.4, aiter d9e5ef7ce). Three of the existing cell statements have expired on that image, two of them being its own declared gaps; the playbook documents what changed and how to check each claim, and the generator turns bench_serving JSON into models.js rows the way gen_cookbook_rows.py does for Kimi-K3. No data change in this commit. The continuity check that makes the rest worth reading: the published recipe carried onto the new image reproduces its own c=16 row to 0.7% (531.7 vs 535.66 output tok/s, TPOT 25.01 vs 25.22), so everything else is a delta measured on one harness rather than two labs. * feat(site): glossary entries for the scheduling and speculative flags verify-cookbook.js fails a launch command containing a flag with no glossary entry, and the three new gfx950 recipes use four it does not have yet: --schedule-policy, --num-continuous-decode-steps, --speculative-num-steps and --speculative-eagle-topk. Also rewrites --kv-cache-dtype, which said bf16 is required with the DSA tilelang path. That is true on CUDA and not on ROCm: _check_tilelang_dsa_fp8_kv raises only when not hip, and the measured gfx950 pool roughly doubles on fp8_e4m3. * data(glm-5.2-fp8): re-measure gfx950 on 0.5.17 -- patches retired, MTP works, fp8 KV legal Replaces the single gfx950 cell with three and closes both gaps it declared. gfx942 is untouched. Rows generated by gen_glm52_mi355x_rows.py; each cell's published launch command is built from the argv.txt of the server that produced its numbers, so a recipe cannot drift from its own table. 1. The two mandatory bpreshuffle patches are retired on ROCm >= 7.2. The cell stated its own exit condition -- 'the CK rewrite ROCm/rocm-libraries#8639 supersedes the disable workaround; not in aiter 7d604afe5'. This image ships aiter d9e5ef7ce and #8639 is in its CK submodule, and sglang 0.5.17 gates rather than disables. A/B on identical argv: GSM8K 0.980 both arms, and the patched arm is slower at every concurrency on both request shapes (-8.8/-8.5/-9.0% at ISL 8192, -7.2/-7.0/-7.0% at 76k), with output-token counts identical to 0.00% in all six pairs. The gotcha is rewritten, not deleted: below ROCm 7.2 the patches are still required. 2. gaps[1] 'MTP not enabled on AMD' -- it is. GlmMoeDsaForCausalLMNextN, draft in the same checkpoint, one flag. Isolated against a control differing by exactly that flag: 2.58-2.76x decode from 8k to 256k for 1.03-1.06x prefill. 3. 'KV-cache dtype must be bfloat16 with DSA tilelang' is a CUDA rule. On ROCm fp8_e4m3 takes the pool 1,654,976 -> 3,194,368 tokens at +0.0pp GSM8K and -0.38% accept length. 4. gaps[0] 'balanced / high-throughput' -- both cells filled. Two things the cells say out loud rather than bury. The balanced recipe ships --mem-fraction-static 0.85, not 0.88: at 0.88 it aborts the process on a 131k prefill with HSA_STATUS_ERROR_OUT_OF_RESOURCES at a KV token usage of 0.04, and 0.85 clears 262k for 5.5% of pool and no measurable speed. And the benchmark rows are --dataset-name random, which this repo already established overstates accept length; degeneracy_probe.py run against GLM-5.2 returns 9 unique tokens of 512 with a top 8-gram repeated 27 times, and accept length is 3.994 on random against 2.969 on ShareGPT and 3.555 on GSM8K. The rows stay for comparability with the cell they replace; the cell says to size from real text. Three honest gaps replace the two closed ones: long-context accuracy under fp8 KV (unverified, with the diff harness to run), AIME25 not re-run on 0.5.17, and DP-attention + DeepEP still untried. * fix(glm-5.2-fp8): put total_tok_s on the convention the site documents DROP THIS COMMIT IF YOU DISAGREE -- it is the only one here that changes a number you published, and it is deliberately last so it lifts cleanly. If it goes, the new gfx950 rows should drop total_tok_s too, or the compare table's 'peak total' column reads a 19x MI300X -> MI355X gap that is a units change. app.js captions the column 'total tok/s adds input tokens', and gen_cookbook_rows.py emits it from bench_serving's total_throughput. The glm-5.2-fp8 rows carry aggregate OUTPUT throughput there instead. Three independent confirmations, no measurement involved: 1. The gfx942 cell holds the same c=1 point twice, once per convention: total_tok_s 47.8 in one row and 430.7 in the other. 47.8 x 9 = 430.2. 2. bench_serving defines total_throughput = (in+out)/duration and output_throughput = out/duration, so the ratio is exactly 1 + isl/osl independent of duration, and exact for --random-range-ratio 1.0. 3. Re-running the gfx950 recipe on a newer image reproduces its published c=16 figure as OUTPUT throughput to 0.7% (531.7 vs 535.66), while that run's actual total_throughput is 4785.6. So nothing measured is rewritten: the published figure moves to output_tok_s, where it is correct, and total_tok_s / tok_s_per_gpu are recomputed. The row that was already input-inclusive is detected and left alone -- which is also the check that the factor is right, since the recomputed twin lands on 430.2 against its 430.7.
Motivation
Fix the CK
gemm_a8w8_blockscale_bpreshufflepath behind:gemm_a8w8_blockscale_bpreshuffleresults on gfx950/MI350/MI355 for fixed inputs and launch shapes.The issue is isolated to the bpreshuffle CK kernel path on gfx950/ROCm 7.2, not general FP8 numerics. This PR updates the CK source so the kernel no longer depends on the removed LLVM illegal-type coercion workaround.
Technical Details
AccDataTypeupdates.Relevant issues:
Test Plan
Validated on gfx950 with ROCm 7.2 using the CK profiler/test path.
Coverage includes:
M=8192,N=512,K={384,640,3968,4096,4224}q_a_proj,q_b_proj,mlp.down_proj, andkv_b_projM={32,48,56,64,72,96,128},N=28672,K=512Test Result
ReportedWkvDeterminismGlm52OutOfAllowlistAccuracyAndDeterminismCOERCE_COUNT 0, confirming the tests ran without-mllvm -amdgpu-coerce-illegal-types=1.+2.47%, with other representative WKV/GLM shapes flat or faster.Submission Checklist