inferencelayer 0.2.4

//! CPU f32 reference LFM2 forward: the exact spec the WGSL kernels are verified against, and itself
//! checked against the transformers gold continuation. Mirrors the candle LFM2 backend op-for-op
//! (ShortConv ring-buffer, GQA attention with per-head QK-RMSNorm + RoPE, SwiGLU MLP, tied lm_head).

use crate::weights::{Lfm2Config, bf16_to_f32};
use anyhow::{Context, Result};
use std::collections::HashMap;
use std::path::Path;

/// A CPU-resident LFM2 with the incremental decode caches.
pub struct RefModel {
    pub cfg: Lfm2Config,
    w: HashMap<String, Vec<f32>>,
    conv_state: Vec<Vec<f32>>, // [layer][hidden*conv_l], oldest tap first
    k_cache: Vec<Vec<f32>>,    // [layer][T*n_kv*head_dim]
    v_cache: Vec<Vec<f32>>,
    t: Vec<usize>, // [layer] kv length
}

fn gemv(w: &[f32], x: &[f32], m: usize, n: usize) -> Vec<f32> {
    (0..m)
        .map(|r| {
            let base = r * n;
            (0..n).map(|j| w[base + j] * x[j]).sum()
        })
        .collect()
}

fn rmsnorm(x: &[f32], w: &[f32], eps: f32) -> Vec<f32> {
    let n = x.len();
    let ms = x.iter().map(|v| v * v).sum::<f32>() / n as f32;
    let inv = 1.0 / (ms + eps).sqrt();
    (0..n).map(|i| x[i] * inv * w[i]).collect()
}

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

impl RefModel {
    pub fn load(dir: impl AsRef<Path>) -> Result<Self> {
        let dir = dir.as_ref();
        let cfg = Lfm2Config::from_json(&std::fs::read(dir.join("config.json"))?)?;
        let bytes = std::fs::read(dir.join("model.safetensors"))?;
        let st = safetensors::SafeTensors::deserialize(&bytes)?;
        let mut w = HashMap::new();
        for (name, view) in st.tensors() {
            w.insert(name.to_string(), bf16_to_f32(view.data()));
        }
        let nl = cfg.n_layers;
        Ok(Self {
            conv_state: vec![vec![0.0; cfg.hidden * cfg.conv_l]; nl],
            k_cache: vec![Vec::new(); nl],
            v_cache: vec![Vec::new(); nl],
            t: vec![0; nl],
            cfg,
            w,
        })
    }

    pub fn reset(&mut self) {
        for s in &mut self.conv_state {
            s.iter_mut().for_each(|v| *v = 0.0);
        }
        self.k_cache.iter_mut().for_each(|c| c.clear());
        self.v_cache.iter_mut().for_each(|c| c.clear());
        self.t.iter_mut().for_each(|v| *v = 0);
    }

    fn g(&self, name: &str) -> &[f32] {
        self.w.get(name).expect(name)
    }

    /// RoPE `(cos, sin)` of length `head_dim` for absolute position `pos`.
    fn rope(&self, pos: usize) -> (Vec<f32>, Vec<f32>) {
        let hd = self.cfg.head_dim;
        let half = hd / 2;
        let theta = self.cfg.rope_theta;
        let mut cos = vec![0.0; hd];
        let mut sin = vec![0.0; hd];
        for i in 0..half {
            let freq = (pos as f32) * theta.powf(-2.0 * i as f32 / hd as f32);
            let (s, c) = freq.sin_cos();
            cos[i] = c;
            cos[i + half] = c;
            sin[i] = s;
            sin[i + half] = s;
        }
        (cos, sin)
    }

    /// Apply RoPE in place to a single head `[head_dim]`: `x_rot = x·cos + rotate_half(x)·sin`.
    fn apply_rope(&self, head: &mut [f32], cos: &[f32], sin: &[f32]) {
        let hd = self.cfg.head_dim;
        let half = hd / 2;
        let orig = head.to_vec();
        for j in 0..hd {
            let rot = if j < half {
                -orig[j + half]
            } else {
                orig[j - half]
            };
            head[j] = orig[j] * cos[j] + rot * sin[j];
        }
    }

    fn conv_step(&mut self, li: usize, x: &[f32]) -> Vec<f32> {
        let h = self.cfg.hidden;
        let l = self.cfg.conv_l;
        let p = format!("model.layers.{li}");
        let bcx = gemv(self.g(&format!("{p}.conv.in_proj.weight")), x, 3 * h, h);
        let conv_w = self.g(&format!("{p}.conv.conv.weight")).to_vec(); // [h*l]
        let out_proj = self.g(&format!("{p}.conv.out_proj.weight")).to_vec();
        let state = &mut self.conv_state[li];
        let mut y = vec![0.0; h];
        for c in 0..h {
            let b = bcx[c];
            let cc = bcx[h + c];
            let xv = bcx[2 * h + c];
            let bx = b * xv;
            // Roll ring buffer: drop oldest (tap 0), shift left, append bx at tap l-1.
            let base = c * l;
            for k in 0..l - 1 {
                state[base + k] = state[base + k + 1];
            }
            state[base + l - 1] = bx;
            // Depthwise conv: Σ_k state[c,k]·w[c,k].
            let mut acc = 0.0;
            for k in 0..l {
                acc += state[base + k] * conv_w[base + k];
            }
            y[c] = cc * acc;
        }
        gemv(&out_proj, &y, h, h)
    }

