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

phase 10: Qwen3-8B support (Milestone ②)

Browse Source 
Qwen3 model (qwen3.rs):
- RMSNorm + QK normalization (per-head q_norm/k_norm)
- GQA: 32 Q heads, 8 KV heads, repeat_kv for attention
- SwiGLU FFN: gate_proj → SiLU → * up_proj → down_proj
- RoPE with transpose for [1,H,S,D] ↔ [S,H,D] layout
- BF16 forward pass, [out,in] weight layout via linear_t
- No attention bias (attention_bias=false)

Tokenizer fixes:
- Fixed unicode_to_byte: shifted bytes now use correct inverse lookup table
- MergeEntry supports both string and array formats
- Both GPT-2 and Qwen3 tokenizers work correctly (English + Chinese)

KVCache refactored:
- Dtype-agnostic: stores raw bytes per-head, works for F32 and BF16
- append_kv_tensor/get_kv_tensors use Tensor directly

CLI updated:
- Auto-detects model type from config.json (gpt2 vs qwen3)
- Supports both GPT-2 (F32) and Qwen3 (BF16)

Verified: Qwen3-8B generates coherent English and Chinese on single RTX 5090.
61/61 tests pass, GPT-2 performance no regression.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-22 00:46:37 +08:00
parent 64084d3489
commit 246ae1c590
7 changed files with 553 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

278
crates/xserv-model/src/qwen3.rs Normal file
View File

