[FlyDSL] [CI] [feat] mega-moe stage2 for gfx1250 by yanboshao · Pull Request #4785 · ROCm/aiter

yanboshao · 2026-08-16T11:08:20Z

Motivation

Enable overlap between Stage 2 GEMM2 and EP combine on gfx1250 to reduce MegaMoE latency. This PR only covers GEMM2/combine overlap; other stages remain unchanged.

Technical Details

Fuse GEMM2 with combine P2P writes.
Split cross-rank synchronization and reduce into separate kernels to balance the performance of both stages.

Test Plan

Simulate an EP+DP workload on four gfx1250 GPUs. Run 61 sequential MoE layers in a single CUDA Graph, covering all operators from dispatch to combine. Layers are connected with RMSNorm and residual connections to maintain stable activations and exercise continuous buffer reuse. Each layer uses randomly generated routing.

Test Result

DSV4 pro

token_per_rank	base/layer (μs)	mega_fused_stage2/layer (μs)	ratio
8	268.7	242.4	9.8%
32	411.3	348.0	15.4%
128	490.2	426.9	12.9%
512	655.5	578.2	11.8%
2048	1569.6	1367.4	12.9%
8192	5536.3	4914.8	11.2%

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

github-actions · 2026-08-16T11:08:50Z

🏷️ 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:gfx1250-ffm-triton`	Run the five-shard gfx1250 FFM Triton test suite
`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 4785 --add-label <label>

github-actions

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

ruff

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

aiter/op_tests/multigpu_tests/test_mega_moe.py

Lines 32 to 46 in cb929e1

    
           import argparse 
        
           import os 
        
           import torch 
        
           import torch.distributed as dist 
        
           import torch.profiler as tprof 
        
           import aiter 
        
           from aiter import dtypes 
        
           from aiter import ActivationType, QuantType, get_gfx 
        
           from aiter.fused_moe import fused_moe 
        
           from aiter.ops.shuffle import shuffle_weight, moe_shuffle_scale 
        
           from aiter.ops.flydsl.moe_common import GateMode 
        
           from aiter.utility import fp4_utils 
        
           from aiter import get_hip_quant, get_torch_quant, pertoken_quant

⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

aiter/op_tests/multigpu_tests/test_mega_moe.py

Line 50 in cb929e1

except Exception: # pragma: no cover

⚠️ [ruff] <RUF059> _{reported by reviewdog 🐶}
Unpacked variable E is never used

aiter/op_tests/multigpu_tests/test_mega_moe.py

Line 160 in cb929e1

E, two_inter = t.shape[:2]

⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

aiter/op_tests/multigpu_tests/test_mega_moe.py

Line 782 in cb929e1

except Exception as _e:

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

Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com>

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

This reverts commit d3b3fb4.

…r (gfx1250 a8w4)" This reverts commit 90b0d59.

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.

Merge origin/yanbo/mega_stage2_gfx1250 onto the branch that already carries PR #4730 MoE GEMM optimizations. Conflicting flydsl paths take the stage2 implementation (ep scatter, MegaMoE dispatch/combine, updated TDM epilogue). Co-authored-by: Cursor <cursoragent@cursor.com>

…orked --dispatch_commu_dtype has had no effect since the file was created in ROCm#4785: resolve_spec wrote transport/prequant/is_fp8/fp8_dtype into the spec and nothing ever read them, quant_tokens_fp8() had no callers, and the base combine hardcodes transport_dtype = bf16. So this is not two overlapping knobs -- it is one live knob and one corpse. --mega_wire takes over the name (--dispatch_wire, matching the kwarg it feeds) and the dead half goes with the flag it belonged to. _FP8_DTYPE stays: the per_Token / per_128x128 weight quant still uses it.

