inferencelayer 0.2.3

//! CPU f32 reference for Qwen3.5 hybrid decode (the arch behind NuExtract-3,
//! `Qwen3_5ForConditionalGeneration`): 24 gated-DeltaNet linear-attention layers + 8 full-attention
//! layers (every 4th), SwiGLU MLP, tied lm_head. This is a pure-CPU decoder — no wgpu/Vulkan — and
//! it is the spec the GPU kernels are checked against, itself validated op-for-op against the
//! transformers gold continuation (see `bin/qwen35-cpu`).
//!
//! Reuses the already-validated gated-delta recurrence [`crate::deltanet::DeltaNetRef::core`] for the
//! linear layers; the full-attention layer (per-head QK-RMSNorm, PARTIAL RoPE on the first
//! `head_dim·partial_rotary_factor` dims, GQA, and the Qwen3.5 attention OUTPUT GATE
//! `o · sigmoid(gate)`) is implemented here. mRoPE collapses to standard RoPE for text-only input
//! (all three position axes equal the token position), which is the decode regime this serves.
//!
//! Numerics that matter for parity: the standard RMSNorms are Gemma-style `x·(1+w)` (weights stored
//! as the delta from 1); only the gated DeltaNet norm uses a plain `w` (handled inside `DeltaNetRef`).

use crate::deltanet::{DeltaNetRef, DeltaNetState};
use crate::weights::bf16_to_f32;
use anyhow::{Context, Result, bail};
use std::collections::HashMap;
use std::path::Path;

const PREFIX: &str = "model.language_model.";

pub struct Qwen35Config {
    pub hidden: usize,
    pub n_layers: usize,
    pub head_dim: usize,
    pub n_heads: usize,
    pub n_kv_heads: usize,
    pub intermediate: usize,
    pub vocab: usize,
    pub eps: f32,
    pub rope_theta: f32,
    pub rotary_dim: usize, // head_dim * partial_rotary_factor
    pub layer_is_full: Vec<bool>,
    // DeltaNet geometry
    pub nk: usize,
    pub nv: usize,
    pub dk: usize,
    pub dv: usize,
    pub conv_kernel: usize,
}

impl Qwen35Config {
    pub fn from_json(bytes: &[u8]) -> Result<Self> {
        let v: serde_json::Value = serde_json::from_slice(bytes)?;
        let t = &v["text_config"];
        let g = |k: &str| -> Result<u64> {
            t[k].as_u64().with_context(|| format!("text_config.{k} missing/not uint"))
        };
        let head_dim = g("head_dim")? as usize;
        let prf = t["partial_rotary_factor"].as_f64().unwrap_or(1.0);
        let layer_types = t["layer_types"]
            .as_array()
            .context("text_config.layer_types missing")?;
        let layer_is_full: Vec<bool> = layer_types
            .iter()
            .map(|x| x.as_str() == Some("full_attention"))
            .collect();
        Ok(Self {
            hidden: g("hidden_size")? as usize,
            n_layers: g("num_hidden_layers")? as usize,
            head_dim,
            n_heads: g("num_attention_heads")? as usize,
            n_kv_heads: g("num_key_value_heads")? as usize,
            intermediate: g("intermediate_size")? as usize,
            vocab: g("vocab_size")? as usize,
            eps: t["rms_norm_eps"].as_f64().unwrap_or(1e-6) as f32,
            rope_theta: t["rope_parameters"]["rope_theta"]
                .as_f64()
                .or_else(|| t["rope_theta"].as_f64())
                .context("rope_theta missing")? as f32,
            rotary_dim: ((head_dim as f64) * prf).round() as usize,
            layer_is_full,
            nk: g("linear_num_key_heads")? as usize,
            nv: g("linear_num_value_heads")? as usize,
            dk: g("linear_key_head_dim")? as usize,
            dv: g("linear_value_head_dim")? as usize,
            conv_kernel: g("linear_conv_kernel_dim")? as usize,
        })
    }
}

