[kawhil-amd]

Adds examples/cco/python/11_flydsl_mega_moe_sdma: a MegaMoE stage2 that writes
its GEMM2 output into a local, rank-sorted send buffer and has comm waves
inside the GEMM kernel ship finished chunks to the peers with cco SDMA, so the
transfer runs on the copy engines instead of the compute units.

The reference it is measured against is aiter's MegaMoEV2, whose fused GEMM2
scatters every output row straight into the origin rank's comb_inp with P2P
stores from the CUs.

Future work: gfx1250 porting, and investigate the cco-sdma use in dispatch + gemm1 fusion (mega-moe stage1).

Files

file	what
mega_moe_gemm2_with_sdma.py	the op (MegaMoESdma) and its three kernels: the plan kernel, GEMM2 with comm waves, the combine kernel; the window layout and the sender side of the SDMA protocol they share. The module docstring is the design document
main.py	test and benchmark against aiter's MegaMoEV2: main.py test (output vs aiter's) and main.py bench (stage2+combine latency)
sdma_bandwidth.py	what one cco SDMA channel moves, with no GEMM in the way

Results

stage2+combine on glm52, 8x MI355, eager launches, fp8 transport,
--chunk-tiles 16, --no-aligned-pair. Output matches aiter bit for bit
(relL2 0.0e+00, 0 mismatched rows) at every size below. Base model config (glm-5)

tokens/rank	128	512	2048	4096	8192	16384
aiter-9b7bb36ad (ms)	0.085	0.159	0.510	0.970	1.831	3.833
this (ms)	0.121	0.205	0.504	0.864	1.483	2.843
speedup	0.71x	0.78x	1.01x	1.12x	1.23x	1.35x

bs is per rank. Single run; the spread is about ±0.05x at 128/512 and ±0.02x
from 2048 up. Below ~2048 the fixed costs of the extra plan kernel and of the
combine outweigh what the copy engine saves, and aiter's P2P scatter wins.
This is a large-batch (prefill) optimisation, not a general one.

Only stage2+combine is compared. That is the part this PR replaces; dispatch and
stage1 are shared and are not timed.

Where the time goes at 8192 (rocprof kernel times, all ranks): the GEMM2 kernel
is 1255-1280 us against aiter's compact stage2 at 1639-1674 us -- but that is "GEMM plus
its communication", since aiter's kernel carries the P2P scatter and this one hands
the transfer to the copy engines. What is paid back is the plan kernel (~18 us), the
combine (133 us of work, plus waiting for the slowest peer) and a GEMM2 that is
still 100-120 us above its own floor (1.12-1.16 ms with the epilogue stores and the
puts removed). The reduce in the combine reads 6.5e4 rows of 6.3 KB and is close to
the memory-bandwidth limit.

Design

Three kernels on one stream, against the baseline's two:

1. plan -- gives every valid sorted row a stable index in its destination's
   run, so a range of sort tiles maps to a contiguous row range per peer.
2. gemm2+sdma -- the aiter GEMM2 mainloop with an epilogue that writes
   send[out_row] instead of scattering, plus comm_wgs leading workgroups
   (one wave per peer) that fill the peer's inv table and then put finished
   chunks while the remaining tiles still compute.
3. combine -- waits for the peers' completion signals and does the top-k reduce.

The contiguity in (1) is what makes (2) cheap: a chunk is one ~1.4 MB SDMA put
rather than a scatter of small ones. The receiver needs, for each of its token
slots, the row to read; the sender writes that table (inv, in the cco window):
the plan kernel for rows that stay on its own rank, the comm waves -- direct stores
into the peer's window -- for the rest. So no index travels with the data and the
combine builds no table and needs no grid barrier.

What makes it fast

Four things, in order of what each was worth. Each is backed by an A/B in the
same harness.

fp8 transport (--quant fp8_blockwise_1x32, default)

bf16 moved twice the bytes and the transfer stopped fitting under the GEMM. A
truncated-shipping probe showed that of the +0.219 ms the comm cost at bs=4096,
+0.205 tracked bytes, not sync. Rows now carry aiter's fp8_blockwise_1x32
in aiter's own row layout -- payload, then the e8m0 scales, inline -- so one
contiguous put still ships a whole row. Reuses _fp8_scale_for_leader and
cvt_scalef32_pk_fp8_f32; the combine dequantises with cvt_pk_f32_fp8.

Write-through epilogue stores

