rlx_runtime/
logit_verify.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Logit / output verification (plan #61).
6//!
7//! Borrowed from MAX's
8//! `tests/integration/accuracy/verify` pattern: every model gets a
9//! parity test that diffs RLX's output vs a reference (HuggingFace
10//! transformers, ONNX Runtime, hand-fused, ...) using cosine
11//! similarity, KL divergence, and absolute tolerance.
12//!
13//! Pure data layer — no HF / ORT integration here. Test code calls
14//! `compare(out, reference, tolerance)` and gets back a structured
15//! report it can `assert!` against. Hooking this up to specific
16//! reference implementations is per-bench wiring.
17
18#[derive(Debug, Clone, Copy)]
19pub struct Tolerance {
20    pub max_abs: f32,
21    pub max_rel: f32,
22    pub min_cosine: f32,
23}
24
25impl Tolerance {
26    /// Strict — for f32-vs-f32 comparisons.
27    pub const STRICT: Self = Self {
28        max_abs: 1e-4,
29        max_rel: 1e-3,
30        min_cosine: 0.9999,
31    };
32    /// Loose — for f16 / bf16 against f32 reference.
33    pub const LOOSE_F16: Self = Self {
34        max_abs: 5e-2,
35        max_rel: 5e-2,
36        min_cosine: 0.999,
37    };
38}
39
40#[derive(Debug, Clone)]
41pub struct Diff {
42    pub n: usize,
43    pub max_abs: f32,
44    pub max_rel: f32,
45    pub mean_abs: f32,
46    pub cosine: f32,
47    pub argmax_diff_index: usize,
48}
49
50#[derive(Debug)]
51pub enum VerifyError {
52    LengthMismatch { got: usize, expected: usize },
53    ToleranceExceeded { diff: Diff, tolerance: Tolerance },
54    Nonfinite,
55}
56
57impl std::fmt::Display for VerifyError {
58    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59        match self {
60            Self::LengthMismatch { got, expected } => {
61                write!(f, "length mismatch: got {got}, expected {expected}")
62            }
63            Self::ToleranceExceeded { diff, tolerance } => write!(
64                f,
65                "tolerance exceeded: max_abs={:.3e} max_rel={:.3e} cos={:.5} \
66                    (limits abs={:.3e}, rel={:.3e}, cos>={:.5})",
67                diff.max_abs,
68                diff.max_rel,
69                diff.cosine,
70                tolerance.max_abs,
71                tolerance.max_rel,
72                tolerance.min_cosine
73            ),
74            Self::Nonfinite => write!(f, "non-finite value (NaN or inf) in output"),
75        }
76    }
77}
78
79impl std::error::Error for VerifyError {}
80
81/// Compute the per-pair diff between `out` and `ref`. NaN/inf in
82/// either side is a hard fail.
83pub fn diff(out: &[f32], reference: &[f32]) -> Result<Diff, VerifyError> {
84    if out.len() != reference.len() {
85        return Err(VerifyError::LengthMismatch {
86            got: out.len(),
87            expected: reference.len(),
88        });
89    }
90    let mut max_abs = 0f32;
91    let mut max_rel = 0f32;
92    let mut sum_abs = 0f32;
93    let mut argmax = 0usize;
94    let mut dot = 0f32;
95    let mut na = 0f32;
96    let mut nb = 0f32;
97    for (i, (&a, &b)) in out.iter().zip(reference).enumerate() {
98        if !a.is_finite() || !b.is_finite() {
99            return Err(VerifyError::Nonfinite);
100        }
101        let d = (a - b).abs();
102        sum_abs += d;
103        if d > max_abs {
104            max_abs = d;
105            argmax = i;
106        }
107        let denom = b.abs().max(1e-12);
108        let rel = d / denom;
109        if rel > max_rel {
110            max_rel = rel;
111        }
112        dot += a * b;
113        na += a * a;
114        nb += b * b;
115    }
116    let n = out.len();
117    let cosine = if na > 0.0 && nb > 0.0 {
118        dot / (na.sqrt() * nb.sqrt())
119    } else {
120        1.0
121    };
122    Ok(Diff {
123        n,
124        max_abs,
125        max_rel,
126        mean_abs: if n > 0 { sum_abs / n as f32 } else { 0.0 },
127        cosine,
128        argmax_diff_index: argmax,
129    })
130}
131
132/// `diff` + tolerance check. Use this in tests / parity harnesses.
133pub fn compare(out: &[f32], reference: &[f32], tol: Tolerance) -> Result<Diff, VerifyError> {
134    let d = diff(out, reference)?;
135    if d.max_abs > tol.max_abs || d.max_rel > tol.max_rel || d.cosine < tol.min_cosine {
136        return Err(VerifyError::ToleranceExceeded {
137            diff: d,
138            tolerance: tol,
139        });
140    }
141    Ok(d)
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn identical_passes() {
150        let a = vec![1.0, 2.0, 3.0];
151        let r = compare(&a, &a, Tolerance::STRICT).unwrap();
152        assert_eq!(r.max_abs, 0.0);
153        assert!((r.cosine - 1.0).abs() < 1e-6);
154    }
155
156    #[test]
157    fn tiny_diff_passes_strict() {
158        let a = [1.0f32, 2.0, 3.0];
159        let b = [1.0 + 1e-6, 2.0 - 1e-6, 3.0];
160        compare(&a, &b, Tolerance::STRICT).unwrap();
161    }
162
163    #[test]
164    fn big_diff_fails() {
165        let a = [1.0f32, 2.0, 3.0];
166        let b = [1.0, 2.0, 99.0];
167        let err = compare(&a, &b, Tolerance::STRICT).unwrap_err();
168        assert!(matches!(err, VerifyError::ToleranceExceeded { .. }));
169    }
170
171    #[test]
172    fn nan_is_hard_fail() {
173        let a = [1.0f32, f32::NAN];
174        let b = [1.0, 0.0];
175        let err = compare(&a, &b, Tolerance::STRICT).unwrap_err();
176        assert!(matches!(err, VerifyError::Nonfinite));
177    }
178}
rlx_runtime/logit_verify.rs

rlx_runtime/
logit_verify.rs
