mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-09 00:27:01 +00:00
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: - https://github.com/ROCm/aiter/issues/3261 - https://github.com/sgl-project/sglang/issues/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.