gfx1250: fix the grouped-MoE expert scan above 512 experts, and compute SiTUv2 instead of SiLU by XiaobingSuper · Pull Request #4482 · ROCm/aiter

XiaobingSuper · 2026-07-31T11:53:48Z

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 alone
makes 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.py ran the tile-aligned prefix scan as a single block of
MAX_EXPERTS_PER_BLOCK = 512 threads, one thread per expert:

in_expert = tid < fx.Uint32(experts)   # tid < 512 always
if in_expert:
    m = buffer_load(m_rsrc, tid, ...)  # expert index == thread index

With E > 512 the experts above the block width are never visited. starts,
psum and contiguous_m are allocated with torch.empty, so those entries keep
whatever was in the allocator's memory; is_last = tid == experts - 1 never
fires, so contiguous_m is not written at all. The remap phase then reads
starts[expert] for the unscanned experts and writes start + slot back as the
route's row, and the next kernel to dereference that row —
moe_fused_quant_preshuffle_routeks_* — walks off the end of the contiguous
buffer:

Memory access fault by GPU node-N ... kernel: moe_fused_quant_preshuffle_routeks_fd3584_r4_fp4_pk8_srctk16

The 512 limit was known for the fused route+psum variant, which is gated on
_FUSED_ROUTE_PSUM_MAX_EXPERTS = 512 in grouped_moe_gfx1250.py.
contiguous_psum / contiguous_psum_remap had no equivalent guard, so instead
of 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:

The carry lives in LDS rather than a register because the chunk loop is a
runtime loop (experts is a kernel argument) — a Python-level accumulator
does not survive it.
Lanes past experts now write 0 into the scan buffer and participate in the
Hillis-Steele steps, instead of being skipped by if in_range. That is needed
so the last lane still holds the chunk total regardless of where experts
falls inside the chunk.
The sweep is written out in both kernels rather than factored into a helper.
@flyc.kernel only AST-transforms the decorated body, so a dynamic for/if
moved 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:

stage1_act = 2 if activation == ActivationType.Swiglu else 1   # Situv2 -> 1 -> silu

ActivationType.Situv2 fell through to the silu code, and situ_beta /
situ_linear_beta were accepted at grouped_gemm_gfx1250_a8w4's signature and
then dropped on the way to _grouped_a8w4_tdm_moe. No error, no fallback — just
the wrong activation. This is the TODO(situv2) left behind when the fused
stage1 epilogue was removed in the MoE refactor (#4394).

On gfx1250 this is the only SiTUv2 path there is: GateMode.SEPARATED routes
SiTUv2 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 operand
in llvm.amdgcn.raw.ptr.buffer.load.lds) — that is a separate issue, not
addressed here.

Fix: add stage1_act = 3 and wire it through both epilogues — the batched
one used by a8w4's fused-quant path, and the element-wise one used by a4w4's
bf16 intermediate.

Implementation notes:

beta / linear_beta are runtime kernel arguments, so every SiTUv2 shape
shares 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, and
the multipliers are hoisted out of the inner loop.
tanh uses the saturating identity 2*sigmoid(2z) - 1 rather than the
(1-e)/(1+e) form used elsewhere in the tree: exp2 of a large positive
argument goes to +inf and rcp(+inf) to 0, so both tails come out right
without an |x| fixup and a sign select. Costs 3 exp2 + 3 rcp per element.
SiTUv2 is bounded by construction, so the swiglu clamp does not apply to it.

Tests

op_tests/test_flydsl_grouped_gemm_gfx1250.py:

New: test_contiguous_psum_matches_cumsum and
test_contiguous_psum_remap_rows_stay_in_bounds, parametrized over
E = 8/256/512/513/896/1024 — either side of MAX_EXPERTS_PER_BLOCK. Counts
come from a real unbalanced random routing, checked against a tile-aligned
torch.cumsum. The second test asserts every remapped row lands inside the
contiguous buffer, which is the property whose violation causes the fault.
Un-skipped: test_grouped_a4w4_situv2_matches_torch_ref, which was marked
skip with "the TDM path runs it as silu".
New: test_grouped_a8w4_situv2_matches_torch_ref, since a8w4 takes the
batched fused-quant epilogue rather than a4w4's element-wise one.
--act situv2 added to the CLI, with --situ-beta / --situ-linear-beta.

The scan tests fail on main exactly where they should, and pass here:

$ pytest op_tests/test_flydsl_grouped_gemm_gfx1250.py -k contiguous_psum   # on main
FAILED test_contiguous_psum_matches_cumsum[513]
FAILED test_contiguous_psum_matches_cumsum[896]
FAILED test_contiguous_psum_matches_cumsum[1024]
FAILED test_contiguous_psum_remap_rows_stay_in_bounds[513]
FAILED test_contiguous_psum_remap_rows_stay_in_bounds[896]
FAILED test_contiguous_psum_remap_rows_stay_in_bounds[1024]
6 failed, 6 passed

Activation correctness against the fp32 torch reference (gate is
logits_diff < 0.01):