The per-unit publish released at system scope. On gfx950 that lowers to
buffer_wbl2 sc0 sc1 -- a write-back of the whole L2, not of the rows just
written -- charged once per (m_block, n_block) unit, 12288 times per launch.
Storing the rows sc0|sc1 instead makes retiring them the release the copy
engine needs, and the fence disappears. Confirmed in the ISA: the per-unit
buffer_wbl2 is gone and all 8 epilogue stores carry sc0 sc1.

2.4x on the GEMM2 kernel. Same recipe as ops/gemm_ar/kernels_sdma.py
(publish="writethrough"), whose docstring explains it.

Note: buf_copy_store's cache_modifier is a BufferCopy enum (0=cached,
2=nt), so sc0|sc1 cannot be reached through it -- the epilogue uses
buffer_store, which passes the value through as the instruction's aux field.

The per-chunk completion counters are also a cache line apart, which is worth
another ~8% on its own.

A per-pair row reservation

cap was the whole token-slot space reserved for every (src, dst) pair --
16x what a pair can use -- which put bs=8192 over the 4 GiB buffer-resource
range. It is now the even share times --row-slack (default 1.5), with the plan
kernel refusing to place a row past it and recording the overflow. This keeps
the handshake-free write into the peer's landing buffer at a statically known
offset.

Window at bs=4096 drops 3.32 GB -> 0.58 GiB; bs=8192 fits at 1.25 GB and bs=16384
at 2.5 GB (fp8).

A plan kernel with no serial chain, and a combine with no build phase

The plan kernel is ~17 us at 8192 tokens/rank (43 us at the start of this round,
66 us at the start of the PR). Each phase was timed on the GPU -- rocprof, or CUDA
graph replay of the kernel alone; timing launches from the host hid them behind the
launch rate:

• Re-decoding the tile for every row (O(SBM) per row, ~12 us). Count and assign
  run one block per tile and one thread per row: each row is decoded once, a wave
  ballot per destination pe plus a popcount gives its rank among the earlier rows
  for that pe, and the per-wave counts cross waves through a few LDS words.
• A Hillis-Steele scan run redundantly in every block (10 barriers, ~15 us). A
  block only needs the prefix of the tiles it owns: one parallel pass over the
  earlier tiles' counts and a short LDS reduction. It runs to the tile's own index,
  so there is no bound on the tile count to get wrong (see below).
• The grid barrier (~11 us). First every block atomically added to one address
  (the atomics queue in L2), then every thread of every block polled every block's
  flag (65k polling loads slowed the stores they waited for). Now it is a two-level
  tree: 32 blocks to a counter, the last block of a group bumps a top counter, and
  each block polls that one address.
• Per-row remote stores of inv (9 us). Moved out of the plan into the GEMM2
  kernel's comm waves, which are idle until the first chunk is written.

The combine runs two blocks per CU now that it has no grid-wide rendezvous: more
row loads in flight for the reduce (8192: 1.528 -> 1.492 ms; 4x and up was slower).

Outputs of the plan kernel are bit-identical to the version before the rewrite on
the same stage1 output.

The comm wave also hangs the copy engine's own trailing ATOMIC on its last put
(remote_signal) instead of draining the queue, so it no longer pins the tail
of the GEMM2 kernel; the peer's combine -- which was going to wait anyway -- polls
the counter with wait_signal. Worth 6-12 us. --signal-mode quiet keeps the
old drain-and-flag path.

Correctness under skewed routing, overflow and call overlap

Review of the first revision found cases that were wrong without any report. They
are fixed and tested (main.py test --hot-bias):

