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

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.

test/CMakeLists.txt

@@ -350,3 +350,4 @@ add_subdirectory(synchronization)
 add_subdirectory(gpu_reference)
 add_subdirectory(util)
 add_subdirectory(gpu_verification)
+add_subdirectory(swiglu_oai_activation)

test/swiglu_oai_activation/CMakeLists.txt

@@ -0,0 +1,7 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+add_gtest_executable(test_swiglu_oai_activation test_swiglu_oai_activation.cpp)
+if(result EQUAL 0)
+    target_link_libraries(test_swiglu_oai_activation PRIVATE utility)
+endif()

test/swiglu_oai_activation/test_swiglu_oai_activation.cpp

@@ -0,0 +1,72 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include <algorithm>
+#include <cmath>
+
+#include "gtest/gtest.h"
+
+// ck::swiglu_oai is the single source of truth shared by the XDL 2-stage MoE epilogue
+// (Activation::swiglu_oai_and_mul) and this host unit test.
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_common.hpp"
+
+namespace {
+
+// Independent fp64 reference for the OAI / gpt-oss SwiGLU activation:
+//   gate * sigmoid(alpha * gate) * (up + 1)
+// with gate clamped to <= limit and up clamped to [-limit, limit]. Computed in double so
+// we are not comparing the fp32 implementation under test against a copy of itself.
+double ref_swiglu_oai(double gate, double up, double limit = 7.0, double alpha = 1.702)
+{
+    gate           = std::min(gate, limit);
+    up             = std::min(std::max(up, -limit), limit);
+    const double s = 1.0 / (1.0 + std::exp(alpha * -gate));
+    return gate * s * (up + 1.0);
+}
+
+constexpr float kTol = 1e-3f;
+
+} // namespace
+
+// Values match the reference when no clamping is active.
+TEST(SwigluOai, MatchesReferenceInRange)
+{
+    const float pts[][2] = {{0.f, 0.f}, {1.f, 2.f}, {-1.5f, 0.5f}, {3.f, -2.f}, {-4.f, -3.f}};
+    for(const auto& p : pts)
+    {
+        const float got = ck::swiglu_oai(p[0], p[1]);
+        const float ref = static_cast<float>(ref_swiglu_oai(p[0], p[1]));
+        EXPECT_NEAR(got, ref, kTol) << "gate=" << p[0] << " up=" << p[1];
+    }
+}
+
+// gate is upper-bounded to limit (7); it is NOT lower-clamped.
+TEST(SwigluOai, GateUpperClamp)
+{
+    EXPECT_NEAR(ck::swiglu_oai(100.f, 1.f), static_cast<float>(ref_swiglu_oai(7.0, 1.0)), kTol);
+    // gate >= limit all saturate to the same value.
+    EXPECT_NEAR(ck::swiglu_oai(7.f, 1.f), ck::swiglu_oai(100.f, 1.f), kTol);
+    // large negative gate passes through unclamped.
+    EXPECT_NEAR(ck::swiglu_oai(-50.f, 1.f), static_cast<float>(ref_swiglu_oai(-50.0, 1.0)), kTol);
+}
+
+// up is symmetric-clamped to [-7, 7].
+TEST(SwigluOai, UpSymmetricClamp)
+{
+    EXPECT_NEAR(ck::swiglu_oai(1.f, 100.f), static_cast<float>(ref_swiglu_oai(1.0, 7.0)), kTol);
+    EXPECT_NEAR(ck::swiglu_oai(1.f, -100.f), static_cast<float>(ref_swiglu_oai(1.0, -7.0)), kTol);
+}
+
+// The "+1" shift on up is part of the OAI form: up == -1 zeroes the output exactly.
+TEST(SwigluOai, UpPlusOneShiftZeroesOutput)
+{
+    EXPECT_FLOAT_EQ(ck::swiglu_oai(2.f, -1.f), 0.f);
+    EXPECT_FLOAT_EQ(ck::swiglu_oai(-3.f, -1.f), 0.f);
+}
+
+// alpha defaults to 1.702 (gpt-oss); passing it explicitly must not change the result.
+TEST(SwigluOai, DefaultAlphaMatchesExplicit)
+{
+    EXPECT_FLOAT_EQ(ck::swiglu_oai(2.f, 3.f), ck::swiglu_oai(2.f, 3.f, 7.0f, 1.702f));
+}