case	logits_diff	rel_l2
a4w4 situv2	3.8e-06	2.8e-03
a8w4 situv2	5.3e-06	3.3e-03
a4w4 silu (regression)	3.7e-06	2.7e-03
a4w4 swiglu (regression)	3.8e-06	2.8e-03

rel_l2 ~3e-3 is 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_l2 above are bit-identical between main and this
branch. 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:

case	tokens	main (us)	this PR (us)
silu, gemm1	128	146.39	145.96
silu, gemm2	128	111.82	113.05
silu, gemm1	1024	149.41	149.45
silu, gemm2	1024	130.98	130.76
silu, end-to-end	128	289.88	292.68
silu, end-to-end	1024	370.51	362.87
swiglu, gemm1	128	146.19	146.65
swiglu, gemm2	128	112.27	112.95
swiglu, gemm1	1024	149.64	149.97
swiglu, gemm2	1024	130.33	131.27
swiglu, end-to-end	128	293.82	290.66
swiglu, end-to-end	1024	360.58	362.80

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:

E	psum main	psum PR	remap main	remap PR
8	2.47	2.64	5.85	5.76
128	2.48	2.69	6.74	6.58
256	2.58	2.67	6.44	6.86
512	2.62	2.77	6.88	6.65
896	faults	4.32	faults	8.20

contiguous_psum_remap — the one the grouped MoE actually calls — is within
noise. contiguous_psum is consistently ~0.15 us slower (+3.5..8.5% on a 2.6 us
kernel), 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:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9591|±  |0.0055|
|     |       |strict-match    |     5|exact_match|↑  |0.9591|±  |0.0055|

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.

github-actions · 2026-07-31T11:54:05Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4482 --add-label <label>

Copilot

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_remap to 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.

github-actions · 2026-07-31T12:11:17Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4482 --add-label <label>

yadaish · 2026-08-03T11:51:43Z

LGTM

coderfeli

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>

…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>

…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), 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>

…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), 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> (cherry picked from commit 0656977)

…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), 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>

* gfx1250 MoE: fused gemm2 P2P scatter-combine (ep_scatter), rebased onto 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), 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> * fix rebase error * combine perf optimization Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com> * feat(ep): push-group fixed-slot dispatch->GEMM1->GEMM2 scatter (gfx1250 a8w4) End-to-end fixed-slot push path for the a8w4 MoE, off by default behind an explicit config switch. Dispatch lands tokens grouped per local expert (no consumer-side gather), a finalize kernel builds the tile schedule, GEMM1 reads A contiguously from the fixed slots, and the GEMM2 TDM epilogue P2P-scatters each route-weighted row into peers' comb_inp via the dispatch-emitted pg_rowmap; combine then just sums. ~4% over the pull (gather) path at DeepSeek and balanced shapes. Squash of the push-group cherry-picks (dispatch base / GEMM1 A-load / GEMM2+combine wiring + parallel finalize + tile_k2) plus these refinements: - Explicit switch instead of env: EpDispatchCombineConfig.push_group (was AITER_EP_PUSH_GROUP); test_mega_moe gains a --push_group flag. - cap_per_expert (was push_group_cap / push_group_safety): 0 => worst-case ws*max_tok_per_rank (never drops out of the box); >0 is a caller-pinned, tile_m-aligned capacity whose correctness the caller owns (overflow dropped in finalize). Auto/safety sizing removed. Padding M-tiles early-exit at the `expert < n_experts` guard (no N/K mainloop), so over-provisioning costs empty workgroup dispatch, not padded compute. - SiTUv2 fixed in the fused GEMM1 path: map Situv2 -> stage1_act=3 and forward situ_beta / situ_linear_beta (previously silently ran as silu). Co-authored-by: XingerZhu <xzhu@amd.com> * opt redandunt ep tokens * updat * rm fill 0 * fix stage1 * update dispatch&combine v2 * Revert "fix stage1" This reverts commit d3b3fb4. * Revert "feat(ep): push-group fixed-slot dispatch->GEMM1->GEMM2 scatter (gfx1250 a8w4)" This reverts commit 90b0d59. * refactor magemoe * update * fix refactor * mega_moe: add a mori HIP dispatch backend MegaMoEGfx1250 owns dispatch -> gemm1 -> gemm2-fused-scatter -> fused combine. This makes the first stage swappable: dispatch_backend="mori" (or MEGA_DISPATCH=mori) runs mori's HIP/JIT gfx1250 TDM dispatch instead of this package's FlyDSL one, and nothing else in the pipeline changes. It works because the two dispatches leave byte-identical state behind: disp_out rows at slot*hidden, out_idx/out_wts at slot*topk+k, the flat dest map as dest_pe*max_recv+slot with null = npes*max_recv, and -- the one field the fused path actually consumes -- recv_to_src_token encoded src_pe*max_tok_per_rank+src_tok, which the GEMM host pass decodes to build ep_rowmap. The recv_num/tok_off signal/ack handshake is the same protocol, and mori's dispatch never touches cross_device_barrier, so the fused combine's phase counter is undisturbed. mori's plan layer takes an arena by duck typing (.handle/.offset), so SymmetricArena is handed over as-is; no extra region and no extra device memory. Two things do differ and are handled here: * geometry comes from mori's own tuning table, not _select_dispatch_config. That table asks for 32 warps above 256 tokens, and mori's dispatch stages a hidden-dim tile per warp in dynamic LDS -- 32*7168*2 = 458 KB against a 320 KB budget. EpCfgIsValid does not check LDS, so it would fail at launch rather than when the plan is built. * total_recv is zeroed here. This package's dispatch clears it in its own Phase 2; mori's only accumulates, and the fused combine never resets it. The recv slot a token lands in does change -- mori's gfx1250 dispatch reserves a block's slots with one atomic and hands them out block-local. Nothing indexes by slot order, but a test diffing arena contents slot-by-slot against the FlyDSL dispatch will see it. Measured on 4x gfx1250, EP4 hidden 7168 topk 6, scatter_fused, against test_mega_moe.py's fp32 reference: 2 layers, 256 tok/rank 945.9 -> 925.8 us/layer logits_diff 0.002174 both 4 layers, 4096 tok/rank 2678.0 -> 2622.2 us/layer logits_diff 0.004344 both Identical diffs: dispatch only moves data, so the deviation from the reference is unchanged. Profile confirms the geometry actually launched is mori's -- 64x8 at 256 tokens, 64x16 at 4096. Needs a mori built with JIT v2 (PR #548 or later) and its libmori_ops_v2.so. * split combine * fix e2e: mori_dispatch + megamoe * codecheck * fix code check * fix test script * clean code * rm src_token_map d2d * refactor --------- Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com> Co-authored-by: zhimding <zhimding@amd.com> Co-authored-by: XingerZhu <xzhu@amd.com> Co-authored-by: jhchouuu <jiahzhou@amd.com> Co-authored-by: valarLip <340077269@qq.com>

…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), 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>

