MoE: correct + faster a16w4 (bf16 A x MXFP4 W) SiTUv2 kernel (replaces old a16w4)

aiter/aot/flydsl/moe.py

@@ -414,7 +414,7 @@ def _make_a_user(a_dtype_user_shape):
         sorted_token_ids, sorted_expert_ids, num_valid_ids = _make_routing()
 
         if stage == 1:
-            if b_dtype in ("fp4", "mxfp4") and a_dtype in ("bf16", "fp8"):
+            if b_dtype == "fp4" and a_dtype in ("bf16", "fp8"):
                 tile_n = resolve_flydsl_stage1_tile_n(inter_dim, tile_n)
 
             a = _make_a_user(_user_a_shape())
@@ -582,9 +582,7 @@ def _make_a_user(a_dtype_user_shape):
                     ),
                     stream=0,
                     swiglu_limit=runtime_swiglu_limit(None, act),
-                    pass_swiglu_limit=not (
-                        a_dtype == "bf16" and b_dtype in ("fp4", "mxfp4")
-                    ),
+                    pass_swiglu_limit=not (a_dtype == "bf16" and b_dtype == "fp4"),
                 )
             else:
                 args = _s1_args_std(
@@ -842,6 +840,102 @@ def _make_a_user(a_dtype_user_shape):
                 )
 
 
