From 7b28e13e8145cc78d5b8f77521b1811f9f447ded Mon Sep 17 00:00:00 2001 From: Alexandra Sidorova Date: Fri, 7 Aug 2026 06:11:13 +0000 Subject: [PATCH] [MoE] Added Gelu with tanh approx for CK XDL 2-stage MoE --- 3rdparty/composable_kernel | 2 +- aiter/ops/quant.py | 1 + aiter/utility/dtypes.py | 8 +++- .../gemm_moe_ck2stages.cu | 5 ++- .../gemm_moe_ck2stages_common.py | 5 ++- .../gen_instances.py | 37 ++++++++++++++++++- csrc/include/aiter_enum.h | 11 +++--- csrc/include/rocm_ops.hpp | 1 + 8 files changed, 59 insertions(+), 11 deletions(-) diff --git a/3rdparty/composable_kernel b/3rdparty/composable_kernel index f33252cebe5..f57307c82bf 160000 --- a/3rdparty/composable_kernel +++ b/3rdparty/composable_kernel @@ -1 +1 @@ -Subproject commit f33252cebe5a52362ec1ee12c124dde7800dda3a +Subproject commit f57307c82bf0f2fb9447843d8ae566d09409504c diff --git a/aiter/ops/quant.py b/aiter/ops/quant.py index 8d88585afde..be8e754e83b 100644 --- a/aiter/ops/quant.py +++ b/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) diff --git a/aiter/utility/dtypes.py b/aiter/utility/dtypes.py index 761ecbb966e..276b031c3ee 100644 --- a/aiter/utility/dtypes.py +++ b/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()) diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu index 007dec29481..ad92fb1fc4c 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu @@ -14,7 +14,8 @@ using MoeKernelMap = std::unordered_map; // 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(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; diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages_common.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages_common.py index fdffe39e6f2..941e95f5790 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages_common.py +++ b/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()} diff --git a/csrc/ck_gemm_moe_2stages_codegen/gen_instances.py b/csrc/ck_gemm_moe_2stages_codegen/gen_instances.py index 513ab461ef1..3a8078b5b8e 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gen_instances.py +++ b/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" diff --git a/csrc/include/aiter_enum.h b/csrc/include/aiter_enum.h index 050a6f8ad43..2328ee18e78 100644 --- a/csrc/include/aiter_enum.h +++ b/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 diff --git a/csrc/include/rocm_ops.hpp b/csrc/include/rocm_ops.hpp index 96e99ffcdb4..27f9fb68468 100644 --- a/csrc/include/rocm_ops.hpp +++ b/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_(m, "MlaVersion") \ .value("V32", MlaVersion::V32) \