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

atom/model_ops/fused_moe/mori_v2_prepare_finalize.py

@@ -17,9 +17,9 @@
 
 Two transports sit behind the same prepare/finalize pair:
 
-  * ATOM_MORI_V2_FUSED=0 -- mori's own v2 op-layer, combine_mode="gather". The
+  * ATOM_MEGA_STAGE2_FUSED=0 -- mori's own v2 op-layer, combine_mode="gather". The
     untouched upstream baseline.
-  * ATOM_MORI_V2_FUSED=1 -- aiter's MegaMoEGfx1250, whose gemm2 epilogue
+  * ATOM_MEGA_STAGE2_FUSED=1 -- aiter's MegaMoEGfx1250, whose gemm2 epilogue
     P2P-writes each weighted (token,k) result straight into the peers' combine
     staging, so combine only barriers + sums. It owns the whole layer
     (dispatch -> expert GEMM -> fused combine), so MoriV2ModularKernel hands it
@@ -41,7 +41,7 @@
 
 import torch
 import torch.distributed as dist
-from aiter import ActivationType, QuantType
+from aiter import ActivationType, QuantType, dtypes
 from aiter.dist.parallel_state import get_dp_group
 from aiter.ops.flydsl.moe_common import GateMode
 
@@ -132,10 +132,10 @@ def _import_v2() -> None:
 
 
 def _resolve_transport() -> str:
-    """ "mega" when ATOM_MORI_V2_FUSED is on, else mori's plain gather op-layer."""
+    """ "mega" when ATOM_MEGA_STAGE2_FUSED is on, else mori's plain gather op-layer."""
     from atom.utils import envs as _atom_envs
 
-    return "mega" if _atom_envs.ATOM_MORI_V2_FUSED else "gather"
+    return "mega" if _atom_envs.ATOM_MEGA_STAGE2_FUSED else "gather"
 
 
 @lru_cache(maxsize=1)
