Conversation
Relands sgl-project#31323, which sgl-project#35105 reverted after it broke test_global_config_read_ratchet.py on main and blocked unrelated PRs. The revert was correct. The alias-form ratchet baseline dropped to 0 on Aug 5, after this branch's last green CI run had passed against a baseline of 12, so the two ServerArgs reads only became a violation once the branch was merged against current main. Three fixes on top of the original change: - _get_aiter_topk_fuse_shared_max_tokens reads chunked_prefill_size and max_prefill_tokens from get_schedule() instead of aliasing the process-global ServerArgs (from sgl-project#35075, which was closed as unmergeable against the reverted tree). - The moe_ep_size read that gates the fused path degrades to the plain append path when the MoE EP group is not initialized, instead of asserting. Both the check in biased_grouped_topk_gpu and its mirror in select_experts go through the same helper so the two stay consistent. - biased_grouped_topk_cpu accepts fused_shared_experts_scaling_factor and rejects a non-None value, matching how fused_topk_cpu handles the parameters it cannot serve. select_experts passes the kwarg unconditionally, so without it the CPU+AMX alias raises TypeError during argument binding even when the value is None.
…tical test The pre-populated buffer had no test: nothing asserted it returns what the plain append would have, and nothing exercised the fallbacks, so the uninitialized-EP-group regression only surfaced by hand on an MI355X. Compare the two spellings on the DeepSeek-V3 / GLM-5 routing shape across a decode step and three prefill chunk sizes, each carrying a different routed scaling factor (unset, 1.0, 1.5, 2.5), under both fused_shared_experts_scaling spellings. Also pin the two conditions that must drop to the plain append: expert parallelism, and reaching the path before the MoE EP group exists.
The three blocks each restated the whole scheme, so the bit-identical claim and the fallback rule appeared three times over. Keep each where it is load-bearing and drop the narration. Also invert the append guard: the empty branch existed only to hold a comment that repeated the one above it.
90d6745 to
a75d427
Compare
|
/rerun-failed-ci |
|
|
@amd-bot ci-status |
CI Status for PR #35113Merge verdict: From this PR's perspective it looks safe — the changed code ( Note This PR's changed code is covered by PR CI. The added test Caution PR CI is incomplete. Changed files: Executed CI failure attribution: AMD: 4 executed failures (0 related) · Others (CPU/NPU): 9 executed failures (0 related). The PR-claimed CPU ratchet fix worked — no AMD Executed Failures (run 32087142949, ROCm 7.2)
( Other Executed Failures
Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
Keep the aiter shared-expert fuse buffer after main's radix4 early-return and fp32-bias aiter path; the two edits occupy the same allocation site.
|
/rerun-failed-ci |
3 participants
Summary
Relands #31323, which #35105 reverted after it failed
base-a-test-cpu (0)on main. That failure is fixed here, along with two latent bugs found while re-testing on MI355X, and the optimization finally gets the test coverage it shipped without.What the optimization does
On GLM-5.2 with the non-EP aiter grouped-topk route, every decoder layer ran a separate kernel to append the shared expert into the top-k ids/weights. Instead, the shared columns are pre-populated once in a persistent top-k buffer and the aiter kernel writes only the routed columns via row stride, so the per-layer append kernel disappears (prefill and decode).
moe_ep_size == 1andnum_fused_shared_experts > 0. Other archs, EP mode and non-applicable routes are untouched.GSM8K on MI355X TP4 (GLM-5.2-MXFP4): 0.931 → 0.933. Throughput and TPOT move under 1% either way — the gain is one fewer kernel launch per layer, not raw speed.
Fixes on top of #31323
1. Read the sizing config through
get_schedule()— the actual revert cause.base-a-test-cpu (0)failed withalias-form process-global config field reads grew: 2 > baseline 0. The ratchet baseline dropped to 0 on Aug 5, after this branch last ran green against a baseline of 12, so the twoServerArgsalias reads only became a violation once merged against current main.2.
biased_grouped_topk_cpuacceptsfused_shared_experts_scaling_factor.select_expertspasses the kwarg unconditionally, so on CPU + AMX the alias raisedTypeErrorduring argument binding even when the value wasNone. A non-Nonevalue is rejected, matching howfused_topk_cpuhandles parameters it cannot serve.New test:
TestAiterGroupedTopkSharedFuse#31323 shipped with nothing exercising the fused path —
fused_shared_experts_scaling_factorhad no non-default coverage anywhere in the repo, which is why both latent bugs above had to be found by hand.test/registered/moe/test_fused_append_remap_per_rank_shared_slots.pynow adds three tests:torch.equal, not merely close. Covered over a decode step and three prefill chunk sizes on the DeepSeek-V3 / GLM-5 routing shape, under bothfused_shared_experts_scaling_factorspellings and four routed scaling factors.Validation
On MI355X everything related to this change is green: the new tests (7 passed, 16 subtests) and the ratchet suite.
CI States
Latest PR Test (Base): ✅ Run #35053826506
Latest PR Test (Extra): ❌ Run #35053826394
Latest PR Test (AMD ROCm 10): ❌ Run #35053826475