Perf/ep disp 1250 ship by zhangfei829 · Pull Request #599 · ROCm/mori

zhangfei829 · 2026-08-24T14:14:37Z

Summary

Six gfx1250 dispatch optimizations that were behind -D switches while being measured are now
unconditional, plus a host-side cache for the launch argument struct.

Results vs `origin/main`

f01-2, EP4 intranode, ct=512 topk=6, MODES=graph COMBINE_IN=inplace, WARMUP=20 ITERS=100,
no -D flags on either arm, 20 interleaved rounds.

dispatch	median	range	one-sided
`origin/main`	41.35us	40.9 - 42.0
this branch	37.60us	37.4 - 37.9
delta	-3.80us (-9.2%)	-4.20 ... -3.20	20/20

Run-to-run spread on this machine is 1.1us, so the effect is 3.5x the resolution.

Per optimization

#	Optimization	Stage	Alone	One-sided
1	payload loop walks a `__ballot` mask instead of all `topk` slots	payload	-2.10us	8/8
2	`destPe`/`keep` kept in registers; null writes deferred to FINALIZE	COUNT + FINALIZE	-0.45us	5/6
3	meta prefetch hoisted above the `s_N` zeroing barrier	FINALIZE	-0.25us	4/6
4	`blkBase`/`blkCount` writes skipped on the tile path	reserve	-0.20us	5/8
5	`__match_any_sync` replaced by `npes` ballots; lane 0 posts the histogram	COUNT	-0.10us	4/6
6	launch argument struct reused across calls	host	host-only

The "Alone" column is not measured against origin/main. Each is an A/B with that one switch off,
on an earlier baseline. They sum to -3.10, which is not -3.80: the bench resolves 0.1us, so these are
two-tick effects and adding them compounds their error. #5 alone is a null result (range -1.10 to
+0.50); it is kept because five switches vs four differ by -0.25.

The four runtime preconditions (_dedupOk, _metapreOk, _bdOk, _blkMapNeeded) are correctness
guards, not tuning knobs, and each still has its fallback arm.

Correctness

CHECK=1 passes 8/8 on this exact code: ct=512,1024,2048,4096 x topk=6,8, both reporting
0 failed, 4/4 points verified. Both topk values are needed -- _etpi > 1 only occurs at topk=8,
so topk=6 alone never enters the fallback arms above.

Notes

~1.45us of the 37.60 is graph replay, not kernel time. The bench captures dispatch and combine as
two separate HIP Graphs, so one pair measurement replays two and pays the per-replay cost twice.
Confirmed by merging them into one graph: pair 69.2 -> 67.5us against a prediction of 67.75, 5/5.
The delta is unaffected (both arms pay it twice), but a real stack captures the layer as one graph,
so the kernel itself is nearer 36.1us.
[Feature]: Add CCL ops #6 is host-side and does not move the device numbers above. It matters in eager mode, where issuing
one dispatch or combine costs 28.6us of host time against a ~36us kernel.
Development scaffolding (timing switches, probe printfs, timestamp macros) is removed, so a default
build is the only build. Verified by disassembly rather than re-benchmarking: for topk=6 the
machine code is instruction-for-instruction identical to the file that produced the -3.80us.

