ferrum_testkit/op_diff/
rms_norm.rs

1//! `rms_norm` op-diff harness — see `crate::op_diff` for the framework.
2
3use super::{random_vec, OpUnderTest, Output};
4
5/// One concrete rms_norm invocation. Inputs:
6///   - `x`: tokens × dim activation
7///   - `w`: dim weight (per-channel scale)
8///   - `eps`: the usual RMSNorm epsilon
9///
10/// Output: tokens × dim. This fixture uses F32 buffers on CPU and Metal;
11/// CUDA's existing default buffers are F16. Host results are returned as f32.
12pub struct RmsNormOp {
13    pub tokens: usize,
14    pub dim: usize,
15    pub eps: f32,
16}
17
18impl RmsNormOp {
19    /// Validate the shared fixture before allocation or backend submission.
20    pub fn validate(&self) -> Result<(), String> {
21        self.expected_output_len().map(|_| ())
22    }
23
24    /// Checked host result size, including the kernels' signed 32-bit row and
25    /// column indices and the final column-loop increment. These are kernel/host
26    /// limits, not a promise that the allocation fits available device memory.
27    pub fn expected_output_len(&self) -> Result<usize, String> {
28        if self.tokens == 0 || self.dim == 0 {
29            return Err("RMSNorm tokens and dim must be nonzero".into());
30        }
31        if !self.eps.is_finite() || self.eps <= 0.0 {
32            return Err("RMSNorm epsilon must be finite and positive".into());
33        }
34        let elements = self
35            .tokens
36            .checked_mul(self.dim)
37            .ok_or("RMSNorm tokens * dim overflows usize")?;
38        let bytes = elements
39            .checked_mul(std::mem::size_of::<f32>())
40            .ok_or("RMSNorm f32 buffer byte size overflows usize")?;
41        if bytes > isize::MAX as usize {
42            return Err("RMSNorm f32 buffer exceeds the host allocation address range".into());
43        }
44        // Metal advances columns by 32. CUDA's step is its block width, capped
45        // at 1024; rounding to a complete warp also covers the current unrounded
46        // launch. Even the thread processing column dim - 1 must safely perform
47        // its final increment before evaluating the loop condition again.
48        let cuda_step = self.dim.min(1024).div_ceil(32) * 32;
49        let max_loop_step = cuda_step.max(32);
50        let final_increment = (self.dim - 1).checked_add(max_loop_step);
51        if elements > i32::MAX as usize
52            || final_increment.is_none_or(|index| index > i32::MAX as usize)
53        {
54            return Err("RMSNorm shape exceeds the kernels' signed 32-bit indexing range".into());
55        }
56        Ok(elements)
57    }
58
59    fn output_len(&self) -> usize {
60        self.expected_output_len().expect("invalid RMSNorm fixture")
61    }
62
63    /// Inputs are derived from seed so per-backend runs see identical x/w.
64    fn build_input(&self, seed: u64) -> (Vec<f32>, Vec<f32>) {
65        let x = random_vec(self.output_len(), -2.0, 2.0, seed);
66        let w = random_vec(self.dim, 0.5, 1.5, seed.wrapping_add(1));
67        (x, w)
68    }
69}
70
71impl OpUnderTest for RmsNormOp {
72    fn name(&self) -> &str {
73        "rms_norm"
74    }
75
76    fn run_cpu(&self, seed: u64) -> Output {
77        use ferrum_kernels::backend::cpu::CpuBackend;
78        use ferrum_kernels::backend::Backend;
79
80        let (x, w) = self.build_input(seed);
81        let mut ctx = CpuBackend::new_context();
82        let x_buf = CpuBackend::from_slice(&x);
83        let w_buf = CpuBackend::from_slice(&w);
84        let mut out = CpuBackend::alloc(self.output_len());
85        CpuBackend::rms_norm(
86            &mut ctx,
87            &x_buf,
88            &w_buf,
89            self.eps,
90            &mut out,
91            self.tokens,
92            self.dim,
93        );
94        CpuBackend::sync(&mut ctx);
95        CpuBackend::to_vec(&out, self.output_len())
96    }
97
98    #[cfg(all(target_os = "macos", feature = "metal"))]
99    fn run_metal(&self, seed: u64) -> Output {
100        use ferrum_kernels::backend::metal::MetalBackend;
101        use ferrum_kernels::backend::Backend;
102
103        let (x, w) = self.build_input(seed);
104        let mut ctx = MetalBackend::new_context();
105        let x_buf = MetalBackend::from_slice(&x);
106        let w_buf = MetalBackend::from_slice(&w);
107        let mut out = MetalBackend::alloc(self.output_len());
108        MetalBackend::rms_norm(
109            &mut ctx,
110            &x_buf,
111            &w_buf,
112            self.eps,
113            &mut out,
114            self.tokens,
115            self.dim,
116        );
117        MetalBackend::sync_checked(&mut ctx)
118            .unwrap_or_else(|error| panic!("RMSNorm Metal completion failed: {error}"));
119        MetalBackend::to_vec(&out, self.output_len())
120    }
121
122    #[cfg(feature = "cuda")]
123    fn run_cuda(&self, seed: u64) -> Output {
124        use ferrum_kernels::backend::cuda::CudaBackend;
125        use ferrum_kernels::backend::Backend;
126
127        let (x, w) = self.build_input(seed);
128        let mut ctx = CudaBackend::new_context();
129        let x_buf = CudaBackend::from_slice(&x);
130        let w_buf = CudaBackend::from_slice(&w);
131        let mut out = CudaBackend::alloc(self.output_len());
132        CudaBackend::rms_norm(
133            &mut ctx,
134            &x_buf,
135            &w_buf,
136            self.eps,
137            &mut out,
138            self.tokens,
139            self.dim,
140        );
141        CudaBackend::sync(&mut ctx);
142        CudaBackend::to_vec(&out, self.output_len())
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149
150    fn operation(tokens: usize, dim: usize, eps: f32) -> RmsNormOp {
151        RmsNormOp { tokens, dim, eps }
152    }
153
154    #[test]
155    fn validates_result_shape_before_allocating_inputs() {
156        let op = operation(3, 33, 1e-6);
157        op.validate().unwrap();
158        assert_eq!(op.expected_output_len().unwrap(), 99);
159        let output = op.run_cpu(7);
160        assert_eq!(output.len(), op.expected_output_len().unwrap());
161        assert!(output.iter().all(|value| value.is_finite()));
162    }
163
164    #[test]
165    fn rejects_empty_shapes_and_nonpositive_or_nonfinite_epsilon() {
166        for (tokens, dim) in [(0, 32), (1, 0), (0, 0)] {
167            assert!(operation(tokens, dim, 1e-6).validate().is_err());
168        }
169        for eps in [0.0, -0.0, -1e-6, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
170            assert!(operation(1, 32, eps).validate().is_err());
171        }
172    }
173
174    #[test]
175    fn rejects_shape_byte_and_kernel_index_overflow_without_allocating() {
176        for (tokens, dim) in [
177            (usize::MAX, 2),
178            (1, usize::MAX / std::mem::size_of::<f32>() + 1),
179            (1, i32::MAX as usize),
180            (2, i32::MAX as usize / 2 + 1),
181        ] {
182            assert!(operation(tokens, dim, 1e-6).expected_output_len().is_err());
183        }
184    }
185
186    #[test]
187    #[cfg(target_pointer_width = "64")]
188    fn validates_final_cuda_loop_increment_without_allocating() {
189        let last_safe_dim = i32::MAX as usize - (1024 - 1);
190        assert_eq!(
191            operation(1, last_safe_dim, 1e-6).expected_output_len(),
192            Ok(last_safe_dim)
193        );
194        // These shapes fit signed element indexing and host address space, but
195        // a CUDA thread's last 1024-wide increment exceeds the signed range.
196        for dim in [last_safe_dim + 1, i32::MAX as usize - 31] {
197            assert!(operation(1, dim, 1e-6).expected_output_len().is_err());
198        }
199    }
200
201    #[test]
202    fn direct_executor_use_rejects_invalid_fixture_before_allocation() {
203        let op = operation(usize::MAX, 2, 1e-6);
204        assert!(std::panic::catch_unwind(|| op.run_cpu(7)).is_err());
205    }
206}
ferrum_testkit/op_diff/rms_norm.rs

ferrum_testkit/op_diff/
rms_norm.rs
