Gahow Wang
603b6eb270
Root-caused the decode non-determinism: matmul's m==1 custom-GEMV fast path reduces over K with a grid-split atomicAdd, whose float accumulation order is non-deterministic. Negligible for attention's stable pre-transposed weights, but for gpt-oss's wide expert GEMMs (K=2880, N up to 5760) over freshly-dequantized MXFP4 weights it produced visibly different results run-to-run (and a wrong argmax). Added gemm::matmul_dense (plain cublasGemmEx, no GEMV shortcut) and route the expert GEMMs through it. Now decode_step (KV cache + GPU sink-attention + MXFP4 experts) is: - deterministic: 3/3 identical runs - correct: top-1 token 12650 = " Paris" for "The capital of France is", MATCH_TOP1 with the host-attention reference forward - end-to-end: gptoss-gen generates 32 tokens at ~6.85 tok/s on one 5090. Removed the temporary A/B debug dumps. gptoss-logits runs both paths and asserts the top-1 match; gptoss-gen times greedy generation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
53 lines
2.3 KiB
Rust
53 lines
2.3 KiB
Rust
//! 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 <mxfp4-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");
|
|
|
|
xserv_cuda::device::set_device(0).unwrap();
|
|
let config = ModelConfig::from_file(&model_dir.join("config.json"));
|
|
eprintln!("[gptoss-logits] loading {} (MXFP4) ...", model_dir.display());
|
|
let (floats, u8s) = loader::load_model_dir_split(&model_dir, Device::Cpu);
|
|
let model = GptOss::from_weights(config, floats, u8s);
|
|
eprintln!("[gptoss-logits] forward over {} tokens", tokens.len());
|
|
|
|
// (1) batched host-attention forward (reference path).
|
|
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!("[forward] top5 next-token (id: logit):");
|
|
for &i in &idx[..5] {
|
|
println!(" {i}: {:.4}", last[i].to_f32());
|
|
}
|
|
|
|
// (2) KV-cache GPU decode path (token-by-token prefill) — must match top-1.
|
|
let mut cache = xserv_model::GpuKVCache::new(&model.config, 512, xserv_tensor::DType::BF16, 0);
|
|
let mut dlog = model.decode_step(tokens[0], &mut cache);
|
|
for &tok in &tokens[1..] {
|
|
dlog = model.decode_step(tok, &mut cache);
|
|
}
|
|
let dh = dlog.to_device(Device::Cpu);
|
|
let dd = dh.as_slice::<bf16>();
|
|
let mut didx: Vec<usize> = (0..vocab).collect();
|
|
didx.sort_by(|&a, &b| dd[b].to_f32().partial_cmp(&dd[a].to_f32()).unwrap());
|
|
println!("[decode ] top5 next-token (id: logit):");
|
|
for &i in &didx[..5] {
|
|
println!(" {i}: {:.4}", dd[i].to_f32());
|
|
}
|
|
println!("MATCH_TOP1: {}", if idx[0] == didx[0] { "YES" } else { "NO" });
|
|
}
|