/// Multi-threaded `[m,n]·[n] -> [m]` matvec (row-major weight). The projections and lm_head dominate
/// per-token cost; naive scalar gemv would make a full generation take minutes.
fn pgemv(w: &[f32], x: &[f32], m: usize, n: usize) -> Vec<f32> {
    debug_assert_eq!(w.len(), m * n);
    debug_assert_eq!(x.len(), n);
    let mut out = vec![0f32; m];
    let threads = std::thread::available_parallelism()
        .map(|v| v.get())
        .unwrap_or(8);
    if threads <= 1 || m < 512 {
        for (r, o) in out.iter_mut().enumerate() {
            let base = r * n;
            *o = (0..n).map(|j| w[base + j] * x[j]).sum();
        }
        return out;
    }
    let chunk = m.div_ceil(threads);
    std::thread::scope(|s| {
        for (ci, oc) in out.chunks_mut(chunk).enumerate() {
            let (w, x) = (&w, &x);
            s.spawn(move || {
                let row0 = ci * chunk;
                for (k, o) in oc.iter_mut().enumerate() {
                    let base = (row0 + k) * n;
                    let mut acc = 0f32;
                    for j in 0..n {
                        acc += w[base + j] * x[j];
                    }
                    *o = acc;
                }
            });
        }
    });
    out
}

/// Gemma-style RMSNorm: `x·rsqrt(mean(x²)+eps)·(1+w)` (Qwen3.5 stores the norm weight as the delta
/// from 1). Used for input/post/q/k/final norms — NOT the gated DeltaNet norm.
fn rmsnorm_1p(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 * (1.0 + w[i])).collect()
}

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

struct FullAttn {
    q_proj: Vec<f32>, // [2·n_heads·head_dim, hidden] (query ‖ gate, per head)
    k_proj: Vec<f32>, // [n_kv·head_dim, hidden]
    v_proj: Vec<f32>,
    q_norm: Vec<f32>, // [head_dim]
    k_norm: Vec<f32>,
    o_proj: Vec<f32>, // [hidden, n_heads·head_dim]
}

struct Mlp {
    gate: Vec<f32>, // [intermediate, hidden]
    up: Vec<f32>,
    down: Vec<f32>, // [hidden, intermediate]
}

enum Mixer {
    Linear(DeltaNetRef),
    Full(FullAttn),
}

struct Layer {
    input_ln: Vec<f32>,
    post_ln: Vec<f32>,
    mixer: Mixer,
    mlp: Mlp,
}

pub struct Qwen35Ref {
    pub cfg: Qwen35Config,
    embed: Vec<f32>, // [vocab, hidden]
    final_norm: Vec<f32>,
    layers: Vec<Layer>,
    // per-layer decode state
    dn_state: Vec<Option<DeltaNetState>>,
    k_cache: Vec<Vec<f32>>, // full layers: [T·n_kv·head_dim]
    v_cache: Vec<Vec<f32>>,
    t: Vec<usize>,
    /// When set, `forward` records each decoder layer's last hidden into `trace`.
    pub capture_trace: bool,
    pub trace: Vec<Vec<f32>>,
}

