phase 12: implement real continuous batching scheduler

Rewrote engine.rs from scratch:
- Scheduler loop: admit → prefill → decode → finish → check new requests
- Multiple sequences run concurrently (max_batch_size configurable)
- Each sequence has independent GpuKVCache
- Non-blocking try_recv() for new requests during decode iterations
- Dynamic join: new requests enter batch immediately, don't wait for others

Verified with concurrent test (tools/test_concurrent.py):
- 3 concurrent requests: wall_time=3.8s, concurrency_ratio=2.82x ✓
- 5 concurrent requests: wall_time=6.1s, concurrency_ratio=4.04x ✓
- All outputs are coherent and correct

Design doc (docs/12-continuous-batching.md) fully rewritten with:
- Detailed scheduler loop pseudocode
- Data structures (Sequence, Scheduler)
- Acceptance criteria with specific test cases
- Clear separation from Phase 13 (HTTP layer)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Cargo.toml

@@ -6,6 +6,7 @@ members = [
     "crates/xserv-kernels",
     "crates/xserv-model",
     "crates/xserv-tokenizer",
+    "crates/xserv-server",
 ]
 
 [workspace.package]
@@ -20,3 +21,6 @@ serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 safetensors = "0.5"
 regex = "1"
+tokio = { version = "1", features = ["full"] }
+axum = "0.8"
+uuid = { version = "1", features = ["v4"] }

crates/xserv-cuda/src/memory.rs

@@ -87,6 +87,20 @@ impl GpuBuffer {
         error::check(unsafe { ffi::cudaMemset(self.ptr, 0, self.len) })
     }
 
