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

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:
2026-05-22 00:46:37 +08:00
parent 64084d3489
commit 246ae1c590
7 changed files with 553 additions and 99 deletions

View File

@@ -1,8 +1,7 @@
use std::io::{self, Write};
use std::path::PathBuf;
use xserv_model::gpt2::{sample_greedy, KVCache};
use xserv_model::{loader, GPT2, ModelConfig};
use xserv_tensor::Device;
use xserv_model::{loader, KVCache, ModelConfig};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
fn main() {
@@ -25,43 +24,59 @@ fn main() {
eprintln!("GPU: {} ({} MB free)", info.name, info.free_memory / 1024 / 1024);
let config = ModelConfig::from_file(&model_dir.join("config.json"));
let model_type = config.model_type.as_deref().unwrap_or("unknown");
eprintln!(
"Model: {:?}, layers={}, hidden={}, heads={}, vocab={}",
config.model_type,
config.num_layers(),
config.hidden(),
config.num_heads(),
config.vocab_size
"Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}",
config.num_layers(), config.hidden(), config.num_heads(),
config.num_kv_heads(), config.vocab_size
);
eprintln!("Loading weights...");
let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
eprintln!("Loaded {} tensors", weights.len());
let model = GPT2::from_weights(config.clone(), weights);
let is_qwen3 = model_type.contains("qwen");
let dtype = if is_qwen3 { DType::BF16 } else { DType::F32 };
// Build model
enum Model {
GPT2(xserv_model::GPT2),
Qwen3(xserv_model::Qwen3),
}
let model = if is_qwen3 {
Model::Qwen3(xserv_model::Qwen3::from_weights(config.clone(), weights))
} else {
Model::GPT2(xserv_model::GPT2::from_weights(config.clone(), weights))
};
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
eprintln!("Ready (KV cache enabled).\n");
eprintln!("Ready (KV cache, dtype={dtype}).\n");
loop {
print!("xserv> ");
io::stdout().flush().unwrap();
let mut input = String::new();
if io::stdin().read_line(&mut input).unwrap() == 0 {
break;
}
if io::stdin().read_line(&mut input).unwrap() == 0 { break; }
let input = input.trim();
if input.is_empty() { continue; }
if input == "quit" || input == "exit" { break; }
let token_ids = tokenizer.encode(input);
let kv_heads = if is_qwen3 { config.num_kv_heads() } else { config.num_heads() };
let mut cache = KVCache::new(
config.num_layers(), config.num_heads(), config.head_dim(),
Device::Cuda(0),
config.num_layers(), kv_heads, config.head_dim(), dtype, Device::Cuda(0),
);
// Prefill
let logits = model.forward_with_cache(&token_ids, &mut cache);
let mut next = sample_greedy(&logits);
// Prefill + decode
let logits = match &model {
Model::GPT2(m) => m.forward_with_cache(&token_ids, &mut cache),
Model::Qwen3(m) => m.forward_with_cache(&token_ids, &mut cache),
};
let mut next = match &model {
Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
};
print!("{input}");
io::stdout().flush().unwrap();
@@ -72,8 +87,14 @@ fn main() {
if tokenizer.eos_token_id() == Some(next) { break; }
let logits = model.forward_with_cache(&[next], &mut cache);
next = sample_greedy(&logits);
let logits = match &model {
Model::GPT2(m) => m.forward_with_cache(&[next], &mut cache),
Model::Qwen3(m) => m.forward_with_cache(&[next], &mut cache),
};
next = match &model {
Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
};
}
println!();
}