feat(moe): enable mega stage1_fused on fp4/fp8 dispatch wire by yanboshao · Pull Request #2384 · ROCm/ATOM

yanboshao · 2026-09-24T06:48:24Z

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

github-actions · 2026-09-24T06:48:48Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2384 --add-label <label>

JiaoliangYu

Found two initialization regressions: an undersized CCO VMM reservation for FP8 compact stage1, and default Mega selection for unsupported FP8/MXFP8 weight methods. Details and CPU-only validation are in the inline comments.

…stage1/stage2 fused switches) Source: #2384 Author: yanboshao (OPEN). Image has PR's first commit (527e581 upstream); 4b3faf2 (required_vmm_bytes sizing) is NOT included. Original commit in image (SHA differs from GitHub): c08275f feat(moe): default-on mega stage1/stage2 fused switches Co-Authored-By: Claude <noreply@anthropic.com>

…20261012 Kimi-K3 dp8ep8 now runs on rocm/fw-bringup-mlperf:gfx1250-atom-20261012 (ROCm 10.1 TheRock wheels, torch 2.11, editable ATOM/aiter/mori in /app), validated on CT15 (node0) + CT14 on 2026-10-13: full GSM8K strict-match 0.9659 +-0.0050, flexible-extract 0.9651, gate pass, ~8 min cold start. - Container and patches: the new image, its private repository, stack and source branches; the DeepSeek-R1 ENV it carries. - Two ATOM patches next to the page (Kimi-K3-455-wideEP-patches/): atom-r1_1012-k3.diff = #2380 (one hunk moved) + #2532 + #2533; atom-r1_1012-mega-vmm.diff = #2522's compact-plan cco arena sizing, needed because this aiter has no MegaMoEGfx1250.required_vmm_bytes() (otherwise ccoMemAlloc failed: -1 at MoE init). #2522 itself is superseded by upstream #2384 (ATOM_MEGA_STAGE1_FUSED); aiter #5431 and #5844 are in the image. - Launch: no venv; MORI_AMD_SMI_PATH (mori dlopens libamd_smi.so, TheRock ships only .so.27, ranks exit 255); ATOM_MLA_QH128_ASM_PAGE_SIZE=0; ATOM_MEGA_STAGE1_FUSED=1; DSR1 decode fast path and pad-row mask off. - Results and known issues updated for this image; AgentX not re-measured. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…k3-20261013; #2522 r3 validated - Image: rocm/fw-bringup:gfx1250-atom-k3-20261013 (sha256:9737256ddbd9860c99fc4f2ff0278d4cd210091ed07adb406e33c329d8362c08), which is rocm/fw-bringup-mlperf:gfx1250-atom-20261012 with the two ATOM patches applied, the launch scripts in /opt/k3 and MORI_AMD_SMI_PATH as ENV default. Full GSM8K on it: strict-match 0.9674 +-0.0049 (CT15 + CT14). Launch with /opt/k3/run_server.sh; the patch files stay for starting from the base image. - ATOM #2522 revision 3 (9cd096a) validated on ATOM main of 2026-10-08 with aiter e965ebfe: picks flydsl_compact for Kimi-K3 with nothing set, GSM8K 0.9621, 14k/500 con32 279.2 tok/s vs 272.2 with MEGA_DISPATCH=mori (+2.6%), median TPOT 94.4 vs 97.2 ms, median TTFT 9.0 vs 8.6 s. This image carries #2384 instead. - ATOM_MEGA_MASK_PAD_ROWS=1 with the compact plan on K3: GSM8K 0.9651, no slowdown. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…355X reference points - Five of the image's eight PRs have merged since (#2532, #2533, #2466, aiter #6121) and #2522 was closed in favour of ATOM #2384; the table gains an upstream-state column. aiter #6438 (K3 grouped-MoE tuning, merged 10-15) is noted as not in the image. - Results: ATOM on MI355X from the InferenceX Kimi-K3 dashboard (2026-10-15 run), concurrency 16, 48 and 56: P90 interactivity (1 / p90 ITL) and throughput per GPU, with the setup differences (TP 8, MTP on). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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>

…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>

yanboshao force-pushed the yanbo/stage1 branch 4 times, most recently from 3967424 to 527e581 Compare September 24, 2026 08:33

JiaoliangYu previously approved these changes Sep 24, 2026

View reviewed changes

JiaoliangYu added the ci:atom label Sep 24, 2026

zufayu requested a review from JiaoliangYu September 26, 2026 00:18

zufayu assigned JiaoliangYu Sep 26, 2026

kvnloo mentioned this pull request Sep 27, 2026

[Cross-feature wave 01] Invariant-driven reliability sweep across ATOM kvnloo/ATOM#9

Open

3 of 32 tasks

JiaoliangYu reviewed Sep 28, 2026

View reviewed changes

Comment thread atom/model_ops/fused_moe/mori_v2_prepare_finalize.py Resolved

Comment thread atom/utils/envs.py Resolved

yanboshao dismissed JiaoliangYu’s stale review via 4b3faf2 September 29, 2026 08:11

JiaoliangYu reviewed Oct 8, 2026

View reviewed changes

Comment thread atom/model_ops/fused_moe/mori_v2_prepare_finalize.py Resolved

yanboshao added 4 commits October 8, 2026 03:07

feat(moe): default-on mega stage1/stage2 fused switches

5b67b3d

fix(moe): size mega cco vmm via aiter, reject non-mxfp4 weights

168ab69

fix(moe): correct stage1_fused fallback hint

6a882b1

docs(dsv4): sync mega stage1 notes with code

1add19b

yanboshao force-pushed the yanbo/stage1 branch from 965081f to 1add19b Compare October 8, 2026 03:09

jhchouuu mentioned this pull request Oct 12, 2026

[Triton/Gluon] [FlyDSL] [CI] [MoE] MegaMoE gfx1250: stage1/stage2 switches, mori combine path, current mori EPv2 ROCm/aiter#6430

Merged

JiaoliangYu approved these changes Oct 14, 2026

View reviewed changes

JiaoliangYu merged commit 4010dca into main Oct 15, 2026

JiaoliangYu deleted the yanbo/stage1 branch October 15, 2026 06:21

zejunchen-zejun mentioned this pull request Oct 15, 2026

[MegaMoE] FlyDSL dispatch on fp8/fp4 wires; Kimi-K3 defaults to the compact plan #2522

Closed

Navigation Menu

Repository navigation

feat(moe): enable mega stage1_fused on fp4/fp8 dispatch wire - #2384

feat(moe): enable mega stage1_fused on fp4/fp8 dispatch wire#2384
JiaoliangYu merged 4 commits into
mainfrom
yanbo/stage1

yanboshao commented Sep 24, 2026

Uh oh!

github-actions Bot commented Sep 24, 2026

Uh oh!

JiaoliangYu left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

yanboshao commented Sep 24, 2026

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Sep 24, 2026

🏷️ CI Guide

Uh oh!

JiaoliangYu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants