[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