SIGN IN SIGN UP

[rocm-libraries] ROCm/rocm-libraries#9396 (commit e1aa8ce)

feat(ck): Added Gelu with Tanh approx to XDL 2-stage MoE
 epilogue (#9396)

## Motivation

Enable the tanh-approximation GELU activation `(gelu_tanh, 0.5 * x * (1
+ tanh(sqrt(2/pi) * (x + 0.044715 * x^3))))` in the Composable Kernel
XDL 2-stage MoE path. The MoE gridwise kernel epilogue currently
supports only `silu/gelu/swiglustep/swiglu_oai`; this adds
`gelu_tanh_and_mul` so models whose MoE experts use the GELU tanh
approximation (e.g. [Gemma-family
MoE](https://huggingface.co/google/gemma-4-26B-A4B/blob/main/config.json))
can use this path.

JIRA ID : ROCM-27619

## Technical Details

- `gridwise_gemm_xdl_cshuffle_common.hpp`: add
`Activation::gelu_tanh_and_mul = 4` to the activation enum.
- `gridwise_moe_gemm.hpp`, `gridwise_moe_gemm_blockscale.hpp`,
`gridwise_moe_mx_gemm_<>.hpp`: wire `gelu_tanh_and_mul` into epilogue
paths, delegating to the existing
`ck::tensor_operation::element_wise::FastGelu` helper (the single source
of truth for the tanh-GELU math, `FastGelu(gate) * up`). Also added
`static_assert` for validation of supported activations
- The activation is applied in fp32 in the epilogue and is orthogonal to
the GEMM compute (MFMA/tile/pipeline untouched) and to quantization
(existing per-token dequant reused). Only the non-blockscale gridwise
kernel is changed.
- Then I plan to port these changes to AITER after
https://github.com/ROCm/aiter/pull/3886 to avoid merge conflicts

## Test Plan

Use `ActOP = 4` in the example `moe_gemm1_xdl_fp8`, rebuild example and
launch ctest

## Test Result

```
ctest -R "^example_moe_gemm1_xdl_fp8$" -V'

Constructing a list of tests
Done constructing a list of tests
Updating test list for fixtures
Added 0 tests to meet fixture requirements
Checking test dependency graph...
Checking test dependency graph end
test 257
    Start 257: example_moe_gemm1_xdl_fp8

257: Test command: example_moe_gemm1_xdl_fp8
257: Working Directory: example/65_gemm_multiply_multiply
257: Test timeout computed to be: 1500
257: a0_t_k: dim 2, lengths {16384, 6144}, strides {6144, 1}
257: b0_e_n_k: dim 3, lengths {8, 6144, 8192}, strides {50331648, 1, 6144}
257: d1_e_n: dim 2, lengths {8, 8192}, strides {8192, 1}
257: d2_e_n: dim 2, lengths {32768, 4096}, strides {1, 0}
257: d0_t_n: dim 2, lengths {16384, 4096}, strides {1, 16384}
257: d2_e_n: dim 2, lengths {32768, 4096}, strides {1, 0}
257: e_t_n: dim 3, lengths {16384, 2, 4096}, strides {8192, 4096, 1}
1/1 Test #257: example_moe_gemm1_xdl_fp8 ........   Passed   83.51 sec

The following tests passed:
        example_moe_gemm1_xdl_fp8

100% tests passed, 0 tests failed out of 1

Label Time Summary:
SMOKE_TEST    =  83.51 sec*proc (1 test)

Total Test time (real) =  83.59 sec
```

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
A
Alexandra Sidorova committed
f57307c82bf0f2fb9447843d8ae566d09409504c
Parent: a39b67b
Committed by assistant-librarian[bot] <assistant-librarian[bot]@users.noreply.github.com> on 8/5/2026, 7:16:09 PM