Add activation mode to com.microsoft::GatedRMSNorm for SiLU/Swish/Sigmoid across CPU, CUDA, and WebGPU

docs/ContribOperators.md

@@ -2325,12 +2325,16 @@ This version of the operator has been available since version 1 of the 'com.micr
 
   Gated RMS normalization as used by Mamba2 / gated DeltaNet attention outputs:
 
-    Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate)
+    Y = X * rsqrt(mean(X^2) + epsilon) * scale * gate_activation(gate)
+
+  where `gate_activation` is one of:
+  - `silu` or `swish`: `z * sigmoid(z)`
+  - `sigmoid`: `sigmoid(z)` (used by Qwen3.8-Flash-Next / qwen4_exp output gating)
 
   The mean of squares is taken over the trailing `C` elements of each row, where `C` is the
   length of `scale`; the input's last dimension must be a multiple of `C`, which lets a
   per-head norm run on a packed (B, T, H * C) tensor without any surrounding Reshape.
-  All arithmetic including SiLU is done in float32 regardless of the tensor type, matching
+  All arithmetic including gate activation is done in float32 regardless of the tensor type, matching
   the reference implementation, so this replaces the exported
   SimplifiedLayerNormalization -> Cast -> Sigmoid -> Mul -> Cast -> Mul -> Cast chain with a
   single launch.
@@ -2342,6 +2346,8 @@ This version of the operator has been available since version 1 of the 'com.micr
 #### Attributes
 
 <dl>
+<dt><tt>activation</tt> : string</dt>
+<dd>Fused gate activation. One of: 'silu', 'swish', 'sigmoid'. 'swish' is an alias of 'silu'.</dd>
 <dt><tt>epsilon</tt> : float</dt>
 <dd>Epsilon added to the mean of squares before the reciprocal square root.</dd>
 </dl>

onnxruntime/contrib_ops/cpu/bert/linear_attention_gates.cc

@@ -4,6 +4,7 @@
 #include "contrib_ops/cpu/bert/linear_attention_gates.h"
 
 #include <cmath>
+#include <string>
 
 #include "core/framework/tensor.h"
 #include "core/mlas/inc/mlas.h"
@@ -57,6 +58,16 @@ inline float SoftplusFloat(float value) {
   return value > 0.0f ? value + std::log(std::exp(-value) + 1.0f) : std::log(std::exp(value) + 1.0f);
 }
 
+inline bool ParseGatedRMSNormActivation(const std::string& activation) {
+  if (activation == "silu" || activation == "swish") {
+    return true;
+  }
+  if (activation == "sigmoid") {
+    return false;
+  }
+  ORT_THROW("activation must be one of: silu, swish, sigmoid");
+}
+
 }  // namespace
 
 template <typename T>
@@ -118,6 +129,7 @@ Status LinearAttentionGate<T>::Compute(OpKernelContext* context) const {
 
 template <typename T>
 GatedRMSNorm<T>::GatedRMSNorm(const OpKernelInfo& info) : OpKernel(info) {
+  use_silu_ = ParseGatedRMSNormActivation(info.GetAttrOrDefault<std::string>("activation", "silu"));
   epsilon_ = info.GetAttrOrDefault<float>("epsilon", 1e-5f);
 }
 
@@ -160,11 +172,20 @@ Status GatedRMSNorm<T>::Compute(OpKernelContext* context) const {
           sum_sq += v * v;
         }
         const float inv_rms = 1.0f / std::sqrt(sum_sq / static_cast<float>(norm_size) + epsilon_);
-        for (int64_t i = 0; i < norm_size; ++i) {
-          const float z = static_cast<float>(gate_data[offset + i]);
-          const float normalized = static_cast<float>(input_data[offset + i]) * inv_rms *
-                                   static_cast<float>(scale_data[i]);
-          output_data[offset + i] = static_cast<T>(normalized * (z * SigmoidFloat(z)));
+        if (use_silu_) {
+          for (int64_t i = 0; i < norm_size; ++i) {
+            const float z = static_cast<float>(gate_data[offset + i]);
+            const float normalized = static_cast<float>(input_data[offset + i]) * inv_rms *
+                                     static_cast<float>(scale_data[i]);
+            output_data[offset + i] = static_cast<T>(normalized * (z * SigmoidFloat(z)));
+          }
+        } else {
+          for (int64_t i = 0; i < norm_size; ++i) {
+            const float z = static_cast<float>(gate_data[offset + i]);
+            const float normalized = static_cast<float>(input_data[offset + i]) * inv_rms *
+                                     static_cast<float>(scale_data[i]);
+            output_data[offset + i] = static_cast<T>(normalized * SigmoidFloat(z));
+          }
         }
       },
       0);