+    /// Copy `count` bytes from `src` buffer at `src_offset` to this buffer at `dst_offset`.
+    pub fn copy_from_device_at(&mut self, src: &GpuBuffer, src_offset: usize, dst_offset: usize, count: usize) -> Result<()> {
+        assert!(src_offset + count <= src.len);
+        assert!(dst_offset + count <= self.len);
+        error::check(unsafe {
+            ffi::cudaMemcpy(
+                self.ptr.add(dst_offset),
+                src.ptr.add(src_offset),
+                count,
+                ffi::CUDA_MEMCPY_D2D,
+            )
+        })
+    }
+
     /// Consume the buffer without freeing GPU memory. Returns the raw pointer and length.
     /// Caller is responsible for eventually calling cudaFree.
     pub fn into_raw(self) -> (*mut u8, usize) {

crates/xserv-kernels/build.rs

@@ -23,6 +23,7 @@ fn main() {
         .file("../../csrc/embedding/embedding.cu")
         .file("../../csrc/embedding/rope.cu")
         .file("../../csrc/attention/causal_mask.cu")
+        .file("../../csrc/embedding/transpose.cu")
         .compile("xserv_kernels");
 
     println!("cargo:rerun-if-changed=../../csrc/");

crates/xserv-kernels/src/activation.rs

@@ -8,43 +8,53 @@ unsafe extern "C" {
     fn launch_silu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
     fn launch_scale_f32(x: *const c_void, out: *mut c_void, scale: f32, n: i32, stream: *mut c_void);
     fn launch_scale_bf16(x: *const c_void, out: *mut c_void, scale: f32, n: i32, stream: *mut c_void);
+    fn launch_add_f32(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_add_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_mul_f32(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_mul_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
 }
 
-pub fn gelu(x: &Tensor) -> Tensor {
-    assert!(x.is_contiguous());
-    assert!(matches!(x.device(), Device::Cuda(_)));
+fn dispatch_unary(x: &Tensor, f32_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
+                  bf16_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void)) -> Tensor {
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
     let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
     let n = x.numel() as i32;
     unsafe {
         match x.dtype() {
-            DType::F32 => launch_gelu_f32(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
-            DType::BF16 => launch_gelu_bf16(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
-            _ => panic!("unsupported dtype for gelu"),
+            DType::F32 => f32_fn(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
+            DType::BF16 => bf16_fn(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
+            _ => panic!("unsupported dtype"),
         }
     }
     xserv_cuda::device::synchronize().unwrap();
     out
 }
 
-pub fn silu(x: &Tensor) -> Tensor {
-    assert!(x.is_contiguous());
-    assert!(matches!(x.device(), Device::Cuda(_)));
-    let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
-    let n = x.numel() as i32;
+fn dispatch_binary(a: &Tensor, b: &Tensor,
+                   f32_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void),
+                   bf16_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void)) -> Tensor {
+    assert_eq!(a.shape(), b.shape());
+    assert!(a.is_contiguous() && b.is_contiguous());
+    assert!(matches!(a.device(), Device::Cuda(_)));
+    assert_eq!(a.dtype(), b.dtype());
+    let out = Tensor::zeros(a.shape(), a.dtype(), a.device());
+    let n = a.numel() as i32;
     unsafe {
-        match x.dtype() {
-            DType::F32 => launch_silu_f32(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
-            DType::BF16 => launch_silu_bf16(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
-            _ => panic!("unsupported dtype for silu"),
+        match a.dtype() {
+            DType::F32 => f32_fn(a.data_ptr() as _, b.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
+            DType::BF16 => bf16_fn(a.data_ptr() as _, b.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
+            _ => panic!("unsupported dtype"),
         }
     }
     xserv_cuda::device::synchronize().unwrap();
     out
 }
 
+pub fn gelu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_gelu_f32, launch_gelu_bf16) }
+pub fn silu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_silu_f32, launch_silu_bf16) }
+
 pub fn scale(x: &Tensor, scale_val: f32) -> Tensor {
-    assert!(x.is_contiguous());
-    assert!(matches!(x.device(), Device::Cuda(_)));
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
     let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
     let n = x.numel() as i32;
     unsafe {
@@ -57,3 +67,6 @@ pub fn scale(x: &Tensor, scale_val: f32) -> Tensor {
     xserv_cuda::device::synchronize().unwrap();
     out
 }
+
+pub fn add(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_add_f32, launch_add_bf16) }
+pub fn mul(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16) }

crates/xserv-kernels/src/lib.rs

@@ -6,8 +6,10 @@ pub mod layernorm;
 pub mod rmsnorm;
 pub mod rope;
 pub mod softmax;
+pub mod transpose;
 
-pub use activation::{gelu, scale, silu};
+pub use activation::{add, gelu, mul, scale, silu};
+pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
 pub use attention::attention;
 pub use embedding::embedding;
 pub use gemm::{batched_matmul, matmul, GemmBackend};

crates/xserv-kernels/src/transpose.rs

@@ -0,0 +1,91 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_reshape_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_merge_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_repeat_kv_bf16(inp: *const c_void, out: *mut c_void, kv_heads: i32, n_rep: i32, seq_len: i32, head_dim: i32, stream: *mut c_void);
+}
+
+/// [S, H*D] → [1, H, S, D] on GPU (BF16)
+pub fn reshape_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
+    let out = Tensor::zeros(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
+    unsafe {
+        launch_reshape_heads_bf16(
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
+        );
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}
+
+/// [1, H, S, D] → [S, H*D] on GPU (BF16)
+pub fn merge_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
+    let hidden = num_heads * head_dim;
+    let out = Tensor::zeros(&[seq_len, hidden], DType::BF16, x.device());
+    unsafe {
+        launch_merge_heads_bf16(
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
+        );
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}
+
+/// [1, H, S, D] → [S, H, D] for RoPE on GPU (BF16)
+pub fn transpose_for_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
+    let out = Tensor::zeros(&[seq_len, num_heads, head_dim], DType::BF16, x.device());
+    unsafe {
+        launch_transpose_hsd_to_shd_bf16(
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
+        );
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}
+
+/// [S, H, D] → [1, H, S, D] after RoPE on GPU (BF16)
+pub fn transpose_from_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
+    let out = Tensor::zeros(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
+    unsafe {
+        launch_transpose_shd_to_hsd_bf16(
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
+        );
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}
+
+/// [1, KV_H, S, D] → [1, KV_H*n_rep, S, D] on GPU (BF16)
+pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
+    if n_rep == 1 { return x.clone(); }
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
+    let kv_heads = x.shape()[1];
+    let seq_len = x.shape()[2];
+    let head_dim = x.shape()[3];
+    let new_heads = kv_heads * n_rep;
+    let out = Tensor::zeros(&[1, new_heads, seq_len, head_dim], DType::BF16, x.device());
+    unsafe {
+        launch_repeat_kv_bf16(
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            kv_heads as i32, n_rep as i32, seq_len as i32, head_dim as i32, std::ptr::null_mut(),
+        );
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}

crates/xserv-model/Cargo.toml

@@ -9,6 +9,7 @@ xserv-tensor = { path = "../xserv-tensor" }
 xserv-kernels = { path = "../xserv-kernels" }
 xserv-tokenizer = { path = "../xserv-tokenizer" }
 half.workspace = true
+smallvec.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 safetensors.workspace = true

crates/xserv-model/src/bin/bench-gpt2.rs

@@ -0,0 +1,198 @@
+use std::path::PathBuf;
+use std::time::Instant;
+use xserv_model::gpt2::{sample_greedy, KVCache};
+use xserv_model::{loader, GPT2, ModelConfig};
+use xserv_tensor::Device;
+use xserv_tokenizer::Tokenizer;
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: bench-gpt2 <model-dir> [--gen-tokens N] [--no-cache]");
+        std::process::exit(1);
+    }
+    let model_dir = PathBuf::from(&args[1]);
+    let gen_tokens: usize = args
+        .iter()
+        .position(|a| a == "--gen-tokens")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(20);
+    let use_cache = !args.iter().any(|a| a == "--no-cache");
+
+    xserv_cuda::device::set_device(0).unwrap();
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
+    let model = GPT2::from_weights(config.clone(), weights);
+    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+
+    // Warmup
+    {
+        let ids = tokenizer.encode("warmup");
+        let _ = model.forward(&ids);
+    }
+
+    eprintln!("mode: {}", if use_cache { "KV cache" } else { "no cache" });
+
+    let prompts: Vec<&str> = vec![
+        "The capital of France is",
+        "Once upon a time in a land far away",
+        "Hello, how are you doing today",
+        "In a shocking finding, scientists discovered a",
+        "The weather today is sunny, so I decided to",
+        "Alan Turing was a British mathematician who",
+        "The best way to learn programming is",
+        "Artificial intelligence will change the world because",
+        "The history of the internet began in the",
+        "A good morning routine starts with",
+        "The stock market crashed because investors",
+        "Deep learning is a subset of machine learning that",
+        "The president of the United States announced",
+        "In the year 2050, humans will",
+        "The secret to happiness is",
+        "When I was a child, I used to",
+        "The most important scientific discovery of the century",
+        "Climate change is caused by",
+        "The recipe for chocolate cake requires",
+        "In conclusion, the evidence suggests that",
+        "The cat sat on the mat and",
+        "According to recent studies, exercise can",
+        "The first step in solving any problem is",
+        "Technology has transformed the way we",
+        "The novel begins with the protagonist",
+        "Education is the most powerful weapon",
+        "The ocean covers more than seventy percent of",
+        "Last night I had a dream about",
+        "The company announced its quarterly earnings",
+        "Music has the power to",
+        "The difference between success and failure is",
+        "In the beginning, there was nothing but",
+        "The doctor told me that I should",
+        "Python is a popular programming language because",
+        "The ancient Romans built roads that",
+        "A balanced diet should include",
+        "The movie received mixed reviews from critics",
+        "Space exploration has led to many",
+        "The teacher asked the students to",
+        "Global warming is one of the most",
+        "The bridge collapsed due to structural",
+        "Quantum computing promises to revolutionize",
+        "The new policy will affect millions of",
+        "During the winter months, it is important to",
+        "The human brain contains approximately",
+        "Democracy depends on the active participation of",
+        "The train arrived at the station exactly",
+        "Researchers at MIT have developed a new",
+        "The smartphone has become an essential part of",
+        "After careful consideration, the committee decided to",
+    ];
+
+    println!("[");
+    for (i, prompt) in prompts.iter().enumerate() {
+        let input_ids = tokenizer.encode(prompt);
+        let input_len = input_ids.len();
+
+        let (generated_ids, ttft_us, token_times_us) = if use_cache {
+            generate_with_cache(&model, &config, &tokenizer, &input_ids, gen_tokens)
+        } else {
+            generate_no_cache(&model, &tokenizer, &input_ids, gen_tokens)
+        };
+
+        let num_generated = generated_ids.len();
+        let generated_text = tokenizer.decode(&generated_ids);
+
+        let tbt_us = if !token_times_us.is_empty() {
+            token_times_us.iter().sum::<u128>() / token_times_us.len() as u128
+        } else { 0 };
+        let total_gen_us: u128 = ttft_us + token_times_us.iter().sum::<u128>();
+        let tpot_us = if num_generated > 0 { total_gen_us / num_generated as u128 } else { 0 };
+
+        let gen_text_escaped = generated_text
+            .replace('\\', "\\\\")
+            .replace('"', "\\\"")
+            .replace('\n', "\\n")
+            .replace('\r', "\\r")
+            .replace('\t', "\\t");
+        let gen_ids_str: Vec<String> = generated_ids.iter().map(|id| id.to_string()).collect();
+
+        print!("  {{\"prompt\": \"{}\", ", prompt.replace('"', "\\\""));
+        print!("\"input_len\": {input_len}, ");
+        print!("\"num_generated\": {num_generated}, ");
+        print!("\"generated_ids\": [{}], ", gen_ids_str.join(", "));
+        print!("\"generated_text\": \"{gen_text_escaped}\", ");
+        print!("\"ttft_us\": {ttft_us}, ");
+        print!("\"tbt_us\": {tbt_us}, ");
+        print!("\"tpot_us\": {tpot_us}}}");
+        if i < prompts.len() - 1 { println!(","); } else { println!(); }
+
+        eprintln!(
+            "[{}/{}] input={input_len}tok gen={num_generated}tok ttft={:.1}ms tbt={:.1}ms | {}",
+            i + 1, prompts.len(),
+            ttft_us as f64 / 1000.0,
+            tbt_us as f64 / 1000.0,
+            &generated_text.replace('\n', " ")[..generated_text.len().min(60)]
+        );
+    }
+    println!("]");
+}
+
+fn generate_with_cache(
+    model: &GPT2, config: &ModelConfig, tokenizer: &Tokenizer,
+    input_ids: &[u32], gen_tokens: usize,
+) -> (Vec<u32>, u128, Vec<u128>) {
+    let mut cache = KVCache::new(
+        config.num_layers(), config.num_heads(), config.head_dim(),
+        xserv_tensor::DType::F32, Device::Cuda(0),
+    );
+
+    // Prefill
+    let t0 = Instant::now();
+    let logits = model.forward_with_cache(input_ids, &mut cache);
+    let first_token = sample_greedy(&logits);
+    let ttft_us = t0.elapsed().as_micros();
+
+    let mut generated = vec![first_token];
+    let mut token_times = Vec::new();
+
+    // Decode
+    for _ in 1..gen_tokens {
+        let last = *generated.last().unwrap();
+        let t_start = Instant::now();
+        let logits = model.forward_with_cache(&[last], &mut cache);
+        let next = sample_greedy(&logits);
+        token_times.push(t_start.elapsed().as_micros());
+        generated.push(next);
+        if tokenizer.eos_token_id() == Some(next) { break; }
+    }
+
+    (generated, ttft_us, token_times)
+}
+
+fn generate_no_cache(
+    model: &GPT2, tokenizer: &Tokenizer,
+    input_ids: &[u32], gen_tokens: usize,
+) -> (Vec<u32>, u128, Vec<u128>) {
+    let mut all_ids = input_ids.to_vec();
+
+    let t0 = Instant::now();
+    let logits = model.forward(&all_ids);
+    let first_token = sample_greedy(&logits);
+    let ttft_us = t0.elapsed().as_micros();
+    all_ids.push(first_token);
+
+    let mut generated = vec![first_token];
+    let mut token_times = Vec::new();
+
+    for _ in 1..gen_tokens {
+        let t_start = Instant::now();
+        let logits = model.forward(&all_ids);
+        let next = sample_greedy(&logits);
+        token_times.push(t_start.elapsed().as_micros());
+        all_ids.push(next);
+        generated.push(next);
+        if tokenizer.eos_token_id() == Some(next) { break; }
+    }
+
+    (generated, ttft_us, token_times)
+}

crates/xserv-model/src/bin/bench-qwen3.rs

@@ -0,0 +1,156 @@
+use std::path::PathBuf;
+use std::time::Instant;
+use xserv_model::qwen3::sample_greedy;
+use xserv_model::{loader, GpuKVCache, ModelConfig, Qwen3};
+use xserv_tensor::{DType, Device};
+use xserv_tokenizer::Tokenizer;
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: bench-qwen3 <model-dir> [--gen-tokens N]");
+        std::process::exit(1);
+    }
+    let model_dir = PathBuf::from(&args[1]);
+    let gen_tokens: usize = args
+        .iter()
+        .position(|a| a == "--gen-tokens")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(20);
+
+    xserv_cuda::device::set_device(0).unwrap();
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    eprintln!("Loading Qwen3-8B weights...");
+    let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
+    eprintln!("Loaded {} tensors", weights.len());
+    let model = Qwen3::from_weights(config.clone(), weights);
+    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+
+    // Warmup
+    {
+        let ids = tokenizer.encode("warmup");
+        let mut cache = GpuKVCache::new(&config, 256, DType::BF16);
+        let _ = model.forward_gpu_cache(&ids, &mut cache);
+    }
+    eprintln!("Warmup done. Running benchmark...");
+
+    let prompts: Vec<&str> = vec![
+        "The capital of France is",
+        "Once upon a time in a land far away",
+        "Hello, how are you doing today",
+        "In a shocking finding, scientists discovered a",
+        "The weather today is sunny, so I decided to",
+        "Alan Turing was a British mathematician who",
+        "The best way to learn programming is",
+        "Artificial intelligence will change the world because",
+        "The history of the internet began in the",
+        "A good morning routine starts with",
+        "The stock market crashed because investors",
+        "Deep learning is a subset of machine learning that",
+        "The president of the United States announced",
+        "In the year 2050, humans will",
+        "The secret to happiness is",
+        "When I was a child, I used to",
+        "The most important scientific discovery of the century",
+        "Climate change is caused by",
+        "The recipe for chocolate cake requires",
+        "In conclusion, the evidence suggests that",
+        "The cat sat on the mat and",
+        "According to recent studies, exercise can",
+        "The first step in solving any problem is",
+        "Technology has transformed the way we",
+        "The novel begins with the protagonist",
+        "Education is the most powerful weapon",
+        "The ocean covers more than seventy percent of",
+        "Last night I had a dream about",
+        "The company announced its quarterly earnings",
+        "Music has the power to",
+        "The difference between success and failure is",
+        "In the beginning, there was nothing but",
+        "The doctor told me that I should",
+        "Python is a popular programming language because",
+        "The ancient Romans built roads that",
+        "A balanced diet should include",
+        "The movie received mixed reviews from critics",
+        "Space exploration has led to many",
+        "The teacher asked the students to",
+        "Global warming is one of the most",
+        "The bridge collapsed due to structural",
+        "Quantum computing promises to revolutionize",
+        "The new policy will affect millions of",
+        "During the winter months, it is important to",
+        "The human brain contains approximately",
+        "Democracy depends on the active participation of",
+        "The train arrived at the station exactly",
+        "Researchers at MIT have developed a new",
+        "The smartphone has become an essential part of",
+        "After careful consideration, the committee decided to",
+    ];
+
+    println!("[");
+    for (i, prompt) in prompts.iter().enumerate() {
+        let input_ids = tokenizer.encode(prompt);
+        let input_len = input_ids.len();
+
+        let mut cache = GpuKVCache::new(&config, 256, DType::BF16);
+
+        // Prefill
+        let t0 = Instant::now();
+        let logits = model.forward_gpu_cache(&input_ids, &mut cache);
+        let first_token = sample_greedy(&logits);
+        let ttft_us = t0.elapsed().as_micros();
+
+        let mut generated = vec![first_token];
+        let mut token_times = Vec::new();
+
+        // Decode
+        for _ in 1..gen_tokens {
+            let last = *generated.last().unwrap();
+            let t_start = Instant::now();
+            let logits = model.forward_gpu_cache(&[last], &mut cache);
+            let next = sample_greedy(&logits);
+            token_times.push(t_start.elapsed().as_micros());
+            generated.push(next);
+            if tokenizer.eos_token_id() == Some(next) { break; }
+        }
+
+        let num_generated = generated.len();
+        let generated_text = tokenizer.decode(&generated);
+        let tbt_us = if !token_times.is_empty() {
+            token_times.iter().sum::<u128>() / token_times.len() as u128
+        } else { 0 };
+        let total_gen_us: u128 = ttft_us + token_times.iter().sum::<u128>();
+        let tpot_us = if num_generated > 0 { total_gen_us / num_generated as u128 } else { 0 };
+
+        let gen_text_escaped = generated_text
+            .replace('\\', "\\\\")
+            .replace('"', "\\\"")
+            .replace('\n', "\\n")
+            .replace('\r', "\\r")
+            .replace('\t', "\\t");
+        let gen_ids_str: Vec<String> = generated.iter().map(|id| id.to_string()).collect();
+
+        print!("  {{\"prompt\": \"{}\", ", prompt.replace('"', "\\\""));
+        print!("\"input_len\": {input_len}, ");
+        print!("\"num_generated\": {num_generated}, ");
+        print!("\"generated_ids\": [{}], ", gen_ids_str.join(", "));
+        print!("\"generated_text\": \"{gen_text_escaped}\", ");
+        print!("\"ttft_us\": {ttft_us}, ");
+        print!("\"tbt_us\": {tbt_us}, ");
+        print!("\"tpot_us\": {tpot_us}}}");
+        if i < prompts.len() - 1 { println!(","); } else { println!(); }
+
+        let display_text = generated_text.replace('\n', " ");
+        let truncated: String = display_text.chars().take(60).collect();
+        eprintln!(
+            "[{}/{}] input={input_len}tok gen={num_generated}tok ttft={:.1}ms tbt={:.1}ms | {}",
+            i + 1, prompts.len(),
+            ttft_us as f64 / 1000.0,
+            tbt_us as f64 / 1000.0,
+            truncated
+        );
+    }
+    println!("]");
+}

crates/xserv-model/src/bin/dump-logits.rs

@@ -0,0 +1,44 @@
+use std::path::PathBuf;
+use xserv_model::{loader, KVCache, ModelConfig, Qwen3};
+use xserv_tensor::{DType, Device};
+use xserv_tokenizer::Tokenizer;
+use half::bf16;
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    let model_dir = PathBuf::from(&args[1]);
+    let prompt = &args[2];
+
+    xserv_cuda::device::set_device(0).unwrap();
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
+    let model = Qwen3::from_weights(config.clone(), weights);
+    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+
+    let token_ids = tokenizer.encode(prompt);
+    eprintln!("Prompt: {prompt}");
+    eprintln!("Token IDs: {token_ids:?}");
+
+    let mut cache = KVCache::new(
+        config.num_layers(), config.num_kv_heads(), config.head_dim(),
+        DType::BF16, Device::Cuda(0),
+    );
+    let logits = model.forward_with_cache(&token_ids, &mut cache);
+    let logits_cpu = logits.to_device(Device::Cpu);
+    let data = logits_cpu.as_slice::<bf16>();
+    let vocab_size = logits.shape()[1];
+    let seq_len = logits.shape()[0];
+
+    // Print top-20 logits for the last position
+    let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
+    let mut indexed: Vec<(usize, f32)> = last_row.iter().enumerate()
+        .map(|(i, v)| (i, v.to_f32()))
+        .collect();
+    indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+
+    println!("Top-20 logits (last position):");
+    for (rank, (id, val)) in indexed.iter().take(20).enumerate() {
+        let tok = tokenizer.decode(&[*id as u32]);
+        println!("  [{rank:>2}] id={id:>6} logit={val:>10.4} token={tok:?}");
+    }
+}

crates/xserv-model/src/bin/xserv-cli.rs

@@ -1,21 +1,19 @@
 use std::io::{self, Write};
 use std::path::PathBuf;
-use xserv_model::{GPT2, ModelConfig};
-use xserv_model::loader;
-use xserv_model::gpt2::sample_greedy;
+use xserv_model::{loader, KVCache, ModelConfig};
+use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
-use xserv_tensor::Device;
 
 fn main() {
     let args: Vec<String> = std::env::args().collect();
     if args.len() < 2 {
         eprintln!("Usage: xserv-cli <model-dir> [--max-tokens N]");
-        eprintln!("  model-dir: path to HF model directory (containing model.safetensors, config.json, tokenizer.json)");
         std::process::exit(1);
     }
 
     let model_dir = PathBuf::from(&args[1]);
-    let max_tokens: usize = args.iter()
+    let max_tokens: usize = args
+        .iter()
         .position(|a| a == "--max-tokens")
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
@@ -25,53 +23,78 @@ fn main() {
     let info = xserv_cuda::device::device_info(0).unwrap();
     eprintln!("GPU: {} ({} MB free)", info.name, info.free_memory / 1024 / 1024);
 
-    // Load config
     let config = ModelConfig::from_file(&model_dir.join("config.json"));
-    eprintln!("Model: {:?}, layers={}, hidden={}, heads={}, vocab={}",
-        config.model_type, config.num_layers(), config.hidden(),
-        config.num_heads(), config.vocab_size);
+    let model_type = config.model_type.as_deref().unwrap_or("unknown");
+    eprintln!(
+        "Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}",
+        config.num_layers(), config.hidden(), config.num_heads(),
+        config.num_kv_heads(), config.vocab_size
+    );
 
-    // Load weights
     eprintln!("Loading weights...");
     let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
     eprintln!("Loaded {} tensors", weights.len());
 
-    // GPT-2 uses weight names without "model." prefix
-    let model = GPT2::from_weights(config, 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))
+    };
 
-    // Load tokenizer
     let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
-    eprintln!("Tokenizer loaded (vocab_size={})", tokenizer.vocab_size());
-    eprintln!("Ready.\n");
+    eprintln!("Ready (KV cache, dtype={dtype}).\n");
 
-    // Interactive loop
     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 mut token_ids = tokenizer.encode(input);
+        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(), kv_heads, config.head_dim(), dtype, Device::Cuda(0),
+        );
+
+        // 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();
 
         for _ in 0..max_tokens {
-            let logits = model.forward(&token_ids);
-            let next = sample_greedy(&logits);
-            token_ids.push(next);
-
             let text = tokenizer.decode(&[next]);
             print!("{text}");
             io::stdout().flush().unwrap();
 
-            if tokenizer.eos_token_id() == Some(next) {
-                break;
-            }
+            if tokenizer.eos_token_id() == Some(next) { break; }
+
+            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!();
     }

crates/xserv-model/src/gpt2.rs

@@ -6,27 +6,112 @@ use crate::config::ModelConfig;
 
 pub struct GPT2 {
     pub config: ModelConfig,
-    wte: Tensor,       // [vocab_size, hidden]
-    wpe: Tensor,       // [max_pos, hidden]
+    wte: Tensor,
+    wpe: Tensor,
     layers: Vec<GPT2Block>,
-    ln_f_g: Tensor,    // [hidden]
-    ln_f_b: Tensor,    // [hidden]
+    ln_f_g: Tensor,
+    ln_f_b: Tensor,
+    lm_head: Tensor, // precomputed wte^T
 }
 
 struct GPT2Block {
     ln_1_g: Tensor,
     ln_1_b: Tensor,
-    // Attention: combined QKV weight + bias, output weight + bias
-    attn_qkv_w: Tensor, // [hidden, 3*hidden]
-    attn_qkv_b: Tensor, // [3*hidden]
-    attn_out_w: Tensor,  // [hidden, hidden]
-    attn_out_b: Tensor,  // [hidden]
+    attn_qkv_w: Tensor,
+    attn_qkv_b: Tensor,
+    attn_out_w: Tensor,
+    attn_out_b: Tensor,
     ln_2_g: Tensor,
     ln_2_b: Tensor,
-    mlp_fc_w: Tensor,   // [hidden, 4*hidden]
-    mlp_fc_b: Tensor,   // [4*hidden]
-    mlp_proj_w: Tensor,  // [4*hidden, hidden]
-    mlp_proj_b: Tensor,  // [hidden]
+    mlp_fc_w: Tensor,
+    mlp_fc_b: Tensor,
+    mlp_proj_w: Tensor,
+    mlp_proj_b: Tensor,
+}
+
+pub struct KVCache {
+    // 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, 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 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 * 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.
+    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 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 * 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_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 {
@@ -39,6 +124,7 @@ impl GPT2 {
         let wpe = take(&mut w, "wpe.weight");
         let ln_f_g = take(&mut w, "ln_f.weight");
         let ln_f_b = take(&mut w, "ln_f.bias");
+        let lm_head = wte.transpose(0, 1).contiguous();
 
         let num_layers = config.num_layers();
         let mut layers = Vec::with_capacity(num_layers);
@@ -60,112 +146,141 @@ impl GPT2 {
             });
         }
 
-        Self { config, wte, wpe, layers, ln_f_g, ln_f_b }
+        Self { config, wte, wpe, layers, ln_f_g, ln_f_b, lm_head }
     }
 
-    /// Full forward pass, returns logits [seq_len, vocab_size].
+    /// Full forward pass without KV cache (for testing / correctness comparison).
     pub fn forward(&self, token_ids: &[u32]) -> Tensor {
         let seq_len = token_ids.len();
         let hidden = self.config.hidden();
         let num_heads = self.config.num_heads();
         let head_dim = self.config.head_dim();
 
-        // Token + position embedding
         let tok_emb = embedding(&self.wte, token_ids);
         let pos_ids: Vec<u32> = (0..seq_len as u32).collect();
         let pos_emb = embedding(&self.wpe, &pos_ids);
         let mut x = add_tensors(&tok_emb, &pos_emb);
 
-        // Transformer layers
         for layer in &self.layers {
-            // Pre-LN attention
-            let residual = x.clone();
-            let normed = layernorm(&x, &layer.ln_1_g, &layer.ln_1_b, self.config.ln_eps());
-
-            // QKV projection: [S, H] @ [H, 3H] + [3H] → [S, 3H]
-            let qkv = linear(&normed, &layer.attn_qkv_w, Some(&layer.attn_qkv_b));
-            // Split into Q, K, V and reshape for multi-head
-            let (q, k, v) = split_qkv(&qkv, num_heads, head_dim, seq_len);
-            // Attention: [1, H, S, D]
-            let attn_out = attention(&q, &k, &v, true);
-            // Merge heads: [1, H, S, D] → [S, hidden]
-            let attn_out = merge_heads(&attn_out, seq_len, hidden);
-            // Output projection
-            let attn_out = linear(&attn_out, &layer.attn_out_w, Some(&layer.attn_out_b));
-            x = add_tensors(&residual, &attn_out);
-
-            // Pre-LN MLP
-            let residual = x.clone();
-            let normed = layernorm(&x, &layer.ln_2_g, &layer.ln_2_b, self.config.ln_eps());
-            let fc = linear(&normed, &layer.mlp_fc_w, Some(&layer.mlp_fc_b));
-            let activated = gelu(&fc);
-            let proj = linear(&activated, &layer.mlp_proj_w, Some(&layer.mlp_proj_b));
-            x = add_tensors(&residual, &proj);
+            x = self.transformer_block(layer, &x, None, 0, seq_len, num_heads, head_dim, hidden);
         }
 
-        // Final layer norm
         let x = layernorm(&x, &self.ln_f_g, &self.ln_f_b, self.config.ln_eps());
+        matmul_2d(&x, &self.lm_head)
+    }
 
-        // LM head (tied with wte): [S, H] @ [H, V] → [S, V]
-        // wte is [V, H], so we need wte^T
-        let lm_head = self.wte.transpose(0, 1).contiguous();
-        matmul_2d(&x, &lm_head)
+    /// Forward pass with KV cache. First call = prefill, subsequent = decode.
+    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 head_dim = self.config.head_dim();
+
+        let tok_emb = embedding(&self.wte, token_ids);
+        let pos_ids: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
+        let pos_emb = embedding(&self.wpe, &pos_ids);
+        let mut x = add_tensors(&tok_emb, &pos_emb);
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            x = self.transformer_block(
+                layer, &x, Some((cache, layer_idx)),
+                pos_offset, new_tokens, num_heads, head_dim, hidden,
+            );
+        }
+
+        let x = layernorm(&x, &self.ln_f_g, &self.ln_f_b, self.config.ln_eps());
+        matmul_2d(&x, &self.lm_head)
+    }
+
+    fn transformer_block(
+        &self,
+        layer: &GPT2Block,
+        x: &Tensor,
+        cache: Option<(&mut KVCache, usize)>,
+        pos_offset: usize,
+        new_tokens: usize,
+        num_heads: usize,
+        head_dim: usize,
+        hidden: usize,
+    ) -> Tensor {
+        let residual = x.clone();
+        let normed = layernorm(x, &layer.ln_1_g, &layer.ln_1_b, self.config.ln_eps());
+
+        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);
+
+        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_tensor(layer_idx, &k_cpu, &v_cpu, new_tokens);
+            cache.get_kv_tensors(layer_idx)
+        } else {
+            (k_new, v_new)
+        };
+
+        let attn_out = attention(&q, &k_full, &v_full, true);
+        let attn_out = merge_heads(&attn_out, new_tokens, hidden);
+        let attn_out = linear(&attn_out, &layer.attn_out_w, Some(&layer.attn_out_b));
+        let x = add_tensors(&residual, &attn_out);
+
+        let residual = x.clone();
+        let normed = layernorm(&x, &layer.ln_2_g, &layer.ln_2_b, self.config.ln_eps());
+        let fc = linear(&normed, &layer.mlp_fc_w, Some(&layer.mlp_fc_b));
+        let activated = gelu(&fc);
+        let proj = linear(&activated, &layer.mlp_proj_w, Some(&layer.mlp_proj_b));
+        add_tensors(&residual, &proj)
     }
 }
 
