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

phase 10: add Qwen3-8B benchmark + performance fix

Browse Source 
Benchmark infrastructure:
- bench-qwen3 binary: 50 prompts × 20 tokens with KV cache
- bench_compare_qwen3.py: comparison against HF transformers (BF16)

Performance fix:
- Precompute transposed weights at model load time (eliminated per-token
  weight transpose CPU round-trip: was 252 transposes × 32MB each = 8GB/token)
- Result: from "infinite" (>10 min/token) to 144ms/token

Results (50 prompts):
- Prefill top-1: 42/50 (84%), top-5: 50/50 (100%) vs HF transformers
- Greedy sequence: 0/50 exact match (BF16 precision drift over 36 layers)
- Performance: TTFT=138ms, TBT=144ms, 6.9 tok/s (HF: 21ms, 45.6 tok/s)
- All outputs are coherent English/Chinese

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-22 10:25:33 +08:00
parent 246ae1c590
commit 268e40d764
4 changed files with 389 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -11,22 +11,22 @@ pub struct Qwen3 {
embed_tokens: Tensor,
layers: Vec<Qwen3Block>,
norm: Tensor,
lm_head: Tensor,
lm_head_t: Tensor, // precomputed transpose
rope_cache: RopeCache,
}
struct Qwen3Block {
input_norm: Tensor, // [hidden]
q_proj_w: Tensor, // [num_heads*head_dim, hidden]
k_proj_w: Tensor, // [num_kv_heads*head_dim, hidden]
v_proj_w: Tensor,
o_proj_w: Tensor, // [hidden, num_heads*head_dim]
q_norm: Tensor, // [head_dim] — per-head QK norm
q_proj_wt: Tensor, // TRANSPOSED: [hidden, num_heads*head_dim]
k_proj_wt: Tensor, // TRANSPOSED: [hidden, num_kv_heads*head_dim]
v_proj_wt: Tensor,
o_proj_wt: Tensor, // TRANSPOSED: [num_heads*head_dim, hidden]
q_norm: Tensor, // [head_dim]
k_norm: Tensor, // [head_dim]
post_norm: Tensor, // [hidden]
gate_proj_w: Tensor, // [intermediate, hidden]
up_proj_w: Tensor,
down_proj_w: Tensor, // [hidden, intermediate]
gate_proj_wt: Tensor, // TRANSPOSED: [hidden, intermediate]
up_proj_wt: Tensor,
down_proj_wt: Tensor, // TRANSPOSED: [intermediate, hidden]
}
impl Qwen3 {
@@ -37,7 +37,7 @@ impl Qwen3 {
let embed_tokens = take(&mut w, "model.embed_tokens.weight");
let norm = take(&mut w, "model.norm.weight");
let lm_head = take(&mut w, "lm_head.weight");
let lm_head_raw = take(&mut w, "lm_head.weight");
let rope_cache = RopeCache::new(
config.max_seq_len().min(8192), // limit for memory
@@ -45,26 +45,33 @@ impl Qwen3 {
config.rope_theta.unwrap_or(1_000_000.0) as f32,
);
// Precompute transposed weights: [out, in] → [in, out] so we can do x @ wt directly
let transpose_w = |t: Tensor| -> Tensor {
t.transpose(0, 1).contiguous()
};
let num_layers = config.num_layers();
let mut layers = Vec::with_capacity(num_layers);
eprintln!("Transposing weights for {} layers...", num_layers);
for i in 0..num_layers {
let p = format!("model.layers.{i}");
layers.push(Qwen3Block {
input_norm: take(&mut w, &format!("{p}.input_layernorm.weight")),
q_proj_w: take(&mut w, &format!("{p}.self_attn.q_proj.weight")),
k_proj_w: take(&mut w, &format!("{p}.self_attn.k_proj.weight")),
v_proj_w: take(&mut w, &format!("{p}.self_attn.v_proj.weight")),
o_proj_w: take(&mut w, &format!("{p}.self_attn.o_proj.weight")),
q_proj_wt: transpose_w(take(&mut w, &format!("{p}.self_attn.q_proj.weight"))),
k_proj_wt: transpose_w(take(&mut w, &format!("{p}.self_attn.k_proj.weight"))),
v_proj_wt: transpose_w(take(&mut w, &format!("{p}.self_attn.v_proj.weight"))),
o_proj_wt: transpose_w(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
q_norm: take(&mut w, &format!("{p}.self_attn.q_norm.weight")),
k_norm: take(&mut w, &format!("{p}.self_attn.k_norm.weight")),
post_norm: take(&mut w, &format!("{p}.post_attention_layernorm.weight")),
gate_proj_w: take(&mut w, &format!("{p}.mlp.gate_proj.weight")),
up_proj_w: take(&mut w, &format!("{p}.mlp.up_proj.weight")),
down_proj_w: take(&mut w, &format!("{p}.mlp.down_proj.weight")),
gate_proj_wt: transpose_w(take(&mut w, &format!("{p}.mlp.gate_proj.weight"))),
up_proj_wt: transpose_w(take(&mut w, &format!("{p}.mlp.up_proj.weight"))),
down_proj_wt: transpose_w(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
});
}
Self { config, embed_tokens, layers, norm, lm_head, rope_cache }
let lm_head_t = transpose_w(lm_head_raw);
Self { config, embed_tokens, layers, norm, lm_head_t, rope_cache }
}
pub fn forward_with_cache(&self, token_ids: &[u32], cache: &mut KVCache) -> Tensor {
@@ -83,10 +90,10 @@ impl Qwen3 {
let residual = x.clone();
let normed = rmsnorm(&x, &layer.input_norm, eps);
// Q/K/V projections (no bias, weight is [out, in])
let q = linear_t(&normed, &layer.q_proj_w);
let k = linear_t(&normed, &layer.k_proj_w);
let v = linear_t(&normed, &layer.v_proj_w);
// Q/K/V projections (pre-transposed weights, x @ wt)
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);
// Reshape to [1, heads, seq, head_dim]
let q = reshape_heads(&q, new_tokens, num_heads, head_dim);
@@ -121,30 +128,31 @@ impl Qwen3 {
// Attention
let attn_out = attention(&q, &k_full, &v_full, true);
let attn_merged = merge_heads_any(&attn_out, new_tokens, hidden);
let attn_proj = linear_t(&attn_merged, &layer.o_proj_w);
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
x = add_any(&residual, &attn_proj);
// SwiGLU FFN
let residual = x.clone();
let normed = rmsnorm(&x, &layer.post_norm, eps);
let gate = linear_t(&normed, &layer.gate_proj_w);
let up = linear_t(&normed, &layer.up_proj_w);
let gate = matmul_2d(&normed, &layer.gate_proj_wt);
let up = matmul_2d(&normed, &layer.up_proj_wt);
let gate_activated = silu(&gate);
let hidden_states = mul_any(&gate_activated, &up);
let down = linear_t(&hidden_states, &layer.down_proj_w);
let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
x = add_any(&residual, &down);
}
let x = rmsnorm(&x, &self.norm, eps);
linear_t(&x, &self.lm_head)
matmul_2d(&x, &self.lm_head_t)
}
}
// --- Helpers ---
fn linear_t(x: &Tensor, weight: &Tensor) -> Tensor {
let w_t = weight.transpose(0, 1).contiguous();
matmul(x, &w_t, GemmBackend::CuBlas)
fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
assert_eq!(a.ndim(), 2);
assert_eq!(b.ndim(), 2);
matmul(a, b, GemmBackend::CuBlas)
}
fn reshape_heads(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {