From 030f49f27086a255520abdd5d20be0fdcb0cf690 Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Wed, 16 Sep 2026 12:01:53 +0200 Subject: [PATCH 01/14] [FlyDSL] Draft FHMoE tuner: catalogue, --fhmoe, and fused_moe smoke Fill untuned_fhmoe.csv with DSV4 I384 shapes and add FhmoeTuner (FMoE as template). generate_fhmoe_data builds routed/shared tensors, pads the dummy expert, preshuffles from gate_mode, and _run_fhmoe calls fused_moe. Candidate timing and tuned CSV write are not wired yet. --- aiter/configs/untuned_fhmoe.csv | 14 +- .../gemm_moe_tune.py | 239 +++++++++++++++++- 2 files changed, 250 insertions(+), 3 deletions(-) diff --git a/aiter/configs/untuned_fhmoe.csv b/aiter/configs/untuned_fhmoe.csv index d1277b4689f..c9a3db35374 100644 --- a/aiter/configs/untuned_fhmoe.csv +++ b/aiter/configs/untuned_fhmoe.csv @@ -1 +1,13 @@ -gfx,cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,q_type,use_g1u1,doweight_stage1 +token,model_dim,inter_dim,expert,topk,shared_expert_id,act_type,dtype,q_dtype_a,q_dtype_w,q_type,use_g1u1,doweight_stage1,hidden_pad,intermediate_pad,gate_mode +1,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +2,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +4,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +8,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +16,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +32,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +64,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +128,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +256,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +512,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +1024,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE +2048,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index d38fa05a290..841ff500031 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -46,6 +46,7 @@ ) from aiter.jit.core import ( AITER_CONFIG_FMOE, + AITER_CONFIG_FHMOE, AITER_CONFIG_GROUPED_FMOE, AITER_CSRC_DIR, AITER_ROOT_DIR, @@ -55,6 +56,7 @@ from aiter.ops.flydsl.moe_common import ( DEFAULT_SITUV2_BETA, DEFAULT_SITUV2_LINEAR_BETA, + GateMode, get_flydsl_activation_name, ) from aiter.ops.flydsl.moe_kernels import ( @@ -492,18 +494,25 @@ def _setup_specific_arguments(self): required=False, help="Only last kernel is tuned, if not, only kernels that are not in the tuned_fmoe.csv are tuned", ) - self.parser.add_argument( + tuner_mode = self.parser.add_mutually_exclusive_group() + tuner_mode.add_argument( "--grouped-gemm", action="store_true", required=False, help="On gfx1250, tune the FlyDSL grouped-GEMM MoE path instead of the normal fmoe tuner.", ) - self.parser.add_argument( + tuner_mode.add_argument( "--mxfp4-flydsl", action="store_true", required=False, help="Tune the FlyDSL mxfp4 a4w4 port as a coupled (g1, g2) unit instead of the normal fmoe tuner.", ) + tuner_mode.add_argument( + "--fhmoe", + action="store_true", + required=False, + help="Tune fused heterogeneous MoE (FHMoE) instead of the normal fmoe tuner.", + ) self.parser.add_argument( "--mxfp4-search-mode", choices=("prune", "full"), @@ -6399,6 +6408,216 @@ def result_to_csv(self, results, file, concat=False): resultdf.to_csv(file, index=False) +class FhmoeTuner(FmoeTuner): + ARG_DEFAULTS: ClassVar[dict[str, Any]] = { + **FmoeTuner.ARG_DEFAULTS, + "tune_file": f"{AITER_CONFIG_FHMOE}", + "untune_file": f"{AITER_ROOT_DIR}/aiter/configs/untuned_fhmoe.csv", + "config_env_name": "AITER_CONFIG_FHMOE", + } + + @staticmethod + def _parse_fhmoe_row(row): + return { + "gfx": str(row["gfx"]), + "cu_num": int(row["cu_num"]), + "token": int(row["token"]), + "model_dim": int(row["model_dim"]), + "inter_dim": int(row["inter_dim"]), + "expert": int(row["expert"]), + "topk": int(row["topk"]), + "shared_expert_id": int(row["shared_expert_id"]), + "act_type": eval(row["act_type"]), + "dtype": eval(row["dtype"]), + "q_dtype_a": eval(row["q_dtype_a"]), + "q_dtype_w": eval(row["q_dtype_w"]), + "q_type": eval(row["q_type"]), + "use_g1u1": int(row["use_g1u1"]), + "doweight_stage1": int(row["doweight_stage1"]), + "hidden_pad": int(row["hidden_pad"]), + "intermediate_pad": int(row["intermediate_pad"]), + "gate_mode": eval(row["gate_mode"]), + } + + @staticmethod + def _pad_dummy_expert(tensor, n_experts=None): + if tensor.ndim == 2: + rows_per_expert = tensor.shape[0] // n_experts + dummy_shape = (rows_per_expert, tensor.shape[1]) + else: + dummy_shape = (1, *tensor.shape[1:]) + dummy = tensor.new_empty(dummy_shape) + dummy.view(torch.uint8).zero_() + return torch.cat([tensor, dummy], dim=0) + + @staticmethod + def generate_fhmoe_data(shape, block_m=32, device="cuda"): + expert = shape["expert"] + topk = shape["topk"] + shared_expert_id = shape["shared_expert_id"] + if topk < 2: + raise ValueError(f"FHMoE topk must include the shared expert; got {topk=}") + + routed_expert = expert - 1 + routed_topk = topk - 1 + data = FmoeTuner.generate_data( + shape["token"], + shape["model_dim"], + shape["inter_dim"], + routed_expert, + routed_topk, + shape["dtype"], + shape["q_dtype_a"], + shape["q_dtype_w"], + shape["q_type"], + shape["use_g1u1"], + block_m, + device, + ) + w1 = FhmoeTuner._pad_dummy_expert(data["w1_qt"]) + w2 = FhmoeTuner._pad_dummy_expert(data["w2_qt"]) + w1_scale = FhmoeTuner._pad_dummy_expert(data["w1_scale"], routed_expert) + w2_scale = FhmoeTuner._pad_dummy_expert(data["w2_scale"], routed_expert) + + shared_device = w1.device + shared_w1_dense = ( + torch.randn( + (1, 2 * shape["inter_dim"], shape["model_dim"]), + dtype=shape["dtype"], + device=shared_device, + ) + / 10 + ) + shared_w2_dense = ( + torch.randn( + (1, shape["model_dim"], shape["inter_dim"]), + dtype=shape["dtype"], + device=shared_device, + ) + / 10 + ) + shared_w1, shared_w1_scale = FmoeTuner.weight_quant( + shared_w1_dense, shape["q_type"], quant_dtype=dtypes.fp8 + ) + shared_w2, shared_w2_scale = FmoeTuner.weight_quant( + shared_w2_dense, shape["q_type"], quant_dtype=dtypes.fp8 + ) + del shared_w1_dense, shared_w2_dense + + shared_ids = torch.full( + (shape["token"], 1), + shared_expert_id, + dtype=data["topk_ids"].dtype, + device=data["topk_ids"].device, + ) + shared_weights = torch.ones( + (shape["token"], 1), + dtype=data["topk_weights"].dtype, + device=data["topk_weights"].device, + ) + topk_ids = torch.cat([data["topk_ids"], shared_ids], dim=1) + topk_weights = torch.cat([data["topk_weights"], shared_weights], dim=1) + + if shape["gate_mode"] == GateMode.INTERLEAVE: + w1 = shuffle_weight_a16w4(w1, 16, True) + w1_scale = shuffle_scale_a16w4(w1_scale, expert, True) + w2 = shuffle_weight_a16w4(w2, 16, False) + w2_scale = shuffle_scale_a16w4(w2_scale, expert, False) + shared_w1 = shuffle_weight_a16w4(shared_w1, 16, True) + shared_w1_scale = shuffle_scale_a16w4(shared_w1_scale, 1, True) + shared_w2 = shuffle_weight_a16w4(shared_w2, 16, False) + shared_w2_scale = shuffle_scale_a16w4(shared_w2_scale, 1, False) + else: + w1 = shuffle_weight(w1, layout=(16, 16)) + w1_scale = shuffle_scale(w1_scale) + w2 = shuffle_weight(w2, layout=(16, 16)) + w2_scale = shuffle_scale(w2_scale) + shared_w1 = shuffle_weight(shared_w1, layout=(16, 16)) + shared_w1_scale = shuffle_scale(shared_w1_scale) + shared_w2 = shuffle_weight(shared_w2, layout=(16, 16)) + shared_w2_scale = shuffle_scale(shared_w2_scale) + w1.is_shuffled = True + w2.is_shuffled = True + + return { + "hidden": data["input"], + "w1": w1, + "w2": w2, + "w1_scale": w1_scale, + "w2_scale": w2_scale, + "shared_w1": shared_w1, + "shared_w2": shared_w2, + "shared_w1_scale": shared_w1_scale, + "shared_w2_scale": shared_w2_scale, + "topk_ids": topk_ids, + "topk_weights": topk_weights, + "a1_qt": data["a1_qt"], + "a1_scale": data["a1_scale"], + "gate_mode": shape["gate_mode"], + "hidden_pad": shape["hidden_pad"], + "intermediate_pad": shape["intermediate_pad"], + "shared_expert_id": shared_expert_id, + } + + @staticmethod + def _run_fhmoe(shape, data): + gate_mode = shape["gate_mode"] + if not isinstance(gate_mode, str): + gate_mode = gate_mode.value + + def _cuda(tensor): + if tensor is None or not torch.is_tensor(tensor): + return tensor + return tensor.cuda() + + w1 = _cuda(data["w1"]) + w2 = _cuda(data["w2"]) + w1.is_shuffled = True + w2.is_shuffled = True + return fused_moe( + _cuda(data["hidden"]), + w1, + w2, + _cuda(data["topk_weights"]), + _cuda(data["topk_ids"]), + activation=shape["act_type"], + quant_type=shape["q_type"], + doweight_stage1=bool(shape["doweight_stage1"]), + w1_scale=_cuda(data["w1_scale"]), + w2_scale=_cuda(data["w2_scale"]), + dtype=shape["dtype"], + hidden_pad=shape["hidden_pad"], + intermediate_pad=shape["intermediate_pad"], + gate_mode=gate_mode, + shared_w1=_cuda(data["shared_w1"]), + shared_w2=_cuda(data["shared_w2"]), + shared_w1_scale=_cuda(data["shared_w1_scale"]), + shared_w2_scale=_cuda(data["shared_w2_scale"]), + shared_expert_id=data["shared_expert_id"], + ) + + def tune(self, untunedf, tunedf, args): + del tunedf, args + for _, row in untunedf.iterrows(): + shape = self._parse_fhmoe_row(row) + print("\nStart tuning FHMoE", [shape[k] for k in self.keys]) + data = self.generate_fhmoe_data(shape) + print( + f"[fhmoe] token={shape['token']} " + f"w1={tuple(data['w1'].shape)} {data['w1'].dtype} " + f"shared_w1={tuple(data['shared_w1'].shape)} {data['shared_w1'].dtype} " + f"topk_ids={tuple(data['topk_ids'].shape)}", + flush=True, + ) + out = self._run_fhmoe(shape, data) + print( + f"[fhmoe] fused_moe out={tuple(out.shape)} {out.dtype}", + flush=True, + ) + del data, out + return [] + + class Mxfp4FlydslTuner(FmoeTuner): """Tune the FlyDSL mxfp4 a4w4 *port* (flydsl_mxmoe_g{1,2}_a4w4_*) as one coupled unit. @@ -7070,6 +7289,12 @@ def _reap(idx): # the stage2 tile optimum), so it must be in the dedup key or the write-back # collapses an ep_fused row into the generic one for the same shape. grouped_key = key + ["gate_mode", "ep_fused"] + fhmoe_key = key + [ + "shared_expert_id", + "hidden_pad", + "intermediate_pad", + "gate_mode", + ] resultList = [ "block_m", "ksplit", @@ -7105,6 +7330,7 @@ def _reap(idx): ] use_grouped = "--grouped-gemm" in sys.argv use_mxfp4_flydsl = "--mxfp4-flydsl" in sys.argv + use_fhmoe = "--fhmoe" in sys.argv if use_grouped: if get_gfx() != "gfx1250": raise SystemExit("--grouped-gemm is only supported on gfx1250") @@ -7118,6 +7344,15 @@ def _reap(idx): tuner = Mxfp4FlydslTuner( "mxfp4FlydslTuner", key, resultList, "mxfp4 a4w4 flydsl port fmoe tuner" ) + elif use_fhmoe: + if get_gfx() != "gfx950": + raise SystemExit("--fhmoe is only supported on gfx950") + tuner = FhmoeTuner( + "fhmoeTuner", + fhmoe_key, + resultList, + "fhmoe tuner", + ) else: tuner = FmoeTuner("fmoeTuner", key, resultList + ["nt"], "fmoe tuner") args = tuner.parse_args() From 1f0509c3d6be32622aace404d15501a7626e64cb Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Wed, 16 Sep 2026 15:48:09 +0200 Subject: [PATCH 02/14] [FlyDSL] Run FHMoE tuner tasks through mp_tuner with host-side weight quant DSV4 dense MXFP4 quant peaks at several GiB; keep FmoeTuner.generate_data on CUDA and build generate_fhmoe_data on CPU so the worker can time fused_moe. --- .../gemm_moe_tune.py | 228 +++++++++++++----- 1 file changed, 169 insertions(+), 59 deletions(-) diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index 841ff500031..0a7079678f8 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -6460,63 +6460,108 @@ def generate_fhmoe_data(shape, block_m=32, device="cuda"): routed_expert = expert - 1 routed_topk = topk - 1 - data = FmoeTuner.generate_data( - shape["token"], - shape["model_dim"], - shape["inter_dim"], - routed_expert, - routed_topk, - shape["dtype"], - shape["q_dtype_a"], - shape["q_dtype_w"], - shape["q_type"], - shape["use_g1u1"], - block_m, - device, + token = shape["token"] + model_dim = shape["model_dim"] + inter_dim = shape["inter_dim"] + dtype = shape["dtype"] + q_type = shape["q_type"] + q_dtype_w = shape["q_dtype_w"] + torch.manual_seed(0) + hidden = torch.randn((token, model_dim), dtype=dtype, device="cpu") / 10 + if shape["use_g1u1"]: + w1_dense = ( + torch.randn( + (routed_expert, inter_dim * 2, model_dim), + dtype=dtype, + device="cpu", + ) + / 10 + ) + else: + w1_dense = torch.randn( + (routed_expert, inter_dim, model_dim), + dtype=dtype, + device="cpu", + ) / 10 + w2_dense = torch.randn( + (routed_expert, model_dim, inter_dim), dtype=dtype, device="cpu" ) - w1 = FhmoeTuner._pad_dummy_expert(data["w1_qt"]) - w2 = FhmoeTuner._pad_dummy_expert(data["w2_qt"]) - w1_scale = FhmoeTuner._pad_dummy_expert(data["w1_scale"], routed_expert) - w2_scale = FhmoeTuner._pad_dummy_expert(data["w2_scale"], routed_expert) - - shared_device = w1.device + w1_shape = w1_dense.shape + w2_shape = w2_dense.shape + with torch.device("cpu"): + w1, w1_scale = FmoeTuner.weight_quant( + w1_dense, q_type, quant_dtype=q_dtype_w + ) + w2, w2_scale = FmoeTuner.weight_quant( + w2_dense, q_type, quant_dtype=q_dtype_w + ) + del w1_dense, w2_dense + if q_dtype_w is not dtypes.fp4x2: + w1 = w1.view(w1_shape) + w2 = w2.view(w2_shape) + else: + w1 = w1.view(w1_shape[0], w1_shape[1], w1_shape[2] // 2) + w2 = w2.view(w2_shape[0], w2_shape[1], w2_shape[2] // 2) + if TUNE_MOE_EXPERT_BALANCE: + score = torch.zeros( + (token, routed_expert), dtype=dtype, device="cpu" + ) + start_col = 0 + end_col = routed_topk + for token_id in range(token): + score[token_id, start_col:end_col] = 1.0 + start_col = end_col % routed_expert + end_col = start_col + routed_topk + else: + routing_seed = os.environ.get("TUNE_MOE_ROUTING_SEED") + if routing_seed is not None: + torch.manual_seed(int(routing_seed)) + score = torch.randn( + (token, routed_expert), dtype=dtype, device="cpu" + ) shared_w1_dense = ( torch.randn( - (1, 2 * shape["inter_dim"], shape["model_dim"]), - dtype=shape["dtype"], - device=shared_device, + (1, 2 * inter_dim, model_dim), dtype=dtype, device="cpu" ) / 10 ) shared_w2_dense = ( - torch.randn( - (1, shape["model_dim"], shape["inter_dim"]), - dtype=shape["dtype"], - device=shared_device, - ) + torch.randn((1, model_dim, inter_dim), dtype=dtype, device="cpu") / 10 ) - shared_w1, shared_w1_scale = FmoeTuner.weight_quant( - shared_w1_dense, shape["q_type"], quant_dtype=dtypes.fp8 - ) - shared_w2, shared_w2_scale = FmoeTuner.weight_quant( - shared_w2_dense, shape["q_type"], quant_dtype=dtypes.fp8 - ) + with torch.device("cpu"): + shared_w1, shared_w1_scale = FmoeTuner.weight_quant( + shared_w1_dense, q_type, quant_dtype=dtypes.fp8 + ) + shared_w2, shared_w2_scale = FmoeTuner.weight_quant( + shared_w2_dense, q_type, quant_dtype=dtypes.fp8 + ) del shared_w1_dense, shared_w2_dense + topk_weights, topk_ids = fused_topk( + hidden.to(device), score.to(device), routed_topk, True + ) + del score + topk_weights = topk_weights.cpu() + topk_ids = topk_ids.cpu() + + w1 = FhmoeTuner._pad_dummy_expert(w1) + w2 = FhmoeTuner._pad_dummy_expert(w2) + w1_scale = FhmoeTuner._pad_dummy_expert(w1_scale, routed_expert) + w2_scale = FhmoeTuner._pad_dummy_expert(w2_scale, routed_expert) shared_ids = torch.full( - (shape["token"], 1), + (token, 1), shared_expert_id, - dtype=data["topk_ids"].dtype, - device=data["topk_ids"].device, + dtype=topk_ids.dtype, + device=topk_ids.device, ) shared_weights = torch.ones( - (shape["token"], 1), - dtype=data["topk_weights"].dtype, - device=data["topk_weights"].device, + (token, 1), + dtype=topk_weights.dtype, + device=topk_weights.device, ) - topk_ids = torch.cat([data["topk_ids"], shared_ids], dim=1) - topk_weights = torch.cat([data["topk_weights"], shared_weights], dim=1) + topk_ids = torch.cat([topk_ids, shared_ids], dim=1) + topk_weights = torch.cat([topk_weights, shared_weights], dim=1) if shape["gate_mode"] == GateMode.INTERLEAVE: w1 = shuffle_weight_a16w4(w1, 16, True) @@ -6540,7 +6585,7 @@ def generate_fhmoe_data(shape, block_m=32, device="cuda"): w2.is_shuffled = True return { - "hidden": data["input"], + "hidden": hidden, "w1": w1, "w2": w2, "w1_scale": w1_scale, @@ -6551,8 +6596,8 @@ def generate_fhmoe_data(shape, block_m=32, device="cuda"): "shared_w2_scale": shared_w2_scale, "topk_ids": topk_ids, "topk_weights": topk_weights, - "a1_qt": data["a1_qt"], - "a1_scale": data["a1_scale"], + "a1_qt": hidden, + "a1_scale": None, "gate_mode": shape["gate_mode"], "hidden_pad": shape["hidden_pad"], "intermediate_pad": shape["intermediate_pad"], @@ -6596,25 +6641,90 @@ def _cuda(tensor): shared_expert_id=data["shared_expert_id"], ) + _FHMOE_MP_DATA_KEYS = ( + "hidden", + "w1", + "w2", + "topk_weights", + "topk_ids", + "w1_scale", + "w2_scale", + "shared_w1", + "shared_w2", + "shared_w1_scale", + "shared_w2_scale", + ) + + @staticmethod + def _run_fhmoe_mp( + hidden, + w1, + w2, + topk_weights, + topk_ids, + w1_scale, + w2_scale, + shared_w1, + shared_w2, + shared_w1_scale, + shared_w2_scale, + shape, + ): + return FhmoeTuner._run_fhmoe( + shape, + { + "hidden": hidden, + "w1": w1, + "w2": w2, + "topk_weights": topk_weights, + "topk_ids": topk_ids, + "w1_scale": w1_scale, + "w2_scale": w2_scale, + "shared_w1": shared_w1, + "shared_w2": shared_w2, + "shared_w1_scale": shared_w1_scale, + "shared_w2_scale": shared_w2_scale, + "shared_expert_id": shape["shared_expert_id"], + }, + ) + def tune(self, untunedf, tunedf, args): - del tunedf, args + del tunedf + tasks = [] for _, row in untunedf.iterrows(): shape = self._parse_fhmoe_row(row) - print("\nStart tuning FHMoE", [shape[k] for k in self.keys]) - data = self.generate_fhmoe_data(shape) - print( - f"[fhmoe] token={shape['token']} " - f"w1={tuple(data['w1'].shape)} {data['w1'].dtype} " - f"shared_w1={tuple(data['shared_w1'].shape)} {data['shared_w1'].dtype} " - f"topk_ids={tuple(data['topk_ids'].shape)}", - flush=True, - ) - out = self._run_fhmoe(shape, data) - print( - f"[fhmoe] fused_moe out={tuple(out.shape)} {out.dtype}", - flush=True, + info = tuple(shape[k] for k in self.keys) + block_m = 32 + print("\nStart tuning FHMoE", info, flush=True) + tasks.append( + ( + (info, "fused", "", block_m), + FhmoeTuner.generate_fhmoe_data, + (shape, block_m), + FhmoeTuner._run_fhmoe_mp, + (list(self._FHMOE_MP_DATA_KEYS), shape), + {}, + None, + ([],), + {}, + None, + 0.01, + 0.01, + ) ) - del data, out + if not tasks: + return [] + rets = mp_tuner( + tasks, + [(len(tasks), ())], + args.mp, + True, + False, + timeout=args.timeout, + verbose=args.verbose, + ) + for info, us, err in rets: + print(f"[fhmoe] mp_tuner us={us} err={err} info={info}", flush=True) return [] From 37317b96095603a744ae2335a166a03040826a36 Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Thu, 17 Sep 2026 13:56:18 +0200 Subject: [PATCH 03/14] =?UTF-8?q?[FlyDSL]=20Sweep=20fused=20FHMoE=20kernel?= =?UTF-8?q?=20pairs=20and=20emit=20native=20tuned=20CSV=20fused=5Fmoe()=20?= =?UTF-8?q?cannot=20inject=20a=20one-row=20candidate=20because=20supports?= =?UTF-8?q?=5Fdsv4=5Fi384=5Ffhmoe=20requires=20every=20padded=20M=20tier.?= =?UTF-8?q?=20Time=20each=20FlyDSL=20s1=C3=97s2=20pair=20through=20=5Ffuse?= =?UTF-8?q?d=5Fmoe=5Fimpl=20and=20keep=20the=20min-us=20row=20with=20block?= =?UTF-8?q?=5Fm/kernelName1/2.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../gemm_moe_tune.py | 302 +++++++++++++++--- 1 file changed, 259 insertions(+), 43 deletions(-) diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index 0a7079678f8..37b21a49689 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -26,7 +26,9 @@ dtype2str_dict, dtypes, ) +from aiter.fhmoe import _use_fhmoe_wrappers from aiter.fused_moe import ( + _fused_moe_impl, _mxfp4_a4w4_stage1_fw, _mxfp4_a4w4_stage2_fw, asm_stage1, @@ -34,6 +36,7 @@ cktile_moe_stage2, fused_moe, fused_topk, + get_padded_M, moe_sorting, torch_moe, torch_moe_stage1, @@ -6604,11 +6607,110 @@ def generate_fhmoe_data(shape, block_m=32, device="cuda"): "shared_expert_id": shared_expert_id, } + BLOCK_MS: ClassVar[tuple[int, ...]] = (16, 32, 64, 128) + + @staticmethod + def _csv_cell(value): + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, GateMode): + return f"GateMode.{value.name}" + if isinstance(value, ActivationType): + return f"ActivationType.{value.name}" + if isinstance(value, QuantType): + return f"QuantType.{value.name}" + return value + + @staticmethod + def _kernel_pairs(shape): + model_dim = shape["model_dim"] + inter_dim = shape["inter_dim"] + if shape["gate_mode"] == GateMode.INTERLEAVE: + a_dtype, b_dtype = "fp8", "fp4" + else: + a_dtype, b_dtype = "fp4", "fp4" + s1_all = get_flydsl_stage1_kernels(a_dtype, b_dtype, "bf16") + s2_all = get_flydsl_stage2_kernels(a_dtype, b_dtype, "bf16") + pairs = [] + for block_m in FhmoeTuner.BLOCK_MS: + s1_names = [] + for kname, kparams in s1_all.items(): + if kparams.get("tile_m") != block_m: + continue + k_batch = kparams.get("k_batch", 1) + k_wave = kparams.get("k_wave", 1) + tile_k = kparams["tile_k"] + if k_batch != 1 or model_dim % k_batch != 0: + continue + k_per_batch = model_dim // k_batch + if k_per_batch % tile_k != 0: + continue + if k_wave > 1 and ( + k_per_batch % k_wave != 0 or (k_per_batch // k_wave) % tile_k != 0 + ): + continue + s1_names.append(kname + ("_fp8" if a_dtype == "fp8" else "_fp4")) + s2_names = [] + for kname, kparams in s2_all.items(): + if kparams.get("tile_m") != block_m: + continue + # Same stage2 cuts as FmoeTuner.gen_flydsl_2stages_task: + # - tile_k must divide inter_dim (K); tile_k=256 on non-256 inter + # (e.g. 384) is parsed verbatim by the wrapper -> OOB/wrong out; + # - LDS is tile_m*(tile_k*2 + tile_n*4); over 160 KiB (gfx950 + # workgroup LDS) the build fails. + if inter_dim % kparams["tile_k"] != 0: + continue + lds = block_m * (kparams["tile_k"] * 2 + kparams["tile_n"] * 4) + if lds > 160 * 1024: + continue + s2_names.append(kname) + for kn1 in s1_names: + for kn2 in s2_names: + pairs.append((block_m, kn1, kn2)) + kernel_regex = os.environ.get("TUNE_MOE_KERNEL_REGEX") + if kernel_regex: + pattern = re.compile(kernel_regex) + pairs = [p for p in pairs if pattern.search(f"{p[1]} {p[2]}")] + return pairs + + @staticmethod + def _write_candidate_csv(path, shape, block_m, kn1, kn2): + gate_mode = shape["gate_mode"] + row = { + "gfx": shape["gfx"], + "cu_num": shape["cu_num"], + "token": get_padded_M(shape["token"]), + "model_dim": shape["model_dim"], + "inter_dim": shape["inter_dim"], + "expert": shape["expert"], + "topk": shape["topk"], + "act_type": FhmoeTuner._csv_cell(shape["act_type"]), + "dtype": FhmoeTuner._csv_cell(shape["dtype"]), + "q_dtype_a": FhmoeTuner._csv_cell(shape["q_dtype_a"]), + "q_dtype_w": FhmoeTuner._csv_cell(shape["q_dtype_w"]), + "q_type": FhmoeTuner._csv_cell(shape["q_type"]), + "use_g1u1": int(shape["use_g1u1"]), + "doweight_stage1": int(shape["doweight_stage1"]), + "shared_expert_id": shape["shared_expert_id"], + "hidden_pad": shape["hidden_pad"], + "intermediate_pad": shape["intermediate_pad"], + "gate_mode": FhmoeTuner._csv_cell(gate_mode), + "block_m": block_m, + "ksplit": 0, + "kernelName1": kn1, + "kernelName2": kn2, + } + pd.DataFrame([row]).to_csv(path, index=False) + @staticmethod - def _run_fhmoe(shape, data): + def _run_fhmoe(shape, data, kn1, kn2, block_m): gate_mode = shape["gate_mode"] if not isinstance(gate_mode, str): gate_mode = gate_mode.value + q_dtype_a = ( + dtypes.fp8 if gate_mode == GateMode.INTERLEAVE.value else dtypes.fp4x2 + ) def _cuda(tensor): if tensor is None or not torch.is_tensor(tensor): @@ -6619,27 +6721,45 @@ def _cuda(tensor): w2 = _cuda(data["w2"]) w1.is_shuffled = True w2.is_shuffled = True - return fused_moe( - _cuda(data["hidden"]), - w1, - w2, - _cuda(data["topk_weights"]), - _cuda(data["topk_ids"]), - activation=shape["act_type"], - quant_type=shape["q_type"], - doweight_stage1=bool(shape["doweight_stage1"]), - w1_scale=_cuda(data["w1_scale"]), - w2_scale=_cuda(data["w2_scale"]), - dtype=shape["dtype"], - hidden_pad=shape["hidden_pad"], - intermediate_pad=shape["intermediate_pad"], - gate_mode=gate_mode, - shared_w1=_cuda(data["shared_w1"]), - shared_w2=_cuda(data["shared_w2"]), - shared_w1_scale=_cuda(data["shared_w1_scale"]), - shared_w2_scale=_cuda(data["shared_w2_scale"]), - shared_expert_id=data["shared_expert_id"], - ) + cfg_fd, cfg_path = tempfile.mkstemp(suffix="_fhmoe_tune.csv") + os.close(cfg_fd) + FhmoeTuner._write_candidate_csv(cfg_path, shape, block_m, kn1, kn2) + import aiter.fused_moe as fmoe_mod + + try: + return _fused_moe_impl( + hidden_states=_cuda(data["hidden"]), + w1=w1, + w2=w2, + topk_weight=_cuda(data["topk_weights"]), + topk_ids=_cuda(data["topk_ids"]), + activation=shape["act_type"].value, + quant_type=shape["q_type"].value, + doweight_stage1=bool(shape["doweight_stage1"]), + w1_scale=_cuda(data["w1_scale"]), + w2_scale=_cuda(data["w2_scale"]), + block_size_M=block_m, + dtype=shape["dtype"], + hidden_pad=shape["hidden_pad"], + intermediate_pad=shape["intermediate_pad"], + gate_mode=gate_mode, + _q_dtype_a=q_dtype_a, + _metadata_transform=_use_fhmoe_wrappers, + _metadata_config_file=cfg_path, + _stage1_extra_args={ + "shared_w1": _cuda(data["shared_w1"]), + "shared_w1_scale": _cuda(data["shared_w1_scale"]), + "shared_expert_id": data["shared_expert_id"], + }, + _stage2_extra_args={ + "shared_w2": _cuda(data["shared_w2"]), + "shared_w2_scale": _cuda(data["shared_w2_scale"]), + "shared_expert_id": data["shared_expert_id"], + }, + ) + finally: + fmoe_mod.cfg_2stages_by_file.pop(cfg_path, None) + os.unlink(cfg_path) _FHMOE_MP_DATA_KEYS = ( "hidden", @@ -6669,6 +6789,9 @@ def _run_fhmoe_mp( shared_w1_scale, shared_w2_scale, shape, + kn1, + kn2, + block_m, ): return FhmoeTuner._run_fhmoe( shape, @@ -6686,6 +6809,9 @@ def _run_fhmoe_mp( "shared_w2_scale": shared_w2_scale, "shared_expert_id": shape["shared_expert_id"], }, + kn1, + kn2, + block_m, ) def tune(self, untunedf, tunedf, args): @@ -6694,24 +6820,28 @@ def tune(self, untunedf, tunedf, args): for _, row in untunedf.iterrows(): shape = self._parse_fhmoe_row(row) info = tuple(shape[k] for k in self.keys) - block_m = 32 - print("\nStart tuning FHMoE", info, flush=True) - tasks.append( - ( - (info, "fused", "", block_m), - FhmoeTuner.generate_fhmoe_data, - (shape, block_m), - FhmoeTuner._run_fhmoe_mp, - (list(self._FHMOE_MP_DATA_KEYS), shape), - {}, - None, - ([],), - {}, - None, - 0.01, - 0.01, - ) + pairs = self._kernel_pairs(shape) + print( + f"\nStart tuning FHMoE {info} candidates={len(pairs)}", + flush=True, ) + for block_m, kn1, kn2 in pairs: + tasks.append( + ( + (info, kn1, kn2, block_m), + FhmoeTuner.generate_fhmoe_data, + (shape,), + FhmoeTuner._run_fhmoe_mp, + (list(self._FHMOE_MP_DATA_KEYS), shape, kn1, kn2, block_m), + {}, + None, + ([],), + {}, + None, + 0.01, + 0.01, + ) + ) if not tasks: return [] rets = mp_tuner( @@ -6723,9 +6853,88 @@ def tune(self, untunedf, tunedf, args): timeout=args.timeout, verbose=args.verbose, ) - for info, us, err in rets: - print(f"[fhmoe] mp_tuner us={us} err={err} info={info}", flush=True) - return [] + return rets + + def post_process(self, results, args, topk=-1, fast_mode=False): + del topk, fast_mode + from collections import defaultdict + + grouped = defaultdict(list) + profile_rows = [] + for info, us, err in results: + key = tuple(info[0]) + kn1, kn2, block_m = info[1], info[2], info[3] + grouped[key].append((kn1, kn2, block_m, us, err)) + profile_rows.append( + { + **dict(zip(self.keys, key)), + "block_m": block_m, + "ksplit": 0, + "kernelName1": kn1, + "kernelName2": kn2, + "us": us, + "err": err, + } + ) + if args.profile_file: + profile_df = pd.DataFrame(profile_rows) + for col in ("act_type", "dtype", "q_dtype_a", "q_dtype_w", "q_type", "gate_mode"): + if col in profile_df.columns: + profile_df[col] = profile_df[col].map(self._csv_cell) + if os.path.exists(args.profile_file): + profile_df = pd.concat( + [pd.read_csv(args.profile_file), profile_df], ignore_index=True + ) + profile_df.to_csv(args.profile_file, index=False) + bests = [] + for key, rets in grouped.items(): + valid = [ + r + for r in rets + if r[3] not in (self.INVALID_TIME, self.INF_TIME) and r[3] >= 0 + ] + if not valid: + print(f"Tuning result for {key} is none", flush=True) + continue + kn1, kn2, block_m, us, err = min(valid, key=lambda r: r[3]) + row = dict(zip(self.keys, key)) + for col in ("act_type", "dtype", "q_dtype_a", "q_dtype_w", "q_type", "gate_mode"): + row[col] = self._csv_cell(row[col]) + row.update( + { + "block_m": block_m, + "ksplit": 0, + "kernelName1": kn1, + "kernelName2": kn2, + "us": us, + } + ) + print( + f"Tuning result for {key} is {block_m, kn1, kn2} {us} us err={err}", + flush=True, + ) + bests.append(row) + if not bests: + return pd.DataFrame(columns=self.columns) + return pd.DataFrame(bests, columns=self.columns) + + def result_to_csv(self, results, file, concat=False): + del concat + old_tunedf = self.get_tuned_gemm_list(file, self.columns) + for col in self.columns: + if col not in old_tunedf.columns: + old_tunedf[col] = "" + valid = results[ + (results["us"] != self.INVALID_TIME) & (results["us"] != self.INF_TIME) + ] + invalid = results[ + (results["us"] == self.INVALID_TIME) | (results["us"] == self.INF_TIME) + ] + resultdf = self.update_tunedf(old_tunedf, valid) + self.success = pd.concat([self.success, valid], ignore_index=True) + self.failed = pd.concat([self.failed, invalid], ignore_index=True) + resultdf = resultdf.astype(str).drop_duplicates(subset=self.keys, keep="last") + resultdf.to_csv(file, index=False) class Mxfp4FlydslTuner(FmoeTuner): @@ -7422,6 +7631,13 @@ def _reap(idx): "tflops", "bw", ] + fhmoe_result_list = [ + "block_m", + "ksplit", + "kernelName1", + "kernelName2", + "us", + ] grouped_result_list = [ "max_m", "tile_m", @@ -7460,7 +7676,7 @@ def _reap(idx): tuner = FhmoeTuner( "fhmoeTuner", fhmoe_key, - resultList, + fhmoe_result_list, "fhmoe tuner", ) else: From 71e77e13e02617556a603e0a068fde2fb6d0117c Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Thu, 17 Sep 2026 17:18:51 +0200 Subject: [PATCH 04/14] [FlyDSL] Wire FHMoE into op_tune.sh and replay fused_moe against a generated CSV python3 op_tests/test_fhmoe.py was exiting 0 without collecting pytest; the tuner times _fused_moe_impl, so a fresh process must prove public fused_moe honors AITER_CONFIG_FHMOE (raise on an empty DSV4 row, launch the generated kernel names otherwise). --- .github/scripts/op_tune.sh | 1 + op_tests/test_fhmoe.py | 4 + op_tests/tuning_tests/test_fhmoe_replay.py | 235 +++++++++++++++++++++ 3 files changed, 240 insertions(+) create mode 100644 op_tests/tuning_tests/test_fhmoe_replay.py diff --git a/.github/scripts/op_tune.sh b/.github/scripts/op_tune.sh index 1d4241e5331..43d0200a7fe 100755 --- a/.github/scripts/op_tune.sh +++ b/.github/scripts/op_tune.sh @@ -24,6 +24,7 @@ declare -a tune_jobs=( "ck_gemm_a8w8_blockscale_bpreshuffle:csrc/ck_gemm_a8w8_blockscale_bpreshuffle:op_tests/test_gemm_a8w8_blockscale.py:python3 csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_tune.py --preshuffle -i aiter/configs/a8w8_blockscale_bpreshuffle_untuned_gemm.csv -o aiter/configs/a8w8_blockscale_bpreshuffle_tuned_gemm.csv" "ck_gemm_a8w8_bpreshuffle:csrc/ck_gemm_a8w8_bpreshuffle:op_tests/test_gemm_a8w8.py:python3 csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py -i aiter/configs/a8w8_bpreshuffle_untuned_gemm.csv -o aiter/configs/a8w8_bpreshuffle_tuned_gemm.csv" "ck_gemm_moe_2stages_codegen:csrc/ck_gemm_moe_2stages_codegen:op_tests/test_moe.py:python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py -i aiter/configs/untuned_fmoe.csv -o aiter/configs/tuned_fmoe.csv" + "fhmoe:csrc/ck_gemm_moe_2stages_codegen:op_tests/test_fhmoe.py:python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py --fhmoe -i aiter/configs/untuned_fhmoe.csv -o aiter/configs/tuned_fhmoe.csv" #"ck_gemm_a4w4_blockscale:csrc/ck_gemm_a4w4_blockscale:op_tests/test_gemm_a4w4_blockscale.py:python3 csrc/ck_gemm_a4w4_blockscale/gemm_a4w4_blockscale_tune.py -i aiter/configs/a4w4_blockscale_untuned_gemm.csv -o aiter/configs/a4w4_blockscale_tuned_gemm.csv" ) diff --git a/op_tests/test_fhmoe.py b/op_tests/test_fhmoe.py index 6062c50110a..a8a11201580 100644 --- a/op_tests/test_fhmoe.py +++ b/op_tests/test_fhmoe.py @@ -1437,3 +1437,7 @@ def test_a4w4_routed_fp8_shared_heterogeneous_path( assert torch.isfinite(actual).all() error = _rel_l2(actual, routed_high + shared_high) assert error <= 5e-2, f"A4W4/FP8 heterogeneous FP32 error: {error:.3e}" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/op_tests/tuning_tests/test_fhmoe_replay.py b/op_tests/tuning_tests/test_fhmoe_replay.py new file mode 100644 index 00000000000..9d3b74a3318 --- /dev/null +++ b/op_tests/tuning_tests/test_fhmoe_replay.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Replay FHMoE's public operator against a generated config in a fresh process. + +The tuner worker times `_fused_moe_impl`. This checks that `fused_moe(...)` +with shared-expert arguments, in a new interpreter, reads only +`AITER_CONFIG_FHMOE` and either raises on an uncovered DSV4 I384 token or +launches the kernel names from that file. +""" + +from __future__ import annotations + +import csv +import os +import subprocess +import sys +import tempfile +import textwrap +import unittest + +AITER_ROOT = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +) + +# Distinct from shipped tuned_fhmoe.csv token=1 (โ€ฆ_atomic without _persist). +_REPLAY_KN1 = "flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w4_gui_kw4_fp8" +_REPLAY_KN2 = "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic_persist" + +_NATIVE_FIELDS = [ + "gfx", + "cu_num", + "token", + "model_dim", + "inter_dim", + "expert", + "topk", + "shared_expert_id", + "act_type", + "dtype", + "q_dtype_a", + "q_dtype_w", + "q_type", + "use_g1u1", + "doweight_stage1", + "hidden_pad", + "intermediate_pad", + "gate_mode", + "block_m", + "ksplit", + "kernelName1", + "kernelName2", +] + + +def _gpu_available(): + try: + import torch + + return torch.cuda.is_available() and torch.cuda.device_count() > 0 + except ImportError: + return False + + +def _write_header_only(path): + with open(path, "w", newline="") as handle: + csv.DictWriter(handle, fieldnames=_NATIVE_FIELDS).writeheader() + + +def _write_generated_row(path, gfx, cu_num): + row = { + "gfx": gfx, + "cu_num": cu_num, + "token": 1, + "model_dim": 7168, + "inter_dim": 384, + "expert": 385, + "topk": 7, + "shared_expert_id": 384, + "act_type": "ActivationType.Silu", + "dtype": "torch.bfloat16", + "q_dtype_a": "torch.float8_e4m3fn", + "q_dtype_w": "torch.float4_e2m1fn_x2", + "q_type": "QuantType.per_1x32", + "use_g1u1": 1, + "doweight_stage1": 0, + "hidden_pad": 0, + "intermediate_pad": 0, + "gate_mode": "GateMode.INTERLEAVE", + "block_m": 32, + "ksplit": 0, + "kernelName1": _REPLAY_KN1, + "kernelName2": _REPLAY_KN2, + } + with open(path, "w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=_NATIVE_FIELDS) + writer.writeheader() + writer.writerow(row) + + +def _child_script(): + return textwrap.dedent( + f"""\ + import sys + sys.path.insert( + 0, + {os.path.join(AITER_ROOT, "csrc", "ck_gemm_moe_2stages_codegen")!r}, + ) + import torch + from aiter import ActivationType, QuantType, dtypes + from aiter.fused_moe import fused_moe + from aiter.jit.utils.chip_info import get_cu_num, get_gfx_runtime + from aiter.ops.flydsl.moe_common import GateMode + from gemm_moe_tune import FhmoeTuner + + shape = {{ + "gfx": get_gfx_runtime(), + "cu_num": int(get_cu_num()), + "token": 1, + "model_dim": 7168, + "inter_dim": 384, + "expert": 385, + "topk": 7, + "shared_expert_id": 384, + "act_type": ActivationType.Silu, + "dtype": torch.bfloat16, + "q_dtype_a": dtypes.fp8, + "q_dtype_w": dtypes.fp4x2, + "q_type": QuantType.per_1x32, + "use_g1u1": 1, + "doweight_stage1": 0, + "hidden_pad": 0, + "intermediate_pad": 0, + "gate_mode": GateMode.INTERLEAVE, + }} + data = FhmoeTuner.generate_fhmoe_data(shape) + + def _cuda(tensor): + if tensor is None or not torch.is_tensor(tensor): + return tensor + return tensor.cuda() + + w1 = _cuda(data["w1"]) + w2 = _cuda(data["w2"]) + w1.is_shuffled = True + w2.is_shuffled = True + try: + out = fused_moe( + _cuda(data["hidden"]), + w1, + w2, + _cuda(data["topk_weights"]), + _cuda(data["topk_ids"]), + activation=ActivationType.Silu, + quant_type=QuantType.per_1x32, + w1_scale=_cuda(data["w1_scale"]), + w2_scale=_cuda(data["w2_scale"]), + dtype=torch.bfloat16, + hidden_pad=0, + intermediate_pad=0, + gate_mode=GateMode.INTERLEAVE.value, + shared_w1=_cuda(data["shared_w1"]), + shared_w2=_cuda(data["shared_w2"]), + shared_w1_scale=_cuda(data["shared_w1_scale"]), + shared_w2_scale=_cuda(data["shared_w2_scale"]), + shared_expert_id=384, + ) + except NotImplementedError as exc: + print("REPLAY_RAISE", str(exc)) + raise SystemExit(0) + print("REPLAY_OK", tuple(out.shape)) + """ + ) + + +@unittest.skipUnless(_gpu_available(), "No GPU available") +class TestFhmoeReplay(unittest.TestCase): + def setUp(self): + from aiter.jit.utils.chip_info import get_gfx + + if get_gfx() != "gfx950": + self.skipTest("FHMoE replay requires gfx950") + + def _run_public_op(self, csv_path, timeout=600): + script_path = os.path.join(os.path.dirname(csv_path), "fhmoe_replay_child.py") + with open(script_path, "w") as handle: + handle.write(_child_script()) + env = os.environ.copy() + env["AITER_CONFIG_FHMOE"] = csv_path + env["PYTHONPATH"] = AITER_ROOT + os.pathsep + env.get("PYTHONPATH", "") + env.pop("AITER_BYPASS_TUNE_CONFIG", None) + return subprocess.run( + [sys.executable, script_path], + capture_output=True, + text=True, + timeout=timeout, + cwd=AITER_ROOT, + env=env, + check=False, + ) + + def test_uncovered_token_raises_in_fresh_process(self): + with tempfile.TemporaryDirectory() as tmp: + csv_path = os.path.join(tmp, "untuned_empty.csv") + _write_header_only(csv_path) + result = self._run_public_op(csv_path) + output = result.stdout + result.stderr + self.assertEqual( + result.returncode, + 0, + f"uncovered replay child failed\n{output[-4000:]}", + ) + self.assertIn("REPLAY_RAISE", result.stdout) + self.assertIn("does not cover this DSV4 I384", result.stdout) + + def test_generated_csv_selected_by_public_fused_moe(self): + from aiter.jit.utils.chip_info import get_cu_num, get_gfx_runtime + + with tempfile.TemporaryDirectory() as tmp: + csv_path = os.path.join(tmp, "generated_fhmoe.csv") + _write_generated_row(csv_path, get_gfx_runtime(), int(get_cu_num())) + result = self._run_public_op(csv_path) + output = result.stdout + result.stderr + self.assertEqual( + result.returncode, + 0, + f"generated replay child failed\n{output[-4000:]}", + ) + self.assertIn("REPLAY_OK", result.stdout) + self.assertIn(_REPLAY_KN1, output) + self.assertIn(_REPLAY_KN2, output) + self.assertNotIn("REPLAY_RAISE", result.stdout) + + +if __name__ == "__main__": + unittest.main() From a7e7f827a5bb5224955d265e5c6fd37d0c286f11 Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Fri, 18 Sep 2026 15:46:07 +0200 Subject: [PATCH 05/14] [FlyDSL] Gate FHMoE tune on torch cosine and time public fused_moe Reject fast-wrong FlyDSL pairs, retune shipped keys with --all, replay the producer CSV through fused_moe, and measure --run_config with shared FP8. --- .github/scripts/op_tune.sh | 2 +- .../gemm_moe_tune.py | 281 +++++++++++++++++- .../tuning_tests/test_fhmoe_accuracy_gate.py | 205 +++++++++++++ op_tests/tuning_tests/test_fhmoe_replay.py | 266 ++++++++++++++--- 4 files changed, 697 insertions(+), 57 deletions(-) create mode 100644 op_tests/tuning_tests/test_fhmoe_accuracy_gate.py diff --git a/.github/scripts/op_tune.sh b/.github/scripts/op_tune.sh index 43d0200a7fe..20483d7f948 100755 --- a/.github/scripts/op_tune.sh +++ b/.github/scripts/op_tune.sh @@ -24,7 +24,7 @@ declare -a tune_jobs=( "ck_gemm_a8w8_blockscale_bpreshuffle:csrc/ck_gemm_a8w8_blockscale_bpreshuffle:op_tests/test_gemm_a8w8_blockscale.py:python3 csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_tune.py --preshuffle -i aiter/configs/a8w8_blockscale_bpreshuffle_untuned_gemm.csv -o aiter/configs/a8w8_blockscale_bpreshuffle_tuned_gemm.csv" "ck_gemm_a8w8_bpreshuffle:csrc/ck_gemm_a8w8_bpreshuffle:op_tests/test_gemm_a8w8.py:python3 csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py -i aiter/configs/a8w8_bpreshuffle_untuned_gemm.csv -o aiter/configs/a8w8_bpreshuffle_tuned_gemm.csv" "ck_gemm_moe_2stages_codegen:csrc/ck_gemm_moe_2stages_codegen:op_tests/test_moe.py:python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py -i aiter/configs/untuned_fmoe.csv -o aiter/configs/tuned_fmoe.csv" - "fhmoe:csrc/ck_gemm_moe_2stages_codegen:op_tests/test_fhmoe.py:python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py --fhmoe -i aiter/configs/untuned_fhmoe.csv -o aiter/configs/tuned_fhmoe.csv" + "fhmoe:csrc/ck_gemm_moe_2stages_codegen:op_tests/test_fhmoe.py:python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py --fhmoe --all -i aiter/configs/untuned_fhmoe.csv -o aiter/configs/tuned_fhmoe.csv" #"ck_gemm_a4w4_blockscale:csrc/ck_gemm_a4w4_blockscale:op_tests/test_gemm_a4w4_blockscale.py:python3 csrc/ck_gemm_a4w4_blockscale/gemm_a4w4_blockscale_tune.py -i aiter/configs/a4w4_blockscale_untuned_gemm.csv -o aiter/configs/a4w4_blockscale_tuned_gemm.csv" ) diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index 37b21a49689..4250af66e02 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -6419,6 +6419,13 @@ class FhmoeTuner(FmoeTuner): "config_env_name": "AITER_CONFIG_FHMOE", } + def _clear_op_caches(self): + super()._clear_op_caches() + import aiter.fused_moe as fmoe_module + + if hasattr(fmoe_module, "cfg_2stages_by_file"): + fmoe_module.cfg_2stages_by_file.clear() + @staticmethod def _parse_fhmoe_row(row): return { @@ -6547,6 +6554,12 @@ def generate_fhmoe_data(shape, block_m=32, device="cuda"): topk_weights = topk_weights.cpu() topk_ids = topk_ids.cpu() + # Unshuffled copies for the torch reference. Shuffle below is kernel ABI. + w1_qt, w2_qt = w1, w2 + w1_scale_qt, w2_scale_qt = w1_scale, w2_scale + shared_w1_qt, shared_w2_qt = shared_w1, shared_w2 + shared_w1_scale_qt, shared_w2_scale_qt = shared_w1_scale, shared_w2_scale + w1 = FhmoeTuner._pad_dummy_expert(w1) w2 = FhmoeTuner._pad_dummy_expert(w2) w1_scale = FhmoeTuner._pad_dummy_expert(w1_scale, routed_expert) @@ -6593,10 +6606,18 @@ def generate_fhmoe_data(shape, block_m=32, device="cuda"): "w2": w2, "w1_scale": w1_scale, "w2_scale": w2_scale, + "w1_qt": w1_qt, + "w2_qt": w2_qt, + "w1_scale_qt": w1_scale_qt, + "w2_scale_qt": w2_scale_qt, "shared_w1": shared_w1, "shared_w2": shared_w2, "shared_w1_scale": shared_w1_scale, "shared_w2_scale": shared_w2_scale, + "shared_w1_qt": shared_w1_qt, + "shared_w2_qt": shared_w2_qt, + "shared_w1_scale_qt": shared_w1_scale_qt, + "shared_w2_scale_qt": shared_w2_scale_qt, "topk_ids": topk_ids, "topk_weights": topk_weights, "a1_qt": hidden, @@ -6774,6 +6795,19 @@ def _cuda(tensor): "shared_w1_scale", "shared_w2_scale", ) + _FHMOE_REF_DATA_KEYS = ( + "hidden", + "w1_qt", + "w2_qt", + "w1_scale_qt", + "w2_scale_qt", + "shared_w1_qt", + "shared_w2_qt", + "shared_w1_scale_qt", + "shared_w2_scale_qt", + "topk_weights", + "topk_ids", + ) @staticmethod def _run_fhmoe_mp( @@ -6814,9 +6848,119 @@ def _run_fhmoe_mp( block_m, ) - def tune(self, untunedf, tunedf, args): - del tunedf + @staticmethod + def _to_ref_device(tensor): + if tensor is None or not torch.is_tensor(tensor): + return tensor + if tensor.is_cuda: + return tensor + if torch.cuda.is_available(): + return tensor.cuda() + return tensor + + @staticmethod + def run_torch_fhmoe( + hidden, + w1_qt, + w2_qt, + w1_scale_qt, + w2_scale_qt, + shared_w1_qt, + shared_w2_qt, + shared_w1_scale_qt, + shared_w2_scale_qt, + topk_weights, + topk_ids, + shape, + ): + """Homogeneous torch refs for routed FP4 and shared FP8, then summed. + + Matches the tuner launch: same tensors, no swiglu_limit (the candidate + path does not pass one). Activation quant is fused in the kernel; + FmoeTuner's a8w4 ref also leaves a1_scale=None. + """ + hidden = FhmoeTuner._to_ref_device(hidden) + w1_qt = FhmoeTuner._to_ref_device(w1_qt) + w2_qt = FhmoeTuner._to_ref_device(w2_qt) + w1_scale_qt = FhmoeTuner._to_ref_device(w1_scale_qt) + w2_scale_qt = FhmoeTuner._to_ref_device(w2_scale_qt) + shared_w1_qt = FhmoeTuner._to_ref_device(shared_w1_qt) + shared_w2_qt = FhmoeTuner._to_ref_device(shared_w2_qt) + shared_w1_scale_qt = FhmoeTuner._to_ref_device(shared_w1_scale_qt) + shared_w2_scale_qt = FhmoeTuner._to_ref_device(shared_w2_scale_qt) + topk_weights = FhmoeTuner._to_ref_device(topk_weights) + topk_ids = FhmoeTuner._to_ref_device(topk_ids) + + routed_ids = topk_ids[:, :-1] + routed_w = topk_weights[:, :-1] + shared_ids = torch.zeros( + (topk_ids.shape[0], 1), + dtype=topk_ids.dtype, + device=topk_ids.device, + ) + shared_w = topk_weights[:, -1:] + dtype = shape["dtype"] + activation = shape["act_type"] + quant_type = shape["q_type"] + doweight_stage1 = bool(shape["doweight_stage1"]) + + ref1 = FmoeTuner.run_torch_moe_stage1( + hidden, + w1_qt, + w2_qt, + routed_w, + routed_ids, + None, + w1_scale_qt, + dtype=dtype, + activation=activation, + quant_type=quant_type, + doweight_stage1=doweight_stage1, + topk=routed_ids.shape[1], + ) + routed = FmoeTuner.run_torch_moe_stage2( + ref1, + w1_qt, + w2_qt, + routed_w, + routed_ids, + None, + w2_scale_qt, + dtype=dtype, + quant_type=quant_type, + doweight_stage1=doweight_stage1, + ) + shared_ref1 = FmoeTuner.run_torch_moe_stage1( + hidden, + shared_w1_qt, + shared_w2_qt, + shared_w, + shared_ids, + None, + shared_w1_scale_qt, + dtype=dtype, + activation=activation, + quant_type=quant_type, + doweight_stage1=doweight_stage1, + topk=1, + ) + shared = FmoeTuner.run_torch_moe_stage2( + shared_ref1, + shared_w1_qt, + shared_w2_qt, + shared_w, + shared_ids, + None, + shared_w2_scale_qt, + dtype=dtype, + quant_type=quant_type, + doweight_stage1=doweight_stage1, + ) + return routed + shared + + def _fhmoe_tune_tasks(self, untunedf): tasks = [] + in_datas = [] for _, row in untunedf.iterrows(): shape = self._parse_fhmoe_row(row) info = tuple(shape[k] for k in self.keys) @@ -6825,6 +6969,8 @@ def tune(self, untunedf, tunedf, args): f"\nStart tuning FHMoE {info} candidates={len(pairs)}", flush=True, ) + if not pairs: + continue for block_m, kn1, kn2 in pairs: tasks.append( ( @@ -6834,26 +6980,136 @@ def tune(self, untunedf, tunedf, args): FhmoeTuner._run_fhmoe_mp, (list(self._FHMOE_MP_DATA_KEYS), shape, kn1, kn2, block_m), {}, - None, - ([],), + FhmoeTuner.run_torch_fhmoe, + (list(self._FHMOE_REF_DATA_KEYS), shape), {}, None, 0.01, 0.01, + cosine_diff_compare, + ) + ) + in_datas.append((len(pairs), ())) + return tasks, in_datas + + def run_config(self, args): + """Public fused_moe + shared FP8, not FmoeTuner's homogeneous path.""" + from aiter.test_common import run_perftest + + results = [] + for _, row in self.untunedf.iterrows(): + shape = self._parse_fhmoe_row(row) + shape_str = ( + f"(token={shape['token']}, dim={shape['model_dim']}, " + f"inter={shape['inter_dim']}, E={shape['expert']}, " + f"topk={shape['topk']}, shared={shape['shared_expert_id']}, " + f"{self._csv_cell(shape['act_type'])}, " + f"{self._csv_cell(shape['dtype'])}, " + f"{self._csv_cell(shape['q_dtype_a'])}, " + f"{self._csv_cell(shape['q_dtype_w'])}, " + f"{self._csv_cell(shape['q_type'])}, " + f"{self._csv_cell(shape['gate_mode'])})" + ) + allowed_err_ratio, allowed_err_ratio_desc = ( + self._get_run_config_err_ratio_limit(row, args) + ) + kernel_us = None + if "us" in row and pd.notna(row["us"]): + try: + kernel_us = float(row["us"]) + except (TypeError, ValueError): + kernel_us = None + try: + data = self.generate_fhmoe_data(shape) + w1 = self._to_ref_device(data["w1"]) + w2 = self._to_ref_device(data["w2"]) + w1.is_shuffled = True + w2.is_shuffled = True + gate_mode = shape["gate_mode"] + if not isinstance(gate_mode, str): + gate_mode = gate_mode.value + out, us = run_perftest( + fused_moe, + self._to_ref_device(data["hidden"]), + w1, + w2, + self._to_ref_device(data["topk_weights"]), + self._to_ref_device(data["topk_ids"]), + activation=shape["act_type"], + quant_type=shape["q_type"], + doweight_stage1=bool(shape["doweight_stage1"]), + w1_scale=self._to_ref_device(data["w1_scale"]), + w2_scale=self._to_ref_device(data["w2_scale"]), + dtype=shape["dtype"], + hidden_pad=shape["hidden_pad"], + intermediate_pad=shape["intermediate_pad"], + gate_mode=gate_mode, + shared_w1=self._to_ref_device(data["shared_w1"]), + shared_w2=self._to_ref_device(data["shared_w2"]), + shared_w1_scale=self._to_ref_device(data["shared_w1_scale"]), + shared_w2_scale=self._to_ref_device(data["shared_w2_scale"]), + shared_expert_id=data["shared_expert_id"], + num_warmup=args.warmup, + num_iters=args.iters, + ) + ref = self.run_torch_fhmoe( + *[data[k] for k in self._FHMOE_REF_DATA_KEYS], + shape, + ) + if out.count_nonzero() == 0 and ref.count_nonzero() > 0: + diag = tensor_compare_diagnostics(ref, out) + status = ( + "error:output is all zeros (kernel produced no output); " + f"{diag}" + ) + else: + err_ratio = cosine_diff_compare( + ref, out, msg=f"run_config {shape_str}" ) + if err_ratio <= allowed_err_ratio: + status = "ok" + else: + diag = tensor_compare_diagnostics(ref, out) + status = ( + f"mismatch:cosine_diff={err_ratio:.6g}" + f"(>{allowed_err_ratio_desc}); {diag}" + ) + results.append( + { + "shape": shape_str, + "e2e_us": us, + "kernel_us": kernel_us, + "status": status, + } ) + except Exception as e: # noqa: BLE001 + results.append( + { + "shape": shape_str, + "e2e_us": -1, + "kernel_us": kernel_us, + "status": f"error:{e}", + } + ) + finally: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return results + + def tune(self, untunedf, tunedf, args): + del tunedf + tasks, in_datas = self._fhmoe_tune_tasks(untunedf) if not tasks: return [] - rets = mp_tuner( + return mp_tuner( tasks, - [(len(tasks), ())], + in_datas, args.mp, True, - False, + True, timeout=args.timeout, verbose=args.verbose, ) - return rets def post_process(self, results, args, topk=-1, fast_mode=False): del topk, fast_mode @@ -6891,10 +7147,15 @@ def post_process(self, results, args, topk=-1, fast_mode=False): valid = [ r for r in rets - if r[3] not in (self.INVALID_TIME, self.INF_TIME) and r[3] >= 0 + if r[3] not in (self.INVALID_TIME, self.INF_TIME) + and r[3] >= 0 + and r[4] <= args.errRatio ] if not valid: - print(f"Tuning result for {key} is none", flush=True) + print( + f"Tuning result for {key} is none, please check errRatio={args.errRatio}", + flush=True, + ) continue kn1, kn2, block_m, us, err = min(valid, key=lambda r: r[3]) row = dict(zip(self.keys, key)) diff --git a/op_tests/tuning_tests/test_fhmoe_accuracy_gate.py b/op_tests/tuning_tests/test_fhmoe_accuracy_gate.py new file mode 100644 index 00000000000..fa3d8070c85 --- /dev/null +++ b/op_tests/tuning_tests/test_fhmoe_accuracy_gate.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""FHMoE tuner must not pick a fast-but-wrong kernel pair.""" + +from __future__ import annotations + +import argparse +import os +import sys +import unittest +from unittest.mock import patch + +import pandas as pd + +AITER_ROOT = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +) +if AITER_ROOT not in sys.path: + sys.path.insert(0, AITER_ROOT) + +try: + from csrc.ck_gemm_moe_2stages_codegen.gemm_moe_tune import ( # noqa: E402 + FhmoeTuner, + FmoeTuner, + cosine_diff_compare, + ) +except ImportError as _import_error: # torch / aiter not in this interpreter + FhmoeTuner = None + FmoeTuner = None + cosine_diff_compare = None + _IMPORT_ERROR = _import_error +else: + _IMPORT_ERROR = None + +_FHMOE_KEYS = [ + "gfx", + "cu_num", + "token", + "model_dim", + "inter_dim", + "expert", + "topk", + "act_type", + "dtype", + "q_dtype_a", + "q_dtype_w", + "q_type", + "use_g1u1", + "doweight_stage1", + "shared_expert_id", + "hidden_pad", + "intermediate_pad", + "gate_mode", +] +_FHMOE_RESULTS = [ + "block_m", + "ksplit", + "kernelName1", + "kernelName2", + "us", +] + + +def _tuner(): + return FhmoeTuner("fhmoeTuner", _FHMOE_KEYS, _FHMOE_RESULTS, "fhmoe tuner") + + +def _info(): + return ( + "gfx950", + 256, + 1, + 7168, + 384, + 385, + 7, + "ActivationType.Silu", + "torch.bfloat16", + "torch.float8_e4m3fn", + "torch.float4_e2m1fn_x2", + "QuantType.per_1x32", + 1, + 0, + 384, + 0, + 0, + "GateMode.INTERLEAVE", + ) + + +def _untuned_row(): + return { + "gfx": "gfx950", + "cu_num": 256, + "token": 1, + "model_dim": 7168, + "inter_dim": 384, + "expert": 385, + "topk": 7, + "shared_expert_id": 384, + "act_type": "ActivationType.Silu", + "dtype": "torch.bfloat16", + "q_dtype_a": "torch.float8_e4m3fn", + "q_dtype_w": "torch.float4_e2m1fn_x2", + "q_type": "QuantType.per_1x32", + "use_g1u1": 1, + "doweight_stage1": 0, + "hidden_pad": 0, + "intermediate_pad": 0, + "gate_mode": "GateMode.INTERLEAVE", + } + + +@unittest.skipIf(FhmoeTuner is None, f"gemm_moe_tune import failed: {_IMPORT_ERROR}") +class TestFhmoeAccuracyGate(unittest.TestCase): + def test_tune_tasks_attach_torch_ref_and_cosine(self): + tuner = _tuner() + df = pd.DataFrame([_untuned_row()]) + with patch.object( + FhmoeTuner, "_kernel_pairs", return_value=[(32, "kn1", "kn2")] + ): + tasks, in_datas = tuner._fhmoe_tune_tasks(df) + self.assertEqual(len(tasks), 1) + self.assertEqual(in_datas, [(1, ())]) + task = tasks[0] + self.assertIs(task[6], FhmoeTuner.run_torch_fhmoe) + self.assertEqual(task[7][0], list(FhmoeTuner._FHMOE_REF_DATA_KEYS)) + self.assertIs(task[12], cosine_diff_compare) + + def test_post_process_rejects_fast_wrong_pair(self): + tuner = _tuner() + info = _info() + results = [ + ((info, "fast_wrong1", "fast_wrong2", 32), 10.0, 0.9), + ((info, "slow_ok1", "slow_ok2", 64), 50.0, 0.01), + ] + args = argparse.Namespace(errRatio=0.1, profile_file="") + out = tuner.post_process(results, args) + self.assertEqual(len(out), 1) + self.assertEqual(out.iloc[0]["kernelName1"], "slow_ok1") + self.assertEqual(out.iloc[0]["kernelName2"], "slow_ok2") + self.assertEqual(int(out.iloc[0]["block_m"]), 64) + + def test_post_process_drops_shape_when_all_pairs_fail_err(self): + tuner = _tuner() + info = _info() + results = [((info, "bad1", "bad2", 32), 10.0, 0.5)] + args = argparse.Namespace(errRatio=0.1, profile_file="") + out = tuner.post_process(results, args) + self.assertTrue(out.empty) + + def test_run_config_is_not_parent_fmoe(self): + self.assertIsNot(FhmoeTuner.run_config, FmoeTuner.run_config) + + def test_run_config_calls_public_fused_moe_with_shared(self): + import torch + + from csrc.ck_gemm_moe_2stages_codegen import gemm_moe_tune as tune_mod + + tuner = _tuner() + tuner.untunedf = pd.DataFrame( + [{**_untuned_row(), "gfx": "gfx950", "cu_num": 256}] + ) + args = argparse.Namespace(warmup=1, iters=1, errRatio=0.1) + fake = torch.ones((1, 8)) + data = { + "hidden": fake, + "w1": fake, + "w2": fake, + "w1_scale": fake, + "w2_scale": fake, + "shared_w1": fake, + "shared_w2": fake, + "shared_w1_scale": fake, + "shared_w2_scale": fake, + "topk_weights": fake, + "topk_ids": fake, + "w1_qt": fake, + "w2_qt": fake, + "w1_scale_qt": fake, + "w2_scale_qt": fake, + "shared_w1_qt": fake, + "shared_w2_qt": fake, + "shared_w1_scale_qt": fake, + "shared_w2_scale_qt": fake, + "shared_expert_id": 384, + } + with patch.object(FhmoeTuner, "generate_fhmoe_data", return_value=data), patch.object( + FhmoeTuner, "run_torch_fhmoe", return_value=fake + ), patch( + "aiter.test_common.run_perftest", return_value=(fake, 40.0) + ) as perf: + out = tuner.run_config(args) + self.assertEqual(len(out), 1) + self.assertEqual(out[0]["status"], "ok") + self.assertEqual(out[0]["e2e_us"], 40.0) + self.assertIn("shared=", out[0]["shape"]) + self.assertIs(perf.call_args[0][0], tune_mod.fused_moe) + self.assertIn("shared_w1", perf.call_args[1]) + self.assertIn("shared_w2", perf.call_args[1]) + self.assertEqual(perf.call_args[1]["shared_expert_id"], 384) + + +if __name__ == "__main__": + unittest.main() diff --git a/op_tests/tuning_tests/test_fhmoe_replay.py b/op_tests/tuning_tests/test_fhmoe_replay.py index 9d3b74a3318..6323411094d 100644 --- a/op_tests/tuning_tests/test_fhmoe_replay.py +++ b/op_tests/tuning_tests/test_fhmoe_replay.py @@ -1,16 +1,18 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Replay FHMoE's public operator against a generated config in a fresh process. +"""Replay FHMoE's public operator against tuner-written config in a fresh process. The tuner worker times `_fused_moe_impl`. This checks that `fused_moe(...)` with shared-expert arguments, in a new interpreter, reads only `AITER_CONFIG_FHMOE` and either raises on an uncovered DSV4 I384 token or -launches the kernel names from that file. +launches the kernel names from a CSV that `FhmoeTuner` actually wrote. +A planted row is not the producer contract. """ from __future__ import annotations import csv +import glob import os import subprocess import sys @@ -21,10 +23,21 @@ AITER_ROOT = os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) +_TUNE_SCRIPT = os.path.join( + AITER_ROOT, "csrc", "ck_gemm_moe_2stages_codegen", "gemm_moe_tune.py" +) -# Distinct from shipped tuned_fhmoe.csv token=1 (โ€ฆ_atomic without _persist). -_REPLAY_KN1 = "flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w4_gui_kw4_fp8" -_REPLAY_KN2 = "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic_persist" +# Two kn2 suffixes, both distinct from shipped token=1 (`โ€ฆ_atomic` with no suffix). +# TUNE_MOE_KERNEL_REGEX matches f"{kn1} {kn2}". +_TUNE_KN1 = "flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w4_gui_kw4_fp8" +_TUNE_KN2_A = "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic_bnt2" +_TUNE_KN2_B = "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic_persist" +_TUNE_KN2S = frozenset((_TUNE_KN2_A, _TUNE_KN2_B)) +_SHIPPED_TOKEN1_KN2 = "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic" +_KERNEL_REGEX = ( + rf"{_TUNE_KN1} {_TUNE_KN2_A}$|{_TUNE_KN1} {_TUNE_KN2_B}$" +) +_REPLAY_TOKENS = (1, 16) _NATIVE_FIELDS = [ "gfx", @@ -51,6 +64,25 @@ "kernelName2", ] +_UNTUNED_FIELDS = [ + "token", + "model_dim", + "inter_dim", + "expert", + "topk", + "shared_expert_id", + "act_type", + "dtype", + "q_dtype_a", + "q_dtype_w", + "q_type", + "use_g1u1", + "doweight_stage1", + "hidden_pad", + "intermediate_pad", + "gate_mode", +] + def _gpu_available(): try: @@ -66,38 +98,50 @@ def _write_header_only(path): csv.DictWriter(handle, fieldnames=_NATIVE_FIELDS).writeheader() -def _write_generated_row(path, gfx, cu_num): - row = { - "gfx": gfx, - "cu_num": cu_num, - "token": 1, - "model_dim": 7168, - "inter_dim": 384, - "expert": 385, - "topk": 7, - "shared_expert_id": 384, - "act_type": "ActivationType.Silu", - "dtype": "torch.bfloat16", - "q_dtype_a": "torch.float8_e4m3fn", - "q_dtype_w": "torch.float4_e2m1fn_x2", - "q_type": "QuantType.per_1x32", - "use_g1u1": 1, - "doweight_stage1": 0, - "hidden_pad": 0, - "intermediate_pad": 0, - "gate_mode": "GateMode.INTERLEAVE", - "block_m": 32, - "ksplit": 0, - "kernelName1": _REPLAY_KN1, - "kernelName2": _REPLAY_KN2, - } +def _write_untuned_rows(path, tokens): with open(path, "w", newline="") as handle: - writer = csv.DictWriter(handle, fieldnames=_NATIVE_FIELDS) + writer = csv.DictWriter(handle, fieldnames=_UNTUNED_FIELDS) writer.writeheader() - writer.writerow(row) + for token in tokens: + writer.writerow( + { + "token": token, + "model_dim": 7168, + "inter_dim": 384, + "expert": 385, + "topk": 7, + "shared_expert_id": 384, + "act_type": "ActivationType.Silu", + "dtype": "torch.bfloat16", + "q_dtype_a": "torch.float8_e4m3fn", + "q_dtype_w": "torch.float4_e2m1fn_x2", + "q_type": "QuantType.per_1x32", + "use_g1u1": 1, + "doweight_stage1": 0, + "hidden_pad": 0, + "intermediate_pad": 0, + "gate_mode": "GateMode.INTERLEAVE", + } + ) -def _child_script(): +def _cleanup_stale_lock_files(): + build_dir = os.path.join(AITER_ROOT, "aiter", "jit", "build") + if not os.path.isdir(build_dir): + return + for pattern in ( + os.path.join(build_dir, "lock_*"), + os.path.join(build_dir, "*", "build", "lock"), + os.path.join(build_dir, "lock_3rdparty_*"), + ): + for lock_file in glob.glob(pattern): + try: + os.remove(lock_file) + except OSError: + pass + + +def _child_script(token): return textwrap.dedent( f"""\ import sys @@ -115,7 +159,7 @@ def _child_script(): shape = {{ "gfx": get_gfx_runtime(), "cu_num": int(get_cu_num()), - "token": 1, + "token": {int(token)}, "model_dim": 7168, "inter_dim": 384, "expert": 385, @@ -180,10 +224,12 @@ def setUp(self): if get_gfx() != "gfx950": self.skipTest("FHMoE replay requires gfx950") - def _run_public_op(self, csv_path, timeout=600): - script_path = os.path.join(os.path.dirname(csv_path), "fhmoe_replay_child.py") + def _run_public_op(self, csv_path, token=1, timeout=600): + script_path = os.path.join( + os.path.dirname(csv_path), f"fhmoe_replay_child_t{token}.py" + ) with open(script_path, "w") as handle: - handle.write(_child_script()) + handle.write(_child_script(token)) env = os.environ.copy() env["AITER_CONFIG_FHMOE"] = csv_path env["PYTHONPATH"] = AITER_ROOT + os.pathsep + env.get("PYTHONPATH", "") @@ -198,6 +244,38 @@ def _run_public_op(self, csv_path, timeout=600): check=False, ) + def _run_fhmoe_tuner(self, untuned_path, tuned_path, timeout=1200): + _cleanup_stale_lock_files() + env = os.environ.copy() + script_dir = os.path.dirname(_TUNE_SCRIPT) + env["PYTHONPATH"] = script_dir + os.pathsep + env.get("PYTHONPATH", "") + env["TUNE_MOE_KERNEL_REGEX"] = _KERNEL_REGEX + env.pop("AITER_BYPASS_TUNE_CONFIG", None) + cmd = [ + sys.executable, + _TUNE_SCRIPT, + "--fhmoe", + "-i", + untuned_path, + "-o", + tuned_path, + "--mp", + "1", + ] + try: + return subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=timeout, + cwd=AITER_ROOT, + env=env, + check=False, + ) + except subprocess.TimeoutExpired: + _cleanup_stale_lock_files() + raise + def test_uncovered_token_raises_in_fresh_process(self): with tempfile.TemporaryDirectory() as tmp: csv_path = os.path.join(tmp, "untuned_empty.csv") @@ -212,23 +290,119 @@ def test_uncovered_token_raises_in_fresh_process(self): self.assertIn("REPLAY_RAISE", result.stdout) self.assertIn("does not cover this DSV4 I384", result.stdout) - def test_generated_csv_selected_by_public_fused_moe(self): + def test_tuner_written_csv_selected_by_public_fused_moe(self): + with tempfile.TemporaryDirectory() as tmp: + untuned_path = os.path.join(tmp, "untuned_fhmoe.csv") + tuned_path = os.path.join(tmp, "tuned_fhmoe.csv") + _write_untuned_rows(untuned_path, _REPLAY_TOKENS) + tune = self._run_fhmoe_tuner(untuned_path, tuned_path) + tune_out = tune.stdout + tune.stderr + self.assertEqual( + tune.returncode, + 0, + f"FhmoeTuner failed\n{tune_out[-4000:]}", + ) + self.assertTrue(os.path.isfile(tuned_path), "tuner did not write -o") + with open(tuned_path, newline="") as handle: + rows = list(csv.DictReader(handle)) + self.assertTrue(rows, f"tuner wrote no rows\n{tune_out[-4000:]}") + self.assertIn("us", rows[0], "producer CSV must include timed us") + by_token = {int(row["token"]): row for row in rows} + self.assertEqual( + set(by_token), + set(_REPLAY_TOKENS), + f"tuner tokens {sorted(by_token)} != {_REPLAY_TOKENS}", + ) + for token, row in by_token.items(): + kn1, kn2 = row["kernelName1"], row["kernelName2"] + self.assertEqual(kn1, _TUNE_KN1, f"token={token} kn1={kn1}") + self.assertIn(kn2, _TUNE_KN2S, f"token={token} kn2={kn2}") + self.assertGreater(float(row["us"]), 0.0) + if token == 1: + self.assertNotEqual( + kn2, + _SHIPPED_TOKEN1_KN2, + "winner matches shipped token=1; cannot prove AITER_CONFIG_FHMOE", + ) + for token, row in by_token.items(): + result = self._run_public_op(tuned_path, token=token) + output = result.stdout + result.stderr + self.assertEqual( + result.returncode, + 0, + f"token={token} replay child failed\n{output[-4000:]}", + ) + self.assertIn("REPLAY_OK", result.stdout) + self.assertIn(row["kernelName1"], output) + self.assertIn(row["kernelName2"], output) + self.assertNotIn("REPLAY_RAISE", result.stdout) + + def test_run_config_times_public_fused_moe(self): from aiter.jit.utils.chip_info import get_cu_num, get_gfx_runtime + shipped = os.path.join(AITER_ROOT, "aiter", "configs", "tuned_fhmoe.csv") + if not os.path.isfile(shipped): + self.skipTest("aiter/configs/tuned_fhmoe.csv is missing") with tempfile.TemporaryDirectory() as tmp: - csv_path = os.path.join(tmp, "generated_fhmoe.csv") - _write_generated_row(csv_path, get_gfx_runtime(), int(get_cu_num())) - result = self._run_public_op(csv_path) + csv_path = os.path.join(tmp, "run_config_fhmoe.csv") + with open(shipped, newline="") as handle: + rows = list(csv.DictReader(handle)) + token1 = next((row for row in rows if int(row["token"]) == 1), None) + if token1 is None: + self.skipTest("shipped tuned_fhmoe.csv has no token=1 row") + token1 = dict(token1) + token1["gfx"] = get_gfx_runtime() + token1["cu_num"] = str(int(get_cu_num())) + with open(csv_path, "w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=token1.keys()) + writer.writeheader() + writer.writerow(token1) + _cleanup_stale_lock_files() + env = os.environ.copy() + script_dir = os.path.dirname(_TUNE_SCRIPT) + env["PYTHONPATH"] = script_dir + os.pathsep + env.get("PYTHONPATH", "") + env.pop("AITER_BYPASS_TUNE_CONFIG", None) + result = subprocess.run( + [ + sys.executable, + _TUNE_SCRIPT, + "--fhmoe", + "--run_config", + csv_path, + "--warmup", + "1", + "--iters", + "2", + "--mp", + "1", + ], + capture_output=True, + text=True, + timeout=600, + cwd=AITER_ROOT, + env=env, + check=False, + ) output = result.stdout + result.stderr self.assertEqual( result.returncode, 0, - f"generated replay child failed\n{output[-4000:]}", + f"--fhmoe --run_config failed\n{output[-4000:]}", + ) + self.assertIn("shared=", output) + lines = [line.strip() for line in output.splitlines()] + self.assertTrue( + any(line.endswith("OK") for line in lines), + f"no OK status line\n{output[-4000:]}", + ) + self.assertFalse( + any(line.endswith("ERROR") for line in lines), + f"run_config ERROR\n{output[-4000:]}", + ) + self.assertFalse( + any(line.endswith("MISMATCH") for line in lines), + f"run_config MISMATCH\n{output[-4000:]}", ) - self.assertIn("REPLAY_OK", result.stdout) - self.assertIn(_REPLAY_KN1, output) - self.assertIn(_REPLAY_KN2, output) - self.assertNotIn("REPLAY_RAISE", result.stdout) if __name__ == "__main__": From 0b2f8cdc2c4dec2429334533169b11f04bddf91a Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Fri, 18 Sep 2026 16:45:03 +0200 Subject: [PATCH 06/14] [FlyDSL] Cover DSV4 M-tiers in the FHMoE producer replay fused_moe(M=16) requires tuned rows for 1, 2, 4, 8, and 16. Tune that ladder, then replay only 1 and 16 through the public op. --- op_tests/tuning_tests/test_fhmoe_replay.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/op_tests/tuning_tests/test_fhmoe_replay.py b/op_tests/tuning_tests/test_fhmoe_replay.py index 6323411094d..07ca4bb7b41 100644 --- a/op_tests/tuning_tests/test_fhmoe_replay.py +++ b/op_tests/tuning_tests/test_fhmoe_replay.py @@ -37,6 +37,8 @@ _KERNEL_REGEX = ( rf"{_TUNE_KN1} {_TUNE_KN2_A}$|{_TUNE_KN1} {_TUNE_KN2_B}$" ) +# fused_moe refuses M=16 unless the CSV has every power-of-two through 16. +_TUNE_TOKENS = (1, 2, 4, 8, 16) _REPLAY_TOKENS = (1, 16) _NATIVE_FIELDS = [ @@ -244,7 +246,7 @@ def _run_public_op(self, csv_path, token=1, timeout=600): check=False, ) - def _run_fhmoe_tuner(self, untuned_path, tuned_path, timeout=1200): + def _run_fhmoe_tuner(self, untuned_path, tuned_path, timeout=1800): _cleanup_stale_lock_files() env = os.environ.copy() script_dir = os.path.dirname(_TUNE_SCRIPT) @@ -294,7 +296,7 @@ def test_tuner_written_csv_selected_by_public_fused_moe(self): with tempfile.TemporaryDirectory() as tmp: untuned_path = os.path.join(tmp, "untuned_fhmoe.csv") tuned_path = os.path.join(tmp, "tuned_fhmoe.csv") - _write_untuned_rows(untuned_path, _REPLAY_TOKENS) + _write_untuned_rows(untuned_path, _TUNE_TOKENS) tune = self._run_fhmoe_tuner(untuned_path, tuned_path) tune_out = tune.stdout + tune.stderr self.assertEqual( @@ -310,8 +312,8 @@ def test_tuner_written_csv_selected_by_public_fused_moe(self): by_token = {int(row["token"]): row for row in rows} self.assertEqual( set(by_token), - set(_REPLAY_TOKENS), - f"tuner tokens {sorted(by_token)} != {_REPLAY_TOKENS}", + set(_TUNE_TOKENS), + f"tuner tokens {sorted(by_token)} != {_TUNE_TOKENS}", ) for token, row in by_token.items(): kn1, kn2 = row["kernelName1"], row["kernelName2"] @@ -324,7 +326,8 @@ def test_tuner_written_csv_selected_by_public_fused_moe(self): _SHIPPED_TOKEN1_KN2, "winner matches shipped token=1; cannot prove AITER_CONFIG_FHMOE", ) - for token, row in by_token.items(): + for token in _REPLAY_TOKENS: + row = by_token[token] result = self._run_public_op(tuned_path, token=token) output = result.stdout + result.stderr self.assertEqual( From 6659ae13de1dbfc889016b04ca96c02375076be0 Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Mon, 21 Sep 2026 10:28:31 +0200 Subject: [PATCH 07/14] [FlyDSL] Skip FHMoE tune on non-gfx950 and drop --all from op_tune.sh Operators Tuning runs the full job list on gfx942; --fhmoe SystemExit must not fail the workflow. --- .github/scripts/op_tune.sh | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/.github/scripts/op_tune.sh b/.github/scripts/op_tune.sh index 20483d7f948..96456ee9da0 100755 --- a/.github/scripts/op_tune.sh +++ b/.github/scripts/op_tune.sh @@ -24,7 +24,7 @@ declare -a tune_jobs=( "ck_gemm_a8w8_blockscale_bpreshuffle:csrc/ck_gemm_a8w8_blockscale_bpreshuffle:op_tests/test_gemm_a8w8_blockscale.py:python3 csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_tune.py --preshuffle -i aiter/configs/a8w8_blockscale_bpreshuffle_untuned_gemm.csv -o aiter/configs/a8w8_blockscale_bpreshuffle_tuned_gemm.csv" "ck_gemm_a8w8_bpreshuffle:csrc/ck_gemm_a8w8_bpreshuffle:op_tests/test_gemm_a8w8.py:python3 csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py -i aiter/configs/a8w8_bpreshuffle_untuned_gemm.csv -o aiter/configs/a8w8_bpreshuffle_tuned_gemm.csv" "ck_gemm_moe_2stages_codegen:csrc/ck_gemm_moe_2stages_codegen:op_tests/test_moe.py:python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py -i aiter/configs/untuned_fmoe.csv -o aiter/configs/tuned_fmoe.csv" - "fhmoe:csrc/ck_gemm_moe_2stages_codegen:op_tests/test_fhmoe.py:python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py --fhmoe --all -i aiter/configs/untuned_fhmoe.csv -o aiter/configs/tuned_fhmoe.csv" + "fhmoe:csrc/ck_gemm_moe_2stages_codegen:op_tests/test_fhmoe.py:python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py --fhmoe -i aiter/configs/untuned_fhmoe.csv -o aiter/configs/tuned_fhmoe.csv" #"ck_gemm_a4w4_blockscale:csrc/ck_gemm_a4w4_blockscale:op_tests/test_gemm_a4w4_blockscale.py:python3 csrc/ck_gemm_a4w4_blockscale/gemm_a4w4_blockscale_tune.py -i aiter/configs/a4w4_blockscale_untuned_gemm.csv -o aiter/configs/a4w4_blockscale_tuned_gemm.csv" ) @@ -47,6 +47,16 @@ for job in "${tune_jobs[@]}"; do continue fi fi + if [ "$shape" = "fhmoe" ]; then + gfx=$(python3 -c "from aiter.jit.utils.chip_info import get_gfx; print(get_gfx())") || gfx="" + if [ "$gfx" != "gfx950" ]; then + echo "============================================================" + echo "Skipping fhmoe: requires gfx950, got ${gfx:-unknown}" + echo "===============================================" + echo + continue + fi + fi echo "============================================================" echo "๐Ÿงช Processing shape: $shape under directory: $dir" echo "------------------------------------------------------------" From 442113bd1c22fc66cdc9ec8adbf841af66cdef9a Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Mon, 21 Sep 2026 16:12:50 +0200 Subject: [PATCH 08/14] [FlyDSL] Cover FHMoE in Tuning Tests like other tuner families csv_validation, a two-pair pipeline smoke, and --run_config on the shipped DSV4 CSV; skip --fhmoe off gfx950 so SystemExit is not a FAIL. --- op_tests/tuning_tests/README.md | 3 +- op_tests/tuning_tests/test_csv_validation.py | 20 +++++ op_tests/tuning_tests/test_run_config.py | 34 ++++++++ op_tests/tuning_tests/test_tune_pipeline.py | 92 +++++++++++++++++++- 4 files changed, 147 insertions(+), 2 deletions(-) diff --git a/op_tests/tuning_tests/README.md b/op_tests/tuning_tests/README.md index 0029a54514f..a4c53f3683b 100644 --- a/op_tests/tuning_tests/README.md +++ b/op_tests/tuning_tests/README.md @@ -38,6 +38,7 @@ runtime budgets. | `batched_a8w8` | `csrc/ck_batched_gemm_a8w8/batched_gemm_a8w8_tune.py` | `a8w8_tuned_batched_gemm.csv` | โœ“ | โœ“ | | `batched_bf16` | `csrc/ck_batched_gemm_bf16/batched_gemm_bf16_tune.py` | `bf16_tuned_batched_gemm.csv` | โœ“ | โœ“ + shape_grouped | | `fmoe` | `csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py` | `tuned_fmoe.csv` + model_configs | โœ“ | โœ“ (bf16/fp8/int8/gelu) | +| `fhmoe` | same + `--fhmoe` | `tuned_fhmoe.csv` | โœ“ | โœ“ (mp1, gfx950) | | `gradlib_bf16` | `gradlib/gradlib/gemm_tuner.py` | `bf16_tuned_gemm.csv` | โœ“ | โœ“ (hipBLASLt/ASM/FlyDSL) | | `gdn_k5_opt` | `csrc/gdn_k5/chunk_gdn_h_opt_tune.py` | `model_configs/*_chunk_gdn_h_opt_tuned.csv` | โœ“ | โœ“ (shape-only varlen smoke) | @@ -125,6 +126,6 @@ python3 -m unittest op_tests.tuning_tests.test_run_config.TestRunConfigCustom -v Available families include `a8w8`, `a8w8_bpreshuffle`, `a8w8_blockscale`, `a8w8_blockscale_bpreshuffle`, `a4w4_blockscale`, `a6w4_asm`, `a4w6_asm`, `batched_a8w8`, `batched_bf16`, -`fmoe`, `gradlib_bf16`, and `gdn_k5_opt`. +`fmoe`, `fhmoe`, `gradlib_bf16`, and `gdn_k5_opt`. The test checks both **exit code** and **per-shape status** โ€” shapes with `ERROR` (kernel crash) or `MISMATCH` (accuracy exceeded errRatio) will fail the test. diff --git a/op_tests/tuning_tests/test_csv_validation.py b/op_tests/tuning_tests/test_csv_validation.py index d1b2a8294cb..641235ed98f 100644 --- a/op_tests/tuning_tests/test_csv_validation.py +++ b/op_tests/tuning_tests/test_csv_validation.py @@ -35,6 +35,7 @@ class TestCSVValidation(unittest.TestCase): "bf16": "bf16_tuned_gemm.csv", "bf16_batched": "bf16_tuned_batched_gemm.csv", "fmoe": "tuned_fmoe.csv", + "fhmoe": "tuned_fhmoe.csv", } def _load_csv(self, name): @@ -173,6 +174,24 @@ def test_flydsl_stage2_sort_block_matches_fmoe_config(self): "FlyDSL stage2 sorting layout mismatches:\n" + "\n".join(mismatches), ) + def test_fhmoe_no_duplicates(self): + self._check_no_duplicates( + "fhmoe", + extra_keys=[ + "act_type", + "dtype", + "q_dtype_a", + "q_dtype_w", + "q_type", + "use_g1u1", + "doweight_stage1", + "shared_expert_id", + "hidden_pad", + "intermediate_pad", + "gate_mode", + ], + ) + def test_no_git_conflict_markers(self): for name, fname in self.TUNED_CSVS.items(): with self.subTest(csv=name): @@ -227,6 +246,7 @@ def test_untuned_csvs_exist(self): "a8w8_untuned_batched_gemm.csv", "bf16_untuned_batched_gemm.csv", "untuned_fmoe.csv", + "untuned_fhmoe.csv", ] for f in untuned_files: with self.subTest(file=f): diff --git a/op_tests/tuning_tests/test_run_config.py b/op_tests/tuning_tests/test_run_config.py index ae21835d2a9..8dd99973d12 100644 --- a/op_tests/tuning_tests/test_run_config.py +++ b/op_tests/tuning_tests/test_run_config.py @@ -44,6 +44,24 @@ def _gpu_available(): return False +def _current_gfx(): + try: + from aiter.jit.utils.chip_info import get_gfx + + return get_gfx() + except Exception: # noqa: BLE001 + return None + + +def _skip_unless_required_gfx(test, name, cfg): + required = cfg.get("required_gfx") + if not required: + return + gfx = _current_gfx() + if gfx != required: + test.skipTest(f"{name} requires {required}, got {gfx or 'unknown'}") + + def _find_tuned_csvs(pattern): """Find all tuned CSVs matching pattern in configs/ and model_configs/.""" found = [] @@ -327,6 +345,17 @@ def _format_failures( "timeout": 3600, "config_property": "AITER_CONFIG_FMOE_FILE", }, + "fhmoe": { + "script": "csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py", + "csv_pattern": "tuned_fhmoe", + "exclude_patterns": ["untuned"], + "extra_args": ["--fhmoe"], + "config_property": "AITER_CONFIG_FHMOE_FILE", + # gemm_moe_tune.py --fhmoe SystemExits in __main__ off gfx950. + "required_gfx": "gfx950", + # 12 DSV4 rows, distinct FlyDSL pairs. 600s died after token=512 OK. + "timeout": 1800, + }, "gradlib_bf16": { "script": "gradlib/gradlib/gemm_tuner.py", "csv_pattern": "bf16_tuned_gemm", @@ -355,6 +384,7 @@ class TestRunConfig(unittest.TestCase): def _test_family(self, name): cfg = TUNER_FAMILIES[name] + _skip_unless_required_gfx(self, name, cfg) timeout = cfg.get("timeout", 600) extra_args = cfg.get("extra_args", None) @@ -444,6 +474,9 @@ def test_batched_bf16(self): def test_fmoe(self): self._test_family("fmoe") + def test_fhmoe(self): + self._test_family("fhmoe") + def test_gradlib_bf16(self): self._test_family("gradlib_bf16") @@ -487,6 +520,7 @@ def test_custom(self): f"Unknown family '{family}'. Available: {list(TUNER_FAMILIES.keys())}", ) cfg = TUNER_FAMILIES[family] + _skip_unless_required_gfx(self, family, cfg) timeout = cfg.get("timeout", 600) extra_args = cfg.get("extra_args", None) diff --git a/op_tests/tuning_tests/test_tune_pipeline.py b/op_tests/tuning_tests/test_tune_pipeline.py index 88f6a312960..c9a615c2b89 100644 --- a/op_tests/tuning_tests/test_tune_pipeline.py +++ b/op_tests/tuning_tests/test_tune_pipeline.py @@ -47,6 +47,23 @@ def _get_platform_dtypes(): return "torch.float8_e4m3fnuz", "QuantType.per_Token" +def _is_gfx950(): + try: + from aiter.jit.utils.chip_info import get_gfx + + return get_gfx() == "gfx950" + except Exception: # noqa: BLE001 + return False + + +# Smoke: two kn1ร—kn2 pairs, not the full FlyDSL cartesian. +_FHMOE_KN1 = "flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w4_gui_kw4_fp8" +_FHMOE_KN2 = "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic" +_FHMOE_KERNEL_REGEX = ( + rf"^{_FHMOE_KN1} {_FHMOE_KN2}_bnt2$|^{_FHMOE_KN1} {_FHMOE_KN2}_persist$" +) + + def _write_csv(path, header, rows): with open(path, "w", newline="") as f: writer = csv.writer(f) @@ -74,7 +91,9 @@ def _cleanup_stale_lock_files(): pass -def _run_tuner(script, untuned, tuned, extra_args=None, timeout=300, mp=1): +def _run_tuner( + script, untuned, tuned, extra_args=None, timeout=300, mp=1, extra_env=None +): _cleanup_stale_lock_files() cmd = [ sys.executable, @@ -95,6 +114,8 @@ def _run_tuner(script, untuned, tuned, extra_args=None, timeout=300, mp=1): env = os.environ.copy() script_dir = os.path.dirname(os.path.join(AITER_ROOT, script)) env["PYTHONPATH"] = script_dir + ":" + env.get("PYTHONPATH", "") + if extra_env: + env.update(extra_env) try: return subprocess.run( cmd, @@ -286,6 +307,69 @@ def setUpClass(cls): "timeout": 1800, "timeout_mp1": 2400, }, + "fhmoe": { + "script": "csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py", + "header": [ + "token", + "model_dim", + "inter_dim", + "expert", + "topk", + "shared_expert_id", + "act_type", + "dtype", + "q_dtype_a", + "q_dtype_w", + "q_type", + "use_g1u1", + "doweight_stage1", + "hidden_pad", + "intermediate_pad", + "gate_mode", + ], + "shapes": [ + ( + 1, + 7168, + 384, + 385, + 7, + 384, + "ActivationType.Silu", + "torch.bfloat16", + "torch.float8_e4m3fn", + "torch.float4_e2m1fn_x2", + "QuantType.per_1x32", + 1, + 0, + 0, + 0, + "GateMode.INTERLEAVE", + ), + ], + "keys": [ + "cu_num", + "token", + "model_dim", + "inter_dim", + "expert", + "topk", + "act_type", + "dtype", + "q_dtype_a", + "q_dtype_w", + "q_type", + "use_g1u1", + "doweight_stage1", + "shared_expert_id", + "hidden_pad", + "intermediate_pad", + "gate_mode", + ], + "extra_args": ["--fhmoe"], + "extra_env": {"TUNE_MOE_KERNEL_REGEX": _FHMOE_KERNEL_REGEX}, + "timeout_mp1": 1200, + }, "gdn_k5_opt": { "script": "csrc/gdn_k5/chunk_gdn_h_opt_tune.py", "header": [ @@ -479,6 +563,7 @@ def _run_one(self, name, mp=1): extra_args=cfg.get("extra_args"), timeout=timeout, mp=mp, + extra_env=cfg.get("extra_env"), ) if result.returncode != 0: print(f"\n=== {name} ({mp_label}) STDOUT ===\n{result.stdout[-2000:]}") @@ -584,6 +669,11 @@ def test_fmoe_mp1(self): def test_fmoe_mp_default(self): self._run_one("fmoe", mp=None) + def test_fhmoe_mp1(self): + if not _is_gfx950(): + self.skipTest("FHMoE requires gfx950") + self._run_one("fhmoe", mp=1) + def test_a6w6_blockscale_mp1(self): self._run_one("a6w6_blockscale", mp=1) From 835633814fcb23b589461ce812c8c6a4d74d3ff0 Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Mon, 21 Sep 2026 17:06:34 +0200 Subject: [PATCH 09/14] [FlyDSL] Drop FHMoE accuracy_gate mock test Cosine vs torch already lives in the tuner and Tuning Tests; the 0-GPU mock lock was not a family pattern. --- .../tuning_tests/test_fhmoe_accuracy_gate.py | 205 ------------------ 1 file changed, 205 deletions(-) delete mode 100644 op_tests/tuning_tests/test_fhmoe_accuracy_gate.py diff --git a/op_tests/tuning_tests/test_fhmoe_accuracy_gate.py b/op_tests/tuning_tests/test_fhmoe_accuracy_gate.py deleted file mode 100644 index fa3d8070c85..00000000000 --- a/op_tests/tuning_tests/test_fhmoe_accuracy_gate.py +++ /dev/null @@ -1,205 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. -"""FHMoE tuner must not pick a fast-but-wrong kernel pair.""" - -from __future__ import annotations - -import argparse -import os -import sys -import unittest -from unittest.mock import patch - -import pandas as pd - -AITER_ROOT = os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -) -if AITER_ROOT not in sys.path: - sys.path.insert(0, AITER_ROOT) - -try: - from csrc.ck_gemm_moe_2stages_codegen.gemm_moe_tune import ( # noqa: E402 - FhmoeTuner, - FmoeTuner, - cosine_diff_compare, - ) -except ImportError as _import_error: # torch / aiter not in this interpreter - FhmoeTuner = None - FmoeTuner = None - cosine_diff_compare = None - _IMPORT_ERROR = _import_error -else: - _IMPORT_ERROR = None - -_FHMOE_KEYS = [ - "gfx", - "cu_num", - "token", - "model_dim", - "inter_dim", - "expert", - "topk", - "act_type", - "dtype", - "q_dtype_a", - "q_dtype_w", - "q_type", - "use_g1u1", - "doweight_stage1", - "shared_expert_id", - "hidden_pad", - "intermediate_pad", - "gate_mode", -] -_FHMOE_RESULTS = [ - "block_m", - "ksplit", - "kernelName1", - "kernelName2", - "us", -] - - -def _tuner(): - return FhmoeTuner("fhmoeTuner", _FHMOE_KEYS, _FHMOE_RESULTS, "fhmoe tuner") - - -def _info(): - return ( - "gfx950", - 256, - 1, - 7168, - 384, - 385, - 7, - "ActivationType.Silu", - "torch.bfloat16", - "torch.float8_e4m3fn", - "torch.float4_e2m1fn_x2", - "QuantType.per_1x32", - 1, - 0, - 384, - 0, - 0, - "GateMode.INTERLEAVE", - ) - - -def _untuned_row(): - return { - "gfx": "gfx950", - "cu_num": 256, - "token": 1, - "model_dim": 7168, - "inter_dim": 384, - "expert": 385, - "topk": 7, - "shared_expert_id": 384, - "act_type": "ActivationType.Silu", - "dtype": "torch.bfloat16", - "q_dtype_a": "torch.float8_e4m3fn", - "q_dtype_w": "torch.float4_e2m1fn_x2", - "q_type": "QuantType.per_1x32", - "use_g1u1": 1, - "doweight_stage1": 0, - "hidden_pad": 0, - "intermediate_pad": 0, - "gate_mode": "GateMode.INTERLEAVE", - } - - -@unittest.skipIf(FhmoeTuner is None, f"gemm_moe_tune import failed: {_IMPORT_ERROR}") -class TestFhmoeAccuracyGate(unittest.TestCase): - def test_tune_tasks_attach_torch_ref_and_cosine(self): - tuner = _tuner() - df = pd.DataFrame([_untuned_row()]) - with patch.object( - FhmoeTuner, "_kernel_pairs", return_value=[(32, "kn1", "kn2")] - ): - tasks, in_datas = tuner._fhmoe_tune_tasks(df) - self.assertEqual(len(tasks), 1) - self.assertEqual(in_datas, [(1, ())]) - task = tasks[0] - self.assertIs(task[6], FhmoeTuner.run_torch_fhmoe) - self.assertEqual(task[7][0], list(FhmoeTuner._FHMOE_REF_DATA_KEYS)) - self.assertIs(task[12], cosine_diff_compare) - - def test_post_process_rejects_fast_wrong_pair(self): - tuner = _tuner() - info = _info() - results = [ - ((info, "fast_wrong1", "fast_wrong2", 32), 10.0, 0.9), - ((info, "slow_ok1", "slow_ok2", 64), 50.0, 0.01), - ] - args = argparse.Namespace(errRatio=0.1, profile_file="") - out = tuner.post_process(results, args) - self.assertEqual(len(out), 1) - self.assertEqual(out.iloc[0]["kernelName1"], "slow_ok1") - self.assertEqual(out.iloc[0]["kernelName2"], "slow_ok2") - self.assertEqual(int(out.iloc[0]["block_m"]), 64) - - def test_post_process_drops_shape_when_all_pairs_fail_err(self): - tuner = _tuner() - info = _info() - results = [((info, "bad1", "bad2", 32), 10.0, 0.5)] - args = argparse.Namespace(errRatio=0.1, profile_file="") - out = tuner.post_process(results, args) - self.assertTrue(out.empty) - - def test_run_config_is_not_parent_fmoe(self): - self.assertIsNot(FhmoeTuner.run_config, FmoeTuner.run_config) - - def test_run_config_calls_public_fused_moe_with_shared(self): - import torch - - from csrc.ck_gemm_moe_2stages_codegen import gemm_moe_tune as tune_mod - - tuner = _tuner() - tuner.untunedf = pd.DataFrame( - [{**_untuned_row(), "gfx": "gfx950", "cu_num": 256}] - ) - args = argparse.Namespace(warmup=1, iters=1, errRatio=0.1) - fake = torch.ones((1, 8)) - data = { - "hidden": fake, - "w1": fake, - "w2": fake, - "w1_scale": fake, - "w2_scale": fake, - "shared_w1": fake, - "shared_w2": fake, - "shared_w1_scale": fake, - "shared_w2_scale": fake, - "topk_weights": fake, - "topk_ids": fake, - "w1_qt": fake, - "w2_qt": fake, - "w1_scale_qt": fake, - "w2_scale_qt": fake, - "shared_w1_qt": fake, - "shared_w2_qt": fake, - "shared_w1_scale_qt": fake, - "shared_w2_scale_qt": fake, - "shared_expert_id": 384, - } - with patch.object(FhmoeTuner, "generate_fhmoe_data", return_value=data), patch.object( - FhmoeTuner, "run_torch_fhmoe", return_value=fake - ), patch( - "aiter.test_common.run_perftest", return_value=(fake, 40.0) - ) as perf: - out = tuner.run_config(args) - self.assertEqual(len(out), 1) - self.assertEqual(out[0]["status"], "ok") - self.assertEqual(out[0]["e2e_us"], 40.0) - self.assertIn("shared=", out[0]["shape"]) - self.assertIs(perf.call_args[0][0], tune_mod.fused_moe) - self.assertIn("shared_w1", perf.call_args[1]) - self.assertIn("shared_w2", perf.call_args[1]) - self.assertEqual(perf.call_args[1]["shared_expert_id"], 384) - - -if __name__ == "__main__": - unittest.main() From e854f4f00cb74ff44722b176da58be1acf856b20 Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Tue, 22 Sep 2026 13:21:09 +0200 Subject: [PATCH 10/14] [FlyDSL] Format FHMoE tuner files for Black and Ruff --- .../gemm_moe_tune.py | 49 +++++++++++-------- op_tests/tuning_tests/test_fhmoe_replay.py | 10 ++-- 2 files changed, 32 insertions(+), 27 deletions(-) diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index 4250af66e02..1cfa6afcb7f 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -48,8 +48,8 @@ rearrange_4bit_elements, ) from aiter.jit.core import ( - AITER_CONFIG_FMOE, AITER_CONFIG_FHMOE, + AITER_CONFIG_FMOE, AITER_CONFIG_GROUPED_FMOE, AITER_CSRC_DIR, AITER_ROOT_DIR, @@ -6488,11 +6488,14 @@ def generate_fhmoe_data(shape, block_m=32, device="cuda"): / 10 ) else: - w1_dense = torch.randn( - (routed_expert, inter_dim, model_dim), - dtype=dtype, - device="cpu", - ) / 10 + w1_dense = ( + torch.randn( + (routed_expert, inter_dim, model_dim), + dtype=dtype, + device="cpu", + ) + / 10 + ) w2_dense = torch.randn( (routed_expert, model_dim, inter_dim), dtype=dtype, device="cpu" ) @@ -6513,9 +6516,7 @@ def generate_fhmoe_data(shape, block_m=32, device="cuda"): w1 = w1.view(w1_shape[0], w1_shape[1], w1_shape[2] // 2) w2 = w2.view(w2_shape[0], w2_shape[1], w2_shape[2] // 2) if TUNE_MOE_EXPERT_BALANCE: - score = torch.zeros( - (token, routed_expert), dtype=dtype, device="cpu" - ) + score = torch.zeros((token, routed_expert), dtype=dtype, device="cpu") start_col = 0 end_col = routed_topk for token_id in range(token): @@ -6526,18 +6527,12 @@ def generate_fhmoe_data(shape, block_m=32, device="cuda"): routing_seed = os.environ.get("TUNE_MOE_ROUTING_SEED") if routing_seed is not None: torch.manual_seed(int(routing_seed)) - score = torch.randn( - (token, routed_expert), dtype=dtype, device="cpu" - ) + score = torch.randn((token, routed_expert), dtype=dtype, device="cpu") shared_w1_dense = ( - torch.randn( - (1, 2 * inter_dim, model_dim), dtype=dtype, device="cpu" - ) - / 10 + torch.randn((1, 2 * inter_dim, model_dim), dtype=dtype, device="cpu") / 10 ) shared_w2_dense = ( - torch.randn((1, model_dim, inter_dim), dtype=dtype, device="cpu") - / 10 + torch.randn((1, model_dim, inter_dim), dtype=dtype, device="cpu") / 10 ) with torch.device("cpu"): shared_w1, shared_w1_scale = FmoeTuner.weight_quant( @@ -7134,7 +7129,14 @@ def post_process(self, results, args, topk=-1, fast_mode=False): ) if args.profile_file: profile_df = pd.DataFrame(profile_rows) - for col in ("act_type", "dtype", "q_dtype_a", "q_dtype_w", "q_type", "gate_mode"): + for col in ( + "act_type", + "dtype", + "q_dtype_a", + "q_dtype_w", + "q_type", + "gate_mode", + ): if col in profile_df.columns: profile_df[col] = profile_df[col].map(self._csv_cell) if os.path.exists(args.profile_file): @@ -7159,7 +7161,14 @@ def post_process(self, results, args, topk=-1, fast_mode=False): continue kn1, kn2, block_m, us, err = min(valid, key=lambda r: r[3]) row = dict(zip(self.keys, key)) - for col in ("act_type", "dtype", "q_dtype_a", "q_dtype_w", "q_type", "gate_mode"): + for col in ( + "act_type", + "dtype", + "q_dtype_a", + "q_dtype_w", + "q_type", + "gate_mode", + ): row[col] = self._csv_cell(row[col]) row.update( { diff --git a/op_tests/tuning_tests/test_fhmoe_replay.py b/op_tests/tuning_tests/test_fhmoe_replay.py index 07ca4bb7b41..543a40cf481 100644 --- a/op_tests/tuning_tests/test_fhmoe_replay.py +++ b/op_tests/tuning_tests/test_fhmoe_replay.py @@ -34,9 +34,7 @@ _TUNE_KN2_B = "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic_persist" _TUNE_KN2S = frozenset((_TUNE_KN2_A, _TUNE_KN2_B)) _SHIPPED_TOKEN1_KN2 = "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic" -_KERNEL_REGEX = ( - rf"{_TUNE_KN1} {_TUNE_KN2_A}$|{_TUNE_KN1} {_TUNE_KN2_B}$" -) +_KERNEL_REGEX = rf"{_TUNE_KN1} {_TUNE_KN2_A}$|{_TUNE_KN1} {_TUNE_KN2_B}$" # fused_moe refuses M=16 unless the CSV has every power-of-two through 16. _TUNE_TOKENS = (1, 2, 4, 8, 16) _REPLAY_TOKENS = (1, 16) @@ -144,8 +142,7 @@ def _cleanup_stale_lock_files(): def _child_script(token): - return textwrap.dedent( - f"""\ + return textwrap.dedent(f"""\ import sys sys.path.insert( 0, @@ -214,8 +211,7 @@ def _cuda(tensor): print("REPLAY_RAISE", str(exc)) raise SystemExit(0) print("REPLAY_OK", tuple(out.shape)) - """ - ) + """) @unittest.skipUnless(_gpu_available(), "No GPU available") From e2f93958a466f1ad450735eecb4252b563089fc0 Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Tue, 22 Sep 2026 14:22:08 +0200 Subject: [PATCH 11/14] [FlyDSL] Publish FHMoE tuned rows under padded token Serving looks up get_padded_M(token), not the catalogue token. post_process now writes the same key _write_candidate_csv already timed with. --- .../gemm_moe_tune.py | 31 ++++++++++++------- 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index 1cfa6afcb7f..75ac3b39de7 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -6690,13 +6690,19 @@ def _kernel_pairs(shape): pairs = [p for p in pairs if pattern.search(f"{p[1]} {p[2]}")] return pairs + @staticmethod + def _lookup_token(token): + # Serving get_2stage_cfgs looks up get_padded_M(token_num), not the + # catalogue token. Time and publish under that same key. + return get_padded_M(int(token)) + @staticmethod def _write_candidate_csv(path, shape, block_m, kn1, kn2): gate_mode = shape["gate_mode"] row = { "gfx": shape["gfx"], "cu_num": shape["cu_num"], - "token": get_padded_M(shape["token"]), + "token": FhmoeTuner._lookup_token(shape["token"]), "model_dim": shape["model_dim"], "inter_dim": shape["inter_dim"], "expert": shape["expert"], @@ -7116,17 +7122,17 @@ def post_process(self, results, args, topk=-1, fast_mode=False): key = tuple(info[0]) kn1, kn2, block_m = info[1], info[2], info[3] grouped[key].append((kn1, kn2, block_m, us, err)) - profile_rows.append( - { - **dict(zip(self.keys, key)), - "block_m": block_m, - "ksplit": 0, - "kernelName1": kn1, - "kernelName2": kn2, - "us": us, - "err": err, - } - ) + profile_row = { + **dict(zip(self.keys, key)), + "block_m": block_m, + "ksplit": 0, + "kernelName1": kn1, + "kernelName2": kn2, + "us": us, + "err": err, + } + profile_row["token"] = self._lookup_token(profile_row["token"]) + profile_rows.append(profile_row) if args.profile_file: profile_df = pd.DataFrame(profile_rows) for col in ( @@ -7161,6 +7167,7 @@ def post_process(self, results, args, topk=-1, fast_mode=False): continue kn1, kn2, block_m, us, err = min(valid, key=lambda r: r[3]) row = dict(zip(self.keys, key)) + row["token"] = self._lookup_token(row["token"]) for col in ( "act_type", "dtype", From e4a78847124b3b8c5ef6feaa161ed54a63d18887 Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Tue, 22 Sep 2026 17:49:50 +0200 Subject: [PATCH 12/14] [FlyDSL] Match FMoE stage1 k_batch filter in FHMoE Drop split-k only when it does not divide K, instead of excluding every k_batch!=1. --- csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index 75ac3b39de7..1685a2ecbea 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -6656,7 +6656,10 @@ def _kernel_pairs(shape): k_batch = kparams.get("k_batch", 1) k_wave = kparams.get("k_wave", 1) tile_k = kparams["tile_k"] - if k_batch != 1 or model_dim % k_batch != 0: + # Same stage1 k_batch cut as FmoeTuner.gen_flydsl_2stages_task: + # drop splits that do not divide K. INTERLEAVE (fp8 A) does not + # emit k_batch>1 today; when fp4 A does, keep legal split-k. + if model_dim % k_batch != 0: continue k_per_batch = model_dim // k_batch if k_per_batch % tile_k != 0: From bdac6e9d9628f309da5cfde4d29f2e7191291162 Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Wed, 23 Sep 2026 16:00:07 +0200 Subject: [PATCH 13/14] [FlyDSL] Honor --warmup/--iters in the FHMoE search mp_tuner was getting empty kwargs, so search always used run_perftest 2/101. --- csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index 1685a2ecbea..24601b5bf30 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -6962,7 +6962,11 @@ def run_torch_fhmoe( ) return routed + shared - def _fhmoe_tune_tasks(self, untunedf): + def _fhmoe_tune_tasks(self, untunedf, args): + measure_kwargs = { + "num_warmup": int(args.warmup), + "num_iters": int(args.iters), + } tasks = [] in_datas = [] for _, row in untunedf.iterrows(): @@ -6983,7 +6987,7 @@ def _fhmoe_tune_tasks(self, untunedf): (shape,), FhmoeTuner._run_fhmoe_mp, (list(self._FHMOE_MP_DATA_KEYS), shape, kn1, kn2, block_m), - {}, + measure_kwargs, FhmoeTuner.run_torch_fhmoe, (list(self._FHMOE_REF_DATA_KEYS), shape), {}, @@ -7102,7 +7106,7 @@ def run_config(self, args): def tune(self, untunedf, tunedf, args): del tunedf - tasks, in_datas = self._fhmoe_tune_tasks(untunedf) + tasks, in_datas = self._fhmoe_tune_tasks(untunedf, args) if not tasks: return [] return mp_tuner( From 6d19303b86f9851a2a9b73180a5c4c456bba5ca0 Mon Sep 17 00:00:00 2001 From: a-canadasruiz Date: Fri, 25 Sep 2026 17:03:14 +0200 Subject: [PATCH 14/14] [FlyDSL] Time both fused and non-fused stage1 quant in FHMoE A fused-only pair list drops shipped DSV4 winners whose kn1 has no trailing _fp8. Emit the same names as FmoeTuner.s1_variants. --- csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index 24601b5bf30..cbac9240182 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -6668,7 +6668,13 @@ def _kernel_pairs(shape): k_per_batch % k_wave != 0 or (k_per_batch // k_wave) % tile_k != 0 ): continue - s1_names.append(kname + ("_fp8" if a_dtype == "fp8" else "_fp4")) + # Trailing _fp8/_fp4 (after _t): fused_moe fuses inter-stage quant + # into stage1; bare kname: quant after. Same as FmoeTuner. + fused = kname + ("_fp8" if a_dtype == "fp8" else "_fp4") + if k_batch > 1: + s1_names.append(fused) + else: + s1_names.extend((kname, fused)) s2_names = [] for kname, kparams in s2_all.items(): if kparams.get("tile_m") != block_m: