[HIP] [CK] Re-submit Relu2 Tier 2 MoE GEMM support (clean CK pin)

aiter/fused_moe.py

@@ -799,6 +799,8 @@ def resolve_activation_dtype(
             q_dtype_a = _bound_split(
                 M, _SWIGLU_MXFP4_BF16_BOUND, dtypes.bf16, dtypes.fp4x2
             )
+        elif activation == ActivationType.Relu2:
+            q_dtype_a = dtypes.bf16
         elif activation == ActivationType.Swiglu or gate_mode == GateMode.INTERLEAVE:
             if gfx != "gfx950":
                 q_dtype_a = dtypes.bf16
@@ -2979,6 +2981,31 @@ def get_2stage_cfgs(
         doweight_stage1=doweight_stage1,
         gate_mode=gate_mode,
     )
+    if activation == ActivationType.Relu2:
+        if use_g1u1:
+            raise NotImplementedError(
+                "ActivationType.Relu2 is gate-only (non-gated) and does not "
+                "support use_g1u1=True; pass use_g1u1=False."
+            )
+        if doweight_stage1:
+            raise NotImplementedError(
+                "ActivationType.Relu2 CK-Tile stage1 instances are generated "
+                "with MulRoutedWeight=False; doweight_stage1=True is not "
+                "supported and would silently drop routed weights."
+            )
+        if has_stage1_bias or has_stage2_bias:
+            raise NotImplementedError(
+                "ActivationType.Relu2 CK-Tile path does not support "
+                "per-expert bias (has_stage1_bias/has_stage2_bias); the "
+                "returned metadata hardcodes has_bias=False."
+            )
+        if gate_mode != GateMode.SEPARATED:
+            raise NotImplementedError(
+                "ActivationType.Relu2 CK-Tile stage1 instance "
+                "(kFFN_gemm1_gate_only) only supports the gate-only/"
+                f"separated weight layout; gate_mode={gate_mode.value!r} "
+                "is not supported."
+            )
     # Configs are keyed on (gfx, cu_num, ...) so archs that share a cu_num
     # (e.g. gfx950 vs gfx1250, both report 256 CU) don't collide. Legacy CSVs
     # without a `gfx` column are backfilled from cu_num at load time via
@@ -3204,6 +3231,18 @@ def _lookup_cfg(c2s):
                 f"[fused_moe] discarding Opus tuned config for unsupported "
                 f"activation {activation}; using default heuristics"
             )
+        elif (
+            kn1.startswith("flydsl_") or kn2.startswith("flydsl_")
+        ) and activation not in (
+            ActivationType.Silu,
+            ActivationType.Swiglu,
+            ActivationType.Situv2,
+        ):
+            cfg = None
+            logger.warning(
+                f"[fused_moe] discarding FlyDSL tuned config for unsupported "
+                f"activation {activation}; using default heuristics"
+            )
         elif _disable_inline_sort and _is_inline_sort_cfg(kn1, kn2):
             cfg = None
             logger.warning("[fused_moe] discarding tuned inline-sort config")
@@ -3539,7 +3578,9 @@ def get_block_m() -> int:
     is_opus2 = _opus_a8w4.is_opus_a8w4_stage2_kernel(kernelName2)
     if is_opus1 or is_flydsl1 or is_flydsl2:
         enable_bias = (
-            _needs_swiglu_bias_support(dtype, q_type) and q_dtype_w == dtypes.fp4x2
+            activation == ActivationType.Swiglu
+            and _needs_swiglu_bias_support(dtype, q_type)
+            and q_dtype_w == dtypes.fp4x2
         )
         _s1_fp4q = is_flydsl1 and "_fp4" in kernelName1.split("_t")[-1]
         if is_opus1:
@@ -3678,6 +3719,40 @@ def get_block_m() -> int:
         and is_shuffled
         and cktile_mxfp4_ok
     )