impl Qwen35Ref {
    pub fn load(dir: impl AsRef<Path>) -> Result<Self> {
        let dir = dir.as_ref();
        let cfg = Qwen35Config::from_json(&std::fs::read(dir.join("config.json"))?)?;
        let bytes = std::fs::read(dir.join("model.safetensors"))
            .context("model.safetensors (single-file text weights) not found")?;
        let st = safetensors::SafeTensors::deserialize(&bytes)?;
        let mut w: HashMap<String, Vec<f32>> = HashMap::new();
        for (name, view) in st.tensors() {
            if let Some(stripped) = name.strip_prefix(PREFIX) {
                w.insert(stripped.to_string(), bf16_to_f32(view.data()));
            }
        }
        let take = |w: &mut HashMap<String, Vec<f32>>, key: String| -> Result<Vec<f32>> {
            w.remove(&key).with_context(|| format!("missing tensor {PREFIX}{key}"))
        };

        let embed = take(&mut w, "embed_tokens.weight".into())?;
        let final_norm = take(&mut w, "norm.weight".into())?;

        let mut layers = Vec::with_capacity(cfg.n_layers);
        for li in 0..cfg.n_layers {
            let p = format!("layers.{li}");
            let input_ln = take(&mut w, format!("{p}.input_layernorm.weight"))?;
            let post_ln = take(&mut w, format!("{p}.post_attention_layernorm.weight"))?;
            let mlp = Mlp {
                gate: take(&mut w, format!("{p}.mlp.gate_proj.weight"))?,
                up: take(&mut w, format!("{p}.mlp.up_proj.weight"))?,
                down: take(&mut w, format!("{p}.mlp.down_proj.weight"))?,
            };
            let mixer = if cfg.layer_is_full[li] {
                Mixer::Full(FullAttn {
                    q_proj: take(&mut w, format!("{p}.self_attn.q_proj.weight"))?,
                    k_proj: take(&mut w, format!("{p}.self_attn.k_proj.weight"))?,
                    v_proj: take(&mut w, format!("{p}.self_attn.v_proj.weight"))?,
                    q_norm: take(&mut w, format!("{p}.self_attn.q_norm.weight"))?,
                    k_norm: take(&mut w, format!("{p}.self_attn.k_norm.weight"))?,
                    o_proj: take(&mut w, format!("{p}.self_attn.o_proj.weight"))?,
                })
            } else {
                let a = format!("{p}.linear_attn");
                Mixer::Linear(DeltaNetRef {
                    nk: cfg.nk,
                    nv: cfg.nv,
                    dk: cfg.dk,
                    dv: cfg.dv,
                    kernel: cfg.conv_kernel,
                    eps: cfg.eps,
                    w_qkv: take(&mut w, format!("{a}.in_proj_qkv.weight"))?,
                    w_z: take(&mut w, format!("{a}.in_proj_z.weight"))?,
                    w_b: take(&mut w, format!("{a}.in_proj_b.weight"))?,
                    w_a: take(&mut w, format!("{a}.in_proj_a.weight"))?,
                    conv_w: take(&mut w, format!("{a}.conv1d.weight"))?, // [conv_dim,1,K] == [conv_dim,K]
                    a_log: take(&mut w, format!("{a}.A_log"))?,
                    dt_bias: take(&mut w, format!("{a}.dt_bias"))?,
                    norm_w: take(&mut w, format!("{a}.norm.weight"))?,
                    w_out: take(&mut w, format!("{a}.out_proj.weight"))?,
                })
            };
            layers.push(Layer { input_ln, post_ln, mixer, mlp });
        }
        if embed.len() != cfg.vocab * cfg.hidden {
            bail!("embed_tokens size {} != vocab*hidden", embed.len());
        }

        let dn_state = layers
            .iter()
            .map(|l| match &l.mixer {
                Mixer::Linear(dn) => Some(dn.fresh_state()),
                Mixer::Full(_) => None,
            })
            .collect();
        let nl = cfg.n_layers;
        Ok(Self {
            cfg,
            embed,
            final_norm,
            layers,
            dn_state,
            k_cache: vec![Vec::new(); nl],
            v_cache: vec![Vec::new(); nl],
            t: vec![0; nl],
            capture_trace: false,
            trace: Vec::new(),
        })
    }

    pub fn reset(&mut self) {
        for (li, l) in self.layers.iter().enumerate() {
            if let Mixer::Linear(dn) = &l.mixer {
                self.dn_state[li] = Some(dn.fresh_state());
            }
            self.k_cache[li].clear();
            self.v_cache[li].clear();
            self.t[li] = 0;
        }
    }

