steeldb/

needle_model.rs

//! Pure-Rust (candle) port of the Cactus **needle-1** Simple Attention Network — the encoder-decoder we
//! CAN faithfully reimplement (its architecture is documented; needle-2's novel MHC/engram/MTP stack is
//! not). This is the model the finetuning harness trains: load the published bf16 safetensors, run
//! teacher-forced, and serve fully in-Rust (no cactus C++ wheel, no Python).
//!
//! Architecture (config.json + `neural/needle_torch.py` reference): d_model 512, 8 heads / 4 KV (GQA),
//! head_dim 64, vocab 8192, RoPE θ=1e4, ZCRMSNorm `(1+scale)·x/rms(x)`. Encoder ×12: pre-norm →
//! GQA self-attn (per-head q/k ZCRMSNorm + RoPE) → sigmoid-gated residual (NO FFN). Decoder ×8: gated
//! self-attn + gated cross-attn (NO FFN). Tied-ish `lm_head` (shipped separately).

use candle_core::{DType, Device, IndexOp, Module, Result, Tensor, D};
use candle_nn::{embedding, ops::softmax, Embedding, Linear, VarBuilder};

pub const D_MODEL: usize = 512;
pub const N_HEADS: usize = 8;
pub const N_KV: usize = 4;
pub const HEAD_DIM: usize = 64;
pub const VOCAB: usize = 8192;
pub const N_ENC: usize = 12;
pub const N_DEC: usize = 8;
const ROPE_THETA: f64 = 10000.0;
const EPS: f64 = 1e-6;

/// Zero-centred RMSNorm: `(1 + scale) · x / rms(x)`, computed in f32.
struct ZCRMSNorm {
    scale: Tensor, // [dim]
}
impl ZCRMSNorm {
    fn load(dim: usize, vb: VarBuilder) -> Result<Self> {
        Ok(Self { scale: vb.get(dim, "weight")? })
    }
    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let dt = x.dtype();
        let x = x.to_dtype(DType::F32)?;
        let rms = (x.sqr()?.mean_keepdim(D::Minus1)? + EPS)?.sqrt()?;
        let scale = (self.scale.to_dtype(DType::F32)? + 1.0)?;
        x.broadcast_div(&rms)?.broadcast_mul(&scale)?.to_dtype(dt)
    }
}

fn no_bias_linear(inp: usize, out: usize, vb: VarBuilder) -> Result<Linear> {
    Ok(Linear::new(vb.get((out, inp), "weight")?, None))
}

/// Precompute RoPE cos/sin: `[T, head_dim/2]`.
fn rope_tables(seq_len: usize, dev: &Device) -> Result<(Tensor, Tensor)> {
    let half = HEAD_DIM / 2;
    let inv: Vec<f32> = (0..half).map(|i| 1f32 / (ROPE_THETA as f32).powf((2 * i) as f32 / HEAD_DIM as f32)).collect();
    let inv = Tensor::from_vec(inv, (1, half), dev)?;
    let t: Vec<f32> = (0..seq_len).map(|i| i as f32).collect();
    let t = Tensor::from_vec(t, (seq_len, 1), dev)?;
    let ang = t.broadcast_mul(&inv)?; // [T, half]
    Ok((ang.cos()?, ang.sin()?))
}

/// Apply split-half RoPE to `x: [B,H,T,head_dim]` with cos/sin `[T, head_dim/2]`.
fn apply_rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result<Tensor> {
    let (_b, _h, t, d) = x.dims4()?;
    let half = d / 2;
    let cos = cos.i((..t, ..))?.reshape((1, 1, t, half))?;
    let sin = sin.i((..t, ..))?.reshape((1, 1, t, half))?;
    let x1 = x.narrow(D::Minus1, 0, half)?;
    let x2 = x.narrow(D::Minus1, half, half)?;
    let r1 = (x1.broadcast_mul(&cos)? - x2.broadcast_mul(&sin)?)?;
    let r2 = (x2.broadcast_mul(&cos)? + x1.broadcast_mul(&sin)?)?;
    Tensor::cat(&[r1, r2], D::Minus1)
}