… fp8 or fp4 wire (#4984) * feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire Dispatch sends bf16 and every receiver quantizes each copy it gets, so a token routed to topk peers is quantized topk times on the same values. Quantizing on the sender does it once per LOCAL token and halves (fp8) or quarters (fp4) what crosses the fabric; only the WMMA scale preshuffle has to stay on the receiver, because its destination is the grouped row that rank assigns. The wire must match what the expert GEMM wants for its A operand -- a8w4 -> fp8, a4w4 -> fp4 -- so it is checked, not inferred: a mismatch is a row-width error that would read into the next token's bytes rather than a slow path. MegaMoEStage2Config.dispatch_wire: bf16 | fp8 | fp4 (mori backend only, which is the one with a channel for the scale row). combine_token_nbytes stays bf16 and separate, so the combine slot stride does not follow the wire down. The e8m0 row is padded to 128 B. Dword is all mori's validator asks for, but TdmWholeOrSplit128 only yields a body for the part of a run that starts 128 B aligned, and at the natural 224 B stride only every 4th token does. mori is told hidden_dim in ELEMENTS at its own element size: fp8 and fp4 both transport as byte8, so an fp4 wire halves the count itself. Receiver side, the route gather learns a prequantized producer: it loads the payload dwords and the e8m0 byte the quant pass would have computed, and the store pass -- the only copy of the preshuffled-scale addressing -- is shared verbatim. No software pipelining needed here: this tree's quant/store split already lets the loads cluster. test_mega_moe_gfx1250.py gains --mega_wire {auto,bf16,fp8,fp4}, allows scatter_fused for a4w4_mxfp4, and drives AITER_FORCE_A8W4 off the quant key -- otherwise -q a4w4_mxfp4 silently measured a8w4. Compile-checked only (no run): 8/8 route-gather variants, 6/6 mori dispatch plans (bf16/fp8/fp4 x 8/16 warps) emit .hsaco, and the four host paths produce the payload/scale widths the GEMM expects. * fix(gfx1250): let a4w4 reach the grouped GEMM, and MegaMoE reach a4w4 Three things kept an fp4 wire from ever being measurable, none of them in the wire itself. grouped_moe_gfx1250: mxfp4 weights arrive either as fp4x2 or as the uint8 view of the same bytes (ATOM's loader keeps them uint8; MegaMoE accepts both). Only the a8w4 arm said so, so a4w4-with-uint8-weights failed the eligibility test and fell through to the 2-stage mxfp4 kernels -- a silent detour to a different kernel family, which on this shape has no tuned config and aborts. The next statement already normalized both spellings for the CSV key, so the asymmetry was an oversight, not a rule. test harness: MegaMoE rejects anything but g1u1 interleave, while resolve_spec gave a4w4 the SEPARATED default, so -q a4w4_mxfp4 --combine scatter_fused could not start. Forcing INTERLEAVE alone then produced uncorrelated output, because shuffle_group still laid the weights out for the 2-stage family (e8m0_shuffle) rather than the grouped one (n32k4 + gguu->gugu rows) -- the MX keys differ only in ACTIVATION dtype, so under MegaMoE both take the grouped prep. run_matrix.sh: never check accuracy in a timed run. The fp32 reference is a per-expert torch loop; at 16384 tokens/rank it pins all four ranks for tens of minutes and reads as a hang. Correctness now runs once per wire at 256 tokens. * perf(mega_moe): hand mori the quantizer's scale rows as they come out The wire's e8m0 rows were being copied into a 128 B-strided buffer before dispatch, because mori laid a row down at exactly the width it was given and the alignment is what makes the transfer fast. mori now derives that stride itself, so the copy -- one extra kernel per dispatch, over every token's scales -- goes away and the quant op's output goes straight onto the wire. scale_nbytes is now what we SEND (hidden/32, packed); scale_dst_nbytes is what ARRIVES (mori's stride), and it is asked of mori rather than recomputed, so a change to the alignment cannot leave the two out of step. * support a4w4 test * revert(test): drop the duplicated a4w4 test rewrite yanbo's 36dac59 was cherry-picked here and then landed upstream separately, in a revised form, as #5052. Keeping both makes the two versions collide; the mega_moe.py half of that commit stays, since it is not on main. * fix(ci): the gfx1250 MegaMoE job passes a --combine value that no longer exists #5052 narrowed --combine to base|fused; the CI line still says scatter_fused, so the job has been exiting on argparse rather than running the test. * refactor(mega_moe): leave the dispatch backend default alone Flipping it to mori was an out-of-scope behaviour change: a quantizing wire already requires the mori backend, and __post_init__ says so with a message naming the kwarg, so the default never had to move. The dataclass field default is unreachable anyway -- MegaMoEGfx1250.__init__ always supplies a value -- but the $MEGA_DISPATCH fallback beside it is live, and it decided the backend for every caller that sets neither. * refactor(test): one dispatch-wire flag, and drop the one that never worked --dispatch_commu_dtype has had no effect since the file was created in #4785: resolve_spec wrote transport/prequant/is_fp8/fp8_dtype into the spec and nothing ever read them, quant_tokens_fp8() had no callers, and the base combine hardcodes transport_dtype = bf16. So this is not two overlapping knobs -- it is one live knob and one corpse. --mega_wire takes over the name (--dispatch_wire, matching the kwarg it feeds) and the dead half goes with the flag it belonged to. _FP8_DTYPE stays: the per_Token / per_128x128 weight quant still uses it. * refactor(mega_moe): say dispatch out loud on everything that means the dispatch wire Half the names this PR added already carried the direction (dispatch_token_nbytes, combine_token_nbytes); the other half did not, and a combine wire is coming. is_quant_wire is the one that mattered: it is a pure dispatch predicate with nine call sites, so 'if config.is_quant_wire' written inside _combine() would take the wrong branch AND RUN -- a bf16 reduce over fp8 bytes, wrong numbers, no error. _DispatchWire is deliberately not generalised: three of its four fields are structurally dispatch-only, so sharing the table would hand combine three dead ones. * refactor(mega_moe): MEGA_WIRE becomes MEGA_DISPATCH_WIRE, and the old name is fatal Combine will want a wire of its own, so the unqualified name had to go while it is still unreleased. The stale name raises rather than falling back: an env var that is silently ignored sends a run that asked for fp4 down the bf16 path and reports nothing, which is exactly the failure a wire benchmark cannot survive. Setting both names to the same value is allowed, so a fleet can be rolled over one script at a time. * fix(gfx1250): stop reusing a1_scale for two different tensors The parameter is the caller's per-token e8m0 rows; forty lines later the quant pass rebinds the same name to the preshuffled grouped scale. Both are uint8 with a plausible shape, so a wrong read is silent. The incoming one now has its own name and the rebinding introduces the grouped meaning exactly once. * refactor(mega_moe): one copy of the stale-env check, and satisfy black Three renamed lines went past 88 columns, which the pre-checks job enforces. The env guard was also true of one caller only: the library evaluates it just when dispatch_wire is None, and the test always passes the kwarg, so the harness -- the launch path most likely to carry a stale MEGA_WIRE -- would have measured bf16 in silence. The helper is now public and the test's argparse default calls it, imported lazily so the mega package is not pulled in before FLYDSL_GPU_ARCH is set. _WIRE_FOR_QUANT picks up the dispatch prefix too; it was the same ambiguity the commit before this one exists to remove. * refactor(mega_moe): the op-level config is not a stage2 config It carries geometry and the dispatch knobs and not one stage2 parameter. Stage2 in this package is the gemm2 epilogue fused into combine, which Stage2ScatterContext already names correctly and which takes nothing from here -- so the config was the only place the word was wrong. Name only: the fields, the flat kwargs and the public surface are unchanged, so no caller moves. --------- Co-authored-by: yanboshao <yashao@amd.com>

github-actions Bot reviewed Aug 16, 2026

View reviewed changes

yanboshao changed the title ~~[feat] mega stage2 of gfx1250~~ [feat] mega-moe stage2 for gfx1250 Aug 17, 2026

yanboshao force-pushed the yanbo/mega_stage2_gfx1250 branch from cabe505 to fd904bb Compare August 17, 2026 08:42

yanboshao marked this pull request as ready for review August 17, 2026 10:06

yanboshao requested a review from a team August 17, 2026 10:06

yanboshao force-pushed the yanbo/mega_stage2_gfx1250 branch from fd904bb to eea42c1 Compare August 17, 2026 12:58

zufayu requested a review from yadaish August 18, 2026 02:40

yanboshao force-pushed the yanbo/mega_stage2_gfx1250 branch from eea42c1 to 641b50a Compare August 18, 2026 05:46

github-actions Bot changed the title ~~[feat] mega-moe stage2 for gfx1250~~ [FlyDSL] [CI] [feat] mega-moe stage2 for gfx1250 Aug 18, 2026

coderfeli reviewed Aug 18, 2026

View reviewed changes

Comment thread aiter/ops/flydsl/grouped_moe_gfx1250.py Resolved

yanboshao force-pushed the yanbo/mega_stage2_gfx1250 branch 2 times, most recently from baecf5d to 50ad7a0 Compare August 19, 2026 03:51

yanboshao and others added 17 commits August 19, 2026 08:09

fix rebase error

a6e7286

combine perf optimization

e0472d9

Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com>

opt redandunt ep tokens

daee06a

updat

12f9b80

rm fill 0

7e81af1

fix stage1

ff911c2

update dispatch&combine v2

9d50571

Revert "fix stage1"

b9de910

This reverts commit d3b3fb4.

Revert "feat(ep): push-group fixed-slot dispatch->GEMM1->GEMM2 scatte…

abb335c

…r (gfx1250 a8w4)" This reverts commit 90b0d59.

refactor magemoe

13980c9

update

aef4957

fix refactor

00ea4d2

split combine

3ab6387

fix e2e: mori_dispatch + megamoe

9fcd14a

yanboshao added 5 commits August 19, 2026 08:11

codecheck

8698ce8

fix code check

879109e

fix test script

3e4628f

clean code

ab0bdfb

rm src_token_map d2d

0bab775

yanboshao force-pushed the yanbo/mega_stage2_gfx1250 branch from 50ad7a0 to 0bab775 Compare August 19, 2026 08:27

yanboshao added 2 commits August 21, 2026 05:29

Merge origin/main into yanbo/mega_stage2_gfx1250

8a4b346

refactor

69f6f9a

yadaish approved these changes Aug 21, 2026

View reviewed changes

junhaha666 approved these changes Aug 21, 2026

View reviewed changes

XingerZhu merged commit ba59a37 into main Aug 21, 2026
57 of 58 checks passed

XingerZhu deleted the yanbo/mega_stage2_gfx1250 branch August 21, 2026 11:02

Navigation Menu

Repository navigation

[FlyDSL] [CI] [feat] mega-moe stage2 for gfx1250 - #4785

[FlyDSL] [CI] [feat] mega-moe stage2 for gfx1250#4785
XingerZhu merged 24 commits into
mainfrom
yanbo/mega_stage2_gfx1250

yanboshao commented Aug 16, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 16, 2026

Uh oh!

github-actions Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

	import argparse
	import os

	import torch
	import torch.distributed as dist
	import torch.profiler as tprof

	import aiter
	from aiter import dtypes
	from aiter import ActivationType, QuantType, get_gfx
	from aiter.fused_moe import fused_moe
	from aiter.ops.shuffle import shuffle_weight, moe_shuffle_scale
	from aiter.ops.flydsl.moe_common import GateMode
	from aiter.utility import fp4_utils
	from aiter import get_hip_quant, get_torch_quant, pertoken_quant

Conversation

yanboshao commented Aug 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Aug 16, 2026

🏷️ CI Guide

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

yanboshao commented Aug 16, 2026 •

edited

Loading