+    if (
+        activation == ActivationType.Relu2
+        and not use_g1u1
+        and q_type == QuantType.per_1x32
+        and q_dtype_w == dtypes.fp4x2
+        and q_dtype_a == dtypes.bf16
+        and dtype == dtypes.bf16
+        and not has_stage2_scatter
+        and not has_activation_scales
+        and bool(opus_weights_shuffled)
+        and get_gfx() == "gfx950"
+    ):
+        _cktile_block_m = 16 if token < 2048 else 32 if token < 16384 else 64
+        return MOEMetadata(
+            functools.partial(
+                cktile_moe_stage1,
+                n_pad_zeros=intermediate_pad // 64 * 64,
+                k_pad_zeros=hidden_pad // 128 * 128,
+                activation=activation,
+                split_k=1,
+                dtype=dtype,
+            ),
+            functools.partial(
+                cktile_moe_stage2,
+                n_pad_zeros=hidden_pad // 64 * 64,
+                k_pad_zeros=intermediate_pad // 128 * 128,
+                activation=activation,
+            ),
+            _cktile_block_m,
+            1,
+            run_1stage,
+            has_bias=False,
+            stage2_has_bias=False,
+        )
     if q_type == QuantType.per_1x32 and q_dtype_w == dtypes.i4x2:
         # Untuned a16wi4 fallback: one shape-safe config on the shared a16w-mix port.
         # Tiles belong in the tuned CSV, not in a heuristic here. ksplit is 0 because
@@ -3849,7 +3924,10 @@ def get_block_m() -> int:
             f"[fused_moe] no tuned FlyDSL config for {keys}, "
             f"using heuristic FlyDSL fallback ({kn1=}, {kn2=})"
         )
-        enable_bias = _needs_swiglu_bias_support(dtype, q_type)
+        enable_bias = (
+            activation == ActivationType.Swiglu
+            and _needs_swiglu_bias_support(dtype, q_type)
+        )
         return MOEMetadata(
             functools.partial(
                 _flydsl_stage1_wrapper,
@@ -4169,7 +4247,8 @@ def fused_moe_2stages(
         and (
             q_dtype_a in [dtypes.bf16, dtypes.fp16]
             and (
-                activation in (ActivationType.Swiglu, ActivationType.Situv2)
+                activation
+                in (ActivationType.Swiglu, ActivationType.Situv2, ActivationType.Relu2)
                 or gate_mode == GateMode.INTERLEAVE
             )
             or (q_dtype_a in [dtypes.fp4x2] and metadata.ksplit > 1 and is_shuffled)
@@ -4409,7 +4488,8 @@ def fused_moe_2stages(
         and w1.dtype == dtypes.fp4x2
         and (
             q_dtype_a in [dtypes.bf16, dtypes.fp16]
-            and activation in (ActivationType.Swiglu, ActivationType.Situv2)
+            and activation
+            in (ActivationType.Swiglu, ActivationType.Situv2, ActivationType.Relu2)
             or (metadata.ksplit > 1 and is_shuffled)
         )
     ):

aiter/ops/quant.py

@@ -709,6 +709,7 @@ def get_torch_act(aType):
         ActivationType.Silu: F.silu,
         ActivationType.Gelu: F.gelu,
         ActivationType.GeluTanh: lambda x: F.gelu(x, approximate="tanh"),
+        ActivationType.Relu2: lambda x: F.relu(x) ** 2,
     }
     return tmp.get(aType, NotImplementedError)
 

csrc/ck_tile_gemm_moe_2stages/gen_instances.py

@@ -101,7 +101,7 @@ def gen_instance(self, k: kernelInstance):
     int KBatch = k_batch.has_value() ? k_batch.value() : 1;
     int stride_A = K;
     int stride_B = K;
-    int stride_C = KBatch > 1 ? N : N / {3 - k.stage}; //gemm1 gate+up need / 2.
+    int stride_C = KBatch > 1 ? N : N / {2 if (k.stage == 1 and k.ActOP != "relu2") else 1}; //gemm1 gate+up need / 2; relu2 is gate-only (non-gated), needs full-width N.
     void *sorted_weights_ptr = topk_weight.has_value() ? topk_weight.value().data_ptr() : nullptr;
 
     {{INSTANCE_CONTENT}}
@@ -190,6 +190,7 @@ def gen_instance(self, k: kernelInstance):
                 ck_tile::tuple<>,
                 row_major,
                 {"ck_tile::MoeFlatmmKind::kFFN_gemm1_split_k" if self.is_split_k else
+                "ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_only" if (k.stage == 1 and self.activation == act_dict["relu2"]) else
                 "ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_up" if k.stage == 1 else
                 "ck_tile::MoeFlatmmKind::kFFN_gemm2"},
                 ck_tile::element_wise::PassThrough,
@@ -283,7 +284,9 @@ 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.activation == 2 and not self.is_split_k else "false"
+                ),
                 "(split_k)": "true" if self.is_split_k else "false",
             }
             format_args = {str(key): value.name for key, value in mapping.items()}
@@ -619,7 +622,7 @@ def generate_common_header(working_path, dispatch_files, manifest_files):
     quant_type = "1x32"
 
     acc_type = "float"
-    act_types = ["silu", "swiglu"]
+    act_types = ["silu", "swiglu", "relu2"]
     c_dtypes = ["bf16"]
     is_split_k_l = [True, False]
 
@@ -650,11 +653,13 @@ def generate_common_header(working_path, dispatch_files, manifest_files):
     for a_type, c_dtype, act_type, is_split_k in itertools.product(
         a_types, c_dtypes, act_types, is_split_k_l
     ):
-        has_bias = act_type == "swiglu"
+        has_bias = act_type == "swiglu" and not is_split_k
 
         # a8w8 do not support
         if a_type in ["fp8", "bf8"] and is_split_k:
             continue
+        if act_type == "relu2" and is_split_k:
+            continue
         codegen = cktile_moe_2stage_gemm_codegen(
             args.working_path,
             a_type,

csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages_common.cuh

@@ -235,7 +235,11 @@ void moe_gemm(const MoeFlatmmHostArgs& args, const ck_stream_config& s)
 
         // TODO: support more act type.
         using FusedAct =
-            std::conditional_t<ActivationOp == 2, ck_tile::moe::Swiglu, ck_tile::moe::MoeSilu>;
+            std::conditional_t<ActivationOp == 2,
+                                ck_tile::moe::Swiglu,
+                                std::conditional_t<ActivationOp == 3,
+                                                    ck_tile::moe::MoeRelu2,
+                                                    ck_tile::moe::MoeSilu>>;
 
         using Kernel = ck_tile::MoeFlatmmKernel<TilePartitioner,
                                                 CodegenFlatmmPipeline,

csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu

@@ -6,6 +6,7 @@
 #include "moe_cktile2stages_name_dispatch.h"
 #include "py_itfs_common.h"
 #include "moe_cktile2stages_heuristic_dispatch_common.h"
+#include "aiter_enum.h"
 #include <cmath>
 
 template <typename ADataType,
@@ -242,7 +243,34 @@ MoeKernel moe_dispatch(int M, int N, int K, int block_m, int activation, bool ha
                                                       false>::dispatch(M, N, K, block_m);
             }
         }
+        else if(activation == static_cast<int>(ActivationType::Relu2) && !has_bias && split_k == 1)
+        {
+            if(stage == 1)
+            {
+                return moe_gemm1_heuristic_dispatcher<ADataType,
+                                                      BDataType,
+                                                      AccDataType,
+                                                      CDataType,
+                                                      3,
+                                                      false,
+                                                      false>::dispatch(M, N, K, block_m);
+            }
+            else
+            {
+                return moe_gemm2_heuristic_dispatcher<ADataType,
+                                                      BDataType,
+                                                      AccDataType,
+                                                      CDataType,
+                                                      3,
+                                                      false,
+                                                      false>::dispatch(M, N, K, block_m);
+            }
+        }
     }
