cortiq-engine 0.8.32

//! CPU forward-pass primitives: RMS-norm, SiLU, sparse SwiGLU FFN.

use crate::pool::Pool;
use cortiq_core::types::NormStyle;

/// SiLU activation function.
#[inline(always)]
/// tanh-approximated GELU (Gemma's GeGLU gate; HF `gelu_pytorch_tanh`).
/// Exact (erf) GELU, `0.5·x·(1 + erf(x/√2))` — HF `hidden_act = "gelu"`.
/// erf by Abramowitz–Stegun 7.1.26 (|error| ≤ 1.5e-7), the same formula the
/// GPU kernels use (WGSL and MSL have no erf).
#[inline]
pub fn gelu_erf(x: f32) -> f32 {
    0.5 * x * (1.0 + erf_f32(x * std::f32::consts::FRAC_1_SQRT_2))
}

/// erf(x) by Abramowitz–Stegun 7.1.26, |error| ≤ 1.5e-7.
#[inline]
pub fn erf_f32(x: f32) -> f32 {
    let a = x.abs();
    let t = 1.0 / (1.0 + 0.327_591_1 * a);
    let y = 1.0
        - (((((1.061_405_4 * t - 1.453_152_1) * t + 1.421_413_8) * t - 0.284_496_74) * t
            + 0.254_829_6)
            * t)
            * (-a * a).exp();
    if x < 0.0 { -y } else { y }
}

pub fn gelu_tanh(x: f32) -> f32 {
    const C: f32 = 0.797_884_6; // √(2/π)
    0.5 * x * (1.0 + (C * (x + 0.044_715 * x * x * x)).tanh())
}

pub fn silu(x: f32) -> f32 {
    x / (1.0 + (-x).exp())
}

/// MiniMax M3's SwiGLU-OAI gated activation.
///
/// The upstream contract is deliberately not ordinary `silu(gate) * up`:
///
/// ```text
/// g = min(gate, limit)                // only an upper clamp
/// u = clamp(up, -limit, limit)
/// out = g * sigmoid(alpha * g) * (u + beta)
/// ```
///
/// MiniMax M3.1 uses `alpha = 1.702`, `beta = 1.0`, and `limit = 7.0`.
/// `None` is useful for source variants which explicitly omit the clamp.
/// Keeping this scalar reference here lets all execution paths share the
/// exact asymmetric clamp rule before an accelerator kernel is admitted.
#[inline]
pub fn swiglu_oai(gate: f32, up: f32, alpha: f32, beta: f32, limit: Option<f32>) -> f32 {
    let (gate, up) = match limit {
        Some(limit) => (gate.min(limit), up.clamp(-limit, limit)),
        None => (gate, up),
    };
    gate * (1.0 / (1.0 + (-alpha * gate).exp())) * (up + beta)
}

/// RMS normalization with explicit weight semantics.
///
/// - `NormStyle::Qwen`  (Llama family): `x̂ · w`
/// - `NormStyle::Gemma`:                `x̂ · (1 + w)`
///
/// Applying the wrong style corrupts every normalization in the
/// forward pass — the style comes from the model arch, never guessed.
pub fn rms_norm(input: &[f32], weight: &[f32], eps: f64, style: NormStyle) -> Vec<f32> {
    let mut out = vec![0.0f32; input.len()];
    rms_norm_into(input, weight, eps, style, &mut out);
    out
}

/// `rms_norm` writing into a caller-owned buffer — the decode hot path
/// calls this twice per layer; the returning variant allocated each time.
pub fn rms_norm_into(input: &[f32], weight: &[f32], eps: f64, style: NormStyle, out: &mut [f32]) {
    let n = input.len();
    debug_assert_eq!(out.len(), n);
    let mean_sq: f64 = input.iter().map(|&x| (x as f64) * (x as f64)).sum::<f64>() / n as f64;
    let inv_rms = 1.0 / (mean_sq + eps).sqrt() as f32;
    match style {
        NormStyle::Qwen => {
            for (o, (&x, &w)) in out.iter_mut().zip(input.iter().zip(weight)) {
                *o = x * inv_rms * w;
            }
        }
        NormStyle::Gemma => {
            for (o, (&x, &w)) in out.iter_mut().zip(input.iter().zip(weight)) {
                *o = x * inv_rms * (1.0 + w);
            }
        }
    }
}