struct Attention {
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    out_proj: Linear,
    q_norm: ZCRMSNorm,
    k_norm: ZCRMSNorm,
}
impl Attention {
    fn load(vb: VarBuilder) -> Result<Self> {
        Ok(Self {
            q_proj: no_bias_linear(D_MODEL, N_HEADS * HEAD_DIM, vb.pp("q_proj"))?,
            k_proj: no_bias_linear(D_MODEL, N_KV * HEAD_DIM, vb.pp("k_proj"))?,
            v_proj: no_bias_linear(D_MODEL, N_KV * HEAD_DIM, vb.pp("v_proj"))?,
            out_proj: no_bias_linear(D_MODEL, D_MODEL, vb.pp("out_proj"))?,
            q_norm: ZCRMSNorm::load(HEAD_DIM, vb.pp("q_norm"))?,
            k_norm: ZCRMSNorm::load(HEAD_DIM, vb.pp("k_norm"))?,
        })
    }

    /// `causal` = apply a causal mask over the (square) self-attention; `rope` = rotate q,k (self-attn).
    fn forward(&self, q_in: &Tensor, kv_in: &Tensor, causal: bool, rope: Option<&(Tensor, Tensor)>) -> Result<Tensor> {
        let (b, tq, _) = q_in.dims3()?;
        let tk = kv_in.dim(1)?;
        let q = self.q_proj.forward(q_in)?.reshape((b, tq, N_HEADS, HEAD_DIM))?.transpose(1, 2)?;
        let k = self.k_proj.forward(kv_in)?.reshape((b, tk, N_KV, HEAD_DIM))?.transpose(1, 2)?;
        let v = self.v_proj.forward(kv_in)?.reshape((b, tk, N_KV, HEAD_DIM))?.transpose(1, 2)?;
        let q = self.q_norm.forward(&q)?;
        let k = self.k_norm.forward(&k)?;
        let (q, k) = match rope {
            Some((cos, sin)) => (apply_rope(&q, cos, sin)?, apply_rope(&k, cos, sin)?),
            None => (q, k),
        };
        // GQA: repeat KV heads to match query heads.
        let rep = N_HEADS / N_KV;
        let k = repeat_kv(&k, rep)?;
        let v = repeat_kv(&v, rep)?;
        let scale = 1.0 / (HEAD_DIM as f64).sqrt();
        let mut att = (q.contiguous()?.matmul(&k.transpose(2, 3)?.contiguous()?)? * scale)?; // [b,h,tq,tk]
        if causal {
            att = att.broadcast_add(&causal_mask(tq, tk, q_in.device())?)?;
        }
        let att = softmax(&att, D::Minus1)?;
        let out = att.matmul(&v.contiguous()?)?; // [b,h,tq,hd]
        let out = out.transpose(1, 2)?.reshape((b, tq, N_HEADS * HEAD_DIM))?;
        self.out_proj.forward(&out)
    }
}

fn repeat_kv(x: &Tensor, rep: usize) -> Result<Tensor> {
    if rep == 1 {
        return Ok(x.clone());
    }
    let (b, kv, t, d) = x.dims4()?;
    x.unsqueeze(2)?.expand((b, kv, rep, t, d))?.reshape((b, kv * rep, t, d))
}

/// Additive causal mask `[1,1,tq,tk]`: 0 on/below the diagonal, -inf above.
fn causal_mask(tq: usize, tk: usize, dev: &Device) -> Result<Tensor> {
    let mut v = vec![0f32; tq * tk];
    for i in 0..tq {
        for j in 0..tk {
            if j > i {
                v[i * tk + j] = f32::NEG_INFINITY;
            }
        }
    }
    Tensor::from_vec(v, (1, 1, tq, tk), dev)
}

fn gate(vb: &VarBuilder, name: &str) -> Result<Tensor> {
    // sigmoid of a learned scalar (shape [1]); returned as a scalar tensor for broadcasting.
    let g = vb.get(1, name)?;
    candle_nn::ops::sigmoid(&g.to_dtype(DType::F32)?)
}

struct EncoderLayer {
    ln: ZCRMSNorm,
    attn: Attention,
    gate: Tensor,
}
impl EncoderLayer {
    fn load(vb: VarBuilder) -> Result<Self> {
        Ok(Self {
            ln: ZCRMSNorm::load(D_MODEL, vb.pp("input_layernorm"))?,
            attn: Attention::load(vb.pp("self_attn"))?,
            gate: gate(&vb, "attn_gate")?,
        })
    }
    fn forward(&self, x: &Tensor, rope: &(Tensor, Tensor)) -> Result<Tensor> {
        let h = self.ln.forward(x)?;
        let a = self.attn.forward(&h, &h, false, Some(rope))?;
        x + a.broadcast_mul(&self.gate)?
    }
}