    fn attn_step(&mut self, li: usize, x: &[f32], pos: usize) -> Vec<f32> {
        let h = self.cfg.hidden;
        let (nh, nkv, hd) = (self.cfg.n_heads, self.cfg.n_kv_heads, self.cfg.head_dim);
        let n_rep = nh / nkv;
        let scaling = (hd as f32).powf(-0.5);
        let p = format!("model.layers.{li}");
        let mut q = gemv(
            self.g(&format!("{p}.self_attn.q_proj.weight")),
            x,
            nh * hd,
            h,
        );
        let mut k = gemv(
            self.g(&format!("{p}.self_attn.k_proj.weight")),
            x,
            nkv * hd,
            h,
        );
        let v = gemv(
            self.g(&format!("{p}.self_attn.v_proj.weight")),
            x,
            nkv * hd,
            h,
        );
        let q_norm = self
            .g(&format!("{p}.self_attn.q_layernorm.weight"))
            .to_vec();
        let k_norm = self
            .g(&format!("{p}.self_attn.k_layernorm.weight"))
            .to_vec();
        let (cos, sin) = self.rope(pos);
        let eps = self.cfg.eps;
        // Per-head QK-RMSNorm then RoPE.
        for head in 0..nh {
            let s = &mut q[head * hd..head * hd + hd];
            let nrm = rmsnorm(s, &q_norm, eps);
            s.copy_from_slice(&nrm);
            self.apply_rope(s, &cos, &sin);
        }
        for head in 0..nkv {
            let s = &mut k[head * hd..head * hd + hd];
            let nrm = rmsnorm(s, &k_norm, eps);
            s.copy_from_slice(&nrm);
            self.apply_rope(s, &cos, &sin);
        }
        // Append to KV cache.
        self.k_cache[li].extend_from_slice(&k);
        self.v_cache[li].extend_from_slice(&v);
        self.t[li] += 1;
        let tt = self.t[li];
        let kc = &self.k_cache[li];
        let vc = &self.v_cache[li];
        let mut out = vec![0.0; nh * hd];
        for qh in 0..nh {
            let kv = qh / n_rep;
            // scores over time.
            let mut scores = vec![0.0; tt];
            for (s, sc) in scores.iter_mut().enumerate() {
                let kbase = (s * nkv + kv) * hd;
                let mut dot = 0.0;
                for j in 0..hd {
                    dot += q[qh * hd + j] * kc[kbase + j];
                }
                *sc = dot * scaling;
            }
            // softmax.
            let mx = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            let mut den = 0.0;
            for s in scores.iter_mut() {
                *s = (*s - mx).exp();
                den += *s;
            }
            // weighted sum of v.
            for (s, &sc) in scores.iter().enumerate() {
                let p = sc / den;
                let vbase = (s * nkv + kv) * hd;
                for j in 0..hd {
                    out[qh * hd + j] += p * vc[vbase + j];
                }
            }
        }
        gemv(
            self.g(&format!("{p}.self_attn.out_proj.weight")),
            &out,
            h,
            nh * hd,
        )
    }

    fn mlp(&self, li: usize, x: &[f32]) -> Vec<f32> {
        let h = self.cfg.hidden;
        let im = self.cfg.intermediate;
        let p = format!("model.layers.{li}");
        let mut g = gemv(self.g(&format!("{p}.feed_forward.w1.weight")), x, im, h);
        let u = gemv(self.g(&format!("{p}.feed_forward.w3.weight")), x, im, h);
        for j in 0..im {
            g[j] = silu(g[j]) * u[j];
        }
        gemv(self.g(&format!("{p}.feed_forward.w2.weight")), &g, h, im)
    }

    /// Single-token forward at absolute position `pos`. Returns logits `[vocab]`.
    pub fn forward(&mut self, token: u32, pos: usize) -> Vec<f32> {
        let h = self.cfg.hidden;
        let embed = self.g("model.embed_tokens.weight");
        let mut hid = embed[token as usize * h..token as usize * h + h].to_vec();
        for li in 0..self.cfg.n_layers {
            let p = format!("model.layers.{li}");
            let normed = rmsnorm(
                &hid,
                self.g(&format!("{p}.operator_norm.weight")),
                self.cfg.eps,
            );
            let op_out = if self.cfg.layer_is_attn[li] {
                self.attn_step(li, &normed, pos)
            } else {
                self.conv_step(li, &normed)
            };
            for i in 0..h {
                hid[i] += op_out[i];
            }
            let ffn_in = rmsnorm(&hid, self.g(&format!("{p}.ffn_norm.weight")), self.cfg.eps);
            let mlp_out = self.mlp(li, &ffn_in);
            for i in 0..h {
                hid[i] += mlp_out[i];
            }
        }
        let hid = rmsnorm(&hid, self.g("model.embedding_norm.weight"), self.cfg.eps);
        // Tied lm_head: logits[v] = Σ_i hid[i]·embed[v,i].
        let embed = self.g("model.embed_tokens.weight");
        gemv(embed, &hid, self.cfg.vocab, h)
    }
}

pub fn argmax(v: &[f32]) -> u32 {
    let mut b = 0usize;
    for (i, &x) in v.iter().enumerate() {
        if x > v[b] {
            b = i;
        }
    }
    b as u32
}

/// Load the gold `{prompt_ids, continuation}` produced from transformers (scratchpad/lfm2_gold.json).
pub fn load_gold(path: impl AsRef<Path>) -> Result<(Vec<u32>, Vec<u32>)> {
    let j: serde_json::Value = serde_json::from_reader(std::fs::File::open(path)?)?;
    let ids = |k: &str| -> Result<Vec<u32>> {
        Ok(j[k]
            .as_array()
            .with_context(|| format!("gold missing {k}"))?
            .iter()
            .map(|x| x.as_u64().unwrap() as u32)
            .collect())
    };
    Ok((ids("prompt_ids")?, ids("continuation")?))
}