• Skewed routing. The plan's prefix was bounded by a tile count sized from the
  send side, but a rank receives rows from every peer, so one rank getting more
  than ~6% over an even share at bs=4096 made the prefix silently wrong on every
  rank. The prefix no longer has such a bound. main.py test --bs 4096 --hot-bias {0.05,0.1,0.3} (rank 0's experts favoured, up to 70% over the even share) is
  bit-exact; the previous revision fails at 0.1.
• Overflow is loud. A row past the per-pair reservation is dropped, which makes
  the output wrong, so MegaMoESdma.check_overflow() raises and main.py calls it.
  tile_base is also clamped to cap, so a chunk that crosses the reservation is
  cut at it instead of reading the next peer's run and writing past this rank's slot
  in the peer's landing buffer. With the default --row-slack 1.5, --hot-bias 0.3
  fails with a message naming row_slack; --row-slack 2 runs and is bit-exact.
• main.py test is stricter. When both sides move fp8 it requires 0 mismatched
  rows, not just relL2 below the tolerance.
• Back-to-back calls. The plan waits for every peer's "consumed" flag before the
  GEMM2 that follows rewrites S, which is what guarantees this rank's own outbound
  puts of the previous call have finished reading it; the comm waves wait for the same
  flag before filling a peer's inv or putting into its landing buffer. close()
  waits for the flags before freeing the window. The reviewer could not trigger the
  unguarded case in 16 back-to-back calls with a delayed rank, and neither did I, so
  this is added from the protocol rather than from a failing run.
• num_valid == 0 writes zeros to cursor and tile_base instead of leaving the
  previous call's values.

Testing

MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 main.py test --bs 4096
MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 main.py bench \
    --bs-list 128,512,2048,4096,8192,16384 --no-aligned-pair

• main.py test is bit-exact (0 mismatched rows) at 2048 / 4096 / 8192 / 16384
  tokens per rank with uniform routing, at 4096 with --hot-bias 0.05 / 0.1 / 0.3 --row-slack 2, and with --signal-mode quiet. At 128/512 aiter picks
  p2p_quant=none (bf16) when it chooses for itself, so the 2.67e-02 there is this
  path's fp8 quantisation error, within the 0.05 tolerance.
• The SDMA output hashes were compared across the refactor (8 configurations x 8
  ranks: bs 128/2048/4096/8192, two skewed routings, quiet signalling, bf16 rows):
  identical.
• 4000 back-to-back calls per size (128 / 2048 / 8192) with the output checked
  against the first call's: no wrong output. A multi-ms stall shows up every few
  thousand calls, with and without this PR's later commits.
• Checked on two aiter versions with the same numbers within noise: aiter
  2247dec34 + FlyDSL 0.3.0, and aiter main 9b7bb36ad (2026-10-28) + FlyDSL 0.4.0
  (speedup at 2048 / 4096 / 8192 / 16384: 1.05x / 1.14x / 1.24x / 1.34-1.35x).
• ISA dumped (FLYDSL_DUMP_IR=1) to confirm the write-through stores and the
  removal of the per-unit buffer_wbl2.

Ruled out by measurement

Recorded so the next person does not repeat them. All measured in the same
harness, at bs=4096 unless noted:

Idea	Result
Shard the chunk counter across n_blocks	No change — the atomic is not contended. (The cache-line padding is load-bearing; the sharding on top of it is not.)
comm_wgs 1 → 8	Within 1%. The comm workgroups stealing CUs is not a factor.
queues > 1, splitting or rotating puts across channels	Slower. KFD recommends a single SDMA engine per peer link on this topology (recommended_sdma_engine_id_mask has one bit per link), so extra channels share it and only add descriptors. Splitting a 1.44 MB run in two cost 3%; rotating whole runs across two channels 2%.
Larger chunk_tiles for link efficiency	Slower. 5.78 MB chunks reach 95% of link peak vs 81% at 1.44 MB, but the longer tail and later first ship cost more.
Deeper GEMM pipeline (persist_cu 240 → 64, or persist_strided)	Uniformly worse on both axes. strided makes the GEMM itself faster (0.615 vs 0.625 ms) but the comm cost jumps 0.139 → 0.383 ms.
Merging a token's same-destination top-k slots	A 33.8% byte cut buys 2.5%. aiter's own aligned_pair is 1.5% slower than without it on this routing. Caveat: measured on uniform random routing, where there are no hot expert pairs. Real skewed routing was not tested and could change this.
Optimising the combine's reduce	It reads 6.5e4 rows x 6.3 KB and writes the output at close to the memory-bandwidth limit; the gain was in removing what surrounded it.
Combine at 3-8 blocks per CU	Slower than 2 per CU at every size.
Tree-barrier group of 8	Slower than 16-64; sleeping between polls made no difference.

What is left. At 8192 the plan is ~18 us plus a launch, the GEMM2 is 100-120 us above
its own floor (the write-through stores and the SDMA traffic competing with the
GEMM), and the combine waits ~60 us for the slowest peer, of which roughly 40 us
is the last chunk's transfer after the GEMM ends (inferred, not measured on its
own). Folding the plan into the GEMM2 kernel is less attractive than it looked: that
grid is cu_num * n_blocks (~5760) workgroups, far more than are resident, so a grid
sync across all of them would deadlock, and dedicated plan workgroups would occupy
CUs for about as long as the launch they save.

Limitations

• gfx950 only, intra-node, 8 ranks, MORI_ENABLE_SDMA=1.
• --queues must be 1 for --signal-mode remote (the default), since the
  signal rides the channel the data went out on.
• The buffer-resource range is 4 GiB: bs=16384 per rank fits with fp8 rows (2.5 GB),
  but not with bf16 rows (--quant none, ~4.8 GB).
