gfx1250: fix the grouped-MoE expert scan above 512 experts, and compute SiTUv2 instead of SiLU - #4482
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
There was a problem hiding this comment.
Pull request overview
This PR fixes two correctness issues in the gfx1250 grouped MoE path: (1) the contiguous-M per-expert prefix scan now correctly supports expert counts > 512 by sweeping in block-sized chunks with an LDS carry, and (2) ActivationType.Situv2 is now computed as SiTUv2 (not accidentally as SiLU) by wiring a new stage1_act=3 through both relevant epilogues.
Changes:
- Rework
moe_contiguous_psum/moe_contiguous_psum_remapto scan experts in chunks and carry offsets across chunks (removing the silent >512 expert drop). - Add SiTUv2 activation support to the gfx1250 TDM epilogue paths, with runtime beta/linear_beta parameters and host-side reciprocal folding.
- Expand tests/CLI to cover SiTUv2 and to regression-test contiguous-psum correctness and remap bounds across expert counts around 512.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| op_tests/test_flydsl_grouped_gemm_gfx1250.py | Adds/updates SiTUv2 tests and adds new contiguous-psum correctness + bounds tests; extends CLI to --act situv2 with beta args. |
| aiter/ops/flydsl/moe_common.py | Updates the SiTUv2 reference docstring to reflect the restored grouped TDM semantics. |
| aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py | Wires SiTUv2 into the gfx1250 TDM stage1 epilogue (batched + element-wise variants) and forwards beta args. |
| aiter/ops/flydsl/kernels/moe_contiguous_psum.py | Implements the chunked expert scan with LDS carry for both psum and psum+remap kernels. |
| aiter/ops/flydsl/kernels/gemm_common_gfx1250.py | Adds shared SiTUv2 math helpers/constants used by gfx1250 GEMM epilogues. |
| aiter/ops/flydsl/grouped_moe_gfx1250.py | Extends stage1 activation encoding to include stage1_act=3 for SiTUv2 and plumbs beta args through the grouped MoE flow. |
| aiter/ops/flydsl/batched_gemm_mxfp4.py | Extends the gfx1250 grouped GEMM launcher to pass SiTUv2 beta args (and their reciprocals) into the TDM kernel. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
…n still capped at 512 Review feedback on ROCm#4482. Drop f32_situ_inv_beta / f32_situ_inv_linear_beta from the TDM kernel and let situv2_consts() take both reciprocals with v_rcp_f32. Both are uniform across the tile, so this is two extra VALU ops per kernel, hoisted out of the inner loop, against two fewer kernel args and no way for a caller to pass a beta and a reciprocal that disagree. Verified numerically identical on the Kimi-K3 betas: a4w4 logits_diff 9.0759e-06 / rel_l2 4.2605e-03 and a8w4 4.4326e-06 / 2.9774e-03 both match the host-reciprocal build to every printed digit (beta=4.0 is a power of two so its rcp is exact; linear_beta=25.0 is not, and its ~1 ulp lands far below the MXFP4 quantisation it feeds). Rewrite the chunked-scan comment to say what it is actually guarding: one Hillis-Steele pass covers one expert per lane, so the old single-pass scan left starts/psum unwritten for every expert past 512, which is how K3's 896 reached the GEMM as garbage offsets and faulted. Sweeping the other single-block scans for the same cap: the route-quant-scatter prefix sum is single-thread serial over E (no cap) and moe_g2l_lut is gated at _G2L_MAX_N with a torch fallback, both fine. moe_route_psum_fused is genuinely capped -- its LDS route counter is one slot per expert, so E>512 needs a wider allocation, not a carry -- and _FUSED_ROUTE_PSUM_MAX_EXPERTS was defined but never enforced. Raise instead of silently dropping experts. The NUMEL companion is left advisory: that sweep is grid-stride, so a larger count is correct, just not worth fusing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
LGTM |
The tile-aligned prefix scan ran one thread per expert in a single MAX_EXPERTS_PER_BLOCK (512) block, so any model with more experts than that silently lost the tail. Kimi-K3 has 896: experts 512..895 were never scanned, `starts`/`psum`/`contiguous_m` kept the uninitialised values of their torch.empty allocation, the masked-to-contiguous row remap turned those into out-of-range rows, and the downstream moe_fused_quant_preshuffle_routeks_* faulted on them (HSA_STATUS_ERROR_MEMORY_FAULT). The cap was known -- the fused route+psum variant is gated on _FUSED_ROUTE_PSUM_MAX_EXPERTS -- but this path had no such guard and corrupted quietly instead. Sweep the experts in block-sized chunks and carry the running offset between them in LDS, so E is no longer bounded by the block width. The carry has to live in LDS rather than a register because the chunk loop is a runtime loop. Lanes past `experts` now feed 0 into the scan so the last lane still holds the chunk total, and the two kernels spell the sweep out separately because @flyc.kernel only AST-transforms the decorated body. Checked against a torch.cumsum reference over the tile-aligned counts for E = 8/256/512/896/1024: exact, where before E=896 mismatched on exactly 384 experts starting at 512.
The TDM stage1 act code only encoded silu and swiglu, so ActivationType.Situv2 fell through to `stage1_act = 1` and was computed as silu -- quietly, since situ_beta/situ_linear_beta were accepted at the grouped entry point and then dropped. Kimi-K3 is `hidden_act="situ"`, so on gfx1250 (where the separated path has no working SiTUv2 kernel) the model had no correct MoE at all. This is the TODO(situv2) left behind when the fused stage1 epilogue was removed. Add stage1_act=3 and wire it through both epilogues: the batched one used by the a8w4 fused-quant path and the element-wise one used by a4w4's bf16 intermediate. beta/linear_beta are runtime kernel arguments, so every SiTUv2 shape shares one compiled kernel; their reciprocals are taken on the host so folding them into the per-element multipliers stays exact, and the multipliers themselves are hoisted out of the inner loop. tanh uses the saturating identity 2*sigmoid(2z)-1 rather than the (1-e)/(1+e) form: exp2 of a large positive argument goes to +inf and rcp(+inf) to 0, so both tails are correct without an |x| fixup or a sign select. SiTUv2 is bounded by construction and takes no swiglu clamp. Un-skip the grouped SiTUv2 test and add the a8w4 case (a separate code path from a4w4); both land at rel_l2 ~3e-3 against the fp32 reference, which is MXFP4 quantisation noise, with silu and swiglu unchanged. Kimi-K3 end to end on 4xMI450 is GSM8K 1319 = 0.9591.
The scan's width is set by the expert count, not the token count, so it was the one part of the grouped-MoE pipeline with no coverage on the axis that actually breaks it. Nothing else in this file varies E far enough to notice: a dropped expert does not show up as a bad number, it shows up as a row index pointing outside the contiguous buffer, and then as a fault in whichever kernel dereferences that row next. Check starts/psum/contiguous_m against a tile-aligned torch.cumsum, and separately check that every remapped route row lands inside the buffer, at E = 8/256/512/513/896/1024 -- either side of MAX_EXPERTS_PER_BLOCK, including Kimi-K3's 896. Counts come from a real unbalanced random routing rather than a uniform split, so the per-expert values differ.
The bounds check sat at the top of flydsl_grouped_gemm_a8w4_masked, so it applied to silu and swiglu launches too -- where the betas are ignored and default to 1.0. A caller that passed a beta of 0 alongside a non-SiTUv2 activation would have been rejected for a parameter the kernel never reads. Gate it on stage1_act == 3.
…n still capped at 512 Review feedback on ROCm#4482. Drop f32_situ_inv_beta / f32_situ_inv_linear_beta from the TDM kernel and let situv2_consts() take both reciprocals with v_rcp_f32. Both are uniform across the tile, so this is two extra VALU ops per kernel, hoisted out of the inner loop, against two fewer kernel args and no way for a caller to pass a beta and a reciprocal that disagree. Verified numerically identical on the Kimi-K3 betas: a4w4 logits_diff 9.0759e-06 / rel_l2 4.2605e-03 and a8w4 4.4326e-06 / 2.9774e-03 both match the host-reciprocal build to every printed digit (beta=4.0 is a power of two so its rcp is exact; linear_beta=25.0 is not, and its ~1 ulp lands far below the MXFP4 quantisation it feeds). Rewrite the chunked-scan comment to say what it is actually guarding: one Hillis-Steele pass covers one expert per lane, so the old single-pass scan left starts/psum unwritten for every expert past 512, which is how K3's 896 reached the GEMM as garbage offsets and faulted. Sweeping the other single-block scans for the same cap: the route-quant-scatter prefix sum is single-thread serial over E (no cap) and moe_g2l_lut is gated at _G2L_MAX_N with a torch fallback, both fine. moe_route_psum_fused is genuinely capped -- its LDS route counter is one slot per expert, so E>512 needs a wider allocation, not a carry -- and _FUSED_ROUTE_PSUM_MAX_EXPERTS was defined but never enforced. Raise instead of silently dropping experts. The NUMEL companion is left advisory: that sweep is grid-stride, so a larger count is correct, just not worth fusing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
a87363b to
01f8c9f
Compare
…to updated main Reapply the gfx1250 fused-MoE ep_scatter feature (combine_mode="scatter_fused": gemm2's TDM epilogue P2P-writes each route-weighted output row into peers' comb_inp, so combine just sums -- no gather-reduce) onto the updated origin/main, which meanwhile landed #4482 (real SiTUv2 + >512-expert chunked psum scan), #4527 (LDS API refactor: lds_*_raw -> make_lds_copy_ops) and #4463 (fused_moe SiTUv2 path). Squashed to a single commit. Per-file resolution: - dispatch_combine_v2/*: vendored cco-LSA v2 intranode dispatch/combine op-layer. - mxfp4_preshuffle_gfx1250_tdm.py: on upstream's make_lds_copy_ops LDS API + SiTUv2 epilogue, add the TDM gather-store ep epilogue (tdm_scatter with in-kernel global_view/lds_view). tdm_scatter is vendored locally (tdm_gather_shim.py, on the stock FlyDSL wheel's low-level TDM intrinsics) so this branch needs no FlyDSL-side patch; route weight hoisted per wm row (_wf_rows). - moe_contiguous_psum.py: adopt upstream's chunked scan (E>512 correct) for the non-EP remap; keep the multi-block grid-stride remap + ep_rowmap kernels for EP. - grouped_moe_gfx1250.py: keep upstream SiTUv2 (stage1_act=3, situ_beta); add the ep_scatter dispatch wiring (ep_rowmap build, _ep_gemm2_kwargs, ep_scatter return). - batched_gemm_mxfp4.py / fused_moe.py: thread both situ_* and ep_* params. - tuned_grouped_fmoe.csv: tuning points (99 rows). Dropped only the full-subtile PF prefetch pipeline (da8d794): loaders and lds_addr_keepalive stay at upstream. lalala-sh's ds-read hoist is kept. gfx1250-only; not run in this environment. Compile-verify the LDS-API migration + ep/SiTUv2 epilogue and re-run test_mega_moe --combine on hardware. Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com> Co-authored-by: zhimding <zhimding@amd.com>
What this fixes
Two independent bugs in the gfx1250 grouped MoE, both found while bringing up
Kimi-K3 (
hidden_act="situ", 896 experts, top-16, MXFP4). Either one alonemakes the model unservable, and neither is K3-specific — the first is wrong for
any model with more than 512 experts, the second for any model using SiTUv2.
They are separate commits and can be reviewed independently.
1. The contiguous-M prefix scan silently dropped every expert past 512
moe_contiguous_psum.pyran the tile-aligned prefix scan as a single block ofMAX_EXPERTS_PER_BLOCK = 512threads, one thread per expert:With
E > 512the experts above the block width are never visited.starts,psumandcontiguous_mare allocated withtorch.empty, so those entries keepwhatever was in the allocator's memory;
is_last = tid == experts - 1neverfires, so
contiguous_mis not written at all. The remap phase then readsstarts[expert]for the unscanned experts and writesstart + slotback as theroute's row, and the next kernel to dereference that row —
moe_fused_quant_preshuffle_routeks_*— walks off the end of the contiguousbuffer:
The 512 limit was known for the fused route+psum variant, which is gated on
_FUSED_ROUTE_PSUM_MAX_EXPERTS = 512ingrouped_moe_gfx1250.py.contiguous_psum/contiguous_psum_remaphad no equivalent guard, so insteadof falling back they corrupted quietly.
Fix: sweep the experts in block-sized chunks and carry the running offset
between chunks in LDS. E is no longer bounded by the block width.
Three details worth a reviewer's attention:
runtime loop (
expertsis a kernel argument) — a Python-level accumulatordoes not survive it.
expertsnow write0into the scan buffer and participate in theHillis-Steele steps, instead of being skipped by
if in_range. That is neededso the last lane still holds the chunk total regardless of where
expertsfalls inside the chunk.
@flyc.kernelonly AST-transforms the decorated body, so a dynamicfor/ifmoved into a plain function stops being traced (it raises
dynamic 'ArithValue' has no Python integer representation).2. SiTUv2 was computed as SiLU
The TDM stage1 epilogue only encoded silu and swiglu:
ActivationType.Situv2fell through to the silu code, andsitu_beta/situ_linear_betawere accepted atgrouped_gemm_gfx1250_a8w4's signature andthen dropped on the way to
_grouped_a8w4_tdm_moe. No error, no fallback — justthe wrong activation. This is the
TODO(situv2)left behind when the fusedstage1 epilogue was removed in the MoE refactor (#4394).
On gfx1250 this is the only SiTUv2 path there is:
GateMode.SEPARATEDroutesSiTUv2 to
flydsl_moe1_afp4_wfp4_bf16_*, which currently fails to build there(
LLVM ERROR: Do not know how to expand this operator's operand, an i64 operandin
llvm.amdgcn.raw.ptr.buffer.load.lds) — that is a separate issue, notaddressed here.
Fix: add
stage1_act = 3and wire it through both epilogues — the batchedone used by a8w4's fused-quant path, and the element-wise one used by a4w4's
bf16 intermediate.
Implementation notes:
beta/linear_betaare runtime kernel arguments, so every SiTUv2 shapeshares one compiled kernel rather than specializing per beta value. Their
reciprocals are taken on the host so folding them into the per-element
multipliers stays exact instead of going through an in-kernel
v_rcp_f32, andthe multipliers are hoisted out of the inner loop.
tanhuses the saturating identity2*sigmoid(2z) - 1rather than the(1-e)/(1+e)form used elsewhere in the tree:exp2of a large positiveargument goes to
+infandrcp(+inf)to0, so both tails come out rightwithout an
|x|fixup and a signselect. Costs 3 exp2 + 3 rcp per element.Tests
op_tests/test_flydsl_grouped_gemm_gfx1250.py:test_contiguous_psum_matches_cumsumandtest_contiguous_psum_remap_rows_stay_in_bounds, parametrized overE = 8/256/512/513/896/1024— either side ofMAX_EXPERTS_PER_BLOCK. Countscome from a real unbalanced random routing, checked against a tile-aligned
torch.cumsum. The second test asserts every remapped row lands inside thecontiguous buffer, which is the property whose violation causes the fault.
test_grouped_a4w4_situv2_matches_torch_ref, which was markedskipwith "the TDM path runs it as silu".test_grouped_a8w4_situv2_matches_torch_ref, since a8w4 takes thebatched fused-quant epilogue rather than a4w4's element-wise one.
--act situv2added to the CLI, with--situ-beta/--situ-linear-beta.The scan tests fail on
mainexactly where they should, and pass here:Activation correctness against the fp32 torch reference (gate is
logits_diff < 0.01):rel_l2 ~3e-3is MXFP4 quantisation noise; it is the same for all four.Performance
The SiTUv2 epilogue is
const_expr-gated, so silu/swiglu codegen is unchanged —their
logits_diff/rel_l2above are bit-identical betweenmainand thisbranch. The two things that are shared are the scan and the four extra runtime
f32 kernel arguments, so both were measured.
Grouped MoE,
E=256 topk=8 model_dim=7168 inter_dim=512, gfx1250:Everything is within ±2% with both signs, i.e. run-to-run noise.
The scan kernels timed on their own, which is where a real cost would show:
contiguous_psum_remap— the one the grouped MoE actually calls — is withinnoise.
contiguous_psumis consistently ~0.15 us slower (+3.5..8.5% on a 2.6 uskernel), which is the two extra barriers and the LDS carry round-trip on a
single-chunk sweep. Reported rather than hidden: it is once per MoE call, so for
a 93-layer model that is ~14 us per forward against a ~27 ms decode step, but a
reviewer may prefer a compile-time single-chunk specialisation instead.
End to end
Kimi-K3 on 4×MI450 (tp4) through ATOM, full 1319-question GSM8K 5-shot,
num_concurrent=8:Same score as the last recorded run of this model on gfx1250, which needed a
model-side workaround (sub-batching the routed MoE to <=128 tokens) to avoid the
fault in 1 — that workaround is removed on the ATOM side now that the scan is
correct.