+    TORCH_CHECK(false,
+	            "moe_dispatch: unsupported (activation=", activation,
+	            ", has_bias=", has_bias, ", split_k=", split_k,
+	            ") combination \u2014 no generated kernel instance covers this case");
 }
 
 torch::Tensor cktile_moe_gemm1(torch::Tensor& XQ,

csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages_common.py

@@ -22,6 +22,7 @@
     "silu": 0,
     # "gelu": 1,
     "swiglu": 2,
+    "relu2": 3,
 }
 
 

csrc/include/aiter_enum.h

@@ -12,6 +12,7 @@ enum class ActivationType : int
     Swiglu   = 2,
     Situv2   = 3,
     GeluTanh = 4,
+    Relu2    = 5,
 };
 
 enum class QuantType : int

csrc/include/rocm_ops.hpp

@@ -33,6 +33,7 @@ namespace py = pybind11;
         .value("Swiglu", ActivationType::Swiglu)                                            \
         .value("Situv2", ActivationType::Situv2)                                            \
         .value("GeluTanh", ActivationType::GeluTanh)                                        \
+        .value("Relu2", ActivationType::Relu2)                                              \
         .export_values();                                                                   \
     pybind11::enum_<MlaVersion>(m, "MlaVersion")                                            \
         .value("V32", MlaVersion::V32)                                                      \