-// --- Helper ops ---
+// --- Helper ops (unchanged) ---
 
 fn linear(x: &Tensor, weight: &Tensor, bias: Option<&Tensor>) -> Tensor {
-    // GPT-2 stores weights as [in, out] (not transposed), so x @ w
     let out = matmul_2d(x, weight);
-    if let Some(b) = bias {
-        add_bias(&out, b)
-    } else {
-        out
-    }
+    if let Some(b) = bias { add_bias(&out, b) } else { out }
 }
 
 fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
-    // a: [S, K], b: [K, N] → [S, N]
     assert_eq!(a.ndim(), 2);
     assert_eq!(b.ndim(), 2);
     matmul(a, b, GemmBackend::CuBlas)
 }
 
 fn add_tensors(a: &Tensor, b: &Tensor) -> Tensor {
-    // Element-wise add on GPU via a simple approach: scale(a, 1.0) + scale(b, 1.0)
-    // TODO: proper add kernel. For now, go through CPU.
-    assert_eq!(a.shape(), b.shape());
-    assert_eq!(a.dtype(), DType::F32);
-    let a_cpu = a.to_device(Device::Cpu);
-    let b_cpu = b.to_device(Device::Cpu);
-    let a_data = a_cpu.as_slice::<f32>();
-    let b_data = b_cpu.as_slice::<f32>();
-    let sum: Vec<f32> = a_data.iter().zip(b_data).map(|(x, y)| x + y).collect();
-    Tensor::from_slice(&sum, a.shape()).to_device(a.device())
+    xserv_kernels::add(a, b)
 }
 
 fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
-    // x: [S, N], bias: [N] → broadcast add
+    // bias: [N], x: [S, N] — broadcast add via reshape
     assert_eq!(x.ndim(), 2);
     assert_eq!(bias.ndim(), 1);
-    assert_eq!(x.shape()[1], bias.shape()[0]);
-    let x_cpu = x.to_device(Device::Cpu);
-    let b_cpu = bias.to_device(Device::Cpu);
-    let x_data = x_cpu.as_slice::<f32>();
-    let b_data = b_cpu.as_slice::<f32>();
     let n = bias.shape()[0];
-    let result: Vec<f32> = x_data.iter().enumerate().map(|(i, &v)| v + b_data[i % n]).collect();
-    Tensor::from_slice(&result, x.shape()).to_device(x.device())
+    assert_eq!(x.shape()[1], n);
+    let rows = x.shape()[0];
+    // Broadcast: tile bias to [S, N] on CPU, then GPU add
+    let b_cpu = bias.to_device(Device::Cpu);
+    match x.dtype() {
+        DType::F32 => {
+            let bd = b_cpu.as_slice::<f32>();
+            let tiled: Vec<f32> = (0..rows).flat_map(|_| bd.iter().copied()).collect();
+            let b_full = Tensor::from_slice(&tiled, x.shape()).to_device(x.device());
+            xserv_kernels::add(x, &b_full)
+        }
+        DType::BF16 => {
+            let bd = b_cpu.as_slice::<half::bf16>();
+            let tiled: Vec<half::bf16> = (0..rows).flat_map(|_| bd.iter().copied()).collect();
+            let b_full = Tensor::from_slice(&tiled, x.shape()).to_device(x.device());
+            xserv_kernels::add(x, &b_full)
+        }
+        _ => panic!("unsupported dtype"),
+    }
 }
 
 fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) -> (Tensor, Tensor, Tensor) {
-    // qkv: [S, 3*H] → Q, K, V each [1, num_heads, S, head_dim]
     let hidden = num_heads * head_dim;
     let qkv_cpu = qkv.to_device(Device::Cpu);
     let data = qkv_cpu.as_slice::<f32>();
 
-    // Split into Q, K, V and directly write in [1, num_heads, S, head_dim] layout
     let mut q_data = vec![0.0f32; num_heads * seq_len * head_dim];
     let mut k_data = vec![0.0f32; num_heads * seq_len * head_dim];
     let mut v_data = vec![0.0f32; num_heads * seq_len * head_dim];
@@ -189,14 +304,11 @@ fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) ->
 }
 
 fn merge_heads(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
-    // [1, num_heads, S, head_dim] → [S, hidden]
     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::<f32>();
 
-    // src layout: [1][num_heads][seq_len][head_dim]
-    // dst layout: [seq_len][hidden] where hidden = num_heads * head_dim
     let mut out = vec![0.0f32; seq_len * hidden];
     for s in 0..seq_len {
         for h in 0..num_heads {
@@ -210,7 +322,7 @@ fn merge_heads(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
 
 /// Greedy sampling: return the argmax token ID from the last position's logits.
 pub fn sample_greedy(logits: &Tensor) -> u32 {
-    assert_eq!(logits.ndim(), 2); // [S, V]
+    assert_eq!(logits.ndim(), 2);
     let logits_cpu = logits.to_device(Device::Cpu);
     let data = logits_cpu.as_slice::<f32>();
     let vocab_size = logits.shape()[1];

crates/xserv-model/src/kv_cache.rs

@@ -0,0 +1,118 @@
+use xserv_cuda::GpuBuffer;
+use xserv_tensor::{DType, Device, Tensor};
+use crate::config::ModelConfig;
+
+/// GPU-resident KV cache. Pre-allocates max_seq_len on GPU,
+/// appends new K/V via D2D copy at offset (no CPU round-trip).
+pub struct GpuKVCache {
+    // Per layer: contiguous GPU buffer for K and V
+    // Layout: [num_kv_heads, max_seq_len, head_dim] — contiguous per head
+    k_bufs: Vec<GpuBuffer>,
+    v_bufs: Vec<GpuBuffer>,
+    seq_len: usize,
+    max_seq_len: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    elem_size: usize,
+    dtype: DType,
+}
+
+impl GpuKVCache {
+    pub fn new(config: &ModelConfig, max_seq_len: usize, dtype: DType) -> Self {
+        let num_layers = config.num_layers();
+        let num_kv_heads = config.num_kv_heads();
+        let head_dim = config.head_dim();
+        let elem_size = dtype.size_bytes();
+        let buf_size = num_kv_heads * max_seq_len * head_dim * elem_size;
+
+        let mut k_bufs = Vec::with_capacity(num_layers);
+        let mut v_bufs = Vec::with_capacity(num_layers);
+        for _ in 0..num_layers {
+            let mut k = GpuBuffer::alloc(buf_size).expect("alloc KV cache K");
+            let mut v = GpuBuffer::alloc(buf_size).expect("alloc KV cache V");
+            k.zero().unwrap();
+            v.zero().unwrap();
+            k_bufs.push(k);
+            v_bufs.push(v);
+        }
+
+        Self { k_bufs, v_bufs, seq_len: 0, max_seq_len, num_kv_heads, head_dim, elem_size, dtype }
+    }
+
+    pub fn seq_len(&self) -> usize { self.seq_len }
+    pub fn max_seq_len(&self) -> usize { self.max_seq_len }
+
+    /// Append new K/V tensors for a given layer.
+    /// k_new, v_new: [1, num_kv_heads, new_tokens, head_dim] on GPU, contiguous.
+    /// `write_pos` is the sequence position to write at (caller manages this).
+    pub fn append(&mut self, layer: usize, k_new: &Tensor, v_new: &Tensor, new_tokens: usize, write_pos: usize) {
+        assert!(write_pos + new_tokens <= self.max_seq_len, "KV cache overflow");
+        let es = self.elem_size;
+        let hd = self.head_dim;
+        let max_s = self.max_seq_len;
+        let nh = self.num_kv_heads;
+
+        let k_src = k_new.storage().gpu_buffer();
+        let v_src = v_new.storage().gpu_buffer();
+
+        for h in 0..nh {
+            let src_off = h * new_tokens * hd * es;
+            let dst_off = (h * max_s + write_pos) * hd * es;
+            let count = new_tokens * hd * es;
+            self.k_bufs[layer].copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
+            self.v_bufs[layer].copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
+        }
+    }
+
+    pub fn advance_seq_len(&mut self, new_tokens: usize) {
+        self.seq_len += new_tokens;
+    }
+
+    /// Get K/V cache tensors for a layer up to `seq_len` tokens: [1, num_kv_heads, seq_len, head_dim]
+    pub fn get_kv(&self, layer: usize) -> (Tensor, Tensor) {
+        let sl = self.seq_len;
+        self.get_kv_len(layer, sl)
+    }
+
+    pub fn get_kv_len(&self, layer: usize, sl: usize) -> (Tensor, Tensor) {
+        let hd = self.head_dim;
+        let nh = self.num_kv_heads;
+        let es = self.elem_size;
+        let max_s = self.max_seq_len;
+
+        // Allocate output tensors [1, nh, sl, hd]
+        let out_size = nh * sl * hd * es;
+        let mut k_out = GpuBuffer::alloc(out_size).expect("alloc k_out");
+        let mut v_out = GpuBuffer::alloc(out_size).expect("alloc v_out");
+
+        // Copy each head's valid portion
+        for h in 0..nh {
+            let src_off = (h * max_s) * hd * es;
+            let dst_off = (h * sl) * hd * es;
+            let count = sl * hd * es;
+            k_out.copy_from_device_at(&self.k_bufs[layer], src_off, dst_off, count).unwrap();
+            v_out.copy_from_device_at(&self.v_bufs[layer], src_off, dst_off, count).unwrap();
+        }
+
+        let shape = &[1usize, nh, sl, hd];
+        let k = unsafe { tensor_from_gpu_buffer(k_out, shape, self.dtype) };
+        let v = unsafe { tensor_from_gpu_buffer(v_out, shape, self.dtype) };
+        (k, v)
+    }
+}
+
+/// Create a Tensor from a GpuBuffer (takes ownership).
+unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType) -> Tensor {
+    use xserv_tensor::storage::Storage;
+    use xserv_tensor::shape::contiguous_strides;
+    use smallvec::SmallVec;
+
+    let storage = Storage::cuda(buf);
+    Tensor::from_storage(
+        storage,
+        SmallVec::from_slice(shape),
+        contiguous_strides(shape),
+        0,
+        dtype,
+    )
+}

crates/xserv-model/src/lib.rs

@@ -1,6 +1,10 @@
 pub mod config;
 pub mod gpt2;
+pub mod kv_cache;
 pub mod loader;
+pub mod qwen3;
 
 pub use config::ModelConfig;
-pub use gpt2::GPT2;
+pub use gpt2::{GPT2, KVCache};
+pub use kv_cache::GpuKVCache;
+pub use qwen3::Qwen3;

crates/xserv-model/src/qwen3.rs

@@ -0,0 +1,340 @@
+use std::collections::HashMap;
+use half::bf16;
+use xserv_kernels::*;
+use xserv_tensor::{DType, Device, Tensor};
+
+use crate::config::ModelConfig;
+use crate::gpt2::KVCache;
+use crate::kv_cache::GpuKVCache;
+
+pub struct Qwen3 {
+    pub config: ModelConfig,
+    embed_tokens: Tensor,
+    layers: Vec<Qwen3Block>,
+    norm: Tensor,
+    lm_head_t: Tensor,  // precomputed transpose
+    rope_cache: RopeCache,
+}
+
+struct Qwen3Block {
+    input_norm: Tensor,       // [hidden]
+    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_wt: Tensor,     // TRANSPOSED: [hidden, intermediate]
+    up_proj_wt: Tensor,
+    down_proj_wt: Tensor,     // TRANSPOSED: [intermediate, hidden]
+}
+
+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_raw = 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,
+        );
+
+        // 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_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_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"))),
+            });
+        }
+
+        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 {
+        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 (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);
+            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 = 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 = 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 = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            x = add_any(&residual, &down);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_2d(&x, &self.lm_head_t)
+    }
+
+    /// Forward with GPU-resident KV cache and GPU transpose/reshape kernels.
+    pub fn forward_gpu_cache(&self, token_ids: &[u32], cache: &mut GpuKVCache) -> 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);
+
+            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);
+
+            // GPU reshape: [S, H*D] → [1, H, S, D]
+            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);
+            let v = xserv_kernels::reshape_heads_gpu(&v, new_tokens, num_kv_heads, head_dim);
+
+            // QK norm (reshape to [H*S, D], rmsnorm, reshape back — stays on GPU)
+            let q = head_rmsnorm(&q, &layer.q_norm, eps);
+            let k = head_rmsnorm(&k, &layer.k_norm, eps);
+
+            // GPU transpose for RoPE: [1, H, S, D] → [S, H, D]
+            let q = xserv_kernels::transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim);
+            let k = xserv_kernels::transpose_for_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
+            rope_inplace(&q, &self.rope_cache, &positions);
+            rope_inplace(&k, &self.rope_cache, &positions);
+            // GPU transpose back: [S, H, D] → [1, H, S, D]
+            let q = xserv_kernels::transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim);
+            let k = xserv_kernels::transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
+
+            // GPU KV cache
+            cache.append(layer_idx, &k, &v, new_tokens, pos_offset);
+            let (k_full, v_full) = cache.get_kv_len(layer_idx, pos_offset + new_tokens);
+
+            // GPU repeat KV for GQA
+            let n_rep = num_heads / num_kv_heads;
+            let k_full = xserv_kernels::repeat_kv_gpu(&k_full, n_rep);
+            let v_full = xserv_kernels::repeat_kv_gpu(&v_full, n_rep);
+
+            let attn_out = attention(&q, &k_full, &v_full, true);
+            // GPU merge_heads: [1, H, S, D] → [S, H*D]
+            let attn_merged = xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+            x = add_any(&residual, &attn_proj);
+
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.post_norm, eps);
+            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 = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            x = add_any(&residual, &down);
+        }
+
+        cache.advance_seq_len(new_tokens);
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_2d(&x, &self.lm_head_t)
+    }
+}
+
+// --- Helpers ---
+
+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 {
+    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 {
+    xserv_kernels::add(a, b)
+}
+
+fn mul_any(a: &Tensor, b: &Tensor) -> Tensor {
+    xserv_kernels::mul(a, b)
+}
+
+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()
+}

