moe rebase and refactor

aiter/aot/flydsl/grouped_moe.py

@@ -37,49 +37,6 @@ def _align_max_m(raw_max_m: int, warp_tile_m: int) -> int:
     return max(int(warp_tile_m), _align_up(raw_max_m, warp_tile_m))
 
 
-def _preshuffled_scale_shape(
-    rows: int, k_dim: int, warp_tile: int, tile_k: int = _TILE_K
-) -> tuple[int, int]:
-    """Mirror moe_grouped_gemm_mxscale_gfx1250._preshuffled_scale_shape.
-
-    The grouped GEMM launchers validate an exact preshuffled E8M0 scale layout
-    (see tests.kernels.test_gemm_mxscale_gfx1250.preshuffle_e8m0_scale), so the
-    AOT dummy tensors must use the same shape, not the plain (rows, k//32) one.
-    """
-    k_scale = int(k_dim) // 32
-    scale_k_per_tile = int(tile_k) // 32
-    if k_scale % scale_k_per_tile != 0:
-        raise ValueError(
-            f"K scale columns must be divisible by tile_k/32, got {k_scale} and {scale_k_per_tile}"
-        )
-    wmma_rep = int(warp_tile) // 16
-    if wmma_rep < 1:
-        raise ValueError(f"warp_tile must be >= 16, got {warp_tile}")
-    if int(rows) % wmma_rep != 0:
-        raise ValueError(
-            f"scale rows must be divisible by wmma_rep={wmma_rep}, got {rows}"
-        )
-    return int(rows) // wmma_rep, k_scale * wmma_rep
-
-
-def _preshuffled_b_scale_shape(rows: int, k_dim: int) -> tuple[int, int]:
-    """Mirror moe_grouped_gemm_mxscale_gfx1250._preshuffled_b_scale_shape.
-
-    Weight (B) scale uses the n32k4 layout (different from the activation/A
-    layout above): a 32-row super-block folds into the column dim, so 32 N-rows
-    collapse to one row and each k_scale column expands x32. The grouped GEMM
-    launchers validate scale_w against THIS shape, so the AOT dummy must match.
-    """
-    k_scale = int(k_dim) // 32
-    if k_scale % 4 != 0:
-        raise ValueError(
-            f"B-scale k columns (K//32) must be divisible by 4 (K%128==0), got {k_scale}"
-        )
-    if int(rows) % 32 != 0:
-        raise ValueError(f"B-scale rows must be divisible by 32, got {rows}")
-    return int(rows) // 32, k_scale * 32
-
-
 def _as_bool(value, default: bool = False) -> bool:
     if value is None or str(value).strip() == "":
         return default
@@ -99,7 +56,7 @@ def _as_float(value, default: float) -> float:
 
 
 def _scheduler_variants(row, base_job):
-    # Production dispatch (grouped_moe_gfx1250._maybe_grouped_gfx1250_a8w4_moe)
+    # Production dispatch (grouped_moe_gfx1250.grouped_gemm_gfx1250_a8w4)
     # hardcodes grouped_persistent_m=False and expert_sched_mode=False; the only
     # runtime axis is dense vs DeepGEMM contiguous-M (auto-enabled for large token
     # counts). Mirror exactly that set so AOT never compiles GEMM variants the
@@ -173,7 +130,6 @@ def parse_csv(csv_path: str):
                 "split_k2": int(row.get("split_k2") or 1),
                 "out_dtype": "bf16" if row.get("dtype") == "torch.bfloat16" else "f16",
                 "persistent_workers": _as_int(row.get("persistent_workers"), None),
-                "stage1_weight_layout": row.get("stage1_weight_layout") or "gguu",
                 "act": act,
                 "situ_beta": _as_float(row.get("situ_beta"), 1.0),
                 "situ_linear_beta": _as_float(row.get("situ_linear_beta"), 1.0),
