[CK] [bug fix] Fix wrong cktile 2stages moe_gemm1_heuristic_dispatcher dispatch in case a template is not built, add template specialization for (Swiglu, HasBias=false, split_k=1)

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

csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages.h

@@ -53,7 +53,19 @@ template <typename ADataType,
           int split_k>
 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_gt_1=",
+                    split_k,
+                    ". Consider adding the specialization in "
+                    "csrc/ck_tile_gemm_moe_2stages/gen_instances.py");
+        return {};
+    }
 };
 
 template <typename ADataType,
@@ -65,7 +77,19 @@ template <typename ADataType,
           int split_k>
 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_gt_1=",
+                    split_k,
+                    ". Consider adding the specialization in "
+                    "csrc/ck_tile_gemm_moe_2stages/gen_instances.py");
+        return {};
+    }
 };
 
 __attribute__((visibility("default"))) torch::Tensor

csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu

@@ -243,6 +243,14 @@ 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);
+    return {};
 }
 
 torch::Tensor cktile_moe_gemm1(torch::Tensor& XQ,

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,113 @@ 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, atol=0.1
+        )
+    finally:
+        get_2stage_cfgs.cache_clear()
+
+
 def _undo_shuffle_scale(scale):
     sm, sn = scale.shape
     return (