crates/xserv-server/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "xserv-server"
+version.workspace = true
+edition.workspace = true
+
+[[bin]]
+name = "xserv-server"
+path = "src/main.rs"
+
+[dependencies]
+xserv-cuda = { path = "../xserv-cuda" }
+xserv-tensor = { path = "../xserv-tensor" }
+xserv-kernels = { path = "../xserv-kernels" }
+xserv-model = { path = "../xserv-model" }
+xserv-tokenizer = { path = "../xserv-tokenizer" }
+half.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+tokio.workspace = true
+axum.workspace = true
+uuid.workspace = true

crates/xserv-server/src/api.rs

@@ -0,0 +1,115 @@
+use axum::Extension;
+use axum::Json;
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use uuid::Uuid;
+
+use crate::engine::{GenerateEvent, GenerateRequest};
+use crate::AppState;
+
+#[derive(Deserialize)]
+pub struct ChatRequest {
+    #[serde(default)]
+    pub model: Option<String>,
+    pub messages: Vec<Message>,
+    #[serde(default = "default_max_tokens")]
+    pub max_tokens: usize,
+}
+
+#[derive(Deserialize)]
+pub struct Message {
+    pub role: String,
+    pub content: String,
+}
+
+fn default_max_tokens() -> usize { 256 }
+
+#[derive(Serialize)]
+pub struct ModelsResponse {
+    object: &'static str,
+    data: Vec<ModelInfo>,
+}
+
+#[derive(Serialize)]
+pub struct ModelInfo {
+    id: String,
+    object: &'static str,
+    owned_by: &'static str,
+}
+
+pub async fn health() -> &'static str { "ok" }
+
+pub async fn list_models(Extension(state): Extension<Arc<AppState>>) -> Json<ModelsResponse> {
+    Json(ModelsResponse {
+        object: "list",
+        data: vec![ModelInfo {
+            id: state.model_name.clone(),
+            object: "model",
+            owned_by: "xserv",
+        }],
+    })
+}
+
+pub async fn chat_completions(
+    Extension(state): Extension<Arc<AppState>>,
+    Json(req): Json<ChatRequest>,
+) -> Json<serde_json::Value> {
+    let id = format!("chatcmpl-{}", Uuid::new_v4());
+    let model_name = state.model_name.clone();
+    let created = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+
+    // Prepare prompt tokens (MutexGuard scoped)
+    let prompt = build_prompt(&req.messages);
+    let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
+
+    // Create channel and submit request (MutexGuard scoped)
+    let (tx, mut rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
+    let gen_req = GenerateRequest {
+        prompt_tokens,
+        max_tokens: req.max_tokens,
+        sender: tx,
+    };
+    state.engine_sender.lock().unwrap().send(gen_req).expect("engine channel closed");
+
+    // Now await — no MutexGuards held here
+    let mut content = String::new();
+    let mut finish_reason = "length".to_string();
+    while let Some(event) = rx.recv().await {
+        match event {
+            GenerateEvent::Token { text, .. } => content.push_str(&text),
+            GenerateEvent::Done { finish_reason: fr } => { finish_reason = fr; break; }
+        }
+    }
+
+    Json(serde_json::json!({
+        "id": id,
+        "object": "chat.completion",
+        "created": created,
+        "model": model_name,
+        "choices": [{
+            "index": 0,
+            "message": { "role": "assistant", "content": content },
+            "finish_reason": finish_reason,
+        }],
+        "usage": {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 0
+        }
+    }))
+}
+
+fn build_prompt(messages: &[Message]) -> String {
+    let mut prompt = String::new();
+    for msg in messages {
+        match msg.role.as_str() {
+            "system" => { prompt.push_str(&msg.content); prompt.push('\n'); }
+            "user" | "assistant" => { prompt.push_str(&msg.content); }
+            _ => {}
+        }
+    }
+    prompt
+}

crates/xserv-server/src/engine.rs

@@ -0,0 +1,161 @@
+use std::collections::VecDeque;
+use std::path::Path;
+use std::sync::mpsc;
+use xserv_model::{GpuKVCache, ModelConfig, Qwen3};
+use xserv_model::loader;
+use xserv_model::qwen3::sample_greedy;
+use xserv_tensor::{DType, Device};
+use xserv_tokenizer::Tokenizer;
+
+pub struct Engine {
+    model: Qwen3,
+    config: ModelConfig,
+    tokenizer: Tokenizer,
+    max_batch_size: usize,
+    max_seq_len: usize,
+}
+
+pub struct GenerateRequest {
+    pub prompt_tokens: Vec<u32>,
+    pub max_tokens: usize,
+    pub sender: tokio::sync::mpsc::Sender<GenerateEvent>,
+}
+
+pub enum GenerateEvent {
+    Token { id: u32, text: String },
+    Done { finish_reason: String },
+}
+
+struct Sequence {
+    id: u64,
+    prompt_tokens: Vec<u32>,
+    generated_tokens: Vec<u32>,
+    max_tokens: usize,
+    kv_cache: GpuKVCache,
+    sender: tokio::sync::mpsc::Sender<GenerateEvent>,
+    prefilled: bool,
+}
+
+impl Engine {
+    pub fn load(model_dir: &Path, max_batch_size: usize) -> Self {
+        xserv_cuda::device::set_device(0).unwrap();
+        let config = ModelConfig::from_file(&model_dir.join("config.json"));
+        eprintln!("[engine] Loading weights...");
+        let weights = loader::load_model_dir(model_dir, Device::Cuda(0));
+        eprintln!("[engine] Loaded {} tensors", weights.len());
+        let model = Qwen3::from_weights(config.clone(), weights);
+        let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+        let max_seq_len = 256;
+        eprintln!("[engine] Ready (max_batch_size={max_batch_size}, max_seq_len={max_seq_len})");
+        Self { model, config, tokenizer, max_batch_size, max_seq_len }
+    }
+
+    pub fn tokenizer(&self) -> &Tokenizer { &self.tokenizer }
+
+    /// Main scheduler loop. Receives requests from channel, manages concurrent sequences.
+    pub fn run(&self, rx: mpsc::Receiver<GenerateRequest>) {
+        let mut waiting: VecDeque<Sequence> = VecDeque::new();
+        let mut running: Vec<Sequence> = Vec::new();
+        let mut next_id: u64 = 0;
+
+        eprintln!("[scheduler] Listening for requests...");
+
+        loop {
+            // Step 1: Remove finished sequences
+            running.retain(|seq| !is_finished(seq));
+
+            // Step 2: Admit new sequences from waiting queue
+            while running.len() < self.max_batch_size {
+                if let Some(seq) = waiting.pop_front() {
+                    running.push(seq);
+                } else {
+                    break;
+                }
+            }
+
+            // Step 3: If nothing to do, blocking wait for new request
+            if running.is_empty() {
+                match rx.recv() {
+                    Ok(req) => {
+                        let seq = self.make_sequence(req, &mut next_id);
+                        running.push(seq);
+                    }
+                    Err(_) => break, // channel closed
+                }
+            }
+
+            // Step 4: Process one iteration for all running sequences
+            for seq in running.iter_mut() {
+                if !seq.prefilled {
+                    // Prefill
+                    let logits = self.model.forward_gpu_cache(&seq.prompt_tokens, &mut seq.kv_cache);
+                    let next = sample_greedy(&logits);
+                    seq.generated_tokens.push(next);
+                    seq.prefilled = true;
+                    self.emit_token(seq, next);
+                } else {
+                    // Decode one token
+                    let last = *seq.generated_tokens.last().unwrap();
+                    let logits = self.model.forward_gpu_cache(&[last], &mut seq.kv_cache);
+                    let next = sample_greedy(&logits);
+                    seq.generated_tokens.push(next);
+                    self.emit_token(seq, next);
+                }
+            }
+
+            // Step 5: Check for newly arrived requests (non-blocking)
+            loop {
+                match rx.try_recv() {
+                    Ok(req) => {
+                        let seq = self.make_sequence(req, &mut next_id);
+                        waiting.push_back(seq);
+                    }
+                    Err(mpsc::TryRecvError::Empty) => break,
+                    Err(mpsc::TryRecvError::Disconnected) => return,
+                }
+            }
+        }
+    }
+
+    fn make_sequence(&self, req: GenerateRequest, next_id: &mut u64) -> Sequence {
+        let id = *next_id;
+        *next_id += 1;
+        let kv_cache = GpuKVCache::new(&self.config, self.max_seq_len, DType::BF16);
+        Sequence {
+            id,
+            prompt_tokens: req.prompt_tokens,
+            generated_tokens: Vec::new(),
+            max_tokens: req.max_tokens,
+            kv_cache,
+            sender: req.sender,
+            prefilled: false,
+        }
+    }
+
+    fn emit_token(&self, seq: &Sequence, token_id: u32) {
+        let text = self.tokenizer.decode(&[token_id]);
+
+        if self.tokenizer.eos_token_id() == Some(token_id) {
+            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
+            let _ = seq.sender.blocking_send(GenerateEvent::Done {
+                finish_reason: "stop".to_string(),
+            });
+        } else if seq.generated_tokens.len() >= seq.max_tokens {
+            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
+            let _ = seq.sender.blocking_send(GenerateEvent::Done {
+                finish_reason: "length".to_string(),
+            });
+        } else {
+            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
+        }
+    }
+}
+
+fn is_finished(seq: &Sequence) -> bool {
+    if seq.generated_tokens.is_empty() { return false; }
+    let last = *seq.generated_tokens.last().unwrap();
+    if seq.generated_tokens.len() >= seq.max_tokens { return true; }
+    // Check EOS — need tokenizer info. Use a simple heuristic:
+    // If sender is closed (receiver dropped), also consider finished.
+    seq.sender.is_closed() || last == 151645 // Qwen3 EOS token ID (hardcoded for now)
+}

crates/xserv-server/src/main.rs

@@ -0,0 +1,66 @@
+mod api;
+mod engine;
+
+use axum::{routing::{get, post}, Extension, Router};
+use std::path::PathBuf;
+use std::sync::{mpsc, Arc, Mutex};
+use engine::GenerateRequest;
+
+pub struct AppState {
+    pub model_name: String,
+    pub engine_sender: Mutex<mpsc::Sender<GenerateRequest>>,
+    pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
+}
+
+#[tokio::main]
+async fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N]");
+        std::process::exit(1);
+    }
+
+    let model_dir = PathBuf::from(&args[1]);
+    let port: u16 = args.iter()
+        .position(|a| a == "--port")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(8080);
+    let max_batch: usize = args.iter()
+        .position(|a| a == "--max-batch")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(4);
+
+    let model_name = model_dir.file_name()
+        .map(|n| n.to_string_lossy().to_string())
+        .unwrap_or_else(|| "unknown".to_string());
+
+    let tokenizer = xserv_tokenizer::Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+
+    // Unbounded channel: allows multiple requests to queue up
+    let (tx, rx) = mpsc::channel::<GenerateRequest>();
+
+    let model_dir_clone = model_dir.clone();
+    std::thread::spawn(move || {
+        let engine = engine::Engine::load(&model_dir_clone, max_batch);
+        engine.run(rx);
+    });
+
+    let state = Arc::new(AppState {
+        model_name,
+        engine_sender: Mutex::new(tx),
+        engine_tokenizer: Mutex::new(tokenizer),
+    });
+
+    let app = Router::new()
+        .route("/health", get(api::health))
+        .route("/v1/models", get(api::list_models))
+        .route("/v1/chat/completions", post(api::chat_completions))
+        .layer(Extension(state));
+
+    let addr = format!("0.0.0.0:{port}");
+    eprintln!("[server] Listening on {addr} (max_batch={max_batch})");
+    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
+    axum::serve(listener, app).await.unwrap();
+}

crates/xserv-tensor/src/lib.rs

@@ -4,5 +4,6 @@ pub mod storage;
 pub mod tensor;
 
 pub use dtype::{DType, TensorDType};
-pub use storage::Device;
+pub use shape::Dims;
+pub use storage::{Device, Storage};
 pub use tensor::Tensor;

crates/xserv-tensor/src/tensor.rs