• No aligned_pair equivalent, so bucket 8192 needs --no-aligned-pair to
  compare like for like.
• Routing is assumed to stay within --row-slack of an even share per (src, dst)
  pair. Beyond it rows are dropped and check_overflow() raises; raise --row-slack.
  Library users must call check_overflow() themselves -- run() does not sync.
• main.py bench builds one MegaMoEV2 sized for the largest bs, so at 128/512
  aiter runs fp8 transport, where on its own it would pick bf16.
• MegaMoESdma reads private attributes of aiter's MegaMoEV2 (_s1_op, _s1_nvm,
  _g2v2_*, ...), so it follows aiter's internals.

🤖 Generated with Claude Code

[jhchouuu]

Ran this on 8x MI355X (ROCm 7.14, FlyDSL 0.4.0, aiter main d5dc194f) with the branch merged into current main. With uniform routing your numbers reproduce, but with even mildly skewed routing the output is wrong on every rank and nothing reports it.

Reproduced (uniform routing)

• main.py. Bit-exact against aiter at bs 2048, 4096 and 8192 (8192 run with --no-aligned-pair): 0 mismatched rows on every rank and overflow 0. bs 128 and 512 give 2.67e-2, as you describe.

• bench.py.

  	128	512	2048	4096	8192
  stage2+combine	0.65x	0.73x	0.99x	1.11x	1.21x
  e2e	0.88x	0.85x	1.02x	1.05x	1.10x

  --graph agrees with eager.

• sdmabw.py. 47.3 GB/s per link at 1.44 MB per put, which fits your 5.8 µs + bytes/60.9 GB/s model.

• Back-to-back calls. I ran 16 calls with no host sync and rank 3 delayed by torch.cuda._sleep (on every call, and on alternate calls) at bs 2048, 4096 and 8192. Every call is bit-exact.

1. Skewed routing gives wrong results on every rank, with overflow 0 (blocking)

I added a bias to the router scores of rank 0's experts in _make_inputs (scores[:, :32] += b) and ran main.py at bs 4096, otherwise with defaults:

bias	rank 0's rows vs the even share	mismatched rows per rank (of 4096)	relL2	overflow
0.05	+9.8%	2496-2648 (and it varies from call to call)	0.36	0
0.1	+21%	2972-3091	0.40	0
0.3	+70%	3497-3583	0.48	0

Rerunning bias 0.3 with --row-slack 8 gives the same result, so the cause is not the per-pair reservation. It is the plan's scan buffer, which is sized from the send side:

• TILES_MAX = ceil(slot_cap/SBM) + local_experts + 1 (sorted_combine.py:231).
• But the scan runs over num_valid, the rows this rank receives from all peers. That can be up to npes * slot_cap, and aiter sizes it with num_valid_max.
• At bs 4096, once expert padding is counted, that leaves only about 6% headroom over the even share.
• Nothing checks tiles <= TILES_MAX, so the scan indexes past its half of the LDS buffer.
• Every rank ends up wrong because the hot rank's GEMM2 output feeds all of them.

Real routing is rarely within 10% of even, so I think this needs a fix before anything else. Size the scan from num_valid_max or tile it. At the least, fail when tiles > TILES_MAX.

2. Per-pair overflow is silent and not contained

• Nothing reads self.overflow. It is allocated and passed to the plan, but neither main.py nor bench.py checks it. With --row-slack 1.03, rank 3 overflowed 72 rows and 18 rows of its output were wrong (relL2 2.45e-2), yet main.py printed PASS and SUCCESS.
• Only the final per-peer total is clamped to cap (sorted_combine.py:313-315). The intermediate tile_base entries are not.
• So a chunk that crosses cap ships past this rank's slot. The comm wave puts [tile_base[t0], tile_base[t1]) (moe_gemm2_localsort.py:596-602), so that chunk:
  • reads the next peer's run in S;
  • writes past this rank's slot in the peer's landing buffer, and past R for the last rank.
• Suggested fix. Clamp lo/hi to cap in the comm wave, and raise when overflow is non-zero.
• A tighter test would also catch this. At bs >= 2048 the reference is bit-exact, so main.py could assert 0 mismatched rows there instead of relL2 < 0.05.

3. Smaller points

