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

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

@@ -30,61 +30,90 @@ struct GPT2Block {
}
pub struct KVCache {
// Per layer, per head: k[layer][head] has seq_len * head_dim floats
k: Vec<Vec<Vec<f32>>>, // [num_layers][num_heads][seq_len * head_dim]
v: Vec<Vec<Vec<f32>>>,
// Per layer, per head: raw bytes (works for both f32 and bf16)
k: Vec<Vec<Vec<u8>>>, // [num_layers][num_heads][seq_len * head_dim * elem_size]
v: Vec<Vec<Vec<u8>>>,
len: usize,
num_heads: usize,
head_dim: usize,
elem_size: usize,
dtype: DType,
device: Device,
}
impl KVCache {
pub fn new(num_layers: usize, num_heads: usize, head_dim: usize, device: Device) -> Self {
pub fn new(num_layers: usize, num_heads: usize, head_dim: usize, dtype: DType, device: Device) -> Self {
Self {
k: (0..num_layers).map(|_| vec![vec![]; num_heads]).collect(),
v: (0..num_layers).map(|_| vec![vec![]; num_heads]).collect(),
len: 0,
num_heads,
head_dim,
elem_size: dtype.size_bytes(),
dtype,
device,
}
}
pub fn seq_len(&self) -> usize { self.len }
/// Append new K/V data. k_new is in [1, H, new_tokens, D] layout (flat).
fn append_kv(&mut self, layer: usize, k_new: &[f32], v_new: &[f32], new_tokens: usize) {
/// Append from a CPU tensor with shape [1, H, new_tokens, D].
pub fn append_kv_tensor(&mut self, layer: usize, k_cpu: &Tensor, v_cpu: &Tensor, new_tokens: usize) {
let hd = self.head_dim;
let es = self.elem_size;
let k_bytes = k_cpu.storage().as_cpu_bytes();
let v_bytes = v_cpu.storage().as_cpu_bytes();
let chunk = new_tokens * hd * es;
for h in 0..self.num_heads {
let off = h * new_tokens * hd;
self.k[layer][h].extend_from_slice(&k_new[off..off + new_tokens * hd]);
self.v[layer][h].extend_from_slice(&v_new[off..off + new_tokens * hd]);
let off = h * chunk;
self.k[layer][h].extend_from_slice(&k_bytes[off..off + chunk]);
self.v[layer][h].extend_from_slice(&v_bytes[off..off + chunk]);
}
if layer == 0 {
self.len += new_tokens;
}
}
/// Reconstruct [1, H, seq_len, D] tensors from per-head cache.
fn get_kv_tensors(&self, layer: usize) -> (Tensor, Tensor) {
/// Reconstruct [1, H, seq_len, D] tensors.
pub fn get_kv_tensors(&self, layer: usize) -> (Tensor, Tensor) {
let sl = self.len;
let hd = self.head_dim;
let nh = self.num_heads;
let mut k_data = vec![0.0f32; nh * sl * hd];
let mut v_data = vec![0.0f32; nh * sl * hd];
let es = self.elem_size;
let head_bytes = sl * hd * es;
let total = nh * head_bytes;
let mut k_data = vec![0u8; total];
let mut v_data = vec![0u8; total];
for h in 0..nh {
let off = h * sl * hd;
k_data[off..off + sl * hd].copy_from_slice(&self.k[layer][h]);
v_data[off..off + sl * hd].copy_from_slice(&self.v[layer][h]);
let off = h * head_bytes;
k_data[off..off + head_bytes].copy_from_slice(&self.k[layer][h]);
v_data[off..off + head_bytes].copy_from_slice(&self.v[layer][h]);
}
let shape = &[1, nh, sl, hd];
let k = Tensor::from_slice(&k_data, shape).to_device(self.device);
let v = Tensor::from_slice(&v_data, shape).to_device(self.device);
let k = tensor_from_raw_bytes(&k_data, shape, self.dtype).to_device(self.device);
let v = tensor_from_raw_bytes(&v_data, shape, self.dtype).to_device(self.device);
(k, v)
}
}
fn tensor_from_raw_bytes(bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
match dtype {
DType::F32 => {
let data: &[f32] = unsafe {
std::slice::from_raw_parts(bytes.as_ptr() as *const f32, bytes.len() / 4)
};
Tensor::from_slice(data, shape)
}
DType::BF16 => {
let data: &[half::bf16] = unsafe {
std::slice::from_raw_parts(bytes.as_ptr() as *const half::bf16, bytes.len() / 2)
};
Tensor::from_slice(data, shape)
}
_ => panic!("unsupported dtype for KV cache"),
}
}
impl GPT2 {
pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
@@ -181,11 +210,10 @@ impl GPT2 {
let qkv = linear(&normed, &layer.attn_qkv_w, Some(&layer.attn_qkv_b));
let (q, k_new, v_new) = split_qkv(&qkv, num_heads, head_dim, new_tokens);
// KV cache: append new K/V, use full cached K/V for attention
let (k_full, v_full) = if let Some((cache, layer_idx)) = cache {
let k_cpu = k_new.to_device(Device::Cpu);
let v_cpu = v_new.to_device(Device::Cpu);
cache.append_kv(layer_idx, k_cpu.as_slice::<f32>(), v_cpu.as_slice::<f32>(), new_tokens);
cache.append_kv_tensor(layer_idx, &k_cpu, &v_cpu, new_tokens);
cache.get_kv_tensors(layer_idx)
} else {
(k_new, v_new)