[MoE] Add gelu_tanh activation to MoE stage1 GEMM kernels

[jonahbernard]

Motivation

We need MXFP4 MoE kernels that use the gelu tanh activation to run Gemma-based models with high performance in vLLM.

Technical Details

I added 2 functions:

1. _gelu_tanh_mul_vec4 in mixed_moe_gemm_2stage.py, which is called inside def _act_vec4(gate_v4, up_v4)
2. gelu_tanh(x) in moe_gemm_2stage.py and it is called in 2 places: 1) write_row_to_lds and 2) _stage1_store_row. write_row_to_lds and _stage1_store_row are the two stage1 epilogue variants.

For split-K support, in kernels/silu_and_mul_fq.py, I added a compile-time branch that gates on act == gelu_tanh to the build_silu_and_mul_fq_module and skips the silu/swiglu-specific code via continue at the end of that gelu tanh specific branch.

Test Plan

Parameterized the test_moe_gemm_2stage PyTest over silu and gelu_tanh. Additionally, added a branch in the PyTorch reference implementation in torch_moe_gemm1 that runs PyTorch's gelu_tanh when torch_moe_gemm1 is called with activation=gelu_tanh.

To test the split-K path, I added a new test test_silu_and_mul_fq that runs silu_and_mul_fq parameterized over:
activation: silu, swiglu, gelu_tanh
quant_mode: none, fp8, fp4
enable_bias: False, True
gui_layout: False, True
shape: (token_num, topk, inter_dim)

When I was adding this test, I discovered a bug in the kernel: a plain if enable_bias: instead of if const_expr(enable_bias):, so the flyc closure transform didn't capture the bias vars into the nested scf.IfOp blocks

Test Result

Command to run tests:

  pytest tests/kernels/test_moe_gemm.py -k "test_moe_gemm_2stage and gelu_tanh" -v

Results:

== 684 passed, 1620 skipped, 2344 deselected, 1 warning in 150.78s (0:02:30) ===

Command to run split-k tests:

pytest tests/kernels/test_silu_and_mul_fq.py -v

Results:

=== 144 passed, 1 warning in 16.00s ==

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[coderfeli]

Could you help to move the act out of this moe gemm kernel file?

[jonahbernard]

hi @coderfeli !

1. do you want both silu and gelu_tanh moved out of this moe gemm kernel file?
2. which directory and file do you want them moved to?
3. do you want the act kernels moved out of mixed_moe_gemm_2stage.py as well?

thank you!