@@ -219,6 +175,7 @@ def _compile_grouped_moe_aux_kernels(job, *, dtype, quant_mode, wmma_rep, contig
     i32 = torch.int32
     u8 = torch.uint8
     bf16 = torch.bfloat16
+    f32 = torch.float32
     E = job["experts"]
     topk = job["topk"]
     model_dim = job["model_dim"]
@@ -233,24 +190,28 @@ def _compile_grouped_moe_aux_kernels(job, *, dtype, quant_mode, wmma_rep, contig
     def _route_ksplit(feat_dim, source_topk, out_e, out_m):
         # build_moe_fused_quant_preshuffle_route_ksplit_module; runtime never
         # sets remap_rows on the grouped MoE fast path (row_starts stays None).
-        launch = build_moe_fused_quant_preshuffle_route_ksplit_module(
-            feat_dim=feat_dim,
-            wmma_rep=wmma_rep,
-            quant_mode=quant_mode,
-            source_topk=source_topk,
-            remap_rows=False,
-        )
-        launch(
-            ptr_arg(torch.empty(0, dtype=bf16, device=dev)),
-            ptr_arg(torch.empty(0, dtype=u8, device=dev)),
-            ptr_arg(torch.empty(0, dtype=u8, device=dev)),
-            ptr_arg(torch.empty(0, dtype=i32, device=dev)),
-            ptr_arg(torch.empty(0, dtype=i32, device=dev)),
-            1,
-            numel,
-            grid,
-            stream=0,
-        )
+        # Precompile both ksplit=True (small token) and ksplit=False (large token).
+        for ks in (True, False):
+            launch = build_moe_fused_quant_preshuffle_route_ksplit_module(
+                feat_dim=feat_dim,
+                wmma_rep=wmma_rep,
+                quant_mode=quant_mode,
+                source_topk=source_topk,
+                remap_rows=False,
+                ksplit=ks,
+            )
+            launch(
+                ptr_arg(torch.empty(0, dtype=bf16, device=dev)),
+                ptr_arg(torch.empty(0, dtype=u8, device=dev)),
+                ptr_arg(torch.empty(0, dtype=u8, device=dev)),
+                ptr_arg(torch.empty(0, dtype=i32, device=dev)),
+                ptr_arg(torch.empty(0, dtype=i32, device=dev)),
+                1,
+                numel,
+                ptr_arg(torch.empty(0, dtype=i32, device=dev)),
+                grid,
+                stream=0,
+            )
 
     def _plain_preshuffle(feat_dim, out_e, out_m, skip_padding):
         launch = build_moe_fused_quant_preshuffle_module(
@@ -303,6 +264,7 @@ def _topids_to_rows():
             E,
             max_m,
             tile_m,
+            ptr_arg(torch.empty(0, dtype=i32, device=dev)),
             stream=0,
         )
 
@@ -339,7 +301,7 @@ def _topids_to_rows():
                 use_expert_row_base=False,
                 max_m=max_m,
             )
-            launch(
+            launch_args = [
                 ptr_arg(torch.empty(0, dtype=i32, device=dev)),
                 ptr_arg(torch.empty(0, dtype=i32, device=dev)),
                 ptr_arg(torch.empty(0, dtype=i32, device=dev)),
@@ -348,9 +310,16 @@ def _topids_to_rows():
                 ptr_arg(torch.empty(0, dtype=u8, device=dev)),
                 ptr_arg(torch.empty(0, dtype=i32, device=dev)),
                 numel,
-                grid,
-                stream=0,
-            )
+            ]
+            if not use_st_ksplit:
+                # Generic path carries the (unused-by-default) EP g2l fusion ABI.
+                launch_args += [
+                    ptr_arg(torch.empty(0, dtype=i32, device=dev)),  # g2l_lut
+                    ptr_arg(torch.empty(0, dtype=f32, device=dev)),  # weight_in
+                    ptr_arg(torch.empty(0, dtype=bf16, device=dev)),  # gather_w
+                    0,  # n_buckets
+                ]
+            launch(*launch_args, grid, stream=0)
         a2_out_e, a2_out_m = E, max_m
 
     # --- Stage2 activation prep (a2): fused grouped quant + preshuffle ---
@@ -394,6 +363,7 @@ def _topids_to_rows():
             ptr_arg(torch.empty(0, dtype=dtype, device=dev)),
             token_num,
             a2_out_e * a2_out_m * (model_dim // 2),
+            ptr_arg(torch.empty(0, dtype=i32, device=dev)),
             stream=0,
         )
 
@@ -402,13 +372,6 @@ def compile_one_config(**job):
     import torch
     from torch._subclasses.fake_tensor import FakeTensorMode
 
-    from aiter.ops.flydsl.kernels.moe_grouped_gemm_mxscale_gfx1250 import (
-        compile_moe_grouped_gemm1_a8w4_masked,
-        compile_moe_grouped_gemm1_mxfp4_masked,
-        compile_moe_grouped_gemm2_a8w4_masked,
-        compile_moe_grouped_gemm2_mxfp4_masked,
-    )
-
     aot_arch = job.pop("gfx", "") or GROUPED_MOE_AOT_ARCH_DEFAULT
     shape_str = (
         # Use .get() so a missing key can't raise here, outside the try below:
@@ -422,179 +385,21 @@ def compile_one_config(**job):
 
     t0 = time.time()
     try:
-        dev = torch.device("cpu")
+        # TODO(aot): only the auxiliary (non-GEMM) kernels are precompiled here.
+        # The grouped GEMM itself moved to the felix 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.
         dtype = torch.bfloat16 if job["out_dtype"] == "bf16" else torch.float16
-        pack = 2 if job["data_format"] == "fp4" else 1
         # Fused prep kernels quantize activations to MXFP8 for a8w4 weights.
         quant_mode = "fp4" if job["data_format"] == "fp4" else "fp8"
-        compiler1 = (
-            compile_moe_grouped_gemm1_mxfp4_masked
-            if job["data_format"] == "fp4"
-            else compile_moe_grouped_gemm1_a8w4_masked
-        )
-        compiler2 = (
-            compile_moe_grouped_gemm2_mxfp4_masked
-            if job["data_format"] == "fp4"
-            else compile_moe_grouped_gemm2_a8w4_masked
-        )
         warp_tile_m = job["tile_m"] // job["m_warp"]
         contiguous = bool(job.get("grouped_contiguous_m", False))
-        common = {
-            "model_dim": job["model_dim"],
-            "inter_dim": job["inter_dim"],
-            "experts": job["experts"],
-            "max_m": job["max_m"],
-            "tile_m": job["tile_m"],
-            "tile_n": job["tile_n"],
-            "tile_k": job["tile_k"],
-            "m_warp": job["m_warp"],
-            "n_warp": job["n_warp"],
-            "out_dtype": job["out_dtype"],
-            "num_buffers": job["num_buffers"],
-            "grouped_persistent_m": job["grouped_persistent_m"],
-            "grouped_contiguous_m": contiguous,
-            "persistent_workers": job["persistent_workers"],
-            "expert_sched_mode": job["expert_sched_mode"],
-        }
-        if contiguous:
-            act_lead = 1
-            ub = job["token_num"] * job["topk"] + job["experts"] * (job["tile_m"] - 1)
-            rows = max(job["tile_m"], _align_up(ub, job["tile_m"]))
-        else:
-            act_lead = job["experts"]
-            rows = job["max_m"]
         with (
             compile_only_env(),
             override_env("FLYDSL_GPU_ARCH", aot_arch),
             FakeTensorMode(),
         ):
-            masked_m = torch.full(
-                (job["experts"],), job["max_m"], dtype=torch.int32, device=dev
-            )
-            # Contiguous-M layout tensor (mirrors runtime psum_t); None otherwise.
-            contiguous_layout = (
-                torch.empty((job["experts"],), dtype=torch.int32, device=dev)
-                if contiguous
-                else None
-            )
-            y1 = torch.empty((act_lead, rows, job["inter_dim"]), dtype=dtype)
-            x1 = torch.empty(
-                (act_lead, rows, job["model_dim"] // pack), dtype=torch.uint8
-            )
-            w1 = torch.empty(
-                (job["experts"], 2 * job["inter_dim"], job["model_dim"] // 2),
-                dtype=torch.uint8,
-            )
-            sx1 = torch.empty(
-                (
-                    act_lead,
-                    *_preshuffled_scale_shape(rows, job["model_dim"], warp_tile_m),
-                ),
-                dtype=torch.uint8,
-            )
-            sw1 = torch.empty(
-                (
-                    job["experts"],
-                    *_preshuffled_b_scale_shape(2 * job["inter_dim"], job["model_dim"]),
-                ),
-                dtype=torch.uint8,
-            )
-            y2 = torch.empty((act_lead, rows, job["model_dim"]), dtype=dtype)
-            x2 = torch.empty(
-                (act_lead, rows, job["inter_dim"] // pack), dtype=torch.uint8
-            )
-            w2 = torch.empty(
-                (job["experts"], job["model_dim"], job["inter_dim"] // 2),
-                dtype=torch.uint8,
-            )
-            sx2 = torch.empty(
-                (
-                    act_lead,
-                    *_preshuffled_scale_shape(rows, job["inter_dim"], warp_tile_m),
-                ),
-                dtype=torch.uint8,
-            )
-            sw2 = torch.empty(
-                (
-                    job["experts"],
-                    *_preshuffled_b_scale_shape(job["model_dim"], job["inter_dim"]),
-                ),
-                dtype=torch.uint8,
-            )
-            exe1 = compiler1(
-                act=job["act"],
-                situ_beta=job.get("situ_beta", 1.0),
-                situ_linear_beta=job.get("situ_linear_beta", 1.0),
-                stage1_weight_layout=job["stage1_weight_layout"],
-                split_k=job["split_k1"],
-                **common,
-            )
-            exe1(
-                y1,
-                x1,
-                w1,
-                sx1,
-                sw1,
-                masked_m,
-                job["max_m"],
-                job["inter_dim"],
-                job["model_dim"],
-                job["experts"],
-                stream=0,
-                _m_tile_map=contiguous_layout,
-            )
-            # Bias-epilogue variant: runtime calls stage1(..., bias=...) when the model
-            # carries per-expert bias (e.g. gpt-oss), which triggers a distinct compiled
-            # kernel (gemm1_bias_* / finalize_act_bias). Precompile it alongside the
-            # bias-free kernel so neither path JITs at first inference.
-            bias1 = torch.empty((job["experts"], 2 * job["inter_dim"]), dtype=dtype)
-            exe1(
-                y1,
-                x1,
-                w1,
-                sx1,
-                sw1,
-                masked_m,
-                job["max_m"],
-                job["inter_dim"],
-                job["model_dim"],
-                job["experts"],
-                stream=0,
-                _m_tile_map=contiguous_layout,
-                bias=bias1,
-            )
-            exe2 = compiler2(split_k=job["split_k2"], **common)
-            exe2(
-                y2,
-                x2,
-                w2,
-                sx2,
-                sw2,
-                masked_m,
-                job["max_m"],
-                job["model_dim"],
-                job["inter_dim"],
-                job["experts"],
-                stream=0,
-                _m_tile_map=contiguous_layout,
-            )
-            # Bias-epilogue variant for stage2 (gemm2_bias_*); see stage1 note above.
-            bias2 = torch.empty((job["experts"], job["model_dim"]), dtype=dtype)
-            exe2(
-                y2,
-                x2,
-                w2,
-                sx2,
-                sw2,
-                masked_m,
-                job["max_m"],
-                job["model_dim"],
-                job["inter_dim"],
-                job["experts"],
-                stream=0,
-                _m_tile_map=contiguous_layout,
-                bias=bias2,
-            )
             # Non-GEMM auxiliary kernels the run-only fast path launches around
             # the GEMMs (fused route+quant+scatter, grouped quant+preshuffle,
             # contiguous prefix-sum(+remap), gather-reduce). These were fused in