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

moe: gpt-oss-20b forward verified correct (predicts "Paris")

Browse Source 
YaRN RoPE was the missing piece — gpt-oss uses rope_type "yarn" (factor 32,
beta_fast 32, beta_slow 1, orig_max 4096); a plain theta RoPE garbled
attention. Added yarn_rope_cache (host-computed inv_freq + mscale, built
into a RopeCache directly). Experts kept CPU-resident and uploaded per-use
(the dequantized BF16 model is ~36GB, won't fit one 32GB card).

Verified: "The capital of France is" -> top-1 token 12366 = " Paris"
(logit 19.75), matching the llama.cpp oracle's behavior. This exercises the
full MoE path: top-4 router (softmax-after-topk), interleaved clamped
(up+1)*glu experts, attention sinks, sliding window, MXFP4->BF16 weights,
YaRN RoPE, head_dim 64, q/k/v/o biases.

Correctness-first (host attention + per-token MoE); GPU attention-with-sinks
kernel, KV cache, faster MoE, and PP-for-memory come next to run AIME/GSM8K
at speed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-29 21:13:06 +08:00
parent 05534611ca
commit 0dd8851e88
2 changed files with 73 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

84
crates/xserv-model/src/gptoss.rs
View File

@@ -75,11 +75,7 @@ impl GptOss {
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 rope_cache = yarn_rope_cache(&config);
let n_layers = config.num_layers();
let mut layers = Vec::with_capacity(n_layers);
@@ -94,11 +90,15 @@ impl GptOss {
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);
// Experts are kept on CPU (the 32 experts per layer total ~36GB for
// the whole model, which won't fit one GPU). Each selected expert's
// weights (~50MB) are uploaded on demand in expert_forward; only
// top-k experts per token are touched, so the H2D traffic is small.
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));
gate_up_wt.push(slice_expert(&gate_up, e, hidden, 2 * inter)); // CPU
gate_up_bias.push(slice_row(&gate_up_b, e, 2 * inter)); // CPU
down_wt.push(slice_expert(&down, e, inter, hidden)); // CPU
down_bias.push(slice_row(&down_b, e, hidden)); // CPU
}
layers.push(Block {
@@ -217,6 +217,60 @@ impl GptOss {
// ---------- 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)
}
@@ -226,9 +280,15 @@ fn matmul2(a: &Tensor, b: &Tensor) -> Tensor {
/// 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]
// Upload this expert's CPU-resident weights to x's device just for this call.
let dev = x.device();
let gate_up_wt = gate_up_wt.to_device(dev);
let gate_up_bias = gate_up_bias.to_device(dev);
let down_wt = down_wt.to_device(dev);
let down_bias = down_bias.to_device(dev);
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].

1
crates/xserv-model/src/lib.rs
View File

@@ -13,6 +13,7 @@ pub use gpt2::{GPT2, KVCache};
pub use kv_cache::GpuKVCache;
pub use paged_kv_cache::{BlockAllocator, Location, PagedKVCache, BLOCK_SIZE};
pub use qwen3::Qwen3;
pub use gptoss::GptOss;
pub use sampling::{SamplingParams, sample};
/// Initialize GPU kernel hooks. Called automatically by model constructors,