[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.

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

mega_moe: add a mori HIP dispatch backend

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:

yanboshao and others added 18 commits August 13, 2026 17:12

fix rebase error

89e8391

combine perf optimization

98d36a2

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

opt redandunt ep tokens

e8ee2fc

updat

7720a67

rm fill 0

67b48dd

fix stage1

83ce25f

update dispatch&combine v2

ee7925d

Revert "fix stage1"

e047c10

This reverts commit d3b3fb4.

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

8a1e474

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

refactor magemoe

5430a2b

update

1bafd52

fix refactor

ce98d2c

Merge pull request #4757 from jhchouuu/jhchouuu/mega-moe-mori-dispatch

fa5abb0

mega_moe: add a mori HIP dispatch backend

split combine

19e7e78

fix e2e: mori_dispatch + megamoe

cb929e1

github-actions Bot reviewed Aug 16, 2026

View reviewed changes

codecheck

b8f8e69

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

fix code check

eea42c1

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

Navigation Menu

Uh oh!

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

[feat] mega-moe stage2 for gfx1250#4785
yanboshao wants to merge 20 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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 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

Uh oh!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yanboshao commented Aug 16, 2026 •

edited

Loading