ferrum_testkit/op_diff/
mod.rs

1//! Cross-backend op-diff harness — PLAYBOOK § 3 L1.
2//!
3//! Runs the same op on CPU (reference) and on each available accelerator
4//! (Metal / CUDA), then reports the **NMSE** (normalized mean-squared
5//! error) of each accelerator's output relative to CPU's. Modelled on
6//! llama.cpp's `tests/test-backend-ops.cpp` `NMSE = mse(a,b) / mse(a,0)`
7//! comparison rather than naive max-abs-diff: NMSE is invariant to the
8//! magnitude of the reference output, so a well-tuned kernel will sit
9//! at the same NMSE regardless of input scaling.
10//!
11//! # Usage
12//!
13//! ```ignore
14//! use ferrum_testkit::op_diff::{compare_backends, NMSE_FP16_TOL, rms_norm::RmsNormOp};
15//!
16//! let report = compare_backends(&RmsNormOp { tokens: 4, dim: 4096, eps: 1e-6 }, 42);
17//! if let Some(nmse) = report.metal_nmse {
18//!     assert!(nmse < NMSE_FP16_TOL, "metal rms_norm NMSE {nmse} exceeds fp16 tol");
19//! }
20//! ```
21//!
22//! # Tolerance buckets (PLAYBOOK § 3.1)
23//!
24//! - `NMSE_FP32_TOL = 1e-7` — fp32 kernels must agree with CPU below this.
25//! - `NMSE_FP16_TOL = 1e-6` — fp16 kernels (Metal default storage).
26//!
27//! Tighter bucketing per op is welcome — define op-specific constants
28//! once empirical baselines are stable.
29
30pub mod activation_bridge;
31pub mod argmax_rows;
32pub mod embedding_lookup;
33pub mod flash_attention;
34pub mod fused_add_rms_norm;
35pub mod gemm;
36pub mod kv_cache_append;
37pub mod marlin_matmul; // stub — see file docs
38pub mod metal_context;
39pub mod paged_varlen_attn; // stub — see file docs
40pub mod qk_norm_rope;
41pub mod required;
42pub mod residual_add;
43pub mod rms_norm;
44pub mod silu_mul;
45pub mod split_qkv;
46pub mod transpose_head_to_token;
47
48/// fp32 kernels — should agree with CPU below this.
49pub const NMSE_FP32_TOL: f64 = 1e-7;
50/// fp16 storage / Metal accumulation — slightly larger tol.
51pub const NMSE_FP16_TOL: f64 = 1e-6;
52
53pub use ferrum_bench_core::release_regression::numerics::nmse;
54
55/// Output of a single op invocation. Each backend produces its own
56/// `Vec<f32>` after `to_vec()`-ing its buffer to host.
57pub type Output = Vec<f32>;
58
59/// A single op-under-test: knows how to run itself on each backend.
60pub trait OpUnderTest {
61    /// Display name (used in test failure messages).
62    fn name(&self) -> &str;
63
64    /// Run on CPU (reference). Always available.
65    fn run_cpu(&self, seed: u64) -> Output;
66
67    /// Run on Metal. Only available with `cfg(all(target_os = "macos", feature = "metal"))`.
68    #[cfg(all(target_os = "macos", feature = "metal"))]
69    fn run_metal(&self, seed: u64) -> Output;
70
71    /// Run on CUDA. Only available with `cfg(feature = "cuda")`.
72    #[cfg(feature = "cuda")]
73    fn run_cuda(&self, seed: u64) -> Output;
74}
75
76/// Cross-backend comparison result.
77///
78/// `cpu` is the reference output. `metal_nmse` / `cuda_nmse` are `None`
79/// on builds that don't include that backend.
80#[derive(Debug)]
81pub struct NmseReport {
82    pub op: String,
83    pub seed: u64,
84    pub cpu: Vec<f32>,
85    pub metal_nmse: Option<f64>,
86    pub cuda_nmse: Option<f64>,
87}
88
89impl NmseReport {
90    /// True if every available accelerator matches CPU below `tol`.
91    pub fn within_tol(&self, tol: f64) -> bool {
92        self.metal_nmse.map_or(true, |n| n < tol) && self.cuda_nmse.map_or(true, |n| n < tol)
93    }
94}
95
96/// Run `op` on every backend the current build supports and assemble
97/// the comparison report.
98pub fn compare_backends(op: &dyn OpUnderTest, seed: u64) -> NmseReport {
99    let cpu = op.run_cpu(seed);
100    let metal_nmse = run_metal_nmse(op, &cpu, seed);
101    let cuda_nmse = run_cuda_nmse(op, &cpu, seed);
102    NmseReport {
103        op: op.name().to_string(),
104        seed,
105        cpu,
106        metal_nmse,
107        cuda_nmse,
108    }
109}
110
111#[cfg(all(target_os = "macos", feature = "metal"))]
112fn run_metal_nmse(op: &dyn OpUnderTest, cpu: &[f32], seed: u64) -> Option<f64> {
113    Some(nmse(cpu, &op.run_metal(seed)))
114}
115
116#[cfg(not(all(target_os = "macos", feature = "metal")))]
117fn run_metal_nmse(_op: &dyn OpUnderTest, _cpu: &[f32], _seed: u64) -> Option<f64> {
118    None
119}
120
121#[cfg(feature = "cuda")]
122fn run_cuda_nmse(op: &dyn OpUnderTest, cpu: &[f32], seed: u64) -> Option<f64> {
123    Some(nmse(cpu, &op.run_cuda(seed)))
124}
125
126#[cfg(not(feature = "cuda"))]
127fn run_cuda_nmse(_op: &dyn OpUnderTest, _cpu: &[f32], _seed: u64) -> Option<f64> {
128    None
129}
130
131/// Convenience: deterministic uniform-random `Vec<f32>` in `[lo, hi)`.
132pub fn random_vec(n: usize, lo: f32, hi: f32, seed: u64) -> Vec<f32> {
133    use rand::{Rng, SeedableRng};
134    let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
135    (0..n).map(|_| rng.random_range(lo..hi)).collect()
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[test]
143    fn random_vec_determinism() {
144        let a = random_vec(100, -1.0, 1.0, 42);
145        let b = random_vec(100, -1.0, 1.0, 42);
146        assert_eq!(a, b);
147    }
148}
ferrum_testkit/op_diff/mod.rs

ferrum_testkit/op_diff/
mod.rs
