xserv/crates/xserv-model/src/gptoss.rs

//! gpt-oss-20b (MoE) forward pass — Phase 19.
//!
//! Correctness-first, in xserv's own style (reuses our kernels; llama.cpp is only
//! a numerical oracle, not a code source). Differences from Qwen3 handled here:
//!   - MoE FFN: per-token top-4 router (softmax after top-k) + clamped-SwiGLU experts
//!   - attention sinks: a per-head learned logit column added to the softmax then
//!     dropped (so attention probabilities do not sum to 1)
//!   - alternating sliding-window attention (window from config on flagged layers)
//!   - q/k/v/o projection biases; head_dim 64; no q/k norm; rotate_half RoPE (θ=150000)
//!
//! Expert weights stay MXFP4-resident on GPU (~10GB for the whole model, fits
//! one 32GB card; BF16 would be ~38GB). Each selected expert is dequantized to a
//! BF16 scratch with our `dequant_mxfp4` kernel right before its GEMM. Dense
//! weights (attn/router/norms/embed/lm_head + expert biases) are BF16.
//!
//! v1 is a self-contained non-paged forward (contiguous KV built per call) used to
//! validate next-token agreement with llama.cpp. Paged-cache + PP + server wiring
//! come after numerical correctness is established.

use std::collections::HashMap;
use half::bf16;
use xserv_cuda::GpuBuffer;
use xserv_kernels::*;
use xserv_tensor::{Device, Tensor};

use crate::config::ModelConfig;

pub struct GptOss {
    pub config: ModelConfig,
    embed_tokens: Tensor, // [vocab, hidden]
    layers: Vec<Block>,
    norm: Tensor,         // [hidden]
    lm_head_t: Tensor,    // [hidden, vocab] (pre-transposed)
    rope_cache: RopeCache,
}

struct Block {
    input_norm: Tensor,   // [hidden]
    post_norm: Tensor,    // [hidden]
    // Attention (weights pre-transposed to [in, out]; biases [out]).
    q_proj_wt: Tensor,    // [hidden, n_heads*head_dim]
    q_bias: Tensor,
    k_proj_wt: Tensor,    // [hidden, n_kv*head_dim]
    k_bias: Tensor,
    v_proj_wt: Tensor,
    v_bias: Tensor,
    o_proj_wt: Tensor,    // [n_heads*head_dim, hidden]
    o_bias: Tensor,
    sinks: Tensor,        // [n_heads] (f32 on host)
    sliding: bool,
    // MoE. Experts stay MXFP4 on GPU; dequantized per use (see expert_forward).
    router_wt: Tensor,    // [hidden, n_experts]
    router_bias: Tensor,  // [n_experts]
    gate_up_blocks: Vec<GpuBuffer>, // per expert: U8 [2*inter, nblk, 16]
    gate_up_scales: Vec<GpuBuffer>, // per expert: U8 [2*inter, nblk]
    gate_up_bias: Vec<Tensor>,      // [2*inter] BF16
    down_blocks: Vec<GpuBuffer>,    // per expert: U8 [hidden, nblk, 16]
    down_scales: Vec<GpuBuffer>,    // per expert: U8 [hidden, nblk]
    down_bias: Vec<Tensor>,         // [hidden] BF16
    gate_up_out: usize,             // 2*inter (dequant out_dim)
    gate_up_nblk: usize,            // hidden/32
    down_out: usize,                // hidden
    down_nblk: usize,               // inter/32
}

