atlas_core/kimi_k3/
situ.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! SiTU-GLU CPU reference.
4//!
5//! Official law (K3 tech report Eq. 12), β₁ = `activation_situ_beta` (4),
6//! β₂ = `activation_situ_linear_beta` (25):
7//!
8//! ```text
9//! softcap(x, β) = β * tanh(x / β)
10//! SiTU-GLU(g, u) = (softcap(g, β1) * sigmoid(g)) ⊙ softcap(u, β2)
11//! ```
12//!
13//! Not SwiGLU. A β=0 mutant is the known-bad: it is unbounded SwiGLU and
14//! must diverge.
15
16#[inline]
17pub fn sigmoid(x: f32) -> f32 {
18    1.0 / (1.0 + (-x).exp())
19}
20
21/// SiLU: `x * sigmoid(x)`. Used by KDA short-conv. Decay `f_a` is a plain linear.
22#[inline]
23pub fn silu(x: f32) -> f32 {
24    x * sigmoid(x)
25}
26
27/// Smooth cap. `beta` must be finite and non-zero.
28#[inline]
29pub fn softcap(x: f32, beta: f32) -> f32 {
30    beta * (x / beta).tanh()
31}
32
33/// Gate branch: `softcap(x, beta) * sigmoid(x)`. Bound `|·| < |beta|`.
34#[inline]
35pub fn situ_gate(x: f32, beta: f32) -> f32 {
36    softcap(x, beta) * sigmoid(x)
37}
38
39/// Up branch: `softcap(x, beta_lin)`. Bound `|·| < |beta_lin|`.
40#[inline]
41pub fn situ_up(x: f32, beta_lin: f32) -> f32 {
42    softcap(x, beta_lin)
43}
44
45/// One SiTU-GLU coordinate. Product is bounded by `|β1 * β2|`.
46#[inline]
47pub fn situ_glu(gate: f32, up: f32, beta: f32, beta_lin: f32) -> f32 {
48    situ_gate(gate, beta) * situ_up(up, beta_lin)
49}
50
51/// Elementwise SiTU-GLU over paired gate/up vectors.
52pub fn situ_glu_vec(gate: &[f32], up: &[f32], beta: f32, beta_lin: f32) -> Vec<f32> {
53    assert_eq!(gate.len(), up.len(), "SiTU-GLU gate/up length mismatch");
54    gate.iter()
55        .zip(up)
56        .map(|(g, u)| situ_glu(*g, *u, beta, beta_lin))
57        .collect()
58}
59
60/// SwiGLU mutant used as the known-bad: `silu(g) * u` (no tanh cap).
61#[inline]
62pub fn swiglu_mutant(gate: f32, up: f32) -> f32 {
63    gate * sigmoid(gate) * up
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    const B1: f32 = 4.0;
71    const B2: f32 = 25.0;
72    const TOL: f32 = 1e-6;
73
74    #[test]
75    fn situ_glu_matches_closed_form_vector() {
76        // Hand-evaluated at (g, u) = (4, 25): both branches sit on tanh(1).
77        let tanh1 = 1.0f32.tanh();
78        let want_gate = B1 * tanh1 * sigmoid(4.0);
79        let want_up = B2 * tanh1;
80        let want = want_gate * want_up;
81        let got = situ_glu(4.0, 25.0, B1, B2);
82        assert!((got - want).abs() < TOL, "closed form {want} vs impl {got}");
83
84        let g = [0.0, 4.0, -2.0, 1.5];
85        let u = [0.0, 25.0, 3.0, -8.0];
86        let out = situ_glu_vec(&g, &u, B1, B2);
87        for i in 0..g.len() {
88            let closed = (B1 * (g[i] / B1).tanh() * sigmoid(g[i])) * (B2 * (u[i] / B2).tanh());
89            assert!(
90                (out[i] - closed).abs() < TOL,
91                "idx {i}: {} vs {closed}",
92                out[i]
93            );
94        }
95        // Zero input is exactly zero (gate branch vanishes).
96        assert_eq!(out[0], 0.0);
97        // Bound: |z| < |β1 β2| = 100.
98        assert!(out.iter().all(|z| z.abs() < 100.0));
99    }
100
101    #[test]
102    fn situ_beta_zero_mutant_diverges_from_situ() {
103        // Known-bad: treating β=0 as "no cap" / SwiGLU. Instrument must fail
104        // this comparison before a green SiTU result is trusted.
105        let g = 4.0f32;
106        let u = 25.0f32;
107        let situ = situ_glu(g, u, B1, B2);
108        let mutant = swiglu_mutant(g, u);
109        assert!(
110            (situ - mutant).abs() > 1.0,
111            "SiTU {situ} must diverge from SwiGLU mutant {mutant}"
112        );
113        // Production betas are never zero; a β=0 call is the planted defect.
114        assert_ne!(B1, 0.0);
115        assert_ne!(B2, 0.0);
116    }
117}
atlas_core/kimi_k3/situ.rs

atlas_core/kimi_k3/
situ.rs
