qualia_core_db/solvers/
feed_forward.rs

1//! Feed-forward network (SwiGLU) — the STEM definition of the transformer FFN block as the
2//! composition it is:
3//!
4//! ```text
5//! FFN(x) = W_down · ( SiLU(W_gate · x) ⊙ (W_up · x) )
6//! ```
7//!
8//! That is three matrix–vector products ([`super::linear_algebra::gemm::matvec`]), one
9//! activation ([`super::activation::silu`]), and one Hadamard product
10//! ([`super::linear_algebra::vector::hadamard_assign`]). Nothing proprietary — the LLM "FFN
11//! block" is exactly this. The runtime's `dispatch_ffn_block_pre_norm` is a backend computing
12//! this same function over quantized weights on the GPU.
13//!
14//! Caller-owned, zero internal allocation (gate/up scratch buffers are supplied).
15
16use crate::solvers::activation::silu;
17use crate::solvers::linear_algebra::gemm::{matvec, Transpose};
18use crate::solvers::linear_algebra::vector::hadamard_assign;
19use crate::solvers::SolversError;
20
21/// Compute `out = W_down · ( SiLU(W_gate · x) ⊙ (W_up · x) )`, row-major, caller-owned.
22///
23/// - `x`: input, length `d_model`
24/// - `w_gate`, `w_up`: `d_ff × d_model` each
25/// - `w_down`: `d_model × d_ff`
26/// - `gate_buf`, `up_buf`: scratch, length `d_ff` each (overwritten)
27/// - `out`: result, length `d_model` (overwritten)
28///
29/// Fails closed ([`SolversError::InvalidDimension`]) on any shape mismatch.
30#[allow(clippy::too_many_arguments)]
31pub fn swiglu_ffn(
32    d_model: usize,
33    d_ff: usize,
34    x: &[f64],
35    w_gate: &[f64],
36    w_up: &[f64],
37    w_down: &[f64],
38    gate_buf: &mut [f64],
39    up_buf: &mut [f64],
40    out: &mut [f64],
41) -> Result<(), SolversError> {
42    if x.len() != d_model
43        || w_gate.len() != d_ff * d_model
44        || w_up.len() != d_ff * d_model
45        || w_down.len() != d_model * d_ff
46        || gate_buf.len() != d_ff
47        || up_buf.len() != d_ff
48        || out.len() != d_model
49    {
50        return Err(SolversError::InvalidDimension);
51    }
52    // gate = W_gate · x   ;   up = W_up · x
53    matvec(Transpose::No, d_ff, d_model, w_gate, x, gate_buf)?;
54    matvec(Transpose::No, d_ff, d_model, w_up, x, up_buf)?;
55    // gate = SiLU(gate) ⊙ up
56    silu(gate_buf);
57    hadamard_assign(gate_buf, up_buf)?;
58    // out = W_down · gate
59    matvec(Transpose::No, d_model, d_ff, w_down, gate_buf, out)?;
60    Ok(())
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    fn silu_scalar(z: f64) -> f64 {
68        z / (1.0 + (-z).exp())
69    }
70
71    #[test]
72    fn matches_hand_computed_swiglu() {
73        // d_model = 2, d_ff = 2. Known weights; verify against an independent hand computation.
74        let x = [1.0, 2.0];
75        // W_gate (2×2): rows [1,0],[0,1]  ⇒ gate = [x0, x1] = [1, 2]
76        let w_gate = [1.0, 0.0, 0.0, 1.0];
77        // W_up (2×2): rows [1,1],[1,-1]   ⇒ up = [x0+x1, x0-x1] = [3, -1]
78        let w_up = [1.0, 1.0, 1.0, -1.0];
79        // W_down (2×2): rows [1,0],[0,1]  ⇒ out = h (identity)
80        let w_down = [1.0, 0.0, 0.0, 1.0];
81
82        let mut gate = [0.0; 2];
83        let mut up = [0.0; 2];
84        let mut out = [0.0; 2];
85        swiglu_ffn(
86            2, 2, &x, &w_gate, &w_up, &w_down, &mut gate, &mut up, &mut out,
87        )
88        .unwrap();
89
90        // Expected: h = silu([1,2]) ⊙ [3,-1]; out = h.
91        let h0 = silu_scalar(1.0) * 3.0;
92        let h1 = silu_scalar(2.0) * -1.0;
93        assert!((out[0] - h0).abs() < 1e-12, "out0 = {} != {}", out[0], h0);
94        assert!((out[1] - h1).abs() < 1e-12, "out1 = {} != {}", out[1], h1);
95    }
96
97    #[test]
98    fn zero_input_gives_zero_output() {
99        // SiLU(0)=0 ⇒ gate=0 ⇒ Hadamard 0 ⇒ out 0, for any weights.
100        let x = [0.0, 0.0, 0.0];
101        let w_gate = [1.0; 9];
102        let w_up = [2.0; 9];
103        let w_down = [3.0; 9];
104        let mut gate = [0.0; 3];
105        let mut up = [0.0; 3];
106        let mut out = [0.0; 3];
107        swiglu_ffn(
108            3, 3, &x, &w_gate, &w_up, &w_down, &mut gate, &mut up, &mut out,
109        )
110        .unwrap();
111        for &v in &out {
112            assert!(v.abs() < 1e-12);
113        }
114    }
115
116    #[test]
117    fn rejects_bad_dims() {
118        let x = [1.0, 2.0];
119        let w = [1.0, 0.0, 0.0, 1.0];
120        let mut gate = [0.0; 2];
121        let mut up = [0.0; 2];
122        let mut out = [0.0; 1]; // wrong: should be d_model = 2
123        assert!(matches!(
124            swiglu_ffn(2, 2, &x, &w, &w, &w, &mut gate, &mut up, &mut out),
125            Err(SolversError::InvalidDimension)
126        ));
127    }
128}
qualia_core_db/solvers/feed_forward.rs

qualia_core_db/solvers/
feed_forward.rs