@@ -18,6 +18,11 @@ pub struct Tensor {
 impl Tensor {
     // --- Creation ---
 
+    /// Create a tensor from raw components (for advanced use like GPU KV cache).
+    pub fn from_storage(storage: Storage, shape: Dims, strides: Dims, offset: usize, dtype: DType) -> Self {
+        Self { storage, shape, strides, offset, dtype }
+    }
+
     pub fn from_slice<T: TensorDType>(data: &[T], shape: &[usize]) -> Self {
         let numel: usize = shape.iter().product();
         assert_eq!(data.len(), numel, "data length mismatch with shape");

crates/xserv-tokenizer/src/bpe.rs

@@ -8,9 +8,11 @@ pub struct Tokenizer {
     decoder: Vec<Vec<u8>>,
     merge_ranks: HashMap<(u32, u32), usize>,
     special_tokens: HashMap<String, u32>,
+    #[allow(dead_code)]
     special_token_ids: HashMap<u32, String>,
     pre_tokenize_re: Regex,
     eos_token_id: Option<u32>,
+    byte_fallback: bool,
 }
 
 #[derive(Deserialize)]
@@ -23,7 +25,16 @@ struct TokenizerJson {
 #[derive(Deserialize)]
 struct ModelSection {
     vocab: HashMap<String, u32>,
-    merges: Vec<String>,
+    merges: Vec<MergeEntry>,
+    #[serde(default)]
+    byte_fallback: bool,
+}
+
+#[derive(Deserialize)]
+#[serde(untagged)]
+enum MergeEntry {
+    Str(String),
+    Pair(Vec<String>),
 }
 
 #[derive(Deserialize)]
@@ -40,7 +51,10 @@ impl Tokenizer {
         let tj: TokenizerJson = serde_json::from_str(&data)
             .unwrap_or_else(|e| panic!("failed to parse tokenizer.json: {e}"));
 
+        let byte_fallback = tj.model.byte_fallback;
+
         // Build encoder: token bytes → ID
+        // All HF tokenizers use GPT-2 byte-to-unicode mapping for vocab keys.
         let mut encoder = HashMap::new();
         for (token_str, &id) in &tj.model.vocab {
             let bytes = token_str_to_bytes(token_str);
@@ -56,13 +70,23 @@ impl Tokenizer {
             decoder[id as usize] = token_str_to_bytes(token_str);
         }
 
-        // Parse merges
+        // Parse merges (supports both "a b" string format and ["a", "b"] array format)
+        let byte_fallback = tj.model.byte_fallback;
         let mut merge_ranks = HashMap::new();
-        for (rank, merge_line) in tj.model.merges.iter().enumerate() {
-            let parts: Vec<&str> = merge_line.splitn(2, ' ').collect();
-            if parts.len() != 2 { continue; }
-            let a_bytes = token_str_to_bytes(parts[0]);
-            let b_bytes = token_str_to_bytes(parts[1]);
+        for (rank, entry) in tj.model.merges.iter().enumerate() {
+            let (a_str, b_str) = match entry {
+                MergeEntry::Str(s) => {
+                    let parts: Vec<&str> = s.splitn(2, ' ').collect();
+                    if parts.len() != 2 { continue; }
+                    (parts[0].to_string(), parts[1].to_string())
+                }
+                MergeEntry::Pair(v) => {
+                    if v.len() != 2 { continue; }
+                    (v[0].clone(), v[1].clone())
+                }
+            };
+            let a_bytes = token_str_to_bytes(&a_str);
+            let b_bytes = token_str_to_bytes(&b_str);
             if let (Some(&a_id), Some(&b_id)) = (encoder.get(&a_bytes), encoder.get(&b_bytes)) {
                 merge_ranks.insert((a_id, b_id), rank);
             }
@@ -84,13 +108,14 @@ impl Tokenizer {
             }
         }
 
-        // GPT-2 pre-tokenization regex.
-        // The original uses (?!\S) lookahead which Rust regex doesn't support.
-        // Simplified: collapse trailing whitespace into one match. Functionally equivalent
-        // for BPE since each whitespace chunk gets encoded independently anyway.
-        let pre_tokenize_re = Regex::new(
-            r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+"
-        ).unwrap();
+        // Pre-tokenization regex
+        let pre_tokenize_re = if byte_fallback {
+            // Qwen-style: split on whitespace boundaries, keep Unicode words/numbers
+            Regex::new(r"[\p{L}\p{N}]+|[^\s\p{L}\p{N}]|\s+").unwrap()
+        } else {
+            // GPT-2 style
+            Regex::new(r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+").unwrap()
+        };
 
         Self {
             encoder,
@@ -100,6 +125,7 @@ impl Tokenizer {
             special_token_ids,
             pre_tokenize_re,
             eos_token_id,
+            byte_fallback,
         }
     }
 
@@ -137,10 +163,16 @@ impl Tokenizer {
     fn encode_ordinary(&self, text: &str, out: &mut Vec<u32>) {
         for mat in self.pre_tokenize_re.find_iter(text) {
             let word = mat.as_str();
+            // Try to encode the whole word first
+            if let Some(&id) = self.encoder.get(word.as_bytes()) {
+                out.push(id);
+                continue;
+            }
+            // Fall back to per-byte encoding
             let word_bytes: Vec<u8> = word.bytes().collect();
             let mut token_ids: Vec<u32> = word_bytes.iter().map(|&b| {
                 *self.encoder.get(&vec![b]).unwrap_or_else(|| {
-                    panic!("byte {b} not in vocab")
+                    panic!("byte {b} (0x{b:02X}) not in vocab")
                 })
             }).collect();
 
@@ -204,48 +236,32 @@ fn token_str_to_bytes(s: &str) -> Vec<u8> {
     s.chars().map(|c| unicode_to_byte(c)).collect()
 }
 
+/// Convert a Unicode char back to the byte it represents in GPT-2 encoding.
 fn unicode_to_byte(c: char) -> u8 {
-    let u = c as u32;
-    // GPT-2 byte encoder: maps bytes 0-255 to specific Unicode code points.
-    // Printable ASCII bytes map to themselves. Others are shifted to 256+.
-    match u {
-        0x21..=0x7E => u as u8,        // '!' to '~'
-        0xA1..=0xAC => u as u8,        // '¡' to '¬'
-        0xAE..=0xFF => u as u8,        // '®' to 'ÿ'
-        // Shifted bytes: 0x100 + original_byte for bytes not in the above ranges
-        0x100..=0x1FF => (u - 0x100) as u8 + {
-            // The shift mapping: byte values 0..=32, 127..=160, 173
-            // are shifted to 256..=288, 289+, etc.
-            0
-        },
-        _ => {
-            // Fallback: for the GPT-2 byte encoder, specific mappings
-            byte_from_unicode_gpt2(c)
+    // Build the inverse map on first use
+    use std::sync::OnceLock;
+    static INV_MAP: OnceLock<HashMap<u32, u8>> = OnceLock::new();
+
+    let map = INV_MAP.get_or_init(|| {
+        let mut m = HashMap::new();
+        // Build GPT-2's bytes_to_unicode forward map, then invert
+        let mut n = 0u32;
+        for b in 0..=255u16 {
+            let byte = b as u8;
+            let unicode = match byte {
+                0x21..=0x7E | 0xA1..=0xAC | 0xAE..=0xFF => byte as u32,
+                _ => {
+                    let u = 256 + n;
+                    n += 1;
+                    u
+                }
+            };
+            m.insert(unicode, byte);
         }
-    }
-}
-
-fn byte_from_unicode_gpt2(c: char) -> u8 {
-    // Build the inverse of GPT-2's bytes_to_unicode mapping.
-    // The mapping assigns printable chars to themselves and shifts unprintable bytes.
-    let u = c as u32;
-    // Direct ASCII printable + Latin-1 supplement printable ranges map identity
-    if (0x21..=0x7E).contains(&u) { return u as u8; }
-    if (0xA1..=0xAC).contains(&u) { return u as u8; }
-    if (0xAE..=0xFF).contains(&u) { return u as u8; }
-
-    // Shifted range: the remaining 68 bytes (0-32, 127-160, 173) get mapped to 256..=323
-    static SHIFTED_BYTES: &[u8] = &[
-        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
-        24, 25, 26, 27, 28, 29, 30, 31, 32, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136,
-        137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
-        154, 155, 156, 157, 158, 159, 160, 173,
-    ];
-    let shifted_start = 256u32;
-    if u >= shifted_start && u < shifted_start + SHIFTED_BYTES.len() as u32 {
-        return SHIFTED_BYTES[(u - shifted_start) as usize];
-    }
-
-    // Shouldn't reach here for valid GPT-2 tokenizer
-    c as u8
+        m
+    });
+
+    *map.get(&(c as u32)).unwrap_or_else(|| {
+        panic!("unmapped unicode char U+{:04X} in tokenizer", c as u32)
+    })
 }

csrc/activation/activations.cu

@@ -45,6 +45,26 @@ __global__ void scale_bf16_kernel(const __nv_bfloat16* x, __nv_bfloat16* out, fl
     if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(x[idx]) * scale);
 }
 
+// Element-wise add: out = a + b
+__global__ void add_f32_kernel(const float* a, const float* b, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = a[idx] + b[idx];
+}
+__global__ void add_bf16_kernel(const __nv_bfloat16* a, const __nv_bfloat16* b, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(a[idx]) + __bfloat162float(b[idx]));
+}
+
+// Element-wise mul: out = a * b
+__global__ void mul_f32_kernel(const float* a, const float* b, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = a[idx] * b[idx];
+}
+__global__ void mul_bf16_kernel(const __nv_bfloat16* a, const __nv_bfloat16* b, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(a[idx]) * __bfloat162float(b[idx]));
+}
+
 extern "C" {
 
 void launch_gelu_f32(const void* x, void* out, int n, void* stream) {
@@ -87,4 +107,29 @@ void launch_scale_bf16(const void* x, void* out, float scale, int n, void* strea
         (const __nv_bfloat16*)x, (__nv_bfloat16*)out, scale, n);
 }
 
+void launch_add_f32(const void* a, const void* b, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    add_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)a, (const float*)b, (float*)out, n);
+}
+void launch_add_bf16(const void* a, const void* b, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    add_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
+}
+void launch_mul_f32(const void* a, const void* b, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    mul_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)a, (const float*)b, (float*)out, n);
+}
+void launch_mul_bf16(const void* a, const void* b, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    mul_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
+}
+
 }

csrc/embedding/transpose.cu

@@ -0,0 +1,161 @@
+#include <cuda_bf16.h>
+
+// Transpose between [S, H, D] and [H, S, D] layouts (used for RoPE and attention).
+// Also handles [S, H*D] → [H, S, D] (reshape_heads) and reverse (merge_heads).
+
+// reshape_heads: [S, H*D] → [1, H, S, D]
+// Input layout:  element at [s, h*D + d] = flat[s * H*D + h*D + d]
+// Output layout: element at [0, h, s, d] = flat[h * S*D + s*D + d]
+__global__ void reshape_heads_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len, int num_heads, int head_dim
+) {
+    int hidden = num_heads * head_dim;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = seq_len * hidden;
+    if (idx >= total) return;
+
+    int s = idx / hidden;
+    int rem = idx % hidden;
+    int h = rem / head_dim;
+    int d = rem % head_dim;
+
+    int out_idx = h * seq_len * head_dim + s * head_dim + d;
+    out[out_idx] = in[idx];
+}
+
+// merge_heads: [1, H, S, D] → [S, H*D]
+// Input layout:  element at [0, h, s, d] = flat[h * S*D + s*D + d]
+// Output layout: element at [s, h*D + d] = flat[s * H*D + h*D + d]
+__global__ void merge_heads_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len, int num_heads, int head_dim
+) {
+    int hidden = num_heads * head_dim;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = seq_len * hidden;
+    if (idx >= total) return;
+
+    // idx is output index: [s, h*D + d]
+    int s = idx / hidden;
+    int rem = idx % hidden;
+    int h = rem / head_dim;
+    int d = rem % head_dim;
+
+    int in_idx = h * seq_len * head_dim + s * head_dim + d;
+    out[idx] = in[in_idx];
+}
+
+// transpose_for_rope: [1, H, S, D] → [S, H, D]
+// Input:  [h, s, d] at h*S*D + s*D + d
+// Output: [s, h, d] at s*H*D + h*D + d
+__global__ void transpose_hsd_to_shd_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len, int num_heads, int head_dim
+) {
+    int total = seq_len * num_heads * head_dim;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total) return;
+
+    // idx = output flat index: s*H*D + h*D + d
+    int s = idx / (num_heads * head_dim);
+    int rem = idx % (num_heads * head_dim);
+    int h = rem / head_dim;
+    int d = rem % head_dim;
+
+    int in_idx = h * seq_len * head_dim + s * head_dim + d;
+    out[idx] = in[in_idx];
+}
+
+// transpose_from_rope: [S, H, D] → [1, H, S, D]
+// Input:  [s, h, d] at s*H*D + h*D + d
+// Output: [h, s, d] at h*S*D + s*D + d
+__global__ void transpose_shd_to_hsd_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len, int num_heads, int head_dim
+) {
+    int total = seq_len * num_heads * head_dim;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total) return;
+
+    // idx = output flat index: h*S*D + s*D + d
+    int h = idx / (seq_len * head_dim);
+    int rem = idx % (seq_len * head_dim);
+    int s = rem / head_dim;
+    int d = rem % head_dim;
+
+    int in_idx = s * num_heads * head_dim + h * head_dim + d;
+    out[idx] = in[in_idx];
+}
+
+// repeat_kv: [1, KV_H, S, D] → [1, KV_H * n_rep, S, D]
+__global__ void repeat_kv_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int kv_heads, int n_rep, int seq_len, int head_dim
+) {
+    int total_heads = kv_heads * n_rep;
+    int total = total_heads * seq_len * head_dim;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total) return;
+
+    int out_h = idx / (seq_len * head_dim);
+    int rem = idx % (seq_len * head_dim);
+    int kv_h = out_h / n_rep;
+
+    int in_idx = kv_h * seq_len * head_dim + rem;
+    out[idx] = in[in_idx];
+}
+
+extern "C" {
+
+void launch_reshape_heads_bf16(const void* in, void* out,
+                               int seq_len, int num_heads, int head_dim, void* stream) {
+    int total = seq_len * num_heads * head_dim;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    reshape_heads_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+}
+
+void launch_merge_heads_bf16(const void* in, void* out,
+                             int seq_len, int num_heads, int head_dim, void* stream) {
+    int total = seq_len * num_heads * head_dim;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    merge_heads_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+}
+
+void launch_transpose_hsd_to_shd_bf16(const void* in, void* out,
+                                       int seq_len, int num_heads, int head_dim, void* stream) {
+    int total = seq_len * num_heads * head_dim;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    transpose_hsd_to_shd_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+}
+
+void launch_transpose_shd_to_hsd_bf16(const void* in, void* out,
+                                       int seq_len, int num_heads, int head_dim, void* stream) {
+    int total = seq_len * num_heads * head_dim;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    transpose_shd_to_hsd_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+}
+
+void launch_repeat_kv_bf16(const void* in, void* out,
+                           int kv_heads, int n_rep, int seq_len, int head_dim, void* stream) {
+    int total = kv_heads * n_rep * seq_len * head_dim;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    repeat_kv_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, kv_heads, n_rep, seq_len, head_dim);
+}
+
+}

docs/09-kv-cache.md

@@ -0,0 +1,67 @@
+# Phase 9: KV Cache + Autoregressive Generation — Design Document
+
+## Goal
+
+实现 KV Cache，将 decode 从每步 full forward (O(S²)) 降为增量计算 (O(S))。这是最大的单点性能提升。
+
+## 核心变化
+
+### Before (no cache)
+```
+每生成一个 token:
+  forward(all_tokens)  → 重新计算所有层的 Q/K/V/attention
+  开销: O(S²) attention per step, S 递增
+```
+
+### After (with cache)
+```
+Prefill:
+  forward(prompt_tokens)  → 计算并缓存所有层的 K/V
+  
+Decode (per token):
+  forward(last_token_only) → 只计算新 token 的 Q/K/V
+  Q: [1, H, 1, D]          → 新 token 的 query
+  K: append to cache        → cache 变为 [1, H, S+1, D]
+  V: append to cache
+  attention: Q @ K_cache^T  → [1, H, 1, S+1], O(S) not O(S²)
+```
+
+## KVCache 数据结构
+
+```rust
+pub struct KVCache {
+    k: Vec<Tensor>,  // per layer, shape [1, num_heads, current_len, head_dim]
+    v: Vec<Tensor>,
+    len: usize,      // current sequence length
+}
+```
+
+## Forward Pass 变化
+
+模型需要两种 forward 模式:
+1. **prefill(tokens)**: 处理完整 prompt，填充 KV cache
+2. **decode(token, cache)**: 处理单个 token，读写 KV cache
+
+## 实现策略
+
+为了最小化改动，在 GPT-2 forward 中加入可选的 `&mut KVCache` 参数：
+- cache=None → 现有行为（full forward）
+- cache=Some → prefill 或 decode 模式
+
+CPU round-trip 问题暂不修复（Phase 15），先让 KV cache 逻辑正确。
+
+## Test Plan
+
+- [x] KV cache vs no-cache: 50/50 bit-identical output
+- [x] Benchmark: 18x decode speedup (407ms → 22ms TBT)
+- [x] 50 prompt validation: 40/50 vs HF (10 are FP divergence, gap 0.04-0.56)
+
+## Takeaways
+
+1. **KV cache 数据布局是核心难点**：初始实现直接 append flat bytes 导致 head 维度交错错误。正确做法：per-head 独立存储，reconstruct 时按 `[1, H, S, D]` layout 组装。这是一个非常容易犯的 layout bug，调试时输出看起来"几乎对"但不完全对。
+
+2. **18x 提速 > 理论预期**：理论上 KV cache 将 decode 从 O(S²) 降到 O(S)，对 S=20-25 的序列预期 ~20x 提速。实测 18x 符合预期。TTFT 也从 400ms 降到 24ms，因为 prefill 只跑一次而不是每步重跑。
+
+3. **xserv vs HF 的 10 个 mismatch 不是 bug**：logit gap 仅 0.04-0.56（在 -80 到 -140 的 logit 值上），是不同 CUDA kernel 实现间的浮点累积误差导致 argmax 翻转。重要验证：**xserv KV-cache vs xserv no-cache 是 50/50 完全一致的**——证明 KV cache 实现本身无误。
+
+4. **CPU round-trip 仍是主要瓶颈**：KV cache 的 per-head 数据存在 CPU Vec 中，每步 decode 都要重新组装成 GPU tensor。这意味着每步仍有 24 次 GPU→CPU→GPU 传输（12 层 × 2 KV）。Phase 15 需要将 KV cache 直接放在 GPU 上。

