needle_core/
kernels.rs

1//! Primitives shared by more than one model generation.
2//!
3//! A function lives here only when every caller uses it with identical
4//! semantics and it bakes in no geometry. Anything conditional on generation
5//! stays with that generation.
6
7extern crate alloc;
8
9use crate::math;
10
11/// Epsilon used by both `_rms_unit` and `_zcrms`, inside the square root.
12pub const EPS: f32 = 1e-6;
13/// Sinkhorn iterations upstream runs.
14pub const SINKHORN_ITERS: usize = 20;
15
16/// `decode._rms_unit` — RMS normalisation with no learned scale.
17///
18/// Distinct from [`crate::norm::zc_rms_norm_vec`], which applies `(1 + γ)`.
19/// The mHC router and the Engram gate both use the unscaled form, in every
20/// generation.
21#[inline]
22pub fn rms_unit(x: &mut [f32]) {
23    let mean_sq = x.iter().map(|v| v * v).sum::<f32>() / x.len() as f32;
24    let inv = 1.0 / math::sqrt(mean_sq + EPS);
25    for v in x.iter_mut() {
26        *v *= inv;
27    }
28}
29
30/// `rms_unit` into a separate destination.
31#[inline]
32pub fn rms_unit_to(x: &[f32], out: &mut [f32]) {
33    debug_assert_eq!(x.len(), out.len());
34    let mean_sq = x.iter().map(|v| v * v).sum::<f32>() / x.len() as f32;
35    let inv = 1.0 / math::sqrt(mean_sq + EPS);
36    for (o, &v) in out.iter_mut().zip(x.iter()) {
37        *o = v * inv;
38    }
39}
40
41/// `architecture._sinkhorn` — alternating log-domain row/column normalisation of
42/// an `n x n` matrix, returning `exp(log_K)`.
43///
44/// Row-major `m[i * n + j]`; `axis=-1` is a row (fixed `i`), `axis=-2` a column
45/// (fixed `j`). Operates in place, leaving the exponentiated result.
46pub fn sinkhorn(m: &mut [f32], n: usize) {
47    debug_assert_eq!(m.len(), n * n);
48    for _ in 0..SINKHORN_ITERS {
49        for i in 0..n {
50            let row = &mut m[i * n..(i + 1) * n];
51            let lse = logsumexp(row);
52            for v in row.iter_mut() {
53                *v -= lse;
54            }
55        }
56        for j in 0..n {
57            let mut max = f32::NEG_INFINITY;
58            for i in 0..n {
59                max = max.max(m[i * n + j]);
60            }
61            let mut sum = 0.0f32;
62            for i in 0..n {
63                sum += math::exp(m[i * n + j] - max);
64            }
65            let lse = max + math::ln(sum);
66            for i in 0..n {
67                m[i * n + j] -= lse;
68            }
69        }
70    }
71    for v in m.iter_mut() {
72        *v = math::exp(*v);
73    }
74}
75
76/// Numerically stable `log(sum(exp(x)))`.
77#[inline]
78fn logsumexp(x: &[f32]) -> f32 {
79    let max = x.iter().copied().fold(f32::NEG_INFINITY, f32::max);
80    if !max.is_finite() {
81        return max;
82    }
83    let sum: f32 = x.iter().map(|&v| math::exp(v - max)).sum();
84    max + math::ln(sum)
85}
needle_core/kernels.rs

needle_core/
kernels.rs