• Own outbound puts are never waited on before reuse (remote mode). Nothing waits for this rank's own puts to finish before the next call's plan and GEMM rewrite idx_send and S. The consumed handshake runs in the comm waves, after the plan has already written. I could not trigger it in the runs above, but no ordering guarantees it. The same applies to close() right after a stream sync.
• tiles == 0 reads at a negative index. When num_valid == 0, sc(final, (tiles - 1) * npes + tid) (sorted_combine.py:311) reads at a negative index.
• The e2e comparison is not like for like.
  • The baseline moe(x) includes aiter's input quantisation, while the new path's stage1() starts from the pre-quantised x_q.
  • At 8192, base_e2e still uses aiter's aligned-pair stage2 even with --no-aligned-pair.
• The 128/512 baseline is not aiter's own choice. bench.py builds one MegaMoEV2 sized for the largest bs, so at those sizes aiter runs fp8 transport. On its own, aiter picks bf16 there (main.py shows p2p_quant=none).
• The Communicator is never destroyed. SortedCombine.close() only calls close_all_resources(); comm.destroy() would also free it. sdmabw.py's comm.finalize() does not exist, so that line is a no-op.
• The encoding comments point to the old location. They reference mori/ops/gemm_ar/_compat.py. Since #817 the encoding lives in mori.cco.device.flydsl.buffer, and CM_SC0_SC1 could replace the literal 1 | 16.

[kawhil-amd]

Thanks for the thorough run, and for the skewed-routing repro in particular. Fixes are in bea7f03 (plus ac0a306 for the plan kernel speed-up); the PR description is updated to match.

1. Skewed routing (blocking) -- fixed. Confirmed. The plan's prefix was bounded by a tile count sized from the send side (slot_cap), and nothing checked tiles against it, so past ~6% over an even share at bs=4096 the prefix was silently truncated. The prefix now sums up to each tile's own index, so there is no bound to exceed. I added main.py --hot-bias (adds a bias to rank 0's experts' router scores, as in your repro). At bs=4096 the previous revision fails at --hot-bias 0.1; the new one is bit-exact (0 mismatched rows) at 0.05, 0.1 and 0.3 (0.3 with --row-slack 2, since rank 0 then receives ~70% over an even share).

2. Overflow -- fixed.

• SortedCombine.check_overflow() raises if rows were dropped, and main.py / bench.py call it. With the default --row-slack 1.5, --hot-bias 0.3 now fails with a message that names row_slack instead of printing SUCCESS.
• tile_base is clamped to cap when the plan writes it. I did it there rather than in the comm wave so that nothing reading tile_base can see an unclamped boundary; a chunk that crosses the reservation is cut at it.
• main.py now requires 0 mismatched rows when both sides move fp8 (your suggestion), in addition to the relL2 tolerance.
• run() itself still does not sync, so library users have to call check_overflow(); this is in the Limitations section.

3. Smaller points

• Own outbound puts before reuse -- guarded. I agree nothing ordered it. The plan now waits for every peer's "consumed" flag before it rewrites idx_send (and the GEMM2 after it rewrites S). A peer sets that flag after its combine has read my data, and the copy engine's queue is FIFO, so by then my puts have finished reading. close() waits for the same flags before freeing the window. Like you, I could not make the unguarded case fail, so this comes from the protocol, not from a failing run.
• tiles == 0 -- fixed. After the plan rewrite this would have left cursor / tile_base holding the previous call's values rather than reading a negative index; it now writes zeros. I did not build a test that routes no rows at all.
• e2e not like for like -- not changed in code; the description now says so. moe(x) includes aiter's input quantisation and, at 8192, still takes aiter's aligned-pair stage2 under --no-aligned-pair (I see that kernel in a rocprof trace). The stage2+combine rows are the like-for-like ones.
• 128/512 baseline -- not changed; also noted in Limitations. bench.py builds one MegaMoEV2 for the largest bs, so aiter runs fp8 there.
• Communicator never destroyed, sdmabw.py comm.finalize(), and the _compat.py / CM_SC0_SC1 comments -- not done in this round. They are not correctness issues, and I have not yet checked the comm.destroy() and mori.cco.device.flydsl.buffer APIs you point to. I will do them in a follow-up commit.

Re-measured on 8x MI355 after the fixes: stage2+combine vs aiter is 0.66x / 0.75x / 0.98x / 1.07x / 1.20x / 1.31x at 128 / 512 / 2048 / 4096 / 8192 / 16384, bit-exact in every case. Your 1.11x / 1.21x at 4096 / 8192 are in line with mine; the small-batch ratios move by about ±0.05x between runs.

[kawhil-amd]

cc @yanboshao @XingerZhu

[yangyuhuiling]

Thanks for the detailed write-up. I tried a few follow-ups on top of 0ad1c7b; two small ones help, and two more are worth adding to your "ruled out" table.