Six changes that were behind -D switches while they were being measured are now just how the kernel works. Measured against origin/main, ct=512 topk=6 graph inplace EP4 on f01-2: dispatch 41.35 -> 37.60us, median -3.80us over 20 interleaved rounds, 20 of 20 one-sided. The run-to-run spread on that machine is 1.1us, so the effect is 3.5x the resolution and does not rest on a statistical argument. What each one does, and what it measured on its own: payload ballot walk -2.10us (8/8) the payload loop walked all topk slots and probed each with __shfl(validMe, l), so every warp paid 6 iterations no matter how many destinations it had. It now walks __ballot(validMe) and visits only the live ones. Same stores, same order. cached destPe/keep -0.45us (5/6) FINALIZE recomputed what COUNT already + deferred null write had; both are carried in registers now. The ~1374 scattered 4B null writes move out of COUNT to FINALIZE, which already branches on the same `keep`. meta prefetch -0.25us (4/6) the gather re-read tokenIndices/weightsBuf for values a lane already holds. The load is hoisted above the s_N zeroing barrier. skip blkBase/blkCount -0.20us (5/8) only the degenerate no-tile meta path reads them; the tile path reads LDS. blkBase sat between the cross-card fetch_add and the barrier the whole block waits on. ballot dedup -0.10us (4/6) with _etpi == 1 the (token, destPe) match key degenerates to destPe, so npes ballots reproduce __match_any_sync exactly, and lane 0 posts every histogram increment. Do not add these up: the five singles come to -3.10, which is neither the -3.80 the set measures together nor a coincidence worth trusting. The bench resolves 0.1us, so a 0.2us effect is two ticks and five weak numbers compound their error. The runtime preconditions stay: _dedupOk / _metapreOk / _bdOk / _blkMapNeeded are correctness guards, not tuning knobs, and every one of them still has its fallback arm. CHECK=1 passes on ct=512/topk6, 4096/topk8, 8/topk6 and 1/topk6 -- more than one shape is needed because _etpi > 1 only occurs at topk=8 and only ct=1 leaves warps without a token. The measurement scaffolding is not part of this change: the timing switches, the probe printfs and the timestamp macros they fed are all gone, so a default build is the only build. Machine code was compared before and after removing them, for topk=6 ct=512 and topk=8 ct=4096, and the fast-path geometry is instruction for instruction identical -- which is why the number above still applies to this file. Comments outside the license header are gone from this file. Co-authored-by: Cursor <cursoragent@cursor.com>

A serving loop calls launch() with the same argument NAMES on every iteration, so the ctypes struct was refilled from scratch each time for the sake of a few values that actually differ. On f01-2 that cost 10-20us per launch against a ~36us kernel, i.e. at small token counts the host, not the GPU, set the pace. The struct is cached and keyed on the tuple of argument names, so a call passing a different set cannot inherit fields from the previous shape. Only what can change is re-read: the keyword arguments themselves, and any pinned default holding a tensor, since its storage may have been reallocated. A pinned default that is an int is cached, because bind() is the only way to change one and it drops the cache. Co-authored-by: Cursor <cursoragent@cursor.com>

#599 reworked the gfx1250 dispatch (six perf changes made unconditional, -233 lines) and touched plan_api's argument binding, which is where both halves of this branch live. Five conflicts in ep_intranode_1250x.hpp, all in the same two categories: * the scale transport itself -- the staging pool, the 4th metadata field, the scale row in perTokM -- kept, since main cannot have it. #599's own edit to perTokM only removed a "(no scale)" comment, and its tokCap constant 384 is 128*3, which this branch already generalised to 128*kMetaFields. * comments #599 deleted that this branch still carried only by not touching those lines. Taken from main. plan_api merged cleanly: the optional-region binding and the new args-struct cache are in different parts of the same loop.

…ormat's reach The TDM header pulls in driver_types.h, which uses hipMemoryType without declaring it, so it only compiles after hip_runtime.h. It used to sit in its own block, which is what kept it there -- #599 removed the comment that separated them, leaving all three includes in one sortable group, and the next clang-format run put it first alphabetically. Every gfx1250 JIT compile then failed with four "unknown type name 'hipMemoryType'" in a ROCm header. Nothing on main is broken today; it is one formatter run away from being broken, which is why this is a blank line and a reason rather than a reordering. Found because a JIT cache wipe turned it into a compile failure -- with a warm cache the kernels never get recompiled and the tree looks fine. Verified from a cold cache: asym fp8 and fp4 with a 224 B scale row, backend parity, and test_op bf16, all on 4x gfx1250.

MegaMoE's mori dispatch launcher zeroed total_recv before every dispatch because mori's dispatch used to accumulate into it. Since ROCm/mori#599 the gfx1250 dispatch stores the count (`*args.totalRecvTokenNum = myRecv` in ep_intranode_1250x.hpp), so the reset is a dead kernel per MoE layer. It is now skipped when the source mori JIT-compiles shows the store and no accumulate; anything else keeps it, as does a caller that asks for it with reset_recv.

