Fused rms silu

docs/src/python/fast.rst

@@ -9,6 +9,7 @@ Fast
   :toctree: _autosummary
 
   rms_norm
+  fused_rms_silu
   layer_norm
   rope
   scaled_dot_product_attention

mlx/backend/metal/kernels/rms_norm.metal

@@ -8,6 +8,7 @@
 using namespace metal;
 
 constant bool has_w [[function_constant(20)]];
+constant bool has_silu [[function_constant(21)]];
 
 template <typename T, int N_READS = RMS_N_READS>
 [[kernel]] void rms_single_row(
@@ -67,14 +68,22 @@ template <typename T, int N_READS = RMS_N_READS>
   out += gid * size_t(axis_size) + lid * N_READS;
   if (lid * N_READS + N_READS <= axis_size) {
     for (int i = 0; i < N_READS; i++) {
-      out[i] =
-          w[w_stride * i] * static_cast<T>(thread_x[i] * local_inv_mean[0]);
+      float norm_val = (has_w ? static_cast<float>(w[w_stride * i]) : 1.0f) *
+          (thread_x[i] * local_inv_mean[0]);
+      if (has_silu) {
+        norm_val = norm_val / (1.0f + metal::fast::exp(-norm_val));
+      }
+      out[i] = static_cast<T>(norm_val);
     }
   } else {
     for (int i = 0; i < N_READS; i++) {
       if ((lid * N_READS + i) < axis_size) {
-        out[i] =
-            w[w_stride * i] * static_cast<T>(thread_x[i] * local_inv_mean[0]);
+        float norm_val = (has_w ? static_cast<float>(w[w_stride * i]) : 1.0f) *
+            (thread_x[i] * local_inv_mean[0]);
+        if (has_silu) {
+          norm_val = norm_val / (1.0f + metal::fast::exp(-norm_val));
+        }
+        out[i] = static_cast<T>(norm_val);
       }
     }
   }
@@ -142,14 +151,24 @@ template <typename T, int N_READS = RMS_N_READS>
   for (uint r = 0; r < axis_size; r += lsize * N_READS) {
     if (r + lid * N_READS + N_READS <= axis_size) {
       for (int i = 0; i < N_READS; i++) {
-        out[r + i] = w[w_stride * (i + r)] *
-            static_cast<T>(x[r + i] * local_inv_mean[0]);
+        float norm_val =
+            (has_w ? static_cast<float>(w[w_stride * (i + r)]) : 1.0f) *
+            (x[r + i] * local_inv_mean[0]);
+        if (has_silu) {
+          norm_val = norm_val / (1.0f + metal::fast::exp(-norm_val));
+        }
+        out[r + i] = static_cast<T>(norm_val);
       }
     } else {
       for (int i = 0; i < N_READS; i++) {
         if ((r + lid * N_READS + i) < axis_size) {
-          out[r + i] = w[w_stride * (i + r)] *
-              static_cast<T>(x[r + i] * local_inv_mean[0]);
+          float norm_val =
+              (has_w ? static_cast<float>(w[w_stride * (i + r)]) : 1.0f) *
+              (x[r + i] * local_inv_mean[0]);
+          if (has_silu) {
+            norm_val = norm_val / (1.0f + metal::fast::exp(-norm_val));
+          }
+          out[r + i] = static_cast<T>(norm_val);
         }
       }
     }

mlx/backend/metal/normalization.cpp

@@ -60,9 +60,108 @@ void RMSNorm::eval_gpu(
     op_name += "_looped";
   }
   op_name += type_to_name(out);
+
+  bool has_w = w.ndim() != 0;
+  bool has_silu = false;
+  std::string hash_name = op_name + (has_w ? "_w" : "_now") + "_nosilu";
+  metal::MTLFCList func_consts = {
+      {&has_w, MTL::DataType::DataTypeBool, 20},
+      {&has_silu, MTL::DataType::DataTypeBool, 21},
+  };
+
   auto& compute_encoder = metal::get_command_encoder(s);
   {
-    auto kernel = d.get_kernel(op_name);
+    auto kernel = d.get_kernel(op_name, hash_name, func_consts);
+
+    MTL::Size grid_dims, group_dims;
+    if (axis_size <= looped_limit) {
+      size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
+      size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
+      size_t threadgroup_size = simd_size * simds_needed;
+      assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    } else {
+      size_t threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    }
+
+    uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(x, 0);
+    compute_encoder.set_input_array(w, 1);
+    compute_encoder.set_output_array(out, 2);
+    compute_encoder.set_bytes(eps_, 3);
+    compute_encoder.set_bytes(axis_size, 4);
+    compute_encoder.set_bytes(w_stride, 5);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+bool FusedRMSSiLU::use_fallback(Stream s) {
+  return s.device == Device::cpu;
+}
+
+void FusedRMSSiLU::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& out = outputs[0];
+
+  auto set_output = [&s, &out](const array& x) {
+    bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1;
+    if (no_copy && x.ndim() > 1) {
+      auto s = x.strides()[x.ndim() - 2];
+      no_copy &= (s == 0 || s == x.shape().back() || x.shape(-2) == 1);
+    }
+    if (no_copy) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            allocator::malloc(x.data_size() * x.itemsize()),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
+    }
+  };
+
+  const array x = set_output(inputs[0]);
+  const array& w = inputs[1];
+
+  auto axis_size = static_cast<uint32_t>(x.shape().back());
+  int n_rows = x.data_size() / axis_size;
+
+  const int simd_size = 32;
+  const int n_reads = RMS_N_READS;
+  const int looped_limit = RMS_LOOPED_LIMIT;
+  std::string op_name = "rms";
+  if (axis_size > looped_limit) {
+    op_name += "_looped";
+  }
+  op_name += type_to_name(out);
+
+  bool has_w = w.ndim() != 0;
+  bool has_silu = true;
+  std::string hash_name = op_name + (has_w ? "_w" : "_now") + "_silu";
+  metal::MTLFCList func_consts = {
+      {&has_w, MTL::DataType::DataTypeBool, 20},
+      {&has_silu, MTL::DataType::DataTypeBool, 21},
+  };
+
+  auto& compute_encoder = metal::get_command_encoder(s);
+  {
+    auto kernel = d.get_kernel(op_name, hash_name, func_consts);
 
     MTL::Size grid_dims, group_dims;
     if (axis_size <= looped_limit) {

mlx/fast.cpp

@@ -120,6 +120,76 @@ array rms_norm(
   return fallback({x, passed_weight})[0];
 }
 
+array fused_rms_silu(
+    const array& x,
+    const std::optional<array>& weight,
+    float eps,
+    StreamOrDevice s_ /* = {} */) {
+  bool has_weight = weight.has_value();
+
+  if (x.ndim() == 0) {
+    std::ostringstream msg;
+    msg << "[fused_rms_silu] Input must have at least 1 dimension but got input with "
+           "0 dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (has_weight) {
+    if ((*weight).ndim() != 1) {
+      std::ostringstream msg;
+      msg << "[fused_rms_silu] (*weight) must have 1 dimension but has "
+          << (*weight).ndim() << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    if ((*weight).size() != x.shape(-1)) {
+      std::ostringstream msg;
+      msg << "[fused_rms_silu] (*weight) must have the same size as the last dimension of"
+             " x but has "
+          << (*weight).size() << " elements.";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+
+  auto out_type = (weight.has_value()) ? result_type(x, (*weight)) : x.dtype();
+  if (!issubdtype(out_type, floating)) {
+    std::ostringstream msg;
+    msg << "[fused_rms_silu] Received unsupported type " << out_type << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto s = to_stream(s_);
+  auto fallback =
+      [has_weight, eps, out_type, s](const std::vector<array>& inputs) {
+        auto x = astype(inputs[0], float32, s);
+        x = multiply(
+            x,
+            rsqrt(
+                add(mean(square(x, s), -1, /* keepdims */ true, s),
+                    array(eps, float32),
+                    s),
+                s),
+            s);
+
+        if (has_weight) {
+          x = multiply(x, inputs[1], s);
+        }
+        x = multiply(x, sigmoid(x, s), s);
+
+        return std::vector<array>{astype(x, out_type, s)};
+      };
+
+  auto passed_weight =
+      (has_weight) ? astype(*weight, out_type, s) : array(1, out_type);
+
+  if (!FusedRMSSiLU::use_fallback(s)) {
+    return array(
+        x.shape(),
+        out_type,
+        std::make_shared<FusedRMSSiLU>(s, fallback, eps),
+        {astype(x, out_type, s), passed_weight});
+  }
+  return fallback({x, passed_weight})[0];
+}
+
 std::vector<array> RMSNorm::vjp(
     const std::vector<array>& primals,
     const std::vector<array>& cotangents,
@@ -182,6 +252,11 @@ bool RMSNorm::is_equivalent(const Primitive& other) const {
   return eps_ == a_other.eps_;
 }
 
+bool FusedRMSSiLU::is_equivalent(const Primitive& other) const {
+  const FusedRMSSiLU& a_other = static_cast<const FusedRMSSiLU&>(other);
+  return eps_ == a_other.eps_;
+}
+
 bool RMSNormVJP::is_equivalent(const Primitive& other) const {
   const RMSNormVJP& a_other = static_cast<const RMSNormVJP&>(other);
   return eps_ == a_other.eps_;

mlx/fast.h

@@ -17,6 +17,12 @@ MLX_API array rms_norm(
     float eps,
     StreamOrDevice s = {});
 
+MLX_API array fused_rms_silu(
+    const array& x,
+    const std::optional<array>& weight,
+    float eps,
+    StreamOrDevice s = {});
+
 MLX_API array layer_norm(
     const array& x,
     const std::optional<array>& weight,

mlx/fast_primitives.h

@@ -72,6 +72,35 @@ class RMSNorm : public Custom {
   float eps_;
 };
 
+class FusedRMSSiLU : public Custom {
+ public:
+  FusedRMSSiLU(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, std::move(fallback)), eps_(eps) {}
+
+  static bool use_fallback(Stream stream);
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(FusedRMSSiLU)
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE()
+
+  auto state() const {
+    return std::make_pair(nullptr, eps_);
+  }
+
+ private:
+  float eps_;
+};
+
 class RMSNormVJP : public Custom {
  public:
   RMSNormVJP(

python/src/fast.cpp

@@ -144,6 +144,30 @@ void init_fast(nb::module_& parent_module) {
             array: The output array.
       )pbdoc");
 
+  m.def(
+      "fused_rms_silu",
+      &mx::fast::fused_rms_silu,
+      "x"_a,
+      "weight"_a.none(),
+      "eps"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def fused_rms_silu(x: array, weight: array | None, eps: float, *, stream: StreamOrDevice = None) -> array"),
+      R"pbdoc(
+        Fused Root Mean Square normalization (RMS norm) and SiLU (Swish) activation.
+
+        Applies RMSNorm followed immediately by SiLU activation in a single GPU pass.
+
+        Args:
+            x (array): Input array.
+            weight (array, optional): A multiplicative weight to scale the result by.
+            eps (float): A small additive constant for numerical stability.
+
+        Returns:
+            array: The output array.
+      )pbdoc");
+
   m.def(
       "layer_norm",
       &mx::fast::layer_norm,

python/tests/test_fast.py

@@ -477,6 +477,32 @@ def test_rms_norm(self):
         rx_fast = mx.fast.rms_norm(x, weight, eps)
         self.assertLess(mx.abs(rx - rx_fast).max(), 1e-6)
 
+    def test_fused_rms_silu(self):
+        def ref_fused_rms_silu(x, weight, eps):
+            w = weight if weight is not None else mx.ones((x.shape[-1],), dtype=x.dtype)
+            h = rms_norm(x, w, eps)
+            return h * mx.sigmoid(h)
+
+        tolerances = {mx.float32: 1e-5, mx.float16: 5e-3, mx.bfloat16: 5e-2}
+        dtypes = [mx.float32, mx.float16, mx.bfloat16]
+        epss = [1e-3, 1e-5]
+        dimss = [31, 32, 4096]
+
+        for dtype in dtypes:
+            for eps in epss:
+                for dims in dimss:
+                    x = mx.random.uniform(shape=(2, dims)).astype(dtype)
+                    weight = mx.random.uniform(shape=(dims,)).astype(dtype)
+                    expected = ref_fused_rms_silu(x, weight, eps)
+                    actual = mx.fast.fused_rms_silu(x, weight, eps)
+                    self.assertLess(mx.abs(expected - actual).max(), tolerances[dtype])
+
+                    expected_now = ref_fused_rms_silu(x, None, eps)
+                    actual_now = mx.fast.fused_rms_silu(x, None, eps)
+                    self.assertLess(
+                        mx.abs(expected_now - actual_now).max(), tolerances[dtype]
+                    )
+
         # Wrong size w raises
         with self.assertRaises(ValueError):
             x = mx.random.uniform(shape=(1, 5))