docs/10-qwen3.md

@@ -0,0 +1,109 @@
+# Phase 10: Qwen3-7B Support — Design Document (Milestone ②)
+
+## Goal
+
+扩展模型定义支持 Qwen3-7B 架构，验证输出正确性。与 GPT-2 的关键差异：RMSNorm、RoPE、GQA、SwiGLU、不共享 embedding。
+
+## 架构差异 (GPT-2 → Qwen3)
+
+| 特性 | GPT-2 | Qwen3-7B |
+|------|-------|----------|
+| Norm | LayerNorm(gamma, beta) | RMSNorm(gamma only) |
+| Position | Learned absolute (wpe) | RoPE (no params) |
+| Attention | MHA (12 Q = 12 KV heads) | GQA (32 Q, 8 KV heads) |
+| QKV projection | Combined c_attn [H, 3H] | Separate q/k/v_proj [H, Hq/Hk/Hv] |
+| FFN | 2 Linear (fc, proj) + GELU | 3 Linear (gate, up, down) + SwiGLU |
+| Weight layout | [in, out] (Conv1D style) | [out, in] (standard Linear) |
+| Tied embeddings | Yes | No (separate lm_head) |
+| hidden_size | 768 | 3584 |
+| num_layers | 12 | 28 |
+| head_dim | 64 | 128 |
+
+## Weight Names (HuggingFace)
+
+```
+model.embed_tokens.weight                           [151936, 3584]
+model.layers.{i}.input_layernorm.weight             [3584]
+model.layers.{i}.self_attn.q_proj.weight            [3584, 3584]    (32 heads × 112 dim? or 28 heads)
+model.layers.{i}.self_attn.q_proj.bias              [3584]
+model.layers.{i}.self_attn.k_proj.weight            [512, 3584]     (4 KV heads × 128 dim)
+model.layers.{i}.self_attn.k_proj.bias              [512]
+model.layers.{i}.self_attn.v_proj.weight            [512, 3584]
+model.layers.{i}.self_attn.v_proj.bias              [512]
+model.layers.{i}.self_attn.o_proj.weight            [3584, 3584]
+model.layers.{i}.post_attention_layernorm.weight    [3584]
+model.layers.{i}.mlp.gate_proj.weight               [18944, 3584]
+model.layers.{i}.mlp.up_proj.weight                 [18944, 3584]
+model.layers.{i}.mlp.down_proj.weight               [3584, 18944]
+model.norm.weight                                    [3584]
+lm_head.weight                                      [151936, 3584]
+```
+
+**注意**: Qwen3 权重是 [out, in] layout，`x @ W^T` 而不是 `x @ W`。
+
+## GQA (Grouped Query Attention)
+
+```
+num_heads = 28, num_kv_heads = 4, head_dim = 128
+Q: [B, 28, S, 128]
+K: [B, 4, S, 128]   ← 每个 KV head 服务 28/4 = 7 个 Q head
+V: [B, 4, S, 128]
+
+attention 时需要 repeat K/V:
+K_expanded: [B, 28, S, 128]  ← repeat_interleave(K, 7, dim=1)
+```
+
+实现：在 CPU 侧 split_qkv 时直接做 repeat。
+
+## SwiGLU FFN
+
+```
+gate = gate_proj(x)     # [S, 3584] @ [3584, 18944]^T → [S, 18944]
+up   = up_proj(x)       # [S, 3584] @ [3584, 18944]^T → [S, 18944]
+out  = silu(gate) * up   # element-wise
+out  = down_proj(out)    # [S, 18944] @ [18944, 3584]^T → [S, 3584]
+```
+
+## 显存预算 (BF16, 单卡 5090)
+
+```
+权重: 7B × 2B = ~14 GB (BF16)
+        7B × 4B = ~28 GB (FP32) — 不够! 必须用 BF16
+KV cache (S=256, B=1): ~0.1 GB
+总计: ~14 GB (BF16), 单卡可运行
+```
+
+**关键**: Qwen3-7B 必须用 BF16 才能在单张 5090 (32GB) 上运行。当前 GPT-2 用 FP32，需要支持 BF16 forward pass。
+
+## Implementation Plan
+
+1. 下载 Qwen3-7B 模型 (BF16, ~14GB)
+2. 实现 Qwen3 模型结构 (qwen3.rs)
+3. 支持 BF16 forward pass (linear_transpose for [out, in] weights)
+4. 实现 GQA (K/V repeat in split)
+5. 集成 RoPE + RMSNorm + SwiGLU
+6. 验证输出
+
+## Test Plan
+
+- [x] 加载 Qwen3-8B BF16 权重 (399 tensors, ~15.5GB) 到单张 5090
+- [x] 英文: "The meaning of life is" → "to be happy"
+- [x] 中文: "请用中文回答：1+1等于几？" → "1加1"
+- [x] 61/61 单元测试无回归
+- [x] GPT-2 benchmark 性能无回归
+
+## Takeaways
+
+1. **Qwen3 实际是 8B，不是 7B**：modelscope 上的 `Qwen/Qwen3-8B` 有 36 层 × hidden 4096 × 32 heads，参数量约 8B。BF16 权重 ~15.5GB，单张 5090 (32GB) 可以运行。
+
+2. **QK Normalization 是 Qwen3 的新特性**：每层有 `q_norm` 和 `k_norm` (shape [head_dim])，对 Q 和 K 做 per-head RMSNorm。这在 attention score 的数值稳定性上很重要——没有 QK norm 会导致 attention score 爆炸。
+
+3. **attention_bias=false**：Qwen3 的 Q/K/V/O projection 没有 bias。这和 GPT-2 (有 bias) 不同。需要在模型代码中条件处理。
+
+4. **Tokenizer 的 byte-to-unicode 映射 bug**：GPT-2 和 Qwen3 都使用同一套 byte-to-unicode 映射（printable ASCII identity，其余 68 bytes shifted to U+0100+）。初始实现中 `unicode_to_byte` 的 shifted 范围转换错误（直接 `u - 0x100` 而非查表），导致中文输入时 UTF-8 bytes 无法正确映射。修复：用 `OnceLock` 缓存反向映射表。
+
+5. **Weight layout [out, in] vs [in, out]**：GPT-2 的 Conv1D 存为 [in, out]，计算 `x @ W`；Qwen3 的 Linear 存为 [out, in]，计算 `x @ W^T`。`linear_t` 函数通过 `weight.transpose(0,1).contiguous()` 处理。
+
+6. **RoPE 的 tensor layout 不匹配**：RoPE kernel 期望 [S, H, D]，但 attention 需要 [1, H, S, D]。需要在 RoPE 前后做 transpose。这引入了额外的 CPU round-trip（因为 transpose+contiguous 经过 CPU）。
+
+7. **GQA repeat_kv 的实现**：每个 KV head 服务 `num_heads/num_kv_heads` 个 Q head。在 CPU 上做数据复制（repeat），简单但每步 decode 都要做。后续应在 attention kernel 中直接支持 GQA 索引，避免数据复制。

docs/11-paged-attention.md

@@ -0,0 +1,59 @@
+# Phase 11: Paged Attention + KV Cache Manager — Design Document
+
+## Goal
+
+将 KV cache 从 CPU Vec 迁移到 GPU，使用 block-based paging 管理显存。消除每步 decode 的 CPU round-trip（当前 KV cache 最大性能瓶颈之一）。
+
+## 当前问题
+
+每步 decode 的 KV cache 路径：
+```
+GPU tensor (K_new) → CPU (per-head Vec append) → reconstruct → CPU tensor → GPU tensor
+```
+这涉及 2 次 GPU↔CPU 拷贝 × 36 层 × 2(K,V) = 144 次 transfer/token。
+
+## 目标设计
+
+KV cache 直接存在 GPU 上，decode 时只做 GPU→GPU append：
+```
+GPU tensor (K_new) → GPU KV cache (in-place append, no CPU)
+```
+
+## 实现方案
+
+### GPU KV Cache（简化版，非 paged）
+
+先实现连续分配的 GPU KV cache（预分配 max_seq_len），消除 CPU round-trip。Paged allocation 留待后续优化。
+
+```rust
+pub struct GpuKVCache {
+    // 预分配: [num_layers, 2, num_kv_heads, max_seq_len, head_dim] on GPU
+    k_caches: Vec<Tensor>,  // per layer: [1, num_kv_heads, max_seq_len, head_dim]
+    v_caches: Vec<Tensor>,
+    seq_len: usize,         // 当前已填充的长度
+    max_seq_len: usize,
+}
+```
+
+### Append 操作
+
+用 cudaMemcpy D2D 将新 K/V 写入 cache 的正确偏移位置：
+```
+k_cache[layer][0, :, seq_len:seq_len+new, :] = k_new[0, :, :, :]
+```
+
+### 读取操作
+
+不需要拷贝——直接用 view/slice 返回 [0, :, 0:seq_len, :] 的 GPU tensor。
+
+## 需要的新功能
+
+1. Tensor slice 支持（view into sub-range of a dimension）
+2. GPU D2D copy at offset（写入 cache 指定位置）
+3. 去掉 Qwen3/GPT-2 forward 中的 CPU round-trip KV cache 路径
+
+## Test Plan
+
+- [ ] GPU KV cache 输出与 CPU KV cache bit-identical
+- [ ] Benchmark: TBT 应显著降低（消除 144 次 CPU round-trip）
+- [ ] 50-prompt correctness re-validation

docs/12-continuous-batching.md

@@ -0,0 +1,153 @@
+# Phase 12: Continuous Batching + Request Scheduler — Design Document
+
+## Goal
+
+实现 iteration-level 请求调度，支持多个请求并发生成 token。核心能力：同时发 N 个请求，N 个请求同时产出 token，新请求可以在 mid-generation 加入 batch。
+
+## 为什么需要 Continuous Batching
+
+**当前问题（串行）**：
+```
+时间 →  [req1 prefill][req1 decode x 100][req2 prefill][req2 decode x 50]...
+GPU利用: ████████████████████████████████████████████████████████████████████
+         req2 等了 100 个 token 的时间才开始
+```
+
+**目标（continuous batching）**：
+```
+时间 →  [req1+req2 prefill][req1+req2 decode][req1 done, req3 加入][req2+req3 decode]...
+GPU利用: ████████████████████████████████████████████████████████████████████
+         req2 和 req1 同时推理，req3 在 req1 完成后立即加入
+```
+
+## 核心设计
+
+### 数据结构
+
+```rust
+pub struct Sequence {
+    pub id: u64,
+    pub prompt_tokens: Vec<u32>,
+    pub generated_tokens: Vec<u32>,
+    pub status: SeqStatus,
+    pub max_tokens: usize,
+    pub kv_cache: GpuKVCache,          // 每个 seq 独立的 KV cache
+    pub output_tx: mpsc::Sender<GenerateEvent>,
+}
+
+pub enum SeqStatus {
+    Waiting,     // 在队列中等待被 admit
+    Running,     // 正在参与 batch forward
+    Finished,    // EOS 或 max_tokens 达到
+}
+
+pub struct Scheduler {
+    waiting: VecDeque<Sequence>,
+    running: Vec<Sequence>,
+    max_batch_size: usize,             // 最大并发请求数
+    next_seq_id: u64,
+}
+```
+
+### 调度循环（Engine 主循环）
+
+```rust
+loop {
+    // Step 1: 回收已完成的 sequence
+    running.retain(|seq| seq.status != Finished);
+
+    // Step 2: Admit 新请求（如果 running < max_batch_size）
+    while running.len() < max_batch_size {
+        if let Some(seq) = waiting.pop_front() {
+            running.push(seq);
+        } else {
+            break;
+        }
+    }
+
+    if running.is_empty() {
+        // 没有任何工作，等待新请求
+        let new_req = request_rx.recv();  // blocking wait
+        waiting.push_back(new_req);
+        continue;
+    }
+
+    // Step 3: 分类 — 哪些需要 prefill，哪些需要 decode
+    let to_prefill: 新加入的 seq（generated_tokens 为空）
+    let to_decode: 已在运行的 seq
+
+    // Step 4: 执行
+    for seq in to_prefill {
+        // Prefill: 完整 prompt 一次 forward
+        model.forward_gpu_cache(&seq.prompt_tokens, &mut seq.kv_cache);
+        seq.status = Running;
+    }
+
+    // Decode: 每个 seq 独立做一步（当前不做 batch forward，留待优化）
+    for seq in to_decode {
+        let last_token = seq.last_generated_token();
+        let logits = model.forward_gpu_cache(&[last_token], &mut seq.kv_cache);
+        let next = sample_greedy(&logits);
+        seq.generated_tokens.push(next);
+        // 发送 token 给客户端
+        seq.output_tx.blocking_send(Token { id: next, text: decode(next) });
+        // 检查完成
+        if next == eos || seq.generated_tokens.len() >= seq.max_tokens {
+            seq.output_tx.blocking_send(Done);
+            seq.status = Finished;
+        }
+    }
+
+    // Step 5: 检查是否有新请求到达（non-blocking）
+    while let Ok(new_req) = request_rx.try_recv() {
+        waiting.push_back(new_req);
+    }
+}
+```
+
+### 关键设计决策
+
+1. **每个 seq 独立 KV cache**：当前不做 batch forward（需要对齐 seq_len），而是每个 seq 独立调用 model.forward_gpu_cache。未来优化为 batched forward。
+
+2. **Prefill 和 Decode 混合**：新加入的 seq 先 prefill（一次 forward），然后下一轮加入 decode batch。
+
+3. **Non-blocking request receive**：decode 循环中用 `try_recv()` 检查新请求，不阻塞推理。
+
+4. **max_batch_size**：受限于 GPU 显存（每个 seq 的 KV cache 占用）。Qwen3-8B 单卡 32GB，每个 seq 的 KV cache 约 256 tokens × 8 heads × 128 dim × 2(KV) × 2B = 1MB。可以并发 ~100 seq。实际受限于推理速度。
+
+## 与 Phase 13 (HTTP API) 的接口
+
+```
+HTTP Handler                              Engine Thread
+     │                                          │
+     │  ──── GenerateRequest ────────►          │
+     │       (prompt_tokens, max_tokens,        │
+     │        output_tx)                        │
+     │                                          │
+     │  ◄──── GenerateEvent (Token/Done) ────   │
+     │       (via tokio::sync::mpsc)            │
+     │                                          │
+```
+
+多个 HTTP handler 可以同时提交请求。Engine 线程内部通过 Scheduler 管理并发。
+
+## 验收测试
+
+必须通过以下测试才算 Phase 12 完成：
+
+1. **并发 3 请求测试**：同时发 3 个请求，验证 3 个请求同时产出 token（不是串行等待）
+2. **吞吐量测试**：并发请求的总 token 吞吐量应接近单请求（因为单个 seq 的 decode 是串行的）
+3. **动态加入测试**：先发 1 个请求开始生成，过 2 秒再发第 2 个，验证第 2 个立即开始（不等第 1 个完成）
+4. **正确性测试**：并发请求的输出内容应与单独跑每个请求一致
+
+## 实现计划
+
+1. 重构 Engine：从 `while recv → generate` 改为 scheduler loop
+2. 每个 Sequence 持有独立的 GpuKVCache
+3. 调度循环实现 admit + prefill + decode + finish
+4. HTTP API 侧改为 unbounded channel（允许多请求同时提交）
+5. 编写并发测试脚本
+
+## 当前状态
+
+**未实现**。当前是 FIFO 串行，一次只处理一个请求。本文档是实现的设计规格。

