[yangyuhuiling]

Summary

Adds examples/cco/python/13_flydsl_mega_moe_stage1_sdma: aiter's MegaMoEV2 Stage1 (dispatch + GEMM1) with the dispatch moved onto the SDMA copy engines. It is the Stage1 counterpart of example 11, which does the same for Stage2 + combine.

This is a draft. With its own 8-wave GEMM1 the SDMA Stage1 is slower than aiter's Stage1 with the open ROCm/aiter#7001. With --gemm1 7001 (#7001's own GEMM1 tiles, the rows packed in the same launch; commit 39f7618) it is even with #7001 at 4096 and 8192 tokens/rank and about 0.05 ms slower at 2048. It is posted as a reference for how cco SDMA can feed a persistent GEMM, and records what bounds SDMA here.

How it works

• pack.py: one block per dispatch segment.
  • Copies this rank's routed rows (fp8 + e8m0) into a send slab in the cco window, grouped by (destination, segment). These are exactly the rows of aiter's production producer.
  • This rank's own rows go straight to the landing buffer.
  • The route headers go to the destination over P2P, as aiter does.
• quant_pack.py (--fuse-quant, MegaMoEStage1Sdma(fuse_quant=True)): replaces pack by folding its copy into the quantisation. aiter's prepare then only plans; its quant workgroups are not launched.
  • route: walks the same rows as pack, one lane per row, without moving payload. It records where each token's rows go and writes the route headers as pack does.
  • qpack: one block per token. It quantises the bf16 row once with aiter's per-1x32 MX fp8 arithmetic (same bytes), stages it in LDS, and stores it to each of the token's slab / landing rows with 16 contiguous bytes per lane.
• gemm1_7001_sdma.py (--gemm1 7001, MegaMoEStage1Sdma(gemm1="7001")): one launch built on #7001's compact Stage1 kernel, its CU dispatch replaced. Workgroups 1..32 pack this rank's rows segment by segment in expert order and flag each one; workgroup 0 puts a peer's segment with SDMA as soon as it is flagged; then every workgroup takes #7001's GEMM1 work items (its tuned compact FP8 tiles), each waiting for the peers' signals and this rank's own segments of its expert. No pack / qpack launch runs before it. See the section below.
• gemm1_8wave_sdma.py: one persistent fused kernel.
  • Workgroup 0: one lane per peer puts the slab segments to the peer's landing buffer at aiter's payload row numbers. Each (peer, segment) gets one data put plus one signalled scale put, in expert order; empty segments get a 4-byte put, so the signal count equals the number of finished segments.
  • Other workgroups: an 8-wave 256x256 GEMM1 + SwiGLU with dynamic work items. An expert's M-tiles start as soon as its rows have landed.
  • Two ways to get the A scales:
    • shuffler (4096/8192 buckets): a few workgroups put the landed scales into the contiguous shuffle_scale layout once per M-tile. A work item then waits on their flags only.
    • ldsw (2048 bucket): every work item waits on the 7 peer signals and stages its scales into LDS. Its loads are issued ahead of the prologue's A/B loads.
• mega_moe_stage1_sdma.py: MegaMoEStage1Sdma.
  • Owns the cco window, DevComm and buffers.
  • install() swaps MegaMoEV2._s1_mega / _s1_mega_bundle. aiter's prepare, Stage2 and combine are used unchanged.
  • scale_mode="auto" picks the mode from the bucket's sort_block_m.
• main.py:
  • test: 4 back-to-back calls with alternating inputs; requires the layer output to match aiter's byte for byte.
  • bench: CUDA Graph timing of Stage1 (incl. prepare) and the whole layer, against aiter.

Results

Setup: 8x MI355X, DeepSeek-V4-Pro a8w4 (v4_pro), EP8, MTPR 8192, CUDA Graph, max over ranks. Stage1 includes prepare. All outputs are byte-identical to aiter's.

Against aiter main f0b96f49e (main.py bench):

tokens/rank	scale mode	Stage1 aiter	Stage1 SDMA	layer aiter	layer SDMA
2048	ldsw	0.933	0.958 (+0.025)	1.572	1.622 (+0.050)
4096	shuffler	1.439	1.369 (-0.070)	2.466	2.443 (-0.023)
8192	shuffler	2.431	2.212 (-0.219)	4.548	4.422 (-0.126)

Against aiter main + ROCm/aiter#7001: same node and container, interleaved #7001 / main / #7001, each aiter tree with its own FlyDSL cache.

tokens/rank	aiter main	aiter + #7001	SDMA Stage1	SDMA - #7001
2048	0.935	0.820	0.950-0.961	+0.13
4096	1.425	1.177	1.367-1.372	+0.19
8192	2.419	2.040	2.210-2.219	+0.18

With --fuse-quant (against aiter main + #7001, main.py bench; two interleaved rounds against pack, plus one run of the committed code):

tokens/rank	aiter + #7001	SDMA, pack	SDMA, fuse-quant	fuse-quant - #7001	layer #7001 / fuse-quant
2048	0.811-0.821	0.957-0.958	0.930-0.937	+0.12	1.44 / 1.58
4096	1.174-1.178	1.364-1.365	1.318-1.331	+0.15	2.22-2.23 / 2.39-2.42
8192	2.003-2.050	2.210-2.225	2.128-2.139	+0.09-0.11	4.21-4.23 / 4.34-4.35

At 8192, eager: route + qpack take 0.033 + 0.083 ms where pack took 0.188 ms. The first qpack stored each lane's row piece straight from registers (two 16-byte pieces 32 bytes apart); that ran at 0.26 ms, because the cco window is uncached and strided partial lines write slowly. Staging the row in LDS fixed it.

Why it loses: once the CU dispatch is as cheap as #7001 makes it, the extra pack copy costs more than the CU time the copy engines save. At 8192 that copy is about 0.19 ms, about 360 MB read and written per rank at ~3.8 TB/s. A variant that ran the pack inside the fused kernel to overlap it with GEMM1 was no faster: both compete for HBM bandwidth. --fuse-quant takes most of that copy away: the send-side preparation is now about 0.12 ms. The remaining gap is in the fused kernel itself, about 1.85 ms at 8192 against about 2.02 ms for #7001's whole Stage1 including prepare: reading A from the uncached window and the per-work-item waits for landed data. Dropping the pack entirely and putting the token-ordered input as it is would not help either: an M-tile's rows come from almost the whole token range of their source rank, so GEMM1 could not start until nearly all of x had landed.

What was measured along the way (fused kernel alone, 8192, no puts)

variant	ms
ldsw staging per work item + 7-signal wait per work item	2.02
ldsw loads issued ahead of the prologue	1.87
shuffler, 4 workgroups	1.83
lower bound: no wait, pre-shuffled scales	1.79
same rows gathered contiguous, contiguous GEMM	1.78

Reading A from the cco window (uncached VMM) vs a torch copy differed by only 0.02-0.06 ms. An earlier 0.22 ms gap turned out to be a measurement-order artefact.

Per-item fixes from example 14 (commit 1d2b90d)

#919 (example 14, the whole layer in one persistent kernel) cut the per-work-item overhead of the same 8-wave GEMM1 and batched the producers' row copy of its SDMA dispatch. This commit ports what applies here:

• tile_find: a ballot + popcount and v_readlane instead of a butterfly max over ds_bpermute (six dependent LDS permutes per work item).
• Readiness in one round trip: a GEMM1 work item reads all shuffler flags (shuffler) or all 7 peer signal counts (ldsw) at once and spins only on a short one; the shufflers do the same per expert. The signal counts are one array indexed by the sender, so its address and stride are looked up once per workgroup. Looking up all seven per item added VGPR spills (1 -> 7 in the ldsw build).
• pack.py: a wave writes the route headers of 64 rows lane-parallel, then copies 4 rows per step with every load of the step before its stores. vmcnt retires loads and stores in order, so one store-ack wait covers 4 rows instead of one.
• bench: per replay the max over ranks, then the median over 30 replays (was the mean of 20). A host hiccup on one rank stretches the replays it delays; the median ignores those.

Not ported:

• The SwiGLU v_rcp_f32 of feat(cco): example 14 -- the whole EP MoE layer in one persistent kernel, Stage2 on SDMA #919: the output must stay byte-identical to aiter's, and aiter's sigmoid uses an IEEE division.
• One acquire per loop instead of per item: this kernel has no per-item acquire. Its per-item L1 invalidate (shuffler) is needed for correctness.

Same node (8x MI355X, rocm/sgl-dev ROCm 10 image), aiter f0b96f49e + the files of ROCm/aiter#7196 and ROCm/aiter#7001, interleaved runs (two each), Stage1 ms including prepare, output byte-identical to aiter's:

tokens/rank	aiter + #7001	pack, before	pack, after	fuse-quant, before	fuse-quant, after
2048	0.800-0.807	0.945-0.946	0.931-0.937	0.915-0.921	0.901-0.908
4096	1.163-1.179	1.383-1.386	1.365-1.377	1.324-1.354	1.323-1.326
8192	2.017-2.050	2.228-2.236	2.210-2.216	2.149-2.161	2.150-2.152

About 0.01-0.02 ms (1%), much less than in #919. That kernel spilled the tile_find lane addresses and invalidated L2 at every work item; this one does neither (0 VGPR spills with shuffler, 1 with ldsw). The gap to #7001 is unchanged, and it is the GEMM1: measured on the same inputs, see the next section.

Benchmark note. On a host whose page cache leaves a NUMA node little free memory, kernel compaction keeps migrating pages. ROCr's system memory (hipHostMalloc, kernel arguments, signals) is userptr-backed by default, so each migration makes KFD evict every queue of the owning process, and with cross-rank waits one rank's eviction stalls all of them (20-580 ms per call seen with example 14). HSA_USERPTR_FOR_PAGED_MEM=0 (GTT buffers instead of userptr) removes it; the numbers in this section were measured with it.

aiter#7001's GEMM1 tiles (--gemm1 7001, commit 39f7618)

Where the gap was. Each fused kernel run alone, without its dispatch (#7001's producers skipping the row copies; this kernel's items not waiting), in one run on the same node and inputs, ms:

tokens/rank	#7001 GEMM1 alone	8-wave GEMM1 alone	#7001 fused (with its CU dispatch)
2048	0.60	0.68-0.69	0.67
4096	0.81	1.04-1.06	0.99-1.06
8192	1.40	1.81-1.82	1.93

The 8-wave GEMM1 is 0.08 / 0.25 / 0.42 ms slower than #7001's tiles. #7001 spends about 0.07 / 0.2 / 0.5 ms of its fused kernel on the CU dispatch, which the SDMA dispatch does not take from the GEMM workgroups, so at 8192 the two fused kernels took about the same time and the Stage1 gap was the pack / qpack launch before this one. (The commit message quotes 1.89-1.95 ms and 0.37 ms from an earlier run whose #7001-alone time was 1.56 ms; the values above come from one run, and a second run gave #7001 alone 1.40 ms again.)

What --gemm1 7001 does. gemm1_7001_sdma is #7001's compact Stage1 kernel with the producers' P2P copies replaced:

• workgroups 1..32 pack this rank's rows, one (destination, segment) at a time in expert order (group segment first), into the send slab (this rank's own rows into its landing rows), write the route headers to the destination over P2P, and flag the segment;
• workgroup 0, one wave per peer, puts each of the peer's segments as soon as it is flagged (the fp8 rows, then the e8m0 with the remote signal) and joins the GEMM queue;
• every workgroup takes #7001's GEMM1 work items (build_fused_gemm1, its tuned compact route-major FP8 tiles); an item waits until every peer has signalled the group segment and experts 0..e of its M-tile and this rank's own segments of e are flagged, where #7001 waits for its producers' tile-ready counters.

The Stage1 output matches #7001's byte for byte. Packing inside the launch removes the pack / qpack launch (0.12-0.18 ms at 8192) and starts the copies about that much earlier; the packing workgroups cost the GEMM about 0.08 ms back.

Same node (crsuse2 v3-091, rocm/sgl-dev ROCm 10 image), aiter f0b96f49e + the files of ROCm/aiter#7196 and #7001, interleaved runs (two each), CUDA Graph, ms:

tokens/rank	Stage1 aiter + #7001	Stage1 --fuse-quant (8-wave)	Stage1 --gemm1 7001	layer aiter + #7001	layer --gemm1 7001
2048	0.813-0.820	0.911-0.915	0.872-0.874	1.432-1.447	1.492-1.499
4096	1.167-1.174	1.333-1.335	1.171-1.173	2.215-2.244	2.269-2.275
8192	2.022-2.052	2.142-2.149	2.010-2.014	4.219-4.244	4.191-4.192

What bounds it now: the copy engines. At 8192 the puts alone take 1.19 ms and #7001's GEMM1 alone 1.39 ms; running together they share HBM and the fused kernel takes about 1.77 ms. At 2048 / 4096 the puts alone (0.63 / 0.82 ms) already take longer than the GEMM1 alone (0.56 / 0.77 ms). KFD recommends one XGMI SDMA engine per peer link on this node, so each peer's rows go through one engine. Tried and dropped:

• Sending the e8m0 rows over P2P from the quantisation so the copy engine moves only the fp8 rows: no faster (the engines are bound by bytes, not by the number of puts), and the per-token remote stores made route + qpack 0.12 -> 0.83 ms.
• A second engine per peer: a queue on the sibling (odd) XGMI engine hangs, as mori records for MI308X; queues on the general engines 0/1 fail in KFD with NO_MEMORY, even one per process; two queues on the same engine run slower than one.
• Fewer bytes with #7001's source-indexed payload (one row per source token and destination rank, about 26% fewer bytes): its GEMM1 without the tuned route-major bodies is 0.42 ms slower at 8192, more than the copies would save.
• Dropping #7001's per-item acquire (the landing rows are read from the cco window) and a cached window (CCO_UNCACHED_WINDOW=0): both correct, neither measurably faster.

Requirements

• aiter with MegaMoEV2's compact prepare. Measured with aiter f0b96f49e plus the files of [FlyDSL] MegaMoEV2: opt-in Stage2 + combine over SDMA with an 8-wave GEMM2 aiter#7196; the Stage1 path itself is aiter main's. --gemm1 7001 also needs the files of [FlyDSL] [Perf] MegaMoE: faster compact Stage1 producer and FP8 GEMM1 tiles (gfx950) aiter#7001 (mega_moe/gemm1.py), and is compared against that aiter's own Stage1.
• FlyDSL 0.4.1.
• mori built with BUILD_CCO_SDMA=ON, run with MORI_ENABLE_SDMA=1.
• A gfx950 node with 8 GPUs.

MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \
  examples/cco/python/13_flydsl_mega_moe_stage1_sdma/main.py test --bs 8192
MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \
  examples/cco/python/13_flydsl_mega_moe_stage1_sdma/main.py bench --bs-list 2048,4096,8192

Test plan

• test --bs 8192 --scale-mode shuffler, --scale-mode ldsw, test --bs 2048, test --bs 4096: PASS, 0 mismatched rows.
• bench --bs-list 2048,4096,8192: byte-identical before and after the graph replays.
• --fuse-quant: test --bs 8192, --bs 4096, --bs 2048: PASS, 0 mismatched rows; bench --bs-list 2048,4096,8192: 0 mismatched rows. Default mode test --bs 8192 re-run: PASS.
• After 1d2b90d: test --bs 8192 (shuffler), --bs 4096, --bs 2048 (ldsw), and --fuse-quant at 8192 and 4096: PASS, 0 mismatched rows. bench with and without --fuse-quant (which also checks every token count, 2048 included): two runs each, interleaved with the previous commit, 0 mismatched rows. The committed files re-ran test --bs 8192, --bs 2048, --bs 4096 --fuse-quant: PASS.
• After 39f7618 (aiter + #7001 files): test --gemm1 7001 at 8192 / 4096 / 2048 and test --bs 8192 --fuse-quant, test --bs 2048 (8-wave): PASS, 0 mismatched rows. bench --gemm1 7001 and bench --fuse-quant two runs each, interleaved, 0 mismatched rows. The committed files are the ones that ran.
• pre-commit: black / ruff / license hooks pass. The black reformat leaves the AST unchanged against the code that was run.

🤖 Generated with Claude Code