Setup. 8x MI355X, aiter 9b7bb36ad, FlyDSL 0.4.0, ROCm 10 / torch 2.11, glm52, fp8 transport, --chunk-tiles 16 --no-aligned-pair, eager. Each variant was timed in interleaved rounds (base, variant, base, ...) and I report the median: 3 rounds per size, 5 for the combined row. Every variant is bit-exact against aiter (0 mismatched rows) at 2048 / 4096 / 8192 / 16384, and at 4096 with --hot-bias 0.1 --row-slack 2.

1. Data puts post a local signal that nothing waits for

sdma_put_run calls sdma.put(...) without a signal keyword. Under sdma_signals(), a put that names no signal keeps the old default, which is a local signal. So every data put trails a 32 B ATOMIC into this rank's own signal pool. Nothing reads that pool: completion is the remote signal under remote, and quiet_queue under quiet. Passing local_signal=False drops it:

sdma.put(peer, win, dst_off, win, src_off, nbytes, queue,
         coop=cco.CoopScope.THREAD, local_signal=False)

2. The completion signal can ride the last data put

Under remote, sdma_finish_peer posts a separate 4 B put just to carry the signal. The final chunk's put can carry remote_signal=True instead. A peer whose final chunk is empty still needs the separate signal put, which is the only case that keeps it.

tokens/rank	2048	4096	8192	16384
no local signal	−3.6%	−0.6%	−0.5%	−0.55%
signal on last put	−0.4%	−0.4%	−0.3%	−0.5%
both		−1.1%	−0.5%

The gain is largest at 2048, where per-put fixed cost is the biggest share of the total.

Tried and ruled out

Idea	Result
Smaller chunks only at the tail (last 16–64 tiles in chunks of 4–8, head kept at CT=16)	Alone: +0.9% at 4096, +0.75% at 8192, +3.6% at 2048. Combined with "no local signal" it is 0.2–1% faster, which is within run-to-run spread.
Bigger head chunks (CT=32) plus a small tail	1.3% slower at 8192 and about 3.5% slower at 4096, with or without the tail. This agrees with your finding on chunk size.
One packed 8-byte store for a row's e8m0 scales (readlane from the block leaders) instead of the per-leader byte store	3.5–4.1% slower at ≥4096. The original byte store is already a single instruction across the leader lanes; the readlanes just add work to the epilogue loop.

The patch for (1) and (2) is a few lines. Happy to push it as a commit or open a PR against your branch, whichever you prefer.

🤖 Generated with Claude Code

[kawhil-amd]

@yangyuhuiling nice idea, i will try it.

[yangyuhuiling]

@kawhil-amd I pushed the two changes from my previous comment as b6e4404 on your branch (maintainer edits were enabled). Please pull it before you push anything else, so your local copy stays in sync.

What it changes, all in mega_moe_gemm2_with_sdma.py:

• sdma_put_run passes local_signal=False, which drops the unused trailing ATOMIC on every data put.
• The comm-wave loop passes remote_signal=True on the last chunk's put under signal_mode="remote".
• sdma_finish_peer(signalled=True) then skips the separate 4-byte signal put. A peer whose last chunk is empty (or a call with no chunks) still gets that put.
• quiet mode is unchanged apart from the dropped local signal.

Correctness. Re-verified on 8x MI355X at b6e4404 with the same aiter/FlyDSL as before. main.py test passes at:

• bs 4096, 8192 (--no-aligned-pair), 4096 with --hot-bias 0.1 --row-slack 2, and 4096 with --signal-mode quiet: bit-exact, 0 mismatched rows;
• bs 128: within tolerance as before (relL2 2.7e-2, the fp8-vs-bf16 difference). Some peers have an empty last chunk at this size, so it also exercises the fallback signal put, and nothing hangs.

Performance. main.py bench, 3 interleaved rounds against 0ad1c7b, medians of stage2+combine:

tokens/rank	0ad1c7b	b6e4404
2048	0.4695 ms	0.4652 ms	−0.9%
4096	0.8676 ms	0.8611 ms	−0.75%
8192	1.4791 ms	1.4720 ms	−0.5%

These gains are a bit smaller than the ones in my previous comment, which came from a quieter run (2048 was −3.6% there).

A few 20-call bench means at 8192 came out 6–8% high. To check that this was not a new stall, I timed 400 individual calls, 3 runs per side at 8192:

	0ad1c7b	b6e4404
median	1.4725 ms	1.4637 ms
p99	1.517–1.525 ms	1.499–1.513 ms
calls > 1.2x median	0	0

So the high means were bench noise; the new commit does not stall more often.

🤖 Generated with Claude Code

[kawhil-amd]