+def _precompile_a16w4_to_cache(
+    *,
+    stage: int,
+    model_dim: int,
+    inter_dim: int,
+    experts: int,
+    topk: int,
+    tile_m: int,
+    tile_n: int,
+    tile_k: int,
+    act: str = "silu",
+    b_nt: int = 2,
+    xcd_swizzle: int = 0,
+    k_wave: int = 1,
+    **kwargs,
+):
+    """AOT for the folded a16w4 (bf16 A x fp4 W) port (moe_2stage_a16wmix).
+
+    The port launch ABI (raw fx.Int64 device pointers) differs from the generic MX
+    gemm (``_s1_args_fp4``), so it can't reuse ``_precompile_to_cache``'s arg
+    builders.  Instead drive the SAME runtime launchers (``flydsl_a16w4_gemm{1,2}``)
+    the fused-MoE op uses, under ``COMPILE_ONLY=1`` — the cache key then matches
+    runtime by construction (``waves_per_eu=None``, ``persist=False``,
+    ``w_layout="standard"``, g2 tile downgrade are all applied inside the
+    launcher).  The compiled artifact is keyed only on the kernel's constexpr
+    params (shapes/tiles/topk/act), never on the launch pointers, grid, or
+    ``n_tokens``, so a single 1-elem real placeholder covers every buffer (real,
+    not fake, because the launcher calls ``.data_ptr()``; COMPILE_ONLY guarantees
+    it is never dereferenced).
+    """
+    import torch
+
+    from aiter.ops.flydsl.kernels.moe_2stage_a16wmix import (
+        flydsl_a16w4_gemm1,
+        flydsl_a16w4_gemm2,
+    )
+
+    z = torch.zeros(1, dtype=torch.int32, device="cpu")
+    common = {
+        "sorted_expert_ids": z,
+        "cumsum_tensor": z,
+        "NE": experts,
+        "D_HIDDEN": model_dim,
+        "D_INTER": inter_dim,
+        "topk": topk,
+        "tile_m": tile_m,
+        "b_nt": b_nt,
+        "xcd_swizzle": xcd_swizzle,
+        "waves_per_eu": None,
+        "stream": 0,
+    }
+    with compile_only_env():
+        if stage == 1:
+            flydsl_a16w4_gemm1(
+                a_bf16=z,
+                w1_u8=z,
+                w1_scale_u8=z,
+                m_indices=z,
+                inter_sorted_bf16=z,
+                n_tokens=1,
+                tile_n=tile_n,
+                tile_k=tile_k,
+                k_wave=k_wave,
+                act=("situv2" if act in ("situv2", "situ") else act),
+                w_layout="standard",
+                **common,
+            )
+        else:
+            # Mirror the runtime stage2 g2 tile downgrade (moe_kernels
+            # _flydsl_moe_stage2_impl a16w4 branch).
+            g2_tile_n = (
+                tile_n
+                if model_dim % tile_n == 0
+                else (256 if model_dim % 256 == 0 else 128)
+            )
+            g2_tile_k = (
+                tile_k
+                if inter_dim % tile_k == 0
+                else (128 if inter_dim % 128 == 0 else 64)
+            )
+            flydsl_a16w4_gemm2(
+                inter_sorted_bf16=z,
+                w2_u8=z,
+                w2_scale_u8=z,
+                sorted_token_ids=z,
+                sorted_weights=z,
+                flat_out=z,
+                M_logical=1,
+                max_sorted=1,
+                tile_n=g2_tile_n,
+                tile_k=g2_tile_k,
+                w_dtype="fp4",
+                **common,
+            )
+
+
 def _precompile_epilogue_to_cache(act: str, inter_dim: int, topk: int):
     """Precompile the CK-Tile split-K post-activation epilogue kernel.
 
@@ -941,37 +1035,58 @@ def compile_one_config(
 
     from torch._subclasses.fake_tensor import FakeTensorMode
 
+    # The folded a16w4 (bf16 A x fp4 W) port takes raw .data_ptr() device
+    # pointers, which FakeTensors don't have; drive its dedicated precompile with
+    # real (COMPILE_ONLY) tensors outside FakeTensorMode.
+    is_a16w4 = (
+        not is_epilogue
+        and kwargs.get("shared_expert_id", -1) < 0
+        and kwargs.get("a_dtype") == "bf16"
+        and kwargs.get("b_dtype") == "fp4"
+    )
+
     t0 = time.time()
     try:
-        with (
-            override_env("FLYDSL_GPU_ARCH", aot_arch),
-            FakeTensorMode(),
-        ):
-            if is_epilogue:
-                _precompile_epilogue_to_cache(
-                    act=kwargs.get("act", "silu"),
-                    inter_dim=inter_dim,
-                    topk=topk,
-                )
-            else:
-                shared_expert_id = kwargs.pop("shared_expert_id", -1)
-                if shared_expert_id >= 0:
-                    from aiter.aot.flydsl.fhmoe import (
-                        precompile_fhmoe_to_cache,
-                    )
-
-                    precompile = precompile_fhmoe_to_cache
-                    kwargs["shared_expert_id"] = shared_expert_id
-                else:
-                    precompile = _precompile_to_cache
-                precompile(
+        if is_a16w4:
+            with override_env("FLYDSL_GPU_ARCH", aot_arch):
+                _precompile_a16w4_to_cache(
                     model_dim=model_dim,
                     inter_dim=inter_dim,
                     experts=experts,
                     topk=topk,
                     cu_num=cu_num,
                     **kwargs,
                 )
+        else:
+            with (
+                override_env("FLYDSL_GPU_ARCH", aot_arch),
+                FakeTensorMode(),
+            ):
+                if is_epilogue:
+                    _precompile_epilogue_to_cache(
+                        act=kwargs.get("act", "silu"),
+                        inter_dim=inter_dim,
+                        topk=topk,
+                    )
+                else:
+                    shared_expert_id = kwargs.pop("shared_expert_id", -1)
+                    if shared_expert_id >= 0:
+                        from aiter.aot.flydsl.fhmoe import (
+                            precompile_fhmoe_to_cache,
+                        )
+
+                        precompile = precompile_fhmoe_to_cache
+                        kwargs["shared_expert_id"] = shared_expert_id
+                    else:
+                        precompile = _precompile_to_cache
+                    precompile(
+                        model_dim=model_dim,
+                        inter_dim=inter_dim,
+                        experts=experts,
+                        topk=topk,
+                        cu_num=cu_num,
+                        **kwargs,
+                    )
         elapsed = time.time() - t0
         result["compile_time"] = elapsed
         print(f"  [OK] compile  {elapsed:6.1f}s  {shape_str}  arch={aot_arch}")