tract_core/ops/nn/
rms_norm.rs

1use crate::internal::*;
2use crate::ops::binary::{BinMiniOp, TypedBinOp};
3use crate::ops::element_wise::ElementWiseOp;
4use crate::ops::math::{Add, Mul, Rsqrt};
5use crate::ops::nn::{Reduce, Reducer};
6use tract_itertools::Itertools;
7
8#[derive(Clone, Debug, Hash, PartialEq, Eq)]
9pub struct RmsNorm {
10    pub axis: usize,
11    pub eps: Arc<Tensor>,
12}
13
14impl Op for RmsNorm {
15    fn name(&self) -> StaticName {
16        "RmsNorm".to_string().into()
17    }
18    fn info(&self) -> TractResult<Vec<String>> {
19        Ok(vec![format!("axis: {:?}, eps: {:?}", self.axis, self.eps)])
20    }
21    op_as_typed_op!();
22}
23
24impl EvalOp for RmsNorm {
25    fn is_stateless(&self) -> bool {
26        true
27    }
28
29    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
30        let input = args_1!(inputs);
31        let in_dt = input.datum_type();
32
33        // Fast path: F32 or F16 input where the normalised axis is the last
34        // (contiguous) one. Use the fused tract_linalg::rms_norm_f32 kernel
35        // (AVX-512 when available; scalar fallback otherwise) instead of the
36        // 4-call MeanOfSquares + Add + Rsqrt + Mul composition below. ~16-18x
37        // faster on Cascade Lake AVX-512, ~equivalent on the scalar fallback
38        // since the composition is also memory-bandwidth bound.
39        if matches!(in_dt, DatumType::F32 | DatumType::F16)
40            && input.rank() > 0
41            && self.axis == input.rank() - 1
42        {
43            let eps_f32: f32 = self.eps.cast_to_scalar::<f32>()?;
44            let already_f32 = in_dt == DatumType::F32;
45            let mut buf = if already_f32 {
46                input.into_tensor()
47            } else {
48                input.cast_to::<f32>()?.into_owned()
49            };
50            let row_len = buf.shape()[self.axis];
51            if row_len > 0 {
52                let data = unsafe { buf.as_slice_mut_unchecked::<f32>() };
53                let rms_norm = &tract_linalg::ops().rms_norm_f32;
54                let total = data.len();
55                tract_linalg::multithread::par_chunks_mut(data, row_len, total, |_, chunk| {
56                    for row in chunk.chunks_mut(row_len) {
57                        rms_norm(row, eps_f32);
58                    }
59                    Ok(())
60                })?;
61            }
62            if already_f32 {
63                return Ok(tvec![buf.into_tvalue()]);
64            }
65            return Ok(tvec![buf.cast_to_dt(in_dt)?.into_owned().into()]);
66        }
67
68        // Slow path: original 4-call composition (kept for non-contiguous axes).
69        let already_f32 = in_dt == DatumType::F32;
70        let input_f32 =
71            if already_f32 { input.into_tensor() } else { input.cast_to::<f32>()?.into_owned() };
72        // eps inherits the input dtype from the declutter pattern (F16 when the
73        // surrounding LayerNorm chain is F16). The MeanOfSquares + Add + Rsqrt
74        // + Mul chain below all runs at F32, so eps must be cast to match —
75        // otherwise the Add::eval call below panics with
76        //   "tensor is F32, accessed as F16"
77        // when input is F16.
78        let eps = self.eps.cast_to::<f32>()?.into_owned();
79        let a1 = Reducer::MeanOfSquares.reduce(&[self.axis], &input_f32)?;
80        let mut a2 = Add.eval(a1.into_tvalue(), eps.into_tvalue(), DatumType::F32)?;
81        Rsqrt {}.eval_in_place(&mut a2, None)?;
82        let a3 = Mul.eval(a2.into_tvalue(), input_f32.into_tvalue(), DatumType::F32)?;
83        if already_f32 {
84            return Ok(tvec![a3.into_tvalue()]);
85        }
86        Ok(tvec![a3.cast_to_dt(in_dt)?.into_owned().into()])
87    }
88}
89
90impl TypedOp for RmsNorm {
91    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
92        ensure!(self.eps.rank() == 0, "RmsNorm: eps must be a rank-0 tensor");
93        ensure!(
94            self.axis < inputs[0].rank(),
95            "RmsNorm: axis {} is out of bounds for input rank {}",
96            self.axis,
97            inputs[0].rank()
98        );
99        let dt = inputs[0].datum_type;
100        let fact = dt.fact(inputs[0].shape.clone());
101        Ok(tvec!(fact))
102    }
103
104    fn input_roi(
105        &self,
106        model: &TypedModel,
107        node: &TypedNode,
108    ) -> TractResult<Option<TVec<Option<TDim>>>> {
109        crate::optim::propagate_roi::bubble_roi(model, node)
110    }
111
112    fn axes_mapping(
113        &self,
114        inputs: &[&TypedFact],
115        _outputs: &[&TypedFact],
116    ) -> TractResult<AxesMapping> {
117        let rank = inputs[0].rank();
118        let mut letters = 'a'..;
119        let axes = (0..rank)
120            .map(|ix| {
121                Axis::new(letters.next().unwrap(), inputs.len(), 1).input(0, ix).output(0, ix)
122            })
123            .collect_vec();
124        AxesMapping::new(1, 1, axes)
125    }
126
127    fn change_axes(
128        &self,
129        model: &TypedModel,
130        node: &TypedNode,
131        _io: InOut,
132        change: &AxisOp,
133    ) -> TractResult<Option<AxisChangeConsequence>> {
134        if let Some(axis) = change.transform_axis(self.axis) {
135            let op = Some(Box::new(RmsNorm { axis, eps: self.eps.clone() }) as _);
136            Ok(Some(AxisChangeConsequence::new(model, node, op, change)))
137        } else {
138            Ok(None)
139        }
140    }
141
142    fn slice(
143        &self,
144        patch: &mut TypedModelPatch,
145        _model: &TypedModel,
146        node: &TypedNode,
147        _prefix: &str,
148        inputs: &[OutletId],
149        output_axis: usize,
150        _start: &TDim,
151        _end: &TDim,
152    ) -> TractResult<Option<TVec<OutletId>>> {
153        rule_if!(output_axis != self.axis);
154        patch.wire_node(&node.name, self.clone(), inputs).map(Some)
155    }
156
157    fn cost(&self, inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
158        let dt = inputs[0].datum_type;
159        let count: TDim = inputs[0].shape.iter().product();
160        // per element: square + accumulate + mul by rsqrt ≈ 3 FMA
161        // per reduction group: 1 div (rsqrt)
162        let groups: TDim = inputs[0]
163            .shape
164            .iter()
165            .enumerate()
166            .filter(|(i, _)| *i != self.axis)
167            .map(|(_, d)| d)
168            .product();
169        Ok(tvec!((Cost::FMA(dt), count * 3), (Cost::Div(dt), groups)))
170    }
171
172    as_op!();
173}
174
175/// Search pattern => A = A * RSQRT(MEAN_OF_SQUARES(A) + EPS)
176pub fn detect_rms_norm(
177    op: &Reduce,
178    model: &TypedModel,
179    node: &TypedNode,
180) -> TractResult<Option<TypedModelPatch>> {
181    rule_if!(op.reducer == Reducer::MeanOfSquares);
182    rule_if!(op.axes.len() == 1);
183    let axis = op.axes[0];
184
185    let in_fact = model.node_input_facts(node.id)?[0];
186    let dt = in_fact.datum_type;
187
188    // Only F16 and F32 is supported.
189    rule_if!(matches!(dt, DatumType::F32 | DatumType::F16));
190
191    // Identify Add operator
192    rule_if_some!(add_succ = model.single_succ(node.id)?);
193    rule_if_some!(add_succ_op = add_succ.op_as::<TypedBinOp>());
194    rule_if!(add_succ_op.0.is::<Add>());
195
196    // Retrieve epsilon
197    let add_consts = model.collect_const_inputs(add_succ);
198    rule_if!(add_consts.len() == 1);
199    let eps = add_consts[0].val().clone();
200    rule_if!(eps.len() == 1);
201    rule_if!(eps.datum_type() == dt);
202    let eps = eps.into_tensor().into_shape(&[])?.into_arc_tensor();
203
204    // Identify Rsqrt
205    rule_if_some!(rsqrt_succ = model.single_succ(add_succ.id)?);
206    rule_if_some!(rsqrt_succ_op = rsqrt_succ.op_as::<ElementWiseOp>());
207    rule_if!(rsqrt_succ_op.0.is::<Rsqrt>());
208
209    // Identify Mul: RSQRT(...) * A
210    rule_if_some!(mul_succ = model.find_succ_bin_with_outlet::<Mul>(rsqrt_succ, &node.inputs[0]));
211
212    let mut patch = TypedModelPatch::default();
213    let rsm_input = patch.taps(model, &node.inputs)?;
214    let out =
215        patch.wire_node(format!("{}.rms_norm", node.name), RmsNorm { axis, eps }, &rsm_input)?;
216
217    patch.shunt_outside(model, mul_succ.id.into(), out[0])?;
218    Ok(Some(patch))
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224    use crate::ops::nn::RmsNorm;
225
226    /// Regression: the declutter pattern (`detect_rms_norm`) stores `eps` with
227    /// the input dtype (F16 when the surrounding LayerNorm chain is F16) — see
228    /// `rule_if!(eps.datum_type() == dt)` above. The eval path runs at F32, so
229    /// it must cast `self.eps` to F32 before using it. Without the cast in
230    /// `RmsNorm::eval`, this test panics with "tensor is F32, accessed as F16".
231    #[test]
232    fn eval_with_f16_eps_and_f16_input() {
233        let to_h = |x: f32| f16::from_f32(x);
234        let input = tensor1(&[to_h(1.0), to_h(2.0), to_h(3.0), to_h(4.0)]);
235        let eps = tensor0(to_h(1e-5)).into_arc_tensor();
236        let op = RmsNorm { axis: 0, eps };
237        let out = op.eval(tvec!(input.clone().into())).expect("eval should not panic");
238        let out = out.into_iter().next().unwrap().into_tensor();
239        assert_eq!(out.datum_type(), DatumType::F16);
240        assert_eq!(out.shape(), &[4]);
241        // Reference: rms = sqrt((1+4+9+16)/4 + eps) = sqrt(7.5 + 1e-5) ≈ 2.7386
242        // normalised: [1, 2, 3, 4] / 2.7386 ≈ [0.365, 0.730, 1.095, 1.461]
243        let got = unsafe { out.as_slice_unchecked::<f16>() };
244        let expected = [0.365_f32, 0.730, 1.095, 1.461];
245        for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
246            let diff = (g.to_f32() - e).abs();
247            assert!(diff < 0.01, "lane {i}: got {} expected {}", g.to_f32(), e);
248        }
249    }
250
251    /// Slow path: when the normalised axis is NOT the trailing one, the fast
252    /// path in `eval` (which dispatches to `tract_linalg::ops().rms_norm_f32`)
253    /// is skipped and the original 4-call `MeanOfSquares` + `Add` + `Rsqrt` +
254    /// `Mul` composition runs. Asserts the result is identical to a hand-
255    /// computed reference, so the slow path stays correct after the fast-path
256    /// addition.
257    #[test]
258    fn eval_with_non_trailing_axis_f32() {
259        // 2x3 input, axis=0 means we normalise across the 2 rows for each
260        // column independently:
261        //   col 0: [1, 4] → mean_sq = (1 + 16) / 2 =  8.5 → 1/√8.5
262        //   col 1: [2, 5] → mean_sq = (4 + 25) / 2 = 14.5 → 1/√14.5
263        //   col 2: [3, 6] → mean_sq = (9 + 36) / 2 = 22.5 → 1/√22.5
264        let input = tensor2(&[[1.0_f32, 2.0, 3.0], [4.0, 5.0, 6.0]]);
265        let eps = tensor0(0.0_f32).into_arc_tensor();
266        let op = RmsNorm { axis: 0, eps };
267        let out = op.eval(tvec!(input.into())).expect("eval should not panic");
268        let out = out.into_iter().next().unwrap().into_tensor();
269        assert_eq!(out.datum_type(), DatumType::F32);
270        assert_eq!(out.shape(), &[2, 3]);
271        let got = unsafe { out.as_slice_unchecked::<f32>() };
272        let inv = |ms: f32| ms.sqrt().recip();
273        let expected: [f32; 6] = [
274            1.0 * inv(8.5),
275            2.0 * inv(14.5),
276            3.0 * inv(22.5),
277            4.0 * inv(8.5),
278            5.0 * inv(14.5),
279            6.0 * inv(22.5),
280        ];
281        for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
282            let diff = (g - e).abs();
283            assert!(diff < 1e-5, "lane {i}: got {g}, want {e}, diff {diff}");
284        }
285    }
286}
tract_core/ops/nn/rms_norm.rs

tract_core/ops/nn/
rms_norm.rs