impl GptOss {
    /// Load gpt-oss from the original MXFP4 HF dir. `floats` holds the BF16
    /// tensors, `u8s` the MXFP4 expert `_blocks`/`_scales` (both from
    /// `loader::load_model_dir_split`). Experts are sliced per-expert and kept on
    /// GPU as MXFP4; dense weights are uploaded as BF16.
    pub fn from_weights(
        config: ModelConfig,
        mut w: HashMap<String, Tensor>,
        mut u8s: HashMap<String, (Vec<u8>, Vec<usize>)>,
    ) -> Self {
        crate::init_kernels();
        let dev = Device::Cuda(0);
        let take = |w: &mut HashMap<String, Tensor>, n: &str| -> Tensor {
            w.remove(n).unwrap_or_else(|| panic!("missing weight: {n}"))
        };
        let take_u8 = |u: &mut HashMap<String, (Vec<u8>, Vec<usize>)>, n: &str| -> (Vec<u8>, Vec<usize>) {
            u.remove(n).unwrap_or_else(|| panic!("missing MXFP4 tensor: {n}"))
        };
        let repl = |t: Tensor| t.to_device(dev);
        // pre-transpose a [out, in] linear weight to [in, out] for x@wt.
        let wt = |t: Tensor| t.to_device(dev).transpose(0, 1).contiguous();

        let hidden = config.hidden();
        let n_experts = config.num_experts();
        let inter = config.intermediate_size.expect("intermediate_size");
        // dequant dims: gate_up out=2*inter, in=hidden (nblk=hidden/32);
        //               down     out=hidden,  in=inter  (nblk=inter/32).
        let (gate_up_out, gate_up_nblk) = (2 * inter, hidden / 32);
        let (down_out, down_nblk) = (hidden, inter / 32);

        let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight"));
        let norm = repl(take(&mut w, "model.norm.weight"));
        let lm_head_t = wt(take(&mut w, "lm_head.weight"));

        let rope_cache = yarn_rope_cache(&config);

        let n_layers = config.num_layers();
        let mut layers = Vec::with_capacity(n_layers);
        for i in 0..n_layers {
            let p = format!("model.layers.{i}");
            // MXFP4 expert tensors: blocks [E, OUT, nblk, 16], scales [E, OUT, nblk].
            let (gu_blk, _) = take_u8(&mut u8s, &format!("{p}.mlp.experts.gate_up_proj_blocks"));
            let (gu_scl, _) = take_u8(&mut u8s, &format!("{p}.mlp.experts.gate_up_proj_scales"));
            let (dn_blk, _) = take_u8(&mut u8s, &format!("{p}.mlp.experts.down_proj_blocks"));
            let (dn_scl, _) = take_u8(&mut u8s, &format!("{p}.mlp.experts.down_proj_scales"));
            let gate_up_b = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj_bias")); // [E, 2*inter]
            let down_b = take(&mut w, &format!("{p}.mlp.experts.down_proj_bias"));       // [E, hidden]

            // per-expert byte spans
            let gu_blk_pe = gate_up_out * gate_up_nblk * 16;
            let gu_scl_pe = gate_up_out * gate_up_nblk;
            let dn_blk_pe = down_out * down_nblk * 16;
            let dn_scl_pe = down_out * down_nblk;

            let mut gate_up_blocks = Vec::with_capacity(n_experts);
            let mut gate_up_scales = Vec::with_capacity(n_experts);
            let mut down_blocks = Vec::with_capacity(n_experts);
            let mut down_scales = Vec::with_capacity(n_experts);
            let mut gate_up_bias = Vec::with_capacity(n_experts);
            let mut down_bias = Vec::with_capacity(n_experts);
            for e in 0..n_experts {
                gate_up_blocks.push(upload_u8(&gu_blk[e * gu_blk_pe..(e + 1) * gu_blk_pe]));
                gate_up_scales.push(upload_u8(&gu_scl[e * gu_scl_pe..(e + 1) * gu_scl_pe]));
                down_blocks.push(upload_u8(&dn_blk[e * dn_blk_pe..(e + 1) * dn_blk_pe]));
                down_scales.push(upload_u8(&dn_scl[e * dn_scl_pe..(e + 1) * dn_scl_pe]));
                gate_up_bias.push(slice_row(&gate_up_b, e, gate_up_out).to_device(dev));
                down_bias.push(slice_row(&down_b, e, down_out).to_device(dev));
            }

            layers.push(Block {
                input_norm: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))),
                post_norm: repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight"))),
                q_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.q_proj.weight"))),
                q_bias: repl(take(&mut w, &format!("{p}.self_attn.q_proj.bias"))),
                k_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.k_proj.weight"))),
                k_bias: repl(take(&mut w, &format!("{p}.self_attn.k_proj.bias"))),
                v_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.v_proj.weight"))),
                v_bias: repl(take(&mut w, &format!("{p}.self_attn.v_proj.bias"))),
                o_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
                o_bias: repl(take(&mut w, &format!("{p}.self_attn.o_proj.bias"))),
                sinks: take(&mut w, &format!("{p}.self_attn.sinks")).to_device(Device::Cpu),
                sliding: config.layer_uses_sliding_window(i),
                router_wt: wt(take(&mut w, &format!("{p}.mlp.router.weight"))),
                router_bias: repl(take(&mut w, &format!("{p}.mlp.router.bias"))),
                gate_up_blocks, gate_up_scales, gate_up_bias,
                down_blocks, down_scales, down_bias,
                gate_up_out, gate_up_nblk, down_out, down_nblk,
            });
        }

        Self { config, embed_tokens, layers, norm, lm_head_t, rope_cache }
    }

    /// Full prefill forward over `token_ids`; returns logits [seq_len, vocab].
    pub fn forward(&self, token_ids: &[u32]) -> Tensor {
        let t = token_ids.len();
        let hidden = self.config.hidden();
        let n_heads = self.config.num_heads();
        let n_kv = self.config.num_kv_heads();
        let head_dim = self.config.head_dim();
        let eps = self.config.rms_norm_eps.unwrap_or(1e-5) as f32;
        let positions: Vec<u32> = (0..t as u32).collect();

        let mut x = embedding(&self.embed_tokens, token_ids); // [T, hidden]

        for layer in &self.layers {
            let residual = x.clone();
            let normed = rmsnorm(&x, &layer.input_norm, eps);

            // Q/K/V projections + bias.
            let q = add_bias(&matmul2(&normed, &layer.q_proj_wt), &layer.q_bias); // [T, n_heads*hd]
            let k = add_bias(&matmul2(&normed, &layer.k_proj_wt), &layer.k_bias); // [T, n_kv*hd]
            let v = add_bias(&matmul2(&normed, &layer.v_proj_wt), &layer.v_bias);

            // RoPE (rotate_half, same convention xserv uses for Qwen3): reshape to
            // [1,H,T,D] -> [T,H,D] -> rope -> back.
            let q = reshape_heads_gpu(&q, t, n_heads, head_dim);
            let k = reshape_heads_gpu(&k, t, n_kv, head_dim);
            let q = transpose_for_rope_gpu(&q, t, n_heads, head_dim);
            let k = transpose_for_rope_gpu(&k, t, n_kv, head_dim);
            rope_inplace(&q, &self.rope_cache, &positions);
            rope_inplace(&k, &self.rope_cache, &positions);
            let q = transpose_from_rope_gpu(&q, t, n_heads, head_dim); // [1,H,T,D]
            let k = transpose_from_rope_gpu(&k, t, n_kv, head_dim);
            let v = reshape_heads_gpu(&v, t, n_kv, head_dim);          // [1,H_kv,T,D]

            // Naive attention with sinks (CPU softmax for correctness).
            let attn = attention_with_sinks(
                &q, &k, &v, &layer.sinks, n_heads, n_kv, head_dim, t,
                if layer.sliding { self.config.sliding_window() } else { None },
            ); // [T, hidden]
            let attn_proj = add_bias(&matmul2(&attn, &layer.o_proj_wt), &layer.o_bias);
            x = add(&residual, &attn_proj);

            // MoE FFN.
            let residual = x.clone();
            let normed = rmsnorm(&x, &layer.post_norm, eps);
            let moe = self.moe_ffn(&normed, layer, hidden);
            x = add(&residual, &moe);
        }

        let x = rmsnorm(&x, &self.norm, eps);
        matmul2(&x, &self.lm_head_t) // [T, vocab]
    }

    /// MoE FFN over [T, hidden]: router top-k softmax, per-token weighted sum of
    /// its top-k experts' clamped-SwiGLU outputs. Correctness-first (per-token).
    fn moe_ffn(&self, x: &Tensor, layer: &Block, hidden: usize) -> Tensor {
        let t = x.shape()[0];
        let top_k = self.config.experts_per_tok();
        let n_experts = self.config.num_experts();
        let limit = self.config.swiglu_limit();

        // router logits [T, n_experts] on host.
        let logits = add_bias(&matmul2(x, &layer.router_wt), &layer.router_bias);
        let logits_h = logits.to_device(Device::Cpu);
        let lg = logits_h.as_slice::<bf16>();

        // Per-token top-k indices + softmax weights (over the chosen k).
        let mut out_rows: Vec<Tensor> = Vec::with_capacity(t);
        for ti in 0..t {
            let row = &lg[ti * n_experts..(ti + 1) * n_experts];
            let mut idx: Vec<usize> = (0..n_experts).collect();
            idx.sort_by(|&a, &b| row[b].to_f32().partial_cmp(&row[a].to_f32()).unwrap());
            let top = &idx[..top_k];
            let maxv = row[top[0]].to_f32();
            let exps: Vec<f32> = top.iter().map(|&e| (row[e].to_f32() - maxv).exp()).collect();
            let sum: f32 = exps.iter().sum();
            let weights: Vec<f32> = exps.iter().map(|w| w / sum).collect();

            // x row as [1, hidden].
            let xr = row_view(x, ti);
            let mut acc: Option<Tensor> = None;
            for (j, &e) in top.iter().enumerate() {
                let y = expert_forward(&xr, layer, e, limit); // [1, hidden]
                let yw = scale_tensor(&y, weights[j]);
                acc = Some(match acc { Some(a) => add(&a, &yw), None => yw });
            }
            out_rows.push(acc.unwrap_or_else(|| zeros_row(hidden)));
        }
        concat_rows(&out_rows) // [T, hidden]
    }
}

