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

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:
2026-05-29 21:05:47 +08:00
parent c7d0750c32
commit 05534611ca
2 changed files with 391 additions and 0 deletions

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());
}
}