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)`
 
   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>
@@ -7683,4 +7689,3 @@ No versioning maintained for experimental ops.
 <dd>Constrain input and output types to float32 tensors.</dd>
 </dl>
 
-

onnxruntime/contrib_ops/bert/linear_attention_gates_common.h

@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cstdint>
+#include <string_view>
+
+#include "core/common/common.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+enum class GatedRMSNormActivation : uint8_t {
+  kSilu = 0,
+  kSigmoid = 1,
+};
+
+inline Status ParseGatedRMSNormActivation(std::string_view activation, GatedRMSNormActivation& out) {
+  if (activation == "silu" || activation == "swish") {
+    out = GatedRMSNormActivation::kSilu;
+    return Status::OK();
+  }
+
+  if (activation == "sigmoid") {
+    out = GatedRMSNormActivation::kSigmoid;
+    return Status::OK();
+  }
+
+  return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "activation must be one of: silu, swish, sigmoid");
+}
+
+inline GatedRMSNormActivation ParseGatedRMSNormActivationOrThrow(std::string_view activation) {
+  GatedRMSNormActivation parsed_activation = GatedRMSNormActivation::kSilu;
+  ORT_THROW_IF_ERROR(ParseGatedRMSNormActivation(activation, parsed_activation));
+  return parsed_activation;
+}
+
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cpu/bert/linear_attention_gates.cc

@@ -45,8 +45,6 @@
 
 #undef REGISTER_KERNEL_TYPED
 
-namespace {
-
 inline float SigmoidFloat(float value) {
   float output;
   MlasComputeLogistic(&value, &output, 1);
@@ -57,8 +55,6 @@
   return value > 0.0f ? value + std::log(std::exp(-value) + 1.0f) : std::log(std::exp(value) + 1.0f);
 }
 
-}  // namespace
-
 template <typename T>
 Status LinearAttentionGate<T>::Compute(OpKernelContext* context) const {
   const Tensor* a = context->Input<Tensor>(0);
@@ -118,6 +114,7 @@
 
 template <typename T>
 GatedRMSNorm<T>::GatedRMSNorm(const OpKernelInfo& info) : OpKernel(info) {
+  activation_ = ParseGatedRMSNormActivationOrThrow(info.GetAttrOrDefault<std::string>("activation", "silu"));
   epsilon_ = info.GetAttrOrDefault<float>("epsilon", 1e-5f);
 }
 
@@ -160,11 +157,20 @@
           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 (activation_ == GatedRMSNormActivation::kSilu) {
+          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

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include "contrib_ops/bert/linear_attention_gates_common.h"
 #include "core/common/common.h"
 #include "core/framework/op_kernel.h"
 
@@ -17,14 +18,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:
+  GatedRMSNormActivation activation_;
   float epsilon_;
 };
 

onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.cc

@@ -93,6 +93,7 @@
 
 template <typename T>
 GatedRMSNorm<T>::GatedRMSNorm(const OpKernelInfo& info) : CudaKernel(info) {
+  activation_ = ParseGatedRMSNormActivationOrThrow(info.GetAttrOrDefault<std::string>("activation", "silu"));
   epsilon_ = info.GetAttrOrDefault<float>("epsilon", 1e-5f);
 }
 
@@ -128,7 +129,9 @@
       reinterpret_cast<const CudaT*>(gate->Data<T>()),
       num_rows,
       static_cast<int>(norm_size),
-      epsilon_);
+      epsilon_,
+      activation_,
+      GetDeviceProp().maxThreadsPerBlock);
 }
 
 template class LinearAttentionGate<float>;

onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.h

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include "contrib_ops/bert/linear_attention_gates_common.h"
 #include "core/common/common.h"
 #include "core/providers/cuda/cuda_kernel.h"
 
@@ -18,14 +19,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:
+  GatedRMSNormActivation activation_;
   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,52 @@ Status LaunchGatedRMSNormKernel(
     const T* gate,
     int64_t num_rows,
     int norm_size,
-    float epsilon) {
+    float epsilon,
+    GatedRMSNormActivation activation,
+    int max_threads_per_block) {
   if (num_rows == 0) {
     return Status::OK();
   }
 
   ORT_RETURN_IF_NOT(num_rows <= std::numeric_limits<int>::max(),
                     "GatedRMSNorm launch requires too many blocks");
+  ORT_RETURN_IF_NOT(max_threads_per_block >= 64, "GatedRMSNorm requires maxThreadsPerBlock >= 64");
   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);
-  } else if (norm_size <= 128) {
-    LAUNCH_GATED_RMS_NORM(128);
-  } else if (norm_size <= 256) {
-    LAUNCH_GATED_RMS_NORM(256);
-  } else if (norm_size <= 512) {
-    LAUNCH_GATED_RMS_NORM(512);
+  const int max_threads = max_threads_per_block;
+  if (norm_size <= 64 || max_threads < 128) {
+    if (activation == GatedRMSNormActivation::kSilu) {
+      LAUNCH_GATED_RMS_NORM(64, true);
+    } else {
+      LAUNCH_GATED_RMS_NORM(64, false);
+    }
+  } else if (norm_size <= 128 || max_threads < 256) {
+    if (activation == GatedRMSNormActivation::kSilu) {
+      LAUNCH_GATED_RMS_NORM(128, true);
+    } else {
+      LAUNCH_GATED_RMS_NORM(128, false);
+    }
+  } else if (norm_size <= 256 || max_threads < 512) {
+    if (activation == GatedRMSNormActivation::kSilu) {
+      LAUNCH_GATED_RMS_NORM(256, true);
+    } else {
+      LAUNCH_GATED_RMS_NORM(256, false);
+    }
+  } else if (norm_size <= 512 || max_threads < 1024) {
+    if (activation == GatedRMSNormActivation::kSilu) {
+      LAUNCH_GATED_RMS_NORM(512, true);
+    } else {
+      LAUNCH_GATED_RMS_NORM(512, false);
+    }
   } else {
-    LAUNCH_GATED_RMS_NORM(1024);
+    if (activation == GatedRMSNormActivation::kSilu) {
+      LAUNCH_GATED_RMS_NORM(1024, true);
+    } else {
+      LAUNCH_GATED_RMS_NORM(1024, false);
+    }
   }
 #undef LAUNCH_GATED_RMS_NORM
 
@@ -168,7 +197,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, GatedRMSNormActivation, int);
 
 INSTANTIATE_LINEAR_ATTENTION_GATES(float)
 INSTANTIATE_LINEAR_ATTENTION_GATES(half)

onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.h

@@ -4,6 +4,7 @@
 #pragma once
 
 #include <cuda_fp16.h>
+#include "contrib_ops/bert/linear_attention_gates_common.h"
 #include "core/providers/cuda/cuda_common.h"
 
 namespace onnxruntime {
@@ -24,7 +25,7 @@ Status LaunchLinearAttentionGateKernel(
     int64_t num_tokens,
     int num_heads);
 
-// Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate), reduced over groups of
+// Y = X * rsqrt(mean(X^2) + epsilon) * scale * gate_activation(gate), reduced over groups of
 // `norm_size` contiguous elements, with all arithmetic in float32.
 template <typename T>
 Status LaunchGatedRMSNormKernel(
@@ -35,7 +36,9 @@ Status LaunchGatedRMSNormKernel(
     const T* gate,
     int64_t num_rows,
     int norm_size,
-    float epsilon);
+    float epsilon,
+    GatedRMSNormActivation activation,
+    int max_threads_per_block);
 
 }  // namespace cuda
 }  // namespace contrib

onnxruntime/contrib_ops/webgpu/bert/linear_attention_gates.cc

@@ -162,13 +162,18 @@
       << "    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 * ") +
+                                            (activation_ == GatedRMSNormActivation::kSilu ? "(z * stable_sigmoid(z))" : "stable_sigmoid(z)") +
+                                            ")")
+      << "\n"
       << "  }\n";
 
   return Status::OK();
 }
 
 GatedRMSNorm::GatedRMSNorm(const OpKernelInfo& info) : WebGpuKernel(info) {
+  activation_ = ParseGatedRMSNormActivationOrThrow(info.GetAttrOrDefault<std::string>("activation", "silu"));
   epsilon_ = info.GetAttrOrDefault<float>("epsilon", 1e-5f);
 }
 
@@ -199,7 +204,8 @@
                                   : norm_size <= 128 ? 128
                                                      : 256;
 
-  GatedRMSNormProgram program{};
+  GatedRMSNormProgram program{activation_};
+  program.CacheHint(static_cast<int>(activation_));
   program.AddInputs({{input, ProgramTensorMetadataDependency::Type},
                      {scale, ProgramTensorMetadataDependency::Type},
                      {gate, ProgramTensorMetadataDependency::Type}})

onnxruntime/contrib_ops/webgpu/bert/linear_attention_gates.h

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include "contrib_ops/bert/linear_attention_gates_common.h"
 #include "core/providers/webgpu/program.h"
 #include "core/providers/webgpu/webgpu_kernel.h"
 
@@ -32,13 +33,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(GatedRMSNormActivation activation) : Program{"GatedRMSNorm"}, activation_(activation) {}
   Status GenerateShaderCode(ShaderHelper& sh) const override;
   WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"norm_size", ProgramUniformVariableDataType::Uint32},
                                           {"epsilon", ProgramUniformVariableDataType::Float32});
+
+ private:
+  GatedRMSNormActivation activation_;
 };
 
 class GatedRMSNorm final : public WebGpuKernel {
@@ -47,6 +51,7 @@ class GatedRMSNorm final : public WebGpuKernel {
   Status ComputeInternal(ComputeContext& context) const override;
 
  private:
+  GatedRMSNormActivation activation_;
   float epsilon_;
 };