[FlyDSL] [CI] [feat] mega-moe stage2 for gfx1250

.github/scripts/aiter_test.sh

@@ -106,6 +106,28 @@ for file in "${sharded_files[@]}"; do
                 _ "$file"
             )
             ;;
+        op_tests/multigpu_tests/test_mega_moe_gfx1250.py)
+            {
+                echo "Running gfx1250 MegaMoE fused-scatter accuracy on 8 GPUs when supported"
+            } | tee -a latest_test.log
+            test_cmd=(
+                timeout 60m
+                bash -c '
+                    set -euo pipefail
+                    test_file=$1
+                    arch=$(python3 -c \
+                        "from aiter.jit.utils.chip_info import get_gfx; print(get_gfx())")
+                    if [[ "$arch" != "gfx1250" ]]; then
+                        echo "Skipping $test_file: requires gfx1250, got $arch"
+                        exit 0
+                    fi
+                    exec env MORI_SHMEM_HEAP_SIZE=40G \
+                        torchrun --standalone --nproc_per_node=8 "$test_file" \
+                        --combine scatter_fused --layers 2 --acc_verify 1
+                '
+                _ "$file"
+            )
+            ;;
         op_tests/multigpu_tests/test_mega_moe_v2.py)
             {
                 echo "Running MegaMoEV2 v4_pro fixed-slot and compact coverage on 8 GPUs"

aiter/aot/flydsl/grouped_moe.py

@@ -386,7 +386,7 @@ def compile_one_config(**job):
     t0 = time.time()
     try:
         # TODO(aot): only the auxiliary (non-GEMM) kernels are precompiled here.
-        # The grouped GEMM itself moved to the felix TDM batched kernel
+        # The grouped GEMM itself moved to the TDM batched kernel
         # (aiter.ops.flydsl.batched_gemm_mxfp4), which has no AOT wiring yet, so
         # it JIT-compiles on first use. Add TDM GEMM jobs here to restore the
         # coverage the deleted moe_grouped_gemm_mxscale path used to provide.

aiter/fused_moe.py

@@ -31,6 +31,7 @@
     gfx_from_cu_num,
 )
 from aiter.jit.utils.torch_guard import torch_compile_guard
+from aiter.ops.flydsl.kernels.mega_moe_gfx1250.types import Stage2ScatterContext
 
 try:
     from aiter.ops.flydsl.moe_common import GateMode