// ---------- helpers ----------

/// Build a YaRN-scaled RoPE cos/sin cache (gpt-oss uses rope_type "yarn").
/// Mirrors HF `_compute_yarn_parameters`: per-dim interpolation/extrapolation
/// ramp between the scaled (theta*factor) and unscaled frequencies, plus a global
/// attention scaling (mscale) folded into cos/sin. Cache layout matches xserv's
/// rope kernel: f32 [max_seq, half_dim], cos[pos*half+i] = cos(pos*invfreq[i])*mscale.
fn yarn_rope_cache(config: &ModelConfig) -> RopeCache {
    use std::f64::consts::PI;
    let head_dim = config.head_dim();
    let half = head_dim / 2;
    let max_seq = config.max_seq_len();
    let base = config.rope_theta.unwrap_or(150000.0);
    // gpt-oss rope_scaling: yarn, factor 32, beta_fast 32, beta_slow 1, orig 4096,
    // truncate false (keep correction range as floats).
    let factor = 32.0f64;
    let (beta_fast, beta_slow) = (32.0f64, 1.0f64);
    let orig_max = 4096.0f64;
    let dim = head_dim as f64;

    let find_dim = |num_rot: f64| (dim * (orig_max / (num_rot * 2.0 * PI)).ln()) / (2.0 * base.ln());
    let low = find_dim(beta_fast).max(0.0);
    let high = find_dim(beta_slow).min(dim - 1.0);
    let denom = (high - low).max(1e-3);

    let mut inv_freq = vec![0f64; half];
    for i in 0..half {
        let pos_freq = base.powf((2 * i) as f64 / dim);
        let extrap = 1.0 / pos_freq;                 // unscaled (extrapolation)
        let interp = 1.0 / (factor * pos_freq);      // scaled   (interpolation)
        let ramp = ((i as f64 - low) / denom).clamp(0.0, 1.0);
        let mask = 1.0 - ramp;                       // extrapolation factor
        inv_freq[i] = interp * (1.0 - mask) + extrap * mask;
    }
    // mscale: 0.1*ln(factor)+1 for factor>1.
    let mscale = (0.1 * factor.ln() + 1.0) as f64;

    let mut cos = vec![0f32; max_seq * half];
    let mut sin = vec![0f32; max_seq * half];
    for p in 0..max_seq {
        for i in 0..half {
            let ang = p as f64 * inv_freq[i];
            cos[p * half + i] = (ang.cos() * mscale) as f32;
            sin[p * half + i] = (ang.sin() * mscale) as f32;
        }
    }
    let bytes = max_seq * half * std::mem::size_of::<f32>();
    let mut cos_buf = xserv_cuda::GpuBuffer::alloc(bytes).expect("alloc yarn cos");
    let mut sin_buf = xserv_cuda::GpuBuffer::alloc(bytes).expect("alloc yarn sin");
    let cb = unsafe { std::slice::from_raw_parts(cos.as_ptr() as *const u8, bytes) };
    let sb = unsafe { std::slice::from_raw_parts(sin.as_ptr() as *const u8, bytes) };
    cos_buf.copy_from_host(cb).unwrap();
    sin_buf.copy_from_host(sb).unwrap();
    RopeCache { cos: cos_buf, sin: sin_buf, max_seq_len: max_seq, half_dim: half }
}

