diff --git a/crates/xserv-model/src/gptoss.rs b/crates/xserv-model/src/gptoss.rs index 66fb21d..2cd03c7 100644 --- a/crates/xserv-model/src/gptoss.rs +++ b/crates/xserv-model/src/gptoss.rs @@ -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::(); + 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]. diff --git a/crates/xserv-model/src/lib.rs b/crates/xserv-model/src/lib.rs index 9f8df15..b85e8aa 100644 --- a/crates/xserv-model/src/lib.rs +++ b/crates/xserv-model/src/lib.rs @@ -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,