[FlyDSL] [CI] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire

.github/scripts/aiter_test.sh

@@ -101,7 +101,7 @@ for file in "${sharded_files[@]}"; do
                     fi
                     exec env MORI_SHMEM_HEAP_SIZE=40G \
                         torchrun --standalone --nproc_per_node=8 "$test_file" \
-                        --combine scatter_fused --layers 2 --acc_verify 1
+                        --combine fused --layers 2 --acc_verify 1
                 '
                 _ "$file"
             )

aiter/aot/flydsl/grouped_moe.py

@@ -209,6 +209,11 @@ def _route_ksplit(feat_dim, source_topk, out_e, out_m):
                 1,
                 numel,
                 ptr_arg(torch.empty(0, dtype=i32, device=dev)),
+                # src_scale: read only by the prequantized build, which this job
+                # does not emit; the pointer still has to be passed. Missing it
+                # would not fail loudly -- compile_one_config swallows the
+                # TypeError and the routeks kernels just stop being precompiled.
+                ptr_arg(torch.empty(0, dtype=u8, device=dev)),
                 grid,
                 stream=0,
             )

aiter/fused_moe.py

@@ -854,6 +854,9 @@ def _fused_moe_impl(
                 doweight_stage1=doweight_stage1,
                 w1_scale=w1_scale,
                 w2_scale=w2_scale,
+                # A quantizing EP dispatch puts the caller's e8m0 row here;
+                # the gfx1250 path uses it to skip a quant it would redo.
+                a1_scale=a1_scale,
                 expert_mask=expert_mask,
                 hidden_pad=hidden_pad,
                 intermediate_pad=intermediate_pad,

aiter/ops/flydsl/grouped_moe_gfx1250.py

@@ -446,6 +446,7 @@ def _grouped_a8w4_tdm_moe(
     stage2_scatter: Stage2ScatterContext | None = None,
     situ_beta=1.0,
     situ_linear_beta=1.0,
+    a1_scale=None,
 ):
     import functools
 
@@ -603,8 +604,37 @@ def _grouped_a8w4_tdm_moe(
     _quant_mode = "fp4" if _is_fp4 else "fp8"
     _a_is_fp4 = 1 if _is_fp4 else 0
 
+    # Bound once, because the quant pass below rebinds a1_scale to the
+    # PRESHUFFLED GROUPED scale. Both are uint8 and both have a plausible
+    # shape, so nothing downstream could tell which one it was handed.
+    src_a1_scale = a1_scale
+
+    # Pre-quantized activation: an MX payload plus its e8m0 row is what a
+    # quantizing EP dispatch delivers, and it is also aiter's standing meaning
+    # for this pair.
+    _prequantized = src_a1_scale is not None and hidden_states.dtype in (
+        dtypes.fp8,
+        torch.uint8,
+        dtypes.fp4x2,
+    )
+    if src_a1_scale is not None and not _prequantized:
+        # Loud rather than silently re-quantizing something already quantized.
+        assert hidden_states.dtype == dtype, (
+            f"a1_scale given with hidden_states dtype {hidden_states.dtype}: "
+            "expected packed MX bytes (pre-quantized) or the model dtype (ignored)"
+        )
+    # Checked, not inferred: a wire disagreeing with the GEMM's data_format would
+    # read into the next token's bytes and produce plausible garbage.
+    _src_width = hidden_states.shape[-1]
+    if _prequantized:
+        _want_width = model_dim // 2 if _is_fp4 else model_dim
+        assert _src_width == _want_width, (
+            f"prequantized {_quant_mode} payload should be {_want_width} B per "
+            f"row at model_dim {model_dim}, got {_src_width}"
+        )
+
     a1_payload, a1_scale = flydsl_moe_fused_quant_preshuffle(
-        hidden_states.reshape(1, token_num, model_dim),
+        hidden_states.reshape(1, token_num, _src_width),
         1,
         contiguous_m,
         wmma_rep=wmma_rep,
@@ -613,6 +643,7 @@ def _grouped_a8w4_tdm_moe(
         topids_to_rows=topids_to_rows,
         source_topk=topk,
         num_valid_routes=_ep_nvr,
+        prequantized_scale=src_a1_scale if _prequantized else None,
     )
 
     # Fuse gemm1 activation + MX quantization + scale preshuffle into the
@@ -900,6 +931,7 @@ def grouped_gemm_gfx1250_a8w4(
     situ_beta: float = 1.0,
     situ_linear_beta: float = 1.0,
     stage2_scatter: Stage2ScatterContext | None = None,
+    a1_scale: torch.Tensor | None = None,
 ):
     """Grouped a8w4/a4w4 MoE on the TDM batched GEMM (gfx1250).
 
@@ -988,19 +1020,17 @@ def _fmt(v):
     ):
         _grouped_dbg("unsupported activation")
         return None
-    is_grouped_a4w4 = q_dtype_a == dtypes.fp4x2 and q_dtype_w == dtypes.fp4x2
-    is_grouped_a8w4 = q_dtype_a == dtypes.fp8 and (
-        q_dtype_w == dtypes.fp4x2 or w1.dtype == torch.uint8
-    )
+    # mxfp4 weights arrive as fp4x2 or as the uint8 view of the same bytes --
+    # ATOM's loader keeps them uint8, and MegaMoE accepts both. Requiring the
+    # packed dtype on the a4w4 arm alone silently routed a4w4-with-uint8-weights
+    # to the 2-stage fallback.
+    w_is_mxfp4 = q_dtype_w == dtypes.fp4x2 or w1.dtype == torch.uint8
+    is_grouped_a4w4 = q_dtype_a == dtypes.fp4x2 and w_is_mxfp4
+    is_grouped_a8w4 = q_dtype_a == dtypes.fp8 and w_is_mxfp4
     if not (is_grouped_a4w4 or is_grouped_a8w4):
         return None
     data_format = "fp4" if is_grouped_a4w4 else "a8w4"
-    # Normalize uint8-viewed fp4 weights back to fp4x2 for CSV key matching.
-    q_dtype_w_key = (
-        dtypes.fp4x2
-        if (q_dtype_w == dtypes.fp4x2 or w1.dtype == torch.uint8)
-        else q_dtype_w
-    )
+    q_dtype_w_key = dtypes.fp4x2 if w_is_mxfp4 else q_dtype_w
     _grouped_dbg(f"eligible data_format={data_format}")
     if w1_scale is None or w2_scale is None:
         return None
@@ -1169,6 +1199,7 @@ def _tdm_env(name):
             stage2_scatter=stage2_scatter,
             situ_beta=situ_beta,
             situ_linear_beta=situ_linear_beta,
+            a1_scale=a1_scale,
             **_tdm_kw,
         )
 

aiter/ops/flydsl/kernels/mega_moe_gfx1250/__init__.py

@@ -5,11 +5,11 @@
 
 import importlib
 
-__all__ = ["MegaMoEGfx1250"]
+__all__ = ["MegaMoEGfx1250", "read_dispatch_wire_env"]
 
 
 def __getattr__(name):
-    if name != "MegaMoEGfx1250":
+    if name not in __all__:
         raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
     return getattr(importlib.import_module(f"{__name__}.mega_moe"), name)