/// Fused residual addition and RMSNorm into a caller-owned buffer:
/// `h += delta`, then `out = rms_norm(h, weight)`.
/// This avoids a separate pass over `h` for the residual addition and sum of squares.
pub fn add_rmsnorm_fused_into(
    h: &mut [f32],
    delta: &[f32],
    weight: &[f32],
    eps: f64,
    style: NormStyle,
    out: &mut [f32],
) {
    let n = h.len();
    debug_assert_eq!(delta.len(), n);
    debug_assert_eq!(out.len(), n);
    let mut mean_sq = 0.0f64;
    for (hv, &dv) in h.iter_mut().zip(delta) {
        let val = *hv + dv;
        *hv = val;
        mean_sq += (val as f64) * (val as f64);
    }
    mean_sq /= n as f64;
    let inv_rms = 1.0 / (mean_sq + eps).sqrt() as f32;
    match style {
        NormStyle::Qwen => {
            for (o, (&x, &w)) in out.iter_mut().zip(h.iter().zip(weight)) {
                *o = x * inv_rms * w;
            }
        }
        NormStyle::Gemma => {
            for (o, (&x, &w)) in out.iter_mut().zip(h.iter().zip(weight)) {
                *o = x * inv_rms * (1.0 + w);
            }
        }
    }
}

/// Sparse SwiGLU FFN: compute only `active_indices` neurons.
///
/// `out = down_projᵀ[·, active] · (silu(gate_proj[active, ·]·h) ⊙ up_proj[active, ·]·h)`
///
/// With a full index list this is bit-identical to the dense path —
/// masking is an execution schedule, not an approximation.
pub fn sparse_ffn_forward(
    hidden_states: &[f32],
    gate_proj_full: &[f32], // [intermediate_size, hidden_size]
    up_proj_full: &[f32],   // [intermediate_size, hidden_size]
    down_proj_full: &[f32], // [hidden_size, intermediate_size]
    hidden_size: usize,
    intermediate_size: usize,
    active_indices: &[u16],
    pool: Option<&Pool>,
) -> Vec<f32> {
    let n_active = active_indices.len();
    if n_active == 0 {
        return vec![0.0; hidden_states.len()];
    }

    let seq_len = hidden_states.len() / hidden_size;
    let mut output = vec![0.0f32; seq_len * hidden_size];

    for s in 0..seq_len {
        let h = &hidden_states[s * hidden_size..(s + 1) * hidden_size];

        // Gather active rows of gate/up, fuse silu(gate)·up. Each active
        // neuron is independent → row-parallel, bit-identical to serial.
        let mut act = vec![0.0f32; n_active];
        let neuron_act = |ai: usize| -> f32 {
            let idx = active_indices[ai];
            let row = idx as usize * hidden_size;
            if row + hidden_size > gate_proj_full.len() {
                return 0.0;
            }
            let mut gate_sum = 0.0f32;
            let mut up_sum = 0.0f32;
            for k in 0..hidden_size {
                gate_sum += gate_proj_full[row + k] * h[k];
                up_sum += up_proj_full[row + k] * h[k];
            }
            silu(gate_sum) * up_sum
        };
        match pool {
            Some(pool) if n_active >= 256 => {
                let act_ptr = SendMut(act.as_mut_ptr());
                pool.run(&move |widx, n| {
                    let chunk = n_active.div_ceil(n);
                    let start = widx * chunk;
                    let end = (start + chunk).min(n_active);
                    for ai in start..end {
                        // SAFETY: disjoint index ranges per worker.
                        unsafe { *act_ptr.at(ai) = neuron_act(ai) };
                    }
                });
            }
            _ => {
                for (ai, dst) in act.iter_mut().enumerate() {
                    *dst = neuron_act(ai);
                }
            }
        }

        // Scatter through the active columns of down_proj.
        let out = &mut output[s * hidden_size..(s + 1) * hidden_size];
        for (ai, &idx) in active_indices.iter().enumerate() {
            let val = act[ai];
            if val.abs() < 1e-12 {
                continue;
            }
            for k in 0..hidden_size {
                out[k] += down_proj_full[k * intermediate_size + idx as usize] * val;
            }
        }
    }

    output
}

