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 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..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,7 +53,19 @@ 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_gt_1=", + split_k, + ". Consider adding the specialization in " + "csrc/ck_tile_gemm_moe_2stages/gen_instances.py"); + return {}; + } }; 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_gt_1=", + split_k, + ". Consider adding the specialization in " + "csrc/ck_tile_gemm_moe_2stages/gen_instances.py"); + 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..3a3d55c52be 100644 --- a/csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu +++ b/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, diff --git a/op_tests/test_moe_mxfp4_inter_dim_dispatch.py b/op_tests/test_moe_mxfp4_inter_dim_dispatch.py index 17838b4bb31..5e1c68767ee 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,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 (