From b379ee226c8af04353b525c4813bf0f8c81cb8d0 Mon Sep 17 00:00:00 2001 From: lalala-sh Date: Tue, 14 Apr 2026 09:59:42 +0000 Subject: [PATCH 1/5] fix moe splitk aot and jit --- aiter/jit/core.py | 7 +++++-- aiter/ops/moe_op.py | 3 ++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/aiter/jit/core.py b/aiter/jit/core.py index 098af6ce6a..a71ed5e128 100644 --- a/aiter/jit/core.py +++ b/aiter/jit/core.py @@ -1402,8 +1402,11 @@ def wrapper(*args, custom_build_args={}, **kwargs): rebuilded_list.append(md_name) raise ModuleNotFoundError("start rebuild") if module is None: - md = custom_build_args.get("md_name", md_name) - module = get_module(md) + try: + module = get_module(md_name) + except Exception: + md = custom_build_args.get("md_name", md_name) + module = get_module(md) except ModuleNotFoundError: d_args = get_args_of_build(md_name) d_args.update(custom_build_args) diff --git a/aiter/ops/moe_op.py b/aiter/ops/moe_op.py index 8356ea12a6..88bfdb76d4 100755 --- a/aiter/ops/moe_op.py +++ b/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 kernelName is None + is_splitk = is_splitk or "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, From f9134415d858a792c169efd7dc2c87451f900a21 Mon Sep 17 00:00:00 2001 From: lalala-sh Date: Wed, 15 Apr 2026 05:16:07 +0000 Subject: [PATCH 2/5] split moe aot to serveral libs base on tuned_moe configs --- aiter/jit/core.py | 7 +- aiter/jit/utils/moe_recipes.py | 168 +++++++++++++++++++++++++++++++++ setup.py | 25 +++++ 3 files changed, 195 insertions(+), 5 deletions(-) create mode 100644 aiter/jit/utils/moe_recipes.py diff --git a/aiter/jit/core.py b/aiter/jit/core.py index a71ed5e128..098af6ce6a 100644 --- a/aiter/jit/core.py +++ b/aiter/jit/core.py @@ -1402,11 +1402,8 @@ def wrapper(*args, custom_build_args={}, **kwargs): rebuilded_list.append(md_name) raise ModuleNotFoundError("start rebuild") if module is None: - try: - module = get_module(md_name) - except Exception: - md = custom_build_args.get("md_name", md_name) - module = get_module(md) + md = custom_build_args.get("md_name", md_name) + module = get_module(md) except ModuleNotFoundError: d_args = get_args_of_build(md_name) d_args.update(custom_build_args) diff --git a/aiter/jit/utils/moe_recipes.py b/aiter/jit/utils/moe_recipes.py new file mode 100644 index 0000000000..92c12f172d --- /dev/null +++ b/aiter/jit/utils/moe_recipes.py @@ -0,0 +1,168 @@ +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 diff --git a/setup.py b/setup.py index cf976fd760..924b3da4e2 100644 --- a/setup.py +++ b/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 = [] From eeb019df0a8f35e8645ceb9c42183748c42ca953 Mon Sep 17 00:00:00 2001 From: lalala-sh Date: Wed, 15 Apr 2026 05:18:48 +0000 Subject: [PATCH 3/5] update copyrights --- aiter/jit/utils/moe_recipes.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/aiter/jit/utils/moe_recipes.py b/aiter/jit/utils/moe_recipes.py index 92c12f172d..56ed009e5f 100644 --- a/aiter/jit/utils/moe_recipes.py +++ b/aiter/jit/utils/moe_recipes.py @@ -1,3 +1,8 @@ +# 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 From 60e13989950314b02120c11fd6c38cf1ba90478f Mon Sep 17 00:00:00 2001 From: lalala-sh Date: Wed, 15 Apr 2026 06:25:36 +0000 Subject: [PATCH 4/5] fix typo --- aiter/ops/moe_op.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/aiter/ops/moe_op.py b/aiter/ops/moe_op.py index 88bfdb76d4..30c9f787a9 100755 --- a/aiter/ops/moe_op.py +++ b/aiter/ops/moe_op.py @@ -512,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 {{}}" ] From 43b69090bf2266afa6a46e8b1eb7f5f70137164b Mon Sep 17 00:00:00 2001 From: lalala-sh Date: Wed, 15 Apr 2026 09:23:10 +0000 Subject: [PATCH 5/5] test shuffle as default and fix moe split jit --- aiter/ops/moe_op.py | 4 ++-- op_tests/test_moe_2stage.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/aiter/ops/moe_op.py b/aiter/ops/moe_op.py index 30c9f787a9..4ea3a5c2e1 100755 --- a/aiter/ops/moe_op.py +++ b/aiter/ops/moe_op.py @@ -236,8 +236,8 @@ def cmdGenFunc_ck_moe_stage( ): mul_routed_weight_stage = 2 if sorted_weights is None else 1 - is_splitk = splitk > 1 and kernelName is None - is_splitk = is_splitk or "Splitk" in kernelName + 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, diff --git a/op_tests/test_moe_2stage.py b/op_tests/test_moe_2stage.py index 52762c568f..b5e716fb40 100644 --- a/op_tests/test_moe_2stage.py +++ b/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.""",