docs/13-http-api.md

@@ -0,0 +1,133 @@
+# Phase 13: HTTP API + Streaming — Design Document (Milestone ③)
+
+## Goal
+
+提供 OpenAI 兼容的 HTTP API，让 xserv 可以作为一个 serving 后端被任何 OpenAI SDK 调用。
+
+## 职责划分
+
+| 组件 | 职责 |
+|------|------|
+| Phase 12 (Scheduler/Engine) | 模型推理 + 请求调度 + token 生成循环 |
+| **Phase 13 (HTTP API)** | HTTP 请求解析 → 内部格式 → 提交给 engine → 从 channel 接收 token → 编码为 HTTP 响应 |
+
+Phase 13 不关心模型如何推理，只负责 HTTP 协议层。
+
+## 技术栈
+
+- **HTTP framework**: axum 0.8
+- **Async runtime**: tokio
+- **Serialization**: serde_json
+- **Channel**: tokio::sync::mpsc (API ↔ Engine)
+
+## API 端点
+
+```
+GET  /health                    → "ok"
+GET  /v1/models                 → {"data": [{"id": "qwen3-8b", ...}]}
+POST /v1/chat/completions       → JSON response (non-streaming)
+POST /v1/chat/completions       → SSE stream (streaming, TODO)
+```
+
+## Architecture
+
+```
+Client
+  │ HTTP POST /v1/chat/completions
+  ▼
+┌──────────────────────────────┐
+│  axum handler                │
+│  1. Deserialize ChatRequest  │
+│  2. Build prompt text        │
+│  3. Tokenize (Mutex<Tokenizer>)│
+│  4. Create mpsc channel      │
+│  5. Submit GenerateRequest   │
+│  6. await tokens from rx     │
+│  7. Build JSON response      │
+└──────────────────────────────┘
+  │ GenerateRequest via SyncSender
+  ▼
+┌──────────────────────────────┐
+│  Engine thread (Phase 12)    │
+│  - recv() request            │
+│  - model.forward_gpu_cache() │
+│  - blocking_send() tokens    │
+└──────────────────────────────┘
+```
+
+## OpenAI 兼容格式
+
+### Request
+```json
+{
+  "model": "qwen3-8b",
+  "messages": [
+    {"role": "system", "content": "You are helpful."},
+    {"role": "user", "content": "Hello"}
+  ],
+  "max_tokens": 256,
+  "stream": false
+}
+```
+
+### Response (non-streaming)
+```json
+{
+  "id": "chatcmpl-xxx",
+  "object": "chat.completion",
+  "created": 1234567890,
+  "model": "qwen3-8b",
+  "choices": [{
+    "index": 0,
+    "message": {"role": "assistant", "content": "Hi there!"},
+    "finish_reason": "stop"
+  }],
+  "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
+}
+```
+
+### SSE Streaming (TODO)
+```
+data: {"choices":[{"delta":{"content":"Hi"}}]}
+
+data: {"choices":[{"delta":{},"finish_reason":"stop"}]}
+
+data: [DONE]
+```
+
+## 当前实现状态
+
+- [x] `/health` — 健康检查
+- [x] `/v1/models` — 模型列表
+- [x] `/v1/chat/completions` (non-streaming) — JSON response
+- [ ] `/v1/chat/completions` (streaming) — SSE
+- [ ] 完整的 `usage` 统计 (token 计数)
+- [ ] 错误处理 (400 for bad request, etc.)
+- [ ] 多轮对话 chat template
+
+## Key Design Decisions
+
+1. **Extension vs State**: 用 `axum::Extension<Arc<AppState>>` 而不是 `Router::with_state`，因为 `SyncSender` 不是 `Sync`（需要 Mutex 包装）。
+
+2. **Engine 在独立 thread**: GPU 同步操作 block 线程，不能放在 tokio runtime 中。
+
+3. **tokio::sync::mpsc 做 token 传输**: Engine (std thread) 用 `blocking_send()`，API (async) 用 `.recv().await`。跨 async/sync 边界通信。
+
+## Test Plan
+
+- [x] curl /health → "ok"
+- [x] curl /v1/models → JSON model list
+- [x] curl /v1/chat/completions → JSON with generated text
+- [ ] Python OpenAI SDK 兼容性测试
+- [ ] SSE streaming 测试
+- [ ] 多轮对话测试
+
+## Takeaways
+
+1. **axum 0.8 的 Handler trait 对 Send 很严格**：async fn 返回的 Future 必须是 Send。`std::sync::MutexGuard` 不是 Send，必须确保它不活过 await point（用 scope 或显式 drop）。
+
+2. **std::sync::mpsc::SyncSender 不是 Sync**：不能直接放在 `Arc<T>` 中被多个 async task 共享。解决方案：`Mutex<SyncSender>` 或换用 `tokio::sync::mpsc::Sender`（是 Sync 的）。
+
+3. **非 streaming 更简单，先跑通再加 SSE**：SSE streaming 涉及 `Stream` trait、lifetime 问题和复杂的类型推导。先用 collect-all-then-respond 跑通 E2E，streaming 作为增量优化。
+
+4. **Engine 加载时间 ~20s（Qwen3-8B）**：需要在 server 启动后等 engine ready 才接受请求，否则请求会 hang 在 channel send 上。当前靠 sync_channel(1) 的背压天然处理。

docs/benchmarks/phase10-qwen3.md

@@ -0,0 +1,54 @@
+# Phase 10 Benchmark: Qwen3-8B
+
+**Date**: 2026-05-22
+**Hardware**: RTX 5090 (32GB, CC 12.0)
+**Model**: Qwen3-8B (BF16, 36 layers, 4096 hidden, 32/8 GQA heads)
+**Config**: 50 prompts × 20 generated tokens, greedy decoding, KV cache
+
+## Correctness
+
+| Metric | Result |
+|--------|--------|
+| Prefill Top-1 match vs HF | **42/50 (84.0%)** |
+| Prefill Top-5 match vs HF | **50/50 (100.0%)** |
+| Greedy sequence match | 0/50 (expected — BF16 drift over decode) |
+
+The 100% top-5 match confirms the model is computing correctly.
+Greedy sequence divergence is due to BF16 precision (7-bit mantissa)
+accumulating across 36 layers of decode steps. Both xserv and HF
+produce coherent, valid completions — they just pick different
+equally-likely tokens at close-logit decision points.
+
+## Performance
+
+| Metric | xserv | transformers (BF16) | Ratio |
+|--------|-------|--------------------:|-------|
+| TTFT (avg) | 138.5 ms | 21.2 ms | 6.5x slower |
+| TBT (avg) | 144.2 ms | 21.9 ms | 6.6x slower |
+| Throughput | 6.9 tok/s | 45.6 tok/s | 0.15x |
+
+## Remaining Performance Gap
+
+~6.6x slower than HF for an 8B BF16 model. Main bottlenecks:
+1. CPU round-trips for add/mul/reshape/merge_heads (~100 per forward pass)
+2. KV cache stored on CPU (rebuilt as GPU tensor each step)
+3. cuBLAS handle per matmul
+4. No kernel fusion
+5. GQA repeat_kv copies data instead of kernel-level indexing
+
+## Output Quality (Sample)
+
+| Prompt | xserv Output |
+|--------|-------------|
+| "The capital of France is" | "Paris. The capital of France is Paris..." |
+| "Climate change is caused by" | "human activities, and the effects are already being felt..." |
+| "The human brain contains approximately" | "86 billion neurons. Each neuron can form synapses..." |
+| "Python is a popular programming language because" | "it is easy to learn and use..." |
+
+## Tracking
+
+| Phase | Model | TTFT (ms) | TBT (ms) | tok/s | Correctness |
+|-------|-------|-----------|----------|-------|-------------|
+| 8 | GPT-2 FP32 | 400.6 | 407.2 | 2.5 | 50/50 vs HF |
+| 9 | GPT-2 FP32 KV | 24.2 | 22.6 | 44.3 | 50/50 self |
+| 10 | Qwen3-8B BF16 KV | 138.5 | 144.2 | 6.9 | 100% top-5 prefill |

docs/benchmarks/phase8-gpt2-baseline.md

@@ -0,0 +1,35 @@
+# Phase 8 Benchmark: GPT-2 124M Baseline
+
+**Date**: 2026-05-21
+**Hardware**: RTX 5090 (32GB, CC 12.0, 170 SMs)
+**Model**: GPT-2 124M (FP32)
+**Config**: 50 prompts × 20 generated tokens, greedy decoding, no KV cache
+
+## Correctness
+
+| Metric | Result |
+|--------|--------|
+| Prompts tested | 50 |
+| Token-level match vs transformers | **50/50 (100.0%)** |
+| Mismatches | 0 |
+
+## Performance
+
+| Metric | xserv | transformers (PyTorch) | Ratio |
+|--------|-------|----------------------|-------|
+| TTFT (avg) | 400.6 ms | 4.0 ms | 100x slower |
+| TBT (avg) | 407.2 ms | 3.8 ms | 106x slower |
+| Throughput | 2.5 tok/s | 260 tok/s | 0.01x |
+
+## Known Bottlenecks
+
+1. **No KV Cache**: full recompute per token (O(S²) attention every step)
+2. **CPU round-trips**: ~100 GPU→CPU→GPU transfers per forward pass for add/bias/split_qkv/merge_heads
+3. **cuBLAS handle per matmul**: ~50 handle create/destroy per forward pass
+4. **No kernel fusion**: every op is a separate kernel launch + sync
+
+## Tracking
+
+| Phase | TTFT (ms) | TBT (ms) | tok/s | Correctness | Notes |
+|-------|-----------|----------|-------|-------------|-------|
+| 8 (baseline) | 400.6 | 407.2 | 2.5 | 50/50 | No KV cache, CPU round-trips |

docs/benchmarks/phase9-kv-cache.md

@@ -0,0 +1,44 @@
+# Phase 9 Benchmark: KV Cache
+
+**Date**: 2026-05-21
+**Hardware**: RTX 5090 (32GB, CC 12.0)
+**Model**: GPT-2 124M (FP32)
+**Config**: 50 prompts × 20 generated tokens, greedy decoding
+
+## Correctness
+
+| Metric | Result |
+|--------|--------|
+| xserv KV-cache vs xserv no-cache | **50/50 (100.0%)** — bit-identical |
+| xserv vs HF transformers | 40/50 (80.0%) |
+
+The 10 mismatches vs HF are floating point divergence (different CUDA kernels, computation order).
+Logit gap at divergence points: min=0.04, max=0.56, avg=0.20. Not a correctness bug.
+
+## Performance
+
+| Metric | Phase 8 (no cache) | Phase 9 (KV cache) | Improvement | HF transformers |
+|--------|-------------------|--------------------|-----------|-----------------| 
+| TTFT (avg) | 400.6 ms | 24.2 ms | **16.5x** | 4.0 ms |
+| TBT (avg) | 407.2 ms | 22.6 ms | **18.0x** | 3.9 ms |
+| Throughput | 2.5 tok/s | 44.3 tok/s | **17.7x** | 257.7 tok/s |
+| vs HF ratio | 0.01x | 0.17x | | 1.0x |
+
+## Analysis
+
+KV cache delivers **~18x speedup** by eliminating redundant computation:
+- Before: every decode step recomputed all layers for all tokens O(S²)
+- After: decode step only computes 1 new token, reads K/V from cache O(S)
+
+Remaining gap vs HF (~6x slower):
+1. CPU round-trips still present (~100 per forward pass)
+2. cuBLAS handle created per matmul
+3. KV cache stored on CPU (rebuilt as GPU tensor each step)
+4. No kernel fusion
+
+## Tracking
+
+| Phase | TTFT (ms) | TBT (ms) | tok/s | Correctness | Notes |
+|-------|-----------|----------|-------|-------------|-------|
+| 8 (baseline) | 400.6 | 407.2 | 2.5 | 50/50 vs HF | No KV cache |
+| 9 (KV cache) | 24.2 | 22.6 | 44.3 | 50/50 self-consistent | 18x speedup |

tools/analyze_divergence.py

@@ -0,0 +1,40 @@
+import json
+import sys
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+model = GPT2LMHeadModel.from_pretrained(sys.argv[2]).eval().cuda()
+tokenizer = GPT2Tokenizer.from_pretrained(sys.argv[2])
+
+with open(sys.argv[1]) as f:
+    xr = json.load(f)
+
+mismatches = []
+for i in range(len(xr)):
+    ids = tokenizer.encode(xr[i]["prompt"])
+    all_ids = list(ids)
+    xserv_gen = xr[i]["generated_ids"]
+    with torch.no_grad():
+        for j in range(len(xserv_gen)):
+            out = model(torch.tensor([all_ids]).cuda())
+            logits = out.logits[0, -1]
+            hf_next = logits.argmax().item()
+            xs_next = xserv_gen[j]
+            if hf_next != xs_next:
+                xs_logit = logits[xs_next].item()
+                hf_logit = logits[hf_next].item()
+                hf_tok = tokenizer.decode([hf_next])
+                xs_tok = tokenizer.decode([xs_next])
+                gap = hf_logit - xs_logit
+                print(
+                    f'[{i+1}] "{xr[i]["prompt"][:42]}" @ tok {j}: '
+                    f'hf={repr(hf_tok)}({hf_logit:.3f}) xserv={repr(xs_tok)}({xs_logit:.3f}) '
+                    f'gap={gap:.4f}'
+                )
+                mismatches.append(gap)
+                break
+            all_ids.append(hf_next)
+
+print(f"\nTotal: {len(mismatches)}/{len(xr)} mismatches")
+if mismatches:
+    print(f"Logit gaps: min={min(mismatches):.4f} max={max(mismatches):.4f} avg={sum(mismatches)/len(mismatches):.4f}")

tools/bench_compare.py