struct DecoderLayer {
    ln: ZCRMSNorm,
    self_attn: Attention,
    self_gate: Tensor,
    cross_ln: ZCRMSNorm,
    cross_attn: Attention,
    cross_gate: Tensor,
}
impl DecoderLayer {
    fn load(vb: VarBuilder) -> Result<Self> {
        Ok(Self {
            ln: ZCRMSNorm::load(D_MODEL, vb.pp("input_layernorm"))?,
            self_attn: Attention::load(vb.pp("self_attn"))?,
            self_gate: gate(&vb, "self_attn_gate")?,
            cross_ln: ZCRMSNorm::load(D_MODEL, vb.pp("encoder_attn_layer_norm"))?,
            cross_attn: Attention::load(vb.pp("encoder_attn"))?,
            cross_gate: gate(&vb, "cross_attn_gate")?,
        })
    }
    fn forward(&self, x: &Tensor, enc: &Tensor, rope: &(Tensor, Tensor)) -> Result<Tensor> {
        let h = self.ln.forward(x)?;
        let sa = self.self_attn.forward(&h, &h, true, Some(rope))?;
        let x = (x + sa.broadcast_mul(&self.self_gate)?)?;
        let hd = self.cross_ln.forward(&x)?;
        let ca = self.cross_attn.forward(&hd, enc, false, None)?;
        x + ca.broadcast_mul(&self.cross_gate)?
    }
}

pub struct NeedleModel {
    embed: Embedding,
    enc: Vec<EncoderLayer>,
    enc_final: ZCRMSNorm,
    dec: Vec<DecoderLayer>,
    dec_norm: ZCRMSNorm,
    lm_head: Linear,
    device: Device,
}

impl NeedleModel {
    /// Load from the needle-1 bf16 safetensors (HF naming). Runs in f32 on CPU.
    pub fn load(safetensors: &std::path::Path, device: &Device) -> Result<Self> {
        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[safetensors], DType::F32, device)? };
        let m = vb.pp("model");
        let embed = embedding(VOCAB, D_MODEL, m.pp("embed_tokens"))?;
        let enc = (0..N_ENC).map(|i| EncoderLayer::load(m.pp("encoder").pp("layers").pp(i))).collect::<Result<_>>()?;
        let enc_final = ZCRMSNorm::load(D_MODEL, m.pp("encoder").pp("final_norm"))?;
        let dec = (0..N_DEC).map(|i| DecoderLayer::load(m.pp("decoder").pp("layers").pp(i))).collect::<Result<_>>()?;
        let dec_norm = ZCRMSNorm::load(D_MODEL, m.pp("decoder").pp("norm"))?;
        let lm_head = no_bias_linear(D_MODEL, VOCAB, vb.pp("lm_head"))?;
        Ok(Self { embed, enc, enc_final, dec, dec_norm, lm_head, device: device.clone() })
    }

    /// Encode `enc_ids: [B,Te]` → encoder states `[B,Te,D]` (no padding mask; assumes packed sequences).
    pub fn encode(&self, enc_ids: &Tensor) -> Result<Tensor> {
        let t = enc_ids.dim(1)?;
        let rope = rope_tables(t, &self.device)?;
        let mut x = self.embed.forward(enc_ids)?;
        for l in &self.enc {
            x = l.forward(&x, &rope)?;
        }
        self.enc_final.forward(&x)
    }

    /// Decode `dec_ids: [B,Td]` against `enc_out` → logits `[B,Td,VOCAB]` (causal self-attn).
    pub fn decode(&self, dec_ids: &Tensor, enc_out: &Tensor) -> Result<Tensor> {
        let t = dec_ids.dim(1)?;
        let rope = rope_tables(t, &self.device)?;
        let mut x = self.embed.forward(dec_ids)?;
        for l in &self.dec {
            x = l.forward(&x, enc_out, &rope)?;
        }
        let x = self.dec_norm.forward(&x)?;
        self.lm_head.forward(&x)
    }

    /// Full teacher-forced forward: `(enc_ids, dec_ids) → logits [B,Td,VOCAB]`.
    pub fn forward(&self, enc_ids: &Tensor, dec_ids: &Tensor) -> Result<Tensor> {
        let enc = self.encode(enc_ids)?;
        self.decode(dec_ids, &enc)
    }

    pub fn device(&self) -> &Device {
        &self.device
    }
}