/// Fused two-position sparse FFN: gate/up/down weight rows are streamed
/// from memory ONCE for both positions. Bit-identical to two single
/// calls (per-position accumulation order is unchanged).
#[allow(clippy::too_many_arguments)]
pub fn sparse_ffn_forward_pair(
    h1: &[f32],
    h2: &[f32],
    gate_proj_full: &[f32],
    up_proj_full: &[f32],
    down_proj_full: &[f32],
    hidden_size: usize,
    intermediate_size: usize,
    active_indices: &[u16],
    pool: Option<&Pool>,
) -> (Vec<f32>, Vec<f32>) {
    let n_active = active_indices.len();
    let mut out1 = vec![0.0f32; hidden_size];
    let mut out2 = vec![0.0f32; hidden_size];
    if n_active == 0 {
        return (out1, out2);
    }

    let mut act1 = vec![0.0f32; n_active];
    let mut act2 = vec![0.0f32; n_active];
    let neuron_pair = |ai: usize| -> (f32, f32) {
        let idx = active_indices[ai];
        let row = idx as usize * hidden_size;
        if row + hidden_size > gate_proj_full.len() {
            return (0.0, 0.0);
        }
        let (mut g1, mut u1, mut g2, mut u2) = (0.0f32, 0.0f32, 0.0f32, 0.0f32);
        for k in 0..hidden_size {
            let gw = gate_proj_full[row + k];
            let uw = up_proj_full[row + k];
            g1 += gw * h1[k];
            g2 += gw * h2[k];
            u1 += uw * h1[k];
            u2 += uw * h2[k];
        }
        (silu(g1) * u1, silu(g2) * u2)
    };
    match pool {
        Some(pool) if n_active >= 256 => {
            let a1 = SendMut(act1.as_mut_ptr());
            let a2 = SendMut(act2.as_mut_ptr());
            pool.run(&move |widx, n| {
                let chunk = n_active.div_ceil(n);
                let start = widx * chunk;
                let end = (start + chunk).min(n_active);
                for ai in start..end {
                    let (v1, v2) = neuron_pair(ai);
                    // SAFETY: disjoint index ranges per worker.
                    unsafe {
                        *a1.at(ai) = v1;
                        *a2.at(ai) = v2;
                    }
                }
            });
        }
        _ => {
            for ai in 0..n_active {
                let (v1, v2) = neuron_pair(ai);
                act1[ai] = v1;
                act2[ai] = v2;
            }
        }
    }

    // Down scatter: one pass over the active columns for both positions.
    for (ai, &idx) in active_indices.iter().enumerate() {
        let (v1, v2) = (act1[ai], act2[ai]);
        if v1.abs() < 1e-12 && v2.abs() < 1e-12 {
            continue;
        }
        for k in 0..hidden_size {
            let dw = down_proj_full[k * intermediate_size + idx as usize];
            out1[k] += dw * v1;
            out2[k] += dw * v2;
        }
    }
    (out1, out2)
}

#[derive(Clone, Copy)]
struct SendMut(*mut f32);
unsafe impl Send for SendMut {}
unsafe impl Sync for SendMut {}