fn matmul2(a: &Tensor, b: &Tensor) -> Tensor {
    matmul(a, b, GemmBackend::CuBlas)
}

/// One expert `e` of `layer`: clamped SwiGLU. x:[*,hidden] -> [*,hidden].
/// Dequantizes the expert's MXFP4 weights to BF16 scratch on GPU, then:
/// gate_up = x@gate_up + bias; gate=even cols, up=odd cols (interleaved);
/// gate.clamp(max=limit); up.clamp(-limit,limit); h=(up+1)*gate*sigmoid(gate*1.702); h@down+bias.
fn expert_forward(x: &Tensor, layer: &Block, e: usize, limit: f32) -> Tensor {
    let gate_up_w = dequant_mxfp4(&layer.gate_up_blocks[e], &layer.gate_up_scales[e],
                                  layer.gate_up_out, layer.gate_up_nblk, 0); // [hidden, 2*inter]
    let gate_up = add_bias(&matmul2(x, &gate_up_w), &layer.gate_up_bias[e]); // [*, 2*inter]
    let h = clamped_swiglu(&gate_up, limit);                                 // [*, inter]
    let down_w = dequant_mxfp4(&layer.down_blocks[e], &layer.down_scales[e],
                               layer.down_out, layer.down_nblk, 0);          // [inter, hidden]
    add_bias(&matmul2(&h, &down_w), &layer.down_bias[e])                     // [*, hidden]
}