@@ -0,0 +1,278 @@
use std::collections::HashMap;
use half::bf16;
use xserv_kernels::*;
use xserv_tensor::{DType, Device, Tensor};
use crate::config::ModelConfig;
use crate::gpt2::KVCache;
pub struct Qwen3 {
pub config: ModelConfig,
embed_tokens: Tensor,
layers: Vec<Qwen3Block>,
norm: Tensor,
lm_head: Tensor,
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
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]
}
impl Qwen3 {
pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
};
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 rope_cache = RopeCache::new(
config.max_seq_len().min(8192), // limit for memory
config.head_dim(),
config.rope_theta.unwrap_or(1_000_000.0) as f32,
);
let num_layers = config.num_layers();
let mut layers = Vec::with_capacity(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_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")),
});
}
Self { config, embed_tokens, layers, norm, lm_head, rope_cache }
}
pub fn forward_with_cache(&self, token_ids: &[u32], cache: &mut KVCache) -> Tensor {
let new_tokens = token_ids.len();
let pos_offset = cache.seq_len();
let hidden = self.config.hidden();
let num_heads = self.config.num_heads();
let num_kv_heads = self.config.num_kv_heads();
let head_dim = self.config.head_dim();
let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
let mut x = embedding(&self.embed_tokens, token_ids);
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
for (layer_idx, layer) in self.layers.iter().enumerate() {
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);
// Reshape to [1, heads, seq, head_dim]
let q = reshape_heads(&q, new_tokens, num_heads, head_dim);
let k = reshape_heads(&k, new_tokens, num_kv_heads, head_dim);
let v = reshape_heads(&v, new_tokens, num_kv_heads, head_dim);
// QK normalization (per-head RMSNorm)
let q = head_rmsnorm(&q, &layer.q_norm, eps);
let k = head_rmsnorm(&k, &layer.k_norm, eps);
// RoPE — kernel expects [S, H, D], our tensors are [1, H, S, D]
// Transpose to [1, S, H, D] → reshape to [S, H, D] for RoPE
let q = transpose_for_rope(&q, new_tokens, num_heads, head_dim);
let k = transpose_for_rope(&k, new_tokens, num_kv_heads, head_dim);
rope_inplace(&q, &self.rope_cache, &positions);
rope_inplace(&k, &self.rope_cache, &positions);
// Transpose back to [1, H, S, D]
let q = transpose_from_rope(&q, new_tokens, num_heads, head_dim);
let k = transpose_from_rope(&k, new_tokens, num_kv_heads, head_dim);
// KV cache
let k_cpu = k.to_device(Device::Cpu);
let v_cpu = v.to_device(Device::Cpu);
cache.append_kv_tensor(layer_idx, &k_cpu, &v_cpu, new_tokens);
let (k_full, v_full) = cache.get_kv_tensors(layer_idx);
// GQA: repeat K/V
let n_rep = num_heads / num_kv_heads;
let k_full = repeat_kv(&k_full, n_rep);
let v_full = repeat_kv(&v_full, n_rep);
// 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);
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_activated = silu(&gate);
let hidden_states = mul_any(&gate_activated, &up);
let down = linear_t(&hidden_states, &layer.down_proj_w);
x = add_any(&residual, &down);
}
let x = rmsnorm(&x, &self.norm, eps);
linear_t(&x, &self.lm_head)
}
}
// --- Helpers ---
fn linear_t(x: &Tensor, weight: &Tensor) -> Tensor {
let w_t = weight.transpose(0, 1).contiguous();
matmul(x, &w_t, GemmBackend::CuBlas)
}
fn reshape_heads(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
let x_cpu = x.to_device(Device::Cpu);
let hidden = num_heads * head_dim;
let src = x_cpu.as_slice::<bf16>();
let mut out = vec![bf16::ZERO; num_heads * seq_len * head_dim];
for s in 0..seq_len {
for h in 0..num_heads {
let si = s * hidden + h * head_dim;
let di = (h * seq_len + s) * head_dim;
out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
}
}
Tensor::from_slice(&out, &[1, num_heads, seq_len, head_dim]).to_device(x.device())
}
fn merge_heads_any(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
let num_heads = x.shape()[1];
let head_dim = x.shape()[3];
let x_cpu = x.to_device(Device::Cpu);
let src = x_cpu.as_slice::<bf16>();
let mut out = vec![bf16::ZERO; seq_len * hidden];
for s in 0..seq_len {
for h in 0..num_heads {
let si = (h * seq_len + s) * head_dim;
let di = s * hidden + h * head_dim;
out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
}
}
Tensor::from_slice(&out, &[seq_len, hidden]).to_device(x.device())
}
/// Per-head RMSNorm: apply RMSNorm to each [head_dim] slice independently.
/// x: [1, H, S, D], norm_weight: [D]
fn head_rmsnorm(x: &Tensor, norm_weight: &Tensor, eps: f32) -> Tensor {
let num_heads = x.shape()[1];
let seq_len = x.shape()[2];
let head_dim = x.shape()[3];
// Reshape to [H*S, D], apply rmsnorm, reshape back
let total_rows = num_heads * seq_len;
let flat = x.reshape(&[total_rows, head_dim]);
let normed = rmsnorm(&flat, norm_weight, eps);
normed.reshape(&[1, num_heads, seq_len, head_dim])
}
/// [1, H, S, D] → [S, H, D] for RoPE kernel
fn transpose_for_rope(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
let x_cpu = x.to_device(Device::Cpu);
let src = x_cpu.as_slice::<bf16>();
let mut out = vec![bf16::ZERO; seq_len * num_heads * head_dim];
for h in 0..num_heads {
for s in 0..seq_len {
let si = (h * seq_len + s) * head_dim;
let di = (s * num_heads + h) * head_dim;
out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
}
}
Tensor::from_slice(&out, &[seq_len, num_heads, head_dim]).to_device(x.device())
}
/// [S, H, D] → [1, H, S, D] after RoPE
fn transpose_from_rope(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
let x_cpu = x.to_device(Device::Cpu);
let src = x_cpu.as_slice::<bf16>();
let mut out = vec![bf16::ZERO; num_heads * seq_len * head_dim];
for s in 0..seq_len {
for h in 0..num_heads {
let si = (s * num_heads + h) * head_dim;
let di = (h * seq_len + s) * head_dim;
out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
}
}
Tensor::from_slice(&out, &[1, num_heads, seq_len, head_dim]).to_device(x.device())
}
fn repeat_kv(x: &Tensor, n_rep: usize) -> Tensor {
if n_rep == 1 { return x.clone(); }
let kv_heads = x.shape()[1];
let seq_len = x.shape()[2];
let head_dim = x.shape()[3];
let x_cpu = x.to_device(Device::Cpu);
let src = x_cpu.as_slice::<bf16>();
let new_heads = kv_heads * n_rep;
let mut out = vec![bf16::ZERO; new_heads * seq_len * head_dim];
let chunk = seq_len * head_dim;
for kv_h in 0..kv_heads {
for r in 0..n_rep {
let dst_h = kv_h * n_rep + r;
out[dst_h * chunk..(dst_h + 1) * chunk]
.copy_from_slice(&src[kv_h * chunk..(kv_h + 1) * chunk]);
}
}
Tensor::from_slice(&out, &[1, new_heads, seq_len, head_dim]).to_device(x.device())
}
fn add_any(a: &Tensor, b: &Tensor) -> Tensor {
assert_eq!(a.shape(), b.shape());
let a_cpu = a.to_device(Device::Cpu);
let b_cpu = b.to_device(Device::Cpu);
let ad = a_cpu.as_slice::<bf16>();
let bd = b_cpu.as_slice::<bf16>();
let r: Vec<bf16> = ad.iter().zip(bd)
.map(|(x, y)| bf16::from_f32(x.to_f32() + y.to_f32()))
.collect();
Tensor::from_slice(&r, a.shape()).to_device(a.device())
}
fn mul_any(a: &Tensor, b: &Tensor) -> Tensor {
assert_eq!(a.shape(), b.shape());
let a_cpu = a.to_device(Device::Cpu);
let b_cpu = b.to_device(Device::Cpu);
let ad = a_cpu.as_slice::<bf16>();
let bd = b_cpu.as_slice::<bf16>();
let r: Vec<bf16> = ad.iter().zip(bd)
.map(|(x, y)| bf16::from_f32(x.to_f32() * y.to_f32()))
.collect();
Tensor::from_slice(&r, a.shape()).to_device(a.device())
}
pub fn sample_greedy(logits: &Tensor) -> u32 {
assert_eq!(logits.ndim(), 2);
let logits_cpu = logits.to_device(Device::Cpu);
let vocab_size = logits.shape()[1];
let seq_len = logits.shape()[0];
let data = logits_cpu.as_slice::<bf16>();
let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
last.iter().enumerate()
.max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
.map(|(i, _)| i as u32).unwrap()
}