    /// Partial RoPE `(cos, sin)` of length `rotary_dim` for absolute position `pos`.
    fn rope(&self, pos: usize) -> (Vec<f32>, Vec<f32>) {
        let rd = self.cfg.rotary_dim;
        let half = rd / 2;
        let theta = self.cfg.rope_theta;
        let mut cos = vec![0.0; rd];
        let mut sin = vec![0.0; rd];
        for i in 0..half {
            let freq = (pos as f32) * theta.powf(-2.0 * i as f32 / rd 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 partial RoPE in place to one head `[head_dim]`: rotate the first `rotary_dim` dims
    /// (non-interleaved rotate_half), pass the rest through.
    fn apply_rope(&self, head: &mut [f32], cos: &[f32], sin: &[f32]) {
        let rd = self.cfg.rotary_dim;
        let half = rd / 2;
        let orig: Vec<f32> = head[..rd].to_vec();
        for j in 0..rd {
            let rot = if j < half { -orig[j + half] } else { orig[j - half] };
            head[j] = orig[j] * cos[j] + rot * sin[j];
        }
    }

    fn full_attn_step(&mut self, li: usize, x: &[f32], pos: usize) -> Vec<f32> {
        let cfg = &self.cfg;
        let (h, hd, nh, nkv) = (cfg.hidden, cfg.head_dim, cfg.n_heads, cfg.n_kv_heads);
        let n_rep = nh / nkv;
        let scaling = (hd as f32).powf(-0.5);
        let eps = cfg.eps;
        let fa = match &self.layers[li].mixer {
            Mixer::Full(fa) => fa,
            _ => unreachable!(),
        };
        // q_proj -> [nh, 2·hd] per head = [query(hd) ‖ gate(hd)]
        let qraw = pgemv(&fa.q_proj, x, 2 * nh * hd, h);
        let mut q = vec![0f32; nh * hd];
        let mut gate = vec![0f32; nh * hd];
        for head in 0..nh {
            let src = head * 2 * hd;
            q[head * hd..head * hd + hd].copy_from_slice(&qraw[src..src + hd]);
            gate[head * hd..head * hd + hd].copy_from_slice(&qraw[src + hd..src + 2 * hd]);
        }
        let mut k = pgemv(&fa.k_proj, x, nkv * hd, h);
        let v = pgemv(&fa.v_proj, x, nkv * hd, h);
        let (cos, sin) = self.rope(pos);
        // Per-head QK-RMSNorm (Gemma-style 1+w over head_dim) then partial RoPE.
        for head in 0..nh {
            let s = &mut q[head * hd..head * hd + hd];
            let nrm = rmsnorm_1p(s, &fa.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_1p(s, &fa.k_norm, eps);
            s.copy_from_slice(&nrm);
            self.apply_rope(s, &cos, &sin);
        }
        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;
            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;
            }
            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;
            }
            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];
                }
            }
        }
        // Attention output gate: o · sigmoid(gate).
        for i in 0..nh * hd {
            out[i] *= 1.0 / (1.0 + (-gate[i]).exp());
        }
        pgemv(&fa.o_proj, &out, h, nh * hd)
    }

    fn linear_step(&mut self, li: usize, x: &[f32]) -> Vec<f32> {
        let h = self.cfg.hidden;
        let (nv, dv) = (self.cfg.nv, self.cfg.dv);
        let dn = match &self.layers[li].mixer {
            Mixer::Linear(dn) => dn,
            _ => unreachable!(),
        };
        let conv_dim = dn.conv_dim();
        // Parallel projections (the per-layer cost), then the verified gated-delta recurrence.
        let mixed = pgemv(&dn.w_qkv, x, conv_dim, h);
        let z = pgemv(&dn.w_z, x, nv * dv, h);
        let b = pgemv(&dn.w_b, x, nv, h);
        let a = pgemv(&dn.w_a, x, nv, h);
        let st = self.dn_state[li].as_mut().unwrap();
        let core = dn.core(st, &mixed, &z, &b, &a);
        pgemv(&dn.w_out, &core, h, nv * dv)
    }

    fn mlp(&self, li: usize, x: &[f32]) -> Vec<f32> {
        let (h, im) = (self.cfg.hidden, self.cfg.intermediate);
        let m = &self.layers[li].mlp;
        let mut g = pgemv(&m.gate, x, im, h);
        let u = pgemv(&m.up, x, im, h);
        for j in 0..im {
            g[j] = silu(g[j]) * u[j];
        }
        pgemv(&m.down, &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, eps) = (self.cfg.hidden, self.cfg.eps);
        let mut hid = self.embed[token as usize * h..token as usize * h + h].to_vec();
        if self.capture_trace {
            self.trace.clear();
        }
        for li in 0..self.cfg.n_layers {
            let normed = rmsnorm_1p(&hid, &self.layers[li].input_ln, eps);
            let op_out = if self.cfg.layer_is_full[li] {
                self.full_attn_step(li, &normed, pos)
            } else {
                self.linear_step(li, &normed)
            };
            for i in 0..h {
                hid[i] += op_out[i];
            }
            let ffn_in = rmsnorm_1p(&hid, &self.layers[li].post_ln, eps);
            let mlp_out = self.mlp(li, &ffn_in);
            for i in 0..h {
                hid[i] += mlp_out[i];
            }
            if self.capture_trace {
                self.trace.push(hid.clone());
            }
        }
        let hid = rmsnorm_1p(&hid, &self.final_norm, eps);
        pgemv(&self.embed, &hid, self.cfg.vocab, h) // tied lm_head
    }
}

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
}