/// Clamped interleaved SwiGLU on host (correctness-first). [*, 2I] -> [*, I].
fn clamped_swiglu(gate_up: &Tensor, limit: f32) -> Tensor {
    const ALPHA: f32 = 1.702;
    let rows = gate_up.shape()[0];
    let two_i = gate_up.shape()[1];
    let inter = two_i / 2;
    let h = gate_up.to_device(Device::Cpu);
    let s = h.as_slice::<bf16>();
    let mut out = vec![bf16::ZERO; rows * inter];
    for r in 0..rows {
        for i in 0..inter {
            let g = s[r * two_i + 2 * i].to_f32();
            let u = s[r * two_i + 2 * i + 1].to_f32();
            let g = g.min(limit);
            let u = u.clamp(-limit, limit);
            let glu = g * (1.0 / (1.0 + (-(g * ALPHA)).exp()));
            out[r * inter + i] = bf16::from_f32((u + 1.0) * glu);
        }
    }
    Tensor::from_slice(&out, &[rows, inter]).to_device(gate_up.device())
}

/// Naive multi-head attention with per-head sink logits, on host (correctness-first).
/// q:[1,n_heads,T,D] k,v:[1,n_kv,T,D] sinks:[n_heads] (host). Returns [T, n_heads*D].
#[allow(clippy::too_many_arguments)]
fn attention_with_sinks(q: &Tensor, k: &Tensor, v: &Tensor, sinks: &Tensor,
                        n_heads: usize, n_kv: usize, head_dim: usize, t: usize,
                        window: Option<usize>) -> Tensor {
    let scale = (head_dim as f32).powf(-0.5);
    let n_rep = n_heads / n_kv;
    let qh = q.to_device(Device::Cpu); let qd = qh.as_slice::<bf16>();
    let kh = k.to_device(Device::Cpu); let kd = kh.as_slice::<bf16>();
    let vh = v.to_device(Device::Cpu); let vd = vh.as_slice::<bf16>();
    let sh = sinks.to_device(Device::Cpu); let sd = sh.as_slice::<bf16>();
    let hidden = n_heads * head_dim;
    let mut out = vec![bf16::ZERO; t * hidden];
    // index helpers: layout [H, T, D] within each (head) block.
    let qi = |h: usize, i: usize, d: usize| (h * t + i) * head_dim + d;
    let kvi = |h: usize, j: usize, d: usize| (h * t + j) * head_dim + d;
    for h in 0..n_heads {
        let kv = h / n_rep;
        for i in 0..t {
            // scores over valid keys j<=i (causal), and j>i-window (sliding).
            let lo = match window { Some(wn) if i + 1 > wn => i + 1 - wn, _ => 0 };
            let mut scores = vec![0f32; i - lo + 1];
            let mut maxv = sd[h].to_f32(); // sink participates in the max
            for j in lo..=i {
                let mut dot = 0f32;
                for d in 0..head_dim {
                    dot += qd[qi(h, i, d)].to_f32() * kd[kvi(kv, j, d)].to_f32();
                }
                let s = dot * scale;
                scores[j - lo] = s;
                if s > maxv { maxv = s; }
            }
            let mut denom = (sd[h].to_f32() - maxv).exp(); // sink column
            for s in &scores { denom += (*s - maxv).exp(); }
            // weighted sum of v (sink contributes no value -> just inflates denom).
            for d in 0..head_dim {
                let mut acc = 0f32;
                for j in lo..=i {
                    let p = (scores[j - lo] - maxv).exp() / denom;
                    acc += p * vd[kvi(kv, j, d)].to_f32();
                }
                out[i * hidden + h * head_dim + d] = bf16::from_f32(acc);
            }
        }
    }
    Tensor::from_slice(&out, &[t, hidden]).to_device(q.device())
}