@@ -180,8 +180,9 @@ def _cco_per_rank_vmm(
     send all its tokens to one peer -> ws * M recv slots, plus a 2x headroom
     (tokens + combine buffers) and a fixed slack, matching test_moe_layer_ep.py.
 
-    MegaMoE's arena needs strictly less than this (one recv-sized token buffer
-    plus an M*topk combine staging), so the same budget covers both transports.
+    The gather transport's budget. MegaMoE sizes its own arena through
+    MegaMoEGfx1250.required_vmm_bytes(): with stage1_fused its dispatch rows
+    are one per route rather than one per token, which outgrows this.
     """
     tok_bytes = max_num_inp_token_per_rank * hidden_dim * itemsize
     win_bytes = ep_size * tok_bytes * 2 + (1 << 24)
@@ -225,6 +226,22 @@ def _cco_per_rank_vmm(
         "the old name is no longer read"
     )
 _MEGA_DISPATCH_WIRE = os.environ.get("MEGA_DISPATCH_WIRE", "bf16")
+# aiter's stage1_fused requires a quantizing wire and the flydsl dispatch.
+_MEGA_STAGE1_FUSED = envs.ATOM_MEGA_STAGE1_FUSED and _MEGA_DISPATCH_WIRE in (
+    "fp8",
+    "fp4",
+)
+
+
+def _check_mega_weights(layer: torch.nn.Module, quant_method: Any) -> None:
+    """Reject expert weights MegaMoEGfx1250 cannot run: it only takes MXFP4."""
+    w13, w2 = layer.w13_weight, layer.w2_weight
+    if not {w13.dtype, w2.dtype} <= {torch.uint8, dtypes.fp4x2}:
+        raise ValueError(
+            "MegaMoEGfx1250 requires MXFP4 w1/w2 weights (uint8 or fp4x2), got "
+            f"{w13.dtype} and {w2.dtype} from {type(quant_method).__name__}; set "
+            "ATOM_MEGA_STAGE2_FUSED=0 to use the gather transport"
+        )
 
 
 def init_mega_transport(
@@ -280,20 +297,32 @@ def init_mega_transport(
         # built (the staging layout itself no longer depends on it).
         _MEGA_DISPATCH_WIRE,
         combine_quant,
+        _MEGA_STAGE1_FUSED,
     )
     cached = _MEGA_TRANSPORTS.get(key)
     if cached is not None:
         return cached
 
     MegaMoEGfx1250 = _import_mega()
-    comm = _init_cco_comm(
-        ep_size,
-        ep_rank,
-        ep_src_global_rank,
-        _cco_per_rank_vmm(
+    if hasattr(MegaMoEGfx1250, "required_vmm_bytes"):
+        per_rank_vmm = MegaMoEGfx1250.required_vmm_bytes(
+            world_size=ep_size,
+            hidden_dim=hidden_dim,
+            max_tokens_per_rank=max_num_inp_token_per_rank,
+            experts=num_experts,
+            topk=num_experts_per_token,
+            stage1_fused=_MEGA_STAGE1_FUSED,
+        )
+    else:
+        logger.warning(
+            "[MegaMoEGfx1250] aiter has no required_vmm_bytes(); falling back to "
+            "the gather transport's VMM budget, which a stage1_fused arena can "
+            "outgrow"
+        )
+        per_rank_vmm = _cco_per_rank_vmm(
             ep_size, hidden_dim, max_num_inp_token_per_rank, data_type_itemsize
-        ),
-    )
+        )
+    comm = _init_cco_comm(ep_size, ep_rank, ep_src_global_rank, per_rank_vmm)
     mega = MegaMoEGfx1250(
         communicator=comm,
         rank=ep_rank,
@@ -314,15 +343,15 @@ def init_mega_transport(
         # Passed, not left to aiter's own read of the env, so the key and the
         # transport cannot drift.
         dispatch_wire=_MEGA_DISPATCH_WIRE,
-        # Only mori's dispatch carries the scale row, so a quantizing wire has
-        # no other backend to run on. Named here rather than left to
-        # $MEGA_DISPATCH, whose default is flydsl: otherwise asking for fp4 is
-        # rejected at the first MoE layer for a reason the operator did not set.
+        # A quantizing wire runs on mori's dispatch unless stage 1 is fused,
+        # which only the flydsl TDM dispatch implements. Named here rather than
+        # left to $MEGA_DISPATCH so the pairing cannot be misconfigured.
         **(
-            {"dispatch_backend": "mori"}
+            {"dispatch_backend": "flydsl" if _MEGA_STAGE1_FUSED else "mori"}
             if _MEGA_DISPATCH_WIRE in ("fp8", "fp4")
             else {}
         ),
+        stage1_fused=_MEGA_STAGE1_FUSED,
         # Only injected when asked for: an aiter without the combine-quant
         # epilogue has no such kwarg and would raise TypeError on every run.
         **({"combine_quant": combine_quant} if combine_quant != "none" else {}),
@@ -339,9 +368,10 @@ def init_mega_transport(
     comm.barrier()
     _MEGA_TRANSPORTS[key] = mega
     logger.info(
-        "[MORI-V2] Created MegaMoE: ep_rank=%d ep_size=%d hidden=%d inter=%d "
+        "[MegaMoEGfx1250] Created: ep_rank=%d ep_size=%d hidden=%d inter=%d "
         "experts=%d topk=%d M=%d act=%s gate=%s quant=%s pad=(%d,%d) "
-        "swiglu_limit=%s dispatch=%s wire=%s combine_quant=%s force_a8w4=%s",
+        "swiglu_limit=%s dispatch=%s wire=%s combine_quant=%s force_a8w4=%s "
+        "stage1_fused=%s per_rank_vmm=%.2fGiB",
         ep_rank,
         ep_size,
         hidden_dim,
@@ -360,6 +390,8 @@ def init_mega_transport(
         combine_quant,
         # The other half of the pair: logged together so a mismatch is readable.
         os.environ.get("AITER_FORCE_A8W4", "0"),
+        mega._config.stage1_fused,
+        per_rank_vmm / (1 << 30),
     )
     return mega
 
@@ -454,12 +486,14 @@ def bind_mega_transport(self, layer: torch.nn.Module, quant_method: Any) -> None
         """
         if self._mega_geometry is None or self.mega is not None:
             return
