Repository navigation
bench_mega_moe: fix kernel table on multi-stream graphs, add --fused_shared_experts - #7074
Merged
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
yanboshao
changed the title
jhchouuu
added a commit
to jhchouuu/aiter
that referenced
this pull request
Oct 28, 2026
…changes They are bench fixes rather than part of the MegaMoE recipe work, and move to ROCm#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>
zufayu
self-requested a review
October 29, 2026 01:04
JiaoliangYu
reviewed
Oct 29, 2026
NOT READY FOR HUMAN REVIEWPR Self-ReviewRevision Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision. Questions to answer in the PR description
Edit the Missing or invalid answers
Shadow mode reports what would block but keeps the status successful. Set |
JiaoliangYu
approved these changes
Oct 30, 2026
jhchouuu
added a commit
to jhchouuu/aiter
that referenced
this pull request
Oct 30, 2026
- bench_mega_moe: keep this branch's stage1/stage2 configurations and bring in ROCm#7074 on top: the kernel table on multi-stream graphs, --fused_shared_experts, and the split profile/time passes, now taken when any configuration runs stage1_fused. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
JiaoliangYu
pushed a commit
that referenced
this pull request
Oct 30, 2026
…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>
solinzby1
pushed a commit
that referenced
this pull request
Nov 2, 2026
…shared_experts (#7074) * bench_mega_moe: fix kernel table on multi-stream graphs * bench_mega_moe: --fused_shared_experts as ATOM DP+EP runs them * bench_mega_moe: profile stage1_fused on classic path, time on default
solinzby1
pushed a commit
that referenced
this pull request
Nov 2, 2026
…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>
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.
Motivation
--profile_table 1with--stage1_fused 1: re-exec each rank withDEBUG_HIP_GRAPH_CLASSIC_PATH=1so the kernel table names the right kernels, and warn when the profiler drops records.--fused_shared_experts Nreplaces--shared_experts: shared experts run inside the MoE as per-rank replicas, as ATOM DP+EP does (E + N*worldexperts,topk + Ncolumns), for both--combine baseandfused; the fp32 reference includes them.Technical Details
Test Plan
Test Result
Submission Checklist