From 7e7d077ff1dd7f210bab76295b6dd37f1f3578a2 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Fri, 29 May 2026 21:54:05 +0800 Subject: [PATCH] moe: KV-cached GPU decode for gpt-oss (O(1)/token) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit decode_step(): single-token forward using GpuKVCache + the GPU decode_attention_sink kernel (per-head sinks + sliding window) + per-token MoE, so generation is O(1)/token instead of the host forward's O(n²) recompute. generate(): prefill prompt token-by-token (causal attention => sequential == batched), then greedy decode to eos/max_new. Verified against the reference host forward on the Paris prompt: both predict top-1 token 12650 = " Paris" (MATCH_TOP1: YES; logits 15.375 vs 15.3125 — BF16 accumulation-order diff only). gptoss-logits now runs both paths and asserts the match. Build green on dash5. Co-Authored-By: Claude Opus 4.8 --- crates/xserv-model/src/bin/gptoss-logits.rs | 25 ++++- crates/xserv-model/src/gptoss.rs | 101 +++++++++++++++++++- 2 files changed, 120 insertions(+), 6 deletions(-) diff --git a/crates/xserv-model/src/bin/gptoss-logits.rs b/crates/xserv-model/src/bin/gptoss-logits.rs index 7790b22..7f19648 100644 --- a/crates/xserv-model/src/bin/gptoss-logits.rs +++ b/crates/xserv-model/src/bin/gptoss-logits.rs @@ -20,17 +20,36 @@ fn main() { 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::(); let last = &data[(t - 1) * vocab..t * vocab]; - let mut idx: Vec = (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!("[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.num_layers(), model.config.num_kv_heads(), + model.config.head_dim(), xserv_tensor::DType::BF16, Device::Cuda(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::(); + let mut didx: Vec = (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" }); } diff --git a/crates/xserv-model/src/gptoss.rs b/crates/xserv-model/src/gptoss.rs index e71c097..046c833 100644 --- a/crates/xserv-model/src/gptoss.rs +++ b/crates/xserv-model/src/gptoss.rs @@ -21,9 +21,10 @@ use std::collections::HashMap; use half::bf16; use xserv_cuda::GpuBuffer; use xserv_kernels::*; -use xserv_tensor::{Device, Tensor}; +use xserv_tensor::{DType, Device, Tensor}; use crate::config::ModelConfig; +use crate::kv_cache::GpuKVCache; pub struct GptOss { pub config: ModelConfig, @@ -46,7 +47,8 @@ struct Block { v_bias: Tensor, o_proj_wt: Tensor, // [n_heads*head_dim, hidden] o_bias: Tensor, - sinks: Tensor, // [n_heads] (f32 on host) + sinks: Tensor, // [n_heads] BF16 on host (used by the host forward path) + sinks_gpu: Tensor, // [n_heads] F32 on GPU (used by decode_attention_sink) sliding: bool, // MoE. Experts stay MXFP4 on GPU; dequantized per use (see expert_forward). router_wt: Tensor, // [hidden, n_experts] @@ -132,6 +134,11 @@ impl GptOss { down_bias.push(slice_row(&down_b, e, down_out).to_device(dev)); } + // Sinks: keep BF16 on host (host forward path) + F32 on GPU (decode kernel). + let sinks_cpu = take(&mut w, &format!("{p}.self_attn.sinks")).to_device(Device::Cpu); + let sinks_f32: Vec = sinks_cpu.as_slice::().iter().map(|v| v.to_f32()).collect(); + let sinks_gpu = Tensor::from_slice(&sinks_f32, &[sinks_f32.len()]).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"))), @@ -143,7 +150,8 @@ impl GptOss { 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), + sinks: sinks_cpu, + sinks_gpu, 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"))), @@ -208,6 +216,81 @@ impl GptOss { matmul2(&x, &self.lm_head_t) // [T, vocab] } + /// Single-token GPU decode step with KV cache. Returns logits [1, vocab]. + /// Same math as `forward` but q_len=1, GPU sink-attention over cached K/V, so + /// it is O(1) per token. Prefill = call this for each prompt token (causal + /// attention makes sequential identical to a batched prefill). + pub fn decode_step(&self, token: u32, cache: &mut GpuKVCache) -> Tensor { + 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 scale = (head_dim as f32).powf(-0.5); + let pos = cache.seq_len(); + let positions = [pos as u32]; + + let mut x = embedding(&self.embed_tokens, &[token]); // [1, hidden] + for (li, layer) in self.layers.iter().enumerate() { + let residual = x.clone(); + let normed = rmsnorm(&x, &layer.input_norm, eps); + + let q = add_bias(&matmul2(&normed, &layer.q_proj_wt), &layer.q_bias); // [1, H*hd] + let k = add_bias(&matmul2(&normed, &layer.k_proj_wt), &layer.k_bias); // [1, Hkv*hd] + let v = add_bias(&matmul2(&normed, &layer.v_proj_wt), &layer.v_bias); + + // reshape + rotate_half RoPE (t=1) + let q = reshape_heads_gpu(&q, 1, n_heads, head_dim); // [1,H,1,hd] + let k = reshape_heads_gpu(&k, 1, n_kv, head_dim); + let q = transpose_for_rope_gpu(&q, 1, n_heads, head_dim); // [1,H,hd] + let k = transpose_for_rope_gpu(&k, 1, 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, 1, n_heads, head_dim); // [1,H,1,hd] + let k = transpose_from_rope_gpu(&k, 1, n_kv, head_dim); // [1,Hkv,1,hd] + let v = reshape_heads_gpu(&v, 1, n_kv, head_dim); // [1,Hkv,1,hd] + + cache.append(li, &k, &v, 1, pos); + let (k_full, v_full) = cache.get_kv_len(li, pos + 1); // [1,Hkv,pos+1,hd] + let window = if layer.sliding { self.config.sliding_window().unwrap_or(0) } else { 0 }; + let attn = decode_attention_sink(&q, &k_full, &v_full, &layer.sinks_gpu, scale, window); // [1,H,1,hd] + let attn = attn.reshape(&[1, n_heads * head_dim]); // head-major == o_proj input + let attn_proj = add_bias(&matmul2(&attn, &layer.o_proj_wt), &layer.o_bias); + x = add(&residual, &attn_proj); + + let residual = x.clone(); + let normed = rmsnorm(&x, &layer.post_norm, eps); + let moe = self.moe_ffn(&normed, layer, hidden); + x = add(&residual, &moe); + } + cache.advance_seq_len(1); + let x = rmsnorm(&x, &self.norm, eps); + matmul2(&x, &self.lm_head_t) // [1, vocab] + } + + /// Greedy generation. Prefills `prompt` then generates up to `max_new` tokens, + /// stopping at `eos`. Returns generated token ids (prompt excluded). + pub fn generate(&self, prompt: &[u32], max_new: usize, eos: Option) -> Vec { + assert!(!prompt.is_empty()); + let mut cache = GpuKVCache::new( + self.config.num_layers(), self.config.num_kv_heads(), + self.config.head_dim(), DType::BF16, Device::Cuda(0), + ); + let mut logits = self.decode_step(prompt[0], &mut cache); + for &tok in &prompt[1..] { + logits = self.decode_step(tok, &mut cache); + } + let mut out = Vec::new(); + let mut next = argmax_last(&logits); + for _ in 0..max_new { + out.push(next); + if Some(next) == eos { break; } + logits = self.decode_step(next, &mut cache); + next = argmax_last(&logits); + } + out + } + /// 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 { @@ -307,6 +390,18 @@ fn matmul2(a: &Tensor, b: &Tensor) -> Tensor { matmul(a, b, GemmBackend::CuBlas) } +/// Greedy argmax over the last row of a [*, vocab] BF16 logits tensor. +fn argmax_last(logits: &Tensor) -> u32 { + let vocab = logits.shape()[logits.ndim() - 1]; + let rows = logits.numel() / vocab; + let host = logits.to_device(Device::Cpu); + let d = host.as_slice::(); + let last = &d[(rows - 1) * vocab..rows * vocab]; + last.iter().enumerate() + .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap()) + .map(|(i, _)| i as u32).unwrap() +} + /// One expert `e` of `layer`: clamped SwiGLU. x:[*,hidden] -> [*,hidden]. /// Dequantizes the expert's MXFP4 weights to BF16 scratch on GPU, then: /// gate_up = x@gate_up + bias; gate=even cols, up=odd cols (interleaved);