[CK][MoE] Add swiglu_oai (OAI SwiGLU) activation to XDL 2-stage MoE epilogue

[LJ-underdog]

Motivation

Enable the OAI-form SwiGLU activation (swiglu_oai, gate * sigmoid(1.702 * gate) * (up + 1), gpt-oss style) in the Composable Kernel XDL 2-stage MoE path. The MoE gridwise kernel epilogue currently supports only silu/gelu; this adds swiglu_oai so OAI-style MoE models can use this path.

Technical Details

• gridwise_gemm_xdl_cshuffle_common.hpp: add Activation::swiglu_oai_and_mul = 3.
• gridwise_moe_gemm.hpp: add the apply_swiglu_oai_activation helper (gate * sigmoid(1.702 * gate) * (up + 1), clamp gate <= 7 and up in [-7, 7], OAI/gpt-oss form) and wire it into all 4 epilogue paths (quant + non-quant x Run / Run_2Lds).
• 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.
• Consumed by aiter via [MoE] Add swiglu_oai (OAI SwiGLU) for per-token fp8 CK XDL 2-stage MoE aiter#3886 (dispatch + codegen); review/merge together.

Test Plan

Validate the new epilogue branch against a torch fp32 OAI-SwiGLU reference through the aiter per-token fp8 MoE path (op-isolate on gfx942 / MI308X).

Test Result

cos_sim = 0.999993 vs the torch fp32 OAI-SwiGLU reference; no NaN. Confirmed the per-token fp8 path dispatches to this GridwiseMoeGemm kernel (rocprofv3) and runs the swiglu_oai epilogue branch.

Submission Checklist

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

[Copilot]

Pull request overview

Adds support for the OAI-form SwiGLU activation (swiglu_oai_and_mul) to the Composable Kernel XDL 2-stage MoE epilogue so OAI/gpt-oss style MoE models can use this path.

Changes:

• Extend the shared Activation enum with swiglu_oai_and_mul = 3.
• Introduce apply_swiglu_oai_activation(gate, up) implementing gate * sigmoid(1.702 * gate) * (up + 1) with the specified clamps.
• Wire the new activation into the four MoE epilogue paths (quant/non-quant × Run/Run_2Lds).

Reviewed changes

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

File	Description
projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp	Adds the OAI SwiGLU helper and integrates the new activation into all relevant epilogue branches.
projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_common.hpp	Extends the shared activation enum to include the new OAI SwiGLU option.

---

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