From 948b9c72de3123e54367eb7db775c917ff1efcc2 Mon Sep 17 00:00:00 2001 From: Ogilthorp3 Date: Sat, 2 May 2026 15:27:22 -0400 Subject: [PATCH] fast::rms_norm: accept Option<&Array> for weight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Aligns with the existing layer_norm_device pattern in the same file and matches Python's mx.fast.rms_norm(x, None, eps) semantics. The underlying C function (mlx_fast_rms_norm) already supports a null weight pointer — it checks weight.ctx and routes to the no-affine path internally, skipping the per-element multiply kernel pass. Only the Rust wrapper needed updating. Why it matters: when callers want plain RMS normalization without an affine weight (common in attention's q/k norm step where the weight would be all-ones), passing None saves one Metal kernel dispatch per call. For a 40-layer LLM at 1045-token prefill with two q/k norms per attention layer, that's ~80 saved kernel passes per forward. Existing call sites continue to compile via auto-derived From<&'a Array> for Option<&'a Array>. The single owned-Array test caller (test_rms_norm) was changed to pass &weight. Validated: cargo test --release -p mlx-rs --lib fast::tests -- --test-threads=1 test result: ok. 6 passed; 0 failed cargo test --release -p mlx-rs --lib nn::normalization -- --test-threads=1 test result: ok. 5 passed; 0 failed Discovered while investigating a 2x LM-prefill perf gap between mlx-rs and Python's mlx_lm bindings on Apple M4 Max — the qwen3_5 LinearAttention's q/k norm path was forced to construct an ones array and call rms_norm with it, then multiply by a scale array afterward, where Python passes None and applies a scalar multiply inline. Microbenches and per-layer profiling confirmed parity at the per-op level; this signature gap was one of the structural divergences from Python's API. --- mlx-rs/src/fast.rs | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/mlx-rs/src/fast.rs b/mlx-rs/src/fast.rs index 6d9d09358..669fc5990 100644 --- a/mlx-rs/src/fast.rs +++ b/mlx-rs/src/fast.rs @@ -182,22 +182,28 @@ pub fn scaled_dot_product_attention_device<'a>( /// # Params /// /// - x: input array -/// - weight: A multiplicative weight to scale the result by. The `weight` should be one-dimensional with the same size as the last axis of `x`. +/// - weight: A multiplicative weight to scale the result by. The `weight` should be one-dimensional +/// with the same size as the last axis of `x`. If not given, no scaling will occur (matches +/// Python's `mx.fast.rms_norm(x, None, eps)` semantics — the underlying Metal kernel skips +/// the per-element multiply, saving a kernel pass). /// - eps: A small additive constant for numerical stability /// - stream: stream or device to evaluate on #[generate_macro(customize(root = "$crate::fast"))] #[default_device] -pub fn rms_norm_device( - x: impl AsRef, - weight: impl AsRef, - eps: f32, +pub fn rms_norm_device<'a>( + #[named] x: impl AsRef, + #[optional] weight: impl Into>, + #[named] eps: f32, #[optional] stream: impl AsRef, ) -> Result { Array::try_from_op(|res| unsafe { mlx_sys::mlx_fast_rms_norm( res, x.as_ref().as_ptr(), - weight.as_ref().as_ptr(), + weight + .into() + .map(|w| w.as_ptr()) + .unwrap_or_else(|| mlx_sys::mlx_array_new()), eps, stream.as_ref().as_ptr(), ) @@ -308,7 +314,7 @@ mod tests { assert_eq!(a.dtype(), crate::Dtype::Float32); let weight = Array::ones::(&[16]).unwrap(); - let result = rms_norm(a, weight, 1e-5).unwrap(); + let result = rms_norm(a, &weight, 1e-5).unwrap(); assert_eq!(result.shape(), [2, 8, 16]); assert_eq!(result.dtype(), crate::Dtype::Float32); assert_float_eq!(