Skip to content

mega_moe: add a mori HIP dispatch backend - #4757

Merged
yanboshao merged 1 commit into
ROCm:yanbo/mega_stage2_gfx1250from
jhchouuu:jhchouuu/mega-moe-mori-dispatch
Aug 14, 2026
Merged

mega_moe: add a mori HIP dispatch backend#4757
yanboshao merged 1 commit into
ROCm:yanbo/mega_stage2_gfx1250from
jhchouuu:jhchouuu/mega-moe-mori-dispatch

Conversation

@jhchouuu

@jhchouuu jhchouuu commented Aug 14, 2026

Copy link
Copy Markdown

Makes MegaMoEGfx1250's 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. 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:

field encoding, both sides
disp_out rows at slot * hidden
out_idx / out_wts [slot * topk + k], full topk row
flat dest map dest_pe * max_recv + slot, null npes * max_recv
recv_to_src_token src_pe * max_tok_per_rank + src_tok
recv_num / tok_off same signal/ack handshake (wait 0, store count+1, reader restores 0)

The 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 touches
cross_device_barrier (its cross-rank rendezvous is the recv_num
handshake), 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 new region, no extra device memory.

Two things that do differ, 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 on the host. This package's dispatch clears it in
    its 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, against
test_mega_moe.py's fp32 reference:

config FlyDSL mori logits_diff
2 layers, 256 tok/rank 945.9 us/layer 925.8 0.002174 (both)
4 layers, 4096 tok/rank 2678.0 us/layer 2622.2 0.004344 (both)

The diffs are identical because dispatch only moves data — the deviation from
the reference should not move, and it does not. --profile_table 1 confirms the
geometry that actually launches is mori's: ..._64x8 at 256 tokens,
..._64x16 at 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.

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.
@yanboshao
yanboshao merged commit fa5abb0 into ROCm:yanbo/mega_stage2_gfx1250 Aug 14, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants