rms_norm: run f16 rows natively instead of widening the whole tensor

core/src/ops/nn/rms_norm.rs

@@ -41,28 +41,32 @@ impl EvalOp for RmsNorm {
             && self.axis == input.rank() - 1
         {
             let eps_f32: f32 = self.eps.cast_to_scalar::<f32>()?;
-            let already_f32 = in_dt == DatumType::F32;
-            let mut buf = if already_f32 {
-                input.into_tensor()
-            } else {
-                input.cast_to::<f32>()?.into_owned()
-            };
-            let row_len = buf.shape()[self.axis];
+            let row_len = input.shape()[self.axis];
+            let mut buf = input.into_tensor();
             if row_len > 0 {
-                let data = unsafe { buf.as_slice_mut_unchecked::<f32>() };
                 let rms_norm = &tract_linalg::ops().rms_norm_f32;
-                let total = data.len();
-                tract_linalg::multithread::par_chunks_mut(data, row_len, total, |_, chunk| {
-                    for row in chunk.chunks_mut(row_len) {
-                        rms_norm(row, eps_f32);
-                    }
-                    Ok(())
-                })?;
+                if in_dt == DatumType::F32 {
+                    let data = unsafe { buf.as_slice_mut_unchecked::<f32>() };
+                    let total = data.len();
+                    tract_linalg::multithread::par_chunks_mut(data, row_len, total, |_, chunk| {
+                        for row in chunk.chunks_mut(row_len) {
+                            rms_norm(row, eps_f32);
+                        }
+                        Ok(())
+                    })?;
+                } else {
+                    let rms_norm = &tract_linalg::ops().rms_norm_f16;
+                    let data = unsafe { buf.as_slice_mut_unchecked::<f16>() };
+                    let total = data.len();
+                    tract_linalg::multithread::par_chunks_mut(data, row_len, total, |_, chunk| {
+                        for row in chunk.chunks_mut(row_len) {
+                            rms_norm(row, eps_f32);
+                        }
+                        Ok(())
+                    })?;
+                }
             }
-            if already_f32 {
-                return Ok(tvec![buf.into_tvalue()]);
-            }
-            return Ok(tvec![buf.cast_to_dt(in_dt)?.into_owned().into()]);
+            return Ok(tvec![buf.into_tvalue()]);
         }
 
         // Slow path: original 4-call composition (kept for non-contiguous axes).

linalg/src/arm64.rs

@@ -497,6 +497,7 @@ pub fn plug(ops: &mut Ops) {
     ops.sum_f32 = Box::new(|| arm64simd_sum_f32_16n::red());
     ops.mul_by_scalar_f32 = Box::new(|| arm64simd_mul_by_scalar_f32_16n::ew());
     ops.softmax2_fastcompact_f32 = Box::new(|| arm64simd_softmax2_fastcompact_f32_16n::red());
+    ops.rms_norm_f16 = Box::new(arm64simd_rms_norm_f16);
     ops.rms_norm_f32 = Box::new(arm64simd_rms_norm_f32);
     #[cfg(not(feature = "no_fp16"))]
     if has_fp16() {

linalg/src/arm64/arm64simd.rs

@@ -23,6 +23,7 @@ pub use hardswish::arm64simd_hardswish_f32_8n;
 pub use leaky_relu::arm64simd_leaky_relu_f32_8n;
 pub use max::arm64simd_max_f32_16n;
 pub use min::arm64simd_min_f32_16n;
+pub use rms_norm::rms_norm_f16 as arm64simd_rms_norm_f16;
 pub use rms_norm::rms_norm_f32 as arm64simd_rms_norm_f32;
 pub use silu::arm64simd_silu_f32_4n;
 pub use silu_fused::arm64simd_silu_f32_4n_fused;

linalg/src/arm64/arm64simd/rms_norm.rs

@@ -13,6 +13,8 @@
 // dispatcher in `core::ops::nn::RmsNorm::eval` is already arch-neutral and
 // will pick this up automatically.
 
+use tract_data::internal::f16;
+
 #[target_feature(enable = "neon")]
 unsafe fn rms_norm_f32_inner(buf: &mut [f32], eps: f32) {
     use std::arch::aarch64::*;
@@ -102,6 +104,108 @@ pub fn rms_norm_f32(buf: &mut [f32], eps: f32) {
     unsafe { rms_norm_f32_inner(buf, eps) }
 }
 
+// f16 rows, arithmetic unchanged from `rms_norm_f32_inner`: FCVTL widens each
+// 16-element group into the same four f32 registers before the same FMLA chain,
+// and FCVTN rounds to nearest-even on the way out, matching `f16::from_f32`.
+#[target_feature(enable = "neon")]
+unsafe fn rms_norm_f16_inner(buf: &mut [f16], eps: f32) {
+    use std::arch::aarch64::*;
+    let n = buf.len();
+    let chunks = n / 16;
+    let tail_start = chunks * 16;
+    let ptr = buf.as_mut_ptr();
+
+    let mut sum_sq: f32 = 0.0;
+    if chunks > 0 {
+        let p = ptr;
+        let c = chunks;
+        let mut sum_v: float32x4_t = vdupq_n_f32(0.0);
+        unsafe {
+            std::arch::asm!("
+                movi v1.4s, 0
+                movi v2.4s, 0
+                movi v3.4s, 0
+                2:
+                    ld1 {{v16.8h, v17.8h}}, [{p}], 32
+                    fcvtl  v4.4s, v16.4h
+                    fcvtl2 v5.4s, v16.8h
+                    fcvtl  v6.4s, v17.4h
+                    fcvtl2 v7.4s, v17.8h
+                    fmla v0.4s, v4.4s, v4.4s
+                    fmla v1.4s, v5.4s, v5.4s
+                    fmla v2.4s, v6.4s, v6.4s
+                    fmla v3.4s, v7.4s, v7.4s
+                    subs {c}, {c}, 1
+                    bne 2b
+                fadd v0.4s, v0.4s, v1.4s
+                fadd v2.4s, v2.4s, v3.4s
+                fadd v0.4s, v0.4s, v2.4s
+                ",
+                p = inout(reg) p => _,
+                c = inout(reg) c => _,
+                inout("v0") sum_v,
+                out("v1") _, out("v2") _, out("v3") _,
+                out("v4") _, out("v5") _, out("v6") _, out("v7") _,
+                out("v16") _, out("v17") _,
+            );
+        }
+        sum_sq = vaddvq_f32(sum_v);
+    }
+    for i in tail_start..n {
+        let x = unsafe { buf.get_unchecked(i).to_f32() };
+        sum_sq += x * x;
+    }
+
+    let mean_sq = sum_sq / (n as f32);
+    let inv_std = (mean_sq + eps).sqrt().recip();
+
+    if chunks > 0 {
+        let p = ptr;
+        let c = chunks;
+        let inv_v: float32x4_t = vdupq_n_f32(inv_std);
+        unsafe {
+            std::arch::asm!("
+                2:
+                    ld1 {{v16.8h, v17.8h}}, [{p}]
+                    fcvtl  v4.4s, v16.4h
+                    fcvtl2 v5.4s, v16.8h
+                    fcvtl  v6.4s, v17.4h
+                    fcvtl2 v7.4s, v17.8h
+                    fmul v4.4s, v4.4s, v0.4s
+                    fmul v5.4s, v5.4s, v0.4s
+                    fmul v6.4s, v6.4s, v0.4s
+                    fmul v7.4s, v7.4s, v0.4s
+                    fcvtn  v16.4h, v4.4s
+                    fcvtn2 v16.8h, v5.4s
+                    fcvtn  v17.4h, v6.4s
+                    fcvtn2 v17.8h, v7.4s
+                    st1 {{v16.8h, v17.8h}}, [{p}], 32
+                    subs {c}, {c}, 1
+                    bne 2b
+                ",
+                p = inout(reg) p => _,
+                c = inout(reg) c => _,
+                in("v0") inv_v,
+                out("v4") _, out("v5") _, out("v6") _, out("v7") _,
+                out("v16") _, out("v17") _,
+            );
+        }
+    }
+    for i in tail_start..n {
+        unsafe {
+            let x = buf.get_unchecked(i).to_f32();
+            *buf.get_unchecked_mut(i) = f16::from_f32(x * inv_std);
+        }
+    }
+}
+
+pub fn rms_norm_f16(buf: &mut [f16], eps: f32) {
+    if buf.is_empty() {
+        return;
+    }
+    unsafe { rms_norm_f16_inner(buf, eps) }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -160,5 +264,26 @@ mod tests {
         let mut x: Vec<f32> = vec![];
         rms_norm_f32(&mut x, 1e-5);
         assert!(x.is_empty());
+        let mut h: Vec<f16> = vec![];
+        rms_norm_f16(&mut h, 1e-5);
+        assert!(h.is_empty());
+    }
+
+    #[test]
+    fn f16_matches_widening_the_row_to_f32() {
+        for len in [1usize, 7, 15, 16, 17, 33, 64, 129, 1024, 4096] {
+            let row: Vec<f16> =
+                (0..len).map(|i| f16::from_f32((i as f32 * 0.13).sin() * 5.0)).collect();
+
+            let mut got = row.clone();
+            rms_norm_f16(&mut got, 1e-5);
+
+            let mut widened: Vec<f32> = row.iter().map(|x| x.to_f32()).collect();
+            rms_norm_f32(&mut widened, 1e-5);
+            let want: Vec<f16> = widened.into_iter().map(f16::from_f32).collect();
+
+            let mismatch = got.iter().zip(&want).position(|(a, b)| a.to_bits() != b.to_bits());
+            assert_eq!(mismatch, None, "len {len}");
+        }
     }
 }

linalg/src/generic/rms_norm.rs

@@ -1,3 +1,5 @@
+use tract_data::internal::f16;
+
 /// Generic scalar reference implementation of fused row-wise RmsNorm.
 ///   out_i = x_i * rsqrt(mean(x_i²) + eps)
 ///
@@ -17,6 +19,28 @@ pub fn rms_norm_f32(buf: &mut [f32], eps: f32) {
     }
 }
 
+/// f16 counterpart of [`rms_norm_f32`], for hosts with no native f16 kernel.
+///
+/// Widens one row at a time into a reused thread-local buffer and defers to
+/// whichever f32 kernel this host registered, so an arch that has a fast f32
+/// kernel but no f16 one keeps its arithmetic and its speed. f16 -> f32 is
+/// exact, so the result matches widening the row in the caller.
+pub fn rms_norm_f16(buf: &mut [f16], eps: f32) {
+    if buf.is_empty() {
+        return;
+    }
+    thread_local! {
+        static WIDE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
+    }
+    WIDE.with(|wide| {
+        let mut wide = wide.borrow_mut();
+        wide.clear();
+        wide.extend(buf.iter().map(|x| x.to_f32()));
+        (crate::ops().rms_norm_f32)(&mut wide, eps);
+        buf.iter_mut().zip(wide.iter()).for_each(|(h, f)| *h = f16::from_f32(*f));
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

linalg/src/lib.rs

@@ -115,6 +115,7 @@ pub struct Ops {
     /// Replaces a 4-call composition (MeanOfSquares + Add + Rsqrt + Mul) with
     /// a single 2-pass kernel. Called once per row by `core::ops::nn::RmsNorm`
     /// when the input is f32 and the axis is the last (contiguous) one.
+    pub rms_norm_f16: Box<dyn Fn(&mut [f16], f32) + Send + Sync>,
     pub rms_norm_f32: Box<dyn Fn(&mut [f32], f32) + Send + Sync>,
 }
 
@@ -253,6 +254,7 @@ pub fn generic() -> Ops {
         */
         softmax2_fastcompact_f16: Box::new(|| generic::reduce::softmax_l2::HSoftMaxL2::red()),
         softmax2_fastcompact_f32: Box::new(|| generic::reduce::softmax_l2::SSoftMaxL2::red()),
+        rms_norm_f16: Box::new(generic::rms_norm::rms_norm_f16),
         rms_norm_f32: Box::new(generic::rms_norm::rms_norm_f32),
     };
     crate::generic::mmm::plug(&mut ops);