@@ -489,6 +490,7 @@ def fused_moe(
     shared_w1_scale: torch.Tensor | None = None,
     shared_w2_scale: torch.Tensor | None = None,
     shared_expert_id: int = -1,
+    stage2_scatter: Stage2ScatterContext | None = None,
 ):
     if (
         any(
@@ -531,6 +533,8 @@ def fused_moe(
         )
     if not block_size_M:
         block_size_M = -1
+    enable_ep_scatter = stage2_scatter is not None
+    scatter_source_map = stage2_scatter.source_token_map if enable_ep_scatter else None
     return fused_moe_(
         hidden_states=hidden_states,
         w1=w1,
@@ -557,6 +561,18 @@ def fused_moe(
         beta=beta,
         linear_beta=linear_beta,
         gate_mode=gate_mode,
+        ep_arena_handle=stage2_scatter.arena_handle if enable_ep_scatter else 0,
+        ep_combine_input_offset=(
+            stage2_scatter.combine_input_offset if enable_ep_scatter else 0
+        ),
+        ep_slot_stride_bytes=(
+            stage2_scatter.slot_stride_bytes if enable_ep_scatter else 0
+        ),
+        ep_max_tokens_per_rank=(
+            stage2_scatter.max_tokens_per_rank if enable_ep_scatter else 0
+        ),
+        ep_world_size=stage2_scatter.world_size if enable_ep_scatter else 0,
+        ep_source_token_map=scatter_source_map,
     )
 
 
@@ -588,6 +604,12 @@ def fused_moe_fake(
     beta: float | None = None,
     linear_beta: float | None = None,
     gate_mode: str = GateMode.SEPARATED.value,
+    ep_arena_handle: int = 0,
+    ep_combine_input_offset: int = 0,
+    ep_slot_stride_bytes: int = 0,
+    ep_max_tokens_per_rank: int = 0,
+    ep_world_size: int = 0,
+    ep_source_token_map: torch.Tensor | None = None,
 ) -> torch.Tensor:
     device = topk_ids.device
     M, _topk = topk_ids.shape
@@ -626,7 +648,23 @@ def fused_moe_(
     beta: float | None = None,
     linear_beta: float | None = None,
     gate_mode: str = GateMode.SEPARATED.value,
+    ep_arena_handle: int = 0,
+    ep_combine_input_offset: int = 0,
+    ep_slot_stride_bytes: int = 0,
+    ep_max_tokens_per_rank: int = 0,
+    ep_world_size: int = 0,
+    ep_source_token_map: torch.Tensor | None = None,
 ) -> torch.Tensor:
+    stage2_scatter = None
+    if ep_source_token_map is not None:
+        stage2_scatter = Stage2ScatterContext(
+            arena_handle=ep_arena_handle,
+            combine_input_offset=ep_combine_input_offset,
+            slot_stride_bytes=ep_slot_stride_bytes,
+            max_tokens_per_rank=ep_max_tokens_per_rank,
+            world_size=ep_world_size,
+            source_token_map=ep_source_token_map,
+        )
     return _fused_moe_impl(
         hidden_states=hidden_states,
         w1=w1,
@@ -653,6 +691,7 @@ def fused_moe_(
         beta=beta,
         linear_beta=linear_beta,
         gate_mode=gate_mode,
+        stage2_scatter=stage2_scatter,
     )
 
 
@@ -682,6 +721,7 @@ def _fused_moe_impl(
     beta: float | None = None,
     linear_beta: float | None = None,
     gate_mode: str = GateMode.SEPARATED.value,
+    stage2_scatter: Stage2ScatterContext | None = None,
     *,
     _q_dtype_a: torch.dtype | None = None,
     _metadata_transform: Callable | None = None,
@@ -804,6 +844,7 @@ def _fused_moe_impl(
                 num_local_tokens=num_local_tokens,
                 situ_beta=1.0 if beta is None else float(beta),
                 situ_linear_beta=1.0 if linear_beta is None else float(linear_beta),
+                stage2_scatter=stage2_scatter,
             )
 
     if grouped_a8w4_out is not None:

aiter/ops/flydsl/grouped_gemm_mxfp4.py

@@ -9,6 +9,7 @@
 
 import torch
 
+from .kernels.mega_moe_gfx1250.types import Stage2ScatterContext
 from .kernels.tensor_shim import ptr_arg
 
 _SUPPORTED_CLUSTER_N = (4, 3, 2)
@@ -90,6 +91,9 @@ def flydsl_grouped_gemm_a8w4_masked(
     cluster_n=-1,
     waves_per_tensor_tdm=-1,
     next_stage_prefetch=0,
+    stage2_scatter: Stage2ScatterContext | None = None,
+    ep_destination_stride=0,
+    ep_row_map=None,
     situ_beta=1.0,
     situ_linear_beta=1.0,
 ):
@@ -115,6 +119,8 @@ def flydsl_grouped_gemm_a8w4_masked(
             f"[grouped-moe tdm] cluster_n={cluster_n} needs n_tiles={n_tiles} "
             f"(N={N}, tile_n={tile_n}) to be an exact multiple"
         )
+    enable_ep_scatter = stage2_scatter is not None
+    ep_row_map_tensor = ep_row_map if ep_row_map is not None else out
     launch_gemm_a8w4_tdm(
         out,
         ptr_arg(a),
@@ -145,7 +151,18 @@ def flydsl_grouped_gemm_a8w4_masked(
         cluster_n,
         _select_next_stage_prefetch(next_stage_prefetch),
         waves_per_tensor_tdm,
-        float(situ_beta),
-        float(situ_linear_beta),
+        enable_ep_scatter=int(enable_ep_scatter),
+        ep_arena_handle=(int(stage2_scatter.arena_handle) if enable_ep_scatter else 0),
+        ep_combine_input_offset=(
+            int(stage2_scatter.combine_input_offset) if enable_ep_scatter else 0
+        ),
+        ep_slot_stride_bytes=(
+            int(stage2_scatter.slot_stride_bytes) if enable_ep_scatter else 0
+        ),
+        ep_destination_stride=int(ep_destination_stride),
+        ep_world_size=int(stage2_scatter.world_size) if enable_ep_scatter else 0,
+        arg_ep_row_map=ep_row_map_tensor,
+        f32_situ_beta=float(situ_beta),
+        f32_situ_linear_beta=float(situ_linear_beta),
     )
     return out