From 05534611ca1e4971a5ec0e82838d505d3a6a3ddf Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Fri, 29 May 2026 21:05:47 +0800 Subject: [PATCH] moe(wip): gptoss.rs first correctness-first forward + logit-dump bin MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GptOss model in xserv's own style (not derived from llama.cpp): BF16 loader for the dequantized weights, naive sink-attention + per-token top-k MoE FFN on host for correctness-first, GPU matmuls via our kernels. Reuses the Qwen3 forward pattern (rotate_half RoPE θ=150000, head_dim 64, no q/k norm) and adds q/k/v/o + expert biases, clamped (up+1)*glu experts, attention sinks, alternating sliding window. gptoss-logits bin dumps next-token logits for fixed token ids to compare with the llama.cpp oracle. WIP: compiles pending fixes; numerical alignment vs llama.cpp is the next step. Then paged-cache + PP wiring + AIME/GSM8K. Co-Authored-By: Claude Opus 4.8 --- crates/xserv-model/src/bin/gptoss-logits.rs | 35 ++ crates/xserv-model/src/gptoss.rs | 356 ++++++++++++++++++++ 2 files changed, 391 insertions(+) create mode 100644 crates/xserv-model/src/bin/gptoss-logits.rs create mode 100644 crates/xserv-model/src/gptoss.rs diff --git a/crates/xserv-model/src/bin/gptoss-logits.rs b/crates/xserv-model/src/bin/gptoss-logits.rs new file mode 100644 index 0000000..f026694 --- /dev/null +++ b/crates/xserv-model/src/bin/gptoss-logits.rs @@ -0,0 +1,35 @@ +//! Dump gpt-oss next-token logits for a fixed token-id sequence, to compare +//! against the llama.cpp oracle (isolates the model forward from tokenizer +//! differences). Usage: gptoss-logits ... +use std::path::PathBuf; +use half::bf16; +use xserv_model::loader; +use xserv_model::{GptOss, ModelConfig}; +use xserv_tensor::Device; + +fn main() { + let args: Vec = std::env::args().collect(); + let model_dir = PathBuf::from(&args[1]); + let tokens: Vec = args[2..].iter().map(|s| s.parse().expect("token id")).collect(); + assert!(!tokens.is_empty(), "need at least one token id"); + + let config = ModelConfig::from_file(&model_dir.join("config.json")); + eprintln!("[gptoss-logits] loading {} ...", model_dir.display()); + let weights = loader::load_model_dir(&model_dir, Device::Cpu); + let model = GptOss::from_weights(config, weights); + eprintln!("[gptoss-logits] forward over {} tokens", tokens.len()); + + let logits = model.forward(&tokens); // [T, vocab] + let vocab = logits.shape()[1]; + let t = logits.shape()[0]; + let host = logits.to_device(Device::Cpu); + let data = host.as_slice::(); + let last = &data[(t - 1) * vocab..t * vocab]; + + let mut idx: Vec = (0..vocab).collect(); + idx.sort_by(|&a, &b| last[b].to_f32().partial_cmp(&last[a].to_f32()).unwrap()); + println!("top10 next-token (id: logit):"); + for &i in &idx[..10] { + println!(" {i}: {:.4}", last[i].to_f32()); + } +} diff --git a/crates/xserv-model/src/gptoss.rs b/crates/xserv-model/src/gptoss.rs new file mode 100644 index 0000000..66fb21d --- /dev/null +++ b/crates/xserv-model/src/gptoss.rs @@ -0,0 +1,356 @@ +//! 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) +//! +//! Weights are loaded from a plain BF16 safetensors dir (MXFP4 experts are +//! dequantized to BF16 offline by tools/gptoss_dequant.py), so the standard +//! loader feeds us BF16 tensors and this file needs no quantization code. +//! +//! 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_kernels::*; +use xserv_tensor::{Device, Tensor}; + +use crate::config::ModelConfig; + +pub struct GptOss { + pub config: ModelConfig, + embed_tokens: Tensor, // [vocab, hidden] + layers: Vec, + 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. + router_wt: Tensor, // [hidden, n_experts] + router_bias: Tensor, // [n_experts] + gate_up_wt: Vec, // per expert: [hidden, 2*inter] + gate_up_bias: Vec, // [2*inter] + down_wt: Vec, // per expert: [inter, hidden] + down_bias: Vec, // [hidden] +} + +impl GptOss { + /// Load gpt-oss from a BF16 (dequantized) HF-format weight map. + pub fn from_weights(config: ModelConfig, mut w: HashMap) -> Self { + crate::init_kernels(); + let dev = Device::Cuda(0); + let take = |w: &mut HashMap, n: &str| -> Tensor { + w.remove(n).unwrap_or_else(|| panic!("missing weight: {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"); + + 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 = RopeCache::new( + config.max_seq_len(), + config.head_dim(), + config.rope_theta.unwrap_or(150000.0) as f32, + ); + + 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}"); + // Experts are stored fused as [E, in, out]; slice per expert into [in, out]. + let gate_up = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj")); // [E, hidden, 2*inter] + let gate_up_b = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj_bias")); // [E, 2*inter] + let down = take(&mut w, &format!("{p}.mlp.experts.down_proj")); // [E, inter, hidden] + let down_b = take(&mut w, &format!("{p}.mlp.experts.down_proj_bias")); // [E, hidden] + let mut gate_up_wt = Vec::with_capacity(n_experts); + let mut gate_up_bias = Vec::with_capacity(n_experts); + let mut down_wt = Vec::with_capacity(n_experts); + let mut down_bias = Vec::with_capacity(n_experts); + for e in 0..n_experts { + gate_up_wt.push(slice_expert(&gate_up, e, hidden, 2 * inter).to_device(dev)); + gate_up_bias.push(slice_row(&gate_up_b, e, 2 * inter).to_device(dev)); + down_wt.push(slice_expert(&down, e, inter, hidden).to_device(dev)); + down_bias.push(slice_row(&down_b, e, hidden).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_wt, gate_up_bias, down_wt, down_bias, + }); + } + + 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 = (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::(); + + // Per-token top-k indices + softmax weights (over the chosen k). + let mut out_rows: Vec = Vec::with_capacity(t); + for ti in 0..t { + let row = &lg[ti * n_experts..(ti + 1) * n_experts]; + let mut idx: Vec = (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 = top.iter().map(|&e| (row[e].to_f32() - maxv).exp()).collect(); + let sum: f32 = exps.iter().sum(); + let weights: Vec = exps.iter().map(|w| w / sum).collect(); + + // x row as [1, hidden]. + let xr = row_view(x, ti); + let mut acc: Option = None; + for (j, &e) in top.iter().enumerate() { + let y = expert_forward(&xr, &layer.gate_up_wt[e], &layer.gate_up_bias[e], + &layer.down_wt[e], &layer.down_bias[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 ---------- + +fn matmul2(a: &Tensor, b: &Tensor) -> Tensor { + matmul(a, b, GemmBackend::CuBlas) +} + +/// One expert: clamped SwiGLU. x:[*,hidden] -> [*,hidden]. +/// gate_up = x@gate_up_wt + 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_wt+bias. +fn expert_forward(x: &Tensor, gate_up_wt: &Tensor, gate_up_bias: &Tensor, + down_wt: &Tensor, down_bias: &Tensor, limit: f32) -> Tensor { + let gate_up = add_bias(&matmul2(x, gate_up_wt), gate_up_bias); // [*, 2*inter] + let h = clamped_swiglu(&gate_up, limit); // [*, inter] + add_bias(&matmul2(&h, down_wt), down_bias) // [*, 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::(); + 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) -> 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::(); + let kh = k.to_device(Device::Cpu); let kd = kh.as_slice::(); + let vh = v.to_device(Device::Cpu); let vd = vh.as_slice::(); + let sh = sinks.to_device(Device::Cpu); let sd = sh.as_slice::(); + 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) +} + +/// Slice expert `e` out of a fused [E, rows, cols] tensor -> [rows, cols]. +fn slice_expert(t: &Tensor, e: usize, rows: usize, cols: usize) -> Tensor { + let host = t.to_device(Device::Cpu); + let s = host.as_slice::(); + let stride = rows * cols; + Tensor::from_slice(&s[e * stride..(e + 1) * stride], &[rows, cols]) +} + +/// 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::(); + 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::(); + 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::(); + let out: Vec = 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::()); + } + Tensor::from_slice(&out, &[rows.len(), n]).to_device(Device::Cuda(0)) +}