[MoE] Added Gelu with tanh approx for CK XDL 2-stage MoE

aiter/ops/quant.py

@@ -682,6 +682,7 @@ def get_torch_act(aType):
         ActivationType.No: lambda *a, **k: a[0],
         ActivationType.Silu: F.silu,
         ActivationType.Gelu: F.gelu,
+        ActivationType.GeluTanh: lambda x: F.gelu(x, approximate="tanh"),
     }
     return tmp.get(aType, NotImplementedError)
 

aiter/utility/dtypes.py

@@ -163,5 +163,11 @@ def _convert(s):
 
 
 def str2ActivationType(s):
-    """Convert string to ActivationType."""
+    members = getattr(ActivationType, "__members__", None)
+    if members is not None:
+        s_lower = s.lower()
+        for name, member in members.items():
+            if name.lower() == s_lower:
+                return member
+        raise argparse.ArgumentTypeError(f"invalid activation type: {s}")
     return getattr(ActivationType, s.capitalize())

csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu

@@ -14,7 +14,8 @@
 using MoeKernelMap = std::unordered_map<std::string, MoeKernel>;
 
 // Map aiter ActivationType (stage1) to the CK stage1 activation op value.
-// CK act values: gelu=0, silu=1, swiglu(OAI swiglu_oai)=3. Value 2 is reserved.
+// CK act values: gelu=0, silu=1, swiglu(OAI swiglu_oai)=3, gelu_tanh_and_mul=4.
+// Value 2 is reserved.
 static inline int map_activation_to_ck_stage1(int activation)
 {
     switch (static_cast<ActivationType>(activation))
@@ -25,6 +26,8 @@ static inline int map_activation_to_ck_stage1(int activation)
         return 0;
     case ActivationType::Swiglu:
         return 3;
+    case ActivationType::GeluTanh:
+        return 4;
     default:
         TORCH_CHECK(false, "Unsupported activation for ck_moe_stage1: ", activation);
         return -1;

csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages_common.py

@@ -18,8 +18,9 @@
 from chip_info import get_gfx
 
 # CK stage1 activation op values. swiglu (OAI swiglu_oai) maps to CK act value 3;
-# value 2 is intentionally skipped (reserved for a future activation variant).
-ACT_OP_MAP = {"gelu": 0, "silu": 1, "swiglu": 3}
+# gelutanh (gelu_tanh_and_mul) maps to CK act value 4; value 2 is intentionally
+# skipped (reserved for a future activation variant).
+ACT_OP_MAP = {"gelu": 0, "silu": 1, "swiglu": 3, "gelutanh": 4}
 ACT_OP_NAME = {v: k for k, v in ACT_OP_MAP.items()}
 
 

csrc/ck_gemm_moe_2stages_codegen/gen_instances.py

@@ -1074,7 +1074,7 @@ def generate_instance_and_lookUpTable(self):
         default="silu",
         required=False,
         type=str,
-        choices=["silu", "gelu", "swiglu"],
+        choices=["silu", "gelu", "swiglu", "gelutanh"],
         help="select activation",
     )
 
@@ -1243,6 +1243,41 @@ def generate_instance_and_lookUpTable(self):
                 False,  # splitk
             )
             codegen.generate_instance_and_lookUpTable()
+
+        # gelutanh (gelu_tanh_and_mul) plain-f8 quant moe (per_tensor / per_token = PTPC).
+        # Same rationale as the swiglu loop above: the general quant loop uses
+        # acts=["silu","gelu"] and so ships 0 gelutanh x plain-f8 instances in the
+        # AOT/wheel prebuild. Runtime JIT (Path B, gen_func passes -b f8) already
+        # generates them on demand, so this loop only extends the AOT prebuild set
+        # and removes the ~first-run cold compile -- it is NOT required for
+        # correctness. Mirrors the general loop's plain-f8 path (f8 x f8, tag=a8w8,
+        # CDEElementOp=MulABScale, plain gridwise_moe_gemm.hpp). act="gelutanh" ->
+        # CK ActOP=4. preshuffle forced True for plain f8. Targets Gemma-family MoE
+        # experts that use the tanh-approx GELU. Independent loop so existing
+        # silu/gelu/swiglu/no-quant coverage is untouched.
+        gelutanh_plain_c_dtypes = ["f16", "b16"]
+        gelutanh_plain_quant_l = ["per_tensor", "per_token"]
+        for (
+            c_dtype,
+            routed_weight,
+            quant,
+        ) in itertools.product(
+            gelutanh_plain_c_dtypes,
+            routed_weight_l,
+            gelutanh_plain_quant_l,
+        ):
+            codegen = ck_moe_2stage_gemm_codegen(
+                args.working_path,
+                "f8",  # a_dtype
+                "f8",  # b_dtype
+                c_dtype,
+                quant_dict[quant],
+                "gelutanh",
+                routed_weight,
+                True,  # preshuffle (plain f8 forced True, mirrors general loop)
+                False,  # splitk
+            )
+            codegen.generate_instance_and_lookUpTable()
     else:
         for b_dtype in args.b_dtype:
             a_dtype = b_dtype if b_dtype != "i4" else "f8"

csrc/include/aiter_enum.h

@@ -6,11 +6,12 @@
 
 enum class ActivationType : int
 {
-    No     = -1,
-    Silu   = 0,
-    Gelu   = 1,
-    Swiglu = 2,
-    Situv2 = 3,
+    No       = -1,
+    Silu     = 0,
+    Gelu     = 1,
+    Swiglu   = 2,
+    Situv2   = 3,
+    GeluTanh = 4,
 };
 
 enum class QuantType : int

csrc/include/rocm_ops.hpp

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