impl SendMut {
    /// See pool::SendMut::at — captures the Sync wrapper, not the field.
    #[inline]
    fn at(self, i: usize) -> *mut f32 {
        unsafe { self.0.add(i) }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_silu() {
        assert!((silu(0.0) - 0.0).abs() < 1e-6);
        assert!((silu(1.0) - 0.7310586).abs() < 1e-4);
    }

    #[test]
    fn rms_norm_qwen_multiplies_by_w() {
        let input = vec![1.0, 2.0, 3.0, 4.0];
        let weight = vec![1.0; 4]; // identity weight in Qwen semantics
        let out = rms_norm(&input, &weight, 1e-6, NormStyle::Qwen);
        let rms = (30.0_f64 / 4.0).sqrt() as f32;
        assert!((out[0] - 1.0 / rms).abs() < 1e-4);

        // w = 2 doubles the output — x̂·w, not x̂·(1+w).
        let out2 = rms_norm(&input, &[2.0; 4], 1e-6, NormStyle::Qwen);
        assert!((out2[0] - 2.0 / rms).abs() < 1e-4);
    }

    #[test]
    fn rms_norm_gemma_adds_one() {
        let input = vec![1.0, 2.0, 3.0, 4.0];
        let weight = vec![0.0; 4]; // identity weight in Gemma semantics
        let out = rms_norm(&input, &weight, 1e-6, NormStyle::Gemma);
        let rms = (30.0_f64 / 4.0).sqrt() as f32;
        assert!((out[0] - 1.0 / rms).abs() < 1e-4);
    }

    #[test]
    fn minimax_gemma_1p_rms_norm_uses_weight_plus_one() {
        // MiniMax M3 applies Gemma-1p to its residual, Q/K and index-Q/K
        // norms. In particular a stored -1 weight must zero that lane; a
        // usual RMSNorm (`x-hat * weight`) would instead flip its sign.
        let input = [3.0f32, 4.0];
        let weight = [-1.0f32, 0.5];
        let expected_rstd = 1.0 / ((12.5f64 + 1e-6).sqrt() as f32);
        let expected = [0.0, 4.0 * expected_rstd * 1.5];

        let got = rms_norm(&input, &weight, 1e-6, NormStyle::Gemma);
        let mut into = [f32::NAN; 2];
        rms_norm_into(&input, &weight, 1e-6, NormStyle::Gemma, &mut into);

        for ((actual, scratch), want) in got.iter().zip(into).zip(expected) {
            assert!((actual - want).abs() < 1e-6, "{actual} != {want}");
            assert!((scratch - want).abs() < 1e-6, "{scratch} != {want}");
        }
    }

    #[test]
    fn minimax_swiglu_oai_has_asymmetric_gate_clamp() {
        const ALPHA: f32 = 1.702;
        const BETA: f32 = 1.0;
        const LIMIT: f32 = 7.0;

        let sigmoid = |x: f32| 1.0 / (1.0 + (-x).exp());
        let positive = swiglu_oai(10.0, 10.0, ALPHA, BETA, Some(LIMIT));
        let expected_positive = 7.0 * sigmoid(ALPHA * 7.0) * 8.0;
        assert!((positive - expected_positive).abs() < 1e-5);

        // The gate is *not* lower-clamped: only up is symmetric. This
        // catches the tempting but wrong `gate.clamp(-limit, limit)` port.
        let negative = swiglu_oai(-10.0, -10.0, ALPHA, BETA, Some(LIMIT));
        let expected_negative = -10.0 * sigmoid(-10.0 * ALPHA) * -6.0;
        assert!((negative - expected_negative).abs() < 1e-10);
        let symmetric_gate = -7.0 * sigmoid(-7.0 * ALPHA) * -6.0;
        assert!((negative - symmetric_gate).abs() > 1e-6);

        let unclamped = swiglu_oai(10.0, -10.0, ALPHA, BETA, None);
        let expected_unclamped = 10.0 * sigmoid(10.0 * ALPHA) * -9.0;
        assert!((unclamped - expected_unclamped).abs() < 1e-5);
    }

    #[test]
    fn test_sparse_ffn_full_active() {
        let hidden = 4;
        let inter = 8;
        let h = vec![1.0f32; hidden];
        let gate = vec![0.1f32; inter * hidden];
        let up = vec![0.1f32; inter * hidden];
        let down = vec![0.1f32; hidden * inter];
        let active: Vec<u16> = (0..inter as u16).collect();

        let out = sparse_ffn_forward(&h, &gate, &up, &down, hidden, inter, &active, None);
        assert_eq!(out.len(), hidden);
        assert!(out.iter().all(|&v| v.abs() > 1e-6));
    }

    #[test]
    fn test_sparse_ffn_half_active() {
        let hidden = 4;
        let inter = 8;
        let h = vec![1.0f32; hidden];
        let gate = vec![0.1f32; inter * hidden];
        let up = vec![0.1f32; inter * hidden];
        let down = vec![0.1f32; hidden * inter];

        let full: Vec<u16> = (0..inter as u16).collect();
        let half: Vec<u16> = (0..inter as u16 / 2).collect();

        let out_full = sparse_ffn_forward(&h, &gate, &up, &down, hidden, inter, &full, None);
        let out_half = sparse_ffn_forward(&h, &gate, &up, &down, hidden, inter, &half, None);

        let mag_full: f32 = out_full.iter().map(|v| v.abs()).sum();
        let mag_half: f32 = out_half.iter().map(|v| v.abs()).sum();
        assert!(mag_half < mag_full);
        assert!(mag_half > 0.0);
    }
}