ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-09 08:37:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#8944 (commit 7be2dbb)

Browse Source 
feat(ck): Add swiglu_oai (OAI SwiGLU) activation to XDL
 2-stage MoE epilogue. (#8944)

## 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.

JIRA ID : ROCM-27213

## 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 ROCm/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

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
This commit is contained in:
Linjun-AMD
2026-07-01 12:36:31 +00:00
committed by assistant-librarian[bot]
parent a65244b86c
commit f4e6fad973
6 changed files with 189 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_common.hpp
View File

@@ -30,9 +30,28 @@ enum Activation
{
gelu_and_mul = 0,
silu_and_mul = 1,
swiglustep_and_mul = 2
swiglustep_and_mul = 2,
swiglu_oai_and_mul = 3
};
// OAI / gpt-oss SwiGLU activation: gate * sigmoid(alpha * gate) * (up + 1), with a
// pre-activation clamp (gate upper-bounded to limit, up symmetric in [-limit, limit]).
// Defaults limit = 7.0, alpha = 1.702 per gpt-oss. Same math as ck_tile::moe::Swiglu, but
// the sigmoid here uses a plain fp32 division (not __builtin_amdgcn_rcpf), so results match
// to ~1e-6 rather than bit-exact. Single source of truth shared by the MoE kernel epilogue
// (swiglu_oai_and_mul) and its host unit test.
// Not constexpr: it always calls math::exp (non-constexpr), so it can never be evaluated in a
// constant expression. inline keeps it ODR-safe in this header. The kernel epilogue and the host
// unit test only call it at runtime.
__host__ __device__ inline float
swiglu_oai(float gate, float up, float limit = 7.0f, float alpha = 1.702f)
{
gate = math::min(gate, limit); // gate <= limit
up = math::min(math::max(up, -limit), limit); // up in [-limit, limit]
const float sig = 1.0f / (1.0f + math::exp(alpha * -gate)); // sigmoid(alpha * gate)
return gate * sig * (up + 1.0f); // OAI form
}
template <typename ALayout,
typename BLayout,
typename ELayout,

83
include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp
View File

@@ -297,6 +297,25 @@ struct GridwiseMoeGemm : public GridwiseGemm_xdl_cshuffle_base<
up = math::min(math::max(up, -kSwiGluClamp), kSwiGluClamp);
return gate * up;
}
// Clamp limit for swiglu_oai_and_mul (gpt-oss / OAI form): gate clamped to <= L,
// up clamped to [-L, L]; L hardcoded to 7.0. alpha = 1.702 per gpt-oss default.
static constexpr float kSwiGluOaiLimit = 7.0f;
static constexpr float kSwiGluOaiAlpha = 1.702f;
// Helper: apply OAI SwiGLU activation gate*sigmoid(alpha*gate)*(up+1) with pre-activation
// clamp (gate upper-bounded, up symmetric). Same math as ck_tile::moe::Swiglu in
// ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp, but the sigmoid here uses a plain fp32
// division (not __builtin_amdgcn_rcpf), so results match to ~1e-6 rather than bit-exact.
// Distinct from swiglustep (no +1, no alpha).
// Used by all four swiglu_oai_and_mul epilogue paths (quant/non-quant x pipeline-A/B).
__host__ __device__ static constexpr float apply_swiglu_oai_activation(float gate, float up)
{
// Delegate to the shared ck::swiglu_oai (defined in gridwise_gemm_xdl_cshuffle_common.hpp)
// so the kernel epilogue and the host unit test exercise a single implementation.
return ck::swiglu_oai(gate, up, kSwiGluOaiLimit, kSwiGluOaiAlpha);
}
using mfma_selector =
MfmaSelector<ComputeTypeA, MPerXdl, NPerXdl, ComputeTypeB, is_single_rate_mfma>;
@@ -1506,6 +1525,26 @@ struct GridwiseMoeGemm : public GridwiseGemm_xdl_cshuffle_base<
}
c_thread_buf_fp32(cidx) = apply_swiglustep_activation(gate, up);
}
else if constexpr(ActivationOperation ==
Activation::swiglu_oai_and_mul)
{
const float scale_up =
p_scale_b[(n0 * NWave * NPerXdl + problem.N) *
PerTokenQuant];
float gate = scale_a * scale_b * c_thread_buf[cidx];
float up = scale_a * scale_up * c_thread_buf_up[cidx];
if constexpr(MulRoutedWeight)
{
gate = gate * topk_weights.template AsType<float>()[m4];
up = up * topk_weights.template AsType<float>()[m4];
}
if constexpr(is_same_v<remove_cvref_t<BDataType>, pk_i4_t>)
{
gate *= 16;
up *= 16;
}
c_thread_buf_fp32(cidx) = apply_swiglu_oai_activation(gate, up);
}
}
else
{
@@ -1579,6 +1618,18 @@ struct GridwiseMoeGemm : public GridwiseGemm_xdl_cshuffle_base<
}
c_thread_buf_fp32(cidx) = apply_swiglustep_activation(gate, up);
}
else if constexpr(ActivationOperation ==
Activation::swiglu_oai_and_mul)
{
float gate = c_thread_buf[cidx];
float up = c_thread_buf_up[cidx];
if constexpr(MulRoutedWeight)
{
gate = gate * topk_weights.template AsType<float>()[m4];
up = up * topk_weights.template AsType<float>()[m4];
}
c_thread_buf_fp32(cidx) = apply_swiglu_oai_activation(gate, up);
}
}
else
{
@@ -2011,6 +2062,26 @@ struct GridwiseMoeGemm : public GridwiseGemm_xdl_cshuffle_base<
}
c_thread_buf_fp32(cidx) = apply_swiglustep_activation(gate, up);
}
else if constexpr(ActivationOperation ==
Activation::swiglu_oai_and_mul)
{
const float scale_up =
p_scale_b[(n0 * NWave * NPerXdl + problem.N) *
PerTokenQuant];
float gate = scale_a * scale_b * c_thread_buf[cidx];
float up = scale_a * scale_up * c_thread_buf_up[cidx];
if constexpr(MulRoutedWeight)
{
gate = gate * topk_weights.template AsType<float>()[m4];
up = up * topk_weights.template AsType<float>()[m4];
}
if constexpr(is_same_v<remove_cvref_t<BDataType>, pk_i4_t>)
{
gate *= 16;
up *= 16;
}
c_thread_buf_fp32(cidx) = apply_swiglu_oai_activation(gate, up);
}
}
else
{
@@ -2084,6 +2155,18 @@ struct GridwiseMoeGemm : public GridwiseGemm_xdl_cshuffle_base<
}
c_thread_buf_fp32(cidx) = apply_swiglustep_activation(gate, up);
}
else if constexpr(ActivationOperation ==
Activation::swiglu_oai_and_mul)
{
float gate = c_thread_buf[cidx];
float up = c_thread_buf_up[cidx];
if constexpr(MulRoutedWeight)
{
gate = gate * topk_weights.template AsType<float>()[m4];
up = up * topk_weights.template AsType<float>()[m4];
}
c_thread_buf_fp32(cidx) = apply_swiglu_oai_activation(gate, up);
}
}
else
{

6
include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm_blockscale.hpp
View File

@@ -1140,6 +1140,9 @@ struct GridwiseMoeGemmBlockScale
BElementwiseOperation b_element_op,
CElementwiseOperation c_element_op)
{
static_assert(ActivationOperation != Activation::swiglu_oai_and_mul,
"gridwise_moe_gemm_blockscale does not support swiglu_oai_and_mul; use the "
"non-blockscale gridwise_moe_gemm.");
#if defined(__gfx942__) || defined(__gfx950__)
constexpr auto b_coherence_flag = NonTemporalLoadB
? AmdBufferCoherenceEnum::WAVE_NT1
@@ -1694,6 +1697,9 @@ struct GridwiseMoeGemmBlockScale
BElementwiseOperation b_element_op,
CElementwiseOperation c_element_op)
{
static_assert(ActivationOperation != Activation::swiglu_oai_and_mul,
"gridwise_moe_gemm_blockscale does not support swiglu_oai_and_mul; use the "
"non-blockscale gridwise_moe_gemm.");
#if defined(__gfx942__) || defined(__gfx950__)
constexpr auto b_coherence_flag = NonTemporalLoadB
? AmdBufferCoherenceEnum::WAVE_NT1