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

model: fused GPU MoE kernel — eliminate CPU roundtrip

Browse Source 
Replace the per-token CPU-routed MoE forward with an all-GPU path:

  1. moe_topk_softmax: GPU top-k + softmax (was CPU sort + softmax)
  2. moe_replicate: broadcast input to all local experts
  3. cublasGemmStridedBatchedEx: batched expert matmul (was per-expert cuBLAS)
  4. moe_weighted_sum: FP32-accumulated weighted sum on GPU (was GPU→CPU→F32→BF16→GPU)

Expert weights stored as contiguous 3D tensors for strided batched GEMM.
Zero CPU↔GPU transfers per MoE layer (was ~40 per token per layer).

Also: configurable geglu_alpha, LayerNorm bias auto-detect, unused-weight
diagnostic at load time.

GSM8K 30-problem: 11/30 → 23/30 (76.7%) vs llama.cpp 30/30 (100%).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-31 13:22:59 +08:00
parent 377a04b81f
commit 4368e79695
6 changed files with 617 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
crates/xserv-model/src/config.rs
View File

@@ -73,6 +73,10 @@ pub struct ModelConfig {
pub rope_scaling: Option<RopeScaling>,
#[serde(default)]
pub swiglu_limit: Option<f64>,
#[serde(default)]
pub geglu_alpha: Option<f64>,
#[serde(default)]
pub hidden_act: Option<String>,
}
impl ModelConfig {
@@ -144,4 +148,8 @@ impl ModelConfig {
pub fn window_size(&self) -> usize {
self.sliding_window.unwrap_or(0)
}
pub fn geglu_alpha(&self) -> f32 {
self.geglu_alpha.unwrap_or(1.702) as f32
}
}

247
crates/xserv-model/src/gpt_oss.rs
View File

@@ -12,15 +12,18 @@ pub struct GptOss {
embed_tokens: Tensor,
layers: Vec<GptOssBlock>,
norm: Tensor,
norm_bias: Option<Tensor>,
lm_head_t: Tensor,
rope_cache: RopeCache,
tp: Option<std::sync::Arc<xserv_distributed::TpContext>>,
local_num_heads: usize,
local_num_kv_heads: usize,
has_norm_bias: bool,
}
struct GptOssBlock {
input_norm: Tensor,
input_norm_bias: Option<Tensor>,
// Attention (with bias)
q_proj_wt: Tensor,
q_proj_bias: Tensor,
@@ -36,12 +39,15 @@ struct GptOssBlock {
window_size: usize,
// MoE MLP
post_norm: Tensor,
post_norm_bias: Option<Tensor>,
router_wt: Tensor,
router_bias: Tensor,
expert_gate_up_wt: Vec<Tensor>,
expert_gate_up_bias: Vec<Tensor>,
expert_down_wt: Vec<Tensor>,
expert_down_bias: Vec<Tensor>,
// 3D expert weights for batched GEMM (contiguous on GPU)
expert_gate_up_wt: Tensor, // [local_experts, hidden, 2*inter]
expert_gate_up_bias: Tensor, // [local_experts, 2*inter]
expert_down_wt: Tensor, // [local_experts, inter, hidden]
expert_down_bias: Tensor, // [local_experts, hidden]
local_experts: usize,
// Activation params
glu_alpha: f32,
glu_limit: f32,
@@ -80,6 +86,7 @@ impl GptOss {
let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight"));
let norm = repl(take(&mut w, "model.norm.weight"));
let norm_bias = w.remove("model.norm.bias").map(|t| repl(t));
let lm_head_t = repl(take(&mut w, "lm_head.weight")).transpose(0, 1).contiguous();
let head_dim = config.head_dim();
@@ -106,13 +113,13 @@ impl GptOss {
let num_layers = config.num_layers();
let num_experts = config.num_experts();
let glu_alpha = 1.702f32;
let glu_alpha = config.geglu_alpha();
let glu_limit = config.swiglu_limit.unwrap_or(7.0) as f32;
let mut layers = Vec::with_capacity(num_layers);
if rank == 0 {
eprintln!(
"Loading gpt-oss weights: {} layers, {} experts, world={world}...",
"Loading gpt-oss weights: {} layers, {} experts, world={world}, glu_alpha={glu_alpha}...",
num_layers, num_experts
);
}
@@ -152,34 +159,26 @@ impl GptOss {
let local_experts = num_experts / world;
let expert_start = rank * local_experts;
let mut expert_gate_up_wt = Vec::with_capacity(local_experts);
let mut expert_gate_up_bias = Vec::with_capacity(local_experts);
let mut expert_down_wt = Vec::with_capacity(local_experts);
let mut expert_down_bias = Vec::with_capacity(local_experts);
let inter2 = gate_up_3d.shape()[2]; // 2 * intermediate_size
let hidden = gate_up_3d.shape()[1];
let inter = down_3d.shape()[1]; // intermediate_size
for local_e in 0..local_experts {
let e = expert_start + local_e;
let gu_slice = slice_expert_3d(&gate_up_3d, e, hidden, inter2);
expert_gate_up_wt.push(gu_slice.to_device(dev));
let gu_bias = slice_expert_2d(&gate_up_bias_2d, e, inter2);
expert_gate_up_bias.push(gu_bias.to_device(dev));
let d_slice = slice_expert_3d(&down_3d, e, inter, hidden);
expert_down_wt.push(d_slice.to_device(dev));
let d_bias = slice_expert_2d(&down_bias_2d, e, hidden);
expert_down_bias.push(d_bias.to_device(dev));
}
// Slice the rank's range of experts as contiguous 3D tensors on GPU
let expert_gate_up_wt = slice_expert_range_3d(&gate_up_3d, expert_start, local_experts, hidden, inter2).to_device(dev);
let expert_gate_up_bias = slice_expert_range_2d(&gate_up_bias_2d, expert_start, local_experts, inter2).to_device(dev);
let expert_down_wt = slice_expert_range_3d(&down_3d, expert_start, local_experts, inter, hidden).to_device(dev);
let expert_down_bias = slice_expert_range_2d(&down_bias_2d, expert_start, local_experts, hidden).to_device(dev);
xserv_cuda::allocator::cached_trim();
let input_norm = repl(take(&mut w, &format!("{p}.input_layernorm.weight")));
let input_norm_bias = w.remove(&format!("{p}.input_layernorm.bias")).map(|t| repl(t));
let post_norm = repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight")));
let post_norm_bias = w.remove(&format!("{p}.post_attention_layernorm.bias")).map(|t| repl(t));
layers.push(GptOssBlock {
input_norm: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))),
input_norm,
input_norm_bias,
q_proj_wt,
q_proj_bias,
k_proj_wt,
@@ -191,13 +190,15 @@ impl GptOss {
sinks,
is_sliding,
window_size,
post_norm: repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight"))),
post_norm,
post_norm_bias,
router_wt,
router_bias,
expert_gate_up_wt,
expert_gate_up_bias,
expert_down_wt,
expert_down_bias,
local_experts,
glu_alpha,
glu_limit,
});
@@ -206,16 +207,35 @@ impl GptOss {
let local_num_heads = config.num_heads() / world;
let local_num_kv_heads = config.num_kv_heads() / world;
let has_norm_bias = norm_bias.is_some();
if rank == 0 {
if has_norm_bias {
eprintln!("gpt-oss: detected LayerNorm bias — using LayerNorm instead of RMSNorm");
}
}
// Warn about unused weights that the model didn't consume
if rank == 0 && !w.is_empty() {
eprintln!("WARNING: {} unused weight(s) in model:", w.len());
let mut keys: Vec<_> = w.keys().collect();
keys.sort();
for k in &keys {
eprintln!(" {k}");
}
}
Self {
config,
embed_tokens,
layers,
norm,
norm_bias,
lm_head_t,
rope_cache,
tp,
local_num_heads,
local_num_kv_heads,
has_norm_bias,
}
}
@@ -229,6 +249,34 @@ impl GptOss {
}
}
#[inline]
fn norm(x: &Tensor, weight: &Tensor, bias: &Option<Tensor>, eps: f32) -> Tensor {
match bias {
Some(b) => layernorm(x, weight, b, eps),
None => rmsnorm(x, weight, eps),
}
}
#[inline]
fn add_norm(x: &Tensor, residual: &Tensor, weight: &Tensor, bias: &Option<Tensor>, eps: f32) -> (Tensor, Tensor) {
match bias {
Some(b) => {
let sum = xserv_kernels::add(x, residual);
let normed = layernorm(&sum, weight, b, eps);
(normed, sum)
}
None => add_rmsnorm(x, residual, weight, eps),
}
}
fn norm_eps(&self) -> f32 {
if self.has_norm_bias {
self.config.ln_eps()
} else {
self.config.rms_norm_eps.unwrap_or(1e-5) as f32
}
}
/// Paged decode: process one token per sequence using paged KV cache.
pub fn forward_decode_paged(
&self,
@@ -245,7 +293,7 @@ impl GptOss {
let num_heads = self.local_num_heads;
let num_kv_heads = self.local_num_kv_heads;
let head_dim = self.config.head_dim();
let eps = self.config.rms_norm_eps.unwrap_or(1e-5) as f32;
let eps = self.norm_eps();
let kv_lens: Vec<i32> = positions.iter().map(|&p| (p + 1) as i32).collect();
for (b, &slot) in seq_slots.iter().enumerate() {
@@ -263,7 +311,7 @@ impl GptOss {
for (layer_idx, layer) in self.layers.iter().enumerate() {
let residual = x.clone();
let normed = rmsnorm(&x, &layer.input_norm, eps);
let normed = Self::norm(&x, &layer.input_norm, &layer.input_norm_bias, eps);
// Q/K/V projections with bias
let q_all = add_bias(&matmul_2d(&normed, &layer.q_proj_wt), &layer.q_proj_bias);
@@ -304,13 +352,11 @@ impl GptOss {
// Residual + post-norm
let (normed, x_new) = add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let (normed, x_new) = Self::add_norm(&attn_proj, &residual, &layer.post_norm, &layer.post_norm_bias, eps);
let residual = x_new;
let normed = normed.contiguous();
// MoE MLP
let moe_out = self.moe_forward(&normed, layer, batch);
x = xserv_kernels::add(&residual, &moe_out);
@@ -322,7 +368,7 @@ impl GptOss {
}
unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); }
let x = rmsnorm(&x, &self.norm, eps);
let x = Self::norm(&x, &self.norm, &self.norm_bias, eps);
unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); }
let logits = matmul_2d(&x, &self.lm_head_t);
unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); }
@@ -341,7 +387,7 @@ impl GptOss {
let num_heads = self.local_num_heads;
let num_kv_heads = self.local_num_kv_heads;
let head_dim = self.config.head_dim();
let eps = self.config.rms_norm_eps.unwrap_or(1e-5) as f32;
let eps = self.norm_eps();
paged_cache.ensure_capacity(slot, pos_offset + new_tokens);
paged_cache.advance_seq_len(slot, new_tokens);
@@ -351,7 +397,7 @@ impl GptOss {
for (layer_idx, layer) in self.layers.iter().enumerate() {
let residual = x.clone();
let normed = rmsnorm(&x, &layer.input_norm, eps);
let normed = Self::norm(&x, &layer.input_norm, &layer.input_norm_bias, eps);
let q = add_bias(&matmul_2d(&normed, &layer.q_proj_wt), &layer.q_proj_bias);
let k = add_bias(&matmul_2d(&normed, &layer.k_proj_wt), &layer.k_proj_bias);
@@ -381,7 +427,7 @@ impl GptOss {
self.all_reduce(&attn_proj);
let attn_proj = add_bias(&attn_proj, &layer.o_proj_bias);
let (normed, x_new) = add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let (normed, x_new) = Self::add_norm(&attn_proj, &residual, &layer.post_norm, &layer.post_norm_bias, eps);
let residual = x_new;
// MoE MLP
@@ -389,90 +435,65 @@ impl GptOss {
x = xserv_kernels::add(&residual, &moe_out);
}
let x = rmsnorm(&x, &self.norm, eps);
let x = Self::norm(&x, &self.norm, &self.norm_bias, eps);
let logits = matmul_2d(&x, &self.lm_head_t);
unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); }
logits
}
/// MoE forward pass for one layer with expert parallelism.
/// Each rank owns `num_experts / world` experts. Tokens routed to non-local
/// experts get zero contribution from this rank; AllReduce sums all ranks.
/// Input: [tokens, hidden], Output: [tokens, hidden]
/// MoE forward pass — fully on GPU via batched GEMM.
///
/// Each rank owns `local_experts` experts. The input is replicated across all
/// local experts, processed via two batched cuBLAS GEMMs (gate_up and down),
/// and the selected experts' outputs are weighted-summed on GPU. Non-selected
/// experts contribute zero (via the routing weights), so no scatter/gather is
/// needed. AllReduce sums partial results across TP ranks.
fn moe_forward(&self, x: &Tensor, layer: &GptOssBlock, num_tokens: usize) -> Tensor {
let hidden = self.config.hidden();
let num_experts = self.config.num_experts();
let top_k = self.config.experts_per_token();
let world = self.tp.as_ref().map(|tp| tp.world).unwrap_or(1);
let rank = self.tp.as_ref().map(|tp| tp.rank).unwrap_or(0);
let local_experts = num_experts / world;
let local_experts = layer.local_experts;
let expert_start = rank * local_experts;
// Router: [tokens, hidden] @ [hidden, num_experts] + bias → [tokens, num_experts]
unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); }
// 1. Router: [tokens, hidden] @ [hidden, num_experts] + bias → [tokens, num_experts]
let router_logits = add_bias(
&matmul_2d(x, &layer.router_wt),
&layer.router_bias,
);
unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); }
let router_cpu = router_logits.to_device(Device::Cpu);
let router_data = router_cpu.as_slice::<bf16>();
let x_cpu = x.to_device(Device::Cpu);
let x_data = x_cpu.as_slice::<bf16>();
let mut output_acc = vec![0.0f32; num_tokens * hidden];
// 2. GPU top-k + softmax
let (topk_ids, topk_weights) = xserv_kernels::moe::moe_topk_softmax(
&router_logits, num_experts, top_k,
);
for t in 0..num_tokens {
let row = &router_data[t * num_experts..(t + 1) * num_experts];
// 3. Replicate input: [tokens, hidden] → [local_experts, tokens, hidden]
let x_rep = xserv_kernels::moe::moe_replicate(x, local_experts);
// Find top-k expert indices (global)
let mut indices: Vec<(usize, f32)> = row.iter()
.enumerate()
.map(|(i, &v)| (i, v.to_f32()))
.collect();
indices.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
let top_indices: Vec<(usize, f32)> = indices[..top_k].to_vec();
// 4. Batched GEMM gate_up: [E, tokens, hidden] @ [E, hidden, 2*inter] → [E, tokens, 2*inter]
let gate_up = xserv_kernels::moe::batched_gemm_strided(&x_rep, &layer.expert_gate_up_wt);
// Softmax over top-k logits
let max_val = top_indices.iter().map(|x| x.1).fold(f32::NEG_INFINITY, f32::max);
let exp_sum: f32 = top_indices.iter().map(|x| (x.1 - max_val).exp()).sum();
let weights: Vec<f32> = top_indices.iter()
.map(|x| (x.1 - max_val).exp() / exp_sum)
.collect();
// 5. Bias add: gate_up += expert_gate_up_bias (in-place)
xserv_kernels::moe::moe_bias_add_3d(&gate_up, &layer.expert_gate_up_bias);
// Fresh GPU upload of token data — immune to cached allocator buffer reuse
let token_slice = &x_data[t * hidden..(t + 1) * hidden];
let token_tensor = Tensor::from_slice(token_slice, &[1, hidden]).to_device(x.device());
// 6. GLU activation: treat [E * tokens, 2*inter] → [E * tokens, inter]
let inter2 = gate_up.shape()[2];
let flat_rows = local_experts * num_tokens;
let gate_up_flat = gate_up.reshape(&[flat_rows, inter2]);
let activated_flat = gpt_oss_glu(&gate_up_flat, layer.glu_alpha, layer.glu_limit);
let inter = inter2 / 2;
let activated = activated_flat.reshape(&[local_experts, num_tokens, inter]);
// 7. Batched GEMM down: [E, tokens, inter] @ [E, inter, hidden] → [E, tokens, hidden]
let down = xserv_kernels::moe::batched_gemm_strided(&activated, &layer.expert_down_wt);
for (k_idx, &(expert_id, _)) in top_indices.iter().enumerate() {
// Only process experts owned by this rank
if expert_id < expert_start || expert_id >= expert_start + local_experts {
continue;
}
let local_id = expert_id - expert_start;
let weight = weights[k_idx];
// 8. Bias add: down += expert_down_bias (in-place)
xserv_kernels::moe::moe_bias_add_3d(&down, &layer.expert_down_bias);
let gate_up_raw = matmul_2d(&token_tensor, &layer.expert_gate_up_wt[local_id]);
let gate_up = add_bias(&gate_up_raw, &layer.expert_gate_up_bias[local_id]);
let activated = gpt_oss_glu(&gate_up, layer.glu_alpha, layer.glu_limit);
let down_raw = matmul_2d(&activated, &layer.expert_down_wt[local_id]);
let down = add_bias(&down_raw, &layer.expert_down_bias[local_id]);
let down_cpu = down.to_device(Device::Cpu);
let down_data = down_cpu.as_slice::<bf16>();
for d in 0..hidden {
output_acc[t * hidden + d] += weight * down_data[d].to_f32();
}
}
}
// Convert accumulated output to BF16 tensor on GPU
let output_bf16: Vec<bf16> = output_acc.iter().map(|&v| bf16::from_f32(v)).collect();
let moe_out = Tensor::from_slice(&output_bf16, &[num_tokens, hidden]).to_device(x.device());
// 9. Weighted sum across experts → [tokens, hidden]
let moe_out = xserv_kernels::moe::moe_weighted_sum(
&down, &topk_ids, &topk_weights,
expert_start, local_experts, top_k,
);
self.all_reduce(&moe_out);
moe_out
@@ -495,8 +516,15 @@ fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
let cols = x.shape()[1];
assert_eq!(bias.shape()[0], cols, "bias size {} != cols {}", bias.shape()[0], cols);
// Broadcast bias to each row using GPU kernels.
// Tile bias [cols] into [rows, cols] by repeating rows, then add element-wise.
let x_c = x.contiguous();
if rows == 1 {
// Fast path: reshape bias [cols] → [1, cols] (zero-copy), add directly on GPU
let bias_2d = bias.reshape(&[1, cols]);
return xserv_kernels::add(&x_c, &bias_2d);
}
// General path: tile bias to [rows, cols] via CPU, then add on GPU
let bias_cpu = bias.to_device(Device::Cpu);
let bias_data = bias_cpu.as_slice::<bf16>();
let mut tiled = Vec::with_capacity(rows * cols);
@@ -504,7 +532,6 @@ fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
tiled.extend_from_slice(bias_data);
}
let bias_tiled = Tensor::from_slice(&tiled, &[rows, cols]).to_device(x.device());
let x_c = x.contiguous();
xserv_kernels::add(&x_c, &bias_tiled)
}
@@ -554,23 +581,23 @@ fn shard_1d(t: &Tensor, rank: usize, world: usize) -> Tensor {
Tensor::from_slice(&shard, &[local])
}
/// Extract expert `e` from a [num_experts, rows, cols] 3D tensor → [rows, cols] 2D
fn slice_expert_3d(t: &Tensor, e: usize, rows: usize, cols: usize) -> Tensor {
/// Extract experts [start..start+count) from a [num_experts, rows, cols] 3D tensor
fn slice_expert_range_3d(t: &Tensor, start: usize, count: usize, rows: usize, cols: usize) -> Tensor {
assert_eq!(t.ndim(), 3);
let host = t.to_device(Device::Cpu);
let data = host.as_slice::<bf16>();
let stride = rows * cols;
let start = e * stride;
let slice = data[start..start + stride].to_vec();
Tensor::from_slice(&slice, &[rows, cols])
let offset = start * stride;
let slice = data[offset..offset + count * stride].to_vec();
Tensor::from_slice(&slice, &[count, rows, cols])
}
/// Extract expert `e` from a [num_experts, dim] 2D tensor → [dim] 1D
fn slice_expert_2d(t: &Tensor, e: usize, dim: usize) -> Tensor {
/// Extract experts [start..start+count) from a [num_experts, dim] 2D tensor
fn slice_expert_range_2d(t: &Tensor, start: usize, count: usize, dim: usize) -> Tensor {
assert_eq!(t.ndim(), 2);
let host = t.to_device(Device::Cpu);
let data = host.as_slice::<bf16>();
let start = e * dim;
let slice = data[start..start + dim].to_vec();
Tensor::from_slice(&slice, &[dim])
let offset = start * dim;
let slice = data[offset..offset + count * dim].to_vec();
Tensor::from_slice(&slice, &[count, dim])
}