/// Row-broadcast bias add: x:[T,N] + bias:[N] -> [T,N], via ones[T,1]@bias[1,N].
fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
    let t = x.shape()[0];
    let n = x.shape()[1];
    let ones = Tensor::from_slice(&vec![bf16::from_f32(1.0); t], &[t, 1]).to_device(x.device());
    let bias_row = bias.reshape(&[1, n]);
    let broadcast = matmul2(&ones, &bias_row); // [T, N]
    add(x, &broadcast)
}

/// Upload raw U8 bytes (an MXFP4 expert slice) to a GPU buffer.
fn upload_u8(bytes: &[u8]) -> GpuBuffer {
    let mut buf = GpuBuffer::alloc(bytes.len()).expect("alloc expert U8");
    buf.copy_from_host(bytes).expect("upload expert U8");
    buf
}

/// Slice row `e` out of [E, n] -> [n].
fn slice_row(t: &Tensor, e: usize, n: usize) -> Tensor {
    let host = t.to_device(Device::Cpu);
    let s = host.as_slice::<bf16>();
    Tensor::from_slice(&s[e * n..(e + 1) * n], &[n])
}

fn row_view(t: &Tensor, row: usize) -> Tensor {
    let cols = t.shape()[1];
    let host = t.to_device(Device::Cpu);
    let s = host.as_slice::<bf16>();
    Tensor::from_slice(&s[row * cols..(row + 1) * cols], &[1, cols]).to_device(t.device())
}

fn scale_tensor(t: &Tensor, s: f32) -> Tensor {
    let host = t.to_device(Device::Cpu);
    let data = host.as_slice::<bf16>();
    let out: Vec<bf16> = data.iter().map(|v| bf16::from_f32(v.to_f32() * s)).collect();
    Tensor::from_slice(&out, t.shape()).to_device(t.device())
}

fn zeros_row(n: usize) -> Tensor {
    Tensor::from_slice(&vec![bf16::ZERO; n], &[1, n]).to_device(Device::Cuda(0))
}

fn concat_rows(rows: &[Tensor]) -> Tensor {
    let n = rows[0].shape()[1];
    let mut out = Vec::with_capacity(rows.len() * n);
    for r in rows {
        let h = r.to_device(Device::Cpu);
        out.extend_from_slice(h.as_slice::<bf16>());
    }
    Tensor::from_slice(&out, &[rows.len(), n]).to_device(Device::Cuda(0))
}