diff --git a/docs/src/python/fast.rst b/docs/src/python/fast.rst index affeb444f8..096b5ddd1a 100644 --- a/docs/src/python/fast.rst +++ b/docs/src/python/fast.rst @@ -9,6 +9,7 @@ Fast :toctree: _autosummary rms_norm + fused_rms_silu layer_norm rope scaled_dot_product_attention diff --git a/mlx/backend/metal/kernels/rms_norm.metal b/mlx/backend/metal/kernels/rms_norm.metal index a50d4a25c6..29f57106df 100644 --- a/mlx/backend/metal/kernels/rms_norm.metal +++ b/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 [[kernel]] void rms_single_row( @@ -67,14 +68,22 @@ template 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(thread_x[i] * local_inv_mean[0]); + float norm_val = (has_w ? static_cast(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(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(thread_x[i] * local_inv_mean[0]); + float norm_val = (has_w ? static_cast(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(norm_val); } } } @@ -142,14 +151,24 @@ template 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(x[r + i] * local_inv_mean[0]); + float norm_val = + (has_w ? static_cast(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(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(x[r + i] * local_inv_mean[0]); + float norm_val = + (has_w ? static_cast(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(norm_val); } } } diff --git a/mlx/backend/metal/normalization.cpp b/mlx/backend/metal/normalization.cpp index 9a222cdd6c..948ddd28c0 100644 --- a/mlx/backend/metal/normalization.cpp +++ b/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& inputs, + std::vector& 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(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) { diff --git a/mlx/fast.cpp b/mlx/fast.cpp index a668fe9abd..33ee69d0cf 100644 --- a/mlx/fast.cpp +++ b/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& 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& 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{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(s, fallback, eps), + {astype(x, out_type, s), passed_weight}); + } + return fallback({x, passed_weight})[0]; +} + std::vector RMSNorm::vjp( const std::vector& primals, const std::vector& 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(other); + return eps_ == a_other.eps_; +} + bool RMSNormVJP::is_equivalent(const Primitive& other) const { const RMSNormVJP& a_other = static_cast(other); return eps_ == a_other.eps_; diff --git a/mlx/fast.h b/mlx/fast.h index 934fadc2b7..d38a7c5033 100644 --- a/mlx/fast.h +++ b/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& weight, + float eps, + StreamOrDevice s = {}); + MLX_API array layer_norm( const array& x, const std::optional& weight, diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 0d2f861045..4f4d8b27d8 100644 --- a/mlx/fast_primitives.h +++ b/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)> fallback, + float eps) + : Custom(stream, std::move(fallback)), eps_(eps) {} + + static bool use_fallback(Stream stream); + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override { + throw std::runtime_error("NYI"); + } + void eval_gpu(const std::vector& inputs, std::vector& 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( diff --git a/python/src/fast.cpp b/python/src/fast.cpp index e59357bc33..b12185c204 100644 --- a/python/src/fast.cpp +++ b/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, diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 5dacaa605c..beca3a2c41 100644 --- a/python/tests/test_fast.py +++ b/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))