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

phase19: MoE support — gpt-oss-20b end-to-end inference with TP=2

Browse Source 
Add Mixture-of-Experts support for the gpt-oss-20b model (20.9B params,
32 experts × top-4 routing). Key additions:

- ModelConfig: MoE fields (num_local_experts, layer_types, sliding_window,
  attention_bias, explicit head_dim, rope_scaling, swiglu_limit)
- YaRN RoPE: RopeCache::new_yarn() with correct frequency interpolation
  and attention_scaling = 0.1*ln(factor)+1
- Custom GLU kernel: gpt_oss_glu_bf16 (clamped sigmoid gate activation)
- Paged attention with sinks + sliding window kernel variant
- GptOss model struct with expert-parallel TP (split 32 experts across ranks)
- bench-gpt-oss binary for TP inference benchmarking

Verified on dash5 with 2x RTX 5090: 63.6 tok/s decode, ~160ms TTFT.
Model generates topically-coherent output (needs chat template for quality).

Known issues:
- Custom GEMV kernel produces NaN with small N (workaround: pad to M=2)
- Prefill doesn't use attention sinks (uses standard flash attention)
- Output quality requires chat template formatting

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-30 15:18:01 +08:00
parent 46bfb59f30
commit 9ad91a4a92
12 changed files with 1390 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
crates/xserv-model/src/qwen3.rs
View File

@@ -198,17 +198,27 @@ impl Qwen3 {
);
for i in lo..hi {
let p = format!("model.layers.{i}");
let q_proj_wt = wt(take(&mut w, &format!("{p}.self_attn.q_proj.weight")));
let k_proj_wt = wt(take(&mut w, &format!("{p}.self_attn.k_proj.weight")));
let v_proj_wt = wt(take(&mut w, &format!("{p}.self_attn.v_proj.weight")));
let q_dim = q_proj_wt.shape()[1];
let kv_dim = k_proj_wt.shape()[1];
let qkv_proj_wt = cat_cols(&[&q_proj_wt, &k_proj_wt, &v_proj_wt]);
drop((q_proj_wt, k_proj_wt, v_proj_wt));
let gate_proj_wt = wt(take(&mut w, &format!("{p}.mlp.gate_proj.weight")));
let up_proj_wt = wt(take(&mut w, &format!("{p}.mlp.up_proj.weight")));
let gate_up_proj_wt = cat_cols(&[&gate_proj_wt, &up_proj_wt]);
drop((gate_proj_wt, up_proj_wt));
layers.push(Qwen3Block {
input_norm: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))),
q_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.q_proj.weight"))),
k_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.k_proj.weight"))),
v_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.v_proj.weight"))),
qkv_proj_wt,
q_dim,
kv_dim,
o_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
q_norm: repl(take(&mut w, &format!("{p}.self_attn.q_norm.weight"))),
k_norm: repl(take(&mut w, &format!("{p}.self_attn.k_norm.weight"))),
post_norm: repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight"))),
gate_proj_wt: wt(take(&mut w, &format!("{p}.mlp.gate_proj.weight"))),
up_proj_wt: wt(take(&mut w, &format!("{p}.mlp.up_proj.weight"))),
gate_up_proj_wt,
down_proj_wt: wt(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
});
}
@@ -272,9 +282,10 @@ impl Qwen3 {
let residual = x.clone();
let normed = rmsnorm(&x, &layer.input_norm, eps);
let q = matmul_2d(&normed, &layer.q_proj_wt);
let k = matmul_2d(&normed, &layer.k_proj_wt);
let v = matmul_2d(&normed, &layer.v_proj_wt);
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let q = xserv_kernels::reshape_heads_gpu(&q, new_tokens, num_heads, head_dim);
let k = xserv_kernels::reshape_heads_gpu(&k, new_tokens, num_kv_heads, head_dim);
@@ -300,8 +311,10 @@ impl Qwen3 {
let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate = matmul_2d(&normed, &layer.gate_proj_wt);
let up = matmul_2d(&normed, &layer.up_proj_wt);
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
let ffn_dim = gate_up.shape()[1] / 2;
let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
let hidden_states = xserv_kernels::silu_mul(&gate, &up);
let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
x = add_any(&residual, &down);
@@ -340,9 +353,10 @@ impl Qwen3 {
let residual = x.clone();
let normed = rmsnorm(&x, &layer.input_norm, eps);
let q_all = matmul_2d(&normed, &layer.q_proj_wt);
let k_all = matmul_2d(&normed, &layer.k_proj_wt);
let v_all = matmul_2d(&normed, &layer.v_proj_wt);
let qkv_all = matmul_2d(&normed, &layer.qkv_proj_wt);
let q_all = qkv_all.narrow(1, 0, layer.q_dim).contiguous();
let k_all = qkv_all.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v_all = qkv_all.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let mut q_rows: Vec<Tensor> = Vec::with_capacity(batch);
for b in 0..batch {
@@ -390,8 +404,10 @@ impl Qwen3 {
let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate = matmul_2d(&normed, &layer.gate_proj_wt);
let up = matmul_2d(&normed, &layer.up_proj_wt);
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
let ffn_dim = gate_up.shape()[1] / 2;
let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
let hidden_states = xserv_kernels::silu_mul(&gate, &up);
let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
x = add_any(&residual, &down);