From 8c78bbecad9c41b85a2feadb2acc332031e4b058 Mon Sep 17 00:00:00 2001 From: Felix Marty Date: Wed, 30 Sep 2026 12:38:20 -0500 Subject: [PATCH 1/5] proposed fix --- csrc/ck_tile_gemm_moe_2stages/gen_instances.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/csrc/ck_tile_gemm_moe_2stages/gen_instances.py b/csrc/ck_tile_gemm_moe_2stages/gen_instances.py index 429119b83e4..91ea52a4445 100644 --- a/csrc/ck_tile_gemm_moe_2stages/gen_instances.py +++ b/csrc/ck_tile_gemm_moe_2stages/gen_instances.py @@ -39,6 +39,7 @@ def __init__( activation, mul_routed_weight_stage, is_split_k, + has_bias, istune=False, ): self.working_path = working_path @@ -54,6 +55,7 @@ def __init__( self.c_dtype = c_dtype.lower() self.quant_type = quant_type self.is_split_k = is_split_k + self.has_bias = has_bias self.activation = act_dict[activation] self.mul_routed_weight_stage = mul_routed_weight_stage @@ -283,7 +285,7 @@ def validate_and_format(template: str, mapping: dict) -> str: "(acc_data_type)": dtype_dict[self.acc_dtype], "(c_data_type)": dtype_dict[self.c_dtype], "(activation)": self.activation, - "(has_bias)": "true" if self.activation == 2 else "false", + "(has_bias)": "true" if self.has_bias else "false", "(split_k)": "true" if self.is_split_k else "false", } format_args = {str(key): value.name for key, value in mapping.items()} @@ -647,10 +649,11 @@ def generate_common_header(working_path, dispatch_files, manifest_files): # for name-based dispatch header generation name_lookup_entries = [] - for a_type, c_dtype, act_type, is_split_k in itertools.product( - a_types, c_dtypes, act_types, is_split_k_l + for a_type, c_dtype, act_type, is_split_k, has_bias in itertools.product( + a_types, c_dtypes, act_types, is_split_k_l, (False, True) ): - has_bias = act_type == "swiglu" + if act_type != "swiglu" and has_bias: + continue # a8w8 do not support if a_type in ["fp8", "bf8"] and is_split_k: @@ -664,6 +667,7 @@ def generate_common_header(working_path, dispatch_files, manifest_files): act_type, 2, is_split_k, + has_bias, False, ) # gen all instances for gemm1 and gemm2 From 61b5bee72ecd84da0b06b7b6de5838b3b8ca5f79 Mon Sep 17 00:00:00 2001 From: Felix Marty Date: Wed, 30 Sep 2026 13:08:23 -0500 Subject: [PATCH 2/5] Fail explicitly for unsupported CK-Tile MoE dispatch Assisted-by: Codex --- .../include/moe_cktile2stages.h | 40 ++++++++++++++++++- .../moe_cktile2stages.cu | 10 +++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages.h b/csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages.h index b99b4b7dd3a..fb459704cf5 100644 --- a/csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages.h +++ b/csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages.h @@ -53,7 +53,25 @@ template struct moe_gemm1_heuristic_dispatcher { - static MoeKernel dispatch(int M, int N, int K, int block_m) {} + static MoeKernel dispatch(int M, int N, int K, int block_m) + { + TORCH_CHECK(false, + "No CK-Tile MoE GEMM1 specialization: activation=", + activation, + ", has_bias=", + kHasBias, + ", split_k=", + split_k, + ", M=", + M, + ", N=", + N, + ", K=", + K, + ", block_m=", + block_m); + return {}; + } }; template struct moe_gemm2_heuristic_dispatcher { - static MoeKernel dispatch(int M, int N, int K, int block_m) {} + static MoeKernel dispatch(int M, int N, int K, int block_m) + { + TORCH_CHECK(false, + "No CK-Tile MoE GEMM2 specialization: activation=", + activation, + ", has_bias=", + kHasBias, + ", split_k=", + split_k, + ", M=", + M, + ", N=", + N, + ", K=", + K, + ", block_m=", + block_m); + return {}; + } }; __attribute__((visibility("default"))) torch::Tensor diff --git a/csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu b/csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu index 8135639d181..34a27ab50f6 100644 --- a/csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu +++ b/csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu @@ -243,6 +243,16 @@ MoeKernel moe_dispatch(int M, int N, int K, int block_m, int activation, bool ha } } } + TORCH_CHECK(false, + "Unsupported CK-Tile MoE activation: activation=", + activation, + ", has_bias=", + has_bias, + ", split_k=", + split_k, + ", block_m=", + block_m); + return {}; } torch::Tensor cktile_moe_gemm1(torch::Tensor& XQ, From ded5df2ab7df7f69706b032de6f4236b70b3ad93 Mon Sep 17 00:00:00 2001 From: Felix Marty Date: Wed, 30 Sep 2026 13:23:16 -0500 Subject: [PATCH 3/5] fix wrongful fused_moe dispatch --- .../include/moe_cktile2stages.h | 28 ++--- .../moe_cktile2stages.cu | 4 +- op_tests/test_moe_mxfp4_inter_dim_dispatch.py | 118 +++++++++++++++++- 3 files changed, 124 insertions(+), 26 deletions(-) diff --git a/csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages.h b/csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages.h index fb459704cf5..2baefd8f09e 100644 --- a/csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages.h +++ b/csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages.h @@ -53,23 +53,17 @@ template struct moe_gemm1_heuristic_dispatcher { - static MoeKernel dispatch(int M, int N, int K, int block_m) + static MoeKernel dispatch(int, int, int, int) { TORCH_CHECK(false, "No CK-Tile MoE GEMM1 specialization: activation=", activation, ", has_bias=", kHasBias, - ", split_k=", + ", split_k_gt_1=", split_k, - ", M=", - M, - ", N=", - N, - ", K=", - K, - ", block_m=", - block_m); + ". Consider adding the specialization in " + "csrc/ck_tile_gemm_moe_2stages/gen_instances.py"); return {}; } }; @@ -83,23 +77,17 @@ template struct moe_gemm2_heuristic_dispatcher { - static MoeKernel dispatch(int M, int N, int K, int block_m) + static MoeKernel dispatch(int, int, int, int) { TORCH_CHECK(false, "No CK-Tile MoE GEMM2 specialization: activation=", activation, ", has_bias=", kHasBias, - ", split_k=", + ", split_k_gt_1=", split_k, - ", M=", - M, - ", N=", - N, - ", K=", - K, - ", block_m=", - block_m); + ". Consider adding the specialization in " + "csrc/ck_tile_gemm_moe_2stages/gen_instances.py"); return {}; } }; diff --git a/csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu b/csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu index 34a27ab50f6..3a3d55c52be 100644 --- a/csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu +++ b/csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu @@ -249,9 +249,7 @@ MoeKernel moe_dispatch(int M, int N, int K, int block_m, int activation, bool ha ", has_bias=", has_bias, ", split_k=", - split_k, - ", block_m=", - block_m); + split_k); return {}; } diff --git a/op_tests/test_moe_mxfp4_inter_dim_dispatch.py b/op_tests/test_moe_mxfp4_inter_dim_dispatch.py index 17838b4bb31..f26df0081c2 100644 --- a/op_tests/test_moe_mxfp4_inter_dim_dispatch.py +++ b/op_tests/test_moe_mxfp4_inter_dim_dispatch.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Dispatch-only coverage for non-256-aligned MXFP4 MoE shapes.""" +"""Dispatch and correctness coverage for MXFP4 MoE shapes.""" import functools from types import SimpleNamespace @@ -11,15 +11,19 @@ import aiter import aiter.fused_moe as fused_moe_module from aiter import ActivationType, QuantType, dtypes -from aiter.fused_moe import cktile_moe_stage2, get_2stage_cfgs +from aiter.fused_moe import cktile_moe_stage2, fused_moe, get_2stage_cfgs from aiter.jit.utils.chip_info import get_gfx from aiter.ops.flydsl.moe_common import GateMode -from aiter.ops.shuffle import shuffle_scale +from aiter.ops.shuffle import shuffle_scale, shuffle_weight +from aiter.ops.triton.quant import dynamic_mxfp4_quant _SKIP = pytest.mark.skipif( get_gfx() not in ("gfx942", "gfx950"), reason="CDNA (gfx942/gfx950) required for MXFP4 MoE dispatch", ) +_GFX950_ONLY = pytest.mark.skipif( + get_gfx() != "gfx950", reason="gfx950 CK-Tile MXFP4 MoE required" +) MODEL_DIM = 6144 INTER_DIM = 384 @@ -249,6 +253,114 @@ def test_cktile_stage2_allows_aligned_mxfp4_shape(monkeypatch): assert called == [True] +@_GFX950_ONLY +def test_cktile_moe_missing_silu_bias_specialization_raises(): + x = torch.zeros((1, 256), dtype=torch.bfloat16, device="cuda") + w1 = torch.zeros((1, 256, 128), dtype=torch.uint8, device="cuda").view( + torch.float4_e2m1fn_x2 + ) + out = torch.empty((1, 1, 128), dtype=torch.bfloat16, device="cuda") + sorted_ids = torch.zeros(16, dtype=torch.int32, device="cuda") + expert_ids = torch.zeros(1, dtype=torch.int32, device="cuda") + num_valid_ids = torch.ones(1, dtype=torch.int32, device="cuda") + bias = torch.zeros((1, 256), dtype=torch.float32, device="cuda") + + with pytest.raises(RuntimeError) as error: + aiter.moe_cktile2stages_gemm1( + x, + w1, + out, + sorted_ids, + expert_ids, + num_valid_ids, + topk=1, + exp_bias=bias, + activation=ActivationType.Silu, + block_m=16, + split_k=1, + ) + assert "activation=0, has_bias=1, split_k_gt_1=0" in str(error.value) + assert ( + "Consider adding the specialization in " + "csrc/ck_tile_gemm_moe_2stages/gen_instances.py" + ) in str(error.value) + + +@_GFX950_ONLY +def test_mxfp4_bf16_swiglu_interleaved_matches_separated(monkeypatch): + monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "2147483647") + monkeypatch.setenv("GPTOSS_SWIGLU_MXFP4_BF16_BOUND", "2147483647") + get_2stage_cfgs.cache_clear() + model_dim, inter_dim = 1024, 256 + assert _stage_backend( + _dispatch( + model_dim=model_dim, + inter_dim=inter_dim, + gate_mode=GateMode.INTERLEAVE, + ).stage1 + ) == "cktile" + + torch.manual_seed(0) + x = torch.randn((TOKEN, model_dim), dtype=torch.bfloat16, device="cuda") + w13 = torch.randn( + (E, inter_dim * 2, model_dim), dtype=torch.bfloat16, device="cuda" + ).mul_(0.1) + w2 = torch.randn( + (E, model_dim, inter_dim), dtype=torch.bfloat16, device="cuda" + ).mul_(0.05) + topk_ids = ( + torch.arange(TOKEN * TOPK, dtype=torch.int32, device="cuda") + .reshape(TOKEN, TOPK) + .remainder_(E) + ) + topk_weights = torch.full( + (TOKEN, TOPK), 1 / TOPK, dtype=torch.float32, device="cuda" + ) + + def quantize(weight): + cols = weight.shape[-1] + packed, scales = dynamic_mxfp4_quant(weight.reshape(-1, cols)) + packed = packed.view(torch.uint8).reshape( + weight.shape[0], weight.shape[1], cols // 2 + ) + scales = scales.view(torch.uint8).reshape(-1, cols // 32) + return packed.view(torch.float4_e2m1fn_x2), scales + + w13_q, w13_scale = quantize(w13) + w2_q, w2_scale = quantize(w2) + + def run(gate_mode): + interleave = gate_mode == GateMode.INTERLEAVE + return fused_moe( + x, + shuffle_weight( + w13_q, is_guinterleave=interleave, gate_up=True + ), + shuffle_weight( + w2_q, is_guinterleave=interleave, gate_up=False + ), + topk_weights, + topk_ids, + activation=ActivationType.Swiglu, + quant_type=QuantType.per_1x32, + w1_scale=shuffle_scale(w13_scale, E, interleave, True), + w2_scale=shuffle_scale(w2_scale, E, interleave, False), + swiglu_limit=7.0, + gate_mode=gate_mode.value, + ) + + try: + with torch.inference_mode(): + separated = run(GateMode.SEPARATED) + interleaved = run(GateMode.INTERLEAVE) + assert torch.count_nonzero(separated) > 0 + torch.testing.assert_close( + interleaved.float(), separated.float(), rtol=0.02, atol=0.05 + ) + finally: + get_2stage_cfgs.cache_clear() + + def _undo_shuffle_scale(scale): sm, sn = scale.shape return ( From 5cf012de06b6cbffd25bb9cb460abf61dbc541a5 Mon Sep 17 00:00:00 2001 From: Felix Marty Date: Wed, 30 Sep 2026 13:31:38 -0500 Subject: [PATCH 4/5] fix tolerance --- op_tests/test_moe_mxfp4_inter_dim_dispatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/op_tests/test_moe_mxfp4_inter_dim_dispatch.py b/op_tests/test_moe_mxfp4_inter_dim_dispatch.py index f26df0081c2..f6ab903761e 100644 --- a/op_tests/test_moe_mxfp4_inter_dim_dispatch.py +++ b/op_tests/test_moe_mxfp4_inter_dim_dispatch.py @@ -355,7 +355,7 @@ def run(gate_mode): interleaved = run(GateMode.INTERLEAVE) assert torch.count_nonzero(separated) > 0 torch.testing.assert_close( - interleaved.float(), separated.float(), rtol=0.02, atol=0.05 + interleaved.float(), separated.float(), rtol=0, atol=0.1 ) finally: get_2stage_cfgs.cache_clear() From acf58d6f03087c3e56dcdef185d0f82b656b1737 Mon Sep 17 00:00:00 2001 From: Felix Marty Date: Wed, 30 Sep 2026 14:08:19 -0500 Subject: [PATCH 5/5] lint --- op_tests/test_moe_mxfp4_inter_dim_dispatch.py | 25 +++++++++---------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/op_tests/test_moe_mxfp4_inter_dim_dispatch.py b/op_tests/test_moe_mxfp4_inter_dim_dispatch.py index f6ab903761e..5e1c68767ee 100644 --- a/op_tests/test_moe_mxfp4_inter_dim_dispatch.py +++ b/op_tests/test_moe_mxfp4_inter_dim_dispatch.py @@ -292,13 +292,16 @@ def test_mxfp4_bf16_swiglu_interleaved_matches_separated(monkeypatch): monkeypatch.setenv("GPTOSS_SWIGLU_MXFP4_BF16_BOUND", "2147483647") get_2stage_cfgs.cache_clear() model_dim, inter_dim = 1024, 256 - assert _stage_backend( - _dispatch( - model_dim=model_dim, - inter_dim=inter_dim, - gate_mode=GateMode.INTERLEAVE, - ).stage1 - ) == "cktile" + assert ( + _stage_backend( + _dispatch( + model_dim=model_dim, + inter_dim=inter_dim, + gate_mode=GateMode.INTERLEAVE, + ).stage1 + ) + == "cktile" + ) torch.manual_seed(0) x = torch.randn((TOKEN, model_dim), dtype=torch.bfloat16, device="cuda") @@ -333,12 +336,8 @@ def run(gate_mode): interleave = gate_mode == GateMode.INTERLEAVE return fused_moe( x, - shuffle_weight( - w13_q, is_guinterleave=interleave, gate_up=True - ), - shuffle_weight( - w2_q, is_guinterleave=interleave, gate_up=False - ), + shuffle_weight(w13_q, is_guinterleave=interleave, gate_up=True), + shuffle_weight(w2_q, is_guinterleave=interleave, gate_up=False), topk_weights, topk_ids, activation=ActivationType.Swiglu,