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

[a-sidorova]

Motivation

Enable the tanh-approximation GELU activation (gelu_tanh, 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))))) in the AITER Composable Kernel XDL 2-stage MoE path. The MoE gridwise epilogue supports only silu/gelu/swiglu; this ports CK's gelu_tanh_and_mul (ROCm/rocm-libraries#9396) into AITER so models whose MoE experts use the GELU tanh approximation (e.g. Gemma-family MoE) can run on this path.

JIRA ID : ROCM-27619

Technical Details

Wires a new GeluTanh activation end-to-end into the CK 2-stage MoE codegen, following the existing swiglu_oai port (78e45124):

• csrc/include/aiter_enum.h: add ActivationType::GeluTanh = 4.
• csrc/include/rocm_ops.hpp: expose GeluTanh on the pybind ActivationType enum.
• csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu: map ActivationType::GeluTanh -> 4 (CK gelu_tanh_and_mul) in map_activation_to_ck_stage1.
• gemm_moe_ck2stages_common.py: add "gelutanh": 4 to ACT_OP_MAP (drives CK ActOP codegen and kernel-instance naming).
• gen_instances.py: add gelutanh to the -act choices and a targeted plain-f8 AOT prebuild loop (mirrors the swiglu loop; not required for correctness since runtime JIT generates on demand).

CK submodule bump

• Bumps 3rdparty/composable_kernel to the merged feat(ck): Added Gelu with Tanh approx to XDL 2-stage MoE epilogue rocm-libraries#9396, which provides the CK-side gelu with tanh approx epilogue. This PR must be built / merged together with that CK commit.

Test/reference support so the change can be validated:

• aiter/ops/quant.py (get_torch_act): add the GeluTanh reference using the tanh-approx GELU (F.gelu(x, approximate="tanh")) to match CK's FastGelu epilogue.
• aiter/utility/dtypes.py (str2ActivationType): case-insensitive enum lookup so mixed-case members like GeluTanh resolve from CLI (-a gelutanh).

The activation is applied in fp32 in the stage-1 epilogue; GEMM compute and quantization are untouched. CK-side support comes from ROCm/rocm-libraries#9396 (gelu_tanh_and_mul = 4).

Test Plan

Run the CK 2-stage MoE test with the new activation on the plain-f8 (a8w8) path:

python3 op_tests/test_moe_2stage.py -a gelutanh -q 0  # bf16
python3 op_tests/test_moe_2stage.py -a gelutanh -q 1 --no-flydsl-csv   # per_tensor fp8
python3 op_tests/test_moe_2stage.py -a gelutanh -q 2 --no-flydsl-csv   # per_token  fp8

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4620 --add-label <label>

[Copilot]

Pull request overview

This PR adds end-to-end support for the tanh-approximation GELU activation (GeluTanh) in the Composable Kernel (CK) XDL 2-stage MoE path, including C++/pybind enum exposure, CK stage-1 activation mapping, and codegen/test reference wiring.

Changes:

• Introduces ActivationType::GeluTanh = 4 and exposes it via pybind.
• Maps GeluTanh to CK’s stage-1 act_op=4 (gelu_tanh_and_mul) and wires "gelutanh": 4 into CK 2-stage codegen.
• Adds a torch reference implementation (F.gelu(..., approximate="tanh")) and improves CLI activation parsing to be case-insensitive.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
csrc/include/aiter_enum.h	Adds the new ActivationType::GeluTanh enum value.
csrc/include/rocm_ops.hpp	Exposes GeluTanh through the pybind ActivationType enum.
csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu	Maps AITER activation to CK stage-1 act_op, adding GeluTanh -> 4.
csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages_common.py	Extends ACT_OP_MAP/naming to include gelutanh: 4.
csrc/ck_gemm_moe_2stages_codegen/gen_instances.py	Adds gelutanh CLI option and an AOT prebuild loop for plain-f8 quant instances.
aiter/utility/dtypes.py	Makes activation parsing case-insensitive for mixed-case enum members like GeluTanh.
aiter/ops/quant.py	Adds torch reference mapping for GeluTanh using tanh-approx GELU.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.