+        # Before init_mega_transport, which builds the cco communicator and
+        # allocates the arena.
+        _check_mega_weights(layer, quant_method)
         inter_dim = getattr(quant_method, "intermediate_size", 0)
         if inter_dim <= 0:
             raise ValueError(
                 "the fused transport needs the per-partition intermediate size, "
-                f"got {inter_dim}; ATOM_MORI_V2_FUSED=1 requires the a8w4 "
-                "(Mxfp4MoEMethod) quant path."
+                f"got {inter_dim} from {type(quant_method).__name__}"
             )
         self.mega = init_mega_transport(
             **self._mega_geometry,
@@ -705,6 +739,12 @@ def forward(
         # a4w4, mega bf16 + a8w4 all match to the last bit). Off by default
         # until it has run a real serve.
         if mega is not None and triton_experts is not None:
+            if mega._config.stage1_fused:
+                raise RuntimeError(
+                    "triton_mega_moe drives MegaMoE's token-major dispatch and "
+                    "cannot run on the compact stage-1 plan; set "
+                    "ATOM_MEGA_STAGE1_FUSED=0 to fall back to mori dispatch"
+                )
             from atom.model_ops.fused_moe_triton import triton_mega_moe
 
             assert not kwargs.get(

atom/utils/envs.py

@@ -254,18 +254,23 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float:
     # the gfx1250-capable path. Only takes effect when the mori all2all path is
     # active (dp_size>1 + expert-parallel + mori installed).
     "ATOM_MORI_V2": lambda: os.getenv("ATOM_MORI_V2", "0") == "1",
-    # gemm2-fused EP combine: the a8w4 grouped gemm2 epilogue P2P-writes its
-    # weighted per-(token,k) results straight into the peers' combine staging, so
-    # combine only barriers + sums. Requires the a8w4 (fp8 act + mxfp4 weight)
+    # MegaMoE stage 2 (gemm2-fused EP combine): the grouped gemm2 epilogue
+    # P2P-writes its weighted per-(token,k) results straight into the peers'
+    # combine staging, so combine only barriers + sums. Requires the mxfp4 weight
     # path; ignored unless ATOM_MORI_V2 is on.
     # This also selects the transport: 1 binds aiter's MegaMoEGfx1250, which owns
-    # the fused dispatch/combine pair (its dispatch kernel is picked by aiter's
-    # own MEGA_DISPATCH=flydsl|mori), 0 binds mori's v2 op-layer running plain
-    # gather, i.e. the untouched upstream baseline.
-    "ATOM_MORI_V2_FUSED": lambda: os.getenv("ATOM_MORI_V2_FUSED", "0") == "1",
+    # the fused dispatch/combine pair, 0 binds mori's v2 op-layer running plain
+    # gather, i.e. the untouched upstream baseline. Formerly ATOM_MORI_V2_FUSED.
+    "ATOM_MEGA_STAGE2_FUSED": lambda: (os.getenv("ATOM_MEGA_STAGE2_FUSED", "1") == "1"),
     # MegaMoE combine (return-trip) wire: bf16 | fp8 | fp4. Prefill-only; decode
-    # always combines in bf16. Ignored unless ATOM_MORI_V2_FUSED is on.
+    # always combines in bf16. Ignored unless ATOM_MEGA_STAGE2_FUSED is on.
     "ATOM_MEGA_COMBINE_WIRE": lambda: os.getenv("ATOM_MEGA_COMBINE_WIRE", "bf16"),
+    # MegaMoE compact-plan stage 1 (aiter stage1_fused): routing/layout planning
+    # fused with a flydsl TDM dispatch that writes the grouped GEMM's per-expert
+    # rows directly. Needs ATOM_MEGA_STAGE2_FUSED and an fp8/fp4
+    # MEGA_DISPATCH_WIRE; a bf16 wire keeps the token-major path regardless of
+    # this flag.
+    "ATOM_MEGA_STAGE1_FUSED": lambda: (os.getenv("ATOM_MEGA_STAGE1_FUSED", "1") == "1"),
     # Reuse a small MegaMoEV2 instance for native DP-unified small decode/
     # verify/draft forwards on the supported EP8, 48-experts-per-rank layout. Set to 0
     # to keep the configured max_num_batched_tokens capacity for every graph.

docs/environment_variables.md

@@ -150,7 +150,9 @@ combine knob as a quality/throughput tradeoff.
 | Variable | Type | Default | Description |
 |----------|------|---------|-------------|
 | **ATOM_MORI_FP4_DISPATCH** | bool | 0 (false) | If set to `1`, quantize activations to packed FP4 (E2M1, `per_1x32`) before the MoE all2all instead of sending bf16 — a quarter of the bytes on the dispatch wire — which selects `EpDispatchIntraNodeKernel_fp4`. MoRI picks its dispatch kernel from the dtype of the tensor handed to `dispatch()` but sizes its staging buffers from the config built at init, so this also switches `scale_dim` to `hidden_dim/32` and the scale type to e8m0. All three are resolved together by `mori_prepare_finalize.resolve_mori_dispatch()`; never set one without the others, as a mismatch strides the staging scale buffer wrong and faults on the first real batch instead of erroring cleanly. |
-| **ATOM_MEGA_COMBINE_WIRE** | str | `bf16` | MegaMoE (`ATOM_MORI_V2_FUSED=1`) combine wire: `bf16`, `fp8` (mxfp8) or `fp4` (mxfp4). Prefill-only: decode steps always combine in bf16, and the choice is DP-agreed so every rank reduces in the same format. |
+| **ATOM_MEGA_STAGE2_FUSED** | bool | 1 (true) | Under `ATOM_MORI_V2=1`, run the MoE layer on aiter's `MegaMoEGfx1250` (MegaMoE): the gemm2 epilogue P2P-writes each weighted route into the peers' combine staging, so combine only barriers and sums. `0` falls back to mori's v2 op-layer with plain gather combine. Formerly `ATOM_MORI_V2_FUSED`. |
+| **ATOM_MEGA_STAGE1_FUSED** | bool | 1 (true) | MegaMoE compact stage 1 (aiter `stage1_fused`): routing/layout planning fused with the flydsl TDM dispatch, which writes the grouped GEMM's per-expert rows directly. Only takes effect with `ATOM_MEGA_STAGE2_FUSED=1` and `MEGA_DISPATCH_WIRE=fp8`/`fp4`; a bf16 wire keeps the token-major path. Not supported together with the Triton experts. |
+| **ATOM_MEGA_COMBINE_WIRE** | str | `bf16` | MegaMoE (`ATOM_MEGA_STAGE2_FUSED=1`) combine wire: `bf16`, `fp8` (mxfp8) or `fp4` (mxfp4). Prefill-only: decode steps always combine in bf16, and the choice is DP-agreed so every rank reduces in the same format. |
 | **ATOM_MEGA_DECODE_MTPR** | int | 128 | Capacity (rows per rank) of the small MegaMoEV2 instance that DP-unified decode / verify / draft passes reuse (`ATOM_MEGA_DECODE_FAST_PATH=1`); one of 128, 256, 512, 1024. A pass with more rows falls back to the instance sized for `max_num_batched_tokens`. Size it to cover the DP-padded decode graphs (`running_bs * (num_speculative_tokens + 1)` rows); above 128 the instance stays on aiter's fixed-slot path only if `AITER_MEGA_FIXED_SLOT_MAX_MTPR` admits it (511 / 1023 / 2047), otherwise it runs compact. |
 | **ATOM_MEGA_MASK_PAD_ROWS** | bool | 0 (false) | If set to `1`, the MegaMoE backend (`--moe-backend mega`) routes this rank's DP pad rows (past `scheduled_tokens` on a step padded to the DP group's `running_tokens`) to expert `-1`: they are not dispatched or computed and come back as zeros, so identical pad rows also stop piling onto the same experts. Covers target and draft passes, eager and CUDA graph. With an aiter whose MegaMoEV2 masks `-1` slots in combine (`supports_combine_mask`) the op returns the zeros; otherwise each MoE layer pays one extra select. Not used by plugin frontends. |
 | **ATOM_MORI_COMBINE_QUANT** | str | `none` | Combine-side codec passed into the MoRI config. `none` returns bf16; `fp8_blockwise` selects `EpCombineIntraNodeKernel_*_fp8bwq_*`; MoRI also accepts `fp8_direct_cast`. |

recipes/DeepSeek-V4.md

@@ -95,6 +95,37 @@ Notes:
   its private expert-weight layout to the rebalancer via `get_eplb_weight_views`,
   so expert migration operates on the live fused weights.
 
+### MegaMoE on gfx1250 (mori v2 transport, EP + DP attention)
+
+On gfx1250 the EP + DP MoE all2all runs on mori's `dispatch_combine_v2`
+(`ATOM_MORI_V2=1`), and by default each routed-MoE layer is handed to aiter's
+`MegaMoEGfx1250`, which owns dispatch, both grouped GEMMs and combine. Two
+switches control how much of that layer is fused; both default to on:
+
+| Variable | Default | Effect |
+|----------|---------|--------|
+| `ATOM_MEGA_STAGE2_FUSED` | `1` | The gemm2 epilogue P2P-writes each weighted route into the peers' combine staging, so combine only barriers and sums. `0` falls back to mori's plain gather combine. Formerly `ATOM_MORI_V2_FUSED`. |
+| `ATOM_MEGA_STAGE1_FUSED` | `1` | Compact stage 1: routing/layout planning fused with the flydsl TDM dispatch, which writes the grouped GEMM1's per-expert rows directly. Needs stage 2 fused and a quantizing `MEGA_DISPATCH_WIRE` (`fp4`, or `fp8` with `AITER_FORCE_A8W4=1`); on a `bf16` wire it is ignored. When on, dispatch runs on flydsl TDM rather than mori. |
+
+```bash
+ATOM_MORI_V2=1 MORI_GPU_ARCHS=gfx1250 FLYDSL_GPU_ARCH=gfx1250 \
+MEGA_DISPATCH_WIRE=fp4 ATOM_MEGA_COMBINE_WIRE=fp4 \
+ATOM_MOE_GU_ITLV=1 ATOM_DP_LM_HEAD_MODE=allgather \
+python -m atom.entrypoints.openai_server \
+  --model deepseek-ai/DeepSeek-V4-Pro \
+  --kv_cache_dtype fp8 -tp 4 \
+  --enable-expert-parallel --enable-dp-attention
+```
+
+Notes:
+- The log line `[MegaMoEGfx1250] Created: ... dispatch=flydsl wire=fp4 ...
+  stage1_fused=True` confirms the compact stage-1 path is active.
+- Compact stage 1 cannot run with the Triton experts (`ATOM_USE_TRITON_MOE*`);
+  set `ATOM_MEGA_STAGE1_FUSED=0` to fall back to mori dispatch, which they can
+  run on.
+- `ATOM_DP_LM_HEAD_MODE=allgather` avoids RCCL `all_to_all_single`, which
+  cannot be captured into a hipGraph on gfx1250.
+
 ### FP8 on MI308 / gfx942 (V4-Flash-Base, FP8 per-block routed experts)
 
 [DeepSeek-V4-Flash-Base](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Base) ships the same V4 architecture (mHC + CSA + HCA + sparse attn + MTP) as V4-Pro, but **routed experts are FP8 e4m3 per-block 128×128** (instead of V4-Pro's FP4 e2m1 microscaling). This trades a small expert-memory increase for end-to-end ROCm `gfx942` (MI308) compatibility — `aiter`'s FP8 grouped GEMM has been tuned for `gfx942`, while the FP4 path was authored for `gfx950` (MI355X).