[FlyDSL] [CI] MegaMoEV2: fuse Dispatch, GEMM1, GEMM2 and Combine into one kernel by IzacharyI · Pull Request #6881 · ROCm/aiter

IzacharyI · 2026-10-23T16:09:04Z

Title: [FlyDSL] [Perf] MegaMoEV2: fuse dispatch, GEMM1, GEMM2 and combine into one kernel

Add two opt-in paths for the EP8 DeepSeek-V4 a8w4 layout on gfx950:

AITER_MEGA_SMALL_M: fuse quant, dispatch, GEMM1, GEMM2 and combine into one MegaMoE kernel; called through forward_small_m().
AITER_MEGA_LARGE_M: fuse dispatch, GEMM1, GEMM2 and combine into one persistent MegaMoE kernel, with quant as a separate kernel.

Motivation

DeepSeek-V4-Pro served with TP8 DP8(attn) EP8 on MI355X runs every MoE call through one MegaMoEV2 instance, from a few tokens per rank up to 8192. These two paths speed up the small- and large-input ends of that range; everything else keeps the existing path.

Technical Details

Small input (up to 512 tokens per rank): one launch does every step, which removes the separate launches and grid-wide waits that dominate MegaMoEV2 time at this size.
Large input (instance sized for 8192 tokens per rank): one persistent kernel overlaps the stages; output is bit-identical to MegaMoEV2.
Skewed routing (a few experts or one rank get most of the routes) is handled inside the kernels, with no host sync; uniform routing keeps the same code path and speed:
- large input: GEMM1 takes m-tiles in payload completion order when one expert holds many tiles, and GEMM2 claims n-major chunks (longer when an expert is large) when this rank's rows warrant it;
- small input: a rank that receives many rows, or has a large expert, reads the expert weights through the cache and sends its GEMM2 rows to the combine as FP8 (the main path's 1x32 format) instead of BF16; the other ranks keep streaming weights and BF16 rows, and the combine reads both formats.
Both paths skip invalid (-1) top-k slots, accept ranks with no tokens, and run every kernel on the stream passed in. The large-input choice depends on no per-rank input, so ranks never split between paths.
Both are off by default and turn on only for gfx950, EP8, 7168 × 3072, 384 experts, top-6, a8w4, SiLU; any other setup keeps the existing path. Two new env vars in mega_moe/envs.py; existing MegaMoEV2 kernels do not change.
SGLang uses these paths through sgl-project/sglang# (merge this PR first): it calls forward_small_m() when the DP-global max token count fits, and turns the large-input path on with set_large_m_enabled() for extend steps with at least 1024 tokens on some rank.
Three new files in aiter/ops/flydsl/kernels/mega_moe/: mega_moe_small_m.py and mega_moe_large_m.py (one per path: dispatch, kernel, runtime) and mega_moe_fused_gemm.py (GEMM tiles and configs both use). MegaMoEV2 imports them only when a path is enabled; its existing path and arguments do not change.
The new test is in the multi-GPU CI skip list next to test_mega_moe_v2.py ([CI] Skip MegaMoE v2 multi-GPU tests #5003) and runs by hand on 8 GPUs.
Open PR [FlyDSL] Add full-expert DeepSeek V4.1 Flash MegaMoE tuning for gfx950 #6154 changes the same MegaMoEV2 files but targets EP2/EP4, so it does not overlap.

Test Plan

8 × MI355X, ROCm 7.2.0, image lmsysorg/sglang-rocm:v0.5.21-rocm720-mi35x-20261007, aiter main 7061cbd, FlyDSL 0.4.0, MORI_SHMEM_HEAP_SIZE=16G.

# small input: accuracy (uneven / idle / empty / padded ranks, -1 slots, skewed routes), interleaved main calls, CUDA graph, timing vs MegaMoEV2
AITER_MEGA_SMALL_M=1 torchrun --nproc-per-node 8 \
  op_tests/multigpu_tests/test_mega_moe_small_m.py --small-m-mtprs 256,512
AITER_MEGA_SMALL_M=1 AITER_MEGA_LARGE_M=1 torchrun --nproc-per-node 8 \
  op_tests/multigpu_tests/test_mega_moe_small_m.py --alt-steps 400
# large input: accuracy vs reference and vs main (also checks explicit stream, rank-local mask, SiLU gate)
AITER_MEGA_LARGE_M=1 torchrun --nproc-per-node 8 op_tests/multigpu_tests/test_mega_moe_v2.py \
  --max-tok-per-rank 8192 --expect-large-m --compare-main --rtol 0.08 --bs-list 1024,2048,4096,8192 --accuracy-max-bs 8192
# large input under skewed routing (append --zipf 1.3, or --concentrate-experts 1 with --bs-list 1024,4096)
AITER_MEGA_LARGE_M=1 torchrun --nproc-per-node 8 op_tests/multigpu_tests/test_mega_moe_v2.py \
  --max-tok-per-rank 8192 --expect-large-m --compare-main --rtol 0.08 --bs-list 1024,8192 --accuracy-max-bs 8192 --zipf 1.3
# timing table below: main path vs this PR on one MTPR-8192 instance (run once without and once with AITER_MEGA_LARGE_M=1)
AITER_MEGA_SMALL_M=1 torchrun --nproc-per-node 8 op_tests/multigpu_tests/test_mega_moe_small_m.py \
  --time-only --time-bs 7,14,28,56,84,126,168,224,280,336,448,1024,2048,4096,8192
# main path without the env vars
torchrun --nproc-per-node 8 op_tests/multigpu_tests/test_mega_moe_v2.py --bs-list 1,7,64,128

End-to-end setup, the same for every line below unless noted:

Workload: InferenceX AgentX (aiperf --scenario agentx, dataset semianalysis_cc_traces_weka_062126), 3600 s per run, concurrency 128 / 256 / 384.
Model and server: DeepSeek-V4-Pro-0813, SGLang TP8 DP8(attn) EP8 (attention data-parallel on 8 ranks, MoE expert-parallel, 48 experts per rank), MegaMoE a8w4 with MTPR 8192, DSPARK speculative decoding with the InferenceX acceptance simulation (accept length 3.77 on every line), sglang-router cache_aware + dp-aware. Server, router and client flags are those of InferenceX run 37608671586.
Hardware and image: 8 × MI355X, lmsysorg/sglang-rocm:v0.5.21-rocm720-mi35x-20261007.

line	where	MoE
InferenceX official	InferenceX CI (run above)	MegaMoEV2 from the image
MegaMoEV2 (our node)	our node	MegaMoEV2 from the image, unchanged
this PR	same node as the line above	the image's aiter with the small- and large-input paths turned on

The two lines on our node differ only in the MoE kernels. They were measured before this branch was rebased onto aiter main and before the skewed-routing changes; the kernel tables above are from the final code. GSM8K uses the same server with real speculative verification instead of the acceptance simulation.

Test Result

V4-Pro EP8 A8W4; H=7168, I=3072, 384 experts, top-k=6; 8 × gfx950 / 256 CU. Latency in µs, CUDA Graph 10 replays × 3 rounds, median of rank-max latency. Speedup = baseline / this PR.

Uniform routing — MTPR=TPR

tokens/rank	main protocol	this PR path	main (µs)	this PR (µs)	Speedup vs main
1	fixedslot	small-M	101.7	84.5	1.20×
2	fixedslot	small-M	179.4	169.2	1.06×
4	fixedslot	small-M	219.1	189.2	1.16×
8	fixedslot	small-M	255.9	226.2	1.13×
16	fixedslot	small-M	322.0	286.8	1.12×
32	fixedslot	small-M	336.3	295.7	1.14×
64	fixedslot	small-M	346.5	304.6	1.14×
128	fixedslot	small-M	364.3	315.3	1.16×
256	compact	small-M	570.7	380.0	1.50×
512	compact	small-M	694.7	560.5	1.24×
8192	compact	large-M	5186.0	4211.8	1.23×

Uniform routing — MTPR=8192

tokens/rank	main protocol	this PR path	main (µs)	this PR (µs)	Speedup vs main
1	compact	small-M	181.0	83.5	2.17×
2	compact	small-M	265.4	170.6	1.56×
4	compact	small-M	311.1	191.9	1.62×
8	compact	small-M	338.2	225.3	1.50×
16	compact	small-M	412.5	290.5	1.42×
32	compact	small-M	433.7	297.8	1.46×
64	compact	small-M	463.8	303.0	1.53×
128	compact	small-M	474.9	314.7	1.51×
256	compact	small-M	594.1	381.3	1.56×
512	compact	small-M	730.0	564.1	1.29×
1024	compact	large-M	970.2	903.1	1.07×
2048	compact	large-M	1681.1	1374.0	1.22×
4096	compact	large-M	2729.8	2304.3	1.18×
8192	compact	large-M	5186.0	4211.8	1.23×

Skewed routing — speedup vs main, MTPR=8192

hot60: floor(0.6 × TPR) tokens route to experts 0–5; Zipf: exponent 1.3; one_rank: all routes target experts 0–5; boundary: experts 0, 1, 2, 3, 46, 47; rank-balanced-hot: six distinct destination ranks per token, with 60% probability of local expert 0 on each selected rank. Values below 1.00× indicate regressions.

tokens/rank	path	uniform	hot60	Zipf	one_rank	boundary	rank-balanced-hot
1	small-M	2.17×	2.14×	2.47×	2.33×	2.28×	2.21×
2	small-M	1.56×	1.58×	2.38×	2.27×	2.21×	2.11×
4	small-M	1.62×	1.92×	2.74×	2.44×	2.42×	1.82×
8	small-M	1.50×	1.56×	2.25×	1.86×	1.81×	1.78×
16	small-M	1.42×	1.36×	1.95×	1.53×	1.54×	1.47×
32	small-M	1.46×	1.21×	1.86×	1.42×	1.39×	1.41×
64	small-M	1.53×	1.26×	1.70×	1.21×	1.16×	1.37×
128	small-M	1.51×	1.18×	1.40×	1.11×	1.09×	1.34×
256	small-M	1.56×	1.32×	1.48×	1.26×	1.19×	1.49×
512	small-M	1.29×	1.15×	1.17×	1.04×	0.93×	1.38×
1024	large-M	1.07×	1.29×	1.31×	1.55×	1.39×	0.94×
2048	large-M	1.22×	1.42×	1.34×	1.50×	1.39×	1.17×
4096	large-M	1.18×	1.44×	1.14×	1.36×	1.40×	1.10×
8192	large-M	1.23×	1.43×	1.07×	1.39×	1.38×	1.09×

End-to-end A/B on our node, MegaMoEV2 (our node) vs this PR. Total token throughput = all input tokens (including prefix-cache hits) plus output tokens, per second per GPU; run 1 / run 2, gain of the means:

conc	MegaMoEV2	this PR	gain
128	42,810 / 43,060	47,137 / 47,425	+10.1%
256	65,795 / 67,857	73,217 / 73,796	+10.0%
384	76,192 / 73,603	78,880 / 79,776	+5.9%

Against the published InferenceX results:

ep_output_tput_vs_p90_interactivity_no_tpep_v2

Red: InferenceX official. Green: this PR. Violet (dashed): MegaMoEV2 (our node), which matches red, so our node reproduces InferenceX. Lines as in the setup table; c1–c48 of red are the TP8 configuration InferenceX publishes there.
Below c128 InferenceX uses TP8, which is still faster than EP, so these paths do not apply there.

InferenceX official → this PR (first run of each point, as plotted):

conc	output tok/s/chip	P90 interactivity (tok/s/user)
128	372.8 → 404.3 (+8.5%)	31.6 → 36.8 (+16.4%)
256	514.9 → 556.9 (+8.2%)	27.5 → 29.7 (+7.9%)
384	578.8 → 623.0 (+7.6%)	22.5 → 22.9 (+2.2%)

GSM8K (1319 questions, 5-shot): MegaMoEV2 0.932–0.938, this PR 0.926–0.948.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

…to one kernel Add two opt-in paths for the EP8 DeepSeek-V4 a8w4 layout on gfx950: - AITER_MEGA_SMALL_M: fuse quant, dispatch, GEMM1, GEMM2 and combine into one MegaMoE kernel; called through forward_small_m(). - AITER_MEGA_LARGE_M: fuse dispatch, GEMM1, GEMM2 and combine into one persistent MegaMoE kernel, with quant as a separate kernel. Signed-off-by: zhengyao <zayao@amd.com>

github-actions · 2026-10-23T16:09:32Z

🏷️ 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 (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
`ci:triton-355`	Run the full Triton test suite on MI35X, not only the tests the change affects
`ci:windows`	Windows JIT op smoke test on the gfx1201 ROCm runner
`ci:multi-gpu`	Aiter multi-GPU tests on the 8-GPU runner
`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: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
`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 6881 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

Opt-in switches for the two aiter MegaMoEV2 paths (needs ROCm/aiter#6881): - SGLANG_OPT_USE_AITER_MEGA_SMALL_M: steps whose DP-global max token count fits the small-input runtimes call forward_small_m(), one kernel for quant, dispatch, GEMM1, GEMM2 and combine. - SGLANG_OPT_USE_AITER_MEGA_LARGE_M: extend steps with at least SGLANG_OPT_AITER_MEGA_LARGE_M_MIN_TOKENS (default 1024) tokens on some DP rank run the persistent large-input kernel. Each MoE layer gets its own protocol slot on the shared operator, and the choice of path depends only on DP-gathered counts, so every EP rank takes the same path.

…e fused paths Correctness: - The large-input choice no longer depends on a per-rank mask: its combine skips invalid (-1) top-k slots, so all ranks take the same path. The small-input combine skips them too. - The large-input quant kernel runs on the stream passed to forward(). - A small-input rank with no tokens passes a 1-row stand-in to the kernel. - Both paths are gated to SiLU. - Small input serves up to 512 tokens per rank. Skewed routing, decided inside the kernels; uniform routing keeps its code path and speed: - Large input: GEMM1 takes m-tiles in payload completion order when one expert holds many tiles, and GEMM2 picks its claim order and chunk length from the rows this rank received. - Small input: a rank that receives many rows, or has a large expert, reads the expert weights through the cache and sends its GEMM2 rows as FP8 (1x32 E8M0 scales); the other ranks keep streaming weights and BF16 rows, and the combine reads both formats per peer. Tests cover the explicit stream, rank-local mask, SiLU gate, empty ranks, -1 slots and skewed routes. Signed-off-by: zhengyao <zayao@amd.com>

The fused kernels waited on cross-rank flags with mori device helpers, which makes FlyDSL link an external library and skip its disk cache, so they could not be precompiled. They now wait with the comm_ops helpers the main-path kernels use; behavior and timing are unchanged. aot/flydsl/mega_moe.py builds, for every rank of the DeepSeek-V4-Pro EP8 a8w4 profile (48 experts per rank), each config of the small-input runtimes (MTPR 256 and 512) and the large-input kernel (MTPR 8192) with its quant kernel. The kernel configs and specs come from functions the runtimes share, so AOT and runtime keys match; a run-only check (FLYDSL_RUNTIME_RUN_ONLY=1) passes on the main, small and large paths. Signed-off-by: zhengyao <zayao@amd.com>

Resolve the conflict with the MegaMoE AOT rework (ROCm#6362), which builds the cache by replaying MegaMoEV2 itself and calling preload_aot_bundles(): - aiter/aot/flydsl/mega_moe.py and README.md take main's version. The MTPR 8192 jobs build the instance with AITER_MEGA_SMALL_M=1 and AITER_MEGA_LARGE_M=1; the runtime's eligibility checks decide which profiles get the small- and large-input kernels. - MegaMoEV2.preload_aot_bundles() also loads the small- and large-input runtimes the instance built. Each runtime's preload_aot_bundle() runs the same launch code as forward() with a no-dispatch flag: every token bucket's kernel for the small-input runtimes, the quant and persistent kernels for the large-input runtime. - The large-input runtime allocates its token-ready arena through the op's symmetric allocator like its other arenas, so the replay's fake shmem covers it. Signed-off-by: zhengyao <zayao@amd.com>

… binding - Large-input GEMM1: thread 0 broadcast the claimed unit and then has_work through one LDS word that was also the first word of the A tile buffer. A slow wave could read has_work as the unit, or tile data as has_work and skip a tile the rest of its block ran (mismatched barriers). Each value now has its own LDS word outside the tile buffers; the loop keeps its barriers, so the GEMM1 code is unchanged. - forward_small_m() needs max_tokens: a per-rank default could put ranks on different small-input runtimes, which share no protocol state. - The small-input kernel splits its work over 8 XCDs, so it is eligible only on 256-CU (8-XCD) gfx950 parts. - bind_weights() refuses strided weights instead of binding a temporary copy a captured graph would outlive; the constructor keeps its contiguous copies for the instance's lifetime. - bind_weights() and bind_layer() check everything that can refuse (contiguity, the small-input W2 dispatch tables during capture, the layer slot) before changing any binding. - The large-input quant runs on the default stream for a null stream handle: ExternalStream(0) would take a new pool stream. Signed-off-by: zhengyao <zayao@amd.com>

aiter-bot · 2026-10-30T09:42:23Z

NOT READY FOR HUMAN REVIEW

PR Self-Review

Revision 0a85ae56868b · mode shadow

Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision.

Questions to answer in the PR description

What behavior does this PR change, and why is the existing behavior insufficient?
What did you validate, and what remains unverified? Include commands, hardware/configuration and results.
What permissions and untrusted-input boundary does this CI change have, and can fork PR code or text reach a privileged execution path?
For every changed op_tests/ file, what behavior does it cover and why is it not redundant? Write None if no op test changed.

Edit the PR Self-Review section in the PR description. Include the current revision SHA; this comment refreshes automatically.

Missing or invalid answers

FAIL — missing ### Reviewed revision in the PR description
FAIL — missing ### Behavior and intent in the PR description
FAIL — missing ### Validation and unverified scope in the PR description
FAIL — missing ### PR-specific question in the PR description
FAIL — missing ### Changed op tests in the PR description
FAIL — Reviewed revision must contain at least the first 7 characters of the current head SHA
FAIL — Behavior and intent has no substantive answer
FAIL — Validation and unverified scope has no substantive answer
FAIL — PR-specific question has no substantive answer
FAIL — changed op test not explained in Changed op tests: op_tests/multigpu_tests/test_mega_moe_small_m.py
FAIL — changed op test not explained in Changed op tests: op_tests/multigpu_tests/test_mega_moe_v2.py
FAIL — Changed op tests must explain why each changed op test is necessary and non-redundant

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

LargeMRuntime.bind_layer() switched the layer slot before binding the weights, so a refused binding (strided weights) left the new slot with the previous weights. It now checks the slot first, binds the weights, then switches the slot, like MegaMoEV2.bind_layer(); the edge test covers both refusals on the runtime itself. Signed-off-by: zhengyao <zayao@amd.com>

With one state slot the edge check's "other" slot was the bound one, so it could not catch bind_layer switching the slot before refusing the weights. The test now builds two slots by default and reports the check as RANGE-ONLY when run with one. It fails on the old bind_layer order and passes on the current one. Signed-off-by: zhengyao <zayao@amd.com>

…ight rebinding - SmallMRuntime refuses a device without 256 CUs itself: its kernel splits the work over 8 XCDs, so a 4-XCD part would leave half of it unclaimed even when the runtime is built directly. - SmallMRuntime and LargeMRuntime keep the contiguous copies they make of strided construction weights for their lifetime, as MegaMoEV2 does, so a graph captured with them never outlives them. - The main and large-input Stage1 refuse scales that are not 1-byte E8M0 rows of model_dim / 32 on the instance's device; the kernel reads a whole row per token. - The layer-slot stress gives layer B its own W2 values (twice A) and nonzero routing weights, checks three slots per cycle against the main path on the same binding, and checks that layer B changes the output, so a call on the wrong layer's weights now fails it. Signed-off-by: zhengyao <zayao@amd.com>

The Stage1 checks accepted any 1-byte scales, so float8_e4m3fn, int8 or bool scales were read as E8M0 exponents without an error. The main and large-input Stage1 now accept uint8 or float8_e8m0fnu only; the edge test feeds both a float8_e4m3fn view and expects a refusal. Signed-off-by: zhengyao <zayao@amd.com>

mega_moe@1.0 is the expert-parallel MoE of one call on every rank of an EP group: dispatch of each token to the ranks holding its top-k experts, the two MXFP4 expert GEMMs with the clamped SwiGLU, and the combine back on the token's rank. The reference sums over the group's experts in float64; the baseline is aiter MegaMoEV2's main path, and with one rank aiter's local fused_moe on the same gate-up-interleaved experts (MegaMoEV2 needs mori's group). Three flydsl kernels serve DeepSeek-V4-Pro EP8 a8w4 (7168 / 3072, 48 experts a rank, top-6): aiter's single-kernel small-input path at 1-256 and 320-512 tokens a rank, and its persistent large-input path at 1024-8192. Their host code takes the weights with every call, sets up the mori symmetric buffers once in the compile function, and writes the combine straight into out. The sources follow ROCm/aiter#6881 at a7b7e7441, including its fix of the large-input GEMM1 work broadcast (one LDS word per broadcast value, apart from the tile buffers) and the small-input 256-CU requirement. Signed-off-by: zhengyao <zayao@amd.com>

…ract - The 32-bit launch generations of the fused paths may wrap: every wait that relied on them growing now waits for equality. Small-input plan words and the combine format flag compare their epoch field exactly (spin_until_shifted_eq_i64), GEMM2 readiness is a release store instead of an unsigned max that would stick at 0xffffffff, and the large-input launch handshake and combine generation wait for their exact value (no peer can be a launch of the same slot ahead). - weights.py holds the binding contract: bind_weights() on the operator and both runtimes refuses weights that are strided or differ from the construction shape, dtype (byte views of FP4/E8M0 are equivalent) or device, and computes every view before it changes any binding. - forward() of every path refuses inputs on another device; a rank that faulted on a foreign pointer would leave its peers waiting. - SmallMRuntime built directly checks the full layout MegaMoEV2 gates it to (small_m_eligible). - python -m aiter.aot.flydsl.mega_moe --check replays in run-only mode and fails on any kernel the cache lacks. - Tests: launches across preset epochs just below 2**31 and 2**32 (the previous code hangs there), every rank with zero rows, the null stream, and a refused binding leaving every bound view, table and slot of the operator and its runtimes unchanged. Signed-off-by: zhengyao <zayao@amd.com>

Signed-off-by: zhengyao <zayao@amd.com>

#14) * feat(flydsl): add the mega_moe op with DeepSeek-V4 EP8 MegaMoE kernels mega_moe@1.0 is the expert-parallel MoE of one call on every rank of an EP group: dispatch of each token to the ranks holding its top-k experts, the two MXFP4 expert GEMMs with the clamped SwiGLU, and the combine back on the token's rank. The reference sums over the group's experts in float64; the baseline is aiter MegaMoEV2's main path, and with one rank aiter's local fused_moe on the same gate-up-interleaved experts (MegaMoEV2 needs mori's group). Three flydsl kernels serve DeepSeek-V4-Pro EP8 a8w4 (7168 / 3072, 48 experts a rank, top-6): aiter's single-kernel small-input path at 1-256 and 320-512 tokens a rank, and its persistent large-input path at 1024-8192. Their host code takes the weights with every call, sets up the mori symmetric buffers once in the compile function, and writes the combine straight into out. The sources follow ROCm/aiter#6881 at a7b7e7441, including its fix of the large-input GEMM1 work broadcast (one LDS word per broadcast value, apart from the tile buffers) and the small-input 256-CU requirement. Signed-off-by: zhengyao <zayao@amd.com> * fix(mega_moe): run it as a ranks group and build it with FlyDSL 0.3.2 - Call the op as a ranks group (ranks="ep_size"): group.py sets up mori's symmetric heap over the group, 16 GiB so MegaMoEV2's and the kernel's buffers both fit; the reference reads each expert from its owner rank's args and decodes to float32. The baseline takes out=None. - communication_ops_utils.py emits its ordered loads, stores, fences and atomics as LLVM-dialect ops, since FlyDSL 0.3.2 (the baseline AITER pin) lacks fx.generic_load/generic_store/atomic_*/memory_fence; schemas renamed. - CI: FIXME workaround exposing TheRock's libdrm headers to mori's JIT. --------- Signed-off-by: zhengyao <zayao@amd.com> Co-authored-by: zhenhuang12 <Zhen.Huang@amd.com>

Main's gfx1250-strict wheel builds (#7379, #7434) and MegaMoEV2 (#6881) edited the code this branch rewrites; both sides are kept: - moe.py: an arch's replay keeps the jobs of that arch's family, retargets them to it and leaves out the MXFP4 ones on gfx1250-strict; deployment jobs now pass through that filter as rows do. run_aot and moe.py's main run inside fail_build_on_missing_target_feature(), and run_aot still ends with the strict errata check. - mega_moe.py: the MTPR 8192 instances build the opt-in small- and large-input paths, and --check compiles nothing. FUSED_PATHS_MTPR and FUSED_PATHS_ENV stay module constants: they say how the replay builds an instance, not what a deployment is. - README: main's gdn_fused_prefill entries next to the deployment ones. #6464 moved DeepSeek-V4-Pro's TP8 rows at M 32 and 2048 to MXMOE stage 1, which compiles the clamp in, so the clamp the TP8 deployment passes now picks kernels the deployment builds. The merge changes no job list: the 1195 deployment jobs are 88ab025's, the 14374 tuned-row jobs main's (apart from the empty expert_mask_size this branch adds) and the 192 MegaMoE jobs main's. Co-authored-by: Cursor <cursoragent@cursor.com>

IzacharyI requested a review from a team October 23, 2026 16:09

github-actions Bot changed the title ~~[FlyDSL] [Perf] MegaMoEV2: fuse dispatch, GEMM1, GEMM2 and combine into one kernel~~ [FlyDSL] [CI] [Perf] MegaMoEV2: fuse dispatch, GEMM1, GEMM2 and combine into one kernel Oct 23, 2026

github-actions Bot added CI FlyDSL labels Oct 23, 2026

IzacharyI changed the title ~~[FlyDSL] [CI] [Perf] MegaMoEV2: fuse dispatch, GEMM1, GEMM2 and combine into one kernel~~ [FlyDSL] MegaMoEV2: fuse dispatch, GEMM1, GEMM2 and combine into one kernel Oct 23, 2026

IzacharyI changed the title ~~[FlyDSL] MegaMoEV2: fuse dispatch, GEMM1, GEMM2 and combine into one kernel~~ [FlyDSL] MegaMoEV2: fuse Dispatch, GEMM1, GEMM2 and Combine into one kernel Oct 23, 2026

IzacharyI assigned IzacharyI and Bernard-Liu and unassigned IzacharyI Oct 23, 2026

coderfeli requested review from GwilliamHu, Yaowu-Xiong and binding7012 October 25, 2026 13:13

zufayu requested review from Bernard-Liu, TennyWang1223, coderfeli and yadaish October 26, 2026 01:08

chushic mentioned this pull request Oct 27, 2026

[FlyDSL] [Perf] MegaMoE: faster compact Stage1 producer and FP8 GEMM1 tiles (gfx950) #7001

Open

5 of 9 tasks

IzacharyI mentioned this pull request Oct 28, 2026

[Perf] [AMD] dsv4: Use the aiter MegaMoE small- and large-input paths sgl-project/sglang#47065

Open

4 of 5 tasks

github-actions Bot changed the title ~~[FlyDSL] MegaMoEV2: fuse Dispatch, GEMM1, GEMM2 and Combine into one kernel~~ [FlyDSL] [CI] MegaMoEV2: fuse Dispatch, GEMM1, GEMM2 and Combine into one kernel Oct 28, 2026

IzacharyI added 3 commits October 28, 2026 13:43

This was referenced Oct 30, 2026

[MegaMoE] Add MiMo V2.6 small-input profile pbkowalski/aiter#2

Draft

[AMD] Route MiMo V2.6 target verify through AITER MegaMoE sgl-project/sglang#47486

Draft

IzacharyI added 3 commits November 2, 2026 05:53

GwilliamHu mentioned this pull request Nov 2, 2026

[FlyDSL] perf(mega_moe): gfx1250 decode megakernel - dispatch, GEMM1, GEMM2 and combine in one launch #7345

Draft

3 of 4 tasks

IzacharyI mentioned this pull request Nov 3, 2026

feat(flydsl): add the mega_moe op with DeepSeek-V4 EP8 MegaMoE kernels AMD-AGI/AKLR#14

Merged

GwilliamHu previously approved these changes Nov 3, 2026

View reviewed changes

Merge branch 'main' into meagemoe_to_aiter

71887d7

coderfeli dismissed GwilliamHu’s stale review via 71887d7 November 3, 2026 14:16

IzacharyI added 2 commits November 4, 2026 03:07

Merge branch 'main' into meagemoe_to_aiter

0a85ae5

Signed-off-by: zhengyao <zayao@amd.com>

yangyuhuiling mentioned this pull request Nov 5, 2026

feat(cco): example 14 -- the whole EP MoE layer in one persistent kernel, Stage2 on SDMA ROCm/mori#919

Draft

6 tasks done

coderfeli approved these changes Nov 5, 2026

View reviewed changes

coderfeli merged commit a060781 into ROCm:main Nov 5, 2026
59 checks passed

zhiding512 mentioned this pull request Nov 9, 2026

[FlyDSL] [Feature] Precompile MoE deployments (EP and TP), described by model and deployment knobs #7063

Open

Navigation Menu

Repository navigation

[FlyDSL] [CI] MegaMoEV2: fuse Dispatch, GEMM1, GEMM2 and Combine into one kernel - #6881

[FlyDSL] [CI] MegaMoEV2: fuse Dispatch, GEMM1, GEMM2 and Combine into one kernel#6881
coderfeli merged 12 commits into
ROCm:mainfrom
IzacharyI:meagemoe_to_aiter

IzacharyI commented Oct 23, 2026 •

edited by GwilliamHu

Loading

Uh oh!

github-actions Bot commented Oct 23, 2026

Uh oh!

aiter-bot commented Oct 30, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

IzacharyI commented Oct 23, 2026 • edited by GwilliamHu Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Uniform routing — MTPR=TPR

Uniform routing — MTPR=8192

Skewed routing — speedup vs main, MTPR=8192

Submission Checklist

Uh oh!

github-actions Bot commented Oct 23, 2026

🏷️ CI Guide

Uh oh!

aiter-bot commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

NOT READY FOR HUMAN REVIEW

PR Self-Review

Questions to answer in the PR description

Missing or invalid answers

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

IzacharyI commented Oct 23, 2026 •

edited by GwilliamHu

Loading

aiter-bot commented Oct 30, 2026 •

edited

Loading