Two fixes on top of this PR: +0.09 to +0.13x on the stage2+combine speedup

Both are costs the fused stage2 was paying on top of a GEMM that was already running at full speed. Measured EP8 on MI355X, fp8, uniform routing, row_slack 1.5, --no-aligned-pair, chunk_tiles 16. Every number is bit-exact against aiter (relL2 = 0).

Results — stage2+combine vs aiter MegaMoEV2

tokens/rank	GLM-5.2 Config	V4.1-Flash Config	V4-Pro Config
4096	1.15 → 1.27x	1.03 → 1.15x	0.96 → 1.10x
8192	1.26 → 1.38x	1.15 → 1.27x	1.07 → 1.20x
16384	1.39 → 1.52x	1.15 → 1.33x	1.08 → 1.21x

1. The AMDGPU UnclusteredHighRPReschedule stage

Fusing the comm role into this kernel cost 8.5% of GEMM2 at 16384 tokens/rank even when the comm role never ran — a diagnostic build that compiled it in behind a runtime-false guard paid the same 0.175 ms.

Ruled out first:

hypothesis	evidence against
register pressure in the K loop	same opcode multiset either way — one s_nop and one s_waitcnt apart
SGPR spills	13 spills, all in the prologue, all v_writelane filling s_nop slots; private_segment_fixed_size = 0
L1 instruction cache	.text 12.0 → 23.2 KB; L1I is 64 KB
code size	duplicating the comm role 4x took .text to 45.8 KB and moved the cost 2.3%

What actually changes is the instruction order. The comm role's live ranges lift the function's register pressure over the threshold that turns the stage on, and it then breaks the main loop's clustering:

K loop (bs=16384, GLM-5.2)	comm role absent	comm role present
instructions / MFMA / loads / VALU	155 / 32 / 12 / 81	153 / 32 / 12 / 81
position of the first MFMA	57	34
VALU before the first MFMA	41	23
mean load → s_waitcnt distance	20.9	13.6

Same work, 35% less to hide the load latency with, paid on every trip. Compiling with amdgpu-disable-unclustered-high-rp-reschedule restores the original schedule.

One gotcha: FlyDSL emits code on the first launch, not inside compile_moe_gemm2_with_sdma(), so the llvm_options scope has to cover that launch too.

2. persist_cu on the fused path

A comm workgroup holds one of its CU's two workgroup slots for the whole kernel. ktrace at 16384 tokens/rank (V4.1-Flash):

	comm role absent	comm role resident
comm WG residency	508 / 120 ticks	194132 / 193152
GEMM WGs on those two CUs	17 / 18	8 / 8
GEMM WG duration, p50	25396	25092
mean per-CU finish	185856	185798
last CU finish	188804	195072

Every workgroup keeps its normal duration and the mean CU finish does not move — only the tail stretches. So it is not interference: the displaced work has to be redistributed, and at the bucket's tuned persist_cu = 192 one workgroup is ~13% of the kernel span, so the tail quantises badly.

Giving the fused path every CU makes workgroups shorter and more numerous, and the tail absorbs the displacement. Fused kernel cost over plain GEMM2 at 16384 tokens/rank:

	persist_cu 192	persist_cu 256
V4.1-Flash	0.200 ms	0.133 ms
V4-Pro	0.274 ms	0.184 ms

Levers that did not help, so they need not be retried: dropping persist_cu, moving the comm workgroups to the end of the grid, making the comm WG count a multiple of the XCD count, poll backoff (1024x slower polling changes nothing), comm_wgs 1 vs 2, push_inv, and cacheline contention on the cdone counter.

Note on the V4.1-Flash rows

The harness gains the V4.1-Flash shape (5120 / 2304 / 384 experts / top-6) so those rows reproduce. aiter has no tuned profile for it at 48 experts per rank — its v41_flash_ep2 profile wants 192 — and _tune_v4_pro returns early unless apply_fp4. So all three networks here run aiter's generic stage2 config; these numbers are not aiter's tuned operating point for Flash.

[jhchouuu]

Second pass at f5771d11, merged onto current main. Stack: 8x MI355X, ROCm 7.14, torch 2.12, FlyDSL 0.4.0, aiter 9b7bb36ad.

Every first-round point I tested is fixed. I found one new problem: the kernels' JIT cache identity does not include the window layout, and that hangs a later process.

The first-round fixes hold

• Skewed routing. main.py test --bs 4096 --hot-bias {0.05, 0.1, 0.3} --row-slack 2 is bit-exact (0 mismatched rows).

