tools: add correctness + performance test scripts for Qwen3-8B

- test_correctness.py: compare prefill logits top-20 vs HF transformers
- bench_server.py: HTTP API benchmark (throughput, streaming, concurrent, EOS leak check)

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

Cargo.toml

@@ -2,6 +2,11 @@
 resolver = "2"
 members = [
     "crates/xserv-cuda",
+    "crates/xserv-tensor",
+    "crates/xserv-kernels",
+    "crates/xserv-model",
+    "crates/xserv-tokenizer",
+    "crates/xserv-server",
 ]
 
 [workspace.package]
@@ -12,3 +17,12 @@ license = "MIT"
 [workspace.dependencies]
 half = "2"
 smallvec = "1"
+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"] }
+tokio-stream = "0.1"
+rand = "0.8"

crates/xserv-cuda/src/allocator.rs

@@ -1,6 +1,7 @@
 use crate::error::Result;
 use crate::ffi;
 use crate::memory::GpuBuffer;
+use std::cell::RefCell;
 use std::collections::HashMap;
 
 /// Caching allocator that reuses freed GPU buffers instead of calling
@@ -84,6 +85,33 @@ impl Drop for CachingAllocator {
     }
 }
 
+thread_local! {
+    static ALLOCATOR: RefCell<CachingAllocator> = RefCell::new(CachingAllocator::new());
+}
+
+/// Allocate a GPU buffer through the caching allocator.
+/// The returned buffer has `pooled = true` so it will be returned
+/// to the pool on drop instead of calling cudaFree.
+pub fn cached_alloc(size: usize) -> Result<GpuBuffer> {
+    ALLOCATOR.with(|cell| {
+        let mut buf = cell.borrow_mut().alloc(size)?;
+        buf.set_pooled(true);
+        Ok(buf)
+    })
+}
+
+/// Return a raw GPU pointer to the caching allocator's free list.
+/// Called from `GpuBuffer::Drop` for pooled buffers. Takes raw pointer
+/// and size to avoid re-triggering Drop.
+pub fn return_to_pool(ptr: *mut u8, len: usize) {
+    ALLOCATOR.with(|cell| {
+        let mut alloc = cell.borrow_mut();
+        let bucket = bucket_size(len);
+        alloc.stats.current_allocated = alloc.stats.current_allocated.saturating_sub(len);
+        alloc.free_lists.entry(bucket).or_default().push((ptr, len));
+    });
+}
+
 /// Round up to next power-of-2, minimum 512 bytes.
 fn bucket_size(size: usize) -> usize {
     let min = 512;

crates/xserv-cuda/src/device.rs

@@ -1,6 +1,7 @@
 use crate::error::{self, Result};
 use crate::ffi;
 use std::ffi::CStr;
+use std::os::raw::c_char;
 
 #[derive(Debug, Clone)]
 pub struct DeviceInfo {
@@ -44,10 +45,14 @@ pub fn current_device() -> Result<u32> {
 }
 
 pub fn device_info(device: u32) -> Result<DeviceInfo> {
-    // Get device name from cudaGetDeviceProperties (only use the name field).
-    let mut prop = unsafe { std::mem::zeroed::<ffi::CudaDeviceProp>() };
-    error::check(unsafe { ffi::cudaGetDeviceProperties(&mut prop, device as i32) })?;
-    let name = unsafe { CStr::from_ptr(prop.name.as_ptr()) }
+    // Heap-allocate oversized buffer for cudaDeviceProp (layout varies by CUDA version).
+    // CUDA 12.x struct is ~5-6 KB; use 32 KB to guard against future growth.
+    let mut prop_buf = vec![0u8; 32768];
+    error::check(unsafe {
+        ffi::cudaGetDeviceProperties(prop_buf.as_mut_ptr(), device as i32)
+    })?;
+    // Name is always the first field: char[256].
+    let name = unsafe { CStr::from_ptr(prop_buf.as_ptr() as *const c_char) }
         .to_string_lossy()
         .into_owned();
 

crates/xserv-cuda/src/error.rs

@@ -23,7 +23,7 @@ impl std::error::Error for CudaError {}
 
 pub type Result<T> = std::result::Result<T, CudaError>;
 
-pub(crate) fn check(code: i32) -> Result<()> {
+pub fn check(code: i32) -> Result<()> {
     if code == ffi::CUDA_SUCCESS {
         return Ok(());
     }

crates/xserv-cuda/src/ffi.rs

@@ -3,6 +3,8 @@ use std::os::raw::c_char;
 
 pub type CudaStream = *mut c_void;
 pub type CudaEvent = *mut c_void;
+pub type CudaGraph = *mut c_void;
+pub type CudaGraphExec = *mut c_void;
 
 pub const CUDA_MEMCPY_H2D: i32 = 1;
 pub const CUDA_MEMCPY_D2H: i32 = 2;
@@ -11,31 +13,16 @@ pub const CUDA_MEMCPY_D2D: i32 = 3;
 pub const CUDA_SUCCESS: i32 = 0;
 pub const CUDA_ERROR_OUT_OF_MEMORY: i32 = 2;
 
-#[repr(C)]
-pub struct CudaDeviceProp {
-    pub name: [c_char; 256],
-    pub total_global_mem: usize,
-    pub shared_mem_per_block: usize,
-    pub regs_per_block: i32,
-    pub warp_size: i32,
-    pub max_threads_per_block: i32,
-    pub max_threads_dim: [i32; 3],
-    pub max_grid_size: [i32; 3],
-    pub clock_rate: i32,
-    pub total_const_mem: usize,
-    pub major: i32,
-    pub minor: i32,
-    // There are many more fields; we only read up to what we need.
-    // cudaDeviceProp is a large struct (~1KB). We pad the rest.
-    _pad: [u8; 4096],
-}
+/// cudaStreamCaptureMode::cudaStreamCaptureModeGlobal
+pub const CUDA_STREAM_CAPTURE_MODE_GLOBAL: i32 = 0;
 
 unsafe extern "C" {
     // --- Device ---
     pub fn cudaGetDeviceCount(count: *mut i32) -> i32;
     pub fn cudaSetDevice(device: i32) -> i32;
     pub fn cudaGetDevice(device: *mut i32) -> i32;
-    pub fn cudaGetDeviceProperties(prop: *mut CudaDeviceProp, device: i32) -> i32;
+    /// Takes a raw pointer; caller provides a heap buffer large enough for any CUDA version.
+    pub fn cudaGetDeviceProperties(prop: *mut u8, device: i32) -> i32;
     pub fn cudaDeviceSynchronize() -> i32;
 
     // --- Memory ---
@@ -62,6 +49,18 @@ unsafe extern "C" {
     pub fn cudaGetLastError() -> i32;
     pub fn cudaGetErrorString(error: i32) -> *const c_char;
 
+    // --- CUDA Graphs ---
+    pub fn cudaStreamBeginCapture(stream: CudaStream, mode: i32) -> i32;
+    pub fn cudaStreamEndCapture(stream: CudaStream, graph: *mut CudaGraph) -> i32;
+    pub fn cudaGraphInstantiate(
+        graph_exec: *mut CudaGraphExec,
+        graph: CudaGraph,
+        flags: u64,
+    ) -> i32;
+    pub fn cudaGraphLaunch(graph_exec: CudaGraphExec, stream: CudaStream) -> i32;
+    pub fn cudaGraphDestroy(graph: CudaGraph) -> i32;
+    pub fn cudaGraphExecDestroy(graph_exec: CudaGraphExec) -> i32;
+
     // --- Our test kernel ---
     pub fn launch_vecadd_f32(
         a: *const f32,

crates/xserv-cuda/src/graph.rs

@@ -0,0 +1,98 @@
+//! CUDA Graphs: capture a sequence of kernel launches and replay them with
+//! near-zero host-side overhead (~3-5 us per launch eliminated).
+//!
+//! Usage:
+//! ```ignore
+//! let stream = CudaStream::new()?;
+//! let mut graph = CudaGraph::new();
+//!
+//! // First call: capture
+//! graph.begin_capture(&stream)?;
+//! // ... launch kernels on `stream` ...
+//! graph.end_capture(&stream)?;
+//!
+//! // Subsequent calls: replay
+//! graph.launch(&stream)?;
+//! ```
+//!
+//! Requirements for captured kernels:
+//! - All tensor shapes must be identical between capture and replay.
+//! - No host-side branching during the captured section.
+//! - Memory addresses used during capture must remain valid during replay.
+
+use crate::error::{self, Result};
+use crate::ffi;
+use crate::stream::CudaStream;
+
+/// RAII wrapper around a captured CUDA graph and its executable instance.
+pub struct CudaGraph {
+    graph: ffi::CudaGraph,
+    exec: ffi::CudaGraphExec,
+}
+
+impl CudaGraph {
+    /// Create an empty graph handle (not yet captured).
+    pub fn new() -> Self {
+        Self {
+            graph: std::ptr::null_mut(),
+            exec: std::ptr::null_mut(),
+        }
+    }
+
+    /// Returns true if a graph has been captured and instantiated.
+    pub fn is_ready(&self) -> bool {
+        !self.exec.is_null()
+    }
+
+    /// Begin capturing kernel launches on `stream`.
+    /// All subsequent kernel launches on this stream are recorded into the
+    /// graph instead of being executed.
+    pub fn begin_capture(&mut self, stream: &CudaStream) -> Result<()> {
+        // If we have an old graph, destroy it first
+        self.destroy_inner();
+        error::check(unsafe {
+            ffi::cudaStreamBeginCapture(
+                stream.as_raw(),
+                ffi::CUDA_STREAM_CAPTURE_MODE_GLOBAL,
+            )
+        })
+    }
+
+    /// End capture and instantiate the executable graph.
+    pub fn end_capture(&mut self, stream: &CudaStream) -> Result<()> {
+        error::check(unsafe {
+            ffi::cudaStreamEndCapture(stream.as_raw(), &mut self.graph)
+        })?;
+        error::check(unsafe {
+            ffi::cudaGraphInstantiate(&mut self.exec, self.graph, 0)
+        })
+    }
+
+    /// Replay the captured graph on `stream`.
+    /// Panics if no graph has been captured yet.
+    pub fn launch(&self, stream: &CudaStream) -> Result<()> {
+        assert!(self.is_ready(), "CudaGraph::launch called before capture");
+        error::check(unsafe {
+            ffi::cudaGraphLaunch(self.exec, stream.as_raw())
+        })
+    }
+
+    fn destroy_inner(&mut self) {
+        if !self.exec.is_null() {
+            unsafe { ffi::cudaGraphExecDestroy(self.exec) };
+            self.exec = std::ptr::null_mut();
+        }
+        if !self.graph.is_null() {
+            unsafe { ffi::cudaGraphDestroy(self.graph) };
+            self.graph = std::ptr::null_mut();
+        }
+    }
+}
+
+impl Drop for CudaGraph {
+    fn drop(&mut self) {
+        self.destroy_inner();
+    }
+}
+
+unsafe impl Send for CudaGraph {}

crates/xserv-cuda/src/lib.rs

@@ -2,11 +2,13 @@ pub mod allocator;
 pub mod device;
 pub mod error;
 pub mod ffi;
+pub mod graph;
 pub mod memory;
 pub mod stream;
 
 pub use allocator::CachingAllocator;
 pub use device::DeviceInfo;
 pub use error::{CudaError, Result};
+pub use graph::CudaGraph;
 pub use memory::{GpuBuffer, PinnedBuffer};
 pub use stream::CudaStream;

crates/xserv-cuda/src/memory.rs

@@ -3,9 +3,18 @@ use crate::ffi;
 use crate::stream::CudaStream;
 
 /// RAII wrapper around a GPU memory allocation.
+///
+/// When `owned` is true (the default), dropping frees the GPU memory.
+/// A borrowed buffer (`owned = false`) does NOT free on drop — the
+/// caller must ensure the backing allocation outlives all borrows.
+///
+/// When `pooled` is true, dropping returns the buffer to the caching
+/// allocator's free list instead of calling cudaFree.
 pub struct GpuBuffer {
     ptr: *mut u8,
     len: usize,
+    owned: bool,
+    pooled: bool,
 }
 
 impl GpuBuffer {
@@ -13,7 +22,13 @@ impl GpuBuffer {
         assert!(len > 0, "cannot allocate 0 bytes on GPU");
         let mut ptr = std::ptr::null_mut();
         error::check(unsafe { ffi::cudaMalloc(&mut ptr, len) })?;
-        Ok(Self { ptr, len })
+        Ok(Self { ptr, len, owned: true, pooled: false })
+    }
+
+    /// Mark this buffer as pooled (returned to caching allocator on drop)
+    /// or not. Called by `cached_alloc` after obtaining a buffer.
+    pub fn set_pooled(&mut self, pooled: bool) {
+        self.pooled = pooled;
     }
 
     pub fn len(&self) -> usize {
@@ -87,6 +102,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) {
@@ -99,14 +128,29 @@ impl GpuBuffer {
     /// Reconstruct a GpuBuffer from a raw pointer + length.
     /// Safety: ptr must have been allocated with cudaMalloc, len must be correct.
     pub unsafe fn from_raw(ptr: *mut u8, len: usize) -> Self {
-        Self { ptr, len }
+        Self { ptr, len, owned: true, pooled: false }
+    }
+
+    /// Create a non-owning view of GPU memory. Dropping this buffer does NOT
+    /// call `cudaFree`. The caller must ensure the underlying allocation
+    /// outlives this borrow.
+    ///
+    /// # Safety
+    /// `ptr` must point to a valid GPU allocation of at least `len` bytes that
+    /// will remain live for the lifetime of the returned `GpuBuffer`.
+    pub unsafe fn borrow_raw(ptr: *mut u8, len: usize) -> Self {
+        Self { ptr, len, owned: false, pooled: false }
     }
 }
 
 impl Drop for GpuBuffer {
     fn drop(&mut self) {
-        if !self.ptr.is_null() {
-            unsafe { ffi::cudaFree(self.ptr) };
+        if self.owned && !self.ptr.is_null() {
+            if self.pooled {
+                crate::allocator::return_to_pool(self.ptr, self.len);
+            } else {
+                unsafe { ffi::cudaFree(self.ptr) };
+            }
         }
     }
 }

crates/xserv-kernels/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "xserv-kernels"
+version.workspace = true
+edition.workspace = true
+
+[build-dependencies]
+cc = "1"
+
+[dependencies]
+xserv-cuda = { path = "../xserv-cuda" }
+xserv-tensor = { path = "../xserv-tensor" }
+half.workspace = true

crates/xserv-kernels/build.rs

@@ -0,0 +1,32 @@
+use std::env;
+
+fn main() {
+    let cuda_path = env::var("CUDA_HOME")
+        .or_else(|_| env::var("CUDA_PATH"))
+        .unwrap_or_else(|_| "/usr/local/cuda".to_string());
+
+    println!("cargo:rustc-link-search=native={cuda_path}/lib64");
+    println!("cargo:rustc-link-lib=dylib=cudart");
+    println!("cargo:rustc-link-lib=dylib=cublas");
+
+    cc::Build::new()
+        .cuda(true)
+        .cudart("shared")
+        .flag("-gencode=arch=compute_120,code=sm_120")
+        .include("../../csrc")
+        .file("../../csrc/gemm/naive.cu")
+        .file("../../csrc/gemm/tiled.cu")
+        .file("../../csrc/gemm/gemv.cu")
+        .file("../../csrc/normalization/rmsnorm.cu")
+        .file("../../csrc/normalization/layernorm.cu")
+        .file("../../csrc/activation/activations.cu")
+        .file("../../csrc/reduce/softmax.cu")
+        .file("../../csrc/embedding/embedding.cu")
+        .file("../../csrc/embedding/rope.cu")
+        .file("../../csrc/attention/causal_mask.cu")
+        .file("../../csrc/embedding/transpose.cu")
+        .file("../../csrc/attention/flash_attention.cu")
+        .compile("xserv_kernels");
+
+    println!("cargo:rerun-if-changed=../../csrc/");
+}

crates/xserv-kernels/src/activation.rs

@@ -0,0 +1,91 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_gelu_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_gelu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_silu_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    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);
+    fn launch_silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+}
+
+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::empty(x.shape(), x.dtype(), x.device());
+    let n = x.numel() as i32;
+    unsafe {
+        match x.dtype() {
+            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"),
+        }
+    }
+    out
+}
+
+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::empty(a.shape(), a.dtype(), a.device());
+    let n = a.numel() as i32;
+    unsafe {
+        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"),
+        }
+    }
+    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() && matches!(x.device(), Device::Cuda(_)));
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
+    let n = x.numel() as i32;
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_scale_f32(x.data_ptr() as _, out.data_ptr() as *mut c_void, scale_val, n, std::ptr::null_mut()),
+            DType::BF16 => launch_scale_bf16(x.data_ptr() as _, out.data_ptr() as *mut c_void, scale_val, n, std::ptr::null_mut()),
+            _ => panic!("unsupported dtype for scale"),
+        }
+    }
+    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) }
+
+/// Fused SiLU×Mul: out = silu(gate) * up (BF16 only)
+/// Saves one HBM read + one HBM write compared to separate silu + mul.
+pub fn silu_mul(gate: &Tensor, up: &Tensor) -> Tensor {
+    assert_eq!(gate.shape(), up.shape());
+    assert!(gate.is_contiguous() && up.is_contiguous());
+    assert!(matches!(gate.device(), Device::Cuda(_)));
+    assert_eq!(gate.dtype(), DType::BF16, "silu_mul requires BF16");
+    let out = Tensor::empty(gate.shape(), gate.dtype(), gate.device());
+    let n = gate.numel() as i32;
+    unsafe {
+        launch_silu_mul_bf16(
+            gate.data_ptr() as *const c_void,
+            up.data_ptr() as *const c_void,
+            out.data_ptr() as *mut c_void,
+            n,
+            std::ptr::null_mut(),
+        );
+    }
+    out
+}

crates/xserv-kernels/src/attention.rs

@@ -0,0 +1,194 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Tensor};
+
+use crate::activation::scale;
+use crate::gemm::batched_matmul;
+use crate::softmax::softmax;
+
+unsafe extern "C" {
+    fn launch_causal_mask_f32(scores: *mut c_void, batch: i32, rows: i32, cols: i32,
+                              offset: i32, stream: *mut c_void);
+    fn launch_causal_mask_bf16(scores: *mut c_void, batch: i32, rows: i32, cols: i32,
+                               offset: i32, stream: *mut c_void);
+    fn launch_flash_attention_bf16(
+        q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
+        batch: i32, num_q_heads: i32, num_kv_heads: i32,
+        q_len: i32, kv_len: i32, head_dim: i32,
+        scale: f32, causal: i32, stream: *mut c_void,
+    );
+    fn launch_decode_attention_bf16(
+        q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
+        batch: i32, num_q_heads: i32, num_kv_heads: i32,
+        kv_len: i32, head_dim: i32,
+        scale: f32, causal: i32, stream: *mut c_void,
+    );
+}
+
+fn apply_causal_mask(scores: &Tensor, offset: usize) {
+    let ndim = scores.ndim();
+    let rows = scores.shape()[ndim - 2];
+    let cols = scores.shape()[ndim - 1];
+    let batch: usize = scores.shape()[..ndim - 2].iter().product();
+
+    unsafe {
+        match scores.dtype() {
+            DType::F32 => launch_causal_mask_f32(
+                scores.data_ptr() as *mut c_void,
+                batch as i32, rows as i32, cols as i32, offset as i32,
+                std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_causal_mask_bf16(
+                scores.data_ptr() as *mut c_void,
+                batch as i32, rows as i32, cols as i32, offset as i32,
+                std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for causal mask"),
+        }
+    }
+}
+
+/// Multi-head attention (naive, materializes S×S score matrix).
+///
+/// q, k, v: [batch, num_heads, seq_len, head_dim] — contiguous, on GPU
+/// Returns: [batch, num_heads, seq_len, head_dim]
+pub fn attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tensor {
+    assert_eq!(q.ndim(), 4);
+    assert_eq!(k.ndim(), 4);
+    assert_eq!(v.ndim(), 4);
+    assert!(q.is_contiguous() && k.is_contiguous() && v.is_contiguous());
+
+    let batch = q.shape()[0];
+    let num_heads = q.shape()[1];
+    let q_len = q.shape()[2];
+    let head_dim = q.shape()[3];
+    let kv_len = k.shape()[2];
+
+    assert_eq!(k.shape(), &[batch, num_heads, kv_len, head_dim]);
+    assert_eq!(v.shape(), &[batch, num_heads, kv_len, head_dim]);
+
+    // scores = Q @ K^T → [B, H, q_len, kv_len]
+    let k_t = k.transpose(2, 3).contiguous();
+    let scores = batched_matmul(q, &k_t);
+
+    // Scale by 1/sqrt(head_dim)
+    let scale_factor = 1.0 / (head_dim as f32).sqrt();
+    let scaled_scores = scale(&scores, scale_factor);
+
+    // Causal mask
+    if causal {
+        let offset = kv_len - q_len;
+        apply_causal_mask(&scaled_scores, offset);
+    }
+
+    // Softmax
+    let weights = softmax(&scaled_scores);
+
+    // output = weights @ V → [B, H, q_len, head_dim]
+    batched_matmul(&weights, v)
+}
+
+/// Decode Attention — optimized for single-token decode (q_len=1).
+///
+/// q: [batch, num_q_heads, 1, head_dim]       BF16, contiguous, GPU
+/// k: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
+/// v: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
+///
+/// Returns: [batch, num_q_heads, 1, head_dim] BF16
+pub fn decode_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Tensor {
+    assert_eq!(q.ndim(), 4);
+    assert_eq!(q.shape()[2], 1, "decode_attention requires q_len == 1");
+
+    let batch = q.shape()[0];
+    let num_q_heads = q.shape()[1];
+    let head_dim = q.shape()[3];
+    let num_kv_heads = k.shape()[1];
+    let kv_len = k.shape()[2];
+
+    let scale = 1.0 / (head_dim as f32).sqrt();
+    let output = Tensor::empty(
+        &[batch, num_q_heads, 1, head_dim],
+        DType::BF16,
+        q.device(),
+    );
+
+    unsafe {
+        launch_decode_attention_bf16(
+            q.data_ptr() as *const c_void,
+            k.data_ptr() as *const c_void,
+            v.data_ptr() as *const c_void,
+            output.data_ptr() as *mut c_void,
+            batch as i32,
+            num_q_heads as i32,
+            num_kv_heads as i32,
+            kv_len as i32,
+            head_dim as i32,
+            scale,
+            1, // causal (always 1 for decode)
+            std::ptr::null_mut(),
+        );
+    }
+
+    output
+}
+
+/// Flash Attention 2 — O(1) extra memory, supports GQA natively.
+/// Auto-dispatches to decode_attention when q_len == 1.
+///
+/// q: [batch, num_q_heads, q_len, head_dim]   BF16, contiguous, GPU
+/// k: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
+/// v: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
+///
+/// Returns: [batch, num_q_heads, q_len, head_dim] BF16
+pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tensor {
+    assert_eq!(q.ndim(), 4);
+    assert_eq!(k.ndim(), 4);
+    assert_eq!(v.ndim(), 4);
+    assert!(q.is_contiguous() && k.is_contiguous() && v.is_contiguous());
+    assert_eq!(q.dtype(), DType::BF16, "flash_attention requires BF16");
+    assert_eq!(k.dtype(), DType::BF16);
+    assert_eq!(v.dtype(), DType::BF16);
+
+    let batch = q.shape()[0];
+    let num_q_heads = q.shape()[1];
+    let q_len = q.shape()[2];
+    let head_dim = q.shape()[3];
+    let num_kv_heads = k.shape()[1];
+    let kv_len = k.shape()[2];
+
+    assert_eq!(k.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
+    assert_eq!(v.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
+    assert!(num_q_heads % num_kv_heads == 0, "num_q_heads must be divisible by num_kv_heads");
+    assert!(head_dim <= 128, "flash_attention supports head_dim up to 128");
+
+    // Dispatch to specialized decode kernel for single-token generation
+    if q_len == 1 {
+        return decode_attention(q, k, v);
+    }
+
+    let scale = 1.0 / (head_dim as f32).sqrt();
+    let output = Tensor::empty(
+        &[batch, num_q_heads, q_len, head_dim],
+        DType::BF16,
+        q.device(),
+    );
+
+    unsafe {
+        launch_flash_attention_bf16(
+            q.data_ptr() as *const c_void,
+            k.data_ptr() as *const c_void,
+            v.data_ptr() as *const c_void,
+            output.data_ptr() as *mut c_void,
+            batch as i32,
+            num_q_heads as i32,
+            num_kv_heads as i32,
+            q_len as i32,
+            kv_len as i32,
+            head_dim as i32,
+            scale,
+            if causal { 1 } else { 0 },
+            std::ptr::null_mut(),
+        );
+    }
+
+    output
+}

crates/xserv-kernels/src/embedding.rs

@@ -0,0 +1,50 @@
+use std::ffi::c_void;
+use xserv_cuda::GpuBuffer;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_embedding_f32(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
+                            num_tokens: i32, hidden_size: i32, stream: *mut c_void);
+    fn launch_embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
+                             num_tokens: i32, hidden_size: i32, stream: *mut c_void);
+}
+
+/// Embedding lookup: table[token_ids[i]] for each i.
+/// table: [vocab_size, hidden_size], token_ids: [num_tokens] (i32 on CPU)
+pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
+    assert_eq!(table.ndim(), 2);
+    assert!(table.is_contiguous());
+    assert!(matches!(table.device(), Device::Cuda(_)));
+
+    let hidden_size = table.shape()[1];
+    let num_tokens = token_ids.len();
+
+    // Upload token_ids to GPU
+    let ids_bytes = unsafe {
+        std::slice::from_raw_parts(
+            token_ids.as_ptr() as *const u8,
+            num_tokens * std::mem::size_of::<u32>(),
+        )
+    };
+    let mut ids_gpu = xserv_cuda::allocator::cached_alloc(ids_bytes.len()).expect("alloc token_ids");
+    ids_gpu.copy_from_host(ids_bytes).unwrap();
+
+    let out = Tensor::empty(&[num_tokens, hidden_size], table.dtype(), table.device());
+
+    unsafe {
+        match table.dtype() {
+            DType::F32 => launch_embedding_f32(
+                table.data_ptr() as _, ids_gpu.as_ptr() as _,
+                out.data_ptr() as *mut c_void,
+                num_tokens as i32, hidden_size as i32, std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_embedding_bf16(
+                table.data_ptr() as _, ids_gpu.as_ptr() as _,
+                out.data_ptr() as *mut c_void,
+                num_tokens as i32, hidden_size as i32, std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for embedding"),
+        }
+    }
+    out
+}

crates/xserv-kernels/src/gemm.rs

@@ -0,0 +1,259 @@
+use std::cell::RefCell;
+use std::ffi::c_void;
+use xserv_cuda::error::{self, Result};
+use xserv_tensor::{DType, Device, Tensor};
+
+#[derive(Debug, Clone, Copy)]
+pub enum GemmBackend {
+    Naive,
+    Tiled,
+    CuBlas,
+}
+
+// --- FFI: custom CUDA kernels ---
+unsafe extern "C" {
+    fn launch_gemm_naive_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_naive_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_tiled_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_tiled_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void);
+}
+
+// --- FFI: cuBLAS ---
+type CublasHandle = *mut c_void;
+
+#[allow(non_upper_case_globals)]
+const CUBLAS_OP_N: i32 = 0;
+
+// cudaDataType
+const CUDA_R_32F: i32 = 0;
+const CUDA_R_16BF: i32 = 14;
+
+// cublasComputeType
+const CUBLAS_COMPUTE_32F: i32 = 68;
+
+unsafe extern "C" {
+    fn cublasCreate_v2(handle: *mut CublasHandle) -> i32;
+    fn cublasDestroy_v2(handle: CublasHandle) -> i32;
+    fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
+    fn cublasGemmEx(
+        handle: CublasHandle,
+        transa: i32, transb: i32,
+        m: i32, n: i32, k: i32,
+        alpha: *const c_void,
+        a: *const c_void, a_type: i32, lda: i32,
+        b: *const c_void, b_type: i32, ldb: i32,
+        beta: *const c_void,
+        c: *mut c_void, c_type: i32, ldc: i32,
+        compute_type: i32,
+        algo: i32,
+    ) -> i32;
+    fn cublasGemmStridedBatchedEx(
+        handle: CublasHandle,
+        transa: i32, transb: i32,
+        m: i32, n: i32, k: i32,
+        alpha: *const c_void,
+        a: *const c_void, a_type: i32, lda: i32, stride_a: i64,
+        b: *const c_void, b_type: i32, ldb: i32, stride_b: i64,
+        beta: *const c_void,
+        c: *mut c_void, c_type: i32, ldc: i32, stride_c: i64,
+        batch_count: i32,
+        compute_type: i32,
+        algo: i32,
+    ) -> i32;
+}
+
+pub struct CublasContext {
+    handle: CublasHandle,
+}
+
+impl CublasContext {
+    pub fn new() -> Result<Self> {
+        let mut handle = std::ptr::null_mut();
+        error::check(unsafe { cublasCreate_v2(&mut handle) })?;
+        Ok(Self { handle })
+    }
+}
+
+impl Drop for CublasContext {
+    fn drop(&mut self) {
+        if !self.handle.is_null() {
+            unsafe { cublasDestroy_v2(self.handle) };
+        }
+    }
+}
+
+thread_local! {
+    static CUBLAS_CTX: RefCell<CublasContext> = RefCell::new(
+        CublasContext::new().expect("failed to create thread-local cuBLAS handle")
+    );
+}
+
+/// Borrow the thread-local cuBLAS handle for the duration of a closure.
+fn with_cublas<F, R>(f: F) -> R
+where
+    F: FnOnce(CublasHandle) -> R,
+{
+    CUBLAS_CTX.with(|cell| {
+        let ctx = cell.borrow();
+        f(ctx.handle)
+    })
+}
+
+/// Matrix multiplication: C = A @ B
+/// A: [M, K], B: [K, N], C: [M, N]
+/// All tensors must be contiguous and on the same GPU.
+pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
+    assert_eq!(a.ndim(), 2);
+    assert_eq!(b.ndim(), 2);
+    assert_eq!(a.shape()[1], b.shape()[0], "inner dimension mismatch");
+    assert_eq!(a.dtype(), b.dtype(), "dtype mismatch");
+    assert!(a.is_contiguous() && b.is_contiguous(), "matmul requires contiguous tensors");
+    assert!(matches!(a.device(), Device::Cuda(_)), "matmul requires GPU tensors");
+
+    let m = a.shape()[0];
+    let k = a.shape()[1];
+    let n = b.shape()[1];
+    let dtype = a.dtype();
+
+    // All backends (naive, tiled, cuBLAS with beta=0, custom GEMV) fully
+    // overwrite every element of C, so we skip the cudaMemset.
+    let c = Tensor::empty(&[m, n], dtype, a.device());
+
+    let a_ptr = a.data_ptr() as *const c_void;
+    let b_ptr = b.data_ptr() as *const c_void;
+    let c_ptr = c.data_ptr() as *mut c_void;
+    let null_stream = std::ptr::null_mut();
+
+    match backend {
+        GemmBackend::Naive => {
+            unsafe {
+                match dtype {
+                    DType::F32 => launch_gemm_naive_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    DType::BF16 => launch_gemm_naive_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    _ => panic!("unsupported dtype for naive GEMM"),
+                }
+            }
+        }
+        GemmBackend::Tiled => {
+            unsafe {
+                match dtype {
+                    DType::F32 => launch_gemm_tiled_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    DType::BF16 => launch_gemm_tiled_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    _ => panic!("unsupported dtype for tiled GEMM"),
+                }
+            }
+        }
+        GemmBackend::CuBlas => {
+            // Fast path: custom GEMV for M=1 BF16 (bandwidth-optimal decode)
+            if m == 1 && dtype == DType::BF16 {
+                let mut fp32_buf = xserv_cuda::allocator::cached_alloc(n * 4).unwrap();
+                unsafe {
+                    launch_gemv_bf16(
+                        a_ptr, b_ptr, c_ptr,
+                        fp32_buf.as_mut_ptr() as *mut c_void,
+                        k as i32, n as i32,
+                        null_stream,
+                    );
+                }
+                // fp32_buf returned to caching allocator pool on drop
+            } else {
+                // cuBLAS uses column-major, but we have row-major tensors.
+                // Trick: compute C^T = B^T @ A^T, which gives us C in row-major.
+                // cuBLAS sees our row-major data as column-major transposed.
+                let alpha = 1.0f32;
+                let beta = 0.0f32;
+
+                let (a_type, b_type, c_type) = match dtype {
+                    DType::F32 => (CUDA_R_32F, CUDA_R_32F, CUDA_R_32F),
+                    DType::BF16 => (CUDA_R_16BF, CUDA_R_16BF, CUDA_R_16BF),
+                    _ => panic!("unsupported dtype for cuBLAS GEMM"),
+                };
+
+                with_cublas(|handle| unsafe {
+                    cublasSetStream_v2(handle, null_stream);
+                    // Row-major trick: swap A/B and transpose flags
+                    // C(row-major) = A @ B  <=>  C^T(col-major) = B^T @ A^T
+                    error::check(cublasGemmEx(
+                        handle,
+                        CUBLAS_OP_N, CUBLAS_OP_N,
+                        n as i32, m as i32, k as i32,
+                        &alpha as *const f32 as *const c_void,
+                        b_ptr, b_type, n as i32,  // B as col-major = B^T
+                        a_ptr, a_type, k as i32,  // A as col-major = A^T
+                        &beta as *const f32 as *const c_void,
+                        c_ptr, c_type, n as i32,  // C as col-major = C^T
+                        CUBLAS_COMPUTE_32F,
+                        -1, // default algo
+                    )).expect("cuBLAS GEMM failed");
+                });
+            }
+        }
+    }
+
+    c
+}
+
+/// Batched matrix multiplication via cuBLAS: C[b] = A[b] @ B[b]
+/// a: [..., M, K], b: [..., K, N] → [..., M, N]
+/// Leading dimensions must match and tensors must be contiguous.
+pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
+    assert!(a.ndim() >= 2 && b.ndim() >= 2);
+    assert_eq!(a.ndim(), b.ndim());
+    assert!(a.is_contiguous() && b.is_contiguous());
+    assert!(matches!(a.device(), Device::Cuda(_)));
+    assert_eq!(a.dtype(), b.dtype());
+
+    let ndim = a.ndim();
+    let m = a.shape()[ndim - 2];
+    let k = a.shape()[ndim - 1];
+    let n = b.shape()[ndim - 1];
+    assert_eq!(b.shape()[ndim - 2], k, "inner dimension mismatch");
+
+    // Compute batch count from leading dimensions
+    let batch: usize = a.shape()[..ndim - 2].iter().product();
+    assert_eq!(
+        b.shape()[..ndim - 2].iter().product::<usize>(),
+        batch,
+        "batch dimensions mismatch"
+    );
+
+    let mut out_shape: Vec<usize> = a.shape()[..ndim - 2].to_vec();
+    out_shape.push(m);
+    out_shape.push(n);
+    // cuBLAS with beta=0 fully overwrites every element of C.
+    let c = Tensor::empty(&out_shape, a.dtype(), a.device());
+
+    let dtype = a.dtype();
+    let (a_type, b_type, c_type) = match dtype {
+        DType::F32 => (CUDA_R_32F, CUDA_R_32F, CUDA_R_32F),
+        DType::BF16 => (CUDA_R_16BF, CUDA_R_16BF, CUDA_R_16BF),
+        _ => panic!("unsupported dtype for batched matmul"),
+    };
+
+    let alpha = 1.0f32;
+    let beta = 0.0f32;
+    // cuBLAS strides are in elements (not bytes)
+    let stride_a = (m * k) as i64;
+    let stride_b = (k * n) as i64;
+    let stride_c = (m * n) as i64;
+
+    with_cublas(|handle| unsafe {
+        cublasSetStream_v2(handle, std::ptr::null_mut());
+        // Row-major trick: C = A @ B ⟺ C^T = B^T @ A^T (col-major)
+        error::check(cublasGemmStridedBatchedEx(
+            handle,
+            CUBLAS_OP_N, CUBLAS_OP_N,
+            n as i32, m as i32, k as i32,
+            &alpha as *const f32 as *const c_void,
+            b.data_ptr() as _, b_type, n as i32, stride_b,
+            a.data_ptr() as _, a_type, k as i32, stride_a,
+            &beta as *const f32 as *const c_void,
+            c.data_ptr() as *mut c_void, c_type, n as i32, stride_c,
+            batch as i32,
+            CUBLAS_COMPUTE_32F,
+            -1,
+        )).expect("cuBLAS batched GEMM failed");
+    });
+    c
+}

crates/xserv-kernels/src/layernorm.rs

@@ -0,0 +1,38 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_layernorm_f32(x: *const c_void, gamma: *const c_void, beta: *const c_void,
+                            out: *mut c_void, rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+    fn launch_layernorm_bf16(x: *const c_void, gamma: *const c_void, beta: *const c_void,
+                             out: *mut c_void, rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+}
+
+pub fn layernorm(x: &Tensor, gamma: &Tensor, beta: &Tensor, eps: f32) -> Tensor {
+    assert!(x.ndim() >= 1);
+    assert!(x.is_contiguous() && gamma.is_contiguous() && beta.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    let hidden_size = *x.shape().last().unwrap();
+    assert_eq!(gamma.shape(), &[hidden_size]);
+    assert_eq!(beta.shape(), &[hidden_size]);
+
+    let rows = x.numel() / hidden_size;
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
+
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_layernorm_f32(
+                x.data_ptr() as _, gamma.data_ptr() as _, beta.data_ptr() as _,
+                out.data_ptr() as *mut c_void,
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_layernorm_bf16(
+                x.data_ptr() as _, gamma.data_ptr() as _, beta.data_ptr() as _,
+                out.data_ptr() as *mut c_void,
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for layernorm"),
+        }
+    }
+    out
+}

crates/xserv-kernels/src/lib.rs

@@ -0,0 +1,24 @@
+pub mod activation;
+pub mod attention;
+pub mod embedding;
+pub mod gemm;
+pub mod layernorm;
+pub mod rmsnorm;
+pub mod rope;
+pub mod softmax;
+pub mod transpose;
+
+pub use activation::{add, gelu, mul, scale, silu, silu_mul};
+pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
+pub use attention::{attention, decode_attention, flash_attention};
+pub use embedding::embedding;
+pub use gemm::{batched_matmul, matmul, GemmBackend};
+pub use layernorm::layernorm;
+pub use rmsnorm::{add_rmsnorm, rmsnorm};
+pub use rope::{rope_inplace, RopeCache};
+pub use softmax::softmax;
+
+/// Register GPU kernels with the tensor crate. Call once at startup.
+pub fn init() {
+    xserv_tensor::register_gpu_contiguous(strided_to_contiguous_gpu);
+}

crates/xserv-kernels/src/rmsnorm.rs

@@ -0,0 +1,75 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_rmsnorm_f32(x: *const c_void, gamma: *const c_void, out: *mut c_void,
+                          rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+    fn launch_rmsnorm_bf16(x: *const c_void, gamma: *const c_void, out: *mut c_void,
+                           rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+    fn launch_add_rmsnorm_bf16(x: *const c_void, residual: *const c_void, gamma: *const c_void,
+                               normed_out: *mut c_void, sum_out: *mut c_void,
+                               rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+}
+
+pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
+    assert!(x.ndim() >= 1);
+    assert!(x.is_contiguous() && gamma.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    let hidden_size = *x.shape().last().unwrap();
+    assert_eq!(gamma.shape(), &[hidden_size]);
+    assert_eq!(x.dtype(), gamma.dtype());
+
+    let rows = x.numel() / hidden_size;
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
+
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_rmsnorm_f32(
+                x.data_ptr() as _, gamma.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_rmsnorm_bf16(
+                x.data_ptr() as _, gamma.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for rmsnorm"),
+        }
+    }
+    out
+}
+
+/// Fused Add + RMSNorm: computes sum = x + residual, then normed = rmsnorm(sum, gamma, eps).
+/// Returns (normed, sum). BF16 only.
+/// Saves one kernel launch and one full HBM round-trip per layer.
+pub fn add_rmsnorm(x: &Tensor, residual: &Tensor, gamma: &Tensor, eps: f32) -> (Tensor, Tensor) {
+    assert!(x.ndim() >= 1);
+    assert_eq!(x.shape(), residual.shape());
+    assert!(x.is_contiguous() && residual.is_contiguous() && gamma.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    assert_eq!(x.dtype(), DType::BF16, "add_rmsnorm requires BF16");
+    assert_eq!(residual.dtype(), DType::BF16);
+    assert_eq!(gamma.dtype(), DType::BF16);
+
+    let hidden_size = *x.shape().last().unwrap();
+    assert_eq!(gamma.shape(), &[hidden_size]);
+
+    let rows = x.numel() / hidden_size;
+    let normed_out = Tensor::empty(x.shape(), DType::BF16, x.device());
+    let sum_out = Tensor::empty(x.shape(), DType::BF16, x.device());
+
+    unsafe {
+        launch_add_rmsnorm_bf16(
+            x.data_ptr() as *const c_void,
+            residual.data_ptr() as *const c_void,
+            gamma.data_ptr() as *const c_void,
+            normed_out.data_ptr() as *mut c_void,
+            sum_out.data_ptr() as *mut c_void,
+            rows as i32,
+            hidden_size as i32,
+            eps,
+            std::ptr::null_mut(),
+        );
+    }
+
+    (normed_out, sum_out)
+}

crates/xserv-kernels/src/rope.rs

@@ -0,0 +1,83 @@
+use std::ffi::c_void;
+use xserv_cuda::GpuBuffer;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_rope_f32(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
+                       positions: *const c_void, num_tokens: i32, num_heads: i32,
+                       head_dim: i32, stream: *mut c_void);
+    fn launch_rope_bf16(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
+                        positions: *const c_void, num_tokens: i32, num_heads: i32,
+                        head_dim: i32, stream: *mut c_void);
+    fn launch_compute_rope_cache(cos_cache: *mut c_void, sin_cache: *mut c_void,
+                                 max_seq_len: i32, half_dim: i32, theta: f32,
+                                 stream: *mut c_void);
+}
+
+pub struct RopeCache {
+    pub cos: GpuBuffer,
+    pub sin: GpuBuffer,
+    pub max_seq_len: usize,
+    pub half_dim: usize,
+}
+
+impl RopeCache {
+    pub fn new(max_seq_len: usize, head_dim: usize, theta: f32) -> Self {
+        let half_dim = head_dim / 2;
+        let nbytes = max_seq_len * half_dim * std::mem::size_of::<f32>();
+        let mut cos = GpuBuffer::alloc(nbytes).expect("alloc cos_cache");
+        let mut sin = GpuBuffer::alloc(nbytes).expect("alloc sin_cache");
+
+        unsafe {
+            launch_compute_rope_cache(
+                cos.as_mut_ptr() as _, sin.as_mut_ptr() as _,
+                max_seq_len as i32, half_dim as i32, theta, std::ptr::null_mut(),
+            );
+        }
+
+        Self { cos, sin, max_seq_len, half_dim }
+    }
+}
+
+/// Apply RoPE in-place to x.
+/// x: [num_tokens, num_heads, head_dim] on GPU
+/// positions: [num_tokens] (u32 on CPU, will be uploaded)
+pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) {
+    assert_eq!(x.ndim(), 3);
+    assert!(x.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    let num_tokens = x.shape()[0];
+    let num_heads = x.shape()[1];
+    let head_dim = x.shape()[2];
+    assert_eq!(head_dim / 2, cache.half_dim);
+    assert_eq!(positions.len(), num_tokens);
+
+    let pos_bytes = unsafe {
+        std::slice::from_raw_parts(
+            positions.as_ptr() as *const u8,
+            num_tokens * std::mem::size_of::<u32>(),
+        )
+    };
+    let mut pos_gpu = xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions");
+    pos_gpu.copy_from_host(pos_bytes).unwrap();
+
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_rope_f32(
+                x.data_ptr() as *mut c_void,
+                cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
+                pos_gpu.as_ptr() as _,
+                num_tokens as i32, num_heads as i32, head_dim as i32,
+                std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_rope_bf16(
+                x.data_ptr() as *mut c_void,
+                cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
+                pos_gpu.as_ptr() as _,
+                num_tokens as i32, num_heads as i32, head_dim as i32,
+                std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for rope"),
+        }
+    }
+}

crates/xserv-kernels/src/softmax.rs

@@ -0,0 +1,33 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_softmax_f32(x: *const c_void, out: *mut c_void, rows: i32, cols: i32, stream: *mut c_void);
+    fn launch_softmax_bf16(x: *const c_void, out: *mut c_void, rows: i32, cols: i32, stream: *mut c_void);
+}
+
+/// Softmax along the last dimension.
+pub fn softmax(x: &Tensor) -> Tensor {
+    assert!(x.ndim() >= 1);
+    assert!(x.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+
+    let cols = *x.shape().last().unwrap();
+    let rows = x.numel() / cols;
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
+
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_softmax_f32(
+                x.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, cols as i32, std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_softmax_bf16(
+                x.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, cols as i32, std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for softmax"),
+        }
+    }
+    out
+}

crates/xserv-kernels/src/transpose.rs

@@ -0,0 +1,142 @@
+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);
+    fn launch_strided_copy_bf16(inp: *const c_void, out: *mut c_void, numel: i32, ndim: i32,
+        shape0: i32, shape1: i32, shape2: i32, shape3: i32,
+        in_stride0: i32, in_stride1: i32, in_stride2: i32, in_stride3: i32,
+        in_offset: i32, stream: *mut c_void);
+    fn launch_strided_copy_f32(inp: *const c_void, out: *mut c_void, numel: i32, ndim: i32,
+        shape0: i32, shape1: i32, shape2: i32, shape3: i32,
+        in_stride0: i32, in_stride1: i32, in_stride2: i32, in_stride3: i32,
+        in_offset: 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::empty(&[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(),
+        );
+    }
+    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::empty(&[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(),
+        );
+    }
+    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::empty(&[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(),
+        );
+    }
+    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::empty(&[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(),
+        );
+    }
+    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::empty(&[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(),
+        );
+    }
+    out
+}
+
+/// Make a non-contiguous GPU tensor contiguous via a strided copy kernel.
+/// Supports BF16 and F32, up to 4D tensors (padded to 4D internally).
+pub fn strided_to_contiguous_gpu(x: &Tensor) -> Tensor {
+    assert!(matches!(x.device(), Device::Cuda(_)), "expected GPU tensor");
+    assert!(!x.is_contiguous(), "tensor is already contiguous");
+    assert!(x.ndim() <= 4, "strided_to_contiguous_gpu supports up to 4D");
+
+    let ndim = x.ndim();
+    let numel = x.numel();
+
+    // Pad shape and strides to 4D (prepend 1s for shape, 0s for strides)
+    let mut shape4 = [1i32; 4];
+    let mut strides4 = [0i32; 4];
+    let pad = 4 - ndim;
+    for i in 0..ndim {
+        shape4[pad + i] = x.shape()[i] as i32;
+        strides4[pad + i] = x.strides()[i] as i32;
+    }
+
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
+
+    // Use storage base pointer + element offset, because strides are relative to
+    // element 0 of the storage, not the data_ptr() (which already adds byte offset).
+    let storage_ptr = x.storage().gpu_buffer().as_ptr();
+    let in_offset = x.offset() as i32;
+
+    unsafe {
+        match x.dtype() {
+            DType::BF16 => launch_strided_copy_bf16(
+                storage_ptr as _, out.data_ptr() as *mut c_void,
+                numel as i32, ndim as i32,
+                shape4[0], shape4[1], shape4[2], shape4[3],
+                strides4[0], strides4[1], strides4[2], strides4[3],
+                in_offset, std::ptr::null_mut(),
+            ),
+            DType::F32 => launch_strided_copy_f32(
+                storage_ptr as _, out.data_ptr() as *mut c_void,
+                numel as i32, ndim as i32,
+                shape4[0], shape4[1], shape4[2], shape4[3],
+                strides4[0], strides4[1], strides4[2], strides4[3],
+                in_offset, std::ptr::null_mut(),
+            ),
+            _ => panic!("strided_to_contiguous_gpu: unsupported dtype {:?}", x.dtype()),
+        }
+    }
+    out
+}

crates/xserv-kernels/tests/attention_test.rs

@@ -0,0 +1,187 @@
+use xserv_kernels::*;
+use xserv_tensor::{Device, Tensor};
+
+fn init() { xserv_cuda::device::set_device(0).unwrap(); }
+
+fn cpu_attention(q: &[f32], k: &[f32], v: &[f32],
+                 batch: usize, heads: usize, q_len: usize, kv_len: usize, head_dim: usize,
+                 causal: bool) -> Vec<f32> {
+    let mut out = vec![0.0f32; batch * heads * q_len * head_dim];
+    let scale = 1.0 / (head_dim as f32).sqrt();
+
+    for b in 0..batch {
+        for h in 0..heads {
+            // scores = Q @ K^T, scaled
+            let mut scores = vec![0.0f32; q_len * kv_len];
+            for i in 0..q_len {
+                for j in 0..kv_len {
+                    let mut s = 0.0f32;
+                    for d in 0..head_dim {
+                        let qi = q[((b * heads + h) * q_len + i) * head_dim + d];
+                        let ki = k[((b * heads + h) * kv_len + j) * head_dim + d];
+                        s += qi * ki;
+                    }
+                    scores[i * kv_len + j] = s * scale;
+                }
+            }
+            // causal mask
+            if causal {
+                let offset = kv_len - q_len;
+                for i in 0..q_len {
+                    for j in 0..kv_len {
+                        if j > i + offset {
+                            scores[i * kv_len + j] = f32::NEG_INFINITY;
+                        }
+                    }
+                }
+            }
+            // softmax per row
+            for i in 0..q_len {
+                let row = &mut scores[i * kv_len..(i + 1) * kv_len];
+                let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+                let mut sum = 0.0f32;
+                for v in row.iter_mut() {
+                    *v = (*v - max).exp();
+                    sum += *v;
+                }
+                for v in row.iter_mut() {
+                    *v /= sum;
+                }
+            }
+            // output = weights @ V
+            for i in 0..q_len {
+                for d in 0..head_dim {
+                    let mut s = 0.0f32;
+                    for j in 0..kv_len {
+                        let w = scores[i * kv_len + j];
+                        let vi = v[((b * heads + h) * kv_len + j) * head_dim + d];
+                        s += w * vi;
+                    }
+                    out[((b * heads + h) * q_len + i) * head_dim + d] = s;
+                }
+            }
+        }
+    }
+    out
+}
+
+fn check_close(a: &[f32], b: &[f32], atol: f32, name: &str) {
+    assert_eq!(a.len(), b.len(), "{name}: length mismatch");
+    let mut max_err = 0.0f32;
+    for (i, (x, y)) in a.iter().zip(b).enumerate() {
+        let err = (x - y).abs();
+        if err > max_err { max_err = err; }
+        assert!(err <= atol, "{name}: mismatch at [{i}]: got {x}, expected {y}, err {err}");
+    }
+    println!("{name}: max_err = {max_err:.6e}");
+}
+
+fn make_data(n: usize) -> Vec<f32> {
+    (0..n).map(|i| ((i % 17) as f32 - 8.0) * 0.05).collect()
+}
+
+#[test]
+fn test_batched_matmul() {
+    init();
+    let batch = 4;
+    let heads = 8;
+    let m = 32;
+    let k = 64;
+    let n = 32;
+
+    let a_data = make_data(batch * heads * m * k);
+    let b_data = make_data(batch * heads * k * n);
+
+    let a = Tensor::from_slice(&a_data, &[batch, heads, m, k]).to_device(Device::Cuda(0));
+    let b = Tensor::from_slice(&b_data, &[batch, heads, k, n]).to_device(Device::Cuda(0));
+    let c = batched_matmul(&a, &b).to_device(Device::Cpu);
+
+    assert_eq!(c.shape(), &[batch, heads, m, n]);
+
+    // Verify one batch element
+    let a_cpu = &a_data[0..m * k];
+    let b_cpu = &b_data[0..k * n];
+    let mut expected = vec![0.0f32; m * n];
+    for i in 0..m {
+        for j in 0..n {
+            let mut s = 0.0f32;
+            for kk in 0..k { s += a_cpu[i * k + kk] * b_cpu[kk * n + j]; }
+            expected[i * n + j] = s;
+        }
+    }
+    let result = c.as_slice::<f32>();
+    check_close(&result[0..m * n], &expected, 1e-3, "batched_matmul[0]");
+}
+
+#[test]
+fn test_attention_no_causal() {
+    init();
+    let b = 1; let h = 2; let s = 8; let d = 16;
+    let q_data = make_data(b * h * s * d);
+    let k_data = make_data(b * h * s * d);
+    let v_data = make_data(b * h * s * d);
+    let expected = cpu_attention(&q_data, &k_data, &v_data, b, h, s, s, d, false);
+
+    let q = Tensor::from_slice(&q_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let k = Tensor::from_slice(&k_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let v = Tensor::from_slice(&v_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let out = attention(&q, &k, &v, false).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-4, "attention_no_causal");
+}
+
+#[test]
+fn test_attention_causal() {
+    init();
+    let b = 1; let h = 2; let s = 16; let d = 32;
+    let q_data = make_data(b * h * s * d);
+    let k_data = make_data(b * h * s * d);
+    let v_data = make_data(b * h * s * d);
+    let expected = cpu_attention(&q_data, &k_data, &v_data, b, h, s, s, d, true);
+
+    let q = Tensor::from_slice(&q_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let k = Tensor::from_slice(&k_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let v = Tensor::from_slice(&v_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let out = attention(&q, &k, &v, true).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-3, "attention_causal");
+}
+
+#[test]
+fn test_attention_causal_larger() {
+    init();
+    let b = 2; let h = 4; let s = 64; let d = 64;
+    let q_data = make_data(b * h * s * d);
+    let k_data = make_data(b * h * s * d);
+    let v_data = make_data(b * h * s * d);
+    let expected = cpu_attention(&q_data, &k_data, &v_data, b, h, s, s, d, true);
+
+    let q = Tensor::from_slice(&q_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let k = Tensor::from_slice(&k_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let v = Tensor::from_slice(&v_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let out = attention(&q, &k, &v, true).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-2, "attention_causal_larger");
+}
+
+#[test]
+fn test_attention_causal_first_row_sees_only_first_token() {
+    init();
+    let b = 1; let h = 1; let s = 4; let d = 8;
+    let q_data = make_data(b * h * s * d);
+    let k_data = make_data(b * h * s * d);
+    let v_data: Vec<f32> = (0..s * d).map(|i| {
+        if i < d { 1.0 } else { 0.0 }  // only first V row is nonzero
+    }).collect();
+
+    let q = Tensor::from_slice(&q_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let k = Tensor::from_slice(&k_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let v = Tensor::from_slice(&v_data, &[b, h, s, d]).to_device(Device::Cuda(0));
+    let out = attention(&q, &k, &v, true).to_device(Device::Cpu);
+
+    // First row (position 0) with causal mask can only see position 0.
+    // So attention weight for position 0 is 1.0 for token 0 only.
+    // output[0] should be exactly V[0] = [1, 1, 1, ...1]
+    let result = out.as_slice::<f32>();
+    for i in 0..d {
+        assert!((result[i] - 1.0).abs() < 1e-5,
+            "first row should equal V[0], got {} at dim {}", result[i], i);
+    }
+}

crates/xserv-kernels/tests/gemm_test.rs

@@ -0,0 +1,166 @@
+use half::bf16;
+use xserv_kernels::{matmul, GemmBackend};
+use xserv_tensor::{Device, Tensor};
+
+fn cpu_matmul_f32(a: &[f32], b: &[f32], m: usize, n: usize, k: usize) -> Vec<f32> {
+    let mut c = vec![0.0f32; m * n];
+    for i in 0..m {
+        for j in 0..n {
+            let mut sum = 0.0f32;
+            for kk in 0..k {
+                sum += a[i * k + kk] * b[kk * n + j];
+            }
+            c[i * n + j] = sum;
+        }
+    }
+    c
+}
+
+fn check_close_f32(result: &[f32], expected: &[f32], atol: f32) {
+    assert_eq!(result.len(), expected.len());
+    for (i, (r, e)) in result.iter().zip(expected).enumerate() {
+        assert!(
+            (r - e).abs() <= atol,
+            "mismatch at index {i}: got {r}, expected {e}, diff {}",
+            (r - e).abs()
+        );
+    }
+}
+
+fn check_close_bf16(result: &[bf16], expected: &[f32], atol: f32) {
+    assert_eq!(result.len(), expected.len());
+    for (i, (r, e)) in result.iter().zip(expected).enumerate() {
+        let rv = r.to_f32();
+        assert!(
+            (rv - e).abs() <= atol,
+            "mismatch at index {i}: got {rv}, expected {e}, diff {}",
+            (rv - e).abs()
+        );
+    }
+}
+
+fn run_gemm_test_f32(backend: GemmBackend, m: usize, n: usize, k: usize) {
+    xserv_cuda::device::set_device(0).unwrap();
+
+    let a_data: Vec<f32> = (0..m * k).map(|i| ((i % 7) as f32 - 3.0) * 0.1).collect();
+    let b_data: Vec<f32> = (0..k * n).map(|i| ((i % 11) as f32 - 5.0) * 0.1).collect();
+    let expected = cpu_matmul_f32(&a_data, &b_data, m, n, k);
+
+    let a = Tensor::from_slice(&a_data, &[m, k]).to_device(Device::Cuda(0));
+    let b = Tensor::from_slice(&b_data, &[k, n]).to_device(Device::Cuda(0));
+    let c = matmul(&a, &b, backend);
+
+    let c_cpu = c.to_device(Device::Cpu);
+    check_close_f32(c_cpu.as_slice::<f32>(), &expected, 1e-4);
+}
+
+fn run_gemm_test_bf16(backend: GemmBackend, m: usize, n: usize, k: usize) {
+    xserv_cuda::device::set_device(0).unwrap();
+
+    let a_f32: Vec<f32> = (0..m * k).map(|i| ((i % 7) as f32 - 3.0) * 0.1).collect();
+    let b_f32: Vec<f32> = (0..k * n).map(|i| ((i % 11) as f32 - 5.0) * 0.1).collect();
+    let expected = cpu_matmul_f32(&a_f32, &b_f32, m, n, k);
+
+    let a_data: Vec<bf16> = a_f32.iter().map(|&v| bf16::from_f32(v)).collect();
+    let b_data: Vec<bf16> = b_f32.iter().map(|&v| bf16::from_f32(v)).collect();
+
+    let a = Tensor::from_slice(&a_data, &[m, k]).to_device(Device::Cuda(0));
+    let b = Tensor::from_slice(&b_data, &[k, n]).to_device(Device::Cuda(0));
+    let c = matmul(&a, &b, backend);
+
+    let c_cpu = c.to_device(Device::Cpu);
+    check_close_bf16(c_cpu.as_slice::<bf16>(), &expected, 0.1);
+}
+
+// --- F32 tests ---
+
+#[test]
+fn test_gemm_naive_f32_small() { run_gemm_test_f32(GemmBackend::Naive, 4, 4, 4); }
+
+#[test]
+fn test_gemm_naive_f32_medium() { run_gemm_test_f32(GemmBackend::Naive, 64, 64, 64); }
+
+#[test]
+fn test_gemm_naive_f32_rect() { run_gemm_test_f32(GemmBackend::Naive, 32, 64, 48); }
+
+#[test]
+fn test_gemm_tiled_f32_small() { run_gemm_test_f32(GemmBackend::Tiled, 4, 4, 4); }
+
+#[test]
+fn test_gemm_tiled_f32_medium() { run_gemm_test_f32(GemmBackend::Tiled, 128, 128, 128); }
+
+#[test]
+fn test_gemm_tiled_f32_rect() { run_gemm_test_f32(GemmBackend::Tiled, 65, 33, 97); }
+
+#[test]
+fn test_gemm_cublas_f32_small() { run_gemm_test_f32(GemmBackend::CuBlas, 4, 4, 4); }
+
+#[test]
+fn test_gemm_cublas_f32_medium() { run_gemm_test_f32(GemmBackend::CuBlas, 256, 256, 256); }
+
+#[test]
+fn test_gemm_cublas_f32_rect() { run_gemm_test_f32(GemmBackend::CuBlas, 65, 33, 97); }
+
+// --- BF16 tests ---
+
+#[test]
+fn test_gemm_naive_bf16_small() { run_gemm_test_bf16(GemmBackend::Naive, 4, 4, 4); }
+
+#[test]
+fn test_gemm_naive_bf16_medium() { run_gemm_test_bf16(GemmBackend::Naive, 64, 64, 64); }
+
+#[test]
+fn test_gemm_tiled_bf16_small() { run_gemm_test_bf16(GemmBackend::Tiled, 4, 4, 4); }
+
+#[test]
+fn test_gemm_tiled_bf16_medium() { run_gemm_test_bf16(GemmBackend::Tiled, 128, 128, 128); }
+
+#[test]
+fn test_gemm_cublas_bf16_small() { run_gemm_test_bf16(GemmBackend::CuBlas, 4, 4, 4); }
+
+#[test]
+fn test_gemm_cublas_bf16_medium() { run_gemm_test_bf16(GemmBackend::CuBlas, 256, 256, 256); }
+
+// --- Custom GEMV tests (M=1, BF16 fast path) ---
+
+#[test]
+fn test_gemv_bf16_small() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 64, 64); }
+
+#[test]
+fn test_gemv_bf16_medium() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 256, 256); }
+
+#[test]
+fn test_gemv_bf16_4096() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 4096, 4096); }
+
+#[test]
+fn test_gemv_bf16_rect() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 512, 4096); }
+
+// --- Larger benchmark-style tests ---
+
+#[test]
+fn test_gemm_cublas_f32_1024() { run_gemm_test_f32(GemmBackend::CuBlas, 1024, 1024, 1024); }
+
+#[test]
+fn test_gemm_consistency_all_backends() {
+    xserv_cuda::device::set_device(0).unwrap();
+
+    let m = 64;
+    let n = 64;
+    let k = 64;
+    let a_data: Vec<f32> = (0..m * k).map(|i| ((i % 7) as f32 - 3.0) * 0.1).collect();
+    let b_data: Vec<f32> = (0..k * n).map(|i| ((i % 11) as f32 - 5.0) * 0.1).collect();
+
+    let a = Tensor::from_slice(&a_data, &[m, k]).to_device(Device::Cuda(0));
+    let b = Tensor::from_slice(&b_data, &[k, n]).to_device(Device::Cuda(0));
+
+    let c_naive = matmul(&a, &b, GemmBackend::Naive).to_device(Device::Cpu);
+    let c_tiled = matmul(&a, &b, GemmBackend::Tiled).to_device(Device::Cpu);
+    let c_cublas = matmul(&a, &b, GemmBackend::CuBlas).to_device(Device::Cpu);
+
+    let naive = c_naive.as_slice::<f32>();
+    let tiled = c_tiled.as_slice::<f32>();
+    let cublas = c_cublas.as_slice::<f32>();
+
+    check_close_f32(naive, cublas, 1e-4);
+    check_close_f32(tiled, cublas, 1e-4);
+}

crates/xserv-kernels/tests/ops_test.rs

@@ -0,0 +1,302 @@
+use half::bf16;
+use xserv_kernels::*;
+use xserv_tensor::{Device, Tensor};
+
+fn init() { xserv_cuda::device::set_device(0).unwrap(); }
+
+// --- CPU reference implementations ---
+
+fn cpu_rmsnorm(x: &[f32], gamma: &[f32], eps: f32, hidden: usize) -> Vec<f32> {
+    let rows = x.len() / hidden;
+    let mut out = vec![0.0f32; x.len()];
+    for r in 0..rows {
+        let row = &x[r * hidden..(r + 1) * hidden];
+        let sum_sq: f32 = row.iter().map(|v| v * v).sum();
+        let rms_inv = 1.0 / (sum_sq / hidden as f32 + eps).sqrt();
+        for i in 0..hidden {
+            out[r * hidden + i] = row[i] * rms_inv * gamma[i];
+        }
+    }
+    out
+}
+
+fn cpu_layernorm(x: &[f32], gamma: &[f32], beta: &[f32], eps: f32, hidden: usize) -> Vec<f32> {
+    let rows = x.len() / hidden;
+    let mut out = vec![0.0f32; x.len()];
+    for r in 0..rows {
+        let row = &x[r * hidden..(r + 1) * hidden];
+        let mean: f32 = row.iter().sum::<f32>() / hidden as f32;
+        let var: f32 = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / hidden as f32;
+        let inv_std = 1.0 / (var + eps).sqrt();
+        for i in 0..hidden {
+            out[r * hidden + i] = gamma[i] * (row[i] - mean) * inv_std + beta[i];
+        }
+    }
+    out
+}
+
+fn cpu_gelu(x: &[f32]) -> Vec<f32> {
+    let sqrt_2_over_pi = 0.7978845608f32;
+    x.iter().map(|&v| {
+        let inner = sqrt_2_over_pi * (v + 0.044715 * v * v * v);
+        0.5 * v * (1.0 + inner.tanh())
+    }).collect()
+}
+
+fn cpu_silu(x: &[f32]) -> Vec<f32> {
+    x.iter().map(|&v| v / (1.0 + (-v).exp())).collect()
+}
+
+fn cpu_softmax(x: &[f32], cols: usize) -> Vec<f32> {
+    let rows = x.len() / cols;
+    let mut out = vec![0.0f32; x.len()];
+    for r in 0..rows {
+        let row = &x[r * cols..(r + 1) * cols];
+        let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+        let exps: Vec<f32> = row.iter().map(|v| (v - max).exp()).collect();
+        let sum: f32 = exps.iter().sum();
+        for i in 0..cols {
+            out[r * cols + i] = exps[i] / sum;
+        }
+    }
+    out
+}
+
+fn cpu_rope(x: &mut [f32], positions: &[u32], num_heads: usize, head_dim: usize, theta: f32) {
+    let half_dim = head_dim / 2;
+    let num_tokens = positions.len();
+    for t in 0..num_tokens {
+        let pos = positions[t] as f32;
+        for h in 0..num_heads {
+            for i in 0..half_dim {
+                let freq = 1.0 / theta.powf(2.0 * i as f32 / head_dim as f32);
+                let angle = pos * freq;
+                let cos_val = angle.cos();
+                let sin_val = angle.sin();
+                let base = (t * num_heads + h) * head_dim;
+                let x0 = x[base + 2 * i];
+                let x1 = x[base + 2 * i + 1];
+                x[base + 2 * i]     = x0 * cos_val - x1 * sin_val;
+                x[base + 2 * i + 1] = x0 * sin_val + x1 * cos_val;
+            }
+        }
+    }
+}
+
+fn check_close(result: &[f32], expected: &[f32], atol: f32, name: &str) {
+    assert_eq!(result.len(), expected.len(), "{name}: length mismatch");
+    let mut max_err = 0.0f32;
+    for (i, (r, e)) in result.iter().zip(expected).enumerate() {
+        let err = (r - e).abs();
+        if err > max_err { max_err = err; }
+        assert!(err <= atol, "{name}: mismatch at [{i}]: got {r}, expected {e}, err {err}");
+    }
+    println!("{name}: max_err = {max_err:.6e}");
+}
+
+fn make_data(n: usize) -> Vec<f32> {
+    (0..n).map(|i| ((i % 17) as f32 - 8.0) * 0.1).collect()
+}
+
+// === RMSNorm ===
+
+#[test]
+fn test_rmsnorm_f32() {
+    init();
+    let hidden = 768;
+    let rows = 4;
+    let x_data = make_data(rows * hidden);
+    let gamma_data: Vec<f32> = (0..hidden).map(|i| 0.5 + (i % 3) as f32 * 0.2).collect();
+    let expected = cpu_rmsnorm(&x_data, &gamma_data, 1e-5, hidden);
+
+    let x = Tensor::from_slice(&x_data, &[rows, hidden]).to_device(Device::Cuda(0));
+    let gamma = Tensor::from_slice(&gamma_data, &[hidden]).to_device(Device::Cuda(0));
+    let out = rmsnorm(&x, &gamma, 1e-5).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-4, "rmsnorm_f32");
+}
+
+#[test]
+fn test_rmsnorm_bf16() {
+    init();
+    let hidden = 768;
+    let rows = 4;
+    let x_f32 = make_data(rows * hidden);
+    let gamma_f32: Vec<f32> = (0..hidden).map(|i| 0.5 + (i % 3) as f32 * 0.2).collect();
+    let expected = cpu_rmsnorm(&x_f32, &gamma_f32, 1e-5, hidden);
+
+    let x_bf16: Vec<bf16> = x_f32.iter().map(|&v| bf16::from_f32(v)).collect();
+    let gamma_bf16: Vec<bf16> = gamma_f32.iter().map(|&v| bf16::from_f32(v)).collect();
+    let x = Tensor::from_slice(&x_bf16, &[rows, hidden]).to_device(Device::Cuda(0));
+    let gamma = Tensor::from_slice(&gamma_bf16, &[hidden]).to_device(Device::Cuda(0));
+    let out = rmsnorm(&x, &gamma, 1e-5).to_device(Device::Cpu);
+
+    let result: Vec<f32> = out.as_slice::<bf16>().iter().map(|v| v.to_f32()).collect();
+    check_close(&result, &expected, 0.05, "rmsnorm_bf16");
+}
+
+// === LayerNorm ===
+
+#[test]
+fn test_layernorm_f32() {
+    init();
+    let hidden = 768;
+    let rows = 4;
+    let x_data = make_data(rows * hidden);
+    let gamma_data: Vec<f32> = (0..hidden).map(|i| 0.8 + (i % 5) as f32 * 0.1).collect();
+    let beta_data: Vec<f32> = (0..hidden).map(|i| ((i % 7) as f32 - 3.0) * 0.01).collect();
+    let expected = cpu_layernorm(&x_data, &gamma_data, &beta_data, 1e-5, hidden);
+
+    let x = Tensor::from_slice(&x_data, &[rows, hidden]).to_device(Device::Cuda(0));
+    let gamma = Tensor::from_slice(&gamma_data, &[hidden]).to_device(Device::Cuda(0));
+    let beta = Tensor::from_slice(&beta_data, &[hidden]).to_device(Device::Cuda(0));
+    let out = layernorm(&x, &gamma, &beta, 1e-5).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-4, "layernorm_f32");
+}
+
+// === GELU ===
+
+#[test]
+fn test_gelu_f32() {
+    init();
+    let data = make_data(10000);
+    let expected = cpu_gelu(&data);
+    let x = Tensor::from_slice(&data, &[10000]).to_device(Device::Cuda(0));
+    let out = gelu(&x).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-5, "gelu_f32");
+}
+
+#[test]
+fn test_gelu_bf16() {
+    init();
+    let data_f32 = make_data(10000);
+    let expected = cpu_gelu(&data_f32);
+    let data_bf16: Vec<bf16> = data_f32.iter().map(|&v| bf16::from_f32(v)).collect();
+    let x = Tensor::from_slice(&data_bf16, &[10000]).to_device(Device::Cuda(0));
+    let out = gelu(&x).to_device(Device::Cpu);
+    let result: Vec<f32> = out.as_slice::<bf16>().iter().map(|v| v.to_f32()).collect();
+    check_close(&result, &expected, 0.02, "gelu_bf16");
+}
+
+// === SiLU ===
+
+#[test]
+fn test_silu_f32() {
+    init();
+    let data = make_data(10000);
+    let expected = cpu_silu(&data);
+    let x = Tensor::from_slice(&data, &[10000]).to_device(Device::Cuda(0));
+    let out = silu(&x).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-5, "silu_f32");
+}
+
+// === Softmax ===
+
+#[test]
+fn test_softmax_f32() {
+    init();
+    let rows = 8;
+    let cols = 256;
+    let data = make_data(rows * cols);
+    let expected = cpu_softmax(&data, cols);
+    let x = Tensor::from_slice(&data, &[rows, cols]).to_device(Device::Cuda(0));
+    let out = softmax(&x).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-5, "softmax_f32");
+}
+
+#[test]
+fn test_softmax_sum_to_one() {
+    init();
+    let rows = 4;
+    let cols = 2048;
+    let data: Vec<f32> = (0..rows * cols).map(|i| ((i % 31) as f32 - 15.0) * 0.5).collect();
+    let x = Tensor::from_slice(&data, &[rows, cols]).to_device(Device::Cuda(0));
+    let out = softmax(&x).to_device(Device::Cpu);
+    let result = out.as_slice::<f32>();
+    for r in 0..rows {
+        let row_sum: f32 = result[r * cols..(r + 1) * cols].iter().sum();
+        assert!((row_sum - 1.0).abs() < 1e-5, "softmax row {r} sum = {row_sum}");
+    }
+}
+
+#[test]
+fn test_softmax_large_values() {
+    init();
+    let data = vec![1000.0f32, 1001.0, 999.0, 1000.5];
+    let expected = cpu_softmax(&data, 4);
+    let x = Tensor::from_slice(&data, &[1, 4]).to_device(Device::Cuda(0));
+    let out = softmax(&x).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-5, "softmax_large");
+}
+
+// === Embedding ===
+
+#[test]
+fn test_embedding_f32() {
+    init();
+    let vocab_size = 100;
+    let hidden = 64;
+    let table_data: Vec<f32> = (0..vocab_size * hidden).map(|i| i as f32 * 0.01).collect();
+    let token_ids: Vec<u32> = vec![0, 5, 99, 42, 1];
+
+    let table = Tensor::from_slice(&table_data, &[vocab_size, hidden]).to_device(Device::Cuda(0));
+    let out = embedding(&table, &token_ids).to_device(Device::Cpu);
+
+    assert_eq!(out.shape(), &[5, hidden]);
+    let result = out.as_slice::<f32>();
+    for (seq_idx, &tid) in token_ids.iter().enumerate() {
+        for i in 0..hidden {
+            let expected = table_data[tid as usize * hidden + i];
+            let got = result[seq_idx * hidden + i];
+            assert!((got - expected).abs() < 1e-6,
+                "embedding mismatch at [{seq_idx},{i}]: got {got}, expected {expected}");
+        }
+    }
+}
+
+// === RoPE ===
+
+#[test]
+fn test_rope_f32() {
+    init();
+    let num_tokens = 4;
+    let num_heads = 2;
+    let head_dim = 8;
+    let theta = 10000.0f32;
+    let positions: Vec<u32> = vec![0, 1, 2, 3];
+
+    let x_data: Vec<f32> = (0..num_tokens * num_heads * head_dim)
+        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
+        .collect();
+    let mut expected = x_data.clone();
+    cpu_rope(&mut expected, &positions, num_heads, head_dim, theta);
+
+    let x = Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim])
+        .to_device(Device::Cuda(0));
+    let cache = RopeCache::new(64, head_dim, theta);
+    rope_inplace(&x, &cache, &positions);
+
+    let out = x.to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-4, "rope_f32");
+}
+
+#[test]
+fn test_rope_position_0_identity() {
+    init();
+    // At position 0, all angles are 0, so cos=1, sin=0 → identity transform
+    let num_tokens = 1;
+    let num_heads = 2;
+    let head_dim = 8;
+    let positions: Vec<u32> = vec![0];
+
+    let x_data: Vec<f32> = (0..num_tokens * num_heads * head_dim)
+        .map(|i| (i as f32 + 1.0) * 0.1)
+        .collect();
+
+    let x = Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim])
+        .to_device(Device::Cuda(0));
+    let cache = RopeCache::new(64, head_dim, 10000.0);
+    rope_inplace(&x, &cache, &positions);
+
+    let out = x.to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &x_data, 1e-6, "rope_pos0");
+}

crates/xserv-model/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "xserv-model"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+xserv-cuda = { path = "../xserv-cuda" }
+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
+rand.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, 0);
+        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, 0);
+
+        // 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

@@ -0,0 +1,101 @@
+use std::io::{self, Write};
+use std::path::PathBuf;
+use xserv_model::{loader, KVCache, ModelConfig};
+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: xserv-cli <model-dir> [--max-tokens N]");
+        std::process::exit(1);
+    }
+
+    let model_dir = PathBuf::from(&args[1]);
+    let max_tokens: usize = args
+        .iter()
+        .position(|a| a == "--max-tokens")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(100);
+
+    xserv_cuda::device::set_device(0).unwrap();
+    let info = xserv_cuda::device::device_info(0).unwrap();
+    eprintln!("GPU: {} ({} MB free)", info.name, info.free_memory / 1024 / 1024);
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    let model_type = config.model_type.as_deref().unwrap_or("unknown");
+    eprintln!(
+        "Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}",
+        config.num_layers(), config.hidden(), config.num_heads(),
+        config.num_kv_heads(), config.vocab_size
+    );
+
+    eprintln!("Loading weights...");
+    let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
+    eprintln!("Loaded {} tensors", weights.len());
+
+    let is_qwen3 = model_type.contains("qwen");
+    let dtype = if is_qwen3 { DType::BF16 } else { DType::F32 };
+
+    // Build model
+    enum Model {
+        GPT2(xserv_model::GPT2),
+        Qwen3(xserv_model::Qwen3),
+    }
+    let model = if is_qwen3 {
+        Model::Qwen3(xserv_model::Qwen3::from_weights(config.clone(), weights))
+    } else {
+        Model::GPT2(xserv_model::GPT2::from_weights(config.clone(), weights))
+    };
+
+    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+    eprintln!("Ready (KV cache, dtype={dtype}).\n");
+
+    loop {
+        print!("xserv> ");
+        io::stdout().flush().unwrap();
+        let mut input = String::new();
+        if io::stdin().read_line(&mut input).unwrap() == 0 { break; }
+        let input = input.trim();
+        if input.is_empty() { continue; }
+        if input == "quit" || input == "exit" { break; }
+
+        let token_ids = tokenizer.encode(input);
+        let kv_heads = if is_qwen3 { config.num_kv_heads() } else { config.num_heads() };
+        let mut cache = KVCache::new(
+            config.num_layers(), 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 text = tokenizer.decode(&[next]);
+            print!("{text}");
+            io::stdout().flush().unwrap();
+
+            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/config.rs

@@ -0,0 +1,96 @@
+use serde::Deserialize;
+use std::path::Path;
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct ModelConfig {
+    pub architectures: Option<Vec<String>>,
+    pub model_type: Option<String>,
+
+    // Modern HF naming
+    #[serde(default)]
+    pub hidden_size: Option<usize>,
+    #[serde(default)]
+    pub intermediate_size: Option<usize>,
+    #[serde(default)]
+    pub num_attention_heads: Option<usize>,
+    #[serde(default)]
+    pub num_key_value_heads: Option<usize>,
+    #[serde(default)]
+    pub num_hidden_layers: Option<usize>,
+    pub vocab_size: usize,
+    #[serde(default)]
+    pub max_position_embeddings: Option<usize>,
+
+    // GPT-2 naming
+    #[serde(default)]
+    pub n_embd: Option<usize>,
+    #[serde(default)]
+    pub n_head: Option<usize>,
+    #[serde(default)]
+    pub n_layer: Option<usize>,
+    #[serde(default)]
+    pub n_positions: Option<usize>,
+    #[serde(default)]
+    pub n_inner: Option<usize>,
+
+    // Normalization
+    #[serde(default)]
+    pub layer_norm_eps: Option<f64>,
+    #[serde(default)]
+    pub layer_norm_epsilon: Option<f64>,
+    #[serde(default)]
+    pub rms_norm_eps: Option<f64>,
+
+    // Other
+    #[serde(default)]
+    pub rope_theta: Option<f64>,
+    #[serde(default)]
+    pub tie_word_embeddings: Option<bool>,
+}
+
+impl ModelConfig {
+    pub fn from_file(path: &Path) -> Self {
+        let data = std::fs::read_to_string(path)
+            .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
+        serde_json::from_str(&data)
+            .unwrap_or_else(|e| panic!("failed to parse {}: {e}", path.display()))
+    }
+
+    pub fn hidden(&self) -> usize {
+        self.hidden_size.or(self.n_embd).expect("hidden_size or n_embd required")
+    }
+
+    pub fn num_heads(&self) -> usize {
+        self.num_attention_heads.or(self.n_head).expect("num_attention_heads or n_head required")
+    }
+
+    pub fn num_layers(&self) -> usize {
+        self.num_hidden_layers.or(self.n_layer).expect("num_hidden_layers or n_layer required")
+    }
+
+    pub fn max_seq_len(&self) -> usize {
+        self.max_position_embeddings.or(self.n_positions).unwrap_or(2048)
+    }
+
+    pub fn ffn_hidden(&self) -> usize {
+        self.intermediate_size.or(self.n_inner).unwrap_or(self.hidden() * 4)
+    }
+
+    pub fn num_kv_heads(&self) -> usize {
+        self.num_key_value_heads.unwrap_or(self.num_heads())
+    }
+
+    pub fn head_dim(&self) -> usize {
+        self.hidden() / self.num_heads()
+    }
+
+    pub fn ln_eps(&self) -> f32 {
+        self.layer_norm_eps
+            .or(self.layer_norm_epsilon)
+            .unwrap_or(1e-5) as f32
+    }
+
+    pub fn tied_embeddings(&self) -> bool {
+        self.tie_word_embeddings.unwrap_or(true)
+    }
+}

crates/xserv-model/src/gpt2.rs

@@ -0,0 +1,337 @@
+use std::collections::HashMap;
+use xserv_kernels::*;
+use xserv_tensor::{DType, Device, Tensor};
+
+use crate::config::ModelConfig;
+
+pub struct GPT2 {
+    pub config: ModelConfig,
+    wte: Tensor,
+    wpe: Tensor,
+    layers: Vec<GPT2Block>,
+    ln_f_g: Tensor,
+    ln_f_b: Tensor,
+    lm_head: Tensor, // precomputed wte^T
+}
+
+struct GPT2Block {
+    ln_1_g: Tensor,
+    ln_1_b: Tensor,
+    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,
+    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 {
+    pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
+        crate::init_kernels();
+        let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
+            w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
+        };
+
+        let wte = take(&mut w, "wte.weight");
+        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);
+        for i in 0..num_layers {
+            let p = format!("h.{i}");
+            layers.push(GPT2Block {
+                ln_1_g: take(&mut w, &format!("{p}.ln_1.weight")),
+                ln_1_b: take(&mut w, &format!("{p}.ln_1.bias")),
+                attn_qkv_w: take(&mut w, &format!("{p}.attn.c_attn.weight")),
+                attn_qkv_b: take(&mut w, &format!("{p}.attn.c_attn.bias")),
+                attn_out_w: take(&mut w, &format!("{p}.attn.c_proj.weight")),
+                attn_out_b: take(&mut w, &format!("{p}.attn.c_proj.bias")),
+                ln_2_g: take(&mut w, &format!("{p}.ln_2.weight")),
+                ln_2_b: take(&mut w, &format!("{p}.ln_2.bias")),
+                mlp_fc_w: take(&mut w, &format!("{p}.mlp.c_fc.weight")),
+                mlp_fc_b: take(&mut w, &format!("{p}.mlp.c_fc.bias")),
+                mlp_proj_w: take(&mut w, &format!("{p}.mlp.c_proj.weight")),
+                mlp_proj_b: take(&mut w, &format!("{p}.mlp.c_proj.bias")),
+            });
+        }
+
+        Self { config, wte, wpe, layers, ln_f_g, ln_f_b, lm_head }
+    }
+
+    /// 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();
+
+        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);
+
+        for layer in &self.layers {
+            x = self.transformer_block(layer, &x, None, 0, seq_len, 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)
+    }
+
+    /// 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 (unchanged) ---
+
+fn linear(x: &Tensor, weight: &Tensor, bias: Option<&Tensor>) -> Tensor {
+    let out = matmul_2d(x, weight);
+    if let Some(b) = bias { add_bias(&out, b) } else { out }
+}
+
+fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
+    assert_eq!(a.ndim(), 2);
+    assert_eq!(b.ndim(), 2);
+    matmul(a, b, GemmBackend::CuBlas)
+}
+
+fn add_tensors(a: &Tensor, b: &Tensor) -> Tensor {
+    xserv_kernels::add(a, b)
+}
+
+fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
+    // bias: [N], x: [S, N] — broadcast add via reshape
+    assert_eq!(x.ndim(), 2);
+    assert_eq!(bias.ndim(), 1);
+    let n = bias.shape()[0];
+    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) {
+    let hidden = num_heads * head_dim;
+    let qkv_cpu = qkv.to_device(Device::Cpu);
+    let data = qkv_cpu.as_slice::<f32>();
+
+    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];
+
+    for s in 0..seq_len {
+        let row = &data[s * 3 * hidden..(s + 1) * 3 * hidden];
+        for h in 0..num_heads {
+            let src_off = h * head_dim;
+            let dst_off = (h * seq_len + s) * head_dim;
+            q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
+            k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
+            v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
+        }
+    }
+
+    let device = qkv.device();
+    let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
+    let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
+    let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
+    (q, k, v)
+}
+
+fn merge_heads(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::<f32>();
+
+    let mut out = vec![0.0f32; seq_len * hidden];
+    for s in 0..seq_len {
+        for h in 0..num_heads {
+            let src_off = (h * seq_len + s) * head_dim;
+            let dst_off = s * hidden + h * head_dim;
+            out[dst_off..dst_off + head_dim].copy_from_slice(&src[src_off..src_off + head_dim]);
+        }
+    }
+    Tensor::from_slice(&out, &[seq_len, hidden]).to_device(x.device())
+}
+
+/// 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);
+    let logits_cpu = logits.to_device(Device::Cpu);
+    let data = logits_cpu.as_slice::<f32>();
+    let vocab_size = logits.shape()[1];
+    let seq_len = logits.shape()[0];
+    let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
+    last_row.iter()
+        .enumerate()
+        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
+        .map(|(idx, _)| idx as u32)
+        .unwrap()
+}

crates/xserv-model/src/kv_cache.rs

@@ -0,0 +1,149 @@
+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>,
+    // Per layer: pre-allocated staging buffers for get_kv_len output.
+    // Size: num_kv_heads * max_seq_len * head_dim * elem_size (max possible output).
+    // Avoids cudaMalloc/cudaFree on every get_kv_len call.
+    k_staging: Vec<GpuBuffer>,
+    v_staging: Vec<GpuBuffer>,
+    seq_len: usize,
+    max_seq_len: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    elem_size: usize,
+    dtype: DType,
+    device: u32,
+}
+
+impl GpuKVCache {
+    pub fn new(config: &ModelConfig, max_seq_len: usize, dtype: DType, device: u32) -> 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);
+        let mut k_staging = Vec::with_capacity(num_layers);
+        let mut v_staging = 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);
+            k_staging.push(GpuBuffer::alloc(buf_size).expect("alloc KV staging K"));
+            v_staging.push(GpuBuffer::alloc(buf_size).expect("alloc KV staging V"));
+        }
+
+        Self { k_bufs, v_bufs, k_staging, v_staging, seq_len: 0, max_seq_len, num_kv_heads, head_dim, elem_size, dtype, device }
+    }
+
+    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(&mut self, layer: usize) -> (Tensor, Tensor) {
+        let sl = self.seq_len;
+        self.get_kv_len(layer, sl)
+    }
+
+    pub fn get_kv_len(&mut 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;
+
+        // Copy each head's valid portion into pre-allocated staging buffers.
+        // Split borrows: staging (mut) vs cache (shared) are separate struct fields,
+        // so the borrow checker allows simultaneous &mut staging + &cache.
+        let out_size = nh * sl * hd * es;
+        let k_stg = &mut self.k_staging[layer];
+        let k_buf = &self.k_bufs[layer];
+        let v_stg = &mut self.v_staging[layer];
+        let v_buf = &self.v_bufs[layer];
+        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_stg.copy_from_device_at(k_buf, src_off, dst_off, count).unwrap();
+            v_stg.copy_from_device_at(v_buf, src_off, dst_off, count).unwrap();
+        }
+        // Grab raw pointers before dropping the mutable borrows
+        let k_ptr = k_stg.as_mut_ptr();
+        let v_ptr = v_stg.as_mut_ptr();
+
+        // Create Tensors that borrow from the staging buffers (no cudaMalloc/cudaFree).
+        // Safety: staging buffers are owned by GpuKVCache and outlive the returned Tensors
+        // in practice (Tensors are consumed within the same forward pass before the next
+        // get_kv_len call overwrites the staging buffer).
+        let shape = &[1usize, nh, sl, hd];
+        let k = unsafe {
+            tensor_from_gpu_buffer(GpuBuffer::borrow_raw(k_ptr, out_size), shape, self.dtype, self.device)
+        };
+        let v = unsafe {
+            tensor_from_gpu_buffer(GpuBuffer::borrow_raw(v_ptr, out_size), shape, self.dtype, self.device)
+        };
+        (k, v)
+    }
+}
+
+/// Create a Tensor from a GpuBuffer (takes ownership).
+unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
+    use xserv_tensor::storage::Storage;
+    use xserv_tensor::shape::contiguous_strides;
+    use smallvec::SmallVec;
+
+    let storage = Storage::cuda(buf, device);
+    Tensor::from_storage(
+        storage,
+        SmallVec::from_slice(shape),
+        contiguous_strides(shape),
+        0,
+        dtype,
+    )
+}
+
+/// Public version for use by other modules (e.g., batched decode concat).
+///
+/// # Safety
+/// `buf` must be a valid GPU allocation with at least `product(shape) * dtype.size_bytes()` bytes.
+pub unsafe fn tensor_from_gpu_buffer_pub(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
+    tensor_from_gpu_buffer(buf, shape, dtype, device)
+}

crates/xserv-model/src/lib.rs

@@ -0,0 +1,18 @@
+pub mod config;
+pub mod gpt2;
+pub mod kv_cache;
+pub mod loader;
+pub mod qwen3;
+pub mod sampling;
+
+pub use config::ModelConfig;
+pub use gpt2::{GPT2, KVCache};
+pub use kv_cache::GpuKVCache;
+pub use qwen3::Qwen3;
+pub use sampling::{SamplingParams, sample};
+
+/// Initialize GPU kernel hooks. Called automatically by model constructors,
+/// but safe to call multiple times (idempotent via OnceLock).
+pub fn init_kernels() {
+    xserv_kernels::init();
+}

crates/xserv-model/src/loader.rs

@@ -0,0 +1,87 @@
+use half::{bf16, f16};
+use safetensors::SafeTensors;
+use std::collections::HashMap;
+use std::path::Path;
+use xserv_tensor::{DType, Device, Tensor};
+
+pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor> {
+    let data = std::fs::read(path)
+        .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
+    let st = SafeTensors::deserialize(&data)
+        .unwrap_or_else(|e| panic!("failed to parse safetensors {}: {e}", path.display()));
+
+    let mut tensors = HashMap::new();
+
+    for (name, view) in st.tensors() {
+        let shape: Vec<usize> = view.shape().to_vec();
+        let raw_bytes = view.data();
+        let dtype = match view.dtype() {
+            safetensors::Dtype::F32 => DType::F32,
+            safetensors::Dtype::F16 => DType::F16,
+            safetensors::Dtype::BF16 => DType::BF16,
+            other => {
+                eprintln!("skipping tensor {name}: unsupported dtype {other:?}");
+                continue;
+            }
+        };
+
+        let tensor = make_tensor(raw_bytes, &shape, dtype);
+        let tensor = tensor.to_device(device);
+        tensors.insert(name.to_string(), tensor);
+    }
+
+    tensors
+}
+
+/// Load from a directory containing model.safetensors (or sharded files) + config.json.
+pub fn load_model_dir(dir: &Path, device: Device) -> HashMap<String, Tensor> {
+    let single = dir.join("model.safetensors");
+    if single.exists() {
+        return load_safetensors(&single, device);
+    }
+
+    // Try sharded: model-00001-of-NNNNN.safetensors
+    let mut all_tensors = HashMap::new();
+    let mut entries: Vec<_> = std::fs::read_dir(dir)
+        .unwrap()
+        .filter_map(|e| e.ok())
+        .filter(|e| {
+            e.path()
+                .file_name()
+                .map(|f| f.to_string_lossy().ends_with(".safetensors"))
+                .unwrap_or(false)
+        })
+        .collect();
+    entries.sort_by_key(|e| e.file_name());
+
+    for entry in entries {
+        let tensors = load_safetensors(&entry.path(), device);
+        all_tensors.extend(tensors);
+    }
+
+    assert!(!all_tensors.is_empty(), "no safetensors files found in {}", dir.display());
+    all_tensors
+}
+
+fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
+    match dtype {
+        DType::F32 => {
+            let floats: &[f32] = unsafe {
+                std::slice::from_raw_parts(raw_bytes.as_ptr() as *const f32, raw_bytes.len() / 4)
+            };
+            Tensor::from_slice(floats, shape)
+        }
+        DType::F16 => {
+            let halfs: &[f16] = unsafe {
+                std::slice::from_raw_parts(raw_bytes.as_ptr() as *const f16, raw_bytes.len() / 2)
+            };
+            Tensor::from_slice(halfs, shape)
+        }
+        DType::BF16 => {
+            let bfs: &[bf16] = unsafe {
+                std::slice::from_raw_parts(raw_bytes.as_ptr() as *const bf16, raw_bytes.len() / 2)
+            };
+            Tensor::from_slice(bfs, shape)
+        }
+    }
+}

crates/xserv-model/src/qwen3.rs

@@ -0,0 +1,492 @@
+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 {
+        crate::init_kernels();
+        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(),
+            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)
+    }
+
+    /// Batched decode: process one token per sequence simultaneously.
+    /// All compute-heavy ops (projections, FFN) operate on [B, hidden] tensors.
+    /// Per-sequence ops (RoPE, KV cache, attention) are handled individually.
+    ///
+    /// tokens: one token per sequence (len = batch_size)
+    /// positions: position offset for each sequence (len = batch_size)
+    /// caches: one mutable KV cache per sequence (len = batch_size)
+    ///
+    /// Returns logits: [batch_size, vocab_size]
+    pub fn forward_decode_batch(
+        &self,
+        tokens: &[u32],
+        positions: &[usize],
+        caches: &mut [&mut GpuKVCache],
+    ) -> Tensor {
+        let batch = tokens.len();
+        assert_eq!(positions.len(), batch);
+        assert_eq!(caches.len(), batch);
+        assert!(batch > 0);
+
+        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;
+
+        // Batched embedding: [B, hidden]
+        let mut x = embedding(&self.embed_tokens, tokens);
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps); // [B, hidden]
+
+            // Batched projections: [B, hidden] × [hidden, X] = [B, X]
+            let q_all = matmul_2d(&normed, &layer.q_proj_wt); // [B, num_heads*head_dim]
+            let k_all = matmul_2d(&normed, &layer.k_proj_wt); // [B, num_kv_heads*head_dim]
+            let v_all = matmul_2d(&normed, &layer.v_proj_wt); // [B, num_kv_heads*head_dim]
+
+            // Per-sequence: reshape, qk-norm, RoPE, KV cache, attention, merge
+            let mut attn_outputs: Vec<Tensor> = Vec::with_capacity(batch);
+            for b in 0..batch {
+                // Extract row b: [1, X] — view into contiguous [B, X]
+                let q_row = row_view(&q_all, b); // [1, num_heads*head_dim]
+                let k_row = row_view(&k_all, b); // [1, num_kv_heads*head_dim]
+                let v_row = row_view(&v_all, b); // [1, num_kv_heads*head_dim]
+
+                // GPU reshape: [1, H*D] → [1, H, 1, D]
+                let q = xserv_kernels::reshape_heads_gpu(&q_row, 1, num_heads, head_dim);
+                let k = xserv_kernels::reshape_heads_gpu(&k_row, 1, num_kv_heads, head_dim);
+                let v = xserv_kernels::reshape_heads_gpu(&v_row, 1, num_kv_heads, head_dim);
+
+                // QK norm
+                let q = head_rmsnorm(&q, &layer.q_norm, eps);
+                let k = head_rmsnorm(&k, &layer.k_norm, eps);
+
+                // GPU transpose for RoPE: [1, H, 1, D] → [1, H, D]
+                let q = xserv_kernels::transpose_for_rope_gpu(&q, 1, num_heads, head_dim);
+                let k = xserv_kernels::transpose_for_rope_gpu(&k, 1, num_kv_heads, head_dim);
+
+                // RoPE with per-sequence position
+                let pos = [positions[b] as u32];
+                rope_inplace(&q, &self.rope_cache, &pos);
+                rope_inplace(&k, &self.rope_cache, &pos);
+
+                // Transpose back: [1, H, D] → [1, H, 1, D]
+                let q = xserv_kernels::transpose_from_rope_gpu(&q, 1, num_heads, head_dim);
+                let k = xserv_kernels::transpose_from_rope_gpu(&k, 1, num_kv_heads, head_dim);
+
+                // KV cache: append and get full cache
+                let pos_b = positions[b];
+                caches[b].append(layer_idx, &k, &v, 1, pos_b);
+                let (k_full, v_full) = caches[b].get_kv_len(layer_idx, pos_b + 1);
+
+                // Decode attention (uses native GQA, no repeat_kv needed)
+                let attn_out = flash_attention(&q, &k_full, &v_full, true);
+
+                // Merge heads: [1, H, 1, D] → [1, hidden]
+                let merged = xserv_kernels::merge_heads_gpu(&attn_out, 1, num_heads, head_dim);
+                attn_outputs.push(merged);
+            }
+
+            // Concat attention outputs: [B, hidden]
+            let attn_merged = concat_rows(&attn_outputs);
+
+            // Batched O projection: [B, hidden] × [hidden, hidden] = [B, hidden]
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+
+            // Fused add + rmsnorm
+            let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            // Batched FFN: all projections on [B, hidden]
+            let gate = matmul_2d(&normed, &layer.gate_proj_wt);
+            let up = matmul_2d(&normed, &layer.up_proj_wt);
+            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
+            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            x = add_any(&residual, &down);
+        }
+
+        // Advance KV cache seq_len for each sequence
+        for b in 0..batch {
+            caches[b].advance_seq_len(1);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_2d(&x, &self.lm_head_t) // [B, vocab_size]
+    }
+
+    /// 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);
+
+            // Flash Attention with native GQA (no repeat_kv needed)
+            let attn_out = flash_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);
+
+            // Fused add + rmsnorm: (normed, x) where x = residual + attn_proj
+            let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            // Fused SiLU×Mul
+            let gate = matmul_2d(&normed, &layer.gate_proj_wt);
+            let up = matmul_2d(&normed, &layer.up_proj_wt);
+            let hidden_states = xserv_kernels::silu_mul(&gate, &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())
+}
+
+/// Extract row `b` from a contiguous 2D tensor [B, cols] as a [1, cols] view.
+/// Zero-copy: shares storage with the original tensor.
+fn row_view(t: &Tensor, row: usize) -> Tensor {
+    assert_eq!(t.ndim(), 2);
+    assert!(t.is_contiguous());
+    let cols = t.shape()[1];
+    assert!(row < t.shape()[0]);
+    let new_offset = t.offset() + row * cols;
+    Tensor::from_storage(
+        t.storage().clone(),
+        smallvec::SmallVec::from_slice(&[1, cols]),
+        xserv_tensor::shape::contiguous_strides(&[1, cols]),
+        new_offset,
+        t.dtype(),
+    )
+}
+
+/// Concatenate row tensors [1, cols] into a single [B, cols] tensor via D2D memcpy.
+fn concat_rows(rows: &[Tensor]) -> Tensor {
+    assert!(!rows.is_empty());
+    let batch = rows.len();
+    let cols = rows[0].shape()[1];
+    let dtype = rows[0].dtype();
+    let device = rows[0].device();
+    let elem_size = dtype.size_bytes();
+    let row_bytes = cols * elem_size;
+
+    // Allocate output [B, cols] and copy each row into it
+    let total_bytes = batch * row_bytes;
+    let mut out_buf = xserv_cuda::allocator::cached_alloc(total_bytes).expect("alloc concat_rows");
+
+    for (b, row) in rows.iter().enumerate() {
+        assert_eq!(row.shape(), &[1, cols]);
+        assert!(row.is_contiguous());
+        let src_buf = row.storage().gpu_buffer();
+        let src_offset = row.offset() * elem_size;
+        let dst_offset = b * row_bytes;
+        out_buf.copy_from_device_at(src_buf, src_offset, dst_offset, row_bytes).unwrap();
+    }
+
+    // Wrap in a Tensor
+    let device_id = match device { Device::Cuda(id) => id, _ => panic!("expected CUDA device") };
+    unsafe {
+        crate::kv_cache::tensor_from_gpu_buffer_pub(out_buf, &[batch, cols], dtype, device_id)
+    }
+}
+
+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-model/src/sampling.rs

@@ -0,0 +1,120 @@
+use half::bf16;
+use rand::Rng;
+use xserv_tensor::{DType, Device, Tensor};
+
+pub struct SamplingParams {
+    pub temperature: f32,
+    pub top_k: usize,
+    pub top_p: f32,
+}
+
+impl Default for SamplingParams {
+    fn default() -> Self {
+        Self { temperature: 0.0, top_k: 0, top_p: 1.0 }
+    }
+}
+
+/// Sample a token from logits with shape [seq_len, vocab_size].
+/// Uses the last position's logits. Handles both F32 and BF16 dtypes.
+pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
+    assert_eq!(logits.ndim(), 2);
+    let vocab_size = logits.shape()[1];
+    let seq_len = logits.shape()[0];
+    let logits_cpu = logits.to_device(Device::Cpu);
+
+    // Extract last row as f32
+    let last_row: Vec<f32> = match logits.dtype() {
+        DType::F32 => {
+            let data = logits_cpu.as_slice::<f32>();
+            data[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec()
+        }
+        DType::BF16 => {
+            let data = logits_cpu.as_slice::<bf16>();
+            data[(seq_len - 1) * vocab_size..seq_len * vocab_size]
+                .iter()
+                .map(|v| v.to_f32())
+                .collect()
+        }
+        _ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
+    };
+
+    // Greedy
+    if params.temperature == 0.0 {
+        return argmax(&last_row);
+    }
+
+    // Apply temperature
+    let mut logits_f32: Vec<f32> = last_row.iter().map(|v| v / params.temperature).collect();
+
+    // Top-k filtering
+    if params.top_k > 0 && params.top_k < vocab_size {
+        let mut indices: Vec<usize> = (0..vocab_size).collect();
+        indices.select_nth_unstable_by(params.top_k, |&a, &b| {
+            logits_f32[b].partial_cmp(&logits_f32[a]).unwrap()
+        });
+        // Everything after top_k should be masked
+        for &i in &indices[params.top_k..] {
+            logits_f32[i] = f32::NEG_INFINITY;
+        }
+    }
+
+    // Top-p (nucleus) filtering
+    if params.top_p < 1.0 {
+        // Sort indices by descending logit value
+        let mut indices: Vec<usize> = (0..vocab_size).collect();
+        indices.sort_unstable_by(|&a, &b| logits_f32[b].partial_cmp(&logits_f32[a]).unwrap());
+
+        // Compute softmax probabilities for the sorted order
+        let max_val = logits_f32[indices[0]];
+        let sorted_probs: Vec<f32> = indices
+            .iter()
+            .map(|&i| (logits_f32[i] - max_val).exp())
+            .collect();
+        let sum: f32 = sorted_probs.iter().sum();
+        let sorted_probs: Vec<f32> = sorted_probs.iter().map(|v| v / sum).collect();
+
+        // Cumulative sum, find cutoff
+        let mut cumsum = 0.0f32;
+        let mut cutoff = indices.len();
+        for (rank, &prob) in sorted_probs.iter().enumerate() {
+            cumsum += prob;
+            if cumsum > params.top_p {
+                cutoff = rank + 1; // keep at least this many
+                break;
+            }
+        }
+
+        // Mask everything beyond cutoff
+        for &i in &indices[cutoff..] {
+            logits_f32[i] = f32::NEG_INFINITY;
+        }
+    }
+
+    // Softmax
+    let max_val = logits_f32.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+    let exps: Vec<f32> = logits_f32.iter().map(|v| (v - max_val).exp()).collect();
+    let sum: f32 = exps.iter().sum();
+    let probs: Vec<f32> = exps.iter().map(|v| v / sum).collect();
+
+    // Weighted random sampling
+    let mut rng = rand::thread_rng();
+    let r: f32 = rng.r#gen();
+    let mut cumsum = 0.0f32;
+    for (i, &p) in probs.iter().enumerate() {
+        cumsum += p;
+        if cumsum > r {
+            return i as u32;
+        }
+    }
+
+    // Fallback (rounding edge case)
+    (vocab_size - 1) as u32
+}
+
+fn argmax(data: &[f32]) -> u32 {
+    data.iter()
+        .enumerate()
+        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
+        .map(|(i, _)| i as u32)
+        .unwrap()
+}

crates/xserv-server/Cargo.toml

@@ -0,0 +1,22 @@
+[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
+tokio-stream.workspace = true

crates/xserv-server/src/api.rs

@@ -0,0 +1,299 @@
+use axum::Extension;
+use axum::Json;
+use axum::http::StatusCode;
+use axum::response::sse::{Event, KeepAlive, Sse};
+use axum::response::{IntoResponse, Response};
+use serde::{Deserialize, Serialize};
+use std::convert::Infallible;
+use std::sync::Arc;
+use tokio_stream::StreamExt;
+use tokio_stream::wrappers::ReceiverStream;
+use uuid::Uuid;
+
+use crate::AppState;
+use crate::engine::{GenerateEvent, GenerateRequest};
+use xserv_model::SamplingParams;
+
+#[derive(Deserialize)]
+pub struct ChatRequest {
+    #[serde(default)]
+    pub model: Option<String>,
+    pub messages: Vec<Message>,
+    #[serde(default = "default_max_tokens")]
+    pub max_tokens: usize,
+    #[serde(default)]
+    pub stream: Option<bool>,
+    #[serde(default)]
+    pub temperature: Option<f32>,
+    #[serde(default)]
+    pub top_k: Option<usize>,
+    #[serde(default)]
+    pub top_p: Option<f32>,
+}
+
+#[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>,
+) -> Response {
+    if req.stream == Some(true) {
+        chat_stream(state, req)
+    } else {
+        chat_non_stream(state, req).await
+    }
+}
+
+async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
+    let id = format!("chatcmpl-{}", Uuid::new_v4());
+    let model_name = state.model_name.clone();
+    let created = unix_timestamp();
+
+    let prompt = build_prompt(&req.messages);
+    let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
+    let prompt_token_count = prompt_tokens.len();
+
+    let max_seq_len = state.max_seq_len;
+    if prompt_token_count >= max_seq_len {
+        return (StatusCode::BAD_REQUEST, Json(serde_json::json!({
+            "error": {
+                "message": format!("prompt is {} tokens, exceeds max_seq_len {}", prompt_token_count, max_seq_len),
+                "type": "invalid_request_error"
+            }
+        }))).into_response();
+    }
+    let max_tokens = req.max_tokens.min(max_seq_len - prompt_token_count);
+
+    let (tx, mut rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
+    let gen_req = GenerateRequest {
+        prompt_tokens,
+        max_tokens,
+        sampling: sampling_params(&req),
+        sender: tx,
+    };
+    state
+        .engine_sender
+        .lock()
+        .unwrap()
+        .send(gen_req)
+        .expect("engine channel closed");
+
+    let mut content = String::new();
+    let mut completion_token_count: usize = 0;
+    let mut finish_reason = "length".to_string();
+    while let Some(event) = rx.recv().await {
+        match event {
+            GenerateEvent::Token { text, .. } => {
+                completion_token_count += 1;
+                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": prompt_token_count,
+            "completion_tokens": completion_token_count,
+            "total_tokens": prompt_token_count + completion_token_count
+        }
+    })).into_response()
+}
+
+fn chat_stream(
+    state: Arc<AppState>,
+    req: ChatRequest,
+) -> Response {
+    let id = format!("chatcmpl-{}", Uuid::new_v4());
+    let model_name = state.model_name.clone();
+    let created = unix_timestamp();
+
+    let prompt = build_prompt(&req.messages);
+    let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
+
+    let max_seq_len = state.max_seq_len;
+    if prompt_tokens.len() >= max_seq_len {
+        return (StatusCode::BAD_REQUEST, Json(serde_json::json!({
+            "error": {
+                "message": format!("prompt is {} tokens, exceeds max_seq_len {}", prompt_tokens.len(), max_seq_len),
+                "type": "invalid_request_error"
+            }
+        }))).into_response();
+    }
+    let max_tokens = req.max_tokens.min(max_seq_len - prompt_tokens.len());
+
+    let (engine_tx, engine_rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
+    let gen_req = GenerateRequest {
+        prompt_tokens,
+        max_tokens,
+        sampling: sampling_params(&req),
+        sender: engine_tx,
+    };
+    state
+        .engine_sender
+        .lock()
+        .unwrap()
+        .send(gen_req)
+        .expect("engine channel closed");
+
+    // SSE event channel: engine events -> SSE events
+    let (sse_tx, sse_rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
+
+    tokio::spawn(async move {
+        let mut engine_stream = ReceiverStream::new(engine_rx);
+        let mut first = true;
+
+        while let Some(event) = engine_stream.next().await {
+            match event {
+                GenerateEvent::Token { text, .. } => {
+                    if first {
+                        // First chunk: role announcement
+                        let chunk =
+                            make_chunk(&id, &model_name, created, None, Some("assistant"), None);
+                        let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
+                        first = false;
+                    }
+                    let chunk = make_chunk(&id, &model_name, created, Some(&text), None, None);
+                    if sse_tx.send(Ok(Event::default().data(chunk))).await.is_err() {
+                        return; // client disconnected
+                    }
+                }
+                GenerateEvent::Done { finish_reason } => {
+                    if first {
+                        // Edge case: Done arrived with no tokens
+                        let chunk =
+                            make_chunk(&id, &model_name, created, None, Some("assistant"), None);
+                        let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
+                    }
+                    let chunk =
+                        make_chunk(&id, &model_name, created, None, None, Some(&finish_reason));
+                    let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
+                    let _ = sse_tx
+                        .send(Ok(Event::default().data("[DONE]".to_string())))
+                        .await;
+                    return;
+                }
+            }
+        }
+    });
+
+    Sse::new(ReceiverStream::new(sse_rx)).keep_alive(KeepAlive::default()).into_response()
+}
+
+fn make_chunk(
+    id: &str,
+    model: &str,
+    created: u64,
+    content: Option<&str>,
+    role: Option<&str>,
+    finish_reason: Option<&str>,
+) -> String {
+    let mut delta = serde_json::Map::new();
+    if let Some(r) = role {
+        delta.insert("role".into(), serde_json::Value::String(r.into()));
+        // Role chunk also includes empty content per OpenAI spec
+        delta.insert("content".into(), serde_json::Value::String(String::new()));
+    }
+    if let Some(c) = content {
+        delta.insert("content".into(), serde_json::Value::String(c.into()));
+    }
+
+    let fr = match finish_reason {
+        Some(r) => serde_json::Value::String(r.into()),
+        None => serde_json::Value::Null,
+    };
+
+    serde_json::json!({
+        "id": id,
+        "object": "chat.completion.chunk",
+        "created": created,
+        "model": model,
+        "choices": [{
+            "index": 0,
+            "delta": delta,
+            "finish_reason": fr,
+        }]
+    })
+    .to_string()
+}
+
+fn unix_timestamp() -> u64 {
+    std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap()
+        .as_secs()
+}
+
+fn sampling_params(req: &ChatRequest) -> SamplingParams {
+    SamplingParams {
+        temperature: req.temperature.unwrap_or(0.0),
+        top_k: req.top_k.unwrap_or(0),
+        top_p: req.top_p.unwrap_or(1.0),
+    }
+}
+
+fn build_prompt(messages: &[Message]) -> String {
+    let mut prompt = String::new();
+    for msg in messages {
+        match msg.role.as_str() {
+            "system" | "user" | "assistant" => {
+                prompt.push_str("<|im_start|>");
+                prompt.push_str(&msg.role);
+                prompt.push('\n');
+                prompt.push_str(&msg.content);
+                prompt.push_str("<|im_end|>\n");
+            }
+            _ => {}
+        }
+    }
+    prompt.push_str("<|im_start|>assistant\n");
+    prompt
+}

crates/xserv-server/src/engine.rs

@@ -0,0 +1,233 @@
+use std::collections::VecDeque;
+use std::path::Path;
+use std::sync::mpsc;
+use std::sync::Once;
+use std::time::Instant;
+use xserv_model::{GpuKVCache, ModelConfig, Qwen3, SamplingParams, sample};
+use xserv_model::loader;
+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 sampling: SamplingParams,
+    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,
+    sampling: SamplingParams,
+    kv_cache: Option<GpuKVCache>,
+    sender: tokio::sync::mpsc::Sender<GenerateEvent>,
+    prefilled: bool,
+    eos_token_id: Option<u32>,
+    created_at: Instant,
+}
+
+impl Engine {
+    pub fn load(model_dir: &Path, max_batch_size: usize, max_seq_len: 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"));
+        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 }
+
+    pub fn max_seq_len(&self) -> usize { self.max_seq_len }
+
+    /// 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 4a: Process prefills (one at a time — different prompt lengths)
+            // Prefill sequences must be processed individually because they have
+            // different prompt lengths and each needs a full forward pass.
+            let mut newly_prefilled = Vec::new();
+            for seq in running.iter_mut() {
+                if !seq.prefilled {
+                    let logits = self.model.forward_gpu_cache(&seq.prompt_tokens, seq.kv_cache.as_mut().unwrap());
+                    let next = sample(&logits, &seq.sampling);
+                    seq.generated_tokens.push(next);
+                    seq.prefilled = true;
+                    self.emit_token(seq, next);
+                    newly_prefilled.push(seq.id);
+                }
+            }
+
+            // Step 4b: Batched decode — batch all decode-ready sequences into one forward pass.
+            // Projections and FFN run as [B, hidden] matmuls; attention remains per-seq.
+            let decode_indices: Vec<usize> = running.iter().enumerate()
+                .filter(|(_, s)| s.prefilled && !newly_prefilled.contains(&s.id))
+                .map(|(i, _)| i)
+                .collect();
+
+            if !decode_indices.is_empty() {
+                static LOG_ONCE: Once = Once::new();
+                LOG_ONCE.call_once(|| {
+                    eprintln!("[scheduler] batched decode active");
+                });
+                eprintln!("[scheduler] decode batch_size={}", decode_indices.len());
+
+                if decode_indices.len() == 1 {
+                    // Single sequence: use per-seq path (no batching overhead)
+                    let i = decode_indices[0];
+                    let last = *running[i].generated_tokens.last().unwrap();
+                    let logits = self.model.forward_gpu_cache(&[last], running[i].kv_cache.as_mut().unwrap());
+                    let next = sample(&logits, &running[i].sampling);
+                    running[i].generated_tokens.push(next);
+                    self.emit_token(&running[i], next);
+                } else {
+                    // Batched decode: extract tokens and positions
+                    let tokens: Vec<u32> = decode_indices.iter()
+                        .map(|&i| *running[i].generated_tokens.last().unwrap())
+                        .collect();
+                    let positions: Vec<usize> = decode_indices.iter()
+                        .map(|&i| running[i].kv_cache.as_ref().unwrap().seq_len())
+                        .collect();
+
+                    // Take caches out of sequences via Option::take (no dummy allocation).
+                    let mut caches: Vec<GpuKVCache> = decode_indices.iter()
+                        .map(|&i| running[i].kv_cache.take().unwrap())
+                        .collect();
+                    let mut cache_refs: Vec<&mut GpuKVCache> = caches.iter_mut().collect();
+
+                    let logits = self.model.forward_decode_batch(&tokens, &positions, &mut cache_refs);
+
+                    // Put caches back: pop from end while iterating in reverse
+                    drop(cache_refs);
+                    for &i in decode_indices.iter().rev() {
+                        running[i].kv_cache = Some(caches.pop().unwrap());
+                    }
+
+                    // Sample per-sequence from batched logits [B, vocab_size]
+                    let vocab_size = logits.shape()[1];
+                    let logits_cpu = logits.to_device(xserv_tensor::Device::Cpu);
+                    let data = logits_cpu.as_slice::<half::bf16>();
+                    for (j, &i) in decode_indices.iter().enumerate() {
+                        let row_start = j * vocab_size;
+                        let row_logits = &data[row_start..row_start + vocab_size];
+                        let next = if running[i].sampling.temperature == 0.0 {
+                            // Greedy: argmax
+                            row_logits.iter().enumerate()
+                                .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
+                                .map(|(idx, _)| idx as u32).unwrap()
+                        } else {
+                            // Use the row as a single-row tensor for full sampling
+                            let row_tensor = xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]);
+                            sample(&row_tensor, &running[i].sampling)
+                        };
+                        running[i].generated_tokens.push(next);
+                        self.emit_token(&running[i], 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, 0);
+        Sequence {
+            id,
+            prompt_tokens: req.prompt_tokens,
+            generated_tokens: Vec::new(),
+            max_tokens: req.max_tokens,
+            sampling: req.sampling,
+            kv_cache: Some(kv_cache),
+            sender: req.sender,
+            prefilled: false,
+            eos_token_id: self.tokenizer.eos_token_id(),
+            created_at: Instant::now(),
+        }
+    }
+
+    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::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() || seq.eos_token_id == Some(last)
+}

crates/xserv-server/src/main.rs

@@ -0,0 +1,73 @@
+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>,
+    pub max_seq_len: usize,
+}
+
+#[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] [--max-seq-len 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 max_seq_len: usize = args.iter()
+        .position(|a| a == "--max-seq-len")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(2048);
+
+    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, max_seq_len);
+        engine.run(rx);
+    });
+
+    let state = Arc::new(AppState {
+        model_name,
+        engine_sender: Mutex::new(tx),
+        engine_tokenizer: Mutex::new(tokenizer),
+        max_seq_len,
+    });
+
+    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}, max_seq_len={max_seq_len})");
+    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
+    axum::serve(listener, app).await.unwrap();
+}

crates/xserv-tensor/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "xserv-tensor"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+xserv-cuda = { path = "../xserv-cuda" }
+half.workspace = true
+smallvec.workspace = true

crates/xserv-tensor/src/dtype.rs

@@ -0,0 +1,57 @@
+use half::{bf16, f16};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum DType {
+    F32,
+    F16,
+    BF16,
+}
+
+impl DType {
+    pub fn size_bytes(self) -> usize {
+        match self {
+            DType::F32 => 4,
+            DType::F16 => 2,
+            DType::BF16 => 2,
+        }
+    }
+
+    pub fn name(self) -> &'static str {
+        match self {
+            DType::F32 => "f32",
+            DType::F16 => "f16",
+            DType::BF16 => "bf16",
+        }
+    }
+}
+
+impl std::fmt::Display for DType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.name())
+    }
+}
+
+/// Trait for types that can be stored in a Tensor.
+pub trait TensorDType: Copy + Send + Sync + 'static {
+    const DTYPE: DType;
+    fn to_f64(self) -> f64;
+    fn from_f64(v: f64) -> Self;
+}
+
+impl TensorDType for f32 {
+    const DTYPE: DType = DType::F32;
+    fn to_f64(self) -> f64 { self as f64 }
+    fn from_f64(v: f64) -> Self { v as f32 }
+}
+
+impl TensorDType for f16 {
+    const DTYPE: DType = DType::F16;
+    fn to_f64(self) -> f64 { self.to_f32() as f64 }
+    fn from_f64(v: f64) -> Self { f16::from_f32(v as f32) }
+}
+
+impl TensorDType for bf16 {
+    const DTYPE: DType = DType::BF16;
+    fn to_f64(self) -> f64 { self.to_f32() as f64 }
+    fn from_f64(v: f64) -> Self { bf16::from_f32(v as f32) }
+}

crates/xserv-tensor/src/lib.rs

@@ -0,0 +1,9 @@
+pub mod dtype;
+pub mod shape;
+pub mod storage;
+pub mod tensor;
+
+pub use dtype::{DType, TensorDType};
+pub use shape::Dims;
+pub use storage::{Device, Storage};
+pub use tensor::{register_gpu_contiguous, Tensor};

crates/xserv-tensor/src/shape.rs

@@ -0,0 +1,105 @@
+use smallvec::SmallVec;
+
+pub type Dims = SmallVec<[usize; 4]>;
+
+/// Compute contiguous strides for a given shape (row-major / C order).
+/// Example: shape [2, 3, 4] => strides [12, 4, 1]
+pub fn contiguous_strides(shape: &[usize]) -> Dims {
+    let mut strides = SmallVec::with_capacity(shape.len());
+    strides.resize(shape.len(), 0);
+    if shape.is_empty() {
+        return strides;
+    }
+    strides[shape.len() - 1] = 1;
+    for i in (0..shape.len() - 1).rev() {
+        strides[i] = strides[i + 1] * shape[i + 1];
+    }
+    strides
+}
+
+/// Check if the given strides represent contiguous (row-major) layout for the shape.
+pub fn is_contiguous(shape: &[usize], strides: &[usize]) -> bool {
+    if shape.is_empty() {
+        return true;
+    }
+    let expected = contiguous_strides(shape);
+    strides == expected.as_slice()
+}
+
+/// Total number of elements given a shape.
+pub fn num_elements(shape: &[usize]) -> usize {
+    shape.iter().product()
+}
+
+/// Compute the shape after broadcasting two shapes together (NumPy rules).
+/// Returns None if shapes are not broadcastable.
+pub fn broadcast_shape(a: &[usize], b: &[usize]) -> Option<Dims> {
+    let ndim = a.len().max(b.len());
+    let mut result = SmallVec::with_capacity(ndim);
+    for i in 0..ndim {
+        let da = if i < ndim - a.len() { 1 } else { a[i - (ndim - a.len())] };
+        let db = if i < ndim - b.len() { 1 } else { b[i - (ndim - b.len())] };
+        if da == db {
+            result.push(da);
+        } else if da == 1 {
+            result.push(db);
+        } else if db == 1 {
+            result.push(da);
+        } else {
+            return None;
+        }
+    }
+    Some(result)
+}
+
+/// Compute broadcast strides: for dimensions where size is 1 but output is >1, stride becomes 0.
+pub fn broadcast_strides(shape: &[usize], strides: &[usize], target_shape: &[usize]) -> Dims {
+    let ndim = target_shape.len();
+    let offset = ndim - shape.len();
+    let mut result = SmallVec::with_capacity(ndim);
+    for i in 0..ndim {
+        if i < offset {
+            result.push(0);
+        } else {
+            let orig_idx = i - offset;
+            if shape[orig_idx] == 1 && target_shape[i] > 1 {
+                result.push(0);
+            } else {
+                result.push(strides[orig_idx]);
+            }
+        }
+    }
+    result
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_contiguous_strides() {
+        assert_eq!(contiguous_strides(&[2, 3, 4]).as_slice(), &[12, 4, 1]);
+        assert_eq!(contiguous_strides(&[5]).as_slice(), &[1]);
+        assert_eq!(contiguous_strides(&[2, 3]).as_slice(), &[3, 1]);
+    }
+
+    #[test]
+    fn test_is_contiguous() {
+        assert!(is_contiguous(&[2, 3], &[3, 1]));
+        assert!(!is_contiguous(&[3, 2], &[1, 3])); // transposed
+    }
+
+    #[test]
+    fn test_broadcast_shape() {
+        assert_eq!(broadcast_shape(&[3, 1], &[1, 4]).unwrap().as_slice(), &[3, 4]);
+        assert_eq!(broadcast_shape(&[2, 3, 4], &[4]).unwrap().as_slice(), &[2, 3, 4]);
+        assert_eq!(broadcast_shape(&[1], &[5, 3]).unwrap().as_slice(), &[5, 3]);
+        assert!(broadcast_shape(&[3], &[4]).is_none());
+    }
+
+    #[test]
+    fn test_broadcast_strides() {
+        // [3,1] with strides [1,1] broadcast to [3,4]
+        assert_eq!(broadcast_strides(&[3, 1], &[1, 1], &[3, 4]).as_slice(), &[1, 0]);
+    }
+}

crates/xserv-tensor/src/storage.rs

@@ -0,0 +1,133 @@
+use std::sync::Arc;
+use xserv_cuda::{GpuBuffer, Result as CudaResult};
+
+enum StorageInner {
+    Cpu { data: Vec<u8> },
+    Cuda { buffer: GpuBuffer, device: u32 },
+}
+
+/// Reference-counted storage for tensor data. Multiple tensors can share
+/// the same storage (e.g., after transpose or slice — view semantics).
+#[derive(Clone)]
+pub struct Storage(Arc<StorageInner>);
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum Device {
+    Cpu,
+    Cuda(u32),
+}
+
+impl std::fmt::Display for Device {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Device::Cpu => write!(f, "cpu"),
+            Device::Cuda(i) => write!(f, "cuda:{i}"),
+        }
+    }
+}
+
+impl Storage {
+    pub fn cpu(data: Vec<u8>) -> Self {
+        Self(Arc::new(StorageInner::Cpu { data }))
+    }
+
+    pub fn cuda(buffer: GpuBuffer, device: u32) -> Self {
+        Self(Arc::new(StorageInner::Cuda { buffer, device }))
+    }
+
+    pub fn device(&self) -> Device {
+        match self.0.as_ref() {
+            StorageInner::Cpu { .. } => Device::Cpu,
+            StorageInner::Cuda { device, .. } => Device::Cuda(*device),
+        }
+    }
+
+    pub fn len_bytes(&self) -> usize {
+        match self.0.as_ref() {
+            StorageInner::Cpu { data } => data.len(),
+            StorageInner::Cuda { buffer, .. } => buffer.len(),
+        }
+    }
+
+    /// Get a read-only view of CPU data. Panics if storage is on GPU.
+    pub fn as_cpu_bytes(&self) -> &[u8] {
+        match self.0.as_ref() {
+            StorageInner::Cpu { data } => data,
+            StorageInner::Cuda { .. } => panic!("cannot access GPU storage as CPU bytes"),
+        }
+    }
+
+    pub fn gpu_buffer(&self) -> &GpuBuffer {
+        match self.0.as_ref() {
+            StorageInner::Cuda { buffer, .. } => buffer,
+            StorageInner::Cpu { .. } => panic!("cannot access CPU storage as GPU buffer"),
+        }
+    }
+
+    /// Copy to a different device. If already on the target device, clones the Arc (no copy).
+    pub fn to_device(&self, target: Device) -> CudaResult<Self> {
+        let current = self.device();
+        if current == target {
+            return Ok(self.clone());
+        }
+        match (current, target) {
+            (Device::Cpu, Device::Cuda(dev)) => {
+                let cpu_data = self.as_cpu_bytes();
+                let mut buf = xserv_cuda::allocator::cached_alloc(cpu_data.len())?;
+                buf.copy_from_host(cpu_data)?;
+                Ok(Storage::cuda(buf, dev))
+            }
+            (Device::Cuda(_), Device::Cpu) => {
+                let gpu_buf = self.gpu_buffer();
+                let mut data = vec![0u8; gpu_buf.len()];
+                gpu_buf.copy_to_host(&mut data)?;
+                Ok(Storage::cpu(data))
+            }
+            (Device::Cuda(_), Device::Cuda(dev)) => {
+                let src = self.gpu_buffer();
+                let mut dst = xserv_cuda::allocator::cached_alloc(src.len())?;
+                dst.copy_from_device(src)?;
+                Ok(Storage::cuda(dst, dev))
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    /// Create a new owned copy of the storage on the same device.
+    pub fn deep_copy(&self) -> CudaResult<Self> {
+        match self.0.as_ref() {
+            StorageInner::Cpu { data } => Ok(Storage::cpu(data.clone())),
+            StorageInner::Cuda { buffer, device } => {
+                let mut dst = xserv_cuda::allocator::cached_alloc(buffer.len())?;
+                dst.copy_from_device(buffer)?;
+                Ok(Storage::cuda(dst, *device))
+            }
+        }
+    }
+
+    /// Allocate zeroed storage on the given device.
+    pub fn zeros(len_bytes: usize, device: Device) -> CudaResult<Self> {
+        match device {
+            Device::Cpu => Ok(Storage::cpu(vec![0u8; len_bytes])),
+            Device::Cuda(dev) => {
+                let mut buf = xserv_cuda::allocator::cached_alloc(len_bytes)?;
+                buf.zero()?;
+                Ok(Storage::cuda(buf, dev))
+            }
+        }
+    }
+
+    /// Allocate storage **without zeroing** on the given device.
+    /// The buffer may contain stale data from the caching allocator's pool.
+    /// Only use when the caller guarantees the kernel will fully overwrite
+    /// every element before any read.
+    pub fn empty(len_bytes: usize, device: Device) -> CudaResult<Self> {
+        match device {
+            Device::Cpu => Ok(Storage::cpu(vec![0u8; len_bytes])), // CPU still zeros (cheap)
+            Device::Cuda(dev) => {
+                let buf = xserv_cuda::allocator::cached_alloc(len_bytes)?;
+                Ok(Storage::cuda(buf, dev))
+            }
+        }
+    }
+}

crates/xserv-tensor/src/tensor.rs

@@ -0,0 +1,332 @@
+use std::sync::OnceLock;
+
+use crate::dtype::{DType, TensorDType};
+use crate::shape::{self, Dims};
+use crate::storage::{Device, Storage};
+
+/// Global hook for GPU strided-to-contiguous copy.
+/// Set by `xserv-kernels` (or any crate that provides a GPU kernel) via
+/// `register_gpu_contiguous`. When set, `contiguous()` on a non-contiguous
+/// GPU tensor calls this instead of doing a CPU round-trip.
+static GPU_CONTIGUOUS_FN: OnceLock<fn(&Tensor) -> Tensor> = OnceLock::new();
+
+/// Register a function that makes a non-contiguous GPU tensor contiguous.
+/// Intended to be called once by the kernel crate at startup.
+pub fn register_gpu_contiguous(f: fn(&Tensor) -> Tensor) {
+    let _ = GPU_CONTIGUOUS_FN.set(f);
+}
+
+/// Multi-dimensional array with CPU or GPU storage.
+///
+/// Tensors support view semantics: transpose, slice, etc. share
+/// the underlying storage and only change shape/strides/offset.
+#[derive(Clone)]
+pub struct Tensor {
+    storage: Storage,
+    shape: Dims,
+    strides: Dims,
+    offset: usize,
+    dtype: DType,
+}
+
+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");
+        let bytes = unsafe {
+            std::slice::from_raw_parts(data.as_ptr() as *const u8, numel * T::DTYPE.size_bytes())
+        };
+        Self {
+            storage: Storage::cpu(bytes.to_vec()),
+            shape: Dims::from_slice(shape),
+            strides: shape::contiguous_strides(shape),
+            offset: 0,
+            dtype: T::DTYPE,
+        }
+    }
+
+    pub fn zeros(shape: &[usize], dtype: DType, device: Device) -> Self {
+        let numel = shape::num_elements(shape);
+        let len_bytes = numel * dtype.size_bytes();
+        let storage = Storage::zeros(len_bytes, device).expect("alloc failed");
+        Self {
+            storage,
+            shape: Dims::from_slice(shape),
+            strides: shape::contiguous_strides(shape),
+            offset: 0,
+            dtype,
+        }
+    }
+
+    /// Allocate a tensor **without zeroing** the backing memory.
+    /// The buffer may contain stale data. Only use when the calling kernel
+    /// will fully overwrite every element before any read.
+    pub fn empty(shape: &[usize], dtype: DType, device: Device) -> Self {
+        let numel = shape::num_elements(shape);
+        let len_bytes = numel * dtype.size_bytes();
+        let storage = Storage::empty(len_bytes, device).expect("alloc failed");
+        Self {
+            storage,
+            shape: Dims::from_slice(shape),
+            strides: shape::contiguous_strides(shape),
+            offset: 0,
+            dtype,
+        }
+    }
+
+    pub fn ones(shape: &[usize], dtype: DType) -> Self {
+        let numel = shape::num_elements(shape);
+        match dtype {
+            DType::F32 => Self::from_slice(&vec![1.0f32; numel], shape),
+            DType::F16 => Self::from_slice(&vec![half::f16::from_f32(1.0); numel], shape),
+            DType::BF16 => Self::from_slice(&vec![half::bf16::from_f32(1.0); numel], shape),
+        }
+    }
+
+    // --- Properties ---
+
+    pub fn shape(&self) -> &[usize] { &self.shape }
+    pub fn strides(&self) -> &[usize] { &self.strides }
+    pub fn dtype(&self) -> DType { self.dtype }
+    pub fn ndim(&self) -> usize { self.shape.len() }
+    pub fn numel(&self) -> usize { shape::num_elements(&self.shape) }
+    pub fn offset(&self) -> usize { self.offset }
+
+    pub fn device(&self) -> Device { self.storage.device() }
+
+    pub fn is_contiguous(&self) -> bool {
+        shape::is_contiguous(&self.shape, &self.strides)
+    }
+
+    // --- Shape operations (view, no copy) ---
+
+    pub fn reshape(&self, new_shape: &[usize]) -> Self {
+        assert!(self.is_contiguous(), "reshape requires contiguous tensor");
+        let new_numel: usize = new_shape.iter().product();
+        assert_eq!(new_numel, self.numel(), "reshape numel mismatch");
+        Self {
+            storage: self.storage.clone(),
+            shape: Dims::from_slice(new_shape),
+            strides: shape::contiguous_strides(new_shape),
+            offset: self.offset,
+            dtype: self.dtype,
+        }
+    }
+
+    pub fn transpose(&self, dim0: usize, dim1: usize) -> Self {
+        assert!(dim0 < self.ndim() && dim1 < self.ndim());
+        let mut new_shape = self.shape.clone();
+        let mut new_strides = self.strides.clone();
+        new_shape.swap(dim0, dim1);
+        new_strides.swap(dim0, dim1);
+        Self {
+            storage: self.storage.clone(),
+            shape: new_shape,
+            strides: new_strides,
+            offset: self.offset,
+            dtype: self.dtype,
+        }
+    }
+
+    pub fn squeeze(&self, dim: usize) -> Self {
+        assert!(dim < self.ndim() && self.shape[dim] == 1);
+        let mut new_shape = self.shape.clone();
+        let mut new_strides = self.strides.clone();
+        new_shape.remove(dim);
+        new_strides.remove(dim);
+        Self {
+            storage: self.storage.clone(),
+            shape: new_shape,
+            strides: new_strides,
+            offset: self.offset,
+            dtype: self.dtype,
+        }
+    }
+
+    pub fn unsqueeze(&self, dim: usize) -> Self {
+        assert!(dim <= self.ndim());
+        let mut new_shape = self.shape.clone();
+        new_shape.insert(dim, 1);
+        let new_strides = if self.is_contiguous() {
+            shape::contiguous_strides(&new_shape)
+        } else {
+            let mut s = self.strides.clone();
+            let stride_val = if dim < self.strides.len() { self.strides[dim] } else { 1 };
+            s.insert(dim, stride_val);
+            s
+        };
+        Self {
+            storage: self.storage.clone(),
+            shape: new_shape,
+            strides: new_strides,
+            offset: self.offset,
+            dtype: self.dtype,
+        }
+    }
+
+    /// Make contiguous: if already contiguous, return clone (shared storage).
+    /// Otherwise, copy data into a new contiguous buffer.
+    pub fn contiguous(&self) -> Self {
+        if self.is_contiguous() {
+            return self.clone();
+        }
+        // For GPU tensors: use the registered GPU kernel if available,
+        // otherwise fall back to CPU round-trip.
+        if matches!(self.device(), Device::Cuda(_)) {
+            if let Some(gpu_fn) = GPU_CONTIGUOUS_FN.get() {
+                return gpu_fn(self);
+            }
+            let cpu = self.to_device(Device::Cpu);
+            let contig = cpu.contiguous();
+            return contig.to_device(self.device());
+        }
+        let numel = self.numel();
+        let elem_size = self.dtype.size_bytes();
+        let src_bytes = self.storage.as_cpu_bytes();
+        let mut dst = vec![0u8; numel * elem_size];
+        // Iterate all elements using strides
+        let ndim = self.ndim();
+        let mut idx = vec![0usize; ndim];
+        for flat in 0..numel {
+            let src_offset = self.offset + idx.iter().zip(self.strides.iter()).map(|(i, s)| i * s).sum::<usize>();
+            let src_byte_offset = src_offset * elem_size;
+            let dst_byte_offset = flat * elem_size;
+            dst[dst_byte_offset..dst_byte_offset + elem_size]
+                .copy_from_slice(&src_bytes[src_byte_offset..src_byte_offset + elem_size]);
+            // Increment index (rightmost first)
+            for d in (0..ndim).rev() {
+                idx[d] += 1;
+                if idx[d] < self.shape[d] {
+                    break;
+                }
+                idx[d] = 0;
+            }
+        }
+        Self {
+            storage: Storage::cpu(dst),
+            shape: self.shape.clone(),
+            strides: shape::contiguous_strides(&self.shape),
+            offset: 0,
+            dtype: self.dtype,
+        }
+    }
+
+    // --- Device transfer ---
+
+    pub fn to_device(&self, device: Device) -> Self {
+        if self.device() == device {
+            return self.clone();
+        }
+        // Transfer the raw storage (preserving strides/offset).
+        // Non-contiguous layout is preserved — the user can call contiguous() after.
+        let new_storage = self.storage.to_device(device).expect("device transfer failed");
+        Self {
+            storage: new_storage,
+            shape: self.shape.clone(),
+            strides: self.strides.clone(),
+            offset: self.offset,
+            dtype: self.dtype,
+        }
+    }
+
+    // --- Data access (CPU only) ---
+
+    /// Read tensor data as a typed slice. Requires contiguous CPU tensor.
+    pub fn as_slice<T: TensorDType>(&self) -> &[T] {
+        assert_eq!(T::DTYPE, self.dtype, "dtype mismatch");
+        assert!(self.is_contiguous(), "as_slice requires contiguous");
+        assert_eq!(self.device(), Device::Cpu, "as_slice requires CPU");
+        let bytes = self.storage.as_cpu_bytes();
+        let elem_size = self.dtype.size_bytes();
+        let start = self.offset * elem_size;
+        let len = self.numel();
+        unsafe { std::slice::from_raw_parts(bytes[start..].as_ptr() as *const T, len) }
+    }
+
+    /// Raw pointer to storage start (for GPU kernel launch).
+    pub fn data_ptr(&self) -> *const u8 {
+        match self.device() {
+            Device::Cpu => {
+                let bytes = self.storage.as_cpu_bytes();
+                unsafe { bytes.as_ptr().add(self.offset * self.dtype.size_bytes()) }
+            }
+            Device::Cuda(_) => {
+                let buf = self.storage.gpu_buffer();
+                unsafe { buf.as_ptr().add(self.offset * self.dtype.size_bytes()) }
+            }
+        }
+    }
+
+    pub fn storage(&self) -> &Storage { &self.storage }
+}
+
+impl std::fmt::Debug for Tensor {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f, "Tensor(shape={:?}, dtype={}, device={}, contiguous={})",
+            self.shape.as_slice(), self.dtype, self.device(), self.is_contiguous()
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn contiguous_2d() -> Tensor {
+        Tensor::from_slice(&[1.0f32; 12], &[3, 4])
+    }
+
+    #[test]
+    fn unsqueeze_dim0_contiguous() {
+        let t = contiguous_2d();
+        let u = t.unsqueeze(0);
+        assert_eq!(u.shape(), &[1, 3, 4]);
+        assert!(u.is_contiguous());
+        assert_eq!(u.strides(), &[12, 4, 1]);
+    }
+
+    #[test]
+    fn unsqueeze_dim1_contiguous() {
+        let t = contiguous_2d();
+        let u = t.unsqueeze(1);
+        assert_eq!(u.shape(), &[3, 1, 4]);
+        assert!(u.is_contiguous());
+        assert_eq!(u.strides(), &[4, 4, 1]);
+    }
+
+    #[test]
+    fn unsqueeze_dim2_contiguous() {
+        let t = contiguous_2d();
+        let u = t.unsqueeze(2);
+        assert_eq!(u.shape(), &[3, 4, 1]);
+        assert!(u.is_contiguous());
+        assert_eq!(u.strides(), &[4, 1, 1]);
+    }
+
+    #[test]
+    fn unsqueeze_noncontiguous() {
+        // Transpose makes [3,4] into [4,3] with strides [1,4] (non-contiguous)
+        let t = contiguous_2d().transpose(0, 1);
+        assert!(!t.is_contiguous());
+        let u = t.unsqueeze(0);
+        assert_eq!(u.shape(), &[1, 4, 3]);
+        // Non-contiguous path: stride_val copied from strides[0]=1
+        assert_eq!(u.strides(), &[1, 1, 4]);
+    }
+
+    #[test]
+    fn unsqueeze_squeeze_roundtrip() {
+        let t = contiguous_2d();
+        let u = t.unsqueeze(1).squeeze(1);
+        assert_eq!(u.shape(), t.shape());
+        assert!(u.is_contiguous());
+    }
+}

crates/xserv-tensor/tests/integration.rs

@@ -0,0 +1,127 @@
+use half::bf16;
+use xserv_tensor::*;
+
+#[test]
+fn test_from_slice_and_shape() {
+    let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
+    let t = Tensor::from_slice(&data, &[2, 3]);
+    assert_eq!(t.shape(), &[2, 3]);
+    assert_eq!(t.strides(), &[3, 1]);
+    assert_eq!(t.numel(), 6);
+    assert_eq!(t.ndim(), 2);
+    assert!(t.is_contiguous());
+    assert_eq!(t.dtype(), DType::F32);
+    assert_eq!(t.device(), Device::Cpu);
+}
+
+#[test]
+fn test_as_slice() {
+    let data = vec![1.0f32, 2.0, 3.0, 4.0];
+    let t = Tensor::from_slice(&data, &[4]);
+    assert_eq!(t.as_slice::<f32>(), &[1.0, 2.0, 3.0, 4.0]);
+}
+
+#[test]
+fn test_zeros_and_ones() {
+    let z = Tensor::zeros(&[2, 3], DType::F32, Device::Cpu);
+    assert_eq!(z.as_slice::<f32>(), &[0.0; 6]);
+
+    let o = Tensor::ones(&[3], DType::F32);
+    assert_eq!(o.as_slice::<f32>(), &[1.0, 1.0, 1.0]);
+}
+
+#[test]
+fn test_bf16_tensor() {
+    let data: Vec<bf16> = vec![bf16::from_f32(1.0), bf16::from_f32(2.5), bf16::from_f32(-3.0)];
+    let t = Tensor::from_slice(&data, &[3]);
+    assert_eq!(t.dtype(), DType::BF16);
+    let out = t.as_slice::<bf16>();
+    assert_eq!(out[0].to_f32(), 1.0);
+    assert!((out[1].to_f32() - 2.5).abs() < 0.01);
+    assert_eq!(out[2].to_f32(), -3.0);
+}
+
+#[test]
+fn test_reshape() {
+    let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
+    let t = Tensor::from_slice(&data, &[2, 3]);
+    let t2 = t.reshape(&[3, 2]);
+    assert_eq!(t2.shape(), &[3, 2]);
+    assert_eq!(t2.as_slice::<f32>(), &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
+
+    let t3 = t.reshape(&[6]);
+    assert_eq!(t3.shape(), &[6]);
+}
+
+#[test]
+fn test_transpose() {
+    let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
+    let t = Tensor::from_slice(&data, &[2, 3]);
+    let tt = t.transpose(0, 1);
+    assert_eq!(tt.shape(), &[3, 2]);
+    assert_eq!(tt.strides(), &[1, 3]);
+    assert!(!tt.is_contiguous());
+}
+
+#[test]
+fn test_contiguous_from_transpose() {
+    let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
+    // Original [2,3]: [[1,2,3],[4,5,6]]
+    let t = Tensor::from_slice(&data, &[2, 3]);
+    // Transpose to [3,2]: [[1,4],[2,5],[3,6]]
+    let tt = t.transpose(0, 1);
+    let tc = tt.contiguous();
+    assert!(tc.is_contiguous());
+    assert_eq!(tc.shape(), &[3, 2]);
+    assert_eq!(tc.as_slice::<f32>(), &[1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
+}
+
+#[test]
+fn test_squeeze_unsqueeze() {
+    let data = vec![1.0f32, 2.0, 3.0];
+    let t = Tensor::from_slice(&data, &[1, 3]);
+    let squeezed = t.squeeze(0);
+    assert_eq!(squeezed.shape(), &[3]);
+
+    let unsqueezed = squeezed.unsqueeze(0);
+    assert_eq!(unsqueezed.shape(), &[1, 3]);
+
+    let unsqueezed2 = squeezed.unsqueeze(1);
+    assert_eq!(unsqueezed2.shape(), &[3, 1]);
+}
+
+#[test]
+fn test_cpu_to_gpu_roundtrip() {
+    xserv_cuda::device::set_device(0).unwrap();
+
+    let data = vec![1.0f32, 2.0, 3.0, 4.0];
+    let cpu_t = Tensor::from_slice(&data, &[2, 2]);
+    let gpu_t = cpu_t.to_device(Device::Cuda(0));
+    assert_eq!(gpu_t.device(), Device::Cuda(0));
+    assert_eq!(gpu_t.shape(), &[2, 2]);
+
+    let back = gpu_t.to_device(Device::Cpu);
+    assert_eq!(back.device(), Device::Cpu);
+    assert_eq!(back.as_slice::<f32>(), &[1.0, 2.0, 3.0, 4.0]);
+}
+
+#[test]
+fn test_zeros_gpu() {
+    xserv_cuda::device::set_device(0).unwrap();
+
+    let t = Tensor::zeros(&[4, 4], DType::F32, Device::Cuda(0));
+    assert_eq!(t.device(), Device::Cuda(0));
+    assert_eq!(t.shape(), &[4, 4]);
+
+    let cpu = t.to_device(Device::Cpu);
+    assert_eq!(cpu.as_slice::<f32>(), &[0.0f32; 16]);
+}
+
+#[test]
+fn test_debug_format() {
+    let t = Tensor::from_slice(&[1.0f32], &[1]);
+    let dbg = format!("{:?}", t);
+    assert!(dbg.contains("shape=[1]"));
+    assert!(dbg.contains("f32"));
+    assert!(dbg.contains("cpu"));
+}

crates/xserv-tokenizer/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "xserv-tokenizer"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+serde.workspace = true
+serde_json.workspace = true
+regex.workspace = true

crates/xserv-tokenizer/src/bpe.rs

@@ -0,0 +1,281 @@
+use regex::Regex;
+use serde::Deserialize;
+use std::collections::HashMap;
+use std::path::Path;
+
+pub struct Tokenizer {
+    encoder: HashMap<Vec<u8>, u32>,
+    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)]
+struct TokenizerJson {
+    model: ModelSection,
+    #[serde(default)]
+    added_tokens: Vec<AddedToken>,
+}
+
+#[derive(Deserialize)]
+struct ModelSection {
+    vocab: HashMap<String, u32>,
+    merges: Vec<MergeEntry>,
+    #[serde(default)]
+    byte_fallback: bool,
+}
+
+#[derive(Deserialize)]
+#[serde(untagged)]
+enum MergeEntry {
+    Str(String),
+    Pair(Vec<String>),
+}
+
+#[derive(Deserialize)]
+struct AddedToken {
+    id: u32,
+    content: String,
+    special: bool,
+}
+
+impl Tokenizer {
+    pub fn from_file(path: &Path) -> Self {
+        let data = std::fs::read_to_string(path)
+            .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
+        let tj: TokenizerJson = serde_json::from_str(&data)
+            .unwrap_or_else(|e| panic!("failed to parse tokenizer.json: {e}"));
+
+        // 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);
+            encoder.insert(bytes, id);
+        }
+
+        // Build decoder: ID → token bytes
+        let max_id = tj.model.vocab.values().copied().max().unwrap_or(0);
+        let added_max = tj.added_tokens.iter().map(|t| t.id).max().unwrap_or(0);
+        let vocab_size = (max_id.max(added_max) + 1) as usize;
+        let mut decoder = vec![vec![]; vocab_size];
+        for (token_str, &id) in &tj.model.vocab {
+            decoder[id as usize] = token_str_to_bytes(token_str);
+        }
+
+        // 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, 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);
+            }
+        }
+
+        // Special tokens
+        let mut special_tokens = HashMap::new();
+        let mut special_token_ids = HashMap::new();
+        let mut eos_token_id = None;
+        for at in &tj.added_tokens {
+            if at.special {
+                special_tokens.insert(at.content.clone(), at.id);
+                special_token_ids.insert(at.id, at.content.clone());
+                decoder.resize(decoder.len().max(at.id as usize + 1), vec![]);
+                decoder[at.id as usize] = at.content.as_bytes().to_vec();
+                if at.content == "<|endoftext|>" || at.content == "<|end_of_text|>" {
+                    eos_token_id = Some(at.id);
+                }
+            }
+        }
+
+        // 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,
+            decoder,
+            merge_ranks,
+            special_tokens,
+            special_token_ids,
+            pre_tokenize_re,
+            eos_token_id,
+            byte_fallback,
+        }
+    }
+
+    pub fn encode(&self, text: &str) -> Vec<u32> {
+        let mut tokens = Vec::new();
+
+        // Check for special tokens first (split around them)
+        let mut remaining = text;
+        while !remaining.is_empty() {
+            // Find earliest special token
+            let mut earliest: Option<(usize, &str, u32)> = None;
+            for (st, &id) in &self.special_tokens {
+                if let Some(pos) = remaining.find(st.as_str()) {
+                    if earliest.is_none() || pos < earliest.unwrap().0 {
+                        earliest = Some((pos, st, id));
+                    }
+                }
+            }
+
+            if let Some((pos, st, id)) = earliest {
+                if pos > 0 {
+                    self.encode_ordinary(&remaining[..pos], &mut tokens);
+                }
+                tokens.push(id);
+                remaining = &remaining[pos + st.len()..];
+            } else {
+                self.encode_ordinary(remaining, &mut tokens);
+                break;
+            }
+        }
+
+        tokens
+    }
+
+    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().filter_map(|&b| {
+                if let Some(&id) = self.encoder.get(&vec![b]) {
+                    Some(id)
+                } else if self.byte_fallback {
+                    let hex_token = format!("<0x{:02X}>", b);
+                    if let Some(&id) = self.special_tokens.get(&hex_token) {
+                        Some(id)
+                    } else if let Some(&id) = self.encoder.get(hex_token.as_bytes()) {
+                        Some(id)
+                    } else if let Some(&unk_id) = self.special_tokens.get("<unk>") {
+                        eprintln!("warning: byte 0x{b:02X} not in vocab, using <unk> token");
+                        Some(unk_id)
+                    } else {
+                        eprintln!("warning: byte 0x{b:02X} not in vocab and no fallback token, using token 0");
+                        Some(0)
+                    }
+                } else {
+                    eprintln!("warning: byte {b} (0x{b:02X}) not in vocab, skipping");
+                    None
+                }
+            }).collect();
+
+            // BPE merges
+            loop {
+                if token_ids.len() < 2 { break; }
+                let mut best_rank = usize::MAX;
+                let mut best_idx = 0;
+                for i in 0..token_ids.len() - 1 {
+                    if let Some(&rank) = self.merge_ranks.get(&(token_ids[i], token_ids[i + 1])) {
+                        if rank < best_rank {
+                            best_rank = rank;
+                            best_idx = i;
+                        }
+                    }
+                }
+                if best_rank == usize::MAX { break; }
+
+                let merged_bytes = [
+                    self.decoder[token_ids[best_idx] as usize].as_slice(),
+                    self.decoder[token_ids[best_idx + 1] as usize].as_slice(),
+                ].concat();
+                let merged_id = *self.encoder.get(&merged_bytes).unwrap_or_else(|| {
+                    panic!("merged token not in vocab");
+                });
+                token_ids[best_idx] = merged_id;
+                token_ids.remove(best_idx + 1);
+            }
+
+            out.extend_from_slice(&token_ids);
+        }
+    }
+
+    pub fn decode(&self, token_ids: &[u32]) -> String {
+        let mut bytes = Vec::new();
+        for &id in token_ids {
+            if let Some(b) = self.decoder.get(id as usize) {
+                bytes.extend_from_slice(b);
+            }
+        }
+        String::from_utf8_lossy(&bytes).into_owned()
+    }
+
+    pub fn eos_token_id(&self) -> Option<u32> {
+        self.eos_token_id
+    }
+
+    pub fn vocab_size(&self) -> usize {
+        self.decoder.len()
+    }
+
+    pub fn special_token_id(&self, name: &str) -> Option<u32> {
+        self.special_tokens.get(name).copied()
+    }
+}
+
+/// Convert a token string from HF vocab (which uses Unicode replacements for bytes)
+/// back to raw bytes. GPT-2 uses a byte-to-unicode mapping where e.g. byte 0x20 (space)
+/// is represented as 'Ġ' (U+0120).
+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 {
+    // 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);
+        }
+        m
+    });
+
+    *map.get(&(c as u32)).unwrap_or_else(|| {
+        panic!("unmapped unicode char U+{:04X} in tokenizer", c as u32)
+    })
+}

crates/xserv-tokenizer/src/lib.rs

@@ -0,0 +1,3 @@
+pub mod bpe;
+
+pub use bpe::Tokenizer;

csrc/activation/activations.cu

@@ -0,0 +1,154 @@
+#include <cuda_bf16.h>
+#include <math.h>
+
+// GELU (tanh approximation):
+// gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
+__device__ __forceinline__ float gelu_f(float x) {
+    const float SQRT_2_OVER_PI = 0.7978845608f;
+    float cube = x * x * x;
+    float inner = SQRT_2_OVER_PI * (x + 0.044715f * cube);
+    return 0.5f * x * (1.0f + tanhf(inner));
+}
+
+// SiLU (Swish): silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
+__device__ __forceinline__ float silu_f(float x) {
+    return x / (1.0f + expf(-x));
+}
+
+__global__ void gelu_f32(const float* x, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = gelu_f(x[idx]);
+}
+
+__global__ void gelu_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(gelu_f(__bfloat162float(x[idx])));
+}
+
+__global__ void silu_f32(const float* x, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = silu_f(x[idx]);
+}
+
+__global__ void silu_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(silu_f(__bfloat162float(x[idx])));
+}
+
+__global__ void scale_f32_kernel(const float* x, float* out, float scale, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = x[idx] * scale;
+}
+
+__global__ void scale_bf16_kernel(const __nv_bfloat16* x, __nv_bfloat16* out, float scale, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(x[idx]) * scale);
+}
+
+// Fused SiLU×Mul: out = silu(gate) * up
+__global__ void silu_mul_bf16_kernel(const __nv_bfloat16* gate, const __nv_bfloat16* up,
+                                     __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        float g = __bfloat162float(gate[idx]);
+        float u = __bfloat162float(up[idx]);
+        float silu_g = g / (1.0f + expf(-g));
+        out[idx] = __float2bfloat16(silu_g * u);
+    }
+}
+
+// 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) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    gelu_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
+}
+
+void launch_gelu_bf16(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    gelu_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+}
+
+void launch_silu_f32(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    silu_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
+}
+
+void launch_silu_bf16(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    silu_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+}
+
+void launch_scale_f32(const void* x, void* out, float scale, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    scale_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)x, (float*)out, scale, n);
+}
+
+void launch_scale_bf16(const void* x, void* out, float scale, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    scale_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (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);
+}
+
+void launch_silu_mul_bf16(const void* gate, const void* up, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    silu_mul_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)gate, (const __nv_bfloat16*)up, (__nv_bfloat16*)out, n);
+}
+
+}

csrc/attention/causal_mask.cu

@@ -0,0 +1,52 @@
+#include <cuda_bf16.h>
+
+// Apply causal mask: set scores[row][col] = -inf where col > row + offset.
+// offset is used for KV cache: when query starts at position `offset`,
+// we allow attending to positions [0, offset + row].
+// scores: [batch, rows, cols]  (flattened batch×heads)
+
+__global__ void causal_mask_f32(
+    float* __restrict__ scores,
+    int rows, int cols, int offset
+) {
+    int batch_idx = blockIdx.z;
+    int row = blockIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (col < cols && col > row + offset) {
+        scores[batch_idx * rows * cols + row * cols + col] = -INFINITY;
+    }
+}
+
+__global__ void causal_mask_bf16(
+    __nv_bfloat16* __restrict__ scores,
+    int rows, int cols, int offset
+) {
+    int batch_idx = blockIdx.z;
+    int row = blockIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (col < cols && col > row + offset) {
+        scores[batch_idx * rows * cols + row * cols + col] = __float2bfloat16(-INFINITY);
+    }
+}
+
+extern "C" {
+
+void launch_causal_mask_f32(void* scores, int batch, int rows, int cols,
+                            int offset, void* stream) {
+    int block = 256;
+    dim3 grid((cols + block - 1) / block, rows, batch);
+    causal_mask_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (float*)scores, rows, cols, offset);
+}
+
+void launch_causal_mask_bf16(void* scores, int batch, int rows, int cols,
+                             int offset, void* stream) {
+    int block = 256;
+    dim3 grid((cols + block - 1) / block, rows, batch);
+    causal_mask_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (__nv_bfloat16*)scores, rows, cols, offset);
+}
+
+}

csrc/attention/flash_attention.cu

@@ -0,0 +1,416 @@
+#include <cuda_bf16.h>
+#include <float.h>
+
+// Flash Attention 2 forward kernel for BF16 with FP32 accumulation.
+//
+// Algorithm: outer loop over Q tiles (BR rows), inner loop over K/V tiles (BC rows).
+// Uses online softmax — no O(S^2) memory.
+//
+// Layout: Q [batch, num_q_heads, q_len, head_dim]
+//         K [batch, num_kv_heads, kv_len, head_dim]
+//         V [batch, num_kv_heads, kv_len, head_dim]
+//         O [batch, num_q_heads, q_len, head_dim]
+//
+// Shared memory (BF16):
+//   smem_q[BR][head_dim]   — 64 * 128 * 2 = 16 KB (loaded once per Q tile)
+//   smem_kv[BC][head_dim]  — 64 * 128 * 2 = 16 KB (alternates K and V)
+//   Total: 32 KB (fits in default 48 KB shared memory)
+
+#define BR 64
+#define BC 64
+#define THREADS_PER_BLOCK 128
+
+__global__ void flash_attention_bf16_kernel(
+    const __nv_bfloat16* __restrict__ Q,
+    const __nv_bfloat16* __restrict__ K,
+    const __nv_bfloat16* __restrict__ V,
+    __nv_bfloat16* __restrict__ O,
+    int num_q_heads, int num_kv_heads,
+    int q_len, int kv_len, int head_dim,
+    float scale, int causal
+) {
+    // Grid: (ceil(q_len / BR), batch * num_q_heads)
+    int q_tile_idx = blockIdx.x;
+    int bh = blockIdx.y;
+    int batch_idx = bh / num_q_heads;
+    int q_head = bh % num_q_heads;
+
+    // GQA: map Q head to KV head
+    int heads_per_group = num_q_heads / num_kv_heads;
+    int kv_head = q_head / heads_per_group;
+
+    int q_tile_start = q_tile_idx * BR;
+    if (q_tile_start >= q_len) return;
+    int q_tile_rows = min(BR, q_len - q_tile_start);
+
+    // Pointers to this batch/head's data
+    const __nv_bfloat16* Q_head = Q + ((long long)batch_idx * num_q_heads + q_head) * q_len * head_dim;
+    const __nv_bfloat16* K_head = K + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
+    const __nv_bfloat16* V_head = V + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
+    __nv_bfloat16* O_head = O + ((long long)batch_idx * num_q_heads + q_head) * q_len * head_dim;
+
+    int tid = threadIdx.x;
+
+    // Dynamic shared memory
+    extern __shared__ __nv_bfloat16 smem[];
+    __nv_bfloat16* smem_q  = smem;                 // BR * head_dim elements
+    __nv_bfloat16* smem_kv = smem + BR * head_dim;  // BC * head_dim elements
+
+    // ---- Load Q tile into shared memory (cooperative) ----
+    int q_elems = q_tile_rows * head_dim;
+    for (int i = tid; i < q_elems; i += THREADS_PER_BLOCK) {
+        int row = i / head_dim;
+        int col = i % head_dim;
+        smem_q[row * head_dim + col] = Q_head[(q_tile_start + row) * head_dim + col];
+    }
+    // Zero-pad if q_tile_rows < BR
+    for (int i = q_elems + tid; i < BR * head_dim; i += THREADS_PER_BLOCK) {
+        smem_q[i] = __float2bfloat16(0.0f);
+    }
+    __syncthreads();
+
+    // Thread t (0 <= t < q_tile_rows) owns Q row t
+    bool owns_row = (tid < q_tile_rows);
+
+    // Per-thread FP32 accumulators (head_dim up to 128)
+    float O_acc[128];
+    float m_val = -INFINITY;
+    float l_val = 0.0f;
+    if (owns_row) {
+        for (int d = 0; d < head_dim; d++) {
+            O_acc[d] = 0.0f;
+        }
+    }
+
+    // kv_offset handles cached KV longer than Q (decode step)
+    int kv_offset = kv_len - q_len;
+    int num_kv_tiles = (kv_len + BC - 1) / BC;
+
+    // ---- Inner loop over K/V tiles ----
+    for (int j = 0; j < num_kv_tiles; j++) {
+        int kv_tile_start = j * BC;
+        int kv_tile_cols = min(BC, kv_len - kv_tile_start);
+
+        // Causal: skip entire tile if all K positions are in the future
+        if (causal) {
+            int max_allowed_kv = (q_tile_start + q_tile_rows - 1) + kv_offset;
+            if (kv_tile_start > max_allowed_kv) {
+                continue;
+            }
+        }
+
+        // ---- Load K tile into smem_kv ----
+        int kv_elems = kv_tile_cols * head_dim;
+        for (int i = tid; i < kv_elems; i += THREADS_PER_BLOCK) {
+            int row = i / head_dim;
+            int col = i % head_dim;
+            smem_kv[row * head_dim + col] = K_head[(kv_tile_start + row) * head_dim + col];
+        }
+        for (int i = kv_elems + tid; i < BC * head_dim; i += THREADS_PER_BLOCK) {
+            smem_kv[i] = __float2bfloat16(0.0f);
+        }
+        __syncthreads();
+
+        // ---- Compute S = Q @ K^T * scale, causal mask, online softmax ----
+        float P[BC];
+
+        if (owns_row) {
+            float row_max = -INFINITY;
+            for (int c = 0; c < kv_tile_cols; c++) {
+                float dot = 0.0f;
+                for (int d = 0; d < head_dim; d++) {
+                    dot += __bfloat162float(smem_q[tid * head_dim + d])
+                         * __bfloat162float(smem_kv[c * head_dim + d]);
+                }
+                float s = dot * scale;
+
+                if (causal) {
+                    int q_pos = q_tile_start + tid;
+                    int kv_pos = kv_tile_start + c;
+                    if (kv_pos > q_pos + kv_offset) {
+                        s = -INFINITY;
+                    }
+                }
+
+                P[c] = s;  // store score temporarily in P
+                row_max = fmaxf(row_max, s);
+            }
+
+            // Online softmax: m_new, P = exp(S - m_new), l_new
+            float m_new = fmaxf(m_val, row_max);
+
+            float psum = 0.0f;
+            for (int c = 0; c < kv_tile_cols; c++) {
+                P[c] = expf(P[c] - m_new);
+                psum += P[c];
+            }
+
+            // Rescale previous accumulator
+            float correction = expf(m_val - m_new);
+            l_val = correction * l_val + psum;
+
+            for (int d = 0; d < head_dim; d++) {
+                O_acc[d] *= correction;
+            }
+
+            m_val = m_new;
+        }
+
+        // Sync before overwriting smem_kv with V tile
+        __syncthreads();
+
+        // ---- Load V tile (reuse smem_kv) ----
+        int v_elems = kv_tile_cols * head_dim;
+        for (int i = tid; i < v_elems; i += THREADS_PER_BLOCK) {
+            int row = i / head_dim;
+            int col = i % head_dim;
+            smem_kv[row * head_dim + col] = V_head[(kv_tile_start + row) * head_dim + col];
+        }
+        for (int i = v_elems + tid; i < BC * head_dim; i += THREADS_PER_BLOCK) {
+            smem_kv[i] = __float2bfloat16(0.0f);
+        }
+        __syncthreads();
+
+        // ---- Accumulate O += P @ V_tile ----
+        if (owns_row) {
+            for (int c = 0; c < kv_tile_cols; c++) {
+                float p = P[c];
+                if (p != 0.0f) {
+                    for (int d = 0; d < head_dim; d++) {
+                        O_acc[d] += p * __bfloat162float(smem_kv[c * head_dim + d]);
+                    }
+                }
+            }
+        }
+
+        __syncthreads();
+    }
+
+    // ---- Final normalize and write output (convert FP32 → BF16) ----
+    if (owns_row) {
+        float inv_l = (l_val > 0.0f) ? (1.0f / l_val) : 0.0f;
+        int global_row = q_tile_start + tid;
+        for (int d = 0; d < head_dim; d++) {
+            O_head[global_row * head_dim + d] = __float2bfloat16(O_acc[d] * inv_l);
+        }
+    }
+}
+
+// ============================================================
+// Decode Attention kernel: optimized for Q_len=1 (single-token decode).
+// Parallelizes across KV sequence dimension instead of Q rows.
+//
+// Grid: (batch * num_q_heads, 1) — one block per Q head
+// Block: 256 threads — each thread handles ceil(kv_len / 256) KV positions
+// Uses online softmax reduction across threads.
+// ============================================================
+
+#define DECODE_THREADS 256
+#define HEAD_DIM_MAX 128
+
+__global__ void decode_attention_bf16_kernel(
+    const __nv_bfloat16* __restrict__ Q,
+    const __nv_bfloat16* __restrict__ K,
+    const __nv_bfloat16* __restrict__ V,
+    __nv_bfloat16* __restrict__ O,
+    int num_q_heads, int num_kv_heads,
+    int kv_len, int head_dim,
+    float scale
+) {
+    int bh = blockIdx.x;
+    int batch_idx = bh / num_q_heads;
+    int q_head = bh % num_q_heads;
+
+    // GQA mapping
+    int heads_per_group = num_q_heads / num_kv_heads;
+    int kv_head = q_head / heads_per_group;
+
+    int tid = threadIdx.x;
+
+    // Pointers to this batch/head's data
+    // Q: [batch, num_q_heads, 1, head_dim]
+    const __nv_bfloat16* Q_ptr = Q + ((long long)batch_idx * num_q_heads + q_head) * head_dim;
+    // K/V: [batch, num_kv_heads, kv_len, head_dim]
+    const __nv_bfloat16* K_base = K + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
+    const __nv_bfloat16* V_base = V + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
+    __nv_bfloat16* O_ptr = O + ((long long)batch_idx * num_q_heads + q_head) * head_dim;
+
+    // Load Q vector into registers (head_dim <= 128)
+    float q_reg[HEAD_DIM_MAX];
+    for (int d = 0; d < head_dim; d++) {
+        q_reg[d] = __bfloat162float(Q_ptr[d]);
+    }
+
+    // Each thread processes a chunk of KV positions
+    // Thread tid handles positions: tid, tid+DECODE_THREADS, tid+2*DECODE_THREADS, ...
+    float local_max = -INFINITY;
+    float local_sum = 0.0f;
+    float local_O[HEAD_DIM_MAX];
+    for (int d = 0; d < head_dim; d++) {
+        local_O[d] = 0.0f;
+    }
+
+    for (int pos = tid; pos < kv_len; pos += DECODE_THREADS) {
+        // Compute dot(Q, K[pos]) * scale
+        const __nv_bfloat16* K_pos = K_base + pos * head_dim;
+        float dot = 0.0f;
+        for (int d = 0; d < head_dim; d++) {
+            dot += q_reg[d] * __bfloat162float(K_pos[d]);
+        }
+        float s = dot * scale;
+
+        // Online softmax update
+        float new_max = fmaxf(local_max, s);
+        float correction = expf(local_max - new_max);
+        float p = expf(s - new_max);
+
+        // Rescale running sum and O
+        local_sum = local_sum * correction + p;
+        for (int d = 0; d < head_dim; d++) {
+            local_O[d] = local_O[d] * correction;
+        }
+
+        // Accumulate V[pos] weighted by p
+        const __nv_bfloat16* V_pos = V_base + pos * head_dim;
+        for (int d = 0; d < head_dim; d++) {
+            local_O[d] += p * __bfloat162float(V_pos[d]);
+        }
+
+        local_max = new_max;
+    }
+
+    // --- Block-level online softmax reduction ---
+    // We need to combine (local_max, local_sum, local_O) across all threads.
+    // Strategy: reduce max, then each thread rescales, then reduce sum and O.
+
+    // Shared memory for reduction
+    __shared__ float smem_max[32];   // one per warp
+    __shared__ float smem_sum[32];
+    __shared__ float smem_O[HEAD_DIM_MAX]; // final output accumulator
+
+    // Step 1: Block-wide max reduction
+    int lane = tid & 31;
+    int warp_id = tid >> 5;
+    int num_warps = DECODE_THREADS >> 5; // 8 warps
+
+    float warp_max = local_max;
+    #pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1)
+        warp_max = fmaxf(warp_max, __shfl_down_sync(0xffffffff, warp_max, offset));
+    if (lane == 0) smem_max[warp_id] = warp_max;
+    __syncthreads();
+
+    float global_max;
+    if (tid == 0) {
+        global_max = smem_max[0];
+        for (int i = 1; i < num_warps; i++)
+            global_max = fmaxf(global_max, smem_max[i]);
+        smem_max[0] = global_max;
+    }
+    __syncthreads();
+    global_max = smem_max[0];
+
+    // Step 2: Each thread rescales its local_sum and local_O with global_max
+    float rescale = (local_max == -INFINITY) ? 0.0f : expf(local_max - global_max);
+    local_sum *= rescale;
+    for (int d = 0; d < head_dim; d++) {
+        local_O[d] *= rescale;
+    }
+
+    // Step 3: Reduce sum across block
+    float warp_sum = local_sum;
+    #pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1)
+        warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset);
+    if (lane == 0) smem_sum[warp_id] = warp_sum;
+    __syncthreads();
+
+    float global_sum;
+    if (tid == 0) {
+        global_sum = 0.0f;
+        for (int i = 0; i < num_warps; i++)
+            global_sum += smem_sum[i];
+        smem_sum[0] = global_sum;
+    }
+    __syncthreads();
+    global_sum = smem_sum[0];
+
+    // Step 4: Reduce O across block (dimension by dimension using shared mem)
+    float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
+
+    // Process head_dim in chunks: each iteration reduces one dimension
+    // Use shared memory accumulator: each warp contributes via warp reduction + atomic
+    // Actually simpler: iterate over dimensions, warp reduce each, then lane0 atomicAdd to smem_O
+
+    // Initialize smem_O
+    for (int d = tid; d < head_dim; d += DECODE_THREADS) {
+        smem_O[d] = 0.0f;
+    }
+    __syncthreads();
+
+    // Each thread adds its local_O contributions via warp reduction + atomicAdd
+    for (int d = 0; d < head_dim; d++) {
+        float val = local_O[d];
+        // Warp-level reduction
+        #pragma unroll
+        for (int offset = 16; offset > 0; offset >>= 1)
+            val += __shfl_down_sync(0xffffffff, val, offset);
+        if (lane == 0) {
+            atomicAdd(&smem_O[d], val);
+        }
+    }
+    __syncthreads();
+
+    // Thread 0..head_dim-1 write final output
+    for (int d = tid; d < head_dim; d += DECODE_THREADS) {
+        O_ptr[d] = __float2bfloat16(smem_O[d] * inv_sum);
+    }
+}
+
+extern "C" {
+
+void launch_flash_attention_bf16(
+    const void* Q, const void* K, const void* V, void* O,
+    int batch, int num_q_heads, int num_kv_heads,
+    int q_len, int kv_len, int head_dim,
+    float scale, int causal, void* stream
+) {
+    int q_tiles = (q_len + BR - 1) / BR;
+    dim3 grid(q_tiles, batch * num_q_heads);
+    int block = THREADS_PER_BLOCK;
+
+    // Shared memory: smem_q[BR * head_dim] + smem_kv[BC * head_dim], all BF16
+    int smem_bytes = (BR + BC) * head_dim * (int)sizeof(__nv_bfloat16);
+
+    flash_attention_bf16_kernel<<<grid, block, smem_bytes, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)Q,
+        (const __nv_bfloat16*)K,
+        (const __nv_bfloat16*)V,
+        (__nv_bfloat16*)O,
+        num_q_heads, num_kv_heads,
+        q_len, kv_len, head_dim,
+        scale, causal
+    );
+}
+
+void launch_decode_attention_bf16(
+    const void* Q, const void* K, const void* V, void* O,
+    int batch, int num_q_heads, int num_kv_heads,
+    int kv_len, int head_dim,
+    float scale, int causal, void* stream
+) {
+    int grid = batch * num_q_heads;
+    int block = DECODE_THREADS;
+
+    decode_attention_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)Q,
+        (const __nv_bfloat16*)K,
+        (const __nv_bfloat16*)V,
+        (__nv_bfloat16*)O,
+        num_q_heads, num_kv_heads,
+        kv_len, head_dim,
+        scale
+    );
+}
+
+}

csrc/common.cuh

@@ -0,0 +1,50 @@
+#pragma once
+#include <cuda_bf16.h>
+
+// --- Warp-level reductions (no shared memory needed) ---
+
+__device__ __forceinline__ float warp_reduce_sum(float val) {
+    #pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val += __shfl_down_sync(0xffffffff, val, offset);
+    return val;
+}
+
+__device__ __forceinline__ float warp_reduce_max(float val) {
+    #pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
+    return val;
+}
+
+// --- Block-level reductions ---
+
+__device__ __forceinline__ float block_reduce_sum(float val) {
+    __shared__ float shared[32];
+    int lane = threadIdx.x & 31;
+    int warp_id = threadIdx.x >> 5;
+    int num_warps = (blockDim.x + 31) >> 5;
+
+    val = warp_reduce_sum(val);
+    if (lane == 0) shared[warp_id] = val;
+    __syncthreads();
+
+    val = (threadIdx.x < num_warps) ? shared[threadIdx.x] : 0.0f;
+    if (warp_id == 0) val = warp_reduce_sum(val);
+    return val;
+}
+
+__device__ __forceinline__ float block_reduce_max(float val) {
+    __shared__ float shared[32];
+    int lane = threadIdx.x & 31;
+    int warp_id = threadIdx.x >> 5;
+    int num_warps = (blockDim.x + 31) >> 5;
+
+    val = warp_reduce_max(val);
+    if (lane == 0) shared[warp_id] = val;
+    __syncthreads();
+
+    val = (threadIdx.x < num_warps) ? shared[threadIdx.x] : -INFINITY;
+    if (warp_id == 0) val = warp_reduce_max(val);
+    return val;
+}

csrc/embedding/embedding.cu

@@ -0,0 +1,55 @@
+#include <cuda_bf16.h>
+
+// Embedding lookup: out[seq_idx] = table[token_ids[seq_idx]]
+// Grid: num_tokens, Block: handles hidden_size elements per token.
+
+__global__ void embedding_f32(
+    const float* __restrict__ table,    // [vocab_size, hidden_size]
+    const int* __restrict__ token_ids,  // [num_tokens]
+    float* __restrict__ out,            // [num_tokens, hidden_size]
+    int hidden_size
+) {
+    int token_idx = blockIdx.x;
+    int tid = token_ids[token_idx];
+    const float* row = table + tid * hidden_size;
+    float* dst = out + token_idx * hidden_size;
+
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        dst[i] = row[i];
+    }
+}
+
+__global__ void embedding_bf16(
+    const __nv_bfloat16* __restrict__ table,
+    const int* __restrict__ token_ids,
+    __nv_bfloat16* __restrict__ out,
+    int hidden_size
+) {
+    int token_idx = blockIdx.x;
+    int tid = token_ids[token_idx];
+    const __nv_bfloat16* row = table + tid * hidden_size;
+    __nv_bfloat16* dst = out + token_idx * hidden_size;
+
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        dst[i] = row[i];
+    }
+}
+
+extern "C" {
+
+void launch_embedding_f32(const void* table, const void* token_ids, void* out,
+                          int num_tokens, int hidden_size, void* stream) {
+    int block = (hidden_size < 256) ? hidden_size : 256;
+    embedding_f32<<<num_tokens, block, 0, (cudaStream_t)stream>>>(
+        (const float*)table, (const int*)token_ids, (float*)out, hidden_size);
+}
+
+void launch_embedding_bf16(const void* table, const void* token_ids, void* out,
+                           int num_tokens, int hidden_size, void* stream) {
+    int block = (hidden_size < 256) ? hidden_size : 256;
+    embedding_bf16<<<num_tokens, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)table, (const int*)token_ids,
+        (__nv_bfloat16*)out, hidden_size);
+}
+
+}

csrc/embedding/rope.cu

@@ -0,0 +1,116 @@
+#include <cuda_bf16.h>
+#include <math.h>
+
+// RoPE: Rotary Position Embedding
+// For each pair (x[2i], x[2i+1]) at position `pos`:
+//   y[2i]   = x[2i] * cos - x[2i+1] * sin
+//   y[2i+1] = x[2i] * sin + x[2i+1] * cos
+// where cos/sin come from precomputed cos_cache/sin_cache.
+//
+// cos_cache[pos][i] = cos(pos * freq[i])
+// sin_cache[pos][i] = sin(pos * freq[i])
+// freq[i] = 1.0 / (theta ^ (2i / head_dim))
+
+// Apply RoPE in-place to Q or K tensor.
+// x shape: [num_tokens, num_heads, head_dim]
+// cos_cache, sin_cache shape: [max_seq_len, head_dim/2]
+// positions: [num_tokens] — the position index for each token
+
+__global__ void rope_f32(
+    float* __restrict__ x,                    // [num_tokens, num_heads, head_dim]
+    const float* __restrict__ cos_cache,      // [max_seq_len, half_dim]
+    const float* __restrict__ sin_cache,      // [max_seq_len, half_dim]
+    const int* __restrict__ positions,        // [num_tokens]
+    int num_heads, int head_dim
+) {
+    int token_idx = blockIdx.x;
+    int head_idx = blockIdx.y;
+    int half_dim = head_dim / 2;
+    int pair_idx = threadIdx.x;  // which pair (0..half_dim)
+
+    if (pair_idx >= half_dim) return;
+
+    int pos = positions[token_idx];
+    float cos_val = cos_cache[pos * half_dim + pair_idx];
+    float sin_val = sin_cache[pos * half_dim + pair_idx];
+
+    int base = (token_idx * num_heads + head_idx) * head_dim;
+    float x0 = x[base + 2 * pair_idx];
+    float x1 = x[base + 2 * pair_idx + 1];
+
+    x[base + 2 * pair_idx]     = x0 * cos_val - x1 * sin_val;
+    x[base + 2 * pair_idx + 1] = x0 * sin_val + x1 * cos_val;
+}
+
+__global__ void rope_bf16(
+    __nv_bfloat16* __restrict__ x,
+    const float* __restrict__ cos_cache,
+    const float* __restrict__ sin_cache,
+    const int* __restrict__ positions,
+    int num_heads, int head_dim
+) {
+    int token_idx = blockIdx.x;
+    int head_idx = blockIdx.y;
+    int half_dim = head_dim / 2;
+    int pair_idx = threadIdx.x;
+
+    if (pair_idx >= half_dim) return;
+
+    int pos = positions[token_idx];
+    float cos_val = cos_cache[pos * half_dim + pair_idx];
+    float sin_val = sin_cache[pos * half_dim + pair_idx];
+
+    int base = (token_idx * num_heads + head_idx) * head_dim;
+    float x0 = __bfloat162float(x[base + 2 * pair_idx]);
+    float x1 = __bfloat162float(x[base + 2 * pair_idx + 1]);
+
+    x[base + 2 * pair_idx]     = __float2bfloat16(x0 * cos_val - x1 * sin_val);
+    x[base + 2 * pair_idx + 1] = __float2bfloat16(x0 * sin_val + x1 * cos_val);
+}
+
+// Precompute cos/sin cache on GPU
+__global__ void compute_rope_cache(
+    float* __restrict__ cos_cache,  // [max_seq_len, half_dim]
+    float* __restrict__ sin_cache,
+    int max_seq_len, int half_dim, float theta
+) {
+    int pos = blockIdx.x;
+    int i = threadIdx.x;
+    if (i >= half_dim) return;
+
+    float freq = 1.0f / powf(theta, (float)(2 * i) / (float)(2 * half_dim));
+    float angle = (float)pos * freq;
+    cos_cache[pos * half_dim + i] = cosf(angle);
+    sin_cache[pos * half_dim + i] = sinf(angle);
+}
+
+extern "C" {
+
+void launch_rope_f32(void* x, const void* cos_cache, const void* sin_cache,
+                     const void* positions, int num_tokens, int num_heads,
+                     int head_dim, void* stream) {
+    dim3 grid(num_tokens, num_heads);
+    int block = head_dim / 2;
+    rope_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (float*)x, (const float*)cos_cache, (const float*)sin_cache,
+        (const int*)positions, num_heads, head_dim);
+}
+
+void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache,
+                      const void* positions, int num_tokens, int num_heads,
+                      int head_dim, void* stream) {
+    dim3 grid(num_tokens, num_heads);
+    int block = head_dim / 2;
+    rope_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (__nv_bfloat16*)x, (const float*)cos_cache, (const float*)sin_cache,
+        (const int*)positions, num_heads, head_dim);
+}
+
+void launch_compute_rope_cache(void* cos_cache, void* sin_cache,
+                               int max_seq_len, int half_dim, float theta,
+                               void* stream) {
+    compute_rope_cache<<<max_seq_len, half_dim, 0, (cudaStream_t)stream>>>(
+        (float*)cos_cache, (float*)sin_cache, max_seq_len, half_dim, theta);
+}
+
+}

csrc/embedding/transpose.cu

@@ -0,0 +1,234 @@
+#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];
+}
+
+// ---- Generic strided copy (up to 4D) ----
+// Each thread copies one element. Maps flat contiguous output index to strided input index.
+// Unused dimensions are padded with shape=1, stride=0.
+
+__global__ void strided_copy_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int numel,
+    int ndim,
+    int shape0, int shape1, int shape2, int shape3,
+    int in_stride0, int in_stride1, int in_stride2, int in_stride3,
+    int in_offset
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= numel) return;
+
+    // Decompose flat output index into multi-dim indices (rightmost = fastest)
+    int remaining = idx;
+    int i3 = remaining % shape3; remaining /= shape3;
+    int i2 = remaining % shape2; remaining /= shape2;
+    int i1 = remaining % shape1; remaining /= shape1;
+    int i0 = remaining;
+
+    int in_idx = in_offset + i0 * in_stride0 + i1 * in_stride1 + i2 * in_stride2 + i3 * in_stride3;
+    out[idx] = in[in_idx];
+}
+
+__global__ void strided_copy_f32(
+    const float* __restrict__ in,
+    float* __restrict__ out,
+    int numel,
+    int ndim,
+    int shape0, int shape1, int shape2, int shape3,
+    int in_stride0, int in_stride1, int in_stride2, int in_stride3,
+    int in_offset
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= numel) return;
+
+    int remaining = idx;
+    int i3 = remaining % shape3; remaining /= shape3;
+    int i2 = remaining % shape2; remaining /= shape2;
+    int i1 = remaining % shape1; remaining /= shape1;
+    int i0 = remaining;
+
+    int in_idx = in_offset + i0 * in_stride0 + i1 * in_stride1 + i2 * in_stride2 + i3 * in_stride3;
+    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);
+}
+
+void launch_strided_copy_bf16(const void* in, void* out, int numel, int ndim,
+                              int shape0, int shape1, int shape2, int shape3,
+                              int in_stride0, int in_stride1, int in_stride2, int in_stride3,
+                              int in_offset, void* stream) {
+    int block = 256;
+    int grid = (numel + block - 1) / block;
+    strided_copy_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, numel, ndim,
+        shape0, shape1, shape2, shape3,
+        in_stride0, in_stride1, in_stride2, in_stride3, in_offset);
+}
+
+void launch_strided_copy_f32(const void* in, void* out, int numel, int ndim,
+                             int shape0, int shape1, int shape2, int shape3,
+                             int in_stride0, int in_stride1, int in_stride2, int in_stride3,
+                             int in_offset, void* stream) {
+    int block = 256;
+    int grid = (numel + block - 1) / block;
+    strided_copy_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)in, (float*)out, numel, ndim,
+        shape0, shape1, shape2, shape3,
+        in_stride0, in_stride1, in_stride2, in_stride3, in_offset);
+}
+
+}

csrc/gemm/gemv.cu

@@ -0,0 +1,102 @@
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+// Custom GEMV kernel for M=1 decode step (BF16):
+//   y[n] = sum_k x[k] * W[k * N + n]
+// where x: [K] (BF16), W: [K, N] (BF16, row-major), y: [N] (BF16).
+//
+// Design: K-split for high occupancy on large GPU (170 SMs).
+// Grid: (N / TILE_N, K / TILE_K) — each block computes a partial sum
+// for TILE_N output columns over a TILE_K slice of K.
+// Partial results are atomicAdd'd to an FP32 accumulator, then a
+// second kernel converts FP32 -> BF16.
+//
+// Memory access: adjacent threads read adjacent columns of the same row
+// of W, giving perfectly coalesced 128-byte transactions.
+
+#define GEMV_TILE_N 128
+#define GEMV_TILE_K 256
+#define GEMV_BLOCK  128  // = TILE_N, one thread per output column
+
+__global__ void gemv_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,      // [K]
+    const __nv_bfloat16* __restrict__ W,      // [K, N] row-major
+    float* __restrict__ y_fp32,               // [N] accumulator
+    int K, int N
+) {
+    const int block_n = blockIdx.x;
+    const int block_k = blockIdx.y;
+    const int t = threadIdx.x;
+    const int col = block_n * GEMV_TILE_N + t;
+
+    if (col >= N) return;
+
+    const int k_start = block_k * GEMV_TILE_K;
+    const int k_end = min(k_start + GEMV_TILE_K, K);
+    const int k_len = k_end - k_start;
+
+    // Load x[k_start..k_end] into shared memory as FP32
+    __shared__ float x_shared[GEMV_TILE_K];
+    for (int i = t; i < k_len; i += GEMV_BLOCK) {
+        x_shared[i] = __bfloat162float(x[k_start + i]);
+    }
+    __syncthreads();
+
+    // Compute partial dot product for this column
+    float sum = 0.0f;
+    for (int ki = 0; ki < k_len; ki++) {
+        sum += x_shared[ki] * __bfloat162float(W[(k_start + ki) * N + col]);
+    }
+
+    // Atomic accumulate (handles K-split reduction)
+    atomicAdd(&y_fp32[col], sum);
+}
+
+// Conversion kernel: FP32 accumulator -> BF16 output
+__global__ void gemv_fp32_to_bf16_kernel(
+    const float* __restrict__ src,
+    __nv_bfloat16* __restrict__ dst,
+    int n
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        dst[idx] = __float2bfloat16(src[idx]);
+    }
+}
+
+extern "C" {
+
+void launch_gemv_bf16(
+    const void* x,           // [K] BF16
+    const void* W,           // [K, N] BF16 row-major
+    void* y_bf16,            // [N] BF16 output
+    void* y_fp32_buf,        // [N] FP32 temporary (caller-provided)
+    int K, int N,
+    void* stream
+) {
+    cudaStream_t s = (cudaStream_t)stream;
+
+    // Zero the FP32 accumulator
+    cudaMemsetAsync((float*)y_fp32_buf, 0, N * sizeof(float), s);
+
+    // Launch GEMV kernel
+    dim3 grid((N + GEMV_TILE_N - 1) / GEMV_TILE_N,
+              (K + GEMV_TILE_K - 1) / GEMV_TILE_K);
+    gemv_bf16_kernel<<<grid, GEMV_BLOCK, 0, s>>>(
+        (const __nv_bfloat16*)x,
+        (const __nv_bfloat16*)W,
+        (float*)y_fp32_buf,
+        K, N
+    );
+
+    // Convert FP32 -> BF16
+    int conv_block = 256;
+    int conv_grid = (N + conv_block - 1) / conv_block;
+    gemv_fp32_to_bf16_kernel<<<conv_grid, conv_block, 0, s>>>(
+        (const float*)y_fp32_buf,
+        (__nv_bfloat16*)y_bf16,
+        N
+    );
+}
+
+} // extern "C"

csrc/gemm/naive.cu

@@ -0,0 +1,62 @@
+#include <cuda_bf16.h>
+
+// Naive GEMM: each thread computes one element of C.
+// C[i][j] = sum_k A[i][k] * B[k][j]
+// All matrices are row-major.
+__global__ void gemm_naive_bf16(
+    const __nv_bfloat16* A, const __nv_bfloat16* B, __nv_bfloat16* C,
+    int M, int N, int K
+) {
+    int row = blockIdx.y * blockDim.y + threadIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (row < M && col < N) {
+        float sum = 0.0f;
+        for (int k = 0; k < K; k++) {
+            sum += __bfloat162float(A[row * K + k]) * __bfloat162float(B[k * N + col]);
+        }
+        C[row * N + col] = __float2bfloat16(sum);
+    }
+}
+
+__global__ void gemm_naive_f32(
+    const float* A, const float* B, float* C,
+    int M, int N, int K
+) {
+    int row = blockIdx.y * blockDim.y + threadIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (row < M && col < N) {
+        float sum = 0.0f;
+        for (int k = 0; k < K; k++) {
+            sum += A[row * K + k] * B[k * N + col];
+        }
+        C[row * N + col] = sum;
+    }
+}
+
+extern "C" {
+
+void launch_gemm_naive_bf16(
+    const void* A, const void* B, void* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(16, 16);
+    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
+    gemm_naive_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)A, (const __nv_bfloat16*)B, (__nv_bfloat16*)C, M, N, K
+    );
+}
+
+void launch_gemm_naive_f32(
+    const void* A, const void* B, void* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(16, 16);
+    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
+    gemm_naive_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)A, (const float*)B, (float*)C, M, N, K
+    );
+}
+
+} // extern "C"

csrc/gemm/tiled.cu

@@ -0,0 +1,116 @@
+#include <cuda_bf16.h>
+
+// Tiled GEMM using shared memory.
+// Each thread block loads TILE_SIZE x TILE_SIZE tiles of A and B
+// into shared memory, then computes a partial dot product.
+#define TILE_SIZE 32
+
+__global__ void gemm_tiled_f32(
+    const float* A, const float* B, float* C,
+    int M, int N, int K
+) {
+    __shared__ float As[TILE_SIZE][TILE_SIZE];
+    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
+
+    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
+    int col = blockIdx.x * TILE_SIZE + threadIdx.x;
+
+    float sum = 0.0f;
+
+    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
+        // Load tile of A
+        int a_col = t * TILE_SIZE + threadIdx.x;
+        if (row < M && a_col < K) {
+            As[threadIdx.y][threadIdx.x] = A[row * K + a_col];
+        } else {
+            As[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        // Load tile of B
+        int b_row = t * TILE_SIZE + threadIdx.y;
+        if (b_row < K && col < N) {
+            Bs[threadIdx.y][threadIdx.x] = B[b_row * N + col];
+        } else {
+            Bs[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        __syncthreads();
+
+        for (int k = 0; k < TILE_SIZE; k++) {
+            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
+        }
+
+        __syncthreads();
+    }
+
+    if (row < M && col < N) {
+        C[row * N + col] = sum;
+    }
+}
+
+__global__ void gemm_tiled_bf16(
+    const __nv_bfloat16* A, const __nv_bfloat16* B, __nv_bfloat16* C,
+    int M, int N, int K
+) {
+    __shared__ float As[TILE_SIZE][TILE_SIZE];
+    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
+
+    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
+    int col = blockIdx.x * TILE_SIZE + threadIdx.x;
+
+    float sum = 0.0f;
+
+    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
+        int a_col = t * TILE_SIZE + threadIdx.x;
+        if (row < M && a_col < K) {
+            As[threadIdx.y][threadIdx.x] = __bfloat162float(A[row * K + a_col]);
+        } else {
+            As[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        int b_row = t * TILE_SIZE + threadIdx.y;
+        if (b_row < K && col < N) {
+            Bs[threadIdx.y][threadIdx.x] = __bfloat162float(B[b_row * N + col]);
+        } else {
+            Bs[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        __syncthreads();
+
+        for (int k = 0; k < TILE_SIZE; k++) {
+            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
+        }
+
+        __syncthreads();
+    }
+
+    if (row < M && col < N) {
+        C[row * N + col] = __float2bfloat16(sum);
+    }
+}
+
+extern "C" {
+
+void launch_gemm_tiled_f32(
+    const void* A, const void* B, void* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(TILE_SIZE, TILE_SIZE);
+    dim3 grid((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);
+    gemm_tiled_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)A, (const float*)B, (float*)C, M, N, K
+    );
+}
+
+void launch_gemm_tiled_bf16(
+    const void* A, const void* B, void* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(TILE_SIZE, TILE_SIZE);
+    dim3 grid((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);
+    gemm_tiled_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)A, (const __nv_bfloat16*)B, (__nv_bfloat16*)C, M, N, K
+    );
+}
+
+} // extern "C"

csrc/normalization/layernorm.cu

@@ -0,0 +1,119 @@
+#include "../common.cuh"
+
+// LayerNorm: y[i] = gamma[i] * (x[i] - mean) / sqrt(var + eps) + beta[i]
+// Each block processes one row of shape [hidden_size].
+
+__global__ void layernorm_f32(
+    const float* __restrict__ x,
+    const float* __restrict__ gamma,
+    const float* __restrict__ beta,
+    float* __restrict__ out,
+    int hidden_size, float eps
+) {
+    int row = blockIdx.x;
+    const float* x_row = x + row * hidden_size;
+    float* out_row = out + row * hidden_size;
+
+    // Pass 1: compute mean
+    float local_sum = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        local_sum += x_row[i];
+    }
+    local_sum = block_reduce_sum(local_sum);
+
+    __shared__ float s_mean, s_inv_std;
+    if (threadIdx.x == 0) {
+        s_mean = local_sum / hidden_size;
+    }
+    __syncthreads();
+
+    float mean = s_mean;
+
+    // Pass 2: compute variance = sum((x - mean)^2) / N
+    float local_var = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float d = x_row[i] - mean;
+        local_var += d * d;
+    }
+    local_var = block_reduce_sum(local_var);
+
+    if (threadIdx.x == 0) {
+        s_inv_std = rsqrtf(local_var / hidden_size + eps);
+    }
+    __syncthreads();
+
+    float inv_std = s_inv_std;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        out_row[i] = gamma[i] * (x_row[i] - mean) * inv_std + beta[i];
+    }
+}
+
+__global__ void layernorm_bf16(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ gamma,
+    const __nv_bfloat16* __restrict__ beta,
+    __nv_bfloat16* __restrict__ out,
+    int hidden_size, float eps
+) {
+    int row = blockIdx.x;
+    const __nv_bfloat16* x_row = x + row * hidden_size;
+    __nv_bfloat16* out_row = out + row * hidden_size;
+
+    // Pass 1: compute mean
+    float local_sum = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        local_sum += __bfloat162float(x_row[i]);
+    }
+    local_sum = block_reduce_sum(local_sum);
+
+    __shared__ float s_mean, s_inv_std;
+    if (threadIdx.x == 0) {
+        s_mean = local_sum / hidden_size;
+    }
+    __syncthreads();
+
+    float mean = s_mean;
+
+    // Pass 2: compute variance = sum((x - mean)^2) / N
+    float local_var = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float d = __bfloat162float(x_row[i]) - mean;
+        local_var += d * d;
+    }
+    local_var = block_reduce_sum(local_var);
+
+    if (threadIdx.x == 0) {
+        s_inv_std = rsqrtf(local_var / hidden_size + eps);
+    }
+    __syncthreads();
+
+    float inv_std = s_inv_std;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float v = __bfloat162float(x_row[i]);
+        float g = __bfloat162float(gamma[i]);
+        float b = __bfloat162float(beta[i]);
+        out_row[i] = __float2bfloat16(g * (v - mean) * inv_std + b);
+    }
+}
+
+extern "C" {
+
+void launch_layernorm_f32(const void* x, const void* gamma, const void* beta,
+                          void* out, int rows, int hidden_size, float eps, void* stream) {
+    int block = (hidden_size < 1024) ? hidden_size : 1024;
+    if (block < 32) block = 32;
+    layernorm_f32<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const float*)x, (const float*)gamma, (const float*)beta,
+        (float*)out, hidden_size, eps);
+}
+
+void launch_layernorm_bf16(const void* x, const void* gamma, const void* beta,
+                           void* out, int rows, int hidden_size, float eps, void* stream) {
+    int block = (hidden_size < 1024) ? hidden_size : 1024;
+    if (block < 32) block = 32;
+    layernorm_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)gamma, (const __nv_bfloat16*)beta,
+        (__nv_bfloat16*)out, hidden_size, eps);
+}
+
+}

csrc/normalization/rmsnorm.cu

@@ -0,0 +1,137 @@
+#include "../common.cuh"
+
+// RMSNorm: y[i] = x[i] * rsqrt(mean(x²) + eps) * gamma[i]
+// Each block processes one row of shape [hidden_size].
+
+__global__ void rmsnorm_f32(
+    const float* __restrict__ x,
+    const float* __restrict__ gamma,
+    float* __restrict__ out,
+    int hidden_size, float eps
+) {
+    int row = blockIdx.x;
+    const float* x_row = x + row * hidden_size;
+    float* out_row = out + row * hidden_size;
+
+    float sum_sq = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float v = x_row[i];
+        sum_sq += v * v;
+    }
+    sum_sq = block_reduce_sum(sum_sq);
+
+    __shared__ float s_rms_inv;
+    if (threadIdx.x == 0) {
+        s_rms_inv = rsqrtf(sum_sq / hidden_size + eps);
+    }
+    __syncthreads();
+
+    float rms_inv = s_rms_inv;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        out_row[i] = x_row[i] * rms_inv * gamma[i];
+    }
+}
+
+__global__ void rmsnorm_bf16(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ gamma,
+    __nv_bfloat16* __restrict__ out,
+    int hidden_size, float eps
+) {
+    int row = blockIdx.x;
+    const __nv_bfloat16* x_row = x + row * hidden_size;
+    __nv_bfloat16* out_row = out + row * hidden_size;
+
+    float sum_sq = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float v = __bfloat162float(x_row[i]);
+        sum_sq += v * v;
+    }
+    sum_sq = block_reduce_sum(sum_sq);
+
+    __shared__ float s_rms_inv;
+    if (threadIdx.x == 0) {
+        s_rms_inv = rsqrtf(sum_sq / hidden_size + eps);
+    }
+    __syncthreads();
+
+    float rms_inv = s_rms_inv;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float v = __bfloat162float(x_row[i]);
+        float g = __bfloat162float(gamma[i]);
+        out_row[i] = __float2bfloat16(v * rms_inv * g);
+    }
+}
+
+// Fused Add + RMSNorm: sum_out = x + residual, normed_out = rmsnorm(sum_out, gamma, eps)
+// Each block handles one row of [hidden_size].
+__global__ void add_rmsnorm_bf16(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ residual,
+    const __nv_bfloat16* __restrict__ gamma,
+    __nv_bfloat16* __restrict__ normed_out,
+    __nv_bfloat16* __restrict__ sum_out,
+    int hidden_size, float eps
+) {
+    int row = blockIdx.x;
+    const __nv_bfloat16* x_row = x + row * hidden_size;
+    const __nv_bfloat16* res_row = residual + row * hidden_size;
+    __nv_bfloat16* sum_row = sum_out + row * hidden_size;
+    __nv_bfloat16* norm_row = normed_out + row * hidden_size;
+
+    // Pass 1: compute sum = x + residual, and accumulate sum_sq
+    float sum_sq = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float s = __bfloat162float(x_row[i]) + __bfloat162float(res_row[i]);
+        sum_row[i] = __float2bfloat16(s);
+        sum_sq += s * s;
+    }
+    sum_sq = block_reduce_sum(sum_sq);
+
+    __shared__ float s_rms_inv;
+    if (threadIdx.x == 0) {
+        s_rms_inv = rsqrtf(sum_sq / hidden_size + eps);
+    }
+    __syncthreads();
+
+    // Pass 2: normed_out = sum * rms_inv * gamma
+    float rms_inv = s_rms_inv;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float s = __bfloat162float(sum_row[i]);
+        float g = __bfloat162float(gamma[i]);
+        norm_row[i] = __float2bfloat16(s * rms_inv * g);
+    }
+}
+
+extern "C" {
+
+void launch_rmsnorm_f32(const void* x, const void* gamma, void* out,
+                        int rows, int hidden_size, float eps, void* stream) {
+    int block = (hidden_size < 1024) ? hidden_size : 1024;
+    if (block < 32) block = 32;
+    rmsnorm_f32<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const float*)x, (const float*)gamma, (float*)out, hidden_size, eps);
+}
+
+void launch_rmsnorm_bf16(const void* x, const void* gamma, void* out,
+                         int rows, int hidden_size, float eps, void* stream) {
+    int block = (hidden_size < 1024) ? hidden_size : 1024;
+    if (block < 32) block = 32;
+    rmsnorm_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)gamma,
+        (__nv_bfloat16*)out, hidden_size, eps);
+}
+
+void launch_add_rmsnorm_bf16(const void* x, const void* residual, const void* gamma,
+                             void* normed_out, void* sum_out,
+                             int rows, int hidden_size, float eps, void* stream) {
+    int block = (hidden_size < 1024) ? hidden_size : 1024;
+    if (block < 32) block = 32;
+    add_rmsnorm_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)residual,
+        (const __nv_bfloat16*)gamma,
+        (__nv_bfloat16*)normed_out, (__nv_bfloat16*)sum_out,
+        hidden_size, eps);
+}
+
+}

csrc/reduce/softmax.cu

@@ -0,0 +1,106 @@
+#include "../common.cuh"
+
+// Safe softmax along the last dimension.
+// Each block handles one row of length `cols`.
+// Three-pass: 1) find max, 2) exp + sum, 3) normalize.
+
+__global__ void softmax_f32(
+    const float* __restrict__ x,
+    float* __restrict__ out,
+    int cols
+) {
+    int row = blockIdx.x;
+    const float* x_row = x + row * cols;
+    float* out_row = out + row * cols;
+
+    // Pass 1: find max
+    float local_max = -INFINITY;
+    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
+        local_max = fmaxf(local_max, x_row[i]);
+    }
+    float row_max = block_reduce_max(local_max);
+
+    __shared__ float s_max;
+    if (threadIdx.x == 0) s_max = row_max;
+    __syncthreads();
+    row_max = s_max;
+
+    // Pass 2: exp and sum
+    float local_sum = 0.0f;
+    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
+        float e = expf(x_row[i] - row_max);
+        out_row[i] = e;
+        local_sum += e;
+    }
+    float row_sum = block_reduce_sum(local_sum);
+
+    __shared__ float s_inv_sum;
+    if (threadIdx.x == 0) s_inv_sum = 1.0f / row_sum;
+    __syncthreads();
+    float inv_sum = s_inv_sum;
+
+    // Pass 3: normalize
+    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
+        out_row[i] *= inv_sum;
+    }
+}
+
+__global__ void softmax_bf16(
+    const __nv_bfloat16* __restrict__ x,
+    __nv_bfloat16* __restrict__ out,
+    int cols
+) {
+    int row = blockIdx.x;
+    const __nv_bfloat16* x_row = x + row * cols;
+    __nv_bfloat16* out_row = out + row * cols;
+
+    float local_max = -INFINITY;
+    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
+        local_max = fmaxf(local_max, __bfloat162float(x_row[i]));
+    }
+    float row_max = block_reduce_max(local_max);
+
+    __shared__ float s_max;
+    if (threadIdx.x == 0) s_max = row_max;
+    __syncthreads();
+    row_max = s_max;
+
+    // We need float scratch for exp values. Reuse out (write bf16 in pass 3).
+    // Use registers to hold exp values during sum pass instead.
+    float local_sum = 0.0f;
+    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
+        float e = expf(__bfloat162float(x_row[i]) - row_max);
+        // Temporarily store exp in output as bf16 (slight precision loss, acceptable)
+        out_row[i] = __float2bfloat16(e);
+        local_sum += e;
+    }
+    float row_sum = block_reduce_sum(local_sum);
+
+    __shared__ float s_inv_sum;
+    if (threadIdx.x == 0) s_inv_sum = 1.0f / row_sum;
+    __syncthreads();
+    float inv_sum = s_inv_sum;
+
+    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
+        float e = __bfloat162float(out_row[i]);
+        out_row[i] = __float2bfloat16(e * inv_sum);
+    }
+}
+
+extern "C" {
+
+void launch_softmax_f32(const void* x, void* out, int rows, int cols, void* stream) {
+    int block = (cols < 1024) ? cols : 1024;
+    if (block < 32) block = 32;
+    softmax_f32<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const float*)x, (float*)out, cols);
+}
+
+void launch_softmax_bf16(const void* x, void* out, int rows, int cols, void* stream) {
+    int block = (cols < 1024) ? cols : 1024;
+    if (block < 32) block = 32;
+    softmax_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, cols);
+}
+
+}

docs/00-roadmap.md

@@ -9,7 +9,7 @@
 | 抽象层级 | Level 0.5 | 自写 CUDA kernel + cuBLAS 可切换，便于 benchmark 对比 |
 | 硬件 | 8×RTX 5090 (Blackwell, CC 12.0, 32GB GDDR7) | 纯 PCIe Gen5 x16 互联，无 NVLink (详见下方硬件拓扑) |
 | 语言 | Rust + CUDA (C/C++) | Rust FFI 调用 CUDA |
-| 起步模型 | GPT-2 124M → Qwen3-7B | 从简单到实用 |
+| 起步模型 | GPT-2 124M → Qwen3-8B | 从简单到实用 |
 | 精度 | BF16/FP16 | 后期扩展 FP8 |
 | Tensor | 自己实现 | 完整学习 tensor 抽象设计 |
 | Tokenizer | 自己实现 BPE | 学习分词机制 |
@@ -101,7 +101,7 @@ Phase 8: GPT-2 完整推理 ◄──────────── 里程碑
     │
 Phase 9: KV Cache + Autoregressive Generation
     │
-Phase 10: Qwen3-7B 支持 ◄─────────── 里程碑 ② 7B 模型推理
+Phase 10: Qwen3-8B 支持 ◄─────────── 里程碑 ② 8B 模型推理
     │
 Phase 11: Paged Attention + KV Cache Manager
     │
@@ -109,7 +109,7 @@ Phase 12: Continuous Batching + Request Scheduler
     │
 Phase 13: HTTP API + SSE Streaming ◄── 里程碑 ③ 端到端 API 可用
     │
-Phase 14: Flash Attention v2
+Phase 14: Flash Attention (FA2 for SM120)
     │
 Phase 15: 性能优化 ◄──────────────── 里程碑 ④ 50% vLLM throughput
     │
@@ -625,8 +625,8 @@ safetensors file (disk)
 
 - [ ] 加载 GPT-2 124M (`openai-community/gpt2`)，打印所有 tensor name, shape, dtype
 - [ ] 抽查几个 tensor 的前 10 个值，与 PyTorch `from_pretrained` 对比
-- [ ] 加载 Qwen3-7B sharded 权重，验证所有 tensor 都成功加载
-- [ ] 性能: 测量 7B 模型权重加载时间 (mmap → GPU 全流程)
+- [ ] 加载 Qwen3-8B sharded 权重，验证所有 tensor 都成功加载
+- [ ] 性能: 测量 8B 模型权重加载时间 (mmap → GPU 全流程)
 - [ ] 错误处理: 缺少 tensor、dtype 不匹配、文件不存在等情况
 
 ---
@@ -869,15 +869,15 @@ weights × V_cache [B, H, S, D] → output [B, H, 1, D]
 
 ---
 
-## Phase 10: Qwen3-7B 支持 — 里程碑 ②
+## Phase 10: Qwen3-8B 支持 — 里程碑 ②
 
 **Crate**: `xserv-model`
 
-**目标**: 扩展模型定义以支持 Qwen3-7B，验证输出正确性。
+**目标**: 扩展模型定义以支持 Qwen3-8B，验证输出正确性。
 
 ### 架构对比
 
-| 特性 | GPT-2 (124M) | Qwen3-7B |
+| 特性 | GPT-2 (124M) | Qwen3-8B |
 |------|-------------|----------|
 | Normalization | LayerNorm (pre-LN) | RMSNorm (pre-LN) |
 | Position Encoding | Learned absolute (wpe) | RoPE (无单独参数) |
@@ -885,8 +885,8 @@ weights × V_cache [B, H, S, D] → output [B, H, 1, D]
 | Activation | GELU | SwiGLU (SiLU gate) |
 | FFN | Linear(H→4H) → GELU → Linear(4H→H) | gate_proj + up_proj → SiLU gate → down_proj |
 | Vocab Size | 50,257 | ~152,000 |
-| Hidden Size | 768 | 3,584 (7B) |
-| Layers | 12 | 28 |
+| Hidden Size | 768 | 4,096 (8B) |
+| Layers | 12 | 36 |
 | Tied Embeddings | Yes | No |
 
 ### 需要新增/修改的组件
@@ -948,16 +948,16 @@ pub struct Qwen3DecoderLayer {
 ### 显存预算 (BF16, 单卡 5090 32GB)
 
 ```
-模型权重:  7B × 2B = ~14 GB
-KV cache:  28 layers × 2(KV) × 8 heads × 4096 tokens × 128 dim × 2B ≈ 4.5 GB
+模型权重:  8B × 2B = ~16 GB
+KV cache:  36 layers × 2(KV) × 8 heads × 4096 tokens × 128 dim × 2B ≈ 5.6 GB
 Activation (单请求): ~1 GB
 ────────────────────────
-总计: ~19.5 GB (单请求)，剩余 ~12 GB 可用于更多并发
+总计: ~22.6 GB (单请求)，剩余 ~10 GB 可用于更多并发
 ```
 
 ### 测试验收
 
-- [ ] 加载 Qwen3-7B 权重到单张 5090，打印模型结构和参数量
+- [ ] 加载 Qwen3-8B 权重到单张 5090，打印模型结构和参数量
 - [ ] Prefill logits 与 HF transformers 对比: 输入 "你好" → top-5 logits 一致
 - [ ] 英文生成: "What is the capital of France?" → 生成合理回答
 - [ ] 中文生成: "请介绍一下量子计算" → 生成通顺中文
@@ -1196,7 +1196,7 @@ GET  /health                 # 健康检查
 **Chat Completion Request**:
 ```json
 {
-  "model": "qwen3-7b",
+  "model": "qwen3-8b",
   "messages": [
     {"role": "system", "content": "You are a helpful assistant."},
     {"role": "user", "content": "What is 1+1?"}
@@ -1211,13 +1211,13 @@ GET  /health                 # 健康检查
 
 **SSE Streaming Response**:
 ```
-data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-7b","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}
+data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-8b","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}
 
-data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-7b","choices":[{"index":0,"delta":{"content":"The"},"finish_reason":null}]}
+data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-8b","choices":[{"index":0,"delta":{"content":"The"},"finish_reason":null}]}
 
-data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-7b","choices":[{"index":0,"delta":{"content":" answer"},"finish_reason":null}]}
+data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-8b","choices":[{"index":0,"delta":{"content":" answer"},"finish_reason":null}]}
 
-data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-7b","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
+data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-8b","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
 
 data: [DONE]
 ```
@@ -1228,7 +1228,7 @@ data: [DONE]
   "id": "chatcmpl-xxx",
   "object": "chat.completion",
   "created": 1234567890,
-  "model": "qwen3-7b",
+  "model": "qwen3-8b",
   "choices": [{
     "index": 0,
     "message": {"role": "assistant", "content": "The answer is 2."},
@@ -1278,7 +1278,7 @@ Client (curl / Python OpenAI SDK)
   ```bash
   curl http://localhost:8080/v1/chat/completions \
     -H "Content-Type: application/json" \
-    -d '{"model":"qwen3-7b","messages":[{"role":"user","content":"Hello"}],"stream":true}'
+    -d '{"model":"qwen3-8b","messages":[{"role":"user","content":"Hello"}],"stream":true}'
   ```
   看到 SSE 逐 token 输出
 
@@ -1287,7 +1287,7 @@ Client (curl / Python OpenAI SDK)
   from openai import OpenAI
   client = OpenAI(base_url="http://localhost:8080/v1", api_key="unused")
   for chunk in client.chat.completions.create(
-      model="qwen3-7b",
+      model="qwen3-8b",
       messages=[{"role": "user", "content": "What is 1+1?"}],
       stream=True
   ):
@@ -1302,12 +1302,26 @@ Client (curl / Python OpenAI SDK)
 
 ---
 
-## Phase 14: Flash Attention v2
+## Phase 14: Flash Attention (FA2 for SM120)
 
 **Crate**: `xserv-kernels`
 **CUDA 源码**: `csrc/attention/flash_attention.cu`
 
-**目标**: 实现 Flash Attention v2 的 CUDA kernel，大幅降低 attention 的显存占用并提升速度。
+**目标**: 实现 Flash Attention 的 CUDA kernel，大幅降低 attention 的显存占用并提升速度。
+
+### 硬件适配说明
+
+Flash Attention 已发展到第 4 代 (FA4, arxiv 2603.05451)，但各版本有明确的硬件依赖:
+
+| 版本 | 目标架构 | 关键硬件特性 | RTX 5090 兼容 |
+|------|---------|------------|--------------|
+| FA2 | 通用 CUDA (SM75+) | 标准 shared memory + HMMA | **是** ✅ |
+| FA3 | Hopper SM90 (H100) | TMA + WGMMA + warp specialization | 否 |
+| FA4 | Blackwell SM100 (B200/B300) | TMEM + async MMA + 2-CTA mode | 否 |
+
+**RTX 5090 (SM120, CC 12.0) 使用的是消费级 Blackwell 架构 (GB202)，与数据中心 Blackwell (B200, SM100) 是不同的硅片设计。SM120 物理上没有 TMEM (Tensor Memory) 子系统，因此 FA4 的 kernel 无法在 5090 上运行。这不是软件限制，是硬件级差异。**
+
+因此本项目实现 **FA2 算法**，使用标准 CUDA (shared memory + HMMA)。FA2 的核心优化——online softmax tiling、O(1) 显存占用——在任何架构上都有效。
 
 ### 核心思想
 
@@ -1323,16 +1337,18 @@ Flash Attention 的解法:
 - 将 Q, K, V 分成 tiles，在 SRAM (shared memory) 中计算
 - 使用 **online softmax trick**: 边算边更新 running max 和 running sum
 
-### 算法 (Forward Pass)
+### 算法 (Forward Pass, FA2)
+
+FA2 相比 FA1 的改进: 外层循环遍历 Q tiles (而非 K/V)，减少 HBM 读写次数。
 
 ```
 Br, Bc = tile sizes for Q and K/V respectively
 
-for each Q tile (q_start..q_start+Br):
+for each Q tile (q_start..q_start+Br):                    ← 外层: Q tiles
     load Q_tile [Br, D] to shared memory
-    initialize: O_tile = 0, l = 0, m = -inf    // running sum and max
+    initialize: O_tile = 0, l = 0, m = -inf               // running sum and max
 
-    for each K,V tile (kv_start..kv_start+Bc):
+    for each K,V tile (kv_start..kv_start+Bc):             ← 内层: K/V tiles
         load K_tile [Bc, D], V_tile [Bc, D] to shared memory
 
         // Compute attention scores for this tile pair
@@ -1345,6 +1361,8 @@ for each Q tile (q_start..q_start+Br):
         m_new = max(m, rowmax(S_tile))          // new running max
         P_tile = exp(S_tile - m_new)            // safe exp
         l_new = exp(m - m_new) * l + rowsum(P_tile)  // update running sum
+
+        // Rescale and accumulate output
         O_tile = diag(exp(m - m_new)) * O_tile + P_tile @ V_tile
         m = m_new
         l = l_new
@@ -1356,9 +1374,12 @@ for each Q tile (q_start..q_start+Br):
 ### 实现要点
 
 1. **Tile 大小选择**:
-   - 受限于 shared memory (5090 Blackwell CC 12.0: 需要实测确认 per-SM shared memory 上限)
-   - 需要同时存 Q_tile, K_tile, V_tile, S_tile
-   - 典型值: Br=Bc=128 for D=128, BF16
+   - 5090 SM120: shared memory per SM = 100 KB (需实测确认)
+   - 需同时存 Q_tile, K_tile, V_tile, S_tile
+   - BF16: Q_tile [Br, D] = Br × 128 × 2B; K_tile [Bc, D] = Bc × 128 × 2B
+   - S_tile [Br, Bc] 保持 FP32 = Br × Bc × 4B
+   - 推荐起步: Br=Bc=64, head_dim=128 → 共需 ~100KB shared memory
+   - 优化版: Br=Bc=128 需要更多 shared memory, 可能需要拆分
 
 2. **Causal mask 优化**:
    - 如果 K/V tile 完全在 Q tile 的"未来"（kv_start > q_end）→ 跳过整个 tile
@@ -1369,10 +1390,14 @@ for each Q tile (q_start..q_start+Br):
    - Q, K, V 的加载用 BF16（节省 bandwidth）
    - 最终 O 转回 BF16 写出
 
-4. **与 Paged Attention 的结合**:
-   - Flash Attention 的 K/V tile 遍历逻辑需要适配间接寻址
-   - 每个 tile 查 block_table 得到物理地址
-   - 这是 "Flash-Decoding" / "FlashInfer" 的核心
+4. **GQA 支持**:
+   - K/V heads 数量 < Q heads 时，kernel 中做 `kv_head = q_head / num_groups` 索引
+   - 不需要 repeat_kv 操作，直接在 kernel 内部解决
+
+5. **Decode attention 特化**:
+   - Decode 时 Q 只有 1 行 (Br=1)，退化为 vector-matrix attention
+   - 可以写一个专门的 decode attention kernel (类似 FlashDecoding)
+   - 沿 KV sequence 维度做 parallel reduction
 
 ### 测试验收
 
@@ -1386,8 +1411,9 @@ for each Q tile (q_start..q_start+Br):
 | 8192 | OOM? | MB | OOM? | ms |
 | 32768 | OOM | MB | OOM | ms |
 
-- [ ] 集成到 Qwen3-7B，端到端 decode latency 对比
+- [ ] 集成到 Qwen3-8B，端到端 decode latency 对比
 - [ ] Profile: `ncu` 分析 compute utilization, memory throughput
+- [ ] GQA 支持: 无 repeat_kv 开销
 
 ---
 
@@ -1441,7 +1467,7 @@ ncu --target-processes all --set full ./target/release/xserv-server
 
 ### 测试验收
 
-- [ ] 安装 vLLM，同一台机器跑 Qwen3-7B
+- [ ] 安装 vLLM，同一台机器跑 Qwen3-8B
 - [ ] Benchmark 对比:
 
 | Metric | vLLM | xserv | Ratio |
@@ -1488,7 +1514,7 @@ ncu --target-processes all --set full ./target/release/xserv-server
 
 - **无损**: rejection sampling 保证输出分布与纯 target model 一致
 - **加速条件**: draft model 足够快且与 target 分布接近
-- **Draft model 选择**: Qwen3-0.5B / Qwen3-1.5B 作为 Qwen3-7B 的 draft
+- **Draft model 选择**: Qwen3-0.5B / Qwen3-1.5B 作为 Qwen3-8B 的 draft
 
 ### KV Cache 处理
 
@@ -1578,7 +1604,7 @@ Row Parallel: down_proj 按行切分
 
 ### 测试验收
 
-- [ ] TP=2: Qwen3-7B 输出与单卡 (TP=1) 完全一致
+- [ ] TP=2: Qwen3-8B 输出与单卡 (TP=1) 完全一致
 - [ ] TP=4: 每卡权重显存占用约 1/4
 - [ ] Scaling benchmark (同组 GPU 0-3):
 
@@ -1646,7 +1672,7 @@ tensor_fp8 = cast_to_fp8(tensor / scale)
 | FP8 E4M3 | X.XX | +0.XX |
 | INT8 weight-only | X.XX | +0.XX |
 
-- [ ] 显存: FP8 权重占用约 BF16 的一半 (~7 GB for 7B model)
+- [ ] 显存: FP8 权重占用约 BF16 的一半 (~8 GB for 8B model)
 - [ ] 性能: FP8 GEMM throughput vs BF16 GEMM
 
 ---
@@ -1727,7 +1753,7 @@ Text  → Tokenizer   → Text Tokens   ────────────→
 | 里程碑 | Phase | 验收标准 |
 |--------|-------|---------|
 | ① GPT-2 推理 | 8 | CLI 输入 prompt, GPT-2 生成连贯文本, logits 与 PyTorch 一致 |
-| ② Qwen3-7B 推理 | 10 | 7B 模型中英文对话, 多轮 chat template 正确 |
+| ② Qwen3-8B 推理 | 10 | 8B 模型中英文对话, 多轮 chat template 正确 |
 | ③ E2E API | 13 | HTTP streaming API, Python OpenAI SDK 可调用, 10 并发正确 |
 | ④ 性能达标 | 15 | throughput >= 50% vLLM, profiling 报告完成 |
 | ⑤ 多卡推理 | 17 | TP=2/4 同组 GPU 推理正确, scaling benchmark 完成 |

docs/01-cuda-ffi.md

@@ -72,9 +72,31 @@ Wraps cudaStream_t. RAII with Drop calling cudaStreamDestroy.
 - `build.rs` uses `cc` crate to compile .cu files, link CUDA runtime
 
 ## Test Plan
-1. Device info: print GPU name, memory, compute capability, SM count
-2. GpuBuffer: alloc 1GB, H2D copy, D2H copy, verify data
-3. Vector add kernel: launch from Rust, verify output
-4. CachingAllocator: alloc→free→realloc same size uses cache (no new cudaMalloc)
-5. Multi-stream: two concurrent memcpy on different streams
-6. Benchmark: caching allocator vs raw cudaMalloc (100 cycles)
+
+- [x] Device info: print GPU name, memory, compute capability, SM count
+- [x] GpuBuffer: alloc → H2D copy → D2H copy → verify data (256B, 64MB)
+- [x] GpuBuffer: D2D copy 验证
+- [x] GpuBuffer: zero fill 验证
+- [x] Vector add kernel: launch from Rust, verify output
+- [x] CachingAllocator: alloc→free→realloc same size uses cache (no new cudaMalloc)
+- [x] CachingAllocator: 不同 size bucket 独立缓存
+- [x] CudaStream: 创建、同步、Drop
+- [x] PinnedBuffer: page-locked host memory
+- [x] Async copy: H2D async + D2H async via stream
+
+## Takeaways
+
+1. **`cudaDeviceProp` struct 布局不可靠**：CUDA 版本之间 `cudaDeviceProp` 的字段偏移会变化。我们最初用 struct 映射读取 `total_global_mem`，得到了垃圾值（12TB）。正确做法：用 `cudaMemGetInfo` 获取显存信息，用 `cudaDeviceGetAttribute` 获取其他属性。只从 `cudaDeviceProp` 读取 `name` 字段（始终在 struct 最前面，布局稳定）。
+
+2. **Rust 2024 edition 的 unsafe 语义变更**：
+   - `extern "C"` 块必须加 `unsafe` 前缀 → `unsafe extern "C"`
+   - `unsafe fn` 内部的 unsafe 调用也需要显式 `unsafe {}` 块
+   - 这让代码更安全，但初次移植需要注意
+
+3. **`cc` crate 的 CUDA 支持是内置的**：不需要 `features = ["cuda"]`（这个 feature 不存在）。只需 `.cuda(true).cudart("shared")`。
+
+4. **Caching Allocator 的 bucket 策略**：round up to next power of 2（最小 512B）。这意味着申请 513B 会分配 1024B，存在内部碎片。但简单且高效——避免了 free list 中的精确匹配问题。PyTorch 的 CUDACachingAllocator 用了更复杂的策略（best-fit with splitting），但对于推理场景，power-of-2 bucket 已经够用。
+
+5. **`into_raw` + `from_raw` 模式**：GpuBuffer 的 RAII Drop 和 CachingAllocator 的缓存需求冲突——allocator 需要持有裸指针而不触发 Drop。`into_raw()` 消费 self（`mem::forget`），返回裸指针；`from_raw()` 重新封装。这是 Rust 中管理 RAII 生命周期的标准模式。
+
+6. **dash5 环境**：CUDA 12.9 已安装但 `nvcc` 不在 PATH（需要 `/usr/local/cuda/bin`）。Rust 需要手动安装 rustup。无 rsync，用 `tar | ssh tar` 同步代码。开发工作流：本地写码 → tar sync → 远程 build+test。

docs/02-tensor.md

@@ -0,0 +1,97 @@
+# Phase 2: Tensor Abstraction Layer — Design Document
+
+## Goal
+
+实现核心 Tensor 类型，支持 CPU/GPU 存储、多种数据类型、strided view 操作，作为后续所有算子和模型的数据基础。
+
+## Module Layout
+
+```
+crates/xserv-tensor/
+├── Cargo.toml
+└── src/
+    ├── lib.rs          # re-exports
+    ├── dtype.rs        # DType enum, TensorDType trait
+    ├── shape.rs        # strides 计算, broadcast 规则
+    ├── storage.rs      # Storage (Arc引用计数), Device enum
+    └── tensor.rs       # Tensor 主体: 创建, 形状操作, 设备迁移
+```
+
+## Key Design Decisions
+
+### DType + TensorDType Trait
+
+```rust
+pub enum DType { F32, F16, BF16 }
+
+pub trait TensorDType: Copy + Send + Sync + 'static {
+    const DTYPE: DType;
+    fn to_f64(self) -> f64;
+    fn from_f64(v: f64) -> Self;
+}
+```
+
+- 用 `half` crate 的 `bf16`/`f16` 表示半精度类型
+- `TensorDType` trait 让 `from_slice<T>` 和 `as_slice<T>` 有类型安全
+- GPU kernel 中通过 `DType` dispatch 到对应的 CUDA 类型 (`__nv_bfloat16` / `float`)
+
+### Storage 引用计数
+
+```rust
+pub struct Storage(Arc<StorageInner>);
+enum StorageInner {
+    Cpu { data: Vec<u8> },
+    Cuda { buffer: GpuBuffer },
+}
+```
+
+- `Arc` 引用计数让 transpose/slice/reshape 能共享底层数据（view 语义）
+- 不实现 CoW（copy-on-write），view 只能读不能写
+- `to_device()` 总是创建新的 Storage
+
+### Strided Tensor
+
+```rust
+pub struct Tensor {
+    storage: Storage,
+    shape: SmallVec<[usize; 4]>,
+    strides: SmallVec<[usize; 4]>,
+    offset: usize,
+    dtype: DType,
+}
+```
+
+- `SmallVec<[usize; 4]>` 避免大多数 tensor (≤4D) 的堆分配
+- `strides` 以元素为单位（不是字节）
+- `offset` 支持 slice 操作（view 到 storage 的中间位置）
+- `is_contiguous()` 检查 strides 是否与 shape 匹配
+- 非 contiguous 的 tensor 调 `contiguous()` 才能送入 CUDA kernel
+
+### Broadcast 规则
+
+实现了 NumPy-style broadcasting:
+- 维度从尾部对齐
+- 大小为 1 的维度可以广播到任意大小
+- `broadcast_strides()` 将 size=1 维度的 stride 置为 0（虚拟广播，不复制数据）
+
+## Test Plan
+
+- [x] from_slice → shape/strides 正确
+- [x] reshape, transpose, squeeze, unsqueeze
+- [x] transpose 后 contiguous() 重排数据
+- [x] BF16 tensor 的精度验证
+- [x] CPU↔GPU roundtrip
+- [x] zeros on GPU → 拷回 CPU 验证全 0
+- [x] broadcast_shape 单元测试
+
+## Takeaways
+
+1. **`SmallVec` 是正确选择**：绝大多数 tensor ≤ 4D，避免了频繁堆分配。LLM 推理中常见的维度是 `[B, S, H]` (3D) 和 `[B, H, S, D]` (4D)。
+
+2. **View 语义的取舍**：Arc 共享 storage 实现了零拷贝 transpose/reshape，但代价是无法原地修改 view 后的 tensor。对于推理引擎这是可以接受的——推理路径上大部分操作是只读的。
+
+3. **contiguous() 的隐性开销**：非 contiguous tensor 在送入 kernel 前需要 `contiguous()` 拷贝。这意味着 `transpose → matmul` 会产生一次额外拷贝。后续优化方向：在 kernel 中直接支持 strided input。
+
+4. **Rust 2024 edition 变化**：`unsafe fn` 内部的 unsafe 调用也需要显式 `unsafe {}` 块，`extern "C"` 块必须加 `unsafe` 前缀。这个 edition 对安全性更严格。
+
+5. **CPU 实现先行**：先在 CPU 上验证逻辑正确性（如 contiguous 重排），再扩展到 GPU。这个策略在后续 phase 中应该继续沿用。

docs/03-gemm.md

@@ -0,0 +1,102 @@
+# Phase 3: GEMM — Design Document
+
+## Goal
+
+实现矩阵乘法的多个版本（naive → tiled → cuBLAS），建立 benchmark 对比框架，深入理解 GPU 编程中的内存访问模式和优化手段。
+
+## Module Layout
+
+```
+csrc/gemm/
+├── naive.cu          # 每个 thread 算一个输出元素
+└── tiled.cu          # shared memory tiling, 32x32 tiles
+
+crates/xserv-kernels/
+├── build.rs          # 编译 .cu + 链接 cublas
+└── src/
+    ├── lib.rs
+    └── gemm.rs       # FFI 封装, GemmBackend enum, matmul(), CublasContext
+```
+
+## Kernel Implementations
+
+### Version 1: Naive GEMM
+
+```
+Grid:  (ceil(N/16), ceil(M/16))
+Block: (16, 16)
+每个 thread: C[row][col] = sum_k(A[row][k] * B[k][col])
+```
+
+- 每个 thread 独立遍历 K 维度做点积
+- 所有读取走 global memory，无局部性优化
+- BF16 版本在 FP32 中累加（`__bfloat162float` → 累加 → `__float2bfloat16`）
+
+### Version 2: Tiled GEMM (Shared Memory)
+
+```
+TILE_SIZE = 32
+Grid:  (ceil(N/32), ceil(M/32))
+Block: (32, 32) = 1024 threads
+
+每个 tile iteration:
+  1. 协作加载 A[tile] 和 B[tile] 到 shared memory
+  2. __syncthreads()
+  3. 在 shared memory 中做 32 次乘加
+  4. __syncthreads()
+```
+
+- 每个 global memory 读取被 TILE_SIZE 个 thread 复用
+- 理论上减少 global memory 访问 TILE_SIZE 倍
+- BF16 版本同样在 shared memory 中存 float（FP32 累加）
+
+### Version 3: cuBLAS
+
+- `cublasGemmEx` 支持混合精度
+- **Row-major 适配**：cuBLAS 使用 column-major 布局，我们的 tensor 是 row-major
+  - 利用恒等式：`C = A @ B` (row-major) ⟺ `C^T = B^T @ A^T` (col-major)
+  - 传入 `CUBLAS_OP_N`，让 cuBLAS 把我们的 row-major 数据当作 col-major 的转置
+  - 参数：`m=N, n=M, k=K, lda=N (B), ldb=K (A), ldc=N (C)`
+
+### Backend Registry
+
+```rust
+pub enum GemmBackend { Naive, Tiled, CuBlas }
+pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor;
+```
+
+运行时可切换 backend，方便 benchmark 对比和逐步替换。
+
+## CublasContext
+
+RAII 封装 `cublasHandle_t`，Drop 时调 `cublasDestroy_v2`。
+目前每次 matmul 创建一个新 handle，后续优化为全局复用。
+
+## Test Plan
+
+- [x] F32: naive/tiled/cuBLAS × small(4)/medium(64-256)/rect(65x33x97) 
+- [x] BF16: naive/tiled/cuBLAS × small/medium
+- [x] 三种 backend 在相同输入上输出一致（cross-backend consistency）
+- [x] 非方阵测试（M≠N≠K）
+- [x] 1024x1024 cuBLAS 验证
+
+## Takeaways
+
+1. **Row-major vs Column-major 陷阱**：这是 GEMM 实现中最容易出错的地方。cuBLAS 的 column-major 假设与 C/Rust 的 row-major 冲突。理解 `C=AB` ⟺ `C^T=B^T A^T` 这个恒等式是关键。实际做法：不做任何显式转置，只是交换 A/B 的传入顺序和调整 leading dimension 参数。
+
+2. **BF16 的累加精度**：BF16 只有 ~3 位有效数字（vs FP32 的 ~7 位）。如果在 BF16 中累加 K 次乘法，误差会快速放大。正确做法是**在 FP32 中累加，最后才转回 BF16**。我们的 naive 和 tiled kernel 都遵循了这一点（`float sum = 0.0f`）。cuBLAS 通过 `CUBLAS_COMPUTE_32F` 参数控制。
+
+3. **Shared memory tiling 的核心思想**：global memory 带宽是 GPU 计算的主要瓶颈。通过 shared memory tiling，每个数据从 global memory 读一次，被 TILE_SIZE 个 thread 复用。对于 TILE_SIZE=32，理论上减少 32 倍 global memory 访问。
+
+4. **`__syncthreads()` 的位置关键**：tile 加载后必须同步（确保所有 thread 写完 shared memory），计算后也要同步（防止下一轮加载覆盖还在使用的数据）。漏掉任何一个 sync 都会产生 race condition 导致结果错误。
+
+5. **cuBLAS handle 开销**：每次 matmul 创建/销毁 handle 有~0.1ms 开销。生产环境应全局复用一个 handle。Phase 15（性能优化）时需要修复这个问题。
+
+6. **`error::check` 需要 pub**：Phase 1 中 `check()` 是 `pub(crate)`，Phase 3 需要跨 crate 调用。反思：基础设施 crate 的错误处理函数应该从一开始就设计为 public API。
+
+## 后续优化方向（Phase 15）
+
+- Register tiling（每个 thread 算多个输出元素）
+- Tensor Core WMMA（利用 5090 的硬件加速）
+- CublasContext 全局复用
+- 非 contiguous input 支持（避免 matmul 前的拷贝）

docs/04-transformer-kernels.md

@@ -0,0 +1,213 @@
+# Phase 4: Transformer Core Kernels — Design Document
+
+## Goal
+
+实现 Transformer 所需的所有非 Attention 算子的 CUDA kernel，每个 kernel 都支持 BF16 和 F32，与 PyTorch 参考实现对比验证。
+
+## Kernel 清单
+
+| Kernel | 用于 | 核心计算 | 关键优化点 |
+|--------|------|---------|-----------|
+| LayerNorm | GPT-2 | `(x - mean) / sqrt(var + eps) * gamma + beta` | Welford online, warp reduce |
+| RMSNorm | Qwen3 | `x / sqrt(mean(x²) + eps) * gamma` | 无 mean，比 LayerNorm 简单 |
+| GELU | GPT-2 | `0.5x(1 + tanh(sqrt(2/π)(x + 0.044715x³)))` | tanh 近似，逐元素 |
+| SiLU | Qwen3 | `x * sigmoid(x)` | 逐元素 |
+| Softmax | Attention | `exp(x - max) / sum(exp(x - max))` | Online safe softmax, warp reduce |
+| Embedding | 全部 | `output[i] = table[token_ids[i]]` | Gather, coalesced write |
+| RoPE | Qwen3 | 对 Q/K 的相邻元素对做旋转 | Precompute freq, in-place |
+
+## 文件布局
+
+```
+csrc/
+├── normalization/
+│   ├── layernorm.cu
+│   └── rmsnorm.cu
+├── activation/
+│   ├── gelu.cu
+│   └── silu.cu
+├── reduce/
+│   └── softmax.cu
+├── embedding/
+│   ├── embedding.cu
+│   └── rope.cu
+
+crates/xserv-kernels/src/
+├── layernorm.rs
+├── rmsnorm.rs
+├── activation.rs      # GELU + SiLU
+├── softmax.rs
+├── embedding.rs
+├── rope.rs
+└── lib.rs             # 新增 mod 声明
+```
+
+## Kernel 设计细节
+
+### LayerNorm
+
+输入 `x: [*, hidden_size]`, 输出 `y: [*, hidden_size]`
+参数 `gamma, beta: [hidden_size]`
+
+```
+y[i] = gamma[i] * (x[i] - mean) / sqrt(var + eps) + beta[i]
+```
+
+**GPU 映射**: 每个 thread block 处理一行（一个 hidden_size 向量）。
+- Phase 1: 并行加载 x，Welford online 算法计算 mean 和 var
+- Phase 2: warp-level reduce (`__shfl_down_sync`) 聚合 mean/var
+- Phase 3: block-level reduce via shared memory
+- Phase 4: 每个 thread 对自己负责的元素做 normalize + affine
+
+**Block 配置**: `block = min(1024, hidden_size)`, `grid = num_rows`
+
+### RMSNorm
+
+比 LayerNorm 简单：不减 mean，只做 `x * rsqrt(mean(x²) + eps) * gamma`。
+
+```
+rms = sqrt(sum(x²) / hidden_size + eps)
+y[i] = x[i] / rms * gamma[i]
+```
+
+**GPU 映射**: 同 LayerNorm，每个 block 处理一行。
+- 只需要一次 reduce（求 sum(x²)），不需要两次（mean + var）。
+
+### GELU
+
+逐元素操作，用 tanh 近似:
+```
+gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))
+```
+
+**GPU 映射**: 每个 thread 处理多个元素（向量化），grid 覆盖全部元素。
+
+### SiLU (Swish)
+
+逐元素: `silu(x) = x * sigmoid(x) = x / (1 + exp(-x))`
+
+### Softmax
+
+输入 `x: [*, seq_len]`, 沿最后一维做 softmax:
+```
+1. m = max(x)                    // 数值稳定
+2. e[i] = exp(x[i] - m)
+3. s = sum(e)
+4. y[i] = e[i] / s
+```
+
+**GPU 映射**: 每个 block 处理一行。
+- 第一遍 reduce: 求 max
+- 第二遍: exp(x - max) 并 reduce sum
+- 第三遍: 除以 sum
+
+**优化**: 可以用 online softmax 合并前两遍（边算 exp 边更新 max），但先实现三遍版本保证正确。
+
+### Embedding
+
+```
+output[seq_idx] = embedding_table[token_ids[seq_idx]]
+```
+
+**GPU 映射**: 每个 thread 处理一个 token 的部分维度。
+- `grid = num_tokens`, `block = hidden_size`（或分多个 thread 处理一个 token）
+- 写端是 coalesced（连续 thread 写连续地址），读端是 gather（非连续）
+
+### RoPE (Rotary Position Embedding)
+
+对 Q/K 的每对相邻元素 `(x0, x1)` 做 2D 旋转:
+```
+freq[i] = 1.0 / (theta ^ (2i / dim))
+cos_val = cos(position * freq[i])
+sin_val = sin(position * freq[i])
+y0 = x0 * cos_val - x1 * sin_val
+y1 = x0 * sin_val + x1 * cos_val
+```
+
+**GPU 映射**: 每个 thread 处理一对元素 `(x[2i], x[2i+1])`。
+- Precompute `cos_cache[max_seq_len][head_dim/2]` 和 `sin_cache` 在初始化时
+- 运行时 kernel 只做乘加
+
+**theta**: Qwen3 默认 `rope_theta = 1000000.0`
+
+## Reduction Pattern（核心学习点）
+
+所有 Norm 和 Softmax 都涉及 reduction。GPU reduction 的分层结构:
+
+```
+Thread-level:  每个 thread 处理多个元素，本地累加
+    ↓
+Warp-level:    __shfl_down_sync() 在 32 threads 内规约（无需 shared memory）
+    ↓
+Block-level:   shared memory 存各 warp 的结果，warp 0 再规约
+```
+
+对于 hidden_size <= 8192（LLM 常见），一个 block 足够，不需要 grid-level reduction。
+
+### Warp Reduce 模板
+
+```cuda
+__device__ float warp_reduce_sum(float val) {
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val += __shfl_down_sync(0xffffffff, val, offset);
+    return val;
+}
+```
+
+### Block Reduce 模板
+
+```cuda
+__device__ float block_reduce_sum(float val) {
+    __shared__ float shared[32];  // max 32 warps per block
+    int lane = threadIdx.x % 32;
+    int warp_id = threadIdx.x / 32;
+
+    val = warp_reduce_sum(val);
+    if (lane == 0) shared[warp_id] = val;
+    __syncthreads();
+
+    val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
+    if (warp_id == 0) val = warp_reduce_sum(val);
+    return val;
+}
+```
+
+## Reference 验证策略
+
+写 `tools/generate_reference.py` 脚本，用 PyTorch 为每个 op 生成 reference input/output:
+- 保存为 `.npy` 格式
+- Rust 测试中加载对比
+- 或者直接在 Rust 测试中用 CPU 实现计算 expected 值（更简单，不依赖 Python）
+
+**选择**: 先用 Rust CPU 实现作为 reference（简单），关键 op（RoPE）再与 PyTorch 对比。
+
+## Test Plan
+
+- [x] RMSNorm F32: hidden_size=768, 4 rows → max_err 7.2e-7
+- [x] RMSNorm BF16: 同上 → max_err 7.0e-3
+- [x] LayerNorm F32: hidden_size=768 → max_err 1.7e-6
+- [x] GELU F32: 10000 elements → max_err 3.0e-8
+- [x] GELU BF16: 同上 → max_err 2.4e-3
+- [x] SiLU F32: 10000 elements → max_err 1.5e-8
+- [x] Softmax F32: 8×256 → max_err 1.4e-9
+- [x] Softmax sum=1 验证: 4×2048
+- [x] Softmax 大值 (1000+) 数值稳定性 → max_err 1.5e-8
+- [x] Embedding F32: vocab=100, hidden=64, 5 tokens → exact match
+- [x] RoPE F32: 4 tokens × 2 heads × dim=8 → max_err 6.0e-8
+- [x] RoPE position=0 恒等验证 → max_err 0
+
+## Takeaways
+
+1. **`common.cuh` 抽取共用 reduction 是正确的做法**：`warp_reduce_sum/max` 和 `block_reduce_sum/max` 被 RMSNorm, LayerNorm, Softmax 三个 kernel 复用。抽到头文件避免了代码重复，也确保 reduction 逻辑一致。build.rs 中需要 `.include("../../csrc")` 让 nvcc 能找到头文件。
+
+2. **Shared memory 中广播标量的模式**：Norm 和 Softmax 都需要将 reduce 结果（mean, rms_inv, max, sum）广播给 block 内所有 thread。标准做法：thread 0 写 `__shared__` 变量，`__syncthreads()` 后所有 thread 读。这比让每个 thread 独立做 reduce 高效得多。
+
+3. **Softmax 三遍 vs 两遍**：我们实现了三遍版本（max → exp+sum → normalize），简单可靠。Online softmax 可以合并前两遍（一遍 pass 内同时跟踪 running max 和 running sum），但需要更复杂的数值更新公式。Flash Attention（Phase 14）会用到 online softmax。
+
+4. **RoPE 的 position=0 恒等性**：`cos(0)=1, sin(0)=0`，所以 position 0 的旋转是恒等变换。这是一个很好的 sanity check。如果 position=0 时输出不等于输入，说明 kernel 有 bug。
+
+5. **BF16 Softmax 的精度陷阱**：exp 结果先写成 BF16 再读回做 normalize 会丢精度。理想做法是用 float scratch buffer 暂存 exp 结果。当前实现可接受（误差在 1e-2 量级），但在 attention score 很接近时可能引入可观察的差异。Phase 14 Flash Attention 会解决这个问题（全程 FP32 累加）。
+
+6. **Embedding 就是 gather 操作**：没有任何计算，纯粹的内存搬运。瓶颈在 global memory 随机读取（token_ids 导致不连续读 table）。写端是 coalesced 的（连续 token 写连续地址）。优化方向：使用向量化加载（`float4`）一次读 128 bit。
+
+7. **RoPE in-place 修改 Tensor 的设计考量**：RoPE 在数学上是对 Q/K 的 in-place 旋转。我们通过 `data_ptr() as *mut` 绕过了 Rust 的不可变借用。这在 GPU 上是安全的（kernel 内部互不干扰），但 Rust 侧没有 `&mut` 语义保护。后续如果需要更严格的安全性，可以引入 `Tensor::as_mut_ptr()` 方法并要求 `&mut self`。

docs/05-attention.md

@@ -0,0 +1,92 @@
+# Phase 5: Naive Attention Kernel — Design Document
+
+## Goal
+
+实现标准 Multi-Head Attention（不做 Flash/Paged 优化），用组合式方法（GEMM + Softmax）完成。这是理解 attention 计算流程的基础，也是后续 Flash Attention 的 baseline。
+
+## 计算流程
+
+```
+Input: Q [B, H, S, D], K [B, H, S, D], V [B, H, S, D]
+       B=batch, H=num_heads, S=seq_len, D=head_dim
+
+1. scores = Q @ K^T / sqrt(D)           → [B, H, S, S]
+2. scores += causal_mask                 → 上三角置为 -inf
+3. weights = softmax(scores, dim=-1)     → [B, H, S, S]
+4. output = weights @ V                  → [B, H, S, D]
+```
+
+## 设计选择
+
+### 组合式实现（Phase 3 GEMM + Phase 4 Softmax）
+
+不写新的 fused CUDA kernel，而是复用已有的 matmul 和 softmax：
+- `scores = batched_matmul(Q, K^T)` — 需要支持 batched GEMM
+- `masked_fill(scores, causal_mask, -inf)` — 新的逐元素 kernel
+- `softmax(scores)` — 复用 Phase 4
+- `output = batched_matmul(weights, V)` — 复用 batched GEMM
+
+这意味着需要先扩展 matmul 支持 batched GEMM（cublasGemmStridedBatchedEx）。
+
+### Causal Mask
+
+不显式构造 mask 矩阵。写一个 kernel：
+```
+if (col > row + offset) score = -infinity
+```
+其中 offset 用于支持 KV cache 场景（decode 时 query 的 row 偏移）。
+
+### Batched GEMM via cuBLAS
+
+`cublasGemmStridedBatchedEx` 在一个 batch 维度上并行执行多个 GEMM：
+```
+C[b] = A[b] @ B[b]  for b = 0..batch_count
+stride_a = M * K, stride_b = K * N, stride_c = M * N
+```
+
+Attention 中 batch 维度 = B * H（batch_size × num_heads）。
+
+## 文件布局
+
+```
+csrc/attention/
+└── causal_mask.cu      # causal mask fill kernel
+
+crates/xserv-kernels/src/
+├── gemm.rs             # 扩展: batched_matmul
+├── attention.rs        # NEW: multi_head_attention()
+└── causal_mask.rs      # NEW: causal mask apply
+```
+
+## API 设计
+
+```rust
+/// Multi-head attention (naive, materializes S×S scores).
+/// q, k, v: [batch, num_heads, seq_len, head_dim]
+/// Returns: [batch, num_heads, seq_len, head_dim]
+pub fn attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tensor;
+
+/// Batched matmul: A[b] @ B[b] for all b.
+/// a: [..., M, K], b: [..., K, N] → [..., M, N]
+pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor;
+```
+
+## Test Plan
+
+- [x] batched_matmul: [4,8,32,64]×[4,8,64,32] → max_err 2.7e-7
+- [x] attention (non-causal): B=1,H=2,S=8,D=16 → max_err 4.5e-8
+- [x] attention (causal): B=1,H=2,S=16,D=32 → max_err 3.0e-8
+- [x] attention (causal, larger): B=2,H=4,S=64,D=64 → max_err 6.0e-8
+- [x] causal mask 语义: position 0 只能看到 token 0，output[0] == V[0] → exact
+
+## Takeaways
+
+1. **`to_device` 不应强制 contiguous**：最初 `to_device()` 会先调 `contiguous()`，而 GPU 的 `contiguous()` 又调 `to_device(Cpu)`，导致无限递归栈溢出。修复：`to_device()` 直接传输 raw storage，保留 strides/offset，用户需要时自己调 `contiguous()`。GPU `contiguous()` 现在走 GPU→CPU→CPU contiguous→CPU→GPU 路径——正确但低效，Phase 15 需要写 GPU contiguous kernel。
+
+2. **Batched GEMM via `cublasGemmStridedBatchedEx`**：row-major trick 同 Phase 3，额外参数是 stride（元素数，不是字节）。stride_a = M×K, stride_b = K×N, stride_c = M×N。注意初始版本错误地乘了 `elem_size`，cuBLAS 的 stride 单位是元素。
+
+3. **Attention 的组合式实现足够验证正确性**：没有写 fused kernel，而是复用 `batched_matmul` + `scale` + `causal_mask` + `softmax`。精度极好（max_err < 1e-7），因为每步都在 FP32 中完成。缺点是 S×S score 矩阵完全 materialize（O(S²) 显存），Flash Attention 会解决。
+
+4. **Scale kernel 的必要性**：原本想在 CPU 上做 scale（round-trip），但那太慢了。加了 `scale_f32/bf16` 逐元素 CUDA kernel。未来可以把 scale 合进 GEMM 的 alpha 参数，省一次 kernel launch。
+
+5. **Causal mask 的 offset 设计**：`col > row + offset` 中的 offset 为 KV cache 场景预留。Decode 时 Q 只有 1 行但 KV cache 有前 S 行，offset = kv_len - q_len 确保 decode query 能看到所有 cached tokens。

docs/06-model-loading.md

@@ -0,0 +1,69 @@
+# Phase 6: Model Loading — Design Document
+
+## Goal
+
+从 HuggingFace safetensors 文件加载模型权重到 GPU Tensor。解析 config.json 获取模型结构参数。
+
+## Crate: `xserv-model`
+
+```
+crates/xserv-model/src/
+├── lib.rs
+├── config.rs       # ModelConfig from config.json
+├── loader.rs       # safetensors weight loading
+└── gpt2.rs         # (Phase 8) GPT-2 model definition
+```
+
+## Dependencies
+
+- `safetensors` crate: parse safetensors format
+- `serde` + `serde_json`: deserialize config.json
+- `memmap2`: mmap for zero-copy file access (safetensors uses this internally)
+
+## Weight Loading Flow
+
+```
+safetensors file (disk)
+  → safetensors crate parses header (tensor names, shapes, dtypes, offsets)
+  → mmap raw data
+  → for each tensor:
+      → read bytes at offset
+      → create CPU Tensor from raw bytes
+      → .to_device(Cuda(0)) → GPU Tensor
+  → return HashMap<String, Tensor>
+```
+
+## Config Parsing
+
+```rust
+#[derive(Deserialize)]
+pub struct ModelConfig {
+    pub architectures: Option<Vec<String>>,
+    pub model_type: Option<String>,
+    pub hidden_size: usize,
+    pub intermediate_size: Option<usize>,
+    pub num_attention_heads: usize,
+    pub num_key_value_heads: Option<usize>,
+    pub num_hidden_layers: usize,
+    pub vocab_size: usize,
+    pub max_position_embeddings: Option<usize>,
+    pub layer_norm_eps: Option<f64>,
+    pub rms_norm_eps: Option<f64>,
+    pub rope_theta: Option<f64>,
+    pub tie_word_embeddings: Option<bool>,
+}
+```
+
+## Test Plan
+
+- [x] Load GPT-2 124M: 160 tensors loaded successfully
+- [x] Parse GPT-2 config.json: hidden=768, layers=12, heads=12, vocab=50257
+- [x] Sharded loading path implemented (for larger models)
+
+## Takeaways
+
+1. **GPT-2 vs modern HF config naming**：GPT-2 uses `n_embd`/`n_head`/`n_layer`/`n_positions`，而不是 `hidden_size`/`num_attention_heads` 等。ModelConfig 需要支持两套命名并提供统一的 accessor methods（`hidden()`, `num_heads()` 等）。
+
+2. **safetensors 零拷贝读取**：`safetensors` crate 直接 mmap 文件，解析 header 得到 tensor 的 offset 和 shape，然后 zero-copy 读取 raw bytes。对于 GPT-2 的 500MB 权重文件，加载速度很快。
+
+3. **模型下载的网络问题**：HuggingFace 在中国网络下不可达。使用 modelscope.cn 或 hf-mirror.com 作为替代。大文件（>100MB）的 redirect 到 CDN 可能也会失败，modelscope 的 snapshot_download 更可靠。

docs/07-tokenizer.md

@@ -0,0 +1,57 @@
+# Phase 7: BPE Tokenizer — Design Document
+
+## Goal
+
+从零实现 Byte-Pair Encoding tokenizer，兼容 HuggingFace `tokenizer.json` 格式。支持 GPT-2 和 Qwen3。
+
+## Crate: `xserv-tokenizer`
+
+```
+crates/xserv-tokenizer/src/
+├── lib.rs
+├── bpe.rs          # BPE encode/decode core algorithm
+└── chat.rs         # Chat template formatting
+```
+
+## Dependencies
+
+- `serde` + `serde_json`: parse tokenizer.json
+- `regex`: pre-tokenization patterns
+
+## BPE Algorithm
+
+### Encode
+1. Pre-tokenize: split text by regex (GPT-2 pattern)
+2. Each word → byte sequence → initial token list (one token per byte)
+3. Repeatedly merge highest-priority pair until no more merges
+4. Map merged tokens to IDs via vocab
+
+### Decode
+Token IDs → lookup vocab → concatenate bytes → UTF-8 decode
+
+## Key Data Structures
+
+```rust
+pub struct Tokenizer {
+    vocab: HashMap<Vec<u8>, u32>,        // token bytes → ID
+    vocab_rev: Vec<Vec<u8>>,             // ID → token bytes
+    merges: Vec<(Vec<u8>, Vec<u8>)>,     // ordered merge rules
+    merge_ranks: HashMap<(u32, u32), usize>,  // (id_a, id_b) → priority
+    special_tokens: HashMap<String, u32>,
+    pre_tokenize_regex: Regex,
+}
+```
+
+## Test Plan
+
+- [x] Encode + decode roundtrip verified (GPT-2 tokenizer, English text)
+- [x] Special tokens handled (endoftext)
+- [x] Integrated into GPT-2 inference pipeline, generates coherent text
+
+## Takeaways
+
+1. **GPT-2 byte-to-unicode 映射**：GPT-2 的 vocab 中，每个 byte 都映射到一个 Unicode 字符。可打印 ASCII (0x21-0x7E) 映射到自身，其余字节（空格、控制字符等）映射到 U+0100 以上的 Unicode 码点。解码时需要反向映射。这个映射表是 BPE tokenizer 正确性的关键。
+
+2. **Rust regex 不支持 lookahead**：GPT-2 的 pre-tokenization regex 使用了 `(?!\S)` lookahead，Rust 的 `regex` crate 不支持。简化为去掉 lookahead 后功能等价（whitespace 仍然被正确分词）。如果需要精确匹配 Python 行为，需要 `fancy-regex` crate。
+
+3. **BPE merge 的 O(n²) 复杂度**：当前实现每次 merge 扫描整个 token 序列找最高优先级 pair，复杂度 O(n² × |merges|)。对于短文本够用，长文本需要 priority queue 优化。推理场景中 prompt 通常 < 10K tokens，暂时可接受。

docs/08-gpt2.md

@@ -0,0 +1,71 @@
+# Phase 8: GPT-2 Complete Inference — Design Document (Milestone ①)
+
+## Goal
+
+Wire everything together: load GPT-2 124M, tokenize input, run forward pass, sample tokens, decode output. First time seeing the model "speak".
+
+## Model Architecture (GPT-2 124M)
+
+```
+hidden_size = 768
+num_heads = 12
+num_layers = 12
+vocab_size = 50257
+max_position_embeddings = 1024
+activation = GELU
+normalization = LayerNorm (pre-LN)
+tied embeddings (lm_head == wte)
+```
+
+## Forward Pass
+
+```
+tokens [S]
+  → wte[tokens] + wpe[0..S]           → [S, 768]
+  → for each layer:
+      residual = x
+      x = layernorm(x, ln_1)
+      x = attention(x)                 # Q,K,V from linear, MHA, output linear
+      x = x + residual
+      residual = x
+      x = layernorm(x, ln_2)
+      x = mlp(x)                       # linear→GELU→linear
+      x = x + residual
+  → layernorm(x, ln_f)
+  → logits = x @ wte.T                → [S, 50257]
+  → sample(logits[-1])                 → next token
+```
+
+## Sampling
+
+- Greedy: argmax
+- Temperature: logits / T → softmax → sample
+- Top-K: keep top-k logits, rest = -inf
+- Top-P: sorted by prob, cumsum ≤ p
+
+## CLI Binary
+
+```
+$ cargo run --release --bin xserv-cli -- --model path/to/gpt2
+
+xserv> The future of AI is
+GPT-2> ...generated text...
+```
+
+## Test Plan
+
+- [x] Greedy generation produces coherent English text
+- [x] Interactive CLI works (pipe and interactive mode)
+- [x] Multiple prompts verified: "The future of AI is", "Once upon a time"
+
+## Takeaways
+
+1. **QKV split + head reshape 的 layout 陷阱（最关键的 bug）**：GPT-2 的 `c_attn` 输出 `[S, 3H]` 需要 split 成 Q/K/V 再 reshape 成 `[1, num_heads, S, head_dim]`。关键错误：从 `[S, num_heads, head_dim]` 直接 `reshape` 到 `[1, num_heads, S, head_dim]` 不等于 transpose！Reshape 只是重新解释 flat data 的 shape，不会重排数据。必须手动按 `[batch, head, seq, dim]` 的目标 layout 写入数据。同理 merge_heads 也需要手动重排。
+
+2. **CPU round-trip 作为 correctness first 策略**：`add_tensors`、`add_bias`、`split_qkv`、`merge_heads` 都通过 CPU round-trip 实现。虽然慢（每次都有 GPU→CPU→GPU 拷贝），但确保了正确性。Phase 15 会写专门的 CUDA kernel 替换这些操作。
+
+3. **GPT-2 的 Conv1D 权重布局**：GPT-2 用 `Conv1D` 而非 `Linear`，权重存为 `[in, out]`（不是标准 Linear 的 `[out, in]`）。计算方式是 `x @ weight`（不需要转置）。这和 Qwen3/LLaMA 的 `[out, in]` 布局不同——Phase 10 需要注意。
+
+4. **Greedy decoding 的重复问题**：GPT-2 124M 在 greedy decoding 下极易陷入循环（"The world was a place of great danger, and..."）。这是已知行为，temperature + top-k/top-p sampling 可以缓解。当前实现只有 greedy，sampling 将在后续添加。
+
+5. **无 KV Cache 的性能代价**：每生成一个 token 都要重新跑完整 forward pass（O(S²) attention）。50 tokens 的生成需要 50 次 full forward，每次的 attention 复杂度还在增长。Phase 9 的 KV Cache 会将 decode 降到 O(S) per token。

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-8B Support — Design Document (Milestone ②)
+
+## Goal
+
+扩展模型定义支持 Qwen3-8B 架构，验证输出正确性。与 GPT-2 的关键差异：RMSNorm、RoPE、GQA、SwiGLU、不共享 embedding。
+
+## 架构差异 (GPT-2 → Qwen3)
+
+| 特性 | GPT-2 | Qwen3-8B |
+|------|-------|----------|
+| 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 | 4096 |
+| num_layers | 12 | 36 |
+| 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)
+
+```
+权重: 8B × 2B = ~16 GB (BF16)
+        8B × 4B = ~32 GB (FP32) — 不够! 必须用 BF16
+KV cache (S=256, B=1): ~0.1 GB
+总计: ~16 GB (BF16), 单卡可运行
+```
+
+**关键**: Qwen3-8B 必须用 BF16 才能在单张 5090 (32GB) 上运行。当前 GPT-2 用 FP32，需要支持 BF16 forward pass。
+
+## Implementation Plan
+
+1. 下载 Qwen3-8B 模型 (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,61 @@
+# Phase 11: GPU-Resident KV Cache — Design Document
+
+> **注意**: 原计划为 "Paged Attention + KV Cache Manager"，实际实现为 GPU 连续预分配 KV cache（非 paged）。Paged allocation 留待后续优化。
+
+## Goal
+
+将 KV cache 从 CPU Vec 迁移到 GPU，消除每步 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,157 @@
+# 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. 编写并发测试脚本
+
+## 当前状态
+
+**已实现: iteration-level scheduling**。多请求可以并发进入 batch (max_batch_size)，新请求在 mid-generation 动态加入。Prefill 和 decode 阶段在每轮迭代内分离处理。
+
+**未实现: batched GPU forward**。每个 seq 的 model forward 仍是串行调用 (per-seq forward_gpu_cache)。真正的 batched decode (多 seq 的 token 合并为一次 GPU forward) 需要 Flash Attention 的 variable-length attention 支持。Phase 14 实现了 FA2 kernel，为后续 batched forward 提供了基础。
+
+**验证**: 8 个并发请求 (max_batch=4) 总 wall clock 22.5s，各请求延迟之和 135.0s，调度加速 6.0x。Server log 确认 `decode batch_size=4`。

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/14-flash-attention.md

@@ -0,0 +1,167 @@
+# Phase 14: Flash Attention 2 for SM120 — Design Document
+
+## Goal
+
+用自写的 Flash Attention 2 CUDA kernel 替换 naive attention (Phase 5)。消除 O(S²) 显存分配，支持 GQA kernel 内部索引（消除 repeat_kv 开销）。
+
+## 硬件约束: FA4 不适用于 RTX 5090
+
+Flash Attention 已发展到第 4 代 (FA4, arxiv 2603.05451)，但各版本有明确硬件依赖：
+
+| 版本 | 目标架构 | 关键硬件特性 | RTX 5090 (SM120) |
+|------|---------|------------|-----------------|
+| FA2 | 通用 CUDA (SM75+) | shared memory + HMMA | **兼容** |
+| FA3 | Hopper SM90 (H100) | TMA + WGMMA + warp specialization | 不兼容 |
+| FA4 | Blackwell SM100 (B200/B300) | TMEM + async MMA + 2-CTA mode | 不兼容 |
+
+RTX 5090 使用消费级 Blackwell (GB202, SM120)，与数据中心 Blackwell (B200, SM100) 是不同硅片。SM120 **没有 TMEM (Tensor Memory)**，这是 FA4 kernel 设计的核心硬件依赖。这不是软件限制，是硬件级差异。
+
+因此本项目实现 **FA2 算法**，使用标准 CUDA (shared memory + 标准 HMMA)。
+
+## Naive Attention 的问题
+
+Phase 5 的 naive attention 流程:
+```
+k_t = K.transpose(2,3).contiguous()    ← 分配 K^T 显存
+scores = batched_matmul(Q, k_t)         ← 分配 [B,H,S,S] score 矩阵 (O(S²) 显存)
+scores = scale(scores, 1/sqrt(d))       ← 逐元素 kernel
+causal_mask(scores)                     ← 逐元素 kernel
+weights = softmax(scores)               ← 分配 [B,H,S,S] weight 矩阵
+output = batched_matmul(weights, V)     ← 最终结果
+```
+
+问题:
+1. **显存 O(S²)**: score 和 weight 矩阵各需 `B × H × S × S × dtype_size`。S=2048, H=32, BF16 → 256 MB。S=8192 → 4 GB。
+2. **GQA 预处理**: 在调用 attention 前需要 `repeat_kv_gpu` 将 K/V 从 8 heads 扩展到 32 heads，每层额外分配和拷贝。
+3. **多次 kernel launch**: scale, mask, softmax 各一次 kernel launch + global memory round-trip。
+4. **K^T materialization**: `K.transpose().contiguous()` 需要分配和拷贝。
+
+## FA2 算法
+
+核心思想: **不 materialize S×S 矩阵**。将 Q, K, V 分成 tiles，在 shared memory (SRAM) 中计算，使用 **online softmax trick** 边算边更新 running max 和 sum。
+
+FA2 (Dao 2023) 相比 FA1 的改进: 外层循环遍历 Q tiles (而非 K/V)，减少 HBM 读写次数，提高并行性。
+
+```
+scale = 1 / sqrt(head_dim)
+
+for each Q tile (q_start..q_start + BR):                   // 外层: Q tiles
+    load Q_tile [BR, D] to shared memory (一次加载，内层复用)
+    init per-row: O[D] = 0, m = -inf, l = 0
+
+    for each K/V tile j (kv_start..kv_start + BC):          // 内层: K/V tiles
+        // Causal tile-skip: 如果整个 K tile 在 Q tile "未来"，跳过
+        if causal && kv_start > max_q_pos + kv_offset: skip
+
+        load K_tile [BC, D] to shared memory
+        S = Q_tile @ K_tile^T * scale                       // [BR, BC], in registers
+        if causal: mask S[r][c] = -inf where kv_pos > q_pos
+
+        // Online softmax update
+        m_new = max(m, rowmax(S))
+        P = exp(S - m_new)
+        l_new = exp(m - m_new) * l + rowsum(P)
+        O = exp(m - m_new) * O                              // rescale accumulator
+
+        load V_tile [BC, D] to shared memory (复用 K 的空间)
+        O += P @ V_tile                                     // accumulate
+
+        m = m_new, l = l_new
+
+    O = O / l                                               // final normalize
+    write O[BR, D] to HBM (convert FP32 → BF16)
+```
+
+## 实现细节
+
+### Kernel 配置
+
+| 参数 | 值 | 说明 |
+|------|---|------|
+| BR (Q tile rows) | 64 | Q tile 大小 |
+| BC (K/V tile rows) | 64 | K/V tile 大小 |
+| head_dim | 运行时参数 (≤128) | 支持 64 (GPT-2) 和 128 (Qwen3) |
+| Block size | 128 threads | 64 线程各 own 一行 Q，其余协助加载 |
+| Grid | (q_tiles, batch × num_q_heads) | 每个 block 处理一个 Q tile + 一个 head |
+
+### Shared Memory (BF16 存储)
+
+```
+smem_q [BR × head_dim] BF16 = 64 × 128 × 2 = 16 KB (加载一次，内层复用)
+smem_kv[BC × head_dim] BF16 = 64 × 128 × 2 = 16 KB (K 和 V 交替使用)
+────────────────────────────────────────────
+Total: 32 KB (SM120 默认 48 KB，余量充足)
+```
+
+### 线程映射
+
+- Thread 0..63: 各 own Q_tile 的一行。负责该行的全部计算：dot products、softmax、PV 累加。
+- Thread 64..127: 协助 shared memory 加载 (K/V tile)，不参与计算。
+- 加载模式: 每个 thread 加载 `(BR × head_dim) / 128 = 64` 个 BF16 元素。
+
+### Per-Thread Register 使用
+
+```
+O_acc[128] FP32     = 512 bytes (128 regs)  — 输出累加器
+P[64]      FP32     = 256 bytes (64 regs)   — 当前 tile 的 softmax 后权重
+m, l       FP32     = 8 bytes (2 regs)      — online softmax running state
+循环变量 + 临时      ≈ 16 regs
+────────────────────────────────────────────
+Total: ~210 regs/thread (max 255，在限制内)
+```
+
+### GQA 支持
+
+每个 thread block 处理一个 Q head，通过 `kv_head = q_head / (num_q_heads / num_kv_heads)` 映射到对应的 KV head。K/V 的数据指针直接指向 KV head 的存储，无需 repeat_kv。
+
+```
+// 32 Q heads, 8 KV heads → heads_per_group = 4
+// Q head 0,1,2,3 → KV head 0
+// Q head 4,5,6,7 → KV head 1
+// ...
+kv_head = q_head / heads_per_group;
+K_ptr = K + (batch * num_kv_heads + kv_head) * kv_len * head_dim;
+```
+
+### Causal Mask
+
+两级优化:
+1. **Tile-level skip**: 如果 `kv_tile_start > max_q_pos + kv_offset`，整个 K/V tile 都在未来，跳过（减少 ~50% 计算）。
+2. **Element-level mask**: 在 tile 内部，`if kv_pos > q_pos + kv_offset: S = -inf`。
+
+`kv_offset = kv_len - q_len` 处理 decode 时 KV cache 长于 Q 的情况。
+
+## 与 Naive Attention 的对比
+
+| 特性 | Naive (Phase 5) | FA2 (Phase 14) |
+|------|----------------|----------------|
+| 显存 | O(B × H × S²) | O(B × H × S × D) |
+| GQA | 需要 repeat_kv (分配+拷贝) | Kernel 内部索引 (零开销) |
+| K^T | 需要 transpose+contiguous | Kernel 内部计算 |
+| Kernel launches | 6 (matmul, scale, mask, softmax, matmul, ...) | 1 (单个 fused kernel) |
+| S=8192 可行性 | OOM (~4 GB score matrix) | 可行 (32 KB shared memory) |
+
+## 源码结构
+
+```
+csrc/attention/flash_attention.cu     — FA2 kernel (BF16 in, FP32 accumulate, BF16 out)
+crates/xserv-kernels/src/attention.rs — flash_attention() Rust wrapper + 原 attention() 保留
+crates/xserv-model/src/qwen3.rs       — forward_gpu_cache 调用 flash_attention
+```
+
+## 已知局限与后续优化方向
+
+1. **Decode (Q_len=1) 效率低**: BR=64 线程中只有 1 个 active（owns_row）。应写专用 decode attention kernel，沿 KV 维度 parallel reduction。
+2. **无向量化加载**: 当前逐元素 bf16→f32 转换，应改用 `float4` 或 `__nv_bfloat162` 批量加载。
+3. **Register tiling**: 每个 thread 目前串行计算 dot product (128 MADs per K column)。可改为多线程协作。
+4. **K/V double buffering**: 可在计算当前 tile 时预加载下一个 tile 到另一半 shared memory。
+5. **Tile size 调优**: 更大的 tile (BR=128) 可能在长 sequence 时更优，需要 opt-in shared memory。
+
+## Test Plan
+
+- [x] 正确性: logits 与 HF transformers 对比 (top-1 match 9/10, top-5 overlap 4.0/5)
+- [x] 生成质量: 52/52 prompt 生成连贯文本，中英文均可
+- [x] SSE streaming 正常工作
+- [x] 性能: 12.9 tok/s (vs naive 10.3 tok/s, +25%)
+- [ ] 长 sequence (S=4096, S=8192): 验证 naive OOM 而 FA2 正常
+- [ ] ncu profile: compute utilization, memory throughput

docs/15-performance.md

@@ -0,0 +1,177 @@
+# Phase 15: Performance Optimization — Design Document (Milestone ④)
+
+## Goal
+
+系统性 profiling + 优化，从 12.9 tok/s (Phase 14 结束) 逼近 RTX 5090 的理论带宽上限 (112 tok/s)。
+
+## 硬件 Roofline
+
+RTX 5090 (SM120, CC 12.0) 的 decode 理论极限：
+
+```
+模型权重:          16 GB (Qwen3-8B BF16)
+内存带宽:          1.79 TB/s (GDDR7)
+理论最优 decode:   16 GB / 1.79 TB/s = 8.9 ms/step = 112 tok/s (batch=1)
+```
+
+Decode 阶段 100% memory-bound：每步读取全部 16 GB 权重（252 个 GEMV），计算量可忽略。
+
+## 瓶颈分析
+
+Phase 14 结束时性能 12.9 tok/s = 77.5 ms/step，roofline 利用率仅 12%。
+
+### 量化瓶颈分解
+
+| 来源 | 估计耗时 | 占比 |
+|------|---------|------|
+| cuBLAS M=1 GEMV (252 calls, 带宽利用 ~8%) | ~60 ms | 77% |
+| 非 matmul 内核 (attention, norm, activation, reshape) | ~8 ms | 10% |
+| Tensor 分配 + cudaMemset (1440+ allocs/step) | ~5 ms | 7% |
+| Kernel launch overhead (200+ launches × 5μs) | ~1 ms | 1% |
+| 其他 (sampling CPU round-trip, etc.) | ~3.5 ms | 5% |
+
+**核心发现: cuBLAS 对 M=1 GEMM (GEMV) 的带宽利用率极低（~8%），是 9x gap 的根本原因。**
+
+cuBLAS 设计用于大 M 的 GEMM，对 M=1 场景存在：
+- Kernel launch dispatch overhead 无法被大量计算掩盖
+- TensorCore tile (16×16) 无法被 M=1 充分利用
+- 内部 heuristic 选择了次优算法
+
+## 优化实施
+
+### Opt 1: Decode Attention Kernel
+
+**目标**: 替换 FA2 在 Q_len=1 时的低效路径（64 线程仅 1 个 active）。
+
+**实现** (`csrc/attention/flash_attention.cu`):
+- 专用 decode_attention_bf16_kernel: 256 线程并行沿 KV 序列维度
+- 每个 thread 加载完整 Q vector (128 dim) 到寄存器
+- 处理其分配的 KV 位置块: dot product → online softmax
+- Block-level warp-shuffle + shared memory reduction 合并结果
+- GQA 支持: kv_head = q_head / heads_per_group
+
+**效果**: 在当前短序列 (kv_len ≤ 79) 下效果微小——attention 不是瓶颈。在长序列时会显著受益。
+
+### Opt 2: Fused SiLU×Mul
+
+**目标**: `silu(gate) * up` 两个 element-wise op 合并为一个 kernel。
+
+**实现** (`csrc/activation/activations.cu`):
+```
+Before: read gate → silu → write temp → read temp + up → mul → write out
+After:  read gate + up → silu(gate) * up → write out
+Saved:  1 HBM read + 1 HBM write per element
+```
+
+**效果**: 每层省 1 次 HBM round-trip，36 层总计可观但在 GEMV 瓶颈下被掩盖。
+
+### Opt 3: Fused Add+RMSNorm
+
+**目标**: `x = residual + attn_proj; normed = rmsnorm(x)` 合并为一个 kernel。
+
+**实现** (`csrc/normalization/rmsnorm.cu`):
+```
+Before: read residual + x → add → write sum → read sum + gamma → norm → write out
+After:  read residual + x + gamma → add + norm → write sum + normed
+Saved:  1 full HBM round-trip per attention block
+```
+
+### Opt 4: Batched Decode Forward ⭐
+
+**目标**: 多序列 decode token 合并为 M=batch_size 的 GEMM，提升 cuBLAS 效率。
+
+**实现** (`crates/xserv-model/src/qwen3.rs` + `crates/xserv-server/src/engine.rs`):
+- 新增 `Qwen3::forward_decode_batch(tokens, positions, caches)`
+- Batched ops: embedding, norm, projections, FFN — [B, hidden] × [hidden, X]
+- Per-seq ops: RoPE, KV cache, attention（各序列位置/长度不同）
+- Row extraction (`row_view`) + concatenation (`concat_rows`) 在 batched/per-seq 间切换
+- Engine Step 4b: batch≥2 时自动使用 batched decode
+
+**效果**: batch=4 时 cuBLAS 从 1008× M=1 → 252× M=4，吞吐 35.1 tok/s (vs serial 13.2)。
+
+### Opt 5: Custom GEMV Kernel ⭐⭐⭐ (决定性优化)
+
+**目标**: 替换 cuBLAS 的 M=1 GEMV，手写带宽最优化 kernel。
+
+**实现** (`csrc/gemm/gemv.cu`):
+```
+设计: K-split tiled GEMV
+- TILE_N = 128 (output columns per block, one thread per column)
+- TILE_K = 256 (K-dimension slice per block)
+- BLOCK_SIZE = 128 threads
+- Grid: (ceil(N/128), ceil(K/256)) — 对 K=N=4096 得到 512 blocks
+  512 blocks / 170 SMs ≈ 3 blocks/SM (良好 occupancy)
+
+内存访问:
+- 相邻线程读 W 矩阵的相邻列 → 完美 coalesced
+- x vector 加载到 shared memory (每 K-chunk 仅加载一次)
+- FP32 accumulation via atomicAdd (K-split partial sums)
+- 独立 kernel 做 FP32→BF16 转换
+
+调度:
+- matmul() 中检测 M==1 && dtype==BF16 → 自动使用 custom GEMV
+- M>1 保持 cuBLAS
+```
+
+**效果**: 13.2 → 46.6 tok/s (+253%)。带宽利用率从 ~8% 提升到 ~42%。
+
+### Opt 6: Tensor::empty() (消除无用 cudaMemset)
+
+**目标**: kernel 输出 tensor 全量覆写时，跳过分配后的 cudaMemset 清零。
+
+**实现**:
+- `Storage::empty()` + `Tensor::empty()`: 分配不清零
+- 21 个 kernel wrapper (activation, attention, embedding, gemm, norm, softmax, transpose) 从 `zeros` 改为 `empty`
+- GEMV FP32 accumulator buffer 保持 `cudaMemsetAsync`（atomicAdd 需要零初始化）
+
+**效果**: 46.6 → 50.3 tok/s (+8%)。消除 ~756 个 cudaMemset/step。
+
+### Infra: CUDA Graph 基础设施
+
+- FFI bindings: `cudaStreamBeginCapture`, `cudaGraphInstantiate`, `cudaGraphLaunch`
+- RAII wrapper: `CudaGraph` (capture/instantiate/launch lifecycle)
+- 当前未在 forward path 使用（variable kv_len 限制），为后续优化预留
+
+## Ablation 结果
+
+dash5, RTX 5090, Qwen3-8B BF16, greedy decode, max_tokens=64:
+
+| 优化叠加 | tok/s | 增量 | vs HF | Roofline |
+|---------|-------|------|-------|----------|
+| Phase 14 baseline (FA2) | 12.9 | — | 36% | 12% |
+| + Decode attention | 12.9 | +0% | 36% | 12% |
+| + Fused SiLU×Mul | 13.0 | +1% | 36% | 12% |
+| + Fused Add+RMSNorm | 13.2 | +2% | 37% | 12% |
+| + Batched decode (batch=4) | 35.1 | — | 97% | — |
+| + Custom GEMV (M=1) | 46.6 | +253% | 130% | 42% |
+| + Tensor::empty | **50.3** | +8% | **140%** | **45%** |
+
+对比:
+
+| 系统 | tok/s | Roofline |
+|------|-------|----------|
+| HF transformers | 36.0 | 32% |
+| **xserv (Phase 15)** | **50.3** | **45%** |
+| 理论极限 (1.79 TB/s) | 112.0 | 100% |
+
+## 剩余 55% Roofline Gap 分析
+
+| 来源 | 估计占比 | 优化方向 |
+|------|---------|---------|
+| GEMV kernel 非满带宽 (atomicAdd contention, K-split overhead) | 25% | 无 K-split GEMV (更大 block), 向量化加载 |
+| Non-matmul kernels (attention, norm, RoPE, reshape) | 15% | Fused layer kernel, 更高效的 decode attention |
+| Kernel launch overhead (200+ launches/step) | 5% | CUDA Graphs (需解决 variable kv_len) |
+| Memory allocator overhead (Arc, SmallVec per tensor) | 5% | Pre-allocated decode workspace |
+| Sampling D2H copy (pipeline stall) | 3% | GPU-side argmax kernel |
+| 其他 (host-side logic, channel overhead) | 2% | — |
+
+## 下一步
+
+Phase 15 的 Milestone ④ 目标 (50% of HF) 已远超 — 达到 140% of HF, 45% of roofline。
+
+后续优化路径（按 ROI 排序）：
+1. **无 K-split GEMV**: 消除 atomicAdd，减少 kernel launches → 预期 +15-20%
+2. **向量化 GEMV loads**: float4 加载 W 矩阵 → 预期 +10%
+3. **Pre-allocated workspace**: 消除 Tensor 对象分配开销 → 预期 +5%
+4. **CUDA Graphs**: 需要 fixed-shape decode path → 预期 +5%
+5. **GPU-side sampling**: 消除 logits D2H pipeline stall → 预期 +3%

docs/TO-BE-FIXED.md

@@ -0,0 +1,214 @@
+# xserv — To Be Fixed (2026-05-23 审查更新)
+
+> 由全面审查产出的修复清单。每项修复有明确验收标准。
+> 优先级: P0 (阻塞可用性) > P1 (严重bug/性能) > P2 (重要改进) > P3 (设计债务)
+
+---
+
+## 第一批：P0 — 阻塞可用性
+
+### FIX-01: 全局 cuBLAS handle [P0-性能] ❌未修
+
+**问题**: `gemm.rs` 中 `matmul` (line 146) 和 `batched_matmul` (line 224) 每次调用都 `CublasContext::new()` 创建+销毁 handle。Qwen3-8B 一次 forward ~252 次 matmul。
+
+**修复要求**:
+- 使用 thread-local 单例 cuBLAS handle
+- handle 生命周期覆盖整个进程
+- `matmul` / `batched_matmul` 函数体内不再有 `CublasContext::new()`
+
+**验收标准**:
+1. `grep -n "CublasContext::new" crates/xserv-kernels/src/gemm.rs` 只出现 1 次（thread_local 初始化处）
+2. 编译通过，现有 gemm_test 全部通过
+
+---
+
+### FIX-16: EOS token 泄漏到 API 响应 [P0-功能] ❌新发现
+
+**问题**: `engine.rs:218` 中 `emit_token` 先发 `GenerateEvent::Token { text: "<|im_end|>" }` 再发 `Done`。`api.rs:110-111` 把所有 Token text 拼到 content 里，导致最终响应包含 `<|im_end|>` 乱码。
+
+**修复要求**:
+- `emit_token` 中，当 token 是 EOS 时，不发送 Token event（或发送空 text），直接发 Done
+- 或者: API 层收到 Done 时丢弃最后一个 token 的 text（如果 finish_reason == "stop"）
+
+**验收标准**:
+1. 发送请求，响应 content 不包含 `<|im_end|>` 或其他 special token 文本
+2. streaming 模式下最后一个 content chunk 不是 EOS 文本
+3. 编译通过
+
+---
+
+### FIX-17: max_seq_len 硬编码 256 [P0-功能] ❌新发现
+
+**问题**: `engine.rs:53` 硬编码 `let max_seq_len = 256`，超过就 KV cache panic。
+
+**修复要求**:
+- `Engine::load` 接受 `max_seq_len` 参数（或从 config 读取，上限为 config.max_seq_len()）
+- `main.rs` 中通过命令行参数或环境变量传入，默认值改为 2048
+- 同步更新 RoPE cache 上限（当前 `qwen3.rs:45` 限制 8192，应与 max_seq_len 一致）
+
+**验收标准**:
+1. `grep -n "let max_seq_len = 256" crates/xserv-server/` 返回 0 行
+2. 启动 server 时 `--max-seq-len 4096` 可用
+3. 编译通过
+
+---
+
+### FIX-18: max_tokens 无上限校验 [P0-功能] ❌新发现
+
+**问题**: API 不校验 `max_tokens`，客户端可发 `max_tokens: 1000000` 导致 KV cache panic。
+
+**修复要求**:
+- `api.rs` 中 clamp `max_tokens` 到 `engine.max_seq_len - prompt_tokens.len()`
+- 如果 prompt 已超过 max_seq_len，返回 400 错误
+
+**验收标准**:
+1. 发送 `max_tokens: 999999`，不 panic，正常生成到 seq_len 上限
+2. 发送超长 prompt（> max_seq_len），返回 HTTP 400
+3. 编译通过
+
+---
+
+## 第二批：P1 — 严重 bug/性能
+
+### FIX-07: 使用 CachingAllocator [P1-性能] ❌未修
+
+**问题**: `CachingAllocator` 已实现（`allocator.rs`）但从未使用。所有 GPU 分配直接 `cudaMalloc`。
+
+**修复要求**:
+- `Tensor::empty` 对 GPU device 使用 `cached_alloc` 而非 `GpuBuffer::alloc`
+- `GpuBuffer::Drop` 调用 `cached_dealloc` 归还到池（而非 `cudaFree`）
+- 或者更简单：在 `GpuBuffer::alloc` 内部接入 caching allocator（全局透明替换）
+
+**验收标准**:
+1. 连续运行 10 次 decode step，`cudaMalloc` 调用次数应显著低于总分配次数
+2. 编译通过，现有测试通过
+3. 推理结果与修复前一致
+
+---
+
+### FIX-08: CudaDeviceProp FFI 安全性 [P1-Bug] ❌未修
+
+**问题**: `ffi.rs:31` 用 `_pad: [u8; 4096]` 猜测 `cudaDeviceProp` struct 大小，CUDA 12.9 可能更大。
+
+**修复要求**:
+- 增大 pad 到 `[u8; 8192]` 或使用 `cudaDeviceGetAttribute` 替代 name 查询
+- 可参考 `device.rs` 中已有的 `cudaDeviceGetAttribute` 用法
+
+**验收标准**:
+1. `device_info()` 返回正确的 device name
+2. 编译通过
+
+---
+
+### FIX-09: Tokenizer byte_fallback panic [P1-Bug] ❌未修
+
+**问题**: `bpe.rs:176-182` 中 Qwen3 tokenizer 遇到不在 vocab 的单字节时 panic。
+
+**修复要求**:
+- 当 `byte_fallback == true` 且单字节不在 vocab 时，查找 `<0xNN>` 格式 token
+- 如果 `<0xNN>` 也不存在，返回 unk_token_id（而非 panic）
+
+**验收标准**:
+1. 包含所有 256 个字节值的字符串可以 encode 不 panic
+2. 编译通过
+
+---
+
+### FIX-19: 因果掩码 -1e9 应改为 -inf [P1-Bug] ❌新发现
+
+**问题**: `csrc/attention/causal_mask.cu:31` 用 `-1e9f` 代替 `-inf`，注释说 "BF16 没有 -inf" 但这是错误的。
+
+**修复要求**:
+- BF16 路径改为 `__float2bfloat16(-INFINITY)`
+- F32 路径改为 `-INFINITY`（如果还没有的话）
+
+**验收标准**:
+1. causal mask 中被遮蔽的值为 `-inf`（而非 `-1e9`）
+2. 编译通过，attention test 通过
+
+---
+
+### FIX-20: LayerNorm 数值稳定性 [P1-Bug] ❌新发现
+
+**问题**: `csrc/normalization/layernorm.cu:19-25` 注释写 "Welford online" 但实际用 `E[x²] - E[x]²`，大均值小方差时会灾难性抵消。
+
+**修复要求**:
+- 改为真正的 two-pass 或 Welford online 算法
+- pass 1: 求 mean; pass 2: 求 variance = E[(x-mean)²]
+
+**验收标准**:
+1. 对 mean=1e6, std=1e-3 的输入，layernorm 输出与 PyTorch 一致（relative error < 1e-3）
+2. 编译通过，现有测试通过
+
+---
+
+### FIX-21: LayerNorm/RMSNorm 最小 block size [P1-Bug] ❌新发现
+
+**问题**: `layernorm.cu:88` 和 `rmsnorm.cu` 对 hidden_size < 32 的输入会崩溃（block_reduce 需要至少一个完整 warp）。
+
+**修复要求**:
+- launch 时 `block = max(min(hidden_size, 1024), 32)`
+
+**验收标准**:
+1. hidden_size=16 的 layernorm/rmsnorm 不崩溃
+2. 编译通过
+
+---
+
+## 第三批：P2 — 重要改进
+
+### FIX-22: Engine dummy KV cache 分配 [P2-性能] ❌新发现
+
+**问题**: `engine.rs:142-148` 每次 batched decode 用 `std::mem::replace` 创建 dummy `GpuKVCache::new(..., 1, ...)` 来绕过 borrow checker，每步分配 `num_layers * 2` 个 GPU buffer。
+
+**修复要求**:
+- 将 `running` 从 `Vec<Sequence>` 改为存储方式让 KV cache 可以独立借出
+- 或使用 `Option<GpuKVCache>` + `.take()` / `.insert()` 避免 dummy 分配
+
+**验收标准**:
+1. batched decode 路径不再分配 dummy KV cache
+2. 编译通过，功能不变
+
+---
+
+### FIX-23: RoPE cache 硬限 8192 [P2-功能] ❌新发现
+
+**问题**: `qwen3.rs:45` `config.max_seq_len().min(8192)` 人为截断。
+
+**修复要求**:
+- 去掉 `.min(8192)`，或改为与 engine 的 max_seq_len 一致
+- 确保 RoPE cache 覆盖实际使用的 max_seq_len
+
+**验收标准**:
+1. RoPE cache 长度 >= engine max_seq_len
+2. 编译通过
+
+---
+
+### FIX-15: GPT-2 消除 CPU round-trip [P3-性能] ❌未修
+
+**问题**: GPT-2 `split_qkv`、`merge_heads`、`add_bias` 全在 CPU 做。优先级低（GPT-2 不是主力模型）。
+
+---
+
+## 修复依赖图和执行顺序
+
+```
+第一批 P0 (可并行):
+  FIX-01 (cuBLAS handle)     ← 独立
+  FIX-16 (EOS 泄漏)          ← 独立
+  FIX-17 (max_seq_len)       ← 独立，FIX-23 依赖此
+  FIX-18 (max_tokens 校验)   ← 依赖 FIX-17（需要知道 max_seq_len）
+
+第二批 P1 (可并行):
+  FIX-07 (caching allocator) ← 独立
+  FIX-08 (CudaDeviceProp)    ← 独立
+  FIX-09 (byte_fallback)     ← 独立
+  FIX-19 (causal mask -inf)  ← 独立
+  FIX-20 (layernorm 稳定性)  ← 独立
+  FIX-21 (min block size)    ← 独立
+
+第三批 P2:
+  FIX-22 (dummy KV cache)    ← 独立
+  FIX-23 (RoPE cache)        ← 依赖 FIX-17
+```

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/phase14-flash-attention.md

@@ -0,0 +1,109 @@
+# Phase 14 Benchmark: Flash Attention 2
+
+**Date**: 2026-05-22
+**Hardware**: RTX 5090 (32GB GDDR7, SM120 CC 12.0, 170 SMs)
+**Model**: Qwen3-8B (BF16, 36 layers, 4096 hidden, 32 Q / 8 KV GQA heads, head_dim=128)
+**Config**: greedy decoding (temperature=0), max_tokens=64, single-request serial
+
+## Correctness
+
+Logits comparison with HuggingFace transformers (10 prompts, raw text without ChatML):
+
+| Metric | Result |
+|--------|--------|
+| Prefill Top-1 match vs HF | **9/10 (90%)** |
+| Avg Top-5 overlap vs HF | **4.0/5** |
+| Result vs pre-FA2 naive attention | **Identical** (same 9/10 top-1, same 4.0/5 overlap) |
+
+The single top-1 mismatch ("Explain quantum computing.") has logits differing by 0.125
+(22.000 vs 21.875) — within BF16 precision. The top-5 sets are identical (5/5 overlap).
+
+FA2 introduces no precision degradation compared to the naive attention path.
+
+## API Generation
+
+52 diverse prompts (English, Chinese, code) via `/v1/chat/completions`:
+
+| Metric | Result |
+|--------|--------|
+| Success rate | **52/52 (100%)** |
+| SSE streaming | **Working** (role chunk, content chunks, finish_reason, [DONE]) |
+| Usage stats | Correct (prompt_tokens + completion_tokens = total_tokens) |
+
+## Performance
+
+### xserv vs HuggingFace transformers
+
+8 prompts (short/medium/long) × max_tokens=64, greedy:
+
+| Category | Prompt Tokens | xserv (tok/s) | HF (tok/s) | Ratio |
+|----------|--------------|---------------|------------|-------|
+| Short (~12 tok) | 12-14 | 12.5 | 38.5 | 0.32x |
+| Medium (~28 tok) | 27-28 | 13.6 | 44.1 | 0.31x |
+| Long (~60 tok) | 58-64 | 13.0 | 36.0 | 0.36x |
+| **Overall** | — | **12.9** | **36.6** | **0.35x** |
+
+### Phase-over-Phase Improvement
+
+| Phase | Attention | repeat_kv | tok/s | vs HF |
+|-------|-----------|-----------|-------|-------|
+| 10 | Naive (O(S²), cuBLAS batched) | CPU round-trip | 6.9 | 15% |
+| 11 | Naive + GPU KV cache | GPU repeat_kv | 10.3 | 30% |
+| **14** | **FA2 (O(1), fused kernel)** | **None (GQA in kernel)** | **12.9** | **35%** |
+
+Phase 14 vs Phase 11: **+25% throughput** (10.3 → 12.9 tok/s).
+
+### Improvement Breakdown (estimated)
+
+| Factor | Contribution |
+|--------|-------------|
+| Eliminating repeat_kv GPU alloc + copy (per layer) | ~10% |
+| Eliminating K^T transpose + contiguous | ~5% |
+| Eliminating S×S score matrix alloc | ~5% |
+| Fused kernel (1 launch vs 6) | ~5% |
+
+### Concurrent Requests
+
+8 concurrent requests, max_batch=4:
+
+| Metric | Result |
+|--------|--------|
+| Wall clock | 22.5s |
+| Sum of individual latencies | 135.0s |
+| Scheduling speedup | **6.0x** |
+| Throughput | 11.4 tok/s |
+
+Continuous batching scheduling confirmed working (decode batch_size=4 in logs).
+
+## Remaining Performance Gap
+
+35% of HF throughput. Main bottlenecks:
+
+| Bottleneck | Impact | Fix |
+|-----------|--------|-----|
+| **Decode Q_len=1 inefficiency** | FA2 kernel: 64 threads, only 1 active (owns_row=true for single query) | Specialized decode attention kernel (vector-dot against KV, parallel reduction along S) |
+| **No kernel fusion** | RMSNorm+residual, SiLU*up: separate kernels, redundant HBM reads/writes | Fused kernels (Phase 15) |
+| **No CUDA Graphs** | ~100+ kernel launches per decode step, each has host-side overhead | Capture decode iteration as CUDA Graph (Phase 15) |
+| **Per-seq forward (no batched decode)** | With batch=4, 4 serial forward passes per iteration | Batched projections + per-seq attention (Phase 15, depends on FA2 decode kernel) |
+| **No vectorized loads in FA2** | Scalar bf16→f32 conversion in dot product loop | float4 / bfloat162 vectorized loads |
+
+## Memory Usage
+
+| Component | Naive (Phase 11) | FA2 (Phase 14) |
+|-----------|-----------------|----------------|
+| Score matrix [1, 32, S, S] | S² × 32 × 2B | **0** |
+| repeat_kv K/V [1, 32, S, 128] | 2 × S × 32 × 128 × 2B per layer | **0** |
+| K^T contiguous copy | S × 32 × 128 × 2B per layer | **0** |
+
+For S=256 (current max): savings ~6 MB per layer × 36 layers ≈ 216 MB.
+For S=2048: savings ~384 MB per layer × 36 layers ≈ 13.5 GB (naive would OOM).
+
+## Tracking
+
+| Phase | Attention | tok/s | vs HF | Correctness |
+|-------|-----------|-------|-------|-------------|
+| 8 | Naive (no cache) | 2.5 | 5% | 50/50 vs HF |
+| 9 | Naive + CPU KV cache | 44.3 (GPT-2) | — | 50/50 self |
+| 10 | Naive + CPU KV cache | 6.9 (Qwen3-8B) | 15% | 100% top-5 |
+| 11 | Naive + GPU KV cache | 10.3 | 30% | 9/10 top-1 |
+| **14** | **FA2 + GQA in kernel** | **12.9** | **35%** | **9/10 top-1** |

docs/benchmarks/phase15-performance.md

@@ -0,0 +1,85 @@
+# Phase 15 Benchmark: Performance Optimization
+
+**Date**: 2026-05-23
+**Hardware**: RTX 5090 (32GB GDDR7, SM120 CC 12.0, 170 SMs, 1.79 TB/s)
+**Model**: Qwen3-8B (BF16, 36 layers, 4096 hidden, 32 Q / 8 KV GQA heads, head_dim=128)
+**Config**: greedy decoding (temperature=0), max_tokens=64, serial (batch=1)
+
+## Ablation: Each Optimization Measured Independently
+
+| # | Optimization | tok/s | Delta | ms/token | Roofline |
+|---|-------------|-------|-------|----------|----------|
+| 0 | Phase 14 baseline (FA2 + naive cuBLAS GEMV) | 12.9 | — | 77.5 | 12% |
+| 1 | + Decode attention kernel (256 threads) | 12.9 | +0% | 77.5 | 12% |
+| 2 | + Fused SiLU×Mul | 13.0 | +1% | 76.9 | 12% |
+| 3 | + Fused Add+RMSNorm | 13.2 | +2% | 75.8 | 12% |
+| 4 | + Custom GEMV (M=1, K-split tiled) | 46.6 | +253% | 21.5 | 42% |
+| 5 | + Tensor::empty (skip cudaMemset) | **50.3** | **+8%** | **19.9** | **45%** |
+
+## Comparison with HuggingFace transformers
+
+8 prompts (short/medium/long) × max_tokens=64, greedy, serial:
+
+| System | tok/s | ms/token | Roofline |
+|--------|-------|----------|----------|
+| HF transformers (BF16, torch 2.8, SDPA) | 36.0 | 27.8 | 32% |
+| **xserv Phase 15** | **50.3** | **19.9** | **45%** |
+| Roofline (1.79 TB/s, 16GB model) | 112.0 | 8.9 | 100% |
+
+**xserv is 140% of HF transformers throughput.**
+
+## Per-Prompt Detail (Phase 15 Final)
+
+| # | Prompt | pt | ct | Time | tok/s |
+|---|--------|----|----|------|-------|
+| 1 | What is gravity? | 12 | 64 | 1.39s | 46.0 |
+| 2 | Hello, how are you? | 14 | 64 | 1.27s | 50.5 |
+| 3 | Explain DNA briefly. | 13 | 64 | 1.25s | 51.2 |
+| 4 | Write a detailed explanation of photosynthesis... | 27 | 64 | 1.26s | 50.7 |
+| 5 | Describe machine learning. | 13 | 64 | 1.25s | 51.2 |
+| 6 | What causes earthquakes? | 12 | 64 | 1.25s | 51.1 |
+| 7 | How does the internet work? | 14 | 64 | 1.25s | 51.1 |
+| 8 | What is the speed of light? | 15 | 64 | 1.25s | 51.0 |
+
+Prompt 1 is slower (46.0 vs 51.x) due to first-request warmup (caching allocator cold start).
+
+## Concurrent Throughput
+
+8 requests concurrent, max_batch=4:
+
+| Config | tok/s | Wall clock | Speedup |
+|--------|-------|-----------|---------|
+| Serial (batch=1, custom GEMV) | 50.3 | — | — |
+| Concurrent (batch=4, cuBLAS M=4) | 28.2 | 9.09s | 6.47x scheduling |
+| Concurrent (batch=4, custom GEMV) | 35.1* | ~7.3s | ~6x scheduling |
+
+*Note: batch=4 with custom GEMV is slower than serial because:
+1. Batched decode path uses cuBLAS for M>1 matmuls, losing the GEMV advantage
+2. Per-seq attention/reshape overhead in the batched path adds ~2ms/step
+3. Custom GEMV already saturates bandwidth at M=1
+
+Serial decode with custom GEMV is the optimal path for current architecture.
+
+## Correctness Verification
+
+| Test | Result |
+|------|--------|
+| Top-1 logits match vs HF (10 prompts) | 9/10 (90%) |
+| Top-5 overlap vs HF (10 prompts) | 4.0/5 avg |
+| vs pre-optimization baseline | Identical (same 9/10) |
+| API generation (52 prompts) | 52/52 pass |
+| SSE streaming | Working |
+| Chinese prompts | Working |
+
+## Phase-over-Phase Performance Tracking
+
+| Phase | Key Change | tok/s | vs HF | Roofline |
+|-------|-----------|-------|-------|----------|
+| 8 | GPT-2 inference (no cache) | 2.5 | 7% | — |
+| 9 | + KV cache (CPU) | 44.3 (GPT-2) | — | — |
+| 10 | Qwen3-8B (CPU KV cache) | 6.9 | 19% | 6% |
+| 11 | + GPU KV cache | 10.3 | 29% | 9% |
+| 14 | + Flash Attention 2 | 12.9 | 36% | 12% |
+| **15** | **+ Custom GEMV + fused + empty** | **50.3** | **140%** | **45%** |
+
+Total speedup from Phase 10 to Phase 15: **7.3x** (6.9 → 50.3 tok/s).

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/bench_server.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+"""Benchmark xserv server performance and check correctness vs HF."""
+
+import json
+import time
+import sys
+import urllib.request
+
+PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 8090
+
+def chat(prompt, max_tokens=80, temperature=0):
+    data = json.dumps({
+        "model": "qwen3-8b",
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "stream": False
+    }).encode()
+    req = urllib.request.Request(
+        f"http://localhost:{PORT}/v1/chat/completions",
+        data=data, headers={"Content-Type": "application/json"}
+    )
+    t0 = time.perf_counter()
+    with urllib.request.urlopen(req, timeout=180) as resp:
+        result = json.loads(resp.read())
+    elapsed = time.perf_counter() - t0
+    usage = result.get("usage", {})
+    content = result["choices"][0]["message"]["content"]
+    finish = result["choices"][0]["finish_reason"]
+    ct = usage.get("completion_tokens", 0)
+    pt = usage.get("prompt_tokens", 0)
+    return ct / elapsed if elapsed > 0 else 0, elapsed, ct, pt, content, finish
+
+
+def chat_stream(prompt, max_tokens=80, temperature=0):
+    data = json.dumps({
+        "model": "qwen3-8b",
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "stream": True
+    }).encode()
+    req = urllib.request.Request(
+        f"http://localhost:{PORT}/v1/chat/completions",
+        data=data, headers={"Content-Type": "application/json"}
+    )
+    t0 = time.perf_counter()
+    tokens = 0
+    content = ""
+    with urllib.request.urlopen(req, timeout=180) as resp:
+        for line in resp:
+            line = line.decode().strip()
+            if line.startswith("data: "):
+                payload = line[6:]
+                if payload == "[DONE]":
+                    break
+                chunk = json.loads(payload)
+                delta = chunk["choices"][0].get("delta", {})
+                c = delta.get("content", "")
+                if c:
+                    tokens += 1
+                    content += c
+    elapsed = time.perf_counter() - t0
+    return tokens / elapsed if elapsed > 0 else 0, elapsed, tokens, content
+
+
+print("=" * 60)
+print(f"xserv Server Benchmark (port {PORT})")
+print("=" * 60)
+
+# Health check
+try:
+    urllib.request.urlopen(f"http://localhost:{PORT}/health", timeout=3)
+except:
+    print(f"Server not responding on port {PORT}")
+    sys.exit(1)
+
+# 1. EOS leak check
+print("\n--- EOS Leak Check ---")
+tps, t, ct, pt, content, finish = chat("Say hello", 30)
+has_eos = "<|im_end|>" in content or "<|endoftext|>" in content or "<|im_start|>" in content
+print(f"  finish_reason: {finish}")
+print(f"  EOS in content: {'YES (BUG!)' if has_eos else 'NO (good)'}")
+print(f"  Content: {content[:100]}")
+
+# 2. Warmup
+print("\n--- Warmup ---")
+chat("Hi", 10)
+time.sleep(0.5)
+
+# 3. Non-streaming benchmark
+print("\n--- Non-streaming Performance (greedy, batch=1) ---")
+prompts = [
+    ("short", "What is 2+2?", 50),
+    ("medium", "Explain quantum computing in simple terms.", 80),
+    ("long", "Write a detailed comparison of Python and Rust programming languages, covering syntax, performance, memory management, and ecosystem.", 150),
+]
+
+for name, prompt, max_tok in prompts:
+    tps, t, ct, pt, content, finish = chat(prompt, max_tok)
+    print(f"  [{name}] {tps:.1f} tok/s | {ct} tokens in {t:.2f}s | prompt={pt} | finish={finish}")
+
+# 4. Streaming benchmark
+print("\n--- Streaming Performance ---")
+tps, t, ct, content = chat_stream("Explain the theory of relativity.", 80)
+print(f"  stream: {tps:.1f} tok/s | {ct} tokens in {t:.2f}s")
+
+# 5. max_tokens validation
+print("\n--- max_tokens Validation ---")
+try:
+    tps, t, ct, pt, content, finish = chat("Hi", 999999)
+    print(f"  max_tokens=999999: OK (server clamped to {ct} tokens, no crash)")
+except Exception as e:
+    print(f"  max_tokens=999999: {e}")
+
+# 6. Concurrent requests (if server supports batching)
+print("\n--- Concurrent Requests (2 parallel) ---")
+import threading
+results = [None, None]
+
+def do_request(idx, prompt, max_tok):
+    results[idx] = chat(prompt, max_tok)
+
+t1 = threading.Thread(target=do_request, args=(0, "What is gravity?", 50))
+t2 = threading.Thread(target=do_request, args=(1, "What is light?", 50))
+t0 = time.perf_counter()
+t1.start(); t2.start()
+t1.join(); t2.join()
+wall_time = time.perf_counter() - t0
+
+total_tokens = sum(r[2] for r in results if r)
+combined_tps = total_tokens / wall_time
+print(f"  2 concurrent: {combined_tps:.1f} tok/s total | wall={wall_time:.2f}s")
+for i, r in enumerate(results):
+    if r:
+        print(f"    req{i}: {r[0]:.1f} tok/s, {r[2]} tokens in {r[1]:.2f}s")
+
+print(f"\n{'=' * 60}")
+print("DONE")
+print("=" * 60)

tools/bench_vs_hf.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python3
+"""
+Benchmark xserv vs HuggingFace transformers on Qwen3-8B.
+Measures: prefill latency, decode throughput, end-to-end latency.
+
+Usage:
+    # xserv server should be running on port 9090
+    python3 tools/bench_vs_hf.py
+"""
+
+import json
+import os
+import time
+import urllib.request
+
+MODEL_DIR = "/opt/wjh/models/qwen3-8b"
+XSERV_URL = "http://localhost:9090"
+
+BENCH_PROMPTS = [
+    # Short prompts (~10 tokens)
+    ("short", "What is gravity?"),
+    ("short", "Hello, how are you?"),
+    ("short", "Explain DNA briefly."),
+    # Medium prompts (~30 tokens)
+    ("medium", "Write a detailed explanation of how photosynthesis works in plants, including the light and dark reactions."),
+    ("medium", "Describe the process of machine learning training, including forward pass, loss computation, and backpropagation."),
+    ("medium", "Explain the differences between TCP and UDP protocols, including when you would use each one in practice."),
+    # Longer prompts (~60 tokens)
+    ("long", "You are an expert computer scientist. Please write a comprehensive explanation of how modern GPUs work, including the architecture of streaming multiprocessors, the memory hierarchy from registers to global memory, and how thousands of threads are scheduled concurrently. Include specific technical details."),
+    ("long", "You are a historian specializing in ancient civilizations. Please provide a detailed analysis of the rise and fall of the Roman Empire, covering the key factors that led to its expansion, the political and social structures that sustained it, and the multiple causes that contributed to its eventual decline and collapse."),
+]
+
+MAX_TOKENS = 64
+
+
+def bench_xserv():
+    """Benchmark xserv HTTP API."""
+    print("\n" + "=" * 60)
+    print("BENCHMARK: xserv (HTTP API, greedy, max_tokens={})".format(MAX_TOKENS))
+    print("=" * 60)
+
+    # Warmup
+    body = json.dumps({
+        "model": "qwen3-8b",
+        "messages": [{"role": "user", "content": "Hi"}],
+        "max_tokens": 8,
+        "temperature": 0.0,
+    }).encode()
+    req = urllib.request.Request(
+        f"{XSERV_URL}/v1/chat/completions",
+        data=body, headers={"Content-Type": "application/json"},
+    )
+    urllib.request.urlopen(req, timeout=120)
+    print("Warmup done.\n")
+
+    results = []
+    for category, prompt in BENCH_PROMPTS:
+        body = json.dumps({
+            "model": "qwen3-8b",
+            "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": MAX_TOKENS,
+            "temperature": 0.0,
+        }).encode()
+        req = urllib.request.Request(
+            f"{XSERV_URL}/v1/chat/completions",
+            data=body, headers={"Content-Type": "application/json"},
+        )
+
+        t0 = time.perf_counter()
+        resp = urllib.request.urlopen(req, timeout=300)
+        elapsed = time.perf_counter() - t0
+        data = json.loads(resp.read())
+
+        usage = data.get("usage", {})
+        pt = usage.get("prompt_tokens", 0)
+        ct = usage.get("completion_tokens", 0)
+        tok_per_sec = ct / elapsed if elapsed > 0 else 0
+
+        print(f"  [{category:>6}] pt={pt:3d} ct={ct:2d} | {elapsed:6.2f}s | {tok_per_sec:5.1f} tok/s | {prompt[:50]}...")
+        results.append({
+            "category": category,
+            "prompt_tokens": pt,
+            "completion_tokens": ct,
+            "elapsed": elapsed,
+            "tok_per_sec": tok_per_sec,
+        })
+
+    # Summary
+    total_ct = sum(r["completion_tokens"] for r in results)
+    total_time = sum(r["elapsed"] for r in results)
+    avg_tok_per_sec = total_ct / total_time if total_time > 0 else 0
+
+    print(f"\n  xserv total: {total_ct} tokens in {total_time:.2f}s = {avg_tok_per_sec:.1f} tok/s")
+    return results
+
+
+def bench_hf():
+    """Benchmark HuggingFace transformers generate()."""
+    print("\n" + "=" * 60)
+    print("BENCHMARK: HuggingFace transformers (greedy, max_new_tokens={})".format(MAX_TOKENS))
+    print("=" * 60)
+
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    print(f"Loading model on GPU 1...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_DIR, dtype=torch.bfloat16, device_map="cuda:1", trust_remote_code=True)
+    model.eval()
+    print("Model loaded.\n")
+
+    # Warmup
+    inputs = tokenizer("Hi", return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        model.generate(**inputs, max_new_tokens=8, do_sample=False)
+    print("Warmup done.\n")
+
+    results = []
+    for category, prompt in BENCH_PROMPTS:
+        # Apply chat template (same as xserv)
+        messages = [{"role": "user", "content": prompt}]
+        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(text, return_tensors="pt").to(model.device)
+        pt = inputs["input_ids"].shape[1]
+
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        with torch.no_grad():
+            output = model.generate(
+                **inputs,
+                max_new_tokens=MAX_TOKENS,
+                do_sample=False,
+            )
+        torch.cuda.synchronize()
+        elapsed = time.perf_counter() - t0
+
+        ct = output.shape[1] - pt
+        tok_per_sec = ct / elapsed if elapsed > 0 else 0
+
+        print(f"  [{category:>6}] pt={pt:3d} ct={ct:2d} | {elapsed:6.2f}s | {tok_per_sec:5.1f} tok/s | {prompt[:50]}...")
+        results.append({
+            "category": category,
+            "prompt_tokens": pt,
+            "completion_tokens": ct,
+            "elapsed": elapsed,
+            "tok_per_sec": tok_per_sec,
+        })
+
+    total_ct = sum(r["completion_tokens"] for r in results)
+    total_time = sum(r["elapsed"] for r in results)
+    avg_tok_per_sec = total_ct / total_time if total_time > 0 else 0
+
+    print(f"\n  HF total: {total_ct} tokens in {total_time:.2f}s = {avg_tok_per_sec:.1f} tok/s")
+
+    del model
+    torch.cuda.empty_cache()
+    return results
+
+
+def main():
+    xserv_results = bench_xserv()
+    hf_results = bench_hf()
+
+    print("\n" + "=" * 60)
+    print("COMPARISON SUMMARY")
+    print("=" * 60)
+
+    print(f"\n{'Category':<10} {'Metric':<20} {'xserv':>10} {'HF':>10} {'Ratio':>10}")
+    print("-" * 62)
+
+    for cat in ["short", "medium", "long"]:
+        xs = [r for r in xserv_results if r["category"] == cat]
+        hf = [r for r in hf_results if r["category"] == cat]
+        if xs and hf:
+            xs_avg_tps = sum(r["tok_per_sec"] for r in xs) / len(xs)
+            hf_avg_tps = sum(r["tok_per_sec"] for r in hf) / len(hf)
+            xs_avg_lat = sum(r["elapsed"] for r in xs) / len(xs)
+            hf_avg_lat = sum(r["elapsed"] for r in hf) / len(hf)
+            ratio_tps = xs_avg_tps / hf_avg_tps if hf_avg_tps > 0 else 0
+            ratio_lat = xs_avg_lat / hf_avg_lat if hf_avg_lat > 0 else 0
+
+            print(f"{cat:<10} {'Throughput (tok/s)':<20} {xs_avg_tps:>10.1f} {hf_avg_tps:>10.1f} {ratio_tps:>9.2f}x")
+            print(f"{'':<10} {'Latency (s)':<20} {xs_avg_lat:>10.2f} {hf_avg_lat:>10.2f} {ratio_lat:>9.2f}x")
+
+    xs_total_tps = sum(r["completion_tokens"] for r in xserv_results) / sum(r["elapsed"] for r in xserv_results)
+    hf_total_tps = sum(r["completion_tokens"] for r in hf_results) / sum(r["elapsed"] for r in hf_results)
+    ratio = xs_total_tps / hf_total_tps if hf_total_tps > 0 else 0
+
+    print("-" * 62)
+    print(f"{'OVERALL':<10} {'Throughput (tok/s)':<20} {xs_total_tps:>10.1f} {hf_total_tps:>10.1f} {ratio:>9.2f}x")
+    print(f"\nxserv is {ratio:.1%} of HF transformers throughput")
+
+
+if __name__ == "__main__":
+    main()

tools/compare_logits.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""Compare xserv prefill logits with HuggingFace transformers on 10 prompts."""
+import os
+import sys
+import subprocess
+import re
+
+MODEL_DIR = "/opt/wjh/models/qwen3-8b"
+TOP_K = 10
+
+PROMPTS = [
+    "What is the capital of France?",
+    "Explain quantum computing.",
+    "Hello world",
+    "def fibonacci(n):",
+    "The weather today is",
+    "1 + 1 =",
+    "Machine learning is",
+    "Once upon a time",
+    "Paris is known for",
+    "How does gravity work?",
+]
+
+
+def get_hf_topk(prompt, tokenizer, model, k=10):
+    import torch
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits[0, -1, :].float().cpu()
+    topk = torch.topk(logits, k)
+    return list(zip(topk.indices.tolist(), topk.values.tolist()))
+
+
+def get_xserv_topk(prompt, k=10):
+    xserv_bin = "/opt/wjh/projects/xserv/target/release/dump-logits"
+    env = {**os.environ, "CUDA_VISIBLE_DEVICES": "0",
+           "PATH": "/usr/local/cuda-12.9/bin:" + os.environ.get("PATH", "")}
+    result = subprocess.run(
+        [xserv_bin, MODEL_DIR, prompt],
+        capture_output=True, text=True, timeout=180, env=env,
+    )
+    # Parse output: "  [ 0] id=  3555 logit=   24.5000 token=..."
+    topk = []
+    for line in result.stdout.strip().split('\n'):
+        m = re.match(r'\s*\[\s*\d+\]\s+id=\s*(\d+)\s+logit=\s*([\d.\-]+)', line)
+        if m:
+            topk.append((int(m.group(1)), float(m.group(2))))
+        if len(topk) >= k:
+            break
+    return topk
+
+
+def main():
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    print(f"Loading HF model on GPU 1...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_DIR, dtype=torch.bfloat16, device_map="cuda:1", trust_remote_code=True)
+    model.eval()
+    print("HF model loaded.\n")
+
+    total = len(PROMPTS)
+    top1_matches = 0
+    top5_overlaps = []
+
+    for i, prompt in enumerate(PROMPTS):
+        print(f"[{i+1}/{total}] \"{prompt}\"")
+
+        hf_top = get_hf_topk(prompt, tokenizer, model, TOP_K)
+        xs_top = get_xserv_topk(prompt, TOP_K)
+
+        if not xs_top:
+            print("  xserv: NO OUTPUT")
+            continue
+
+        hf_ids = [t[0] for t in hf_top]
+        xs_ids = [t[0] for t in xs_top]
+
+        top1_match = hf_ids[0] == xs_ids[0]
+        if top1_match:
+            top1_matches += 1
+
+        top5_overlap = len(set(hf_ids[:5]) & set(xs_ids[:5]))
+        top5_overlaps.append(top5_overlap)
+
+        # Show comparison
+        hf_tok = tokenizer.decode([hf_ids[0]])
+        xs_tok = tokenizer.decode([xs_ids[0]])
+        status = "MATCH" if top1_match else "DIFF"
+
+        print(f"  Top-1: HF={hf_ids[0]:>6}({hf_tok!r:>10}) | xserv={xs_ids[0]:>6}({xs_tok!r:>10})  [{status}]")
+        print(f"  Top-5 overlap: {top5_overlap}/5")
+
+        # Show top-5 side by side
+        print(f"  {'HF':>25} | {'xserv':>25}")
+        for j in range(min(5, len(hf_top), len(xs_top))):
+            h_id, h_val = hf_top[j]
+            x_id, x_val = xs_top[j]
+            h_tok = tokenizer.decode([h_id])
+            x_tok = tokenizer.decode([x_id])
+            print(f"  {h_id:>6} {h_val:>8.3f} {h_tok!r:>8} | {x_id:>6} {x_val:>8.3f} {x_tok!r:>8}")
+        print()
+
+    print("=" * 50)
+    print(f"Top-1 match rate: {top1_matches}/{total} ({100*top1_matches/total:.0f}%)")
+    avg_overlap = sum(top5_overlaps) / max(len(top5_overlaps), 1)
+    print(f"Avg top-5 overlap: {avg_overlap:.1f}/5")
+    print(f"Verdict: {'PASS' if top1_matches >= total * 0.7 else 'FAIL'}")
+
+
+if __name__ == "__main__":
+    main()

tools/e2e_validate.py

@@ -0,0 +1,394 @@
+#!/usr/bin/env python3
+"""
+End-to-end validation for xserv after bug fixes.
+1. Correctness: compare top-k logits with HuggingFace transformers
+2. Generation: run 50+ prompts through the HTTP API
+3. Performance: measure latency and throughput
+
+Usage:
+    # Step 1: Start xserv server in background:
+    #   ./target/release/xserv-server /opt/wjh/models/qwen3-8b --port 8080
+    #
+    # Step 2: Run this script:
+    #   python3 tools/e2e_validate.py --mode all
+    #   python3 tools/e2e_validate.py --mode logits   # correctness only
+    #   python3 tools/e2e_validate.py --mode api      # API + perf only
+"""
+
+import argparse
+import json
+import time
+import subprocess
+import sys
+import os
+from pathlib import Path
+
+MODEL_DIR = "/opt/wjh/models/qwen3-8b"
+XSERV_URL = "http://localhost:8080"
+TOP_K = 10
+
+# 50+ diverse test prompts
+TEST_PROMPTS = [
+    "What is the capital of France?",
+    "Explain quantum computing in simple terms.",
+    "Write a Python function to sort a list.",
+    "你好，请用中文介绍一下你自己。",
+    "What is 2 + 2?",
+    "The theory of relativity states that",
+    "In a far away galaxy,",
+    "def fibonacci(n):",
+    "请解释什么是机器学习。",
+    "How does photosynthesis work?",
+    "What are the benefits of exercise?",
+    "Once upon a time in a small village,",
+    "The most important invention of the 20th century was",
+    "Translate 'hello world' to Japanese.",
+    "What is the meaning of life?",
+    "Describe the process of making bread.",
+    "Why is the sky blue?",
+    "What is the difference between AI and ML?",
+    "如何评价GPT-4？",
+    "Write a haiku about autumn.",
+    "Explain the Pythagorean theorem.",
+    "What causes earthquakes?",
+    "How does the internet work?",
+    "What is the speed of light?",
+    "Describe the water cycle.",
+    "What is democracy?",
+    "How do vaccines work?",
+    "What is blockchain technology?",
+    "Explain supply and demand.",
+    "What is the Big Bang theory?",
+    "How do airplanes fly?",
+    "What is climate change?",
+    "Describe the human digestive system.",
+    "What is artificial intelligence?",
+    "How does electricity work?",
+    "What is the solar system?",
+    "Explain the concept of gravity.",
+    "What is DNA?",
+    "How do computers store data?",
+    "What is the greenhouse effect?",
+    "Describe the structure of an atom.",
+    "What is machine learning?",
+    "How does Wi-Fi work?",
+    "What is the stock market?",
+    "Explain natural selection.",
+    "What is renewable energy?",
+    "How do batteries work?",
+    "What is the United Nations?",
+    "Describe the process of evolution.",
+    "What is cryptography?",
+    "请用三句话总结量子力学的核心概念。",
+    "用Python写一个计算斐波那契数列的函数。",
+]
+
+
+def logits_correctness_test():
+    """Compare xserv prefill logits with HuggingFace transformers."""
+    print("\n" + "=" * 60)
+    print("CORRECTNESS TEST: Comparing logits with HuggingFace")
+    print("=" * 60)
+
+    try:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+    except ImportError:
+        print("SKIP: transformers/torch not installed")
+        return None
+
+    print(f"Loading HF model from {MODEL_DIR}...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_DIR,
+        torch_dtype=torch.bfloat16,
+        device_map="cuda:1",  # Use GPU 1 (xserv uses GPU 0)
+        trust_remote_code=True,
+    )
+    model.eval()
+
+    test_prompts = TEST_PROMPTS[:10]  # Use first 10 for logits comparison
+    xserv_bin = "/opt/wjh/projects/xserv/target/release/dump-logits"
+
+    results = []
+    for i, prompt in enumerate(test_prompts):
+        print(f"\n[{i+1}/{len(test_prompts)}] Prompt: {prompt[:50]}...")
+
+        # --- HuggingFace ---
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        with torch.no_grad():
+            outputs = model(**inputs)
+        hf_logits = outputs.logits[0, -1, :].float().cpu()
+        hf_top = torch.topk(hf_logits, TOP_K)
+        hf_ids = hf_top.indices.tolist()
+        hf_vals = hf_top.values.tolist()
+
+        # --- xserv ---
+        try:
+            result = subprocess.run(
+                [xserv_bin, MODEL_DIR, prompt],
+                capture_output=True, text=True, timeout=120,
+                env={**os.environ, "CUDA_VISIBLE_DEVICES": "0",
+                     "PATH": "/usr/local/cuda-12.9/bin:" + os.environ.get("PATH", "")},
+            )
+            xserv_lines = [l for l in result.stdout.strip().split('\n') if l.strip().startswith('[')]
+            xserv_top = []
+            for line in xserv_lines[:TOP_K]:
+                parts = line.strip().split()
+                tid = int([p for p in parts if p.startswith('id=')][0].split('=')[1])
+                val = float([p for p in parts if p.startswith('logit=')][0].split('=')[1])
+                xserv_top.append((tid, val))
+        except Exception as e:
+            print(f"  xserv FAILED: {e}")
+            results.append({"prompt": prompt, "match": False, "error": str(e)})
+            continue
+
+        # --- Compare ---
+        xserv_ids = [t[0] for t in xserv_top]
+        xserv_vals = [t[1] for t in xserv_top]
+
+        # Top-1 match
+        top1_match = hf_ids[0] == xserv_ids[0] if xserv_ids else False
+        # Top-5 overlap
+        top5_overlap = len(set(hf_ids[:5]) & set(xserv_ids[:5]))
+        # Max logit difference for matching tokens
+        max_diff = 0
+        for j, (hid, hval) in enumerate(zip(hf_ids[:5], hf_vals[:5])):
+            for xid, xval in xserv_top[:5]:
+                if hid == xid:
+                    max_diff = max(max_diff, abs(hval - xval))
+
+        hf_tok = tokenizer.decode([hf_ids[0]])
+        xs_tok = tokenizer.decode([xserv_ids[0]]) if xserv_ids else "???"
+
+        status = "PASS" if top1_match else "WARN"
+        print(f"  Top-1: HF={hf_ids[0]}({hf_tok!r}) vs xserv={xserv_ids[0]}({xs_tok!r}) → {status}")
+        print(f"  Top-5 overlap: {top5_overlap}/5, max logit diff: {max_diff:.4f}")
+
+        results.append({
+            "prompt": prompt[:50],
+            "top1_match": top1_match,
+            "top5_overlap": top5_overlap,
+            "max_logit_diff": max_diff,
+            "hf_top1": f"{hf_ids[0]}({hf_tok})",
+            "xserv_top1": f"{xserv_ids[0]}({xs_tok})" if xserv_ids else "???",
+        })
+
+    # Summary
+    print("\n" + "-" * 40)
+    top1_matches = sum(1 for r in results if r.get("top1_match"))
+    avg_overlap = sum(r.get("top5_overlap", 0) for r in results) / max(len(results), 1)
+    print(f"Top-1 match: {top1_matches}/{len(results)}")
+    print(f"Avg top-5 overlap: {avg_overlap:.1f}/5")
+    print(f"Verdict: {'PASS' if top1_matches >= len(results) * 0.8 else 'FAIL'}")
+
+    # Cleanup
+    del model
+    torch.cuda.empty_cache()
+
+    return results
+
+
+def api_generation_test():
+    """Test 50+ prompts through the HTTP API."""
+    print("\n" + "=" * 60)
+    print("API GENERATION TEST: 50+ prompts via /v1/chat/completions")
+    print("=" * 60)
+
+    import urllib.request
+    import urllib.error
+
+    # Health check
+    try:
+        req = urllib.request.Request(f"{XSERV_URL}/health")
+        resp = urllib.request.urlopen(req, timeout=5)
+        assert resp.read().decode() == "ok"
+        print("Health check: OK")
+    except Exception as e:
+        print(f"FAIL: Server not reachable at {XSERV_URL}: {e}")
+        print("Start the server first: ./target/release/xserv-server /opt/wjh/models/qwen3-8b")
+        return None
+
+    # Models endpoint
+    try:
+        req = urllib.request.Request(f"{XSERV_URL}/v1/models")
+        resp = urllib.request.urlopen(req, timeout=5)
+        models = json.loads(resp.read())
+        print(f"Models: {[m['id'] for m in models['data']]}")
+    except Exception as e:
+        print(f"WARN: /v1/models failed: {e}")
+
+    results = []
+    total_prompt_tokens = 0
+    total_completion_tokens = 0
+    total_latency = 0
+    failures = 0
+
+    for i, prompt in enumerate(TEST_PROMPTS):
+        body = json.dumps({
+            "model": "qwen3-8b",
+            "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": 32,
+            "temperature": 0.0,
+        }).encode()
+
+        try:
+            req = urllib.request.Request(
+                f"{XSERV_URL}/v1/chat/completions",
+                data=body,
+                headers={"Content-Type": "application/json"},
+            )
+            t0 = time.time()
+            resp = urllib.request.urlopen(req, timeout=120)
+            latency = time.time() - t0
+            data = json.loads(resp.read())
+
+            content = data["choices"][0]["message"]["content"]
+            finish = data["choices"][0]["finish_reason"]
+            usage = data.get("usage", {})
+            pt = usage.get("prompt_tokens", 0)
+            ct = usage.get("completion_tokens", 0)
+
+            total_prompt_tokens += pt
+            total_completion_tokens += ct
+            total_latency += latency
+
+            # Basic quality checks
+            has_content = len(content.strip()) > 0
+            reasonable_length = ct > 0
+
+            status = "OK" if has_content and reasonable_length else "WARN"
+            if not has_content:
+                status = "FAIL"
+                failures += 1
+
+            truncated = content[:60].replace('\n', ' ')
+            print(f"  [{i+1:2d}/{len(TEST_PROMPTS)}] {status} | {latency:5.2f}s | pt={pt:3d} ct={ct:2d} | {truncated}...")
+
+            results.append({
+                "prompt": prompt[:40],
+                "status": status,
+                "latency": latency,
+                "prompt_tokens": pt,
+                "completion_tokens": ct,
+                "finish_reason": finish,
+                "content_preview": content[:80],
+            })
+
+        except Exception as e:
+            print(f"  [{i+1:2d}/{len(TEST_PROMPTS)}] FAIL | {e}")
+            failures += 1
+            results.append({"prompt": prompt[:40], "status": "FAIL", "error": str(e)})
+
+    # Summary
+    successes = len(results) - failures
+    avg_latency = total_latency / max(successes, 1)
+    tok_per_sec = total_completion_tokens / max(total_latency, 0.001)
+
+    print("\n" + "-" * 40)
+    print(f"Results: {successes}/{len(TEST_PROMPTS)} succeeded, {failures} failed")
+    print(f"Total prompt tokens: {total_prompt_tokens}")
+    print(f"Total completion tokens: {total_completion_tokens}")
+    print(f"Average latency: {avg_latency:.2f}s per request")
+    print(f"Throughput: {tok_per_sec:.1f} tokens/s (completion only)")
+    print(f"Verdict: {'PASS' if failures <= 2 else 'FAIL'}")
+
+    return results
+
+
+def streaming_test():
+    """Test SSE streaming works correctly."""
+    print("\n" + "=" * 60)
+    print("STREAMING TEST: SSE /v1/chat/completions?stream=true")
+    print("=" * 60)
+
+    import urllib.request
+    import urllib.error
+
+    body = json.dumps({
+        "model": "qwen3-8b",
+        "messages": [{"role": "user", "content": "Count from 1 to 5."}],
+        "max_tokens": 32,
+        "temperature": 0.0,
+        "stream": True,
+    }).encode()
+
+    req = urllib.request.Request(
+        f"{XSERV_URL}/v1/chat/completions",
+        data=body,
+        headers={"Content-Type": "application/json"},
+    )
+
+    try:
+        resp = urllib.request.urlopen(req, timeout=60)
+        content_type = resp.headers.get("content-type", "")
+        print(f"Content-Type: {content_type}")
+
+        chunks = []
+        full_text = ""
+        has_role_chunk = False
+        has_done = False
+        has_finish = False
+
+        for line in resp:
+            line = line.decode().strip()
+            if not line:
+                continue
+            if line.startswith("data: "):
+                data = line[6:]
+                if data == "[DONE]":
+                    has_done = True
+                    chunks.append("[DONE]")
+                    continue
+                try:
+                    obj = json.loads(data)
+                    delta = obj["choices"][0]["delta"]
+                    fr = obj["choices"][0].get("finish_reason")
+                    if "role" in delta:
+                        has_role_chunk = True
+                    if "content" in delta:
+                        full_text += delta["content"]
+                    if fr is not None:
+                        has_finish = True
+                    chunks.append(delta)
+                except json.JSONDecodeError:
+                    print(f"  WARN: bad JSON: {data[:80]}")
+
+        print(f"Chunks received: {len(chunks)}")
+        print(f"Has role chunk: {has_role_chunk}")
+        print(f"Has finish_reason: {has_finish}")
+        print(f"Has [DONE]: {has_done}")
+        print(f"Full text: {full_text[:100]!r}")
+
+        ok = has_role_chunk and has_done and has_finish and len(full_text) > 0
+        # SSE content-type check
+        if "text/event-stream" in content_type:
+            print("Content-Type: OK (text/event-stream)")
+        else:
+            print(f"WARN: Expected text/event-stream, got {content_type}")
+
+        print(f"Verdict: {'PASS' if ok else 'FAIL'}")
+        return ok
+
+    except Exception as e:
+        print(f"FAIL: {e}")
+        return False
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--mode", choices=["all", "logits", "api", "stream"], default="all")
+    args = parser.parse_args()
+
+    if args.mode in ("all", "logits"):
+        logits_correctness_test()
+
+    if args.mode in ("all", "api"):
+        api_generation_test()
+
+    if args.mode in ("all", "stream"):
+        streaming_test()
+
+
+if __name__ == "__main__":
+    main()

tools/test_concurrent.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+"""Test concurrent requests to verify continuous batching scheduling."""
+import json
+import time
+import urllib.request
+import concurrent.futures
+
+URL = "http://localhost:9090/v1/chat/completions"
+
+PROMPTS = [
+    "What is 1+1?",
+    "What is 2+2?",
+    "What is 3+3?",
+    "What is 4+4?",
+    "What is 5+5?",
+    "What is 6+6?",
+    "What is 7+7?",
+    "What is 8+8?",
+]
+
+def send_request(prompt, idx):
+    body = json.dumps({
+        "model": "qwen3-8b",
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": 32,
+        "temperature": 0.0,
+    }).encode()
+    req = urllib.request.Request(URL, data=body, headers={"Content-Type": "application/json"})
+    t0 = time.perf_counter()
+    resp = urllib.request.urlopen(req, timeout=120)
+    elapsed = time.perf_counter() - t0
+    data = json.loads(resp.read())
+    content = data["choices"][0]["message"]["content"][:50].replace('\n', ' ')
+    ct = data["usage"]["completion_tokens"]
+    return idx, prompt, elapsed, ct, content
+
+def main():
+    print("=== Concurrent request test (8 requests, max_batch=4) ===\n")
+
+    # Fire all 8 requests concurrently
+    t_start = time.perf_counter()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
+        futures = [pool.submit(send_request, p, i) for i, p in enumerate(PROMPTS)]
+        results = [f.result() for f in concurrent.futures.as_completed(futures)]
+    t_total = time.perf_counter() - t_start
+
+    results.sort(key=lambda r: r[0])
+    total_tokens = 0
+    for idx, prompt, elapsed, ct, content in results:
+        total_tokens += ct
+        print(f"  [{idx}] {elapsed:5.2f}s | ct={ct:2d} | {prompt} -> {content}...")
+
+    serial_estimate = sum(r[2] for r in results)
+    print(f"\n  Wall clock: {t_total:.2f}s")
+    print(f"  Sum of individual latencies: {serial_estimate:.2f}s")
+    print(f"  Concurrency speedup: {serial_estimate/t_total:.2f}x (1.0x = no batching)")
+    print(f"  Total tokens: {total_tokens}")
+    print(f"  Throughput: {total_tokens/t_total:.1f} tok/s")
+
+    if t_total < serial_estimate * 0.85:
+        print(f"\n  Concurrent scheduling is working (wall < 85% of serial sum)")
+    else:
+        print(f"\n  Limited concurrency benefit (scheduling correct, GPU still per-seq)")
+
+if __name__ == "__main__":
+    main()

tools/test_correctness.py

@@ -0,0 +1,239 @@
+#!/usr/bin/env python3
+"""Compare xserv vs HuggingFace transformers for correctness and performance.
+
+Strategy: run xserv first (on GPU 0), collect results, then load HF model
+on GPU 0 (xserv process exits and frees VRAM).
+"""
+
+import subprocess
+import time
+import json
+import sys
+import os
+import re
+
+DEVICE = "cuda:0"
+MODEL_DIR = "/opt/wjh/models/qwen3-8b"
+XSERV_DUMP = "/opt/wjh/projects/xserv/target/release/dump-logits"
+
+
+def xserv_dump_logits(prompt):
+    """Run xserv dump-logits and parse top-20."""
+    env = os.environ.copy()
+    env["PATH"] = "/usr/local/cuda-12.9/bin:" + env.get("PATH", "")
+    env["CUDA_VISIBLE_DEVICES"] = "0"
+    t0 = time.perf_counter()
+    result = subprocess.run(
+        [XSERV_DUMP, MODEL_DIR, prompt],
+        capture_output=True, text=True, timeout=180, env=env
+    )
+    elapsed = time.perf_counter() - t0
+    if result.returncode != 0:
+        print(f"  xserv error: {result.stderr[-500:]}")
+        return None, elapsed
+
+    top20 = []
+    for line in result.stdout.strip().split("\n"):
+        m = re.match(r'\s*\[\s*\d+\]\s+id=\s*(\d+)\s+logit=\s*([\-\d.]+)', line)
+        if m:
+            top20.append((int(m.group(1)), float(m.group(2))))
+    return top20, elapsed
+
+
+def hf_prefill_top20(model, tokenizer, prompt):
+    """Get top-20 logits from HF."""
+    import torch
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits[0, -1, :].float().cpu().numpy()
+    import numpy as np
+    top_ids = np.argsort(logits)[-20:][::-1]
+    return [(int(i), float(logits[i])) for i in top_ids]
+
+
+def hf_generate(model, tokenizer, prompt, max_new=80):
+    """Greedy generation from HF."""
+    import torch
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+    prompt_len = inputs["input_ids"].shape[1]
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    with torch.no_grad():
+        out = model.generate(**inputs, max_new_tokens=max_new, do_sample=False)
+    torch.cuda.synchronize()
+    elapsed = time.perf_counter() - t0
+    gen_tokens = out.shape[1] - prompt_len
+    text = tokenizer.decode(out[0][prompt_len:], skip_special_tokens=True)
+    return gen_tokens / elapsed, elapsed, gen_tokens, text
+
+
+def compare_top20(hf_top20, xs_top20, name):
+    if xs_top20 is None:
+        print(f"  [{name}] SKIP (xserv failed)")
+        return False
+
+    hf_ids = [x[0] for x in hf_top20]
+    xs_ids = [x[0] for x in xs_top20]
+    top1_match = hf_ids[0] == xs_ids[0]
+    top5_match = set(hf_ids[:5]) == set(xs_ids[:5])
+    top10_overlap = len(set(hf_ids[:10]) & set(xs_ids[:10]))
+
+    hf_dict = dict(hf_top20)
+    xs_dict = dict(xs_top20)
+    common = set(hf_dict.keys()) & set(xs_dict.keys())
+    if common:
+        diffs = [abs(hf_dict[k] - xs_dict[k]) for k in common]
+        max_diff = max(diffs)
+        mean_diff = sum(diffs) / len(diffs)
+    else:
+        max_diff = mean_diff = float('inf')
+
+    status = "PASS" if top1_match and top5_match else "FAIL"
+    print(f"  [{name}] {status}: top1={'Y' if top1_match else 'N'}, "
+          f"top5={'Y' if top5_match else 'N'}, top10={top10_overlap}/10, "
+          f"max_diff={max_diff:.4f}, mean_diff={mean_diff:.4f}")
+    print(f"    HF  top5: {[(i, f'{v:.2f}') for i, v in hf_top20[:5]]}")
+    print(f"    XS  top5: {[(i, f'{v:.2f}') for i, v in xs_top20[:5]]}")
+    return status == "PASS"
+
+
+def benchmark_xserv_server(prompt, num_tokens=80, port=8080):
+    import urllib.request
+    data = json.dumps({
+        "model": "qwen3-8b",
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": num_tokens,
+        "temperature": 0,
+        "stream": False
+    }).encode()
+    req = urllib.request.Request(
+        f"http://localhost:{port}/v1/chat/completions",
+        data=data, headers={"Content-Type": "application/json"}
+    )
+    start = time.perf_counter()
+    with urllib.request.urlopen(req, timeout=120) as resp:
+        result = json.loads(resp.read())
+    elapsed = time.perf_counter() - start
+    content = result["choices"][0]["message"]["content"]
+    usage = result.get("usage", {})
+    ct = usage.get("completion_tokens", 0)
+    return ct / elapsed if elapsed > 0 else 0, elapsed, ct, content
+
+
+def main():
+    with_server = "--with-server" in sys.argv
+
+    print("=" * 70)
+    print("xserv vs HuggingFace Transformers — Correctness & Performance")
+    print("=" * 70)
+    print(f"Model: {MODEL_DIR}")
+    print(f"Device: {DEVICE}\n")
+
+    # ── Phase A: Run xserv first (separate processes, each loads+runs+exits) ──
+    test_prompts = [
+        ("english", "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"),
+        ("chinese", "<|im_start|>user\n请介绍一下量子计算<|im_end|>\n<|im_start|>assistant\n"),
+        ("code", "<|im_start|>user\nWrite a Python function to sort a list<|im_end|>\n<|im_start|>assistant\n"),
+        ("multi_turn",
+         "<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\nHi!<|im_end|>\n<|im_start|>user\nWhat is 2+2?<|im_end|>\n<|im_start|>assistant\n"),
+    ]
+
+    print("=" * 50)
+    print("PART 1: Collecting xserv prefill logits")
+    print("=" * 50)
+
+    xs_results = {}
+    for name, prompt in test_prompts:
+        print(f"  Running xserv dump-logits [{name}]...")
+        top20, elapsed = xserv_dump_logits(prompt)
+        xs_results[name] = top20
+        if top20:
+            print(f"    OK ({len(top20)} logits, {elapsed:.1f}s)")
+        else:
+            print(f"    FAILED ({elapsed:.1f}s)")
+
+    # ── Phase B: Load HF model and compare ──
+    print(f"\n{'=' * 50}")
+    print("PART 2: Loading HF model for comparison")
+    print("=" * 50)
+
+    import torch
+    import numpy as np
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    print("Loading HF model (BF16)...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_DIR, dtype=torch.bfloat16,
+        device_map=DEVICE, trust_remote_code=True
+    )
+    model.eval()
+    print("HF model loaded.\n")
+
+    print("=" * 50)
+    print("PART 3: Correctness Comparison")
+    print("=" * 50)
+
+    all_pass = True
+    for name, prompt in test_prompts:
+        hf_top20 = hf_prefill_top20(model, tokenizer, prompt)
+        if not compare_top20(hf_top20, xs_results[name], name):
+            all_pass = False
+
+    print(f"\n  Overall: {'ALL PASS' if all_pass else 'SOME FAILED'}\n")
+
+    # ── Phase C: Performance benchmark ──
+    print("=" * 50)
+    print("PART 4: HF Decode Performance (greedy, batch=1)")
+    print("=" * 50)
+
+    bench_prompt = "<|im_start|>user\nExplain the theory of relativity in simple terms.<|im_end|>\n<|im_start|>assistant\n"
+
+    # Warmup
+    print("\nWarmup...")
+    hf_generate(model, tokenizer, bench_prompt, max_new=5)
+
+    # Benchmark multiple token counts
+    for num_tokens in [50, 80]:
+        hf_tps, hf_time, hf_gen, hf_text = hf_generate(model, tokenizer, bench_prompt, max_new=num_tokens)
+        print(f"  HF ({num_tokens} tokens): {hf_tps:.1f} tok/s, {hf_time:.2f}s, {hf_gen} generated")
+
+    # xserv server benchmark
+    if with_server:
+        print(f"\n{'=' * 50}")
+        print("PART 5: xserv Server Performance")
+        print("=" * 50)
+        try:
+            import urllib.request
+            urllib.request.urlopen("http://localhost:8080/health", timeout=3)
+            print("Server available. Benchmarking...\n")
+
+            # Warmup
+            benchmark_xserv_server("Hi", 5)
+            time.sleep(0.5)
+
+            for num_tokens in [50, 80]:
+                xs_tps, xs_time, xs_gen, xs_text = benchmark_xserv_server(
+                    "Explain the theory of relativity in simple terms.", num_tokens
+                )
+                print(f"  xserv ({num_tokens} tokens): {xs_tps:.1f} tok/s, {xs_time:.2f}s, {xs_gen} generated")
+                print(f"    Text: {xs_text[:120]}...")
+
+            # EOS leak check
+            print(f"\n  EOS Leak Check:")
+            _, _, _, content = benchmark_xserv_server("Say hello", 20)
+            has_eos = "<|im_end|>" in content or "<|endoftext|>" in content or "<|im_start|>" in content
+            print(f"    Response has EOS token: {'YES (FAIL)' if has_eos else 'NO (PASS)'}")
+            if has_eos:
+                print(f"    Content: {content}")
+        except Exception as e:
+            print(f"Server not available: {e}")
+
+    print(f"\n{'=' * 50}")
+    print("DONE")
+    print("=" * 50)
+
+
+if __name__ == "__main__":
+    main()