onnxruntime/contrib_ops/cpu/bert/linear_attention_gates.h

@@ -17,14 +17,15 @@ class LinearAttentionGate final : public OpKernel {
   Status Compute(OpKernelContext* context) const override;
 };
 
-// Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate).
+// Y = X * rsqrt(mean(X^2) + epsilon) * scale * gate_activation(gate).
 template <typename T>
 class GatedRMSNorm final : public OpKernel {
  public:
   explicit GatedRMSNorm(const OpKernelInfo& info);
   Status Compute(OpKernelContext* context) const override;
 
  private:
+  bool use_silu_;
   float epsilon_;
 };
 

onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.cc

@@ -7,13 +7,26 @@
 #include "core/providers/cuda/cuda_type_conversion.h"
 
 #include <limits>
+#include <string>
 
 namespace onnxruntime {
 namespace contrib {
 namespace cuda {
 
 using namespace onnxruntime::cuda;  // CudaKernel, OrtToCudaType
 
+namespace {
+bool ParseGatedRMSNormActivation(const std::string& activation) {
+  if (activation == "silu" || activation == "swish") {
+    return true;
+  }
+  if (activation == "sigmoid") {
+    return false;
+  }
+  ORT_THROW("activation must be one of: silu, swish, sigmoid");
+}
+}  // namespace
+
 #define REGISTER_KERNEL_TYPED(Op, T)                                   \
   ONNX_OPERATOR_TYPED_KERNEL_EX(                                       \
       Op,                                                              \
@@ -93,6 +106,7 @@
 
 template <typename T>
 GatedRMSNorm<T>::GatedRMSNorm(const OpKernelInfo& info) : CudaKernel(info) {
+  use_silu_ = ParseGatedRMSNormActivation(info.GetAttrOrDefault<std::string>("activation", "silu"));
   epsilon_ = info.GetAttrOrDefault<float>("epsilon", 1e-5f);
 }
 
@@ -128,7 +142,8 @@
       reinterpret_cast<const CudaT*>(gate->Data<T>()),
       num_rows,
       static_cast<int>(norm_size),
-      epsilon_);
+      epsilon_,
+      use_silu_);
 }
 
 template class LinearAttentionGate<float>;

onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.h

@@ -18,14 +18,15 @@ class LinearAttentionGate final : public onnxruntime::cuda::CudaKernel {
   Status ComputeInternal(OpKernelContext* context) const override;
 };
 
-// Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate).
+// Y = X * rsqrt(mean(X^2) + epsilon) * scale * gate_activation(gate).
 template <typename T>
 class GatedRMSNorm final : public onnxruntime::cuda::CudaKernel {
  public:
   explicit GatedRMSNorm(const OpKernelInfo& info);
   Status ComputeInternal(OpKernelContext* context) const override;
 
  private:
+  bool use_silu_;
   float epsilon_;
 };
 

onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.cu

@@ -29,7 +29,11 @@ namespace {
 // Matches OP_Sigmoid in core/providers/cuda/activation/activations_impl.cu: the branch keeps the
 // exponent argument non-positive so large-magnitude inputs cannot overflow.
 __device__ __forceinline__ float SigmoidFloat(float x) {
-  return x > 0.0f ? 1.0f / (1.0f + expf(-x)) : 1.0f - 1.0f / (1.0f + expf(x));
+  if (x > 0.0f) {
+    return 1.0f / (1.0f + expf(-x));
+  }
+  const float e = expf(x);
+  return e / (1.0f + e);
 }
 
 // Matches OP_Softplus in the same file.
@@ -63,7 +67,7 @@ __global__ void LinearAttentionGateKernel(
 
 // One block per normalization group. The input is read twice (once for the sum of squares, once
 // for the output); the group is a few hundred bytes so the second read is an L1 hit.
-template <typename T, int kThreadsPerBlock>
+template <typename T, int kThreadsPerBlock, bool kUseSilu>
 __global__ void GatedRMSNormKernel(
     T* output,
     const T* input,
@@ -96,7 +100,8 @@ __global__ void GatedRMSNormKernel(
   for (int i = threadIdx.x; i < norm_size; i += kThreadsPerBlock) {
     const float z = to_float<T>(g[i]);
     const float normalized = to_float<T>(x[i]) * inv_rms * to_float<T>(scale[i]);
-    y[i] = from_float<T>(normalized * (z * SigmoidFloat(z)));
+    const float activated_gate = kUseSilu ? (z * SigmoidFloat(z)) : SigmoidFloat(z);
+    y[i] = from_float<T>(normalized * activated_gate);
   }
 }
 
@@ -136,28 +141,49 @@ Status LaunchGatedRMSNormKernel(
     const T* gate,
     int64_t num_rows,
     int norm_size,
-    float epsilon) {
+    float epsilon,
+    bool use_silu) {
   if (num_rows == 0) {
     return Status::OK();
   }
 
   ORT_RETURN_IF_NOT(num_rows <= std::numeric_limits<int>::max(),
                     "GatedRMSNorm launch requires too many blocks");
   const int blocks = static_cast<int>(num_rows);
-#define LAUNCH_GATED_RMS_NORM(threads)                            \
-  GatedRMSNormKernel<T, threads><<<blocks, threads, 0, stream>>>( \
+#define LAUNCH_GATED_RMS_NORM(threads, use_silu_mode)                            \
+  GatedRMSNormKernel<T, threads, use_silu_mode><<<blocks, threads, 0, stream>>>( \
       output, input, scale, gate, norm_size, epsilon)
 
   if (norm_size <= 64) {
-    LAUNCH_GATED_RMS_NORM(64);
+    if (use_silu) {
+      LAUNCH_GATED_RMS_NORM(64, true);
+    } else {
+      LAUNCH_GATED_RMS_NORM(64, false);
+    }
   } else if (norm_size <= 128) {
-    LAUNCH_GATED_RMS_NORM(128);
+    if (use_silu) {
+      LAUNCH_GATED_RMS_NORM(128, true);
+    } else {
+      LAUNCH_GATED_RMS_NORM(128, false);
+    }
   } else if (norm_size <= 256) {
-    LAUNCH_GATED_RMS_NORM(256);
+    if (use_silu) {
+      LAUNCH_GATED_RMS_NORM(256, true);
+    } else {
+      LAUNCH_GATED_RMS_NORM(256, false);
+    }
   } else if (norm_size <= 512) {
-    LAUNCH_GATED_RMS_NORM(512);
+    if (use_silu) {
+      LAUNCH_GATED_RMS_NORM(512, true);
+    } else {
+      LAUNCH_GATED_RMS_NORM(512, false);
+    }
   } else {
-    LAUNCH_GATED_RMS_NORM(1024);
+    if (use_silu) {
+      LAUNCH_GATED_RMS_NORM(1024, true);
+    } else {
+      LAUNCH_GATED_RMS_NORM(1024, false);
+    }
   }
 #undef LAUNCH_GATED_RMS_NORM
 
@@ -168,7 +194,7 @@ Status LaunchGatedRMSNormKernel(
   template Status LaunchLinearAttentionGateKernel<T>(cudaStream_t, T*, T*, const T*, const T*,  \
                                                      const float*, const float*, int64_t, int); \
   template Status LaunchGatedRMSNormKernel<T>(cudaStream_t, T*, const T*, const T*, const T*,   \
-                                              int64_t, int, float);
+                                              int64_t, int, float, bool);
 
 INSTANTIATE_LINEAR_ATTENTION_GATES(float)
 INSTANTIATE_LINEAR_ATTENTION_GATES(half)

onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.h

@@ -24,8 +24,9 @@ Status LaunchLinearAttentionGateKernel(
     int64_t num_tokens,
     int num_heads);
 
-// Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate), reduced over groups of
-// `norm_size` contiguous elements, with all arithmetic in float32.
+// Y = X * rsqrt(mean(X^2) + epsilon) * scale * gate_activation(gate), reduced over groups of
+// `norm_size` contiguous elements, with all arithmetic in float32. `use_silu=true` applies
+// SiLU/Swish (z * Sigmoid(z)); false applies Sigmoid.
 template <typename T>
 Status LaunchGatedRMSNormKernel(
     cudaStream_t stream,
@@ -35,7 +36,8 @@ Status LaunchGatedRMSNormKernel(
     const T* gate,
     int64_t num_rows,
     int norm_size,
-    float epsilon);
+    float epsilon,
+    bool use_silu);
 
 }  // namespace cuda
 }  // namespace contrib

onnxruntime/contrib_ops/webgpu/bert/linear_attention_gates.cc

@@ -7,10 +7,24 @@
 #include "core/providers/webgpu/webgpu_supported_types.h"
 #include "contrib_ops/webgpu/webgpu_contrib_kernels.h"
 
+#include <string>
+
 namespace onnxruntime {
 namespace contrib {
 namespace webgpu {
 
+namespace {
+bool ParseGatedRMSNormActivation(const std::string& activation) {
+  if (activation == "silu" || activation == "swish") {
+    return true;
+  }
+  if (activation == "sigmoid") {
+    return false;
+  }
+  ORT_THROW("activation must be one of: silu, swish, sigmoid");
+}
+}  // namespace
+
 ONNX_OPERATOR_KERNEL_EX(
     LinearAttentionGate,
     kMSDomain,
@@ -162,13 +176,15 @@
       << "    let z = f32(" << gate.GetByOffset("base + i") << ");\n"
       << "    let normalized = f32(" << input.GetByOffset("base + i") << ") * inv_rms * f32("
       << scale.GetByOffset("i") << ");\n"
-      << "    " << output.SetByOffset("base + i", "output_element_t(normalized * (z * stable_sigmoid(z)))") << "\n"
+      << "    " << output.SetByOffset("base + i", std::string("output_element_t(normalized * ") + (use_silu_ ? "(z * stable_sigmoid(z))" : "stable_sigmoid(z)") + ")")
+      << "\n"
       << "  }\n";
 
   return Status::OK();
 }
 
 GatedRMSNorm::GatedRMSNorm(const OpKernelInfo& info) : WebGpuKernel(info) {
+  use_silu_ = ParseGatedRMSNormActivation(info.GetAttrOrDefault<std::string>("activation", "silu"));
   epsilon_ = info.GetAttrOrDefault<float>("epsilon", 1e-5f);
 }
 
@@ -199,7 +215,8 @@
                                   : norm_size <= 128 ? 128
                                                      : 256;
 
-  GatedRMSNormProgram program{};
+  GatedRMSNormProgram program{use_silu_};
+  program.CacheHint(use_silu_);
   program.AddInputs({{input, ProgramTensorMetadataDependency::Type},
                      {scale, ProgramTensorMetadataDependency::Type},
                      {gate, ProgramTensorMetadataDependency::Type}})

onnxruntime/contrib_ops/webgpu/bert/linear_attention_gates.h

@@ -32,13 +32,16 @@ class LinearAttentionGate final : public WebGpuKernel {
   Status ComputeInternal(ComputeContext& context) const override;
 };
 
-// Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate).
+// Y = X * rsqrt(mean(X^2) + epsilon) * scale * gate_activation(gate).
 class GatedRMSNormProgram final : public Program<GatedRMSNormProgram> {
  public:
-  GatedRMSNormProgram() : Program{"GatedRMSNorm"} {}
+  GatedRMSNormProgram(bool use_silu) : Program{"GatedRMSNorm"}, use_silu_(use_silu) {}
   Status GenerateShaderCode(ShaderHelper& sh) const override;
   WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"norm_size", ProgramUniformVariableDataType::Uint32},
                                           {"epsilon", ProgramUniformVariableDataType::Float32});
+
+ private:
+  bool use_silu_;
 };
 
 class GatedRMSNorm final : public WebGpuKernel {
@@ -47,6 +50,7 @@ class GatedRMSNorm final : public WebGpuKernel {
   Status ComputeInternal(ComputeContext& context) const override;
 
  private:
+  bool use_silu_;
   float epsilon_;
 };
 

onnxruntime/core/graph/contrib_ops/bert_defs.cc

@@ -3644,12 +3644,16 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
 constexpr const char* GatedRMSNorm_ver1_doc = R"DOC(
 Gated RMS normalization as used by Mamba2 / gated DeltaNet attention outputs:
 
-  Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate)
+  Y = X * rsqrt(mean(X^2) + epsilon) * scale * gate_activation(gate)
+
+where `gate_activation` is one of:
+- `silu` or `swish`: `z * sigmoid(z)`
+- `sigmoid`: `sigmoid(z)` (used by Qwen3.8-Flash-Next / qwen4_exp output gating)
 
 The mean of squares is taken over the trailing `C` elements of each row, where `C` is the
 length of `scale`; the input's last dimension must be a multiple of `C`, which lets a
 per-head norm run on a packed (B, T, H * C) tensor without any surrounding Reshape.
-All arithmetic including SiLU is done in float32 regardless of the tensor type, matching
+All arithmetic including gate activation is done in float32 regardless of the tensor type, matching
 the reference implementation, so this replaces the exported
 SimplifiedLayerNormalization -> Cast -> Sigmoid -> Mul -> Cast -> Mul -> Cast chain with a
 single launch.
@@ -3663,6 +3667,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
               "Epsilon added to the mean of squares before the reciprocal square root.",
               AttributeProto::FLOAT,
               1e-5f)
+        .Attr("activation",
+              "Fused gate activation. One of: 'silu', 'swish', 'sigmoid'. "
+              "'swish' is an alias of 'silu'.",
+              AttributeProto::STRING,
+              std::string("silu"))
         .Input(0,
                "X",
                "Input tensor with shape (..., H * C). Normalization is applied over each "