fix moe splitk aot and jit

aiter/jit/utils/moe_recipes.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+
+# mypy: allow-untyped-defs
+
+import csv
+from pathlib import Path
+from typing import Dict, List, Set, Tuple
+
+_DTYPE_MAP = {
+    "torch.float8_e4m3fn": "f8",
+    "torch.float8_e4m3fnuz": "f8",
+    "torch.bfloat16": "b16",
+    "torch.float16": "f16",
+    "torch.int8": "i8",
+    "torch.float4_e2m1fn_x2": "fp4x2",
+}
+
+_QUANT_ALIAS = {"per_128x128": "per_1x128"}
+
+
+def _build_moe_variant(
+    aiter_csrc_dir: str,
+    a_dtype: str,
+    b_dtype: str,
+    c_dtype: str,
+    activation: str,
+    quant_type: str,
+    mul_routed_weight_stage: int,
+    preshuffle_mode: bool,
+    is_splitk: bool,
+) -> Tuple[str, list]:
+    parts = [
+        "module_moe_ck2stages",
+        a_dtype,
+        b_dtype,
+        "preshuffle_on" if preshuffle_mode else "preshuffle_off",
+        c_dtype,
+        activation,
+        quant_type,
+        f"mulWeightStage{mul_routed_weight_stage}",
+    ]
+    if is_splitk:
+        parts.append("splitk")
+
+    flags = ""
+    if preshuffle_mode and b_dtype == "fp4x2":
+        flags += " --preshuffle"
+    if is_splitk:
+        flags += " --issplitk"
+
+    md_name = "_".join(parts)
+    blob_gen_cmd = [
+        (
+            f"{aiter_csrc_dir}/ck_gemm_moe_2stages_codegen/gen_instances.py "
+            f"-a {a_dtype} -b {b_dtype} -c {c_dtype} -q {quant_type} "
+            f"-act {activation} -m {mul_routed_weight_stage}{flags} -w {{}}"
+        )
+    ]
+    return md_name, blob_gen_cmd
+
+
+def _normalize_dtype(dtype: str) -> str:
+    return _DTYPE_MAP.get(dtype, dtype)
+
+
+def _normalize_enum_str(s: str) -> str:
+    """'QuantType.per_1x128' -> 'per_1x128', 'ActivationType.Silu' -> 'silu'.
+
+    Same pattern as moe_op.py: ``str(enum).split(".")[-1].lower()``.
+    """
+    return s.split(".")[-1].lower() if "." in s else s.lower()
+
+
+def _normalize_quant_type(quant_type: str) -> str:
+    q = _normalize_enum_str(quant_type)
+    return _QUANT_ALIAS.get(q, q)
+
+
+def _normalize_activation(activation: str) -> str:
+    return _normalize_enum_str(activation)
+
+
+def _infer_preshuffle_modes(b_dtype: str, quant_type: str) -> List[bool]:
+    """Infer preshuffle modes based on runtime behavior.
+
+    - fp4x2: may or may not be pre-shuffled -> both off and on
+    - no-quant (b16/f16): always pre-shuffled at runtime -> preshuffle_on only
+    - other quantized types: always shuffled -> preshuffle_on only
+    """
+    if b_dtype == "fp4x2":
+        return [False, True]
+    return [True]
+
+
+def _should_include_splitk(row: Dict, quant_type: str) -> bool:
+    """splitk only applies to f8/f8 per_1x128 (blockscale) kernels."""
+    if quant_type != "per_1x128":
+        return False
+    ksplit = row.get("ksplit")
+    if not ksplit:
+        return False
+    try:
+        return int(float(ksplit)) > 1
+    except (TypeError, ValueError):
+        return False
+
+
+def _get_mul_weight_stage(row: Dict) -> int:
+    value = row.get("doweight_stage1")
+    if not value:
+        return 2
+    v = str(value).strip().lower()
+    if v in ("1", "true"):
+        return 1
+    return 2
+
+
+def _get_tuned_fmoe_rows() -> List[Dict]:
+    configs_dir = Path(__file__).resolve().parents[2] / "configs"
+    model_configs_dir = configs_dir / "model_configs"
+    tuned_paths = [configs_dir / "tuned_fmoe.csv"]
+    tuned_paths.extend(sorted(model_configs_dir.glob("*tuned_fmoe*.csv")))
+
+    rows: List[Dict] = []
+    for path in tuned_paths:
+        if not path.exists():
+            continue
+        with path.open() as f:
+            rows.extend(csv.DictReader(f))
+    return rows
+
+
+def get_moe_ck2stages_prebuild_variants(aiter_csrc_dir: str) -> List[Dict]:
+    seen: Set[Tuple] = set()
+    results: List[Dict] = []
+    for row in _get_tuned_fmoe_rows():
+        c_dtype = _normalize_dtype(row["dtype"])
+        a_dtype = _normalize_dtype(row.get("q_dtype_a") or row["dtype"])
+        b_dtype = _normalize_dtype(row.get("q_dtype_w") or row["dtype"])
+        quant_type = _normalize_quant_type(row["q_type"])
+        activation = _normalize_activation(row["act_type"])
+        mul_weight_stage = _get_mul_weight_stage(row)
+        need_splitk = _should_include_splitk(row, quant_type)
+
+        for preshuffle in _infer_preshuffle_modes(b_dtype, quant_type):
+            for splitk in [False, True] if need_splitk else [False]:
+                key = (
+                    a_dtype,
+                    b_dtype,
+                    c_dtype,
+                    quant_type,
+                    activation,
+                    mul_weight_stage,
+                    preshuffle,
+                    splitk,
+                )
+                if key in seen:
+                    continue
+                seen.add(key)
+                md_name, blob_gen_cmd = _build_moe_variant(
+                    aiter_csrc_dir=aiter_csrc_dir,
+                    a_dtype=a_dtype,
+                    b_dtype=b_dtype,
+                    c_dtype=c_dtype,
+                    activation=activation,
+                    quant_type=quant_type,
+                    mul_routed_weight_stage=mul_weight_stage,
+                    preshuffle_mode=preshuffle,
+                    is_splitk=splitk,
+                )
+                results.append({"md_name": md_name, "blob_gen_cmd": blob_gen_cmd})
+    return results