mori's gfx1250 dispatch now reserves its slots once per destination (the aggregated reservation, ROCm/mori#745) and hands the push combine its landing offsets through the dispatch. MegaMoE drives the plans directly, so it has to size and pass what mori's op layer does: - dest_pe_token_counter holds 2 * world_size ints: send counts, then per destination the reservation base and the push combine's landing offset. It held world_size; the dispatch wrote the second half into allocator padding. - recv_num signals are 8 bytes a source. The region held 4 bytes a source; the arena's 256-byte alignment hid the overrun. - The combine reads its landing offsets from dest_pe_token_counter, which MegaMoE did not pass, so the mxfp4 combine (stage2_fused=False) got a null pointer there. - MORI_EP_TOKOFF_EXT=1 turns the aggregated reservation off, so a mori push combine is refused with it, as mori's op layer refuses it. - The push landing rows come from mori's push_land_rows(): one per (token, destination rank), max_tokens * min(world_size, topk) rather than world_size * max_recv. - The mori dispatch launcher drops total_recv.zero_(): gfx1250's dispatch overwrites the count (since ROCm/mori#599), and mori's op layer skips the same zero there. Against a mori from before ROCm#745 the landing rows keep the old count and the two larger buffers have unused tails. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…tches, mori combine path, current mori EPv2 (#6430) * MegaMoE: stage1_fused/stage2_fused switches, mori combine path * bench_mega_moe: fix kernel table on multi-stream graphs * MegaMoE: follow mori tokoff-ext and selfFirst defaults * bench_mega_moe: --fused_shared_experts as ATOM DP+EP runs them * MegaMoE: follow mori#726 and size mori's fp4 combine flags Four fixes to the mori dispatch/combine path (stage1_fused/stage2_fused off), each from driving mori's plans without its op layer: - dispatch: since ROCm/mori#726 the plan takes tok_off_ext and derives selfFirst itself; self_first_enabled() is gone, so importing it fails. - tuning lookup: key it on the dispatch wire the way mori's op layer does ("fp4_disp_bf16_comb" for fp4). Under the hard-coded "bf16" key topk 9 matched no entry and both dispatch and the fp4 combine fell back to the single-shot default. - comb_inp: add the overlapped fp4 combine's u64 flag per landing row (four ranks up). Its 24-warp geometries write them past the region, i.e. past the arena: a page fault at 16384 tokens a rank. - close(): synchronize first, so a launch still in flight does not read a window that has been freed. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: zero-copy hand-off from fused_moe to mori's combine With stage2_fused off, fused_moe's rows went into a fresh tensor and then into mori's staging rows (comb_inp): the bf16 combine copied them in its own kernel, the fp4 one through a host-issued copy whenever the recv bound was below max_recv -- every decode step under ATOM. Have fused_moe write them there directly. fused_moe(output=) already existed but the gfx1250 grouped path ignored it and copied at the end; thread it down to the gather-reduce. mori's bf16 combine then sees its input at the staging base and skips the copy, and the fp4 send reads the staging rows in place. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: follow mori#745/#746 on the mori dispatch/combine path mori's gfx1250 dispatch now reserves its slots once per destination (the aggregated reservation, ROCm/mori#745) and hands the push combine its landing offsets through the dispatch. MegaMoE drives the plans directly, so it has to size and pass what mori's op layer does: - dest_pe_token_counter holds 2 * world_size ints: send counts, then per destination the reservation base and the push combine's landing offset. It held world_size; the dispatch wrote the second half into allocator padding. - recv_num signals are 8 bytes a source. The region held 4 bytes a source; the arena's 256-byte alignment hid the overrun. - The combine reads its landing offsets from dest_pe_token_counter, which MegaMoE did not pass, so the mxfp4 combine (stage2_fused=False) got a null pointer there. - MORI_EP_TOKOFF_EXT=1 turns the aggregated reservation off, so a mori push combine is refused with it, as mori's op layer refuses it. - The push landing rows come from mori's push_land_rows(): one per (token, destination rank), max_tokens * min(world_size, topk) rather than world_size * max_recv. - The mori dispatch launcher drops total_recv.zero_(): gfx1250's dispatch overwrites the count (since ROCm/mori#599), and mori's op layer skips the same zero there. Against a mori from before #745 the landing rows keep the old count and the two larger buffers have unused tails. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: bind mori's dispatch staging buffer (ROCm/mori#703) mori#703 moves the gfx1250 dispatch's staging arrays out of each kernel module's .bss into one buffer the caller allocates and binds as staging_base; an unbound plan is refused on the host. MegaMoE drives EpDispatchPlan directly, so it allocates ep_staging_bytes() once and binds it to every dispatch plan (they never run at the same time), and drops it after the plans on close(). A mori without ep_staging_bytes keeps the .bss arrays and is left as is. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: keep accepting dispatch_backend stage1_fused now picks the dispatch, and the dispatch_backend argument and config field went away with that. Callers still pass it (ATOM main passes dispatch_backend="mori" on a quantizing wire; ROCm/ATOM#2384 passes "flydsl" or "mori" alongside stage1_fused) and log mega._config.dispatch_backend, so both would fail on this aiter. Accept the argument again and check it against stage1_fused: "mori" is stage1_fused=False, "flydsl" is the compact plan. A "flydsl" without stage1_fused asked for the FlyDSL token-major dispatch, which no longer exists, so it is refused rather than run on mori. MegaMoEConfig gets a read-only dispatch_backend derived from stage1_fused. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * ci: run the MegaMoE bench on mori's dispatch stage1_fused=false with stage2_fused=true is mori's fp4 dispatch with the fused combine (the bench's auto wire is fp4 for a4w4_mxfp4), the configuration DeepSeek-R1 serves with. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: gfx1250 behaviour tests ahead of the refactor test_mega_moe_gfx1250.py pins what a change to MegaMoEGfx1250's dispatch/combine wiring must keep, through the constructor and members its callers use today: - every dispatch/combine pairing matches the fp32 reference, and a captured graph replays bit for bit what eager produced; - the pairings MegaMoE rejects fail before the symmetric arena is allocated; - mori's combine reads fused_moe's rows in place (zero-copy staging); - close() releases the plans before the window, the window before its memory; - the members ATOM's Triton experts reach (_dispatch, _combine, _arena, _config). It spawns its own ranks and skips (exit 0) off gfx1250, below two GPUs or without mori, so CI's MULTIGPU job can run it on gfx950. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: build from a MegaMoERecipe and a MegaMoEGeometry A caller now says what to run once, as a recipe, and MegaMoE builds it: recipe = MegaMoERecipe(dispatch=DispatchKind.MORI, combine=CombineKind.FUSED_SCATTER, dispatch_wire="fp4", combine_wires=("bf16", "fp4"), activation=..., quant_type=...) geometry = MegaMoEGeometry(rank, world_size, model_dim, inter_dim, experts, topk, max_tokens_per_rank) mega = MegaMoEGfx1250.from_recipe(recipe, geometry, communicator=comm) - MegaMoERecipe holds the transport (dispatch and combine kind, dispatch wire, the combine wires to build) and the expert-GEMM recipe; frozen and hashable, so a caller can key its cache of instances on it. MegaMoEGeometry holds the model and EP shape. - MegaMoERecipe.from_legacy is the one reader of stage1_fused, stage2_fused, dispatch_backend and the $MEGA_DISPATCH_WIRE / $MEGA_COMBINE_QUANT defaults. The keyword constructor goes through it, so both entries build the same thing. - MegaMoEConfig is now that recipe resolved against a geometry: frozen, with the same properties as before (stage1_fused, stage2_fused and dispatch_backend included, for the callers that read them). The dispatch schedule it used to have written into it lives on the instance; the dispatch_block_num / dispatch_warp_num_per_block fallback it carried was never reached. - forward() takes combine_wire= (bf16 | fp8 | fp4), the recipe's vocabulary, alongside combine_quant=. Every check and its message is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: probe mori once, and decide everything before the arena exists _MoriCapabilities gathers what MegaMoE needs to know about the installed mori -- the plan classes, the tuning lookup, and each helper a mori PR added (scale_stride_bytes #593, push_wire_nbytes #719, push_land_rows #745, ep_staging_bytes #703, TokOffExt.wanted #726, self_first_enabled) -- into one cached probe, and the decisions that hang off them into its methods: the tuned schedules, the landing-row sizes, the staging size, where the slot word lives. The ImportError/hasattr branches are gone from the pipeline; the RuntimeError messages for an mori that is too old are unchanged. _plan_transport runs all of it, plus the fused reduce's tiles and mori combine's block/epoch-slot check, before the symmetric arena is allocated. A configuration that cannot run -- mori missing, MORI_EP_TOKOFF_EXT=1 with mori's fp4 combine, no reduce tile, too many combine blocks -- now fails with nothing allocated instead of leaking the arena. TokOffExt.wanted, collective when MORI_EP_TOKOFF_EXT is on, moves ahead of the arena's own collective allocation on every rank alike. The compact dispatch with the fused combine still never loads mori's ops. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: one arena layout, and required_vmm_bytes from it The arena's regions now come from one formula, _arena_layout: the dispatch says what it writes (_dispatch_regions: the routing both kinds share, then its own rows), the combine what it stages through (_combine_regions), and _ArenaLayout.merge lays them down in the order the pipeline always used, so every pairing keeps the offsets it had. SymmetricArena takes the layout (or a region list, as before) and frees its memory if the window registration fails. MegaMoEGfx1250.required_vmm_bytes(recipe, geometry) reports the same layout, rounded the way ccoMemAlloc rounds an allocation -- the per_rank_vmm a caller needs for this MegaMoE, before any communicator exists, and exactly what the constructor then allocates. It is MegaMoE's share only; a communicator that also carries mori's own EP ops needs theirs on top. required_vmm_bytes(world_size=, hidden_dim=, max_tokens_per_rank=, experts=, topk=, stage1_fused=), the keyword form #6525 proposed and ATOM's yanbo/mega already calls, still works: it names neither wires nor combine nor inter_dim, so it returns the largest any recipe with that dispatch takes. _compact_gemm_align_m reads activation, quant type and inter_dim off the config instead of taking them separately. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: dispatch and combine stages under one dispatcher; linear forward The transport is now two stages the recipe picks, over one shared arena: _MoriDispatch / _CompactDispatch tokens out to their experts' ranks _MoriCombine / _FusedScatterCombine results back, summed per token Each declares its arena regions (regions(), pure -- what _arena_layout and required_vmm_bytes merge), plans in __init__ with nothing allocated, and allocates in build(). They meet only through _EpResources (the arena, the routing state both ends read, TokOffExt), the _DispatchBatch the dispatch hands the expert GEMM, and the _ComputeTarget the combine has it write into. _EpDispatcher picks the two from tables, runs every check before the arena is allocated, and frees what was built when a later step fails -- the arena and plans used to leak. close() goes: synchronize, combine plans, dispatch plans, staging, arena, TokOffExt. forward() is now: pick the step's combine wire, check the inputs, dispatch, prepare, compute, finish. The compact plan's machinery -- slots, side stream, per-step geometry (_CompactStep, now immutable) -- lives in _CompactDispatch, and the next layer's plan launches at the end of its dispatch, the same place on the GPU streams as before. The recv views are made once at build instead of per step. The TDM dispatch keeps only its compact arm; the token-major arms were unreachable once stage1_fused=False went through mori. Kept for the callers that reach past forward(), ATOM's Triton experts among them: _config, _arena, _dispatch() and its 5-tuple, _combine(), _recv_dispatch_scales(), prefetch_compact_plan(), warmup_compact_plan(). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: size mori's regions from describe(); the bench sizes cco from the recipe On a mori with Plan.describe() (ROCm/mori#844), the mori stages ask it for their arena regions and the dispatch's staging bytes -- the numbers mori's own kernels and its bind-time size check use -- instead of keeping formulas that mirror them. Each stage builds its plan request in one place (request()), which both describe() and the plan take. describe() compiles nothing and resolves against gfx1250 by name, so required_vmm_bytes stays off the GPU; the slot word's kwargs stay out of it, since where that word lives sizes nothing and asking would be collective. An older mori keeps the formulas. SymmetricArena.size() lets a #844 plan check its regions at bind. bench_mega_moe builds from a MegaMoERecipe and sizes its communicator with required_vmm_bytes(recipe, geometry) instead of a fixed 16 GiB. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: CPU tests for the recipe and the arena it sizes tests/test_mega_moe_gfx1250_recipe.py needs no GPU and skips where MegaMoE does not import. It checks how the legacy flags map to a recipe and that the rejections keep their wording, and lays the arena out over a grid of geometries and every pairing against the region list #6430 built inline -- through a fake mori both with and without Plan.describe(), the describe one following mori's rules, which add no overlapped-reduce flags on a wide EP. required_vmm_bytes is held to the layout as cco rounds it, and its keyword form to at least that. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: bind mori's dispatch staging as a tensor A mori with ROCm/mori#844 checks a staging tensor against the plan's stagingBytes, device and alignment at bind, and holds it; an address is taken on trust. An older mori reads the tensor's address at each launch, as it did the int. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * TEMP: grouped GEMM takes a strided scale view again (fixup for #6946) #6946 flattens the grouped GEMM's scale arguments with .view(-1), which a non-contiguous tensor cannot do. MegaMoE's compact rows hand over exactly that -- each row's e8m0 scales sit inside its wire row -- so every stage1_fused MegaMoE step raises. The kernel only takes the base address and walks the rows by a_scale_row_stride_bytes, so a strided scale view becomes a one-dimensional view over the same base, as it was before #6946. Temporary: drop this commit once main fixes it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * bench_mega_moe: take out the kernel-table and --fused_shared_experts changes They are bench fixes rather than part of the MegaMoE recipe work, and move to #7074. This reverts "bench_mega_moe: fix kernel table on multi-stream graphs" and "bench_mega_moe: --fused_shared_experts as ATOM DP+EP runs them". Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * FlyDSL grouped GEMM: pass a strided scale view at its base address Replaces the temporary workaround for the cached launcher (#6946), which flattens the grouped GEMM's scale arguments with .view(-1). A non-contiguous tensor cannot do that, and MegaMoE's expert-grouped dispatch hands over exactly one: each row's e8m0 scales sit inside its wire row, read by a_scale_row_stride_bytes. The kernel reads these arguments only through their base pointer. tensor_shim gains flat_arg(), beside ptr_arg(): a contiguous tensor is .view(-1), as before, so no other caller changes; a strided one becomes the run of elements from its first to its last, at the same base. tests/test_grouped_gemm_strided_scales.py checks flat_arg and what the launcher hands the compiled kernel, without a GPU. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: stage contracts, mori sizing from describe() only, from_recipe first The recipe - The expert GEMM's fields move into ExpertComputeConfig, nested as MegaMoERecipe.compute; the recipe's own fields are the transport's. - Which pairings run is declared by the stages: a dispatch names the wires it carries and the BatchLayout it delivers (token-major or expert-grouped), a combine the layouts it accepts and the wires it offers. The recipe checks against those, and its errors name recipe fields (dispatch=, combine=, dispatch_wire=, combine_wires=). Each stage's geometry check (check()) runs from MegaMoEConfig. - from_recipe is the entry point. The keyword constructor warns DeprecationWarning; the underscored members ATOM's Triton experts use are documented as kept for them. The facade reads geometry and compute fields through properties rather than copying them. The stages - DispatchStage and CombineStage (Protocols) say what each provides; every stage has regions(config, caps) and check(config), and is built as Stage(config, caps). The compact- and mori-only attributes are gone from the shared base. - Each dispatch stage owns the stride its e8m0 rows arrive at (scale_stride), replacing MegaMoEConfig.dispatch_scale_dst_nbytes. - The batch carries its layout explicitly, and an expert-grouped one an ExpertGrouping; the combine hands the grouped GEMM its dispatch context in the _ComputeTarget, so the expert GEMM does not branch on the layout. - Whether the combine needs mori's aggregated reservation and whether the dispatch makes it are each stage's to say; the dispatcher refuses the pair before the arena exists. - One schedule pick (_pick) serves both mori plans and the compact kernels; _quantize returns the scale rows' address instead of leaving it in state. - _EpDispatcher has dispatch/prepare/combine/prefetch/warmup and the recv views, so forward() does not reach into the stages, and close() releases the rest even when one part raises. mori - _MoriCapabilities needs Plan.describe() (ROCm/mori#844). Regions, staging bytes and the scale stride come from there, so the copies of mori's formulas and the older-mori branches are gone; such a mori also overwrites total_recv in its dispatch, which is why nothing zeroes it. - _XDB_FLAG_SLOTS, TokOffExt and the plan constructors are reached only through _MoriCapabilities. Checks - An out-of-range recv_token_bound (and next_recv_token_bound) raises on the compact dispatch as on mori's, in dispatch, prefetch and warmup. - a1_scale on a quantizing dispatch wire raises ValueError (was an assert). - The keyword form of required_vmm_bytes raises when no candidate recipe runs at the geometry, instead of returning 0. - _MoriCombine.prepare documents why writing fused_moe's rows straight into comb_inp is safe: the next dispatch does not finish until every peer has dispatched, which each does only after its own combine. Tests: the CPU test fakes mori from describe() alone and covers the stage contract messages, the bound checks and the keyword-bound refusal; the multi-GPU test builds through from_recipe and checks the deprecated keyword constructor against it. bench_moe_gemm_a4w4_cudagraph builds through from_recipe. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: split mega_moe.py into recipe, mori_caps, arena, stages, transport A move, no change in behaviour: each block of mega_moe.py lands in one module as it was, and mega_moe.py keeps only MegaMoEGfx1250. - recipe.py: the wires, DispatchKind / CombineKind, ExpertComputeConfig, MegaMoERecipe, MegaMoEGeometry, MegaMoEConfig. - mori_caps.py: _MoriCapabilities and mori's region names. - arena.py: SymmetricArena, the region order, _ArenaLayout, cco's rounding. - stages/base.py: DispatchStage / CombineStage and what the stages share (routing, resources, batch, compute target, schedule pick). - stages/dispatch_mori.py, dispatch_compact.py, combine_fused.py, combine_mori.py: one stage each; stages/__init__.py maps each kind to its stage. - transport.py: _EpDispatcher and the one arena layout. The only edit: recipe.py looks the stages up when asked for (the stages import recipe.py). mega_moe.py still exports the package's public names, and Routing, which a bench imports from there. The CPU test imports the private names it checks from the modules that define them. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: compact dispatch bound and plan fixes, wire-picked GEMM A dtype - Compact dispatch: with no recv_token_bound the step is sized for the arena's capacity; a plan started for another alignment is refused; the plan's segment and wave limits are checked before the arena exists; plans are compiled when created, ahead of any graph capture. - A quantizing dispatch wire sets fused_moe's quant_dtype_a. - bench_mega_moe runs on gfx1250 only. - Tests: fp8 dispatch and fp8 combine cases, rank-local failures raised on every rank, CPU tests skip only without flydsl. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE CPU test: skip without mori as well - The package's kernel modules import mori.cco, so the test skips when mori is not installed instead of failing collection. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: yanboshao <yashao@amd.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

zhangfei829 and others added 2 commits August 24, 2026 21:23

zhangfei829 merged commit 5236e43 into main Aug 24, 2026
4 of 5 checks passed

Navigation Menu

Repository navigation

Perf/ep disp 1250 ship - #599

Perf/ep disp 1250 ship#599
zhangfei829 merged 2 commits into
mainfrom
perf/ep-disp-1250-ship

zhangfei829 commented Aug 24, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

zhangfei829 commented Aug 24, 2026

Summary

Results vs origin/main

Per optimization

Correctness

Notes

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Results vs `origin/main`