Repository navigation
Work around all-zero fused MoE output at block_m=32 - #4451
jatseng-ai wants to merge 1 commit into
Conversation
The CK 2-stage MoE kernel writes nothing at block_m=32: fused_moe returns an identically zero tensor rather than failing or producing NaNs. Reproduced on gfx942 for bf16/QuantType.No and fp8/QuantType.per_Token. The fault depends on block_m alone, not on the token count. Forcing block_m=32 at token=64, a size that is otherwise correct, also yields zeros, while 64 and 128 are correct at every size tested. This is severe in serving because get_block_size_M selects 32 for token <= 32, i.e. every decode step, while prefill takes larger tiles and stays correct. A MoE model then retrieves context correctly and degenerates into repetition while generating, since the routed-expert contribution is silently dropped and only shared experts and attention still contribute. On an internal long-context accuracy suite this showed up as 9/10 degenerate completions at 4k and ~23/50 at 131k, with the retrieval target still found in every case. The clamp is placed after both the heuristic and the tuned-CSV branch assign block_m, because either can select 32 and 469 of the 617 rows in configs/tuned_fmoe.csv pin block_m=32 with run_1stage=0. AITER_MIN_BLOCK_M restores the original selection for anyone bisecting the underlying kernel bug. 1stage and asm paths are untouched. This is a workaround, not a fix. The CK 2-stage sources under csrc/py_itfs_ck and get_block_size_M itself are byte-identical between v0.1.13.post1 and v0.1.18, so the same source produces a working kernel in one build and a no-op in the other, which points at codegen or build configuration rather than the algorithm. Add op_tests/test_moe_blockm_invariance.py asserting that fused_moe results do not depend on block_m. The existing MoE tests miss this because they only exercise whichever tile the heuristic happens to pick, and the failure mode is an all-zero output rather than an exception, so finiteness checks pass. Co-authored-by: Cursor <cursoragent@cursor.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
Review: PR #4451 — clamp block_m ≥64 to mask CK 2-stage all-zero bug
|
|
Root cause found — and it's already fixed upstream in CK; this can be a CK-submodule bump instead of a block_m clamp. The all-zero output isn't a block_m=32 dispatch/tiling issue — it's a miscompile-free copy-paste bug in a CK refactor. Already fixed upstream. ROCm/composable_kernel develop commit f33252ce That commit is a descendant of the CK submodule this branch pins (b675945), so this tree predates the fix — hence the all-zero decode. Suggested change. Bump the CK submodule to include f33252ce (or cherry-pick it) and drop the block_m≥64 clamp. This also covers the fp8-per-token and blockscale gate-up paths, which the clamp only helps incidentally. Verification (gfx950, standalone CK repro of the 32-tile gemm1 gate-up instance): device_absmax goes 0 → 302.25 with the fix; the block_m=128 control is unaffected. |
|
Closing this workaround PR — it is superseded.
Thanks for the review and pointer to the proper fix! |
3 participants
Summary
The CK 2-stage fused-MoE kernel returns an identically zero output tensor at
block_m=32on gfx942. It does not raise, does not produce NaNs, and logs nothing —fused_moesimply returns all zeros.get_block_size_Mselectsblock_m=32for token counts <= 32, which in serving means every decode step. Prefill uses larger tiles and is correct. The result is a MoE model that reads context correctly and then degenerates during generation, because the entire routed-expert contribution is silently dropped while attention and (unfused) shared experts keep the output superficially fluent.This PR clamps
block_mto >= 64 on the 2-stage path and adds a regression test. It is a workaround, not a root-cause fix — see "Open question" below.Reproducer
Standalone, ~2 min, no model download. bf16 /
token | v0.1.18 absmax | state | v0.1.13.post1 absmax | selected block_m -- | -- | -- | -- | -- 1 | 0.000000 | all zero | 0.012573 | 32 8 | 0.000000 | all zero | 0.015259 | 32 32 | 0.000000 | all zero | 0.020142 | 32 33 | 0.022949 | ok | 0.022949 | 64 128 | 0.020752 | ok | 0.020752 | 128QuantType.No,model_dim=4096, inter_dim=512, E=128, topk=8, identical seeds on both versions:Test
op_tests/test_moe_blockm_invariance.pyasserts thatfused_moeresults do not depend onblock_m, which is a scheduling knob and must not change what is computed. It covers bf16 and fp8/per_Tokenat token 1 (decode) and token 64 (to decouple tile size from token count).Two details that matter for anyone extending it:
The failure mode is an all-zero tensor, not an exception or a NaN, so a test that only checks for a successful launch or for finiteness passes straight through it. The test asserts non-zero output explicitly.
Kernel selection is
functools.lru_cached, so changingblock_min-process without clearing those caches silently reuses the first decision and the test passes vacuously. The helper clears them between runs.Open question for maintainers (the actual root cause)
I could not find the source-level cause, and the evidence suggests it is not in the algorithm:
csrc/py_itfs_ck/(includingmoe_ck_2stages_kernel.cuandmoe_ck_2stages_gemm_impl/) is byte-identical between v0.1.13.post1 and v0.1.18 —diff -rqreports no differences for the MoE kernels.get_block_size_Mis byte-identical as well, so both versions select the same tile for the same shape.The
csrc/ck_gemm_moe_2stages_codegen/directory differs only ingemm_moe_tune.py.Identical sources therefore produce a working kernel in one build and a silent no-op in the other, which points at codegen or build configuration rather than the kernel logic. Someone familiar with the CK instantiation path can likely find this quickly; the reproducer above should make it a short loop.
Scope of validation
Stated plainly so reviewers can judge:
gfx942 only. gfx950 untested.
One shape (
E=128, topk=8, model_dim=4096, inter_dim=512). Other shapes untested.block_m=16untested — the clamp raises it to 64 on the assumption that a smaller tile is at least as broken, which is unverified.The clamp gives small batches a larger tile than the cost heuristic chose. A decode-throughput A/B is in progress; numbers will be added here.