aiter/ops/moe_op.py

@@ -236,7 +236,8 @@ def cmdGenFunc_ck_moe_stage(
 ):
 
     mul_routed_weight_stage = 2 if sorted_weights is None else 1
-    is_splitk = splitk > 1
+    is_splitk = splitk > 1 and not kernelName
+    is_splitk = is_splitk or (bool(kernelName) and "Splitk" in kernelName)
     outtype = str2dtype_dict[dst_type] if is_splitk else out.dtype
     md_name, blob_gen_cmd = get_moe_stage_module(
         hidden_states.dtype,
@@ -511,19 +512,19 @@ def get_moe_stage_module(
     act = str(activation).split(".")[-1].lower()
     quant_type = str(quant_type).split(".")[-1].lower()
 
-    md_name = ("_").join(
-        [
-            "module_moe_ck2stages",
-            Adtype,
-            Bdtype,
-            "preshuffle_on" if preshuffle_mode else "preshuffle_off",
-            Cdtype,
-            act,
-            quant_type,
-            f"mulWeightStage{mul_routed_weight_stage}",
-            "splitk" if is_splitk else "",
-        ]
-    )
+    parts = [
+        "module_moe_ck2stages",
+        Adtype,
+        Bdtype,
+        "preshuffle_on" if preshuffle_mode else "preshuffle_off",
+        Cdtype,
+        act,
+        quant_type,
+        f"mulWeightStage{mul_routed_weight_stage}",
+    ]
+    if is_splitk:
+        parts.append("splitk")
+    md_name = "_".join(parts)
     blob_gen_cmd = [
         f"{AITER_CSRC_DIR}/ck_gemm_moe_2stages_codegen/gen_instances.py -a {Adtype} -b {Bdtype} -c {Cdtype} -q {quant_type} -act {act} -m {mul_routed_weight_stage} {preshuffle_str} {splitk_str} -w {{}}"
     ]

op_tests/test_moe_2stage.py

@@ -47,7 +47,7 @@ def test_fmoe(
     doweight_stage1=False,
     hidden_pad=0,
     intermediate_pad=0,
-    preshuffle=False,
+    preshuffle=True,
 ):
     if get_gfx() not in ["gfx950"] and qType == aiter.QuantType.per_1x32:
         return
@@ -391,7 +391,7 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor):
     "--preshuffle",
     type=dtypes.str2bool,
     nargs="*",
-    default=[False, True],
+    default=[True],
     help="""Whether to use pre-shuffle weight mode. Default is [False, True].
     -p f    # False.
     -p t    # True.""",

setup.py

@@ -217,11 +217,36 @@ def get_exclude_ops():
         from jit.utils.mha_recipes import (
             get_mha_varlen_prebuild_variants_by_names,
         )
+        from jit.utils.moe_recipes import get_moe_ck2stages_prebuild_variants
         import glob
 
         exclude_ops = get_exclude_ops()
         all_opts_args_build, _ = core.get_args_of_build("all", exclude=exclude_ops)
 
+        moe_base_args = None
+        filtered_opts_args_build = []
+        for one_opt_args in all_opts_args_build:
+            if one_opt_args["md_name"] == "module_moe_ck2stages":
+                moe_base_args = one_opt_args
+                continue
+            filtered_opts_args_build.append(one_opt_args)
+        all_opts_args_build = filtered_opts_args_build
+
+        if ENABLE_CK and moe_base_args is not None:
+            moe_variants = get_moe_ck2stages_prebuild_variants(core.AITER_CSRC_DIR)
+            for v in moe_variants:
+                all_opts_args_build.append(
+                    {
+                        "md_name": v["md_name"],
+                        "srcs": moe_base_args["srcs"],
+                        "flags_extra_cc": moe_base_args["flags_extra_cc"],
+                        "flags_extra_hip": moe_base_args["flags_extra_hip"],
+                        "extra_include": moe_base_args["extra_include"],
+                        "blob_gen_cmd": v["blob_gen_cmd"],
+                        "third_party": moe_base_args["third_party"],
+                    }
+                )
+
         if PREBUILD_KERNELS == 1 and ENABLE_CK:
             extra_args_build = []