XiaobingSuper requested review from a team and Copilot July 31, 2026 11:53

Copilot started reviewing on behalf of XiaobingSuper July 31, 2026 11:54 View session

XiaobingSuper closed this Jul 31, 2026

XiaobingSuper deleted the xiaobing/gfx1250-k3-situ-moe branch July 31, 2026 11:54

Copilot AI reviewed Jul 31, 2026

View reviewed changes

Comment thread aiter/ops/flydsl/batched_gemm_mxfp4.py Outdated

XiaobingSuper restored the xiaobing/gfx1250-k3-situ-moe branch July 31, 2026 12:06

XiaobingSuper reopened this Jul 31, 2026

XiaobingSuper requested review from coderfeli and yadaish July 31, 2026 12:11

yadaish reviewed Jul 31, 2026

View reviewed changes

Comment thread aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py

yadaish reviewed Jul 31, 2026

View reviewed changes

Comment thread aiter/ops/flydsl/kernels/moe_contiguous_psum.py

yadaish requested review from XingerZhu and lalala-sh July 31, 2026 12:48

XiaobingSuper requested a review from yadaish August 3, 2026 05:03

coderfeli previously approved these changes Aug 3, 2026

View reviewed changes

XiaobingSuper and others added 5 commits August 3, 2026 15:26

zejunchen-zejun dismissed coderfeli’s stale review via 01f8c9f August 3, 2026 15:37

zejunchen-zejun force-pushed the xiaobing/gfx1250-k3-situ-moe branch from a87363b to 01f8c9f Compare August 3, 2026 15:37

XiaobingSuper requested a review from coderfeli August 4, 2026 01:27

coderfeli approved these changes Aug 4, 2026

View reviewed changes

XiaobingSuper merged commit e997844 into ROCm:main Aug 4, 2026
42 checks passed

Navigation Menu

Uh oh!

gfx1250: fix the grouped-MoE expert scan above 512 experts, and compute SiTUv2 instead of SiLU - #4482

gfx1250: fix the grouped-MoE expert scan above 512 experts, and compute SiTUv2 instead of SiLU#4482
XiaobingSuper merged 5 commits into
ROCm:mainfrom
XiaobingSuper:xiaobing/gfx1250-k3-situ-moe

XiaobingSuper commented Jul 31, 2026

Uh oh!

github-actions Bot commented Jul 31, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

github-actions Bot commented Jul 31, 2026

Uh oh!

Uh oh!

Uh oh!

yadaish commented Aug 3, 2026

Uh oh!

coderfeli left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

XiaobingSuper commented Jul 31, 2026

What this fixes

1. The contiguous-M prefix scan silently dropped every expert past 512

2. SiTUv2 was computed as SiLU

Tests

Performance

End to end

Uh oh!

github-actions Bot commented Jul 31, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

github-actions Bot commented Jul 31, 2026

🏷️ CI Guide

Uh oh!

Uh oh!

Uh oh!

yadaish commented Aug 3, 2026

Uh oh!

coderfeli left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants