[HIP] [CK] [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.

[a-sidorova]

Hi @yzhou103! Could you please review the PR?

There are three failed jobs - they seems like sporadic failures since they're not related to my changes. Unfortunately I don't have permissions to relaunch CI jobs. I will be really appreciate for the help with CI!
FYI: I've already tried to rebase to the man branch - a few failed CI jobs are disappeared, but not all of them.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.

Suppressed comments (2)

aiter/utility/dtypes.py:169

• str2ActivationType() still fails to resolve mixed-case enum members (e.g. CLI gelutanh -> GeluTanh) if ActivationType does not provide __members__ (pybind11 enums in this repo typically don’t). In that case the fallback s.capitalize() produces Gelutanh, which won’t match GeluTanh, so argparse will reject a valid activation.

def str2ActivationType(s):
    s = str(s)
    members = getattr(ActivationType, "__members__", None)
    if members is not None:
        s_lower = s.lower()

aiter/ops/quant.py:687

• get_torch_act() currently returns the exception type NotImplementedError when aType isn’t in the map, which is callable and will return an exception instance rather than raising—leading to confusing downstream errors when the caller tries to use it as an activation function. It’s safer to raise immediately with a clear message.

        ActivationType.Silu: F.silu,
        ActivationType.Gelu: F.gelu,
        ActivationType.GeluTanh: lambda x: F.gelu(x, approximate="tanh"),
    }
    return tmp.get(aType, NotImplementedError)

[a-sidorova]

@zufayu @yzhou103 hello! I've managed to green CI after several rebases. Could you please review this PR? Everything is ready for review

[simondanielsson]

Thanks for the good work, we're running a similar thing on a custom aiter wheel with a customer in prod for Gemma4 and would be great to see this PR land asap

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.

Suppressed comments (4)

aiter/fused_moe.py:3847

• ck_moe_stage1's split-k post-activation path does not handle ActivationType.GeluTanh, so per_1x128 splitk runs with gelutanh will raise at runtime even though the CK epilogue supports it for non-splitk. Handle GeluTanh explicitly using aiter.gelu_tanh_and_mul.

        valid_out = tmp_out[: token_num * topk, :]
        if activation == ActivationType.Silu:
            aiter.silu_and_mul(out, valid_out.view(dtypes.fp32))
        elif activation == ActivationType.Gelu:
            aiter.gelu_and_mul(out, valid_out.view(dtypes.fp32))
        else:
            raise ValueError(
                f"Unsupported activation for split-k post-activation: {activation}"
            )

aiter/fused_moe.py:4003

• cktile_moe_stage1 non-interleaved split-k post-activation does not handle ActivationType.GeluTanh, so GeluTanh will raise for this path. Add a GeluTanh branch using aiter.gelu_tanh_and_mul for the (out, valid_out) activation+mul.

            elif activation == ActivationType.Silu:
                aiter.silu_and_mul(out, valid_out)
            elif activation == ActivationType.Swiglu:
                aiter.swiglu_and_mul(out, valid_out)
            elif activation == ActivationType.Gelu:
                aiter.gelu_and_mul(out, valid_out)
            else:
                raise ValueError(
                    f"Unsupported activation for split-k post-activation: {activation}"
                )

aiter/fused_moe.py:3976

• cktile_moe_stage1 interleaved split-k post-activation has no GeluTanh handling, so GeluTanh requests will raise even though the math is straightforward (tanh-approx GELU on the gate, then multiply by up). Add a GeluTanh branch using torch.nn.functional.gelu(..., approximate="tanh").

            elif activation == ActivationType.Gelu:
                NLane = 16
                N0 = inter_dim // NLane
                flat = valid_out.view(-1, N0, 2, NLane)
                gate = flat[:, :, 0, :].reshape(-1, inter_dim)
                up = flat[:, :, 1, :].reshape(-1, inter_dim)
                out.view(-1, inter_dim).copy_(torch.nn.functional.gelu(gate) * up)
            else:
                raise ValueError(
                    f"Unsupported activation for interleaved split-k "
                    f"post-activation: {activation}"
                )

aiter/fused_moe.py:3415

• The split-k post-activation path in asm_stage1 does not handle ActivationType.GeluTanh, so using the new activation with ksplit>0 will raise at runtime. Add an explicit GeluTanh branch and use the existing aiter.gelu_tanh_and_mul implementation.

This issue also appears in the following locations of the same file:

• line 3839
• line 3965
• line 3994

        if activation == ActivationType.Silu:
            aiter.silu_and_mul(out, tmp_out.view(dtypes.fp32))
        elif activation == ActivationType.Swiglu:
            aiter.swiglu_and_mul(out, tmp_out.view(dtypes.fp32))
        elif activation == ActivationType.Gelu:
            aiter.gelu_and_mul(out, tmp_out.view(dtypes.fp32))
        else:
            raise ValueError(
                f"Unsupported activation for split-k post-activation: {activation}"
            )

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

aiter/fused_moe.py:2329

• In the tier-fallback loop, keys_fb is built twice (first from keys, then immediately overwritten from lookup_keys) and logged twice. The first assignment is dead code and the duplicate INFO logs add overhead. Build keys_fb once from lookup_keys and avoid INFO logging inside this hot loop.

                    keys_disabled[:2] + (fallback_tier,) + keys_disabled[3:]
                )
                result = primary.get(keys_fb, None)
                if result is None:

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

[a-sidorova]

@yzhou103 hello! The CI is green, your comments are applied. Could you please take a look again? Thank you in advance

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

[yzhou103]

LGTM

[a-sidorova]

Waiting for additional approve for merging with bumping CK