mega_moe: add a mori HIP dispatch backend - #4757
Merged
yanboshao merged 1 commit intoAug 14, 2026
Merged
Conversation
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 ROCm#548 or later) and its libmori_ops_v2.so.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Makes
MegaMoEGfx1250's first stage swappable:dispatch_backend="mori"(orMEGA_DISPATCH=mori) runs mori's HIP/JIT gfx1250 TDM dispatch instead of thispackage's FlyDSL one. Nothing else in the pipeline changes — gemm1, the gemm2
P2P scatter epilogue and the fused combine are untouched, and the default is
unchanged, so not setting anything is byte-for-byte today's path.
Why it composes
The two dispatches leave identical state in the arena:
disp_outslot * hiddenout_idx/out_wts[slot * topk + k], full topk rowdest_pe * max_recv + slot, nullnpes * max_recvrecv_to_src_tokensrc_pe * max_tok_per_rank + src_tokrecv_num/tok_offThe reverse map is the one field the fused path actually consumes — the GEMM
host pass decodes it to build
ep_rowmap. mori's dispatch also never touchescross_device_barrier(its cross-rank rendezvous is therecv_numhandshake), so the fused combine's phase counter is undisturbed.
mori's plan layer takes an arena by duck typing (
.handle/.offset), soSymmetricArenais handed over as-is: no new region, no extra device memory.Two things that do differ, handled here
_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 KBagainst a320 KB budget.
EpCfgIsValiddoes not check LDS, so it would fail at launchrather than when the plan is built.
total_recvis zeroed on the host. This package's dispatch clears it inits own Phase 2; mori's only accumulates, and the fused combine never resets it.
One visible consequence: the recv slot a token lands in changes, because 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
4x gfx1250, EP4 hidden 7168 topk 6,
--combine scatter_fused, againsttest_mega_moe.py's fp32 reference:The diffs are identical because dispatch only moves data — the deviation from
the reference should not move, and it does not.
--profile_table 1confirms thegeometry that actually launches is mori's:
..._64x8at 256 tokens,..._64x16at 4096.Requires
A mori built with JIT v2 (ROCm/mori#548 or later) and its
libmori_ops_v2.so.Import failure is reported with which half is missing. Selecting the backend
costs ~5 s at construction (three geometries, one JIT compile each), before any
cudagraph capture.