gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

moe(wip): gptoss.rs first correctness-first forward + logit-dump bin

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-29 21:05:47 +08:00
parent c7d0750c32
commit 05534611ca
2 changed files with 391 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
crates/xserv-model/src/bin/gptoss-logits.rs Normal file
View File

@@ -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 <bf16-model-dir> <tok0> <tok1> ...
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<String> = std::env::args().collect();
let model_dir = PathBuf::from(&args[1]);
let tokens: Vec<u32> = 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::<bf16>();
let last = &data[(t - 1) * vocab..t * vocab];
let mut idx: Vec<usize> = (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());
}
}

356
crates/xserv-model/src/gptoss.rs Normal file
View File

@@ -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<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.
router_wt: Tensor, // [hidden, n_experts]
router_bias: Tensor, // [n_experts]
gate_up_wt: Vec<Tensor>, // per expert: [hidden, 2*inter]
gate_up_bias: Vec<Tensor>, // [2*inter]
down_wt: Vec<Tensor>, // per expert: [inter, hidden]
down_bias: Vec<Tensor>, // [hidden]
}
impl GptOss {
/// Load gpt-oss from a BF16 (dequantized) HF-format weight map.
pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> 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 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<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.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::<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)
}
/// 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::<bf16>();
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::<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))
}