@@ -0,0 +1,154 @@
+"""
+Compare xserv GPT-2 output against HuggingFace transformers.
+Reads xserv results from JSON, runs same prompts through transformers, compares token-by-token.
+Also measures transformers timing for performance comparison.
+
+Usage:
+    python3 tools/bench_compare.py <xserv_results.json> <model_dir>
+"""
+
+import json
+import sys
+import time
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+
+def main():
+    if len(sys.argv) < 3:
+        print(f"Usage: {sys.argv[0]} <xserv_results.json> <model_dir>")
+        sys.exit(1)
+
+    xserv_path = sys.argv[1]
+    model_dir = sys.argv[2]
+
+    with open(xserv_path) as f:
+        xserv_results = json.load(f)
+
+    print(f"Loading transformers model from {model_dir}...")
+    model = GPT2LMHeadModel.from_pretrained(model_dir)
+    tokenizer = GPT2Tokenizer.from_pretrained(model_dir)
+    model.eval()
+    model.cuda()
+
+    # Warmup
+    with torch.no_grad():
+        model(torch.tensor([[tokenizer.encode("warmup")[0]]]).cuda())
+    torch.cuda.synchronize()
+
+    total = len(xserv_results)
+    match_count = 0
+    mismatch_count = 0
+    xserv_ttft_sum = 0.0
+    xserv_tbt_sum = 0.0
+    hf_ttft_sum = 0.0
+    hf_tbt_sum = 0.0
+    num_with_tbt = 0
+
+    print(f"\n{'='*100}")
+    print(f"{'#':>3} {'Match':>5} {'Prompt':<45} {'xserv TTFT':>10} {'HF TTFT':>10} {'xserv TBT':>10} {'HF TBT':>10}")
+    print(f"{'='*100}")
+
+    for i, xr in enumerate(xserv_results):
+        prompt = xr["prompt"]
+        gen_tokens = xr["num_generated"]
+        xserv_ids = xr["generated_ids"]
+
+        input_ids = tokenizer.encode(prompt)
+        input_tensor = torch.tensor([input_ids]).cuda()
+
+        # Generate with transformers, measuring timing
+        hf_generated = []
+        hf_token_times = []
+
+        with torch.no_grad():
+            all_ids = input_tensor.clone()
+
+            # TTFT
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            out = model(all_ids)
+            torch.cuda.synchronize()
+            hf_ttft_us = (time.perf_counter() - t0) * 1e6
+            next_id = out.logits[0, -1].argmax().item()
+            hf_generated.append(next_id)
+            all_ids = torch.cat([all_ids, torch.tensor([[next_id]]).cuda()], dim=1)
+
+            # Remaining tokens
+            for _ in range(1, gen_tokens):
+                torch.cuda.synchronize()
+                t_start = time.perf_counter()
+                out = model(all_ids)
+                torch.cuda.synchronize()
+                elapsed = (time.perf_counter() - t_start) * 1e6
+                hf_token_times.append(elapsed)
+                next_id = out.logits[0, -1].argmax().item()
+                hf_generated.append(next_id)
+                all_ids = torch.cat([all_ids, torch.tensor([[next_id]]).cuda()], dim=1)
+
+                eos_id = tokenizer.eos_token_id
+                if eos_id is not None and next_id == eos_id:
+                    break
+
+        hf_tbt_us = sum(hf_token_times) / len(hf_token_times) if hf_token_times else 0
+
+        # Compare
+        match = xserv_ids == hf_generated
+        if match:
+            match_count += 1
+            status = "  OK "
+        else:
+            mismatch_count += 1
+            status = "FAIL!"
+
+        xserv_ttft_ms = xr["ttft_us"] / 1000.0
+        xserv_tbt_ms = xr["tbt_us"] / 1000.0
+        hf_ttft_ms = hf_ttft_us / 1000.0
+        hf_tbt_ms = hf_tbt_us / 1000.0
+
+        prompt_short = prompt[:43] + ".." if len(prompt) > 45 else prompt
+        print(f"{i+1:>3} {status} {prompt_short:<45} {xserv_ttft_ms:>8.1f}ms {hf_ttft_ms:>8.1f}ms {xserv_tbt_ms:>8.1f}ms {hf_tbt_ms:>8.1f}ms")
+
+        if not match:
+            # Show first divergence
+            for j in range(max(len(xserv_ids), len(hf_generated))):
+                x = xserv_ids[j] if j < len(xserv_ids) else None
+                h = hf_generated[j] if j < len(hf_generated) else None
+                if x != h:
+                    x_tok = tokenizer.decode([x]) if x is not None else "<none>"
+                    h_tok = tokenizer.decode([h]) if h is not None else "<none>"
+                    print(f"      ↳ diverge at token {j}: xserv={x}({repr(x_tok)}) vs hf={h}({repr(h_tok)})")
+                    break
+
+        xserv_ttft_sum += xr["ttft_us"]
+        xserv_tbt_sum += xr["tbt_us"]
+        hf_ttft_sum += hf_ttft_us
+        hf_tbt_sum += hf_tbt_us
+        if xr["tbt_us"] > 0:
+            num_with_tbt += 1
+
+    print(f"{'='*100}")
+    print(f"\n=== CORRECTNESS ===")
+    print(f"Total prompts: {total}")
+    print(f"Match:    {match_count}/{total} ({match_count/total*100:.1f}%)")
+    print(f"Mismatch: {mismatch_count}/{total}")
+
+    print(f"\n=== PERFORMANCE (average) ===")
+    print(f"{'Metric':<20} {'xserv':>12} {'transformers':>12} {'ratio':>10}")
+    print(f"{'-'*54}")
+    avg_x_ttft = xserv_ttft_sum / total / 1000
+    avg_h_ttft = hf_ttft_sum / total / 1000
+    avg_x_tbt = xserv_tbt_sum / num_with_tbt / 1000 if num_with_tbt > 0 else 0
+    avg_h_tbt = hf_tbt_sum / num_with_tbt / 1000 if num_with_tbt > 0 else 0
+    print(f"{'TTFT (ms)':<20} {avg_x_ttft:>10.1f}ms {avg_h_ttft:>10.1f}ms {avg_x_ttft/avg_h_ttft:>9.1f}x")
+    print(f"{'TBT (ms)':<20} {avg_x_tbt:>10.1f}ms {avg_h_tbt:>10.1f}ms {avg_x_tbt/avg_h_tbt if avg_h_tbt > 0 else 0:>9.1f}x")
+    xserv_tps = 1000.0 / avg_x_tbt if avg_x_tbt > 0 else 0
+    hf_tps = 1000.0 / avg_h_tbt if avg_h_tbt > 0 else 0
+    print(f"{'Throughput (tok/s)':<20} {xserv_tps:>10.1f}   {hf_tps:>10.1f}   {xserv_tps/hf_tps if hf_tps > 0 else 0:>9.2f}x")
+
+    print(f"\nNote: xserv currently has no KV cache — full recompute per token.")
+    print(f"      transformers also runs without KV cache in this benchmark for fair comparison.")
+
+
+if __name__ == "__main__":
+    main()

tools/bench_compare_qwen3.py

@@ -0,0 +1,137 @@
+"""
+Compare xserv Qwen3 output against HuggingFace transformers.
+Usage: python3 tools/bench_compare_qwen3.py <xserv_results.json> <model_dir>
+"""
+import json
+import sys
+import time
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+def main():
+    if len(sys.argv) < 3:
+        print(f"Usage: {sys.argv[0]} <xserv_results.json> <model_dir>")
+        sys.exit(1)
+
+    xserv_path = sys.argv[1]
+    model_dir = sys.argv[2]
+
+    with open(xserv_path) as f:
+        xserv_results = json.load(f)
+
+    print(f"Loading transformers model from {model_dir}...")
+    model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.bfloat16)
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model.eval()
+    model.cuda()
+
+    # Warmup
+    with torch.no_grad():
+        ids = tokenizer.encode("warmup", return_tensors="pt").cuda()
+        model(ids)
+    torch.cuda.synchronize()
+
+    total = len(xserv_results)
+    match_count = 0
+    mismatch_count = 0
+    xserv_ttft_sum = 0.0
+    xserv_tbt_sum = 0.0
+    hf_ttft_sum = 0.0
+    hf_tbt_sum = 0.0
+    num_with_tbt = 0
+
+    print(f"\n{'='*100}")
+    print(f"{'#':>3} {'Match':>5} {'Prompt':<45} {'xserv TTFT':>10} {'HF TTFT':>10} {'xserv TBT':>10} {'HF TBT':>10}")
+    print(f"{'='*100}")
+
+    for i, xr in enumerate(xserv_results):
+        prompt = xr["prompt"]
+        gen_tokens = xr["num_generated"]
+        xserv_ids = xr["generated_ids"]
+
+        input_ids = tokenizer.encode(prompt, return_tensors="pt").cuda()
+        hf_generated = []
+        hf_token_times = []
+
+        with torch.no_grad():
+            all_ids = input_ids.clone()
+
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            out = model(all_ids)
+            torch.cuda.synchronize()
+            hf_ttft_us = (time.perf_counter() - t0) * 1e6
+            next_id = out.logits[0, -1].argmax().item()
+            hf_generated.append(next_id)
+            all_ids = torch.cat([all_ids, torch.tensor([[next_id]]).cuda()], dim=1)
+
+            for _ in range(1, gen_tokens):
+                torch.cuda.synchronize()
+                t_start = time.perf_counter()
+                out = model(all_ids)
+                torch.cuda.synchronize()
+                elapsed = (time.perf_counter() - t_start) * 1e6
+                hf_token_times.append(elapsed)
+                next_id = out.logits[0, -1].argmax().item()
+                hf_generated.append(next_id)
+                all_ids = torch.cat([all_ids, torch.tensor([[next_id]]).cuda()], dim=1)
+
+                if next_id == tokenizer.eos_token_id:
+                    break
+
+        hf_tbt_us = sum(hf_token_times) / len(hf_token_times) if hf_token_times else 0
+
+        match = xserv_ids == hf_generated
+        if match:
+            match_count += 1
+            status = "  OK "
+        else:
+            mismatch_count += 1
+            status = "FAIL!"
+
+        xserv_ttft_ms = xr["ttft_us"] / 1000.0
+        xserv_tbt_ms = xr["tbt_us"] / 1000.0
+        hf_ttft_ms = hf_ttft_us / 1000.0
+        hf_tbt_ms = hf_tbt_us / 1000.0
+
+        prompt_short = prompt[:43] + ".." if len(prompt) > 45 else prompt
+        print(f"{i+1:>3} {status} {prompt_short:<45} {xserv_ttft_ms:>8.1f}ms {hf_ttft_ms:>8.1f}ms {xserv_tbt_ms:>8.1f}ms {hf_tbt_ms:>8.1f}ms")
+
+        if not match:
+            for j in range(max(len(xserv_ids), len(hf_generated))):
+                x = xserv_ids[j] if j < len(xserv_ids) else None
+                h = hf_generated[j] if j < len(hf_generated) else None
+                if x != h:
+                    x_tok = tokenizer.decode([x]) if x is not None else "<none>"
+                    h_tok = tokenizer.decode([h]) if h is not None else "<none>"
+                    print(f"      diverge@{j}: xserv={x}({repr(x_tok)}) hf={h}({repr(h_tok)})")
+                    break
+
+        xserv_ttft_sum += xr["ttft_us"]
+        xserv_tbt_sum += xr["tbt_us"]
+        hf_ttft_sum += hf_ttft_us
+        hf_tbt_sum += hf_tbt_us
+        if xr["tbt_us"] > 0:
+            num_with_tbt += 1
+
+    print(f"{'='*100}")
+    print(f"\n=== CORRECTNESS ===")
+    print(f"Total: {total}, Match: {match_count}/{total} ({match_count/total*100:.1f}%), Mismatch: {mismatch_count}")
+
+    print(f"\n=== PERFORMANCE ===")
+    print(f"{'Metric':<20} {'xserv':>12} {'transformers':>12} {'ratio':>10}")
+    print(f"{'-'*54}")
+    avg_x_ttft = xserv_ttft_sum / total / 1000
+    avg_h_ttft = hf_ttft_sum / total / 1000
+    avg_x_tbt = xserv_tbt_sum / num_with_tbt / 1000 if num_with_tbt > 0 else 0
+    avg_h_tbt = hf_tbt_sum / num_with_tbt / 1000 if num_with_tbt > 0 else 0
+    print(f"{'TTFT (ms)':<20} {avg_x_ttft:>10.1f}ms {avg_h_ttft:>10.1f}ms {avg_x_ttft/avg_h_ttft if avg_h_ttft>0 else 0:>9.1f}x")
+    print(f"{'TBT (ms)':<20} {avg_x_tbt:>10.1f}ms {avg_h_tbt:>10.1f}ms {avg_x_tbt/avg_h_tbt if avg_h_tbt>0 else 0:>9.1f}x")
+    xserv_tps = 1000.0 / avg_x_tbt if avg_x_tbt > 0 else 0
+    hf_tps = 1000.0 / avg_h_tbt if avg_h_tbt > 0 else 0
+    print(f"{'Throughput (tok/s)':<20} {xserv_tps:>10.1f}   {hf_tps:>10.1f}   {xserv_tps/hf_tps if hf_tps>0 else 0:>9.2f}x")
+
+
+if __name__ == "__main__":
+    main()

tools/test_concurrent.py

@@ -0,0 +1,107 @@
+"""
+Test concurrent request handling.
+Sends N requests simultaneously, verifies they all produce tokens concurrently.
+
+Usage: python3 tools/test_concurrent.py <server_url> [num_requests]
+"""
+import sys
+import time
+import json
+import threading
+import urllib.request
+import urllib.error
+
+
+def send_request(url, prompt, max_tokens, results, idx):
+    """Send a chat completion request and record timing."""
+    body = json.dumps({
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": max_tokens,
+    }).encode()
+
+    req = urllib.request.Request(
+        f"{url}/v1/chat/completions",
+        data=body,
+        headers={"Content-Type": "application/json"},
+    )
+
+    t0 = time.time()
+    try:
+        with urllib.request.urlopen(req, timeout=120) as resp:
+            data = json.loads(resp.read())
+            t1 = time.time()
+            content = data["choices"][0]["message"]["content"]
+            results[idx] = {
+                "status": "ok",
+                "content": content,
+                "duration_s": t1 - t0,
+                "finish_reason": data["choices"][0]["finish_reason"],
+            }
+    except Exception as e:
+        t1 = time.time()
+        results[idx] = {"status": "error", "error": str(e), "duration_s": t1 - t0}
+
+
+def main():
+    url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:9090"
+    n = int(sys.argv[2]) if len(sys.argv) > 2 else 3
+    max_tokens = 10
+
+    prompts = [
+        "What is the capital of France?",
+        "Tell me about quantum computing",
+        "How do airplanes fly?",
+        "What is machine learning?",
+        "Explain gravity in simple terms",
+    ][:n]
+
+    print(f"Sending {n} concurrent requests to {url} (max_tokens={max_tokens})")
+    print("=" * 70)
+
+    results = [None] * n
+    threads = []
+
+    t_start = time.time()
+    for i, prompt in enumerate(prompts):
+        t = threading.Thread(target=send_request, args=(url, prompt, max_tokens, results, i))
+        threads.append(t)
+        t.start()
+
+    for t in threads:
+        t.join()
+    t_total = time.time() - t_start
+
+    print(f"\n{'#':>2} {'Status':>6} {'Duration':>8} {'Content':<50}")
+    print("-" * 70)
+    for i, r in enumerate(results):
+        if r["status"] == "ok":
+            content_short = r["content"].replace("\n", " ")[:48]
+            print(f"{i+1:>2} {'OK':>6} {r['duration_s']:>6.1f}s  {content_short}")
+        else:
+            print(f"{i+1:>2} {'FAIL':>6} {r['duration_s']:>6.1f}s  {r['error'][:48]}")
+
+    print("=" * 70)
+    print(f"Total wall time: {t_total:.1f}s")
+
+    # Analyze concurrency
+    durations = [r["duration_s"] for r in results if r["status"] == "ok"]
+    if len(durations) >= 2:
+        sequential_estimate = sum(durations)
+        actual_wall = t_total
+        concurrency_ratio = sequential_estimate / actual_wall if actual_wall > 0 else 0
+
+        print(f"Sum of individual durations: {sequential_estimate:.1f}s")
+        print(f"Actual wall time: {actual_wall:.1f}s")
+        print(f"Concurrency ratio: {concurrency_ratio:.2f}x")
+
+        if concurrency_ratio > 1.5:
+            print("✓ CONCURRENT: requests are being processed in parallel")
+        else:
+            print("✗ SERIAL: requests appear to be processed sequentially")
+
+    all_ok = all(r["status"] == "ok" for r in results)
+    print(f"\nAll requests succeeded: {all_ok}")
+
+
+if __name__ == "__main__":
+    main()