• Overflow is loud.

  • --hot-bias 0.3 at the default slack fails, and rank 0 reports rows did not fit the per-pair reservation ... raise row_slack.
  • --row-slack 1.03 with uniform routing now fails. In the first round it printed SUCCESS.

• Uniform routing. Bit-exact at bs 2048 / 4096 / 8192 / 16384, with --no-aligned-pair from 8192 up. --signal-mode quiet is bit-exact at 4096.

• Back-to-back calls with data that changes. The existing checks (4 calls in test, and your 4000-call run) reuse the same inputs. So a combine that read a peer's previous rows, or a sender that overwrote a slow peer's landing buffer, would still produce matching output.

  I added a variant that doubles the stage1 output scale in place on every odd call (_s1_osd holds e8m0 bytes, so adding 1 is an exact x2). Call k must then equal ref * 2**(k % 2), and a synced sanity pair confirms the check has teeth: the x2 call matches 2*ref exactly and differs from ref on all 4096 rows.

  I ran 16 calls with no host sync, at bs 4096, with torch.cuda._sleep(50M) before every call (or every other call) on one rank:

  • remote and quiet;
  • --hot-bias 0.1 with the hot rank as the delayed one.

  0 mismatched rows on every rank and every call. It might be worth folding this into main.py test, since it is the only form that can catch the handshake failing.

1. The kernel cache key misses the layout, so a later process can hang (blocking)

In one container, main.py test --bs 64 followed by main.py test --bs 16 makes the bs-16 run hang in its first sc.run (stuck in the synchronize after it, GPUs busy). Each row below is a fresh container:

sequence	bs 64	bs 16
bs 16 alone		SUCCESS
bs 64, then bs 16	SUCCESS	hang (timeout)
bs 64, then bs 16, FLYDSL_RUNTIME_ENABLE_CACHE=0	SUCCESS	SUCCESS

The cause. FlyDSL's persistent cache identifies a kernel by its source plus its scalar closure values. layout (ArenaLayout), comm (CommSpec) and dec (a dict) are not scalars, so none of what they carry reaches the key:

• build_combine reads layout.off_counter / off_flags / off_done inside the kernel (mega_moe_gemm2_with_sdma.py:1320-1324, 1418).
• The comm role reads lay.off_counter / off_done (:1084, :218).
• The plan kernel reaches log2_max_tok and recv_cap only through **dec (:414, :489).

off_sig and every offset after it depend on slot_cap, but the kernel names carry cap and not slot_cap. With row_slack=1.5, max_tok 16, 32 and 64 all give cap=128. The bs-16 process therefore loads the bs-64 binaries and spins on a flag at the wrong offset.

This reaches any process that follows another one with a different max_tok on the same ~/.flydsl cache, a server restart included. Within one process, two MegaMoESdmas with different max_tok may well hit the same thing, but I did not test that.

The fix. Hoist every value a kernel bakes in into a plain int local before the @flyc.kernel, so it becomes a closure scalar and enters the key. That means the offsets, slot_cap, log2_max_tok and recv_cap. Alternatively, put them in the kernel names. ops/gemm_ar/kernels_sdma.py (_phase, phase_specialization) does this for the same reason.

2. The harness hangs on this stack before it reaches the SDMA path

main.py test times the reference with _time(lambda: moe._run_stage2(...)), and bench does the same in base_s2. Both run aiter's _run_stage2 back to back with no host sync and no dispatch in between. Here, two such calls already hang inside aiter, while a synced pair runs fine. So both test and bench hang before the SDMA path runs.

I got through by syncing before each reference call. Which ROCm and torch did you run on? If this is aiter-version dependent, a sync in the reference timing, or a note in the README, would save the next person the hunt.

I did not re-measure performance this round: the shared node picked up another job partway through, so those numbers are discarded.

Smaller points

• close() can free the window early. _wait_peers_consumed returns silently when it times out (30 s), and close() then frees the window anyway. Raising, or at least warning, would turn a stuck peer into an error instead of a use-after-free on the copy engine.
• Leftovers from the first round:
  • The communicator is still only close_all_resources()'d, never destroyed.
  • sdma_bandwidth.py's comm.finalize() is still a no-op behind hasattr.
  • The _compat.py / 1 | 16 comments still point to the old location (mori.cco.device.flydsl.buffer / CM_SC0_SC1).
• The README table predates f5771d11. The README and the PR body show 1.12x / 1.23x / 1.35x; your latest comment has 1.27x / 1.38x / 1.52x.
• Example 11 is also taken by #866. It adds examples/cco/python/11_flydsl_sdma_allgather_engine_wait and edits the same README rows, so whichever lands second needs renumbering.