docs: TP=1/2/4 xserv vs llama.cpp benchmark results

AIME 2025 + GSM8K at TP=1/2/4. Quality on par across engines/TP. Opposite
perf scaling: xserv TPOT improves with TP (21->17->15ms) while llama.cpp
row-split regresses over PCIe (10->19->20ms), crossing over so xserv is faster
at TP=4. Includes the clean same-path bench-tp scaling (58/76/86 tok/s).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

.gitignore

@@ -7,3 +7,16 @@
 **/*.rs.bk
 .env
 *.npy
+
+# llama.cpp baseline (cloned/submoduled by tools/setup-llama-cpp.sh)
+/third_party/llama.cpp/build/
+/third_party/llama.cpp/models/
+*.gguf
+
+# Benchmark output + fetched datasets (transferred to GPU host, not committed)
+/bench-out/
+/tools/bench/data/
+/tools/__pycache__/
+/tools/bench/__pycache__/
+/tools/bench/**/__pycache__/
+

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "third_party/llama.cpp"]
+	path = third_party/llama.cpp
+	url = https://github.com/ggerganov/llama.cpp

Cargo.toml

@@ -7,6 +7,7 @@ members = [
     "crates/xserv-model",
     "crates/xserv-tokenizer",
     "crates/xserv-server",
+    "crates/xserv-distributed",
 ]
 
 [workspace.package]
@@ -24,3 +25,5 @@ regex = "1"
 tokio = { version = "1", features = ["full"] }
 axum = "0.8"
 uuid = { version = "1", features = ["v4"] }
+tokio-stream = "0.1"
+rand = "0.8"

README.md

@@ -0,0 +1,160 @@
+# xserv
+
+> 从零用 **Rust + CUDA** 构建的 LLM 推理引擎，目标是吃透 LLM Serving 全栈技术。
+
+xserv 不依赖 PyTorch / vLLM / TensorRT 等现成框架，自己实现了张量抽象、CUDA kernel、
+分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。当前在单张 RTX 5090 上可以
+跑通 **Qwen3-8B**（BF16），并提供一套与 **llama.cpp** 对比正确性和性能的标准 benchmark。
+
+## 现状一览
+
+- **模型**：GPT-2（124M）、Qwen3-8B（BF16）
+- **性能**（RTX 5090，Qwen3-8B BF16，贪心解码，单流）：约 **56 tok/s**，约为 HF transformers 的 1.4×、llama.cpp 的 ~0.6×
+- **精度**：在 AIME 2025 / GSM8K 上与 llama.cpp 同权重对比基本持平（数值保真度验证通过）
+- **服务**：OpenAI 兼容 `/v1/chat/completions`，支持 SSE 流式输出
+- **关键能力**：自写 GEMM / Flash-Attention 2(SM120) / Paged-Attention kernel、
+  分页 KV cache（含 **CPU 换出/换入** 弹性显存）、连续批处理（continuous batching）、
+  CUDA Graph 解码、按显存自适应的 KV 池
+
+> 这是一个以学习为主的项目，逐 Phase 推进，每步都做数值/端到端验证。
+
+## 架构
+
+```
+xserv/
+├── csrc/                  # CUDA 源码 (.cu/.cuh)
+│   ├── gemm/              #   GEMM (naive / tiled / gemv)
+│   ├── attention/         #   Flash-Attention 2 (SM120)、Paged-Attention、causal mask
+│   ├── normalization/     #   LayerNorm / RMSNorm
+│   ├── activation/        #   GELU / SiLU
+│   ├── embedding/         #   embedding lookup / RoPE / transpose
+│   └── reduce/            #   softmax
+├── crates/
+│   ├── xserv-cuda/        # CUDA FFI、Stream、显存分配器、Pinned 内存、CUDA Graph
+│   ├── xserv-tensor/      # Tensor 类型（strided 布局、BF16/F16/F32、CPU↔GPU）
+│   ├── xserv-kernels/     # kernel registry（自写 kernel + cuBLAS 可切换）
+│   ├── xserv-tokenizer/   # BPE 分词器
+│   ├── xserv-model/       # 模型定义（GPT-2 / Qwen3）、权重加载、KV cache、采样
+│   └── xserv-server/      # tokio + axum HTTP 服务、调度器
+├── tools/                 # 辅助脚本 + benchmark 套件（见下）
+└── docs/                  # 每个 Phase 的设计文档 + benchmark 报告
+```
+
+## 环境要求
+
+- **GPU**：NVIDIA，计算能力 SM120（RTX 5090 / Blackwell）。其它架构需调整 `CUDA_ARCH`。
+- **CUDA Toolkit**：12.9（`nvcc` 需在 `PATH`，构建 `.cu` 依赖它）
+- **Rust**：edition 2024（建议较新的 stable 工具链）
+- **模型**：HuggingFace 目录格式（含 `config.json`、`tokenizer.json`、`*.safetensors`）
+
+## 构建
+
+```bash
+export CUDA_HOME=/usr/local/cuda-12.9
+export PATH=$CUDA_HOME/bin:$PATH
+cargo build --release
+```
+
+如果本地没有 GPU/CUDA，可用远端构建脚本把代码同步到带卡的机器上构建/运行/测试：
+
+```bash
+./tools/sync-and-build.sh build      # 远端 cargo build --release
+./tools/sync-and-build.sh test       # 远端 cargo test
+```
+
+（远端主机、目录、模型路径在 `tools/sync-and-build.sh` 顶部配置。）
+
+## 基本用法
+
+### 1. 启动 HTTP 服务（OpenAI 兼容）
+
+```bash
+./target/release/xserv-server /path/to/qwen3-8b \
+    --port 8080 \
+    --max-batch 4 \
+    --max-seq-len 8192 \
+    --swap-space-gb 8
+```
+
+参数说明：
+
+| 参数 | 含义 | 默认 |
+|------|------|------|
+| `--port` | 监听端口 | 8080 |
+| `--max-batch` | 解码批大小（并发上限） | 4 |
+| `--max-seq-len` | 单序列最大长度 | 2048 |
+| `--swap-space-gb` | KV 换出到 CPU 的 pinned 内存大小（0 关闭） | 8 |
+
+请求示例（流式）：
+
+```bash
+curl http://localhost:8080/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "qwen3-8b",
+    "messages": [{"role": "user", "content": "用一句话解释什么是注意力机制"}],
+    "max_tokens": 256,
+    "temperature": 0,
+    "stream": true
+  }'
+```
+
+其它端点：`GET /health`、`GET /v1/models`。
+
+### 2. 命令行推理
+
+```bash
+# 单轮生成
+cargo run --release --bin xserv-cli -- /path/to/qwen3-8b --max-tokens 256
+
+# 交互式多轮对话
+cargo run --release --bin xserv-chat -- /path/to/qwen3-8b
+```
+
+### 3. 单机性能基准
+
+```bash
+# 输出每个 prompt 的 TTFT / TBT / TPOT（JSON）
+cargo run --release --bin bench-qwen3 -- /path/to/qwen3-8b --gen-tokens 64 [--cuda-graph]
+```
+
+## 与 llama.cpp 对比 benchmark
+
+`tools/bench/` 提供一套一键对比套件，把 xserv 和 **llama.cpp**（同一份 BF16 权重）放在
+相同负载下，黑盒通过 OpenAI API 对比：
+
+- **性能**：TTFT、TPOT、吞吐（单流 + 不同并发）
+- **精度**：AIME 2025、GSM8K（标准数据集，exact-match 评分）
+
+```bash
+# 一次性准备（需联网的机器）：拉取 llama.cpp 子模块 + 下载数据集
+git submodule update --init third_party/llama.cpp        # 固定在 tag b9371
+HF_ENDPOINT=https://hf-mirror.com python3 -m tools.bench.fetch_datasets
+
+# 一键对比（构建 llama.cpp + 转 GGUF + 构建 xserv + 跑两套 + 出报告）
+./tools/sync-and-build.sh bench -- --max-seq-len 8192 --quality-limit 50
+./tools/sync-and-build.sh fetch-bench-out
+# 报告产物：bench-out/comparison-<时间戳>.{md,json}
+```
+
+设计细节见 `docs/16-llama-cpp-comparison.md`，结果报告见 `docs/benchmarks/llama-cpp-comparison.md`。
+
+## 文档
+
+- `docs/00-roadmap.md`：总体路线图与各 Phase 设计
+- `docs/01..15-*.md`：CUDA FFI / Tensor / GEMM / Attention / KV cache / 性能优化等每个 Phase 的设计文档
+- `docs/16-llama-cpp-comparison.md`：llama.cpp 对比基准的设计
+- `docs/benchmarks/`：各阶段的 benchmark 报告
+
+## 路线图（节选）
+
+已完成 Phase 0–15：CUDA 基础设施 → Tensor → GEMM → Transformer kernels → Attention →
+模型加载 → 分词器 → GPT-2 → KV cache → Qwen3-8B → Paged Attention → 连续批处理 →
+HTTP API → Flash Attention 2 → 性能优化；并在此基础上加入了 **llama.cpp 对比基准**
+与 **KV CPU 换出** 等基础设施。
+
+后续方向：投机解码（speculative decoding）、张量并行（TP，多卡）、量化（FP8 / INT8）、多模态。
+
+## 许可
+
+MIT

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).
+    // Heap-allocate oversized buffer for cudaDeviceProp (layout varies by CUDA version).
-    let mut prop = unsafe { std::mem::zeroed::<ffi::CudaDeviceProp>() };
+    // CUDA 12.x struct is ~5-6 KB; use 32 KB to guard against future growth.
-    error::check(unsafe { ffi::cudaGetDeviceProperties(&mut prop, device as i32) })?;
+    let mut prop_buf = vec![0u8; 32768];
-    let name = unsafe { CStr::from_ptr(prop.name.as_ptr()) }
+    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/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)]
+/// cudaStreamCaptureMode::cudaStreamCaptureModeGlobal
-pub struct CudaDeviceProp {
+pub const CUDA_STREAM_CAPTURE_MODE_GLOBAL: i32 = 0;
-    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],
-}
 
 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 ---
@@ -52,6 +39,7 @@ unsafe extern "C" {
         stream: CudaStream,
     ) -> i32;
     pub fn cudaMemset(devptr: *mut u8, value: i32, count: usize) -> i32;
+    pub fn cudaMemsetAsync(devptr: *mut u8, value: i32, count: usize, stream: CudaStream) -> i32;
 
     // --- Stream ---
     pub fn cudaStreamCreate(stream: *mut CudaStream) -> i32;
@@ -62,6 +50,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 {
@@ -101,6 +116,56 @@ impl GpuBuffer {
         })
     }
 
+    /// Async copy `count` bytes from `src` at `src_offset` to `self` at `dst_offset` on `stream`.
+    pub fn copy_from_device_at_async(&mut self, src: &GpuBuffer, src_offset: usize, dst_offset: usize, count: usize, stream: &CudaStream) -> Result<()> {
+        assert!(src_offset + count <= src.len);
+        assert!(dst_offset + count <= self.len);
+        error::check(unsafe {
+            ffi::cudaMemcpyAsync(
+                self.ptr.add(dst_offset),
+                src.ptr.add(src_offset),
+                count,
+                ffi::CUDA_MEMCPY_D2D,
+                stream.as_raw(),
+            )
+        })
+    }
+
+    /// Copy `count` bytes from this GPU buffer at `src_offset` to a host slice (D2H).
+    pub fn copy_to_host_at(&self, dst: &mut [u8], src_offset: usize, count: usize) -> Result<()> {
+        assert!(src_offset + count <= self.len, "src range out of bounds");
+        assert!(count <= dst.len(), "host dst too small");
+        error::check(unsafe {
+            ffi::cudaMemcpy(
+                dst.as_mut_ptr(),
+                self.ptr.add(src_offset),
+                count,
+                ffi::CUDA_MEMCPY_D2H,
+            )
+        })
+    }
+
+    /// Copy `count` bytes from a host slice to this GPU buffer at `dst_offset` (H2D).
+    pub fn copy_from_host_at(&mut self, src: &[u8], dst_offset: usize, count: usize) -> Result<()> {
+        assert!(dst_offset + count <= self.len, "dst range out of bounds");
+        assert!(count <= src.len(), "host src too small");
+        error::check(unsafe {
+            ffi::cudaMemcpy(
+                self.ptr.add(dst_offset),
+                src.as_ptr(),
+                count,
+                ffi::CUDA_MEMCPY_H2D,
+            )
+        })
+    }
+
+    /// Async zero fill on stream.
+    pub fn zero_async(&mut self, stream: &CudaStream) -> Result<()> {
+        error::check(unsafe {
+            ffi::cudaMemsetAsync(self.ptr, 0, self.len, stream.as_raw())
+        })
+    }
+
     /// 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) {
@@ -113,17 +178,32 @@ 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() {
+        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) };
             }
         }
     }
+}
 
 unsafe impl Send for GpuBuffer {}
 

crates/xserv-distributed/Cargo.toml

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

crates/xserv-distributed/build.rs

@@ -0,0 +1,13 @@
+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");
+    // NCCL is typically installed as a system library.
+    println!("cargo:rustc-link-search=native=/usr/lib/x86_64-linux-gnu");
+    println!("cargo:rustc-link-lib=dylib=nccl");
+    println!("cargo:rustc-link-lib=dylib=cudart");
+}

crates/xserv-distributed/src/ffi.rs

@@ -0,0 +1,65 @@
+//! Minimal NCCL FFI bindings (hand-written, like the CUDA bindings).
+//! Only the collectives we need for tensor parallelism.
+
+use std::ffi::c_void;
+use std::os::raw::c_char;
+use xserv_cuda::ffi::CudaStream;
+
+/// Opaque NCCL communicator handle (`ncclComm_t`).
+pub type NcclComm = *mut c_void;
+
+/// `ncclUniqueId` is a 128-byte opaque blob shared from rank 0 to all ranks.
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct NcclUniqueId {
+    pub internal: [c_char; 128],
+}
+
+impl Default for NcclUniqueId {
+    fn default() -> Self {
+        Self { internal: [0; 128] }
+    }
+}
+
+// ncclDataType_t (subset)
+pub const NCCL_FLOAT32: i32 = 7;
+pub const NCCL_BF16: i32 = 9;
+
+// ncclRedOp_t
+pub const NCCL_SUM: i32 = 0;
+
+// ncclResult_t
+pub const NCCL_SUCCESS: i32 = 0;
+
+unsafe extern "C" {
+    pub fn ncclGetUniqueId(uid: *mut NcclUniqueId) -> i32;
+    // ncclUniqueId is passed BY VALUE (a 128-byte struct) per the NCCL ABI.
+    pub fn ncclCommInitRank(comm: *mut NcclComm, nranks: i32, commid: NcclUniqueId, rank: i32) -> i32;
+    pub fn ncclCommDestroy(comm: NcclComm) -> i32;
+    pub fn ncclAllReduce(
+        sendbuff: *const c_void,
+        recvbuff: *mut c_void,
+        count: usize,
+        datatype: i32,
+        op: i32,
+        comm: NcclComm,
+        stream: CudaStream,
+    ) -> i32;
+    pub fn ncclGroupStart() -> i32;
+    pub fn ncclGroupEnd() -> i32;
+    pub fn ncclGetErrorString(result: i32) -> *const c_char;
+}
+
+pub fn err_string(result: i32) -> String {
+    unsafe {
+        let p = ncclGetErrorString(result);
+        if p.is_null() {
+            return format!("nccl error {result}");
+        }
+        std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned()
+    }
+}
+
+pub fn check(result: i32, what: &str) {
+    assert_eq!(result, NCCL_SUCCESS, "{what} failed: {}", err_string(result));
+}

crates/xserv-distributed/src/lib.rs

@@ -0,0 +1,97 @@
+//! Tensor-parallel primitives for xserv.
+//!
+//! Process model: one OS thread per TP rank, each bound to one GPU. NCCL is
+//! used for the collective (AllReduce); a hand-rolled P2P AllReduce may replace
+//! it later as a learning exercise (see docs/17-tensor-parallelism.md).
+
+pub mod ffi;
+
+use std::ffi::c_void;
+
+use ffi::{NcclComm, NcclUniqueId};
+use xserv_cuda::device;
+use xserv_cuda::GpuBuffer;
+
+pub use ffi::NcclUniqueId as UniqueId;
+
+/// The CUDA "null" (default) stream. The model's kernels and cuBLAS calls run
+/// on it, so issuing NCCL on the same stream keeps AllReduce correctly ordered
+/// after the producing matmul and before the consuming kernel — no extra sync.
+const NULL_STREAM: xserv_cuda::ffi::CudaStream = std::ptr::null_mut();
+
+/// Generate a unique id on one rank (typically rank 0) and broadcast the bytes
+/// to all ranks out-of-band (e.g. via a shared variable across threads).
+pub fn get_unique_id() -> NcclUniqueId {
+    let mut id = NcclUniqueId::default();
+    ffi::check(unsafe { ffi::ncclGetUniqueId(&mut id) }, "ncclGetUniqueId");
+    id
+}
+
+/// Per-rank tensor-parallel context: NCCL communicator + a dedicated stream.
+pub struct TpContext {
+    pub rank: usize,
+    pub world: usize,
+    pub device: u32,
+    comm: NcclComm,
+}
+
+// The NCCL communicator is owned by exactly one rank thread.
+unsafe impl Send for TpContext {}
+
+impl TpContext {
+    /// Initialize this rank. Must be called from the thread that will own this
+    /// rank's GPU work; binds the thread to `device` first. All ranks must call
+    /// this concurrently with the same `id` and `world`.
+    pub fn init(rank: usize, world: usize, id: NcclUniqueId, device: u32) -> Self {
+        device::set_device(device).expect("set_device");
+        let mut comm: NcclComm = std::ptr::null_mut();
+        // Wrap the concurrent inits in a group so they rendezvous without deadlock.
+        ffi::check(unsafe { ffi::ncclGroupStart() }, "ncclGroupStart(init)");
+        ffi::check(
+            unsafe { ffi::ncclCommInitRank(&mut comm, world as i32, id, rank as i32) },
+            "ncclCommInitRank",
+        );
+        ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
+        Self { rank, world, device, comm }
+    }
+
+    /// In-place AllReduce(sum) over `count` BF16 elements in `buf`.
+    pub fn all_reduce_sum_bf16(&self, buf: &mut GpuBuffer, count: usize) {
+        self.all_reduce_sum_bf16_ptr(buf.as_mut_ptr() as *mut c_void, count);
+    }
+
+    /// In-place AllReduce(sum) directly on a device pointer (`count` BF16 elems),
+    /// issued on the null stream so it is ordered with the model's kernels.
+    /// Asynchronous: a later sync (e.g. the D2H logits copy) completes it.
+    ///
+    /// # Safety
+    /// `ptr` must point to at least `count` BF16 elements of valid device memory
+    /// on this rank's device. The reduction is in-place (send == recv).
+    pub fn all_reduce_sum_bf16_ptr(&self, ptr: *mut c_void, count: usize) {
+        if self.world == 1 {
+            return; // nothing to reduce
+        }
+        ffi::check(
+            unsafe {
+                ffi::ncclAllReduce(
+                    ptr as *const c_void,
+                    ptr,
+                    count,
+                    ffi::NCCL_BF16,
+                    ffi::NCCL_SUM,
+                    self.comm,
+                    NULL_STREAM,
+                )
+            },
+            "ncclAllReduce",
+        );
+    }
+}
+
+impl Drop for TpContext {
+    fn drop(&mut self) {
+        if !self.comm.is_null() {
+            unsafe { ffi::ncclCommDestroy(self.comm) };
+        }
+    }
+}

crates/xserv-distributed/tests/allreduce.rs

@@ -0,0 +1,50 @@
+//! 2-GPU AllReduce smoke test. Skips if fewer than 2 GPUs are present.
+
+use half::bf16;
+use std::thread;
+use xserv_cuda::{device, GpuBuffer};
+use xserv_distributed::{get_unique_id, TpContext};
+
+#[test]
+fn allreduce_two_gpu_sum() {
+    let world = 2usize;
+    if device::device_count().unwrap_or(0) < world as i32 {
+        eprintln!("skip: need >= {world} GPUs");
+        return;
+    }
+
+    let id = get_unique_id();
+    let n = 4096usize;
+
+    let handles: Vec<_> = (0..world)
+        .map(|rank| {
+            let id = id;
+            thread::spawn(move || {
+                let ctx = TpContext::init(rank, world, id, rank as u32);
+
+                // Rank r fills its buffer with (r + 1).
+                let val = bf16::from_f32((rank + 1) as f32);
+                let host = vec![val; n];
+                let src = unsafe {
+                    std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2)
+                };
+                let mut buf = GpuBuffer::alloc(n * 2).unwrap();
+                buf.copy_from_host(src).unwrap();
+
+                ctx.all_reduce_sum_bf16(&mut buf, n);
+
+                let mut out = vec![0u8; n * 2];
+                buf.copy_to_host(&mut out).unwrap();
+                let res = unsafe { std::slice::from_raw_parts(out.as_ptr() as *const bf16, n) };
+                (res[0].to_f32(), res[n - 1].to_f32())
+            })
+        })
+        .collect();
+
+    // sum over ranks of (r+1) = 1 + 2 = 3
+    for h in handles {
+        let (first, last) = h.join().unwrap();
+        assert_eq!(first, 3.0, "AllReduce(sum) first element");
+        assert_eq!(last, 3.0, "AllReduce(sum) last element");
+    }
+}

crates/xserv-kernels/build.rs

@@ -16,6 +16,7 @@ fn main() {
         .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")
@@ -24,6 +25,8 @@ fn main() {
         .file("../../csrc/embedding/rope.cu")
         .file("../../csrc/attention/causal_mask.cu")
         .file("../../csrc/embedding/transpose.cu")
+        .file("../../csrc/attention/flash_attention.cu")
+        .file("../../csrc/attention/paged_attention.cu")
         .compile("xserv_kernels");
 
     println!("cargo:rerun-if-changed=../../csrc/");

crates/xserv-kernels/src/activation.rs

@@ -12,13 +12,16 @@ unsafe extern "C" {
     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::zeros(x.shape(), x.dtype(), x.device());
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
-    let n = x.numel() as i32;
+    let n = x.numel();
+    assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
+    let n = n 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()),
@@ -26,7 +29,6 @@ fn dispatch_unary(x: &Tensor, f32_fn: unsafe extern "C" fn(*const c_void, *mut c
             _ => panic!("unsupported dtype"),
         }
     }
-    xserv_cuda::device::synchronize().unwrap();
     out
 }
 
@@ -37,8 +39,10 @@ fn dispatch_binary(a: &Tensor, b: &Tensor,
     assert!(a.is_contiguous() && b.is_contiguous());
     assert!(matches!(a.device(), Device::Cuda(_)));
     assert_eq!(a.dtype(), b.dtype());
-    let out = Tensor::zeros(a.shape(), a.dtype(), a.device());
+    let out = Tensor::empty(a.shape(), a.dtype(), a.device());
-    let n = a.numel() as i32;
+    let n = a.numel();
+    assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
+    let n = n 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()),
@@ -46,7 +50,6 @@ fn dispatch_binary(a: &Tensor, b: &Tensor,
             _ => panic!("unsupported dtype"),
         }
     }
-    xserv_cuda::device::synchronize().unwrap();
     out
 }
 
@@ -55,8 +58,10 @@ pub fn silu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_silu_f32, launch_si
 
 pub fn scale(x: &Tensor, scale_val: f32) -> Tensor {
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
-    let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
-    let n = x.numel() as i32;
+    let n = x.numel();
+    assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
+    let n = n 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()),
@@ -64,9 +69,31 @@ pub fn scale(x: &Tensor, scale_val: f32) -> Tensor {
             _ => panic!("unsupported dtype for scale"),
         }
     }
-    xserv_cuda::device::synchronize().unwrap();
     out
 }
 
 pub fn add(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_add_f32, launch_add_bf16) }
 pub fn mul(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16) }
+
+/// 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();
+    assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
+    let n = n 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

@@ -10,6 +10,29 @@ unsafe extern "C" {
                               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 launch_paged_decode_attention_bf16(
+        q: *const c_void,
+        k_cache: *const c_void,
+        v_cache: *const c_void,
+        o: *mut c_void,
+        block_tables: *const i32,
+        context_lens: *const i32,
+        batch: i32, num_q_heads: i32, num_kv_heads: i32,
+        head_dim: i32, max_blocks_per_seq: i32,
+        scale: f32, stream: *mut c_void,
+    );
 }
 
 fn apply_causal_mask(scores: &Tensor, offset: usize) {
@@ -33,7 +56,6 @@ fn apply_causal_mask(scores: &Tensor, offset: usize) {
             _ => panic!("unsupported dtype for causal mask"),
         }
     }
-    xserv_cuda::device::synchronize().unwrap();
 }
 
 /// Multi-head attention (naive, materializes S×S score matrix).
@@ -75,3 +97,164 @@ pub fn attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tensor {
     // 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
+}
+
+/// Paged decode attention.
+///
+/// q: [batch, num_q_heads, 1, head_dim] BF16, contiguous, GPU
+/// k_cache_ptr / v_cache_ptr: pointers to [num_blocks, num_kv_heads, BLOCK_SIZE, head_dim] BF16 pools
+/// block_tables_ptr: i32 [batch, max_blocks_per_seq] (rows already arranged for this batch)
+/// context_lens_ptr: i32 [batch]
+///
+/// Returns: [batch, num_q_heads, 1, head_dim] BF16
+#[allow(clippy::too_many_arguments)]
+pub fn paged_decode_attention(
+    q: &Tensor,
+    k_cache_ptr: *const c_void,
+    v_cache_ptr: *const c_void,
+    block_tables_ptr: *const i32,
+    context_lens_ptr: *const i32,
+    batch: usize,
+    num_q_heads: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    max_blocks_per_seq: usize,
+) -> Tensor {
+    assert_eq!(q.ndim(), 4);
+    assert_eq!(q.shape()[2], 1, "paged_decode_attention requires q_len == 1");
+    assert_eq!(q.dtype(), DType::BF16);
+    assert!(num_q_heads % num_kv_heads == 0, "GQA: num_q_heads must be divisible by num_kv_heads");
+    assert!(head_dim <= 128);
+
+    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_paged_decode_attention_bf16(
+            q.data_ptr() as *const c_void,
+            k_cache_ptr,
+            v_cache_ptr,
+            output.data_ptr() as *mut c_void,
+            block_tables_ptr,
+            context_lens_ptr,
+            batch as i32,
+            num_q_heads as i32,
+            num_kv_heads as i32,
+            head_dim as i32,
+            max_blocks_per_seq as i32,
+            scale,
+            std::ptr::null_mut(),
+        );
+    }
+
+    output
+}

crates/xserv-kernels/src/dispatch.rs

@@ -0,0 +1,118 @@
+//! Low-level kernel dispatchers for CUDA Graph capture.
+//! These functions write to pre-allocated output buffers and accept an explicit stream.
+
+use std::ffi::c_void;
+
+// Re-declare the extern functions we need (same as in the individual modules)
+unsafe extern "C" {
+    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);
+    fn launch_silu_mul_bf16(gate: *const c_void, up: *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_embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
+                             num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void);
+    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_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_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);
+    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,
+    );
+}
+
+/// Raw rmsnorm dispatch: writes to pre-allocated `out`.
+pub unsafe fn 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) {
+    launch_rmsnorm_bf16(x, gamma, out, rows, hidden_size, eps, stream);
+}
+
+/// Raw add_rmsnorm dispatch.
+pub unsafe fn 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) {
+    launch_add_rmsnorm_bf16(x, residual, gamma, normed_out, sum_out, rows, hidden_size, eps, stream);
+}
+
+/// Raw silu_mul dispatch.
+pub unsafe fn silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void) {
+    launch_silu_mul_bf16(gate, up, out, n, stream);
+}
+
+/// Raw add dispatch.
+pub unsafe fn add_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void) {
+    launch_add_bf16(a, b, out, n, stream);
+}
+
+/// Raw embedding dispatch.
+pub unsafe fn embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
+                             num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void) {
+    launch_embedding_bf16(table, token_ids, out, num_tokens, hidden_size, vocab_size, stream);
+}
+
+/// Raw reshape_heads dispatch.
+pub unsafe fn reshape_heads_bf16(inp: *const c_void, out: *mut c_void,
+                                 seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void) {
+    launch_reshape_heads_bf16(inp, out, seq_len, num_heads, head_dim, stream);
+}
+
+/// Raw merge_heads dispatch.
+pub unsafe fn merge_heads_bf16(inp: *const c_void, out: *mut c_void,
+                               seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void) {
+    launch_merge_heads_bf16(inp, out, seq_len, num_heads, head_dim, stream);
+}
+
+/// Raw transpose HSD->SHD dispatch.
+pub unsafe fn 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) {
+    launch_transpose_hsd_to_shd_bf16(inp, out, seq_len, num_heads, head_dim, stream);
+}
+
+/// Raw transpose SHD->HSD dispatch.
+pub unsafe fn 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) {
+    launch_transpose_shd_to_hsd_bf16(inp, out, seq_len, num_heads, head_dim, stream);
+}
+
+/// Raw RoPE dispatch (in-place).
+pub unsafe fn 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) {
+    launch_rope_bf16(x, cos_cache, sin_cache, positions, num_tokens, num_heads, head_dim, stream);
+}
+
+/// Raw GEMV dispatch (BF16, M=1). Caller must provide fp32 accumulator buffer.
+pub unsafe fn 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) {
+    launch_gemv_bf16(x, w, y_bf16, y_fp32_buf, k, n, stream);
+}
+
+/// Raw decode attention dispatch.
+pub unsafe fn 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, stream: *mut c_void) {
+    launch_decode_attention_bf16(q, k, v, o, batch, num_q_heads, num_kv_heads, kv_len, head_dim, scale, 1, stream);
+}
+
+// cuBLAS FFI
+pub type CublasHandle = *mut c_void;
+
+unsafe extern "C" {
+    fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
+}
+
+/// Set cuBLAS stream. Must be called before any cuBLAS operations during graph capture.
+pub unsafe fn set_cublas_stream(handle: CublasHandle, stream: *mut c_void) {
+    cublasSetStream_v2(handle, stream);
+}

crates/xserv-kernels/src/embedding.rs

@@ -4,9 +4,9 @@ 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);
+                            num_tokens: i32, hidden_size: i32, vocab_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);
+                             num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void);
 }
 
 /// Embedding lookup: table[token_ids[i]] for each i.
@@ -18,6 +18,9 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
 
     let hidden_size = table.shape()[1];
     let num_tokens = token_ids.len();
+    let vocab_size = table.shape()[0];
+    assert!(num_tokens <= i32::MAX as usize, "too many tokens for i32 kernel param");
+    assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
 
     // Upload token_ids to GPU
     let ids_bytes = unsafe {
@@ -26,26 +29,29 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
             num_tokens * std::mem::size_of::<u32>(),
         )
     };
-    let mut ids_gpu = GpuBuffer::alloc(ids_bytes.len()).expect("alloc token_ids");
+    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::zeros(&[num_tokens, hidden_size], table.dtype(), table.device());
+    for &tid in token_ids {
+        assert!((tid as usize) < vocab_size, "token_id {tid} out of bounds (vocab_size={vocab_size})");
+    }
+
+    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(),
+                num_tokens as i32, hidden_size as i32, vocab_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(),
+                num_tokens as i32, hidden_size as i32, vocab_size as i32, std::ptr::null_mut(),
             ),
             _ => panic!("unsupported dtype for embedding"),
         }
     }
-    xserv_cuda::device::synchronize().unwrap();
     out
 }

crates/xserv-kernels/src/gemm.rs

@@ -1,3 +1,4 @@
+use std::cell::RefCell;
 use std::ffi::c_void;
 use xserv_cuda::error::{self, Result};
 use xserv_tensor::{DType, Device, Tensor};
@@ -15,10 +16,11 @@ unsafe extern "C" {
     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;
+pub type CublasHandle = *mut c_void;
 
 #[allow(non_upper_case_globals)]
 const CUBLAS_OP_N: i32 = 0;
@@ -81,6 +83,30 @@ impl Drop for CublasContext {
     }
 }
 
+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)
+    })
+}
+
+/// Get the thread-local cuBLAS handle for use with dispatch module.
+pub fn cublas_handle() -> CublasHandle {
+    CUBLAS_CTX.with(|cell| {
+        cell.borrow().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.
@@ -97,7 +123,9 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
     let n = b.shape()[1];
     let dtype = a.dtype();
 
-    let c = Tensor::zeros(&[m, n], dtype, a.device());
+    // 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;
@@ -113,7 +141,6 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
                     _ => panic!("unsupported dtype for naive GEMM"),
                 }
             }
-            xserv_cuda::device::synchronize().unwrap();
         }
         GemmBackend::Tiled => {
             unsafe {
@@ -123,13 +150,24 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
                     _ => panic!("unsupported dtype for tiled GEMM"),
                 }
             }
-            xserv_cuda::device::synchronize().unwrap();
         }
         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 ctx = CublasContext::new().unwrap();
                 let alpha = 1.0f32;
                 let beta = 0.0f32;
 
@@ -139,12 +177,12 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
                     _ => panic!("unsupported dtype for cuBLAS GEMM"),
                 };
 
-            unsafe {
+                with_cublas(|handle| unsafe {
-                cublasSetStream_v2(ctx.handle, null_stream);
+                    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(
-                    ctx.handle,
+                        handle,
                         CUBLAS_OP_N, CUBLAS_OP_N,
                         n as i32, m as i32, k as i32,
                         &alpha as *const f32 as *const c_void,
@@ -155,8 +193,8 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
                         CUBLAS_COMPUTE_32F,
                         -1, // default algo
                     )).expect("cuBLAS GEMM failed");
+                });
             }
-            xserv_cuda::device::synchronize().unwrap();
         }
     }
 
@@ -190,7 +228,8 @@ pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
     let mut out_shape: Vec<usize> = a.shape()[..ndim - 2].to_vec();
     out_shape.push(m);
     out_shape.push(n);
-    let c = Tensor::zeros(&out_shape, a.dtype(), a.device());
+    // 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 {
@@ -206,12 +245,11 @@ pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
     let stride_b = (k * n) as i64;
     let stride_c = (m * n) as i64;
 
-    let ctx = CublasContext::new().unwrap();
+    with_cublas(|handle| unsafe {
-    unsafe {
+        cublasSetStream_v2(handle, std::ptr::null_mut());
-        cublasSetStream_v2(ctx.handle, std::ptr::null_mut());
         // Row-major trick: C = A @ B ⟺ C^T = B^T @ A^T (col-major)
         error::check(cublasGemmStridedBatchedEx(
-            ctx.handle,
+            handle,
             CUBLAS_OP_N, CUBLAS_OP_N,
             n as i32, m as i32, k as i32,
             &alpha as *const f32 as *const c_void,
@@ -223,7 +261,6 @@ pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
             CUBLAS_COMPUTE_32F,
             -1,
         )).expect("cuBLAS batched GEMM failed");
-    }
+    });
-    xserv_cuda::device::synchronize().unwrap();
     c
 }

crates/xserv-kernels/src/layernorm.rs

@@ -17,7 +17,9 @@ pub fn layernorm(x: &Tensor, gamma: &Tensor, beta: &Tensor, eps: f32) -> Tensor
     assert_eq!(beta.shape(), &[hidden_size]);
 
     let rows = x.numel() / hidden_size;
-    let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
+    assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
+    assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
 
     unsafe {
         match x.dtype() {
@@ -34,6 +36,5 @@ pub fn layernorm(x: &Tensor, gamma: &Tensor, beta: &Tensor, eps: f32) -> Tensor
             _ => panic!("unsupported dtype for layernorm"),
         }
     }
-    xserv_cuda::device::synchronize().unwrap();
     out
 }

crates/xserv-kernels/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod activation;
 pub mod attention;
+pub mod dispatch;
 pub mod embedding;
 pub mod gemm;
 pub mod layernorm;
@@ -8,12 +9,17 @@ pub mod rope;
 pub mod softmax;
 pub mod transpose;
 
-pub use activation::{add, gelu, mul, scale, silu};
+pub use activation::{add, gelu, mul, scale, silu, silu_mul};
-pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
+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;
+pub use attention::{attention, decode_attention, flash_attention, paged_decode_attention};
 pub use embedding::embedding;
 pub use gemm::{batched_matmul, matmul, GemmBackend};
 pub use layernorm::layernorm;
-pub use rmsnorm::rmsnorm;
+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

@@ -6,6 +6,9 @@ unsafe extern "C" {
                           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 {
@@ -17,7 +20,9 @@ pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
     assert_eq!(x.dtype(), gamma.dtype());
 
     let rows = x.numel() / hidden_size;
-    let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
+    assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
+    assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
 
     unsafe {
         match x.dtype() {
@@ -32,6 +37,43 @@ pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
             _ => panic!("unsupported dtype for rmsnorm"),
         }
     }
-    xserv_cuda::device::synchronize().unwrap();
     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;
+    assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
+    assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
+    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

@@ -34,7 +34,6 @@ impl RopeCache {
                 max_seq_len as i32, half_dim as i32, theta, std::ptr::null_mut(),
             );
         }
-        xserv_cuda::device::synchronize().unwrap();
 
         Self { cos, sin, max_seq_len, half_dim }
     }
@@ -59,7 +58,7 @@ pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) {
             num_tokens * std::mem::size_of::<u32>(),
         )
     };
-    let mut pos_gpu = GpuBuffer::alloc(pos_bytes.len()).expect("alloc positions");
+    let mut pos_gpu = xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions");
     pos_gpu.copy_from_host(pos_bytes).unwrap();
 
     unsafe {
@@ -81,5 +80,4 @@ pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) {
             _ => panic!("unsupported dtype for rope"),
         }
     }
-    xserv_cuda::device::synchronize().unwrap();
 }

crates/xserv-kernels/src/softmax.rs

@@ -14,7 +14,9 @@ pub fn softmax(x: &Tensor) -> Tensor {
 
     let cols = *x.shape().last().unwrap();
     let rows = x.numel() / cols;
-    let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
+    assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
+    assert!(cols <= i32::MAX as usize, "cols too large for i32 kernel param");
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
 
     unsafe {
         match x.dtype() {
@@ -29,6 +31,5 @@ pub fn softmax(x: &Tensor) -> Tensor {
             _ => panic!("unsupported dtype for softmax"),
         }
     }
-    xserv_cuda::device::synchronize().unwrap();
     out
 }

crates/xserv-kernels/src/transpose.rs

@@ -7,20 +7,27 @@ unsafe extern "C" {
     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::zeros(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
+    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(),
         );
     }
-    xserv_cuda::device::synchronize().unwrap();
     out
 }
 
@@ -29,14 +36,13 @@ pub fn merge_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: u
     assert_eq!(x.dtype(), DType::BF16);
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
     let hidden = num_heads * head_dim;
-    let out = Tensor::zeros(&[seq_len, hidden], DType::BF16, x.device());
+    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(),
         );
     }
-    xserv_cuda::device::synchronize().unwrap();
     out
 }
 
@@ -44,14 +50,13 @@ pub fn merge_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: u
 pub fn transpose_for_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
     assert_eq!(x.dtype(), DType::BF16);
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
-    let out = Tensor::zeros(&[seq_len, num_heads, head_dim], DType::BF16, x.device());
+    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(),
         );
     }
-    xserv_cuda::device::synchronize().unwrap();
     out
 }
 
@@ -59,14 +64,13 @@ pub fn transpose_for_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head
 pub fn transpose_from_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
     assert_eq!(x.dtype(), DType::BF16);
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
-    let out = Tensor::zeros(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
+    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(),
         );
     }
-    xserv_cuda::device::synchronize().unwrap();
     out
 }
 
@@ -79,13 +83,60 @@ pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
     let seq_len = x.shape()[2];
     let head_dim = x.shape()[3];
     let new_heads = kv_heads * n_rep;
-    let out = Tensor::zeros(&[1, new_heads, seq_len, head_dim], DType::BF16, x.device());
+    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(),
         );
     }
-    xserv_cuda::device::synchronize().unwrap();
+    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/gemm_test.rs

@@ -121,6 +121,20 @@ fn test_gemm_cublas_bf16_small() { run_gemm_test_bf16(GemmBackend::CuBlas, 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]

crates/xserv-kernels/tests/ops_test.rs

@@ -74,10 +74,10 @@ fn cpu_rope(x: &mut [f32], positions: &[u32], num_heads: usize, head_dim: usize,
                 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 x0 = x[base + i];
-                let x1 = x[base + 2 * i + 1];
+                let x1 = x[base + i + half_dim];
-                x[base + 2 * i]     = x0 * cos_val - x1 * sin_val;
+                x[base + i] = x0 * cos_val - x1 * sin_val;
-                x[base + 2 * i + 1] = x0 * sin_val + x1 * cos_val;
+                x[base + i + half_dim] = x1 * cos_val + x0 * sin_val;
             }
         }
     }

crates/xserv-model/Cargo.toml

@@ -8,8 +8,10 @@ xserv-cuda = { path = "../xserv-cuda" }
 xserv-tensor = { path = "../xserv-tensor" }
 xserv-kernels = { path = "../xserv-kernels" }
 xserv-tokenizer = { path = "../xserv-tokenizer" }
+xserv-distributed = { path = "../xserv-distributed" }
 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-qwen3.rs

@@ -1,14 +1,14 @@
 use std::path::PathBuf;
 use std::time::Instant;
 use xserv_model::qwen3::sample_greedy;
-use xserv_model::{loader, GpuKVCache, ModelConfig, Qwen3};
+use xserv_model::{loader, DecodeGraphState, 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]");
+        eprintln!("Usage: bench-qwen3 <model-dir> [--gen-tokens N] [--cuda-graph]");
         std::process::exit(1);
     }
     let model_dir = PathBuf::from(&args[1]);
@@ -18,6 +18,7 @@ fn main() {
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
         .unwrap_or(20);
+    let use_cuda_graph = args.iter().any(|a| a == "--cuda-graph");
 
     xserv_cuda::device::set_device(0).unwrap();
 
@@ -31,9 +32,21 @@ fn main() {
     // Warmup
     {
         let ids = tokenizer.encode("warmup");
-        let mut cache = GpuKVCache::new(&config, 256, DType::BF16);
+        let mut cache = GpuKVCache::new(&config, 256, DType::BF16, 0);
         let _ = model.forward_gpu_cache(&ids, &mut cache);
     }
+
+    // CUDA Graph setup
+    let layer_ptrs = model.layer_weight_ptrs();
+    let (norm_w, lm_head, embed, cos, sin) = model.graph_capture_ptrs();
+    let mut decode_graph = if use_cuda_graph {
+        eprintln!("CUDA Graph mode enabled");
+        Some(DecodeGraphState::new(&config))
+    } else {
+        None
+    };
+    let mut graph_captured = false;
+
     eprintln!("Warmup done. Running benchmark...");
 
     let prompts: Vec<&str> = vec![
@@ -94,7 +107,13 @@ fn main() {
         let input_ids = tokenizer.encode(prompt);
         let input_len = input_ids.len();
 
-        let mut cache = GpuKVCache::new(&config, 256, DType::BF16);
+        let mut cache = GpuKVCache::new(&config, 256, DType::BF16, 0);
+
+        // Reset graph state for new prompt
+        graph_captured = false;
+        if let Some(ref mut g) = decode_graph {
+            g.invalidate();
+        }
 
         // Prefill
         let t0 = Instant::now();
@@ -109,8 +128,35 @@ fn main() {
         for _ in 1..gen_tokens {
             let last = *generated.last().unwrap();
             let t_start = Instant::now();
+
+            let next = if let Some(ref mut graph) = decode_graph {
+                if !graph_captured {
+                    // First decode token: run ungraphed, then capture
                     let logits = model.forward_gpu_cache(&[last], &mut cache);
-            let next = sample_greedy(&logits);
+                    graph_captured = true;
+                    graph.capture(&layer_ptrs, norm_w, lm_head, embed, cos, sin);
+                    sample_greedy(&logits)
+                } else {
+                    // Replay captured graphs
+                    let pos = cache.seq_len() as u32;
+                    graph.execute(last, pos, &mut cache, &layer_ptrs, embed, config.vocab_size as i32, config.hidden() as i32);
+                    cache.advance_seq_len(1);
+                    // Read logits from graph buffer
+                    let vocab_size = config.vocab_size;
+                    let mut logits_bytes = vec![0u8; vocab_size * 2];
+                    graph.logits_buffer().copy_to_host(&mut logits_bytes).unwrap();
+                    let logits_data: &[half::bf16] = unsafe {
+                        std::slice::from_raw_parts(logits_bytes.as_ptr() as *const half::bf16, vocab_size)
+                    };
+                    logits_data.iter().enumerate()
+                        .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
+                        .map(|(idx, _)| idx as u32).unwrap()
+                }
+            } else {
+                let logits = model.forward_gpu_cache(&[last], &mut cache);
+                sample_greedy(&logits)
+            };
+
             token_times.push(t_start.elapsed().as_micros());
             generated.push(next);
             if tokenizer.eos_token_id() == Some(next) { break; }

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

@@ -0,0 +1,194 @@
+//! Tensor-parallel E2E benchmark for Qwen3.
+//!
+//! Spawns one thread per TP rank (each bound to one GPU), loads the sharded
+//! model, and runs greedy autoregressive generation. Because lm_head is
+//! replicated and the post-AllReduce hidden state is identical on every rank,
+//! all ranks compute identical logits and pick the same greedy token — so the
+//! rank threads stay in lockstep via the per-layer AllReduces without any
+//! token broadcast. Rank 0 records output + timings.
+//!
+//! Usage: bench-tp <model-dir> [--tp N] [--gen-tokens N] [--devices 0,1,2,3]
+//!
+//! Run with --tp 1 / 2 / 4 and compare the printed text (correctness) and
+//! tok/s (performance).
+
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::thread;
+use std::time::Instant;
+
+use xserv_model::qwen3::sample_greedy;
+use xserv_model::{loader, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
+use xserv_tensor::{DType, Device};
+use xserv_tokenizer::Tokenizer;
+
+struct PromptResult {
+    gen_ids: Vec<u32>,
+    ttft_ms: f64,
+    decode_tok_s: f64,
+}
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: bench-tp <model-dir> [--tp N] [--gen-tokens N] [--devices 0,1,2,3]");
+        std::process::exit(1);
+    }
+    let model_dir = PathBuf::from(&args[1]);
+    let world: usize = arg(&args, "--tp").and_then(|s| s.parse().ok()).unwrap_or(1).max(1);
+    let gen_tokens: usize = arg(&args, "--gen-tokens").and_then(|s| s.parse().ok()).unwrap_or(64);
+    let devices: Vec<u32> = match arg(&args, "--devices") {
+        Some(s) => s.split(',').filter_map(|d| d.trim().parse().ok()).collect(),
+        None => (0..world as u32).collect(),
+    };
+    assert_eq!(devices.len(), world, "--devices count must equal --tp");
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    assert!(
+        config.num_kv_heads() % world == 0,
+        "num_kv_heads {} not divisible by tp {world}",
+        config.num_kv_heads()
+    );
+    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+    let eos = tokenizer.eos_token_id();
+
+    let prompts: Vec<&str> = vec![
+        "The capital of France is",
+        "Explain photosynthesis in one sentence.",
+        "Write a haiku about the ocean.",
+        "List three uses of a hammer.",
+        "What is the speed of light?",
+        "Describe the water cycle briefly.",
+        "Who wrote Romeo and Juliet?",
+        "Translate 'good morning' into Spanish.",
+    ];
+    let prompt_ids: Vec<Vec<u32>> = prompts.iter().map(|p| tokenizer.encode(p)).collect();
+
+    // Tensors are not Send (their Storage holds a raw GPU pointer), so each rank
+    // thread loads its own CPU copy of the weights and shards in-thread. Loading
+    // is not part of the timed region.
+    let id = if world > 1 { Some(xserv_distributed::get_unique_id()) } else { None };
+
+    let handles: Vec<_> = (0..world)
+        .map(|rank| {
+            let model_dir = model_dir.clone();
+            let config = config.clone();
+            let prompt_ids = prompt_ids.clone();
+            let device = devices[rank];
+            thread::spawn(move || {
+                run_rank(rank, world, device, id, config, model_dir, prompt_ids, gen_tokens, eos)
+            })
+        })
+        .collect();
+
+    let mut rank0: Option<Vec<PromptResult>> = None;
+    for (rank, h) in handles.into_iter().enumerate() {
+        let r = h.join().expect("rank thread panicked");
+        if rank == 0 {
+            rank0 = r;
+        }
+    }
+
+    let results = rank0.expect("rank 0 produced no results");
+    println!("\n=== TP={world} (devices {devices:?}) — Qwen3 E2E benchmark ===");
+    println!("{:<45} {:>10} {:>12} {:>8}", "prompt", "TTFT(ms)", "decode tok/s", "gen");
+    let mut tps_sum = 0.0;
+    for (i, r) in results.iter().enumerate() {
+        let text = tokenizer.decode(&r.gen_ids).replace('\n', " ");
+        let short: String = text.chars().take(50).collect();
+        let p: String = prompts[i].chars().take(43).collect();
+        println!(
+            "{:<45} {:>10.1} {:>12.1} {:>8} | {}",
+            p, r.ttft_ms, r.decode_tok_s, r.gen_ids.len(), short
+        );
+        tps_sum += r.decode_tok_s;
+    }
+    println!("--- mean decode throughput: {:.1} tok/s ---", tps_sum / results.len() as f64);
+
+    // Machine-readable line for cross-TP correctness diffing (rank 0 token ids).
+    let all_ids: Vec<String> = results
+        .iter()
+        .map(|r| r.gen_ids.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(","))
+        .collect();
+    println!("CORRECTNESS_IDS tp={world} {}", all_ids.join(" | "));
+}
+
+fn run_rank(
+    rank: usize,
+    world: usize,
+    device: u32,
+    id: Option<xserv_distributed::UniqueId>,
+    config: ModelConfig,
+    model_dir: PathBuf,
+    prompt_ids: Vec<Vec<u32>>,
+    gen_tokens: usize,
+    eos: Option<u32>,
+) -> Option<Vec<PromptResult>> {
+    // Bind this thread to its GPU and set up the TP communicator.
+    let tp = if world > 1 {
+        Some(Arc::new(xserv_distributed::TpContext::init(rank, world, id.unwrap(), device)))
+    } else {
+        xserv_cuda::device::set_device(device).unwrap();
+        None
+    };
+
+    // Load this rank's own CPU copy of the weights and shard in-thread.
+    let weights = loader::load_model_dir(&model_dir, Device::Cpu);
+    let model = Qwen3::from_weights_tp(config.clone(), weights, rank, world, device, tp.clone());
+
+    // Per-rank paged KV cache holds only this rank's local KV heads.
+    let local_kv = config.num_kv_heads() / world;
+    let max_seq = 2048usize;
+    let max_blocks_per_seq = max_seq.div_ceil(BLOCK_SIZE);
+    let total_blocks = max_blocks_per_seq + 8;
+    let mut cache = PagedKVCache::new_tp(
+        &config, local_kv, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, device,
+    );
+
+    // Warmup (init kernels / allocator / NCCL channels) — not timed.
+    cache.register_sequence(0).unwrap();
+    let _ = model.forward_prefill_paged(&[1u32, 2, 3], 0, &mut cache);
+    cache.free_sequence(0);
+
+    let mut out = Vec::new();
+    for ids in &prompt_ids {
+        cache.register_sequence(0).unwrap();
+
+        // Prefill (TTFT).
+        let t0 = Instant::now();
+        let logits = model.forward_prefill_paged(ids, 0, &mut cache);
+        let first = sample_greedy(&logits);
+        let ttft_ms = t0.elapsed().as_secs_f64() * 1000.0;
+
+        let mut generated = vec![first];
+
+        // Decode.
+        let t1 = Instant::now();
+        let mut steps = 0usize;
+        for _ in 1..gen_tokens {
+            let last = *generated.last().unwrap();
+            if eos == Some(last) {
+                break;
+            }
+            let pos = cache.seq_len(0);
+            let logits = model.forward_decode_paged(&[last], &[pos], &[0], &mut cache);
+            let next = sample_greedy(&logits);
+            generated.push(next);
+            steps += 1;
+        }
+        let decode_s = t1.elapsed().as_secs_f64();
+        let decode_tok_s = if steps > 0 && decode_s > 0.0 { steps as f64 / decode_s } else { 0.0 };
+
+        cache.free_sequence(0);
+
+        if rank == 0 {
+            out.push(PromptResult { gen_ids: generated, ttft_ms, decode_tok_s });
+        }
+    }
+
+    if rank == 0 { Some(out) } else { None }
+}
+
+fn arg<'a>(args: &'a [String], flag: &str) -> Option<&'a str> {
+    args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).map(|s| s.as_str())
+}

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

@@ -0,0 +1,419 @@
+use std::io::{self, IsTerminal, Write};
+use std::path::PathBuf;
+
+use xserv_model::{loader, sample, ModelConfig, PagedKVCache, Qwen3, SamplingParams, BLOCK_SIZE};
+use xserv_tensor::{DType, Device};
+use xserv_tokenizer::Tokenizer;
+
+const SLOT: usize = 0;
+
+struct CliOptions {
+    model_dir: PathBuf,
+    max_tokens: usize,
+    max_seq_len: usize,
+    sampling: SamplingParams,
+    system_prompt: Option<String>,
+    enable_thinking: bool,
+    color: bool,
+}
+
+enum Finish {
+    Stop { token_id: u32 },
+    Length,
+}
+
+fn main() {
+    let opts = parse_args();
+
+    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(&opts.model_dir.join("config.json"));
+    let model_type = config.model_type.as_deref().unwrap_or("unknown");
+    if !model_type.contains("qwen") {
+        eprintln!("xserv-chat currently supports Qwen-style ChatML models only; got model_type={model_type}");
+        std::process::exit(2);
+    }
+
+    let max_seq_len = opts.max_seq_len.min(config.max_seq_len()).max(1);
+    eprintln!(
+        "Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}, max_seq_len={}",
+        config.num_layers(),
+        config.hidden(),
+        config.num_heads(),
+        config.num_kv_heads(),
+        config.vocab_size,
+        max_seq_len
+    );
+
+    eprintln!("Loading weights...");
+    let weights = loader::load_model_dir(&opts.model_dir, Device::Cuda(0));
+    eprintln!("Loaded {} tensors", weights.len());
+    let model = Qwen3::from_weights(config.clone(), weights);
+
+    let tokenizer = Tokenizer::from_file(&opts.model_dir.join("tokenizer.json"));
+    let mut cache = new_paged_cache(&config, max_seq_len);
+    cache.register_sequence(SLOT).expect("register chat slot");
+    let use_color = opts.color && io::stdout().is_terminal();
+
+    eprintln!("Ready (paged KV cache, persistent chat slot).");
+    eprintln!("Commands: /exit, /quit, /clear\n");
+
+    loop {
+        print!("user> ");
+        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;
+        }
+
+        match input {
+            "/exit" | "/quit" | "exit" | "quit" => break,
+            "/clear" => {
+                cache.free_sequence(SLOT);
+                cache.register_sequence(SLOT).expect("register chat slot");
+                eprintln!("history and KV cache cleared");
+                continue;
+            }
+            "/help" => {
+                print_help();
+                continue;
+            }
+            _ => {}
+        }
+
+        let include_system = cache.seq_len(SLOT) == 0;
+        let prompt = build_turn_prompt(
+            opts.system_prompt.as_deref(),
+            include_system,
+            input,
+            opts.enable_thinking,
+        );
+        let prompt_tokens = tokenizer.encode(&prompt);
+        if prompt_tokens.is_empty() {
+            continue;
+        }
+
+        let used = cache.seq_len(SLOT);
+        let remaining = max_seq_len.saturating_sub(used);
+        if prompt_tokens.len() >= remaining {
+            eprintln!(
+                "context full: {used}/{max_seq_len} tokens used, new turn needs {} tokens; use /clear",
+                prompt_tokens.len()
+            );
+            continue;
+        }
+        let max_new_tokens = opts.max_tokens.min(remaining - prompt_tokens.len());
+
+        print!("assistant> ");
+        io::stdout().flush().unwrap();
+        let finish = generate_with_paged_cache(
+            &model,
+            &mut cache,
+            &tokenizer,
+            &prompt_tokens,
+            &opts.sampling,
+            max_new_tokens,
+            use_color,
+        );
+        match finish {
+            Finish::Stop { token_id } => {
+                append_after_stop(&model, &mut cache, &tokenizer, max_seq_len, token_id);
+            }
+            Finish::Length => {
+                append_text_to_cache(&model, &mut cache, &tokenizer, max_seq_len, "<|im_end|>\n");
+            }
+        }
+        println!();
+    }
+}
+
+fn parse_args() -> CliOptions {
+    let args: Vec<String> = std::env::args().skip(1).collect();
+    if args.is_empty() || args.iter().any(|a| a == "--help" || a == "-h") {
+        print_usage_and_exit(0);
+    }
+
+    let mut model_dir = None;
+    let mut max_tokens = 256usize;
+    let mut max_seq_len = 2048usize;
+    let mut temperature = 0.0f32;
+    let mut top_k = 0usize;
+    let mut top_p = 1.0f32;
+    let mut system_prompt = None;
+    let mut enable_thinking = false;
+    let mut color = true;
+
+    let mut i = 0;
+    while i < args.len() {
+        match args[i].as_str() {
+            "-m" | "--model" => {
+                i += 1;
+                model_dir = args.get(i).map(PathBuf::from);
+            }
+            "--max-tokens" => {
+                i += 1;
+                max_tokens = parse_value(&args, i, "--max-tokens");
+            }
+            "--max-seq-len" => {
+                i += 1;
+                max_seq_len = parse_value(&args, i, "--max-seq-len");
+            }
+            "--temperature" => {
+                i += 1;
+                temperature = parse_value(&args, i, "--temperature");
+            }
+            "--top-k" => {
+                i += 1;
+                top_k = parse_value(&args, i, "--top-k");
+            }
+            "--top-p" => {
+                i += 1;
+                top_p = parse_value(&args, i, "--top-p");
+            }
+            "--system" => {
+                i += 1;
+                system_prompt = args.get(i).cloned();
+                if system_prompt.is_none() {
+                    eprintln!("missing value for --system");
+                    std::process::exit(2);
+                }
+            }
+            "--think" => {
+                enable_thinking = true;
+            }
+            "--no-color" => {
+                color = false;
+            }
+            arg if arg.starts_with('-') => {
+                eprintln!("unknown option: {arg}");
+                print_usage_and_exit(2);
+            }
+            arg => {
+                if model_dir.is_some() {
+                    eprintln!("unexpected extra argument: {arg}");
+                    print_usage_and_exit(2);
+                }
+                model_dir = Some(PathBuf::from(arg));
+            }
+        }
+        i += 1;
+    }
+
+    CliOptions {
+        model_dir: model_dir.unwrap_or_else(|| {
+            eprintln!("missing model directory");
+            print_usage_and_exit(2);
+        }),
+        max_tokens: max_tokens.max(1),
+        max_seq_len: max_seq_len.max(1),
+        sampling: SamplingParams {
+            temperature,
+            top_k,
+            top_p,
+        },
+        system_prompt,
+        enable_thinking,
+        color,
+    }
+}
+
+fn parse_value<T: std::str::FromStr>(args: &[String], i: usize, name: &str) -> T {
+    args.get(i).and_then(|s| s.parse().ok()).unwrap_or_else(|| {
+        eprintln!("invalid or missing value for {name}");
+        std::process::exit(2);
+    })
+}
+
+fn print_usage_and_exit(code: i32) -> ! {
+    eprintln!(
+        "Usage: xserv-chat <model-dir> [options]\n\
+         \n\
+         Options:\n\
+         \t-m, --model DIR       Model directory\n\
+         \t--max-tokens N        Max generated tokens per turn (default: 256)\n\
+         \t--max-seq-len N       Persistent KV context length (default: 2048)\n\
+         \t--temperature F       Sampling temperature, 0 = greedy (default: 0)\n\
+         \t--top-k N             Top-k sampling, 0 = disabled (default: 0)\n\
+         \t--top-p F             Top-p sampling (default: 1.0)\n\
+         \t--system TEXT         System prompt for the first turn after start or /clear\n\
+         \t--think               Let Qwen3 emit thinking; rendered gray on terminals\n\
+         \t--no-color            Disable ANSI color for thinking output\n\
+         \t-h, --help            Show this help"
+    );
+    std::process::exit(code);
+}
+
+fn print_help() {
+    eprintln!("Commands:");
+    eprintln!("  /clear  clear chat history and free/recreate the paged KV slot");
+    eprintln!("  /exit   quit");
+    eprintln!("  /quit   quit");
+}
+
+fn new_paged_cache(config: &ModelConfig, max_seq_len: usize) -> PagedKVCache {
+    let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
+    let total_blocks = (max_blocks_per_seq + 1).max(2);
+    // Single-slot interactive CLI: no swap pool (cpu_total_blocks = 0).
+    PagedKVCache::new(config, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, 0)
+}
+
+fn build_turn_prompt(
+    system: Option<&str>,
+    include_system: bool,
+    user_input: &str,
+    enable_thinking: bool,
+) -> String {
+    let mut prompt = String::new();
+    if include_system {
+        if let Some(system) = system {
+            if !system.trim().is_empty() {
+                prompt.push_str("<|im_start|>system\n");
+                prompt.push_str(system.trim());
+                prompt.push_str("<|im_end|>\n");
+            }
+        }
+    }
+    prompt.push_str("<|im_start|>user\n");
+    prompt.push_str(user_input);
+    prompt.push_str("<|im_end|>\n");
+    prompt.push_str("<|im_start|>assistant\n");
+    if !enable_thinking {
+        prompt.push_str("<think>\n\n</think>\n\n");
+    }
+    prompt
+}
+
+fn generate_with_paged_cache(
+    model: &Qwen3,
+    cache: &mut PagedKVCache,
+    tokenizer: &Tokenizer,
+    prompt_tokens: &[u32],
+    sampling: &SamplingParams,
+    max_tokens: usize,
+    use_color: bool,
+) -> Finish {
+    let logits = model.forward_prefill_paged(prompt_tokens, SLOT, cache);
+    let mut next = sample(&logits, sampling);
+    let mut decode_buffer = Vec::new();
+    let mut in_thinking = false;
+
+    for _ in 0..max_tokens {
+        let position = cache.seq_len(SLOT);
+        let logits = model.forward_decode_paged(&[next], &[position], &[SLOT], cache);
+        if is_stop_token(tokenizer, next) {
+            print_stream_text(
+                &tokenizer.flush_decode_stream(&mut decode_buffer),
+                in_thinking,
+                use_color,
+            );
+            io::stdout().flush().unwrap();
+            return Finish::Stop { token_id: next };
+        }
+
+        print_generated_token(
+            tokenizer,
+            next,
+            &mut decode_buffer,
+            &mut in_thinking,
+            use_color,
+        );
+        io::stdout().flush().unwrap();
+        next = sample(&logits, sampling);
+    }
+
+    print_stream_text(
+        &tokenizer.flush_decode_stream(&mut decode_buffer),
+        in_thinking,
+        use_color,
+    );
+    io::stdout().flush().unwrap();
+    Finish::Length
+}
+
+fn append_after_stop(
+    model: &Qwen3,
+    cache: &mut PagedKVCache,
+    tokenizer: &Tokenizer,
+    max_seq_len: usize,
+    stop_token_id: u32,
+) {
+    if tokenizer.special_token_id("<|im_end|>") == Some(stop_token_id) {
+        append_text_to_cache(model, cache, tokenizer, max_seq_len, "\n");
+    }
+}
+
+fn append_text_to_cache(
+    model: &Qwen3,
+    cache: &mut PagedKVCache,
+    tokenizer: &Tokenizer,
+    max_seq_len: usize,
+    text: &str,
+) {
+    let tokens = tokenizer.encode(text);
+    if tokens.is_empty() || cache.seq_len(SLOT) + tokens.len() > max_seq_len {
+        return;
+    }
+    let _ = model.forward_prefill_paged(&tokens, SLOT, cache);
+}
+
+fn print_generated_token(
+    tokenizer: &Tokenizer,
+    token_id: u32,
+    decode_buffer: &mut Vec<u8>,
+    in_thinking: &mut bool,
+    use_color: bool,
+) {
+    if tokenizer.special_token_id("<think>") == Some(token_id) {
+        print_stream_text(
+            &tokenizer.flush_decode_stream(decode_buffer),
+            *in_thinking,
+            use_color,
+        );
+        *in_thinking = true;
+        print_stream_text("<think>", true, use_color);
+        return;
+    }
+
+    if tokenizer.special_token_id("</think>") == Some(token_id) {
+        print_stream_text(
+            &tokenizer.flush_decode_stream(decode_buffer),
+            *in_thinking,
+            use_color,
+        );
+        print_stream_text("</think>", true, use_color);
+        *in_thinking = false;
+        return;
+    }
+
+    let text = tokenizer.decode_token_stream(token_id, decode_buffer);
+    print_stream_text(&text, *in_thinking, use_color);
+}
+
+fn print_stream_text(text: &str, in_thinking: bool, use_color: bool) {
+    if text.is_empty() {
+        return;
+    }
+    if in_thinking && use_color {
+        print!("\x1b[90m{text}\x1b[0m");
+    } else {
+        print!("{text}");
+    }
+}
+
+fn is_stop_token(tokenizer: &Tokenizer, token_id: u32) -> bool {
+    tokenizer.eos_token_id() == Some(token_id)
+        || tokenizer.special_token_id("<|im_end|>") == Some(token_id)
+        || tokenizer.special_token_id("<|endoftext|>") == Some(token_id)
+        || tokenizer.special_token_id("<|end_of_text|>") == Some(token_id)
+}

crates/xserv-model/src/decode_graph.rs

@@ -0,0 +1,458 @@
+//! CUDA Graph integration for batch=1 single-sequence decode.
+//!
+//! Uses a per-layer split graph approach:
+//! - Pre-attention graph: RMSNorm + QKV projections + reshape + QK-norm + RoPE
+//! - Ungraphed: KV cache append + decode attention (variable kv_len)
+//! - Post-attention graph: merge_heads + O-proj + add_rmsnorm + FFN + residual
+//! - Final graph: last RMSNorm + lm_head GEMV
+
+use std::ffi::c_void;
+use xserv_cuda::{CudaGraph, CudaStream, GpuBuffer};
+use xserv_kernels::dispatch;
+use xserv_kernels::gemm::cublas_handle;
+
+use crate::config::ModelConfig;
+use crate::kv_cache::GpuKVCache;
+
+/// Pre-allocated intermediate buffers for decode (batch=1).
+/// All buffers have stable GPU addresses for CUDA Graph replay.
+struct DecodeBuffers {
+    // Hidden-size buffers: [1, hidden]
+    x: GpuBuffer,               // running hidden state
+    normed: GpuBuffer,           // rmsnorm output
+    attn_out: GpuBuffer,         // attention output [1, num_heads, 1, head_dim]
+    attn_merged: GpuBuffer,      // merge_heads output [1, hidden]
+    o_proj: GpuBuffer,           // O projection output [1, hidden]
+    normed2: GpuBuffer,          // post-attn norm output [1, hidden]
+    sum_out: GpuBuffer,          // add_rmsnorm sum output [1, hidden]
+    down: GpuBuffer,             // down projection output [1, hidden]
+
+    // QKV projection outputs
+    q_proj: GpuBuffer,           // [1, num_heads * head_dim]
+    k_proj: GpuBuffer,           // [1, num_kv_heads * head_dim]
+    v_proj: GpuBuffer,           // [1, num_kv_heads * head_dim]
+
+    // Reshaped: [1, H, 1, D]
+    q_reshaped: GpuBuffer,
+    k_reshaped: GpuBuffer,
+    v_reshaped: GpuBuffer,
+
+    // After QK-norm (same shape as reshaped)
+    q_normed: GpuBuffer,
+    k_normed: GpuBuffer,
+
+    // RoPE transposed: [1, H, D]
+    q_rope: GpuBuffer,
+    k_rope: GpuBuffer,
+
+    // After RoPE transpose back: [1, H, 1, D]
+    q_final: GpuBuffer,
+    k_final: GpuBuffer,
+
+    // FFN intermediates
+    gate: GpuBuffer,             // [1, intermediate]
+    up: GpuBuffer,               // [1, intermediate]
+    silu_out: GpuBuffer,         // [1, intermediate]
+
+    // GEMV fp32 accumulators (separate per output dimension)
+    fp32_hidden: GpuBuffer,      // for hidden-sized GEMV outputs
+    fp32_q: GpuBuffer,           // for Q projection
+    fp32_kv: GpuBuffer,          // for K/V projection
+    fp32_intermediate: GpuBuffer,// for gate/up projections
+    fp32_vocab: GpuBuffer,       // for lm_head
+
+    // Token ID and position (GPU-resident, updated before replay)
+    token_id_gpu: GpuBuffer,     // 4 bytes (u32)
+    position_gpu: GpuBuffer,     // 4 bytes (u32)
+
+    // Final output
+    logits: GpuBuffer,           // [1, vocab_size]
+}
+
+pub struct DecodeGraphState {
+    stream: CudaStream,
+    buffers: DecodeBuffers,
+
+    // Per-layer graph pairs
+    pre_attn_graphs: Vec<CudaGraph>,
+    post_attn_graphs: Vec<CudaGraph>,
+    final_graph: CudaGraph,
+
+    captured: bool,
+
+    // Model dimensions
+    hidden: usize,
+    num_heads: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    intermediate: usize,
+    vocab_size: usize,
+    num_layers: usize,
+    eps: f32,
+}
+
+impl DecodeGraphState {
+    pub fn new(config: &ModelConfig) -> Self {
+        let hidden = config.hidden();
+        let num_heads = config.num_heads();
+        let num_kv_heads = config.num_kv_heads();
+        let head_dim = config.head_dim();
+        let intermediate = config.ffn_hidden();
+        let vocab_size = config.vocab_size;
+        let num_layers = config.num_layers();
+        let eps = config.rms_norm_eps.unwrap_or(1e-6) as f32;
+        let es = 2usize; // BF16 = 2 bytes
+
+        let stream = CudaStream::new().expect("create CUDA stream for graph");
+
+        let alloc = |size: usize| -> GpuBuffer {
+            GpuBuffer::alloc(size).expect("alloc decode graph buffer")
+        };
+
+        let buffers = DecodeBuffers {
+            x: alloc(hidden * es),
+            normed: alloc(hidden * es),
+            attn_out: alloc(num_heads * head_dim * es),
+            attn_merged: alloc(hidden * es),
+            o_proj: alloc(hidden * es),
+            normed2: alloc(hidden * es),
+            sum_out: alloc(hidden * es),
+            down: alloc(hidden * es),
+
+            q_proj: alloc(num_heads * head_dim * es),
+            k_proj: alloc(num_kv_heads * head_dim * es),
+            v_proj: alloc(num_kv_heads * head_dim * es),
+
+            q_reshaped: alloc(num_heads * head_dim * es),
+            k_reshaped: alloc(num_kv_heads * head_dim * es),
+            v_reshaped: alloc(num_kv_heads * head_dim * es),
+
+            q_normed: alloc(num_heads * head_dim * es),
+            k_normed: alloc(num_kv_heads * head_dim * es),
+
+            q_rope: alloc(num_heads * head_dim * es),
+            k_rope: alloc(num_kv_heads * head_dim * es),
+
+            q_final: alloc(num_heads * head_dim * es),
+            k_final: alloc(num_kv_heads * head_dim * es),
+
+            gate: alloc(intermediate * es),
+            up: alloc(intermediate * es),
+            silu_out: alloc(intermediate * es),
+
+            fp32_hidden: alloc(hidden * 4),
+            fp32_q: alloc(num_heads * head_dim * 4),
+            fp32_kv: alloc(num_kv_heads * head_dim * 4),
+            fp32_intermediate: alloc(intermediate * 4),
+            fp32_vocab: alloc(vocab_size * 4),
+
+            token_id_gpu: alloc(4),
+            position_gpu: alloc(4),
+
+            logits: alloc(vocab_size * es),
+        };
+
+        let pre_attn_graphs = (0..num_layers).map(|_| CudaGraph::new()).collect();
+        let post_attn_graphs = (0..num_layers).map(|_| CudaGraph::new()).collect();
+
+        Self {
+            stream,
+            buffers,
+            pre_attn_graphs,
+            post_attn_graphs,
+            final_graph: CudaGraph::new(),
+            captured: false,
+            hidden,
+            num_heads,
+            num_kv_heads,
+            head_dim,
+            intermediate,
+            vocab_size,
+            num_layers,
+            eps,
+        }
+    }
+
+    pub fn is_captured(&self) -> bool {
+        self.captured
+    }
+
+    /// Capture all per-layer graphs. Called once after the first decode step.
+    pub fn capture(
+        &mut self,
+        layers: &[LayerWeightPtrs],
+        norm_weight: *const c_void,
+        lm_head_wt: *const c_void,
+        _embed_table: *const c_void,
+        rope_cos: *const c_void,
+        rope_sin: *const c_void,
+    ) {
+        let s = self.stream.as_raw();
+        let h = self.hidden as i32;
+        let nh = self.num_heads as i32;
+        let nkv = self.num_kv_heads as i32;
+        let hd = self.head_dim as i32;
+        let inter = self.intermediate as i32;
+        let vocab = self.vocab_size as i32;
+        let eps = self.eps;
+
+        let cublas = cublas_handle();
+
+        // Set cuBLAS to use our stream
+        unsafe { dispatch::set_cublas_stream(cublas, s); }
+
+        for (l, lw) in layers.iter().enumerate() {
+            // === Pre-attention graph ===
+            self.pre_attn_graphs[l].begin_capture(&self.stream).expect("begin pre-attn capture");
+            unsafe {
+                // RMSNorm
+                dispatch::rmsnorm_bf16(
+                    self.buffers.x.as_ptr() as _, lw.input_norm, self.buffers.normed.as_mut_ptr() as _,
+                    1, h, eps, s,
+                );
+
+                // Q projection (GEMV)
+                dispatch::gemv_bf16(
+                    self.buffers.normed.as_ptr() as _, lw.q_proj_wt, self.buffers.q_proj.as_mut_ptr() as _,
+                    self.buffers.fp32_q.as_mut_ptr() as _,
+                    h, nh * hd, s,
+                );
+
+                // K projection (GEMV)
+                dispatch::gemv_bf16(
+                    self.buffers.normed.as_ptr() as _, lw.k_proj_wt, self.buffers.k_proj.as_mut_ptr() as _,
+                    self.buffers.fp32_kv.as_mut_ptr() as _,
+                    h, nkv * hd, s,
+                );
+
+                // V projection (GEMV)
+                dispatch::gemv_bf16(
+                    self.buffers.normed.as_ptr() as _, lw.v_proj_wt, self.buffers.v_proj.as_mut_ptr() as _,
+                    self.buffers.fp32_kv.as_mut_ptr() as _,
+                    h, nkv * hd, s,
+                );
+
+                // Reshape heads: [1, H*D] -> [1, H, 1, D]
+                dispatch::reshape_heads_bf16(self.buffers.q_proj.as_ptr() as _, self.buffers.q_reshaped.as_mut_ptr() as _, 1, nh, hd, s);
+                dispatch::reshape_heads_bf16(self.buffers.k_proj.as_ptr() as _, self.buffers.k_reshaped.as_mut_ptr() as _, 1, nkv, hd, s);
+                dispatch::reshape_heads_bf16(self.buffers.v_proj.as_ptr() as _, self.buffers.v_reshaped.as_mut_ptr() as _, 1, nkv, hd, s);
+
+                // QK norm (head-level rmsnorm: treat [1,H,1,D] as [H, D])
+                dispatch::rmsnorm_bf16(self.buffers.q_reshaped.as_ptr() as _, lw.q_norm, self.buffers.q_normed.as_mut_ptr() as _, nh, hd, eps, s);
+                dispatch::rmsnorm_bf16(self.buffers.k_reshaped.as_ptr() as _, lw.k_norm, self.buffers.k_normed.as_mut_ptr() as _, nkv, hd, eps, s);
+
+                // Transpose for RoPE: [1,H,1,D] -> [1,H,D]
+                dispatch::transpose_hsd_to_shd_bf16(self.buffers.q_normed.as_ptr() as _, self.buffers.q_rope.as_mut_ptr() as _, 1, nh, hd, s);
+                dispatch::transpose_hsd_to_shd_bf16(self.buffers.k_normed.as_ptr() as _, self.buffers.k_rope.as_mut_ptr() as _, 1, nkv, hd, s);
+
+                // RoPE (in-place, reads position_gpu)
+                dispatch::rope_bf16(self.buffers.q_rope.as_mut_ptr() as _, rope_cos, rope_sin, self.buffers.position_gpu.as_ptr() as _, 1, nh, hd, s);
+                dispatch::rope_bf16(self.buffers.k_rope.as_mut_ptr() as _, rope_cos, rope_sin, self.buffers.position_gpu.as_ptr() as _, 1, nkv, hd, s);
+
+                // Transpose back: [1,H,D] -> [1,H,1,D]
+                dispatch::transpose_shd_to_hsd_bf16(self.buffers.q_rope.as_ptr() as _, self.buffers.q_final.as_mut_ptr() as _, 1, nh, hd, s);
+                dispatch::transpose_shd_to_hsd_bf16(self.buffers.k_rope.as_ptr() as _, self.buffers.k_final.as_mut_ptr() as _, 1, nkv, hd, s);
+            }
+            self.pre_attn_graphs[l].end_capture(&self.stream).expect("end pre-attn capture");
+
+            // === Post-attention graph ===
+            self.post_attn_graphs[l].begin_capture(&self.stream).expect("begin post-attn capture");
+            unsafe {
+                // Merge heads: [1,H,1,D] -> [1, hidden]
+                // attn_out is written by ungraphed attention
+                dispatch::merge_heads_bf16(self.buffers.attn_out.as_ptr() as _, self.buffers.attn_merged.as_mut_ptr() as _, 1, nh, hd, s);
+
+                // O projection
+                dispatch::gemv_bf16(
+                    self.buffers.attn_merged.as_ptr() as _, lw.o_proj_wt, self.buffers.o_proj.as_mut_ptr() as _,
+                    self.buffers.fp32_hidden.as_mut_ptr() as _,
+                    nh * hd, h, s,
+                );
+
+                // Fused Add+RMSNorm: normed2 = rmsnorm(o_proj + x), sum_out = o_proj + x
+                dispatch::add_rmsnorm_bf16(
+                    self.buffers.o_proj.as_ptr() as _, self.buffers.x.as_ptr() as _, lw.post_norm,
+                    self.buffers.normed2.as_mut_ptr() as _, self.buffers.sum_out.as_mut_ptr() as _,
+                    1, h, eps, s,
+                );
+
+                // Gate projection
+                dispatch::gemv_bf16(
+                    self.buffers.normed2.as_ptr() as _, lw.gate_proj_wt, self.buffers.gate.as_mut_ptr() as _,
+                    self.buffers.fp32_intermediate.as_mut_ptr() as _,
+                    h, inter, s,
+                );
+
+                // Up projection
+                dispatch::gemv_bf16(
+                    self.buffers.normed2.as_ptr() as _, lw.up_proj_wt, self.buffers.up.as_mut_ptr() as _,
+                    self.buffers.fp32_intermediate.as_mut_ptr() as _,
+                    h, inter, s,
+                );
+
+                // Fused SiLU x Mul
+                dispatch::silu_mul_bf16(self.buffers.gate.as_ptr() as _, self.buffers.up.as_ptr() as _, self.buffers.silu_out.as_mut_ptr() as _, inter, s);
+
+                // Down projection
+                dispatch::gemv_bf16(
+                    self.buffers.silu_out.as_ptr() as _, lw.down_proj_wt, self.buffers.down.as_mut_ptr() as _,
+                    self.buffers.fp32_hidden.as_mut_ptr() as _,
+                    inter, h, s,
+                );
+
+                // x = sum_out + down (residual connection for next layer)
+                dispatch::add_bf16(self.buffers.sum_out.as_ptr() as _, self.buffers.down.as_ptr() as _, self.buffers.x.as_mut_ptr() as _, h, s);
+            }
+            self.post_attn_graphs[l].end_capture(&self.stream).expect("end post-attn capture");
+        }
+
+        // === Final graph: norm + lm_head ===
+        self.final_graph.begin_capture(&self.stream).expect("begin final capture");
+        unsafe {
+            dispatch::rmsnorm_bf16(self.buffers.x.as_ptr() as _, norm_weight, self.buffers.normed.as_mut_ptr() as _, 1, h, eps, s);
+            dispatch::gemv_bf16(
+                self.buffers.normed.as_ptr() as _, lm_head_wt, self.buffers.logits.as_mut_ptr() as _,
+                self.buffers.fp32_vocab.as_mut_ptr() as _,
+                h, vocab, s,
+            );
+        }
+        self.final_graph.end_capture(&self.stream).expect("end final capture");
+
+        // Reset cuBLAS back to null stream
+        unsafe { dispatch::set_cublas_stream(cublas, std::ptr::null_mut()); }
+
+        self.captured = true;
+    }
+
+    /// Execute a single decode step using captured graphs.
+    pub fn execute(
+        &mut self,
+        token_id: u32,
+        position: u32,
+        cache: &mut GpuKVCache,
+        _layers: &[LayerWeightPtrs],
+        embed_table: *const c_void,
+        vocab_size: i32,
+        hidden_size: i32,
+    ) {
+        assert!(self.captured, "must call capture() before execute()");
+        let s = self.stream.as_raw();
+        let nkv = self.num_kv_heads;
+        let nh = self.num_heads;
+        let hd = self.head_dim;
+        let es = 2usize; // BF16
+
+        // Upload token ID and position to fixed GPU buffers
+        self.buffers.token_id_gpu.copy_from_host(&token_id.to_le_bytes()).unwrap();
+        self.buffers.position_gpu.copy_from_host(&position.to_le_bytes()).unwrap();
+
+        // Embedding (outside graph since token_id changes each step)
+        unsafe {
+            dispatch::embedding_bf16(
+                embed_table,
+                self.buffers.token_id_gpu.as_ptr() as _,
+                self.buffers.x.as_mut_ptr() as _,
+                1, hidden_size, vocab_size, s,
+            );
+        }
+
+        for l in 0..self.num_layers {
+            // Pre-attention graph (norm + QKV + reshape + QK-norm + RoPE)
+            self.pre_attn_graphs[l].launch(&self.stream).expect("launch pre-attn graph");
+
+            // Ungraphed: KV cache append
+            // k_final shape: [1, num_kv_heads, 1, head_dim] (after RoPE pipeline)
+            // v_reshaped shape: [1, num_kv_heads, 1, head_dim] (V skips RoPE)
+            let pos = position as usize;
+
+            let k_buf_size = nkv * hd * es;
+            let v_buf_size = nkv * hd * es;
+            let shape = [1usize, nkv, 1, hd];
+
+            // Synchronize before accessing buffers for KV cache append
+            self.stream.synchronize().expect("sync before kv cache");
+
+            let k_view = unsafe {
+                crate::kv_cache::tensor_from_gpu_buffer_pub(
+                    GpuBuffer::borrow_raw(self.buffers.k_final.as_mut_ptr(), k_buf_size),
+                    &shape,
+                    xserv_tensor::DType::BF16,
+                    0,
+                )
+            };
+            let v_view = unsafe {
+                crate::kv_cache::tensor_from_gpu_buffer_pub(
+                    GpuBuffer::borrow_raw(self.buffers.v_reshaped.as_mut_ptr(), v_buf_size),
+                    &shape,
+                    xserv_tensor::DType::BF16,
+                    0,
+                )
+            };
+            cache.append(l, &k_view, &v_view, 1, pos);
+
+            // Ungraphed: get full KV cache and run decode attention
+            let (k_full, v_full) = cache.get_kv_len(l, pos + 1);
+            let kv_len = (pos + 1) as i32;
+            let scale = 1.0 / (hd as f32).sqrt();
+
+            // Attention output written to attn_out (separate from q_final)
+            unsafe {
+                dispatch::decode_attention_bf16(
+                    self.buffers.q_final.as_ptr() as _,
+                    k_full.data_ptr() as _,
+                    v_full.data_ptr() as _,
+                    self.buffers.attn_out.as_mut_ptr() as _,
+                    1, nh as i32, nkv as i32,
+                    kv_len, hd as i32,
+                    scale, s,
+                );
+            }
+
+            // Synchronize before post-attention graph reads attn_out
+            self.stream.synchronize().expect("sync before post-attn");
+
+            // Post-attention graph (merge + O-proj + add_rmsnorm + FFN + residual)
+            self.post_attn_graphs[l].launch(&self.stream).expect("launch post-attn graph");
+        }
+
+        // Final graph (norm + lm_head)
+        self.final_graph.launch(&self.stream).expect("launch final graph");
+
+        // Sync to ensure logits are ready
+        self.stream.synchronize().expect("sync after decode");
+    }
+
+    /// Get the logits buffer (for reading results after execute).
+    pub fn logits_buffer(&self) -> &GpuBuffer {
+        &self.buffers.logits
+    }
+
+    /// Invalidate captured graphs (e.g. when switching sequences).
+    pub fn invalidate(&mut self) {
+        self.captured = false;
+        self.pre_attn_graphs = (0..self.num_layers).map(|_| CudaGraph::new()).collect();
+        self.post_attn_graphs = (0..self.num_layers).map(|_| CudaGraph::new()).collect();
+        self.final_graph = CudaGraph::new();
+    }
+}
+
+unsafe impl Send for DecodeGraphState {}
+
+/// Lightweight struct holding raw pointers to a layer's weight tensors.
+/// Used to avoid passing the full model struct into the graph capture code.
+pub struct LayerWeightPtrs {
+    pub input_norm: *const c_void,
+    pub q_proj_wt: *const c_void,
+    pub k_proj_wt: *const c_void,
+    pub v_proj_wt: *const c_void,
+    pub o_proj_wt: *const c_void,
+    pub q_norm: *const c_void,
+    pub k_norm: *const c_void,
+    pub post_norm: *const c_void,
+    pub gate_proj_wt: *const c_void,
+    pub up_proj_wt: *const c_void,
+    pub down_proj_wt: *const c_void,
+}
+
+unsafe impl Send for LayerWeightPtrs {}
+unsafe impl Sync for LayerWeightPtrs {}

crates/xserv-model/src/gpt2.rs

@@ -116,6 +116,7 @@ fn tensor_from_raw_bytes(bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor
 
 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}"))
         };
@@ -279,12 +280,15 @@ fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
 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 device = qkv.device();
+    let dtype = qkv.dtype();
 
+    match dtype {
+        DType::F32 => {
+            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 {
@@ -295,20 +299,45 @@ fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) ->
                     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)
         }
+        DType::BF16 => {
+            let data = qkv_cpu.as_slice::<half::bf16>();
+            let mut q_data = vec![half::bf16::ZERO; num_heads * seq_len * head_dim];
+            let mut k_data = vec![half::bf16::ZERO; num_heads * seq_len * head_dim];
+            let mut v_data = vec![half::bf16::ZERO; 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 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)
+        }
+        _ => panic!("unsupported dtype {:?} in split_qkv", dtype),
+    }
+}
 
 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 device = x.device();
+    let dtype = x.dtype();
 
+    match dtype {
+        DType::F32 => {
+            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 {
@@ -317,7 +346,22 @@ fn merge_heads(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
                     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())
+            Tensor::from_slice(&out, &[seq_len, hidden]).to_device(device)
+        }
+        DType::BF16 => {
+            let src = x_cpu.as_slice::<half::bf16>();
+            let mut out = vec![half::bf16::ZERO; 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(device)
+        }
+        _ => panic!("unsupported dtype {:?} in merge_heads", dtype),
+    }
 }
 
 /// Greedy sampling: return the argmax token ID from the last position's logits.

crates/xserv-model/src/kv_cache.rs

@@ -9,16 +9,22 @@ pub struct GpuKVCache {
     // 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) -> Self {
+    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();
@@ -27,6 +33,8 @@ impl GpuKVCache {
 
         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");
@@ -34,9 +42,11 @@ impl GpuKVCache {
             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, seq_len: 0, max_seq_len, num_kv_heads, head_dim, elem_size, dtype }
+        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 }
@@ -66,48 +76,63 @@ impl GpuKVCache {
 
     pub fn advance_seq_len(&mut self, new_tokens: usize) {
         self.seq_len += new_tokens;
+        assert!(self.seq_len <= self.max_seq_len, "KV cache seq_len ({}) exceeds max_seq_len ({})", self.seq_len, self.max_seq_len);
     }
 
     /// Get K/V cache tensors for a layer up to `seq_len` tokens: [1, num_kv_heads, seq_len, head_dim]
-    pub fn get_kv(&self, layer: usize) -> (Tensor, Tensor) {
+    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(&self, layer: usize, sl: usize) -> (Tensor, Tensor) {
+    pub fn get_kv_len(&mut self, layer: usize, sl: usize) -> (Tensor, Tensor) {
+        assert!(sl <= self.max_seq_len, "get_kv_len: sl ({sl}) exceeds max_seq_len ({})", self.max_seq_len);
         let hd = self.head_dim;
         let nh = self.num_kv_heads;
         let es = self.elem_size;
         let max_s = self.max_seq_len;
 
-        // Allocate output tensors [1, nh, sl, hd]
+        // 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 mut k_out = GpuBuffer::alloc(out_size).expect("alloc k_out");
+        let k_stg = &mut self.k_staging[layer];
-        let mut v_out = GpuBuffer::alloc(out_size).expect("alloc v_out");
+        let k_buf = &self.k_bufs[layer];
-
+        let v_stg = &mut self.v_staging[layer];
-        // Copy each head's valid portion
+        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_out.copy_from_device_at(&self.k_bufs[layer], src_off, dst_off, count).unwrap();
+            k_stg.copy_from_device_at(k_buf, src_off, dst_off, count).unwrap();
-            v_out.copy_from_device_at(&self.v_bufs[layer], 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(k_out, shape, self.dtype) };
+        let k = unsafe {
-        let v = unsafe { tensor_from_gpu_buffer(v_out, shape, self.dtype) };
+            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) -> Tensor {
+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);
+    let storage = Storage::cuda(buf, device);
     Tensor::from_storage(
         storage,
         SmallVec::from_slice(shape),
@@ -116,3 +141,11 @@ unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType)
         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

@@ -1,10 +1,22 @@
 pub mod config;
+pub mod decode_graph;
 pub mod gpt2;
 pub mod kv_cache;
 pub mod loader;
+pub mod paged_kv_cache;
 pub mod qwen3;
+pub mod sampling;
 
 pub use config::ModelConfig;
+pub use decode_graph::{DecodeGraphState, LayerWeightPtrs};
 pub use gpt2::{GPT2, KVCache};
 pub use kv_cache::GpuKVCache;
+pub use paged_kv_cache::{BlockAllocator, Location, PagedKVCache, BLOCK_SIZE};
 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/paged_kv_cache.rs

@@ -0,0 +1,588 @@
+//! Paged KV cache: vLLM-style block-based KV cache with O(1) allocation
+//! and indirection via per-sequence block tables.
+//!
+//! Physical layout per layer:
+//!   K pool: [total_blocks, num_kv_heads, BLOCK_SIZE, head_dim] BF16
+//!   V pool: same
+//!
+//! Logical view per sequence: a list of physical block ids. Token at logical
+//! position p lives in block_ids[p / BLOCK_SIZE] at slot (p % BLOCK_SIZE).
+
+use crate::config::ModelConfig;
+use xserv_cuda::{GpuBuffer, PinnedBuffer};
+use xserv_tensor::{DType, Tensor};
+
+pub const BLOCK_SIZE: usize = 16;
+
+/// Stack-based block allocator: O(1) alloc/free.
+pub struct BlockAllocator {
+    free_stack: Vec<u32>,
+    total: usize,
+}
+
+impl BlockAllocator {
+    pub fn new(total_blocks: usize) -> Self {
+        // Reserve block 0 as a sentinel "null" block (never allocated).
+        // Free list contains [total-1, total-2, ..., 1] so pop returns 1 first.
+        // total_blocks==0 means "disabled" (e.g. swap off): empty free list.
+        let mut free_stack = Vec::with_capacity(total_blocks.saturating_sub(1));
+        for b in (1..total_blocks).rev() {
+            free_stack.push(b as u32);
+        }
+        Self { free_stack, total: total_blocks }
+    }
+
+    pub fn alloc(&mut self) -> Option<u32> {
+        self.free_stack.pop()
+    }
+
+    pub fn free(&mut self, block: u32) {
+        debug_assert!((block as usize) < self.total && block != 0);
+        self.free_stack.push(block);
+    }
+
+    pub fn free_count(&self) -> usize {
+        self.free_stack.len()
+    }
+
+    pub fn total(&self) -> usize {
+        self.total
+    }
+
+    pub fn can_alloc(&self, n: usize) -> bool {
+        self.free_stack.len() >= n
+    }
+}
+
+/// Where a sequence's KV blocks currently live.
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub enum Location {
+    Gpu,
+    Cpu,
+}
+
+/// Per-sequence state held in the cache.
+#[derive(Clone)]
+pub struct SeqState {
+    /// Block ids into the GPU pool when `location == Gpu`, or into the CPU
+    /// (pinned host) pool when `location == Cpu`.
+    pub block_ids: Vec<u32>,
+    pub seq_len: usize,
+    pub location: Location,
+}
+
+pub struct PagedKVCache {
+    // [layer]: GpuBuffer of size total_blocks * nkv * BLOCK_SIZE * hd * elem_size
+    k_pools: Vec<GpuBuffer>,
+    v_pools: Vec<GpuBuffer>,
+
+    // CPU (pinned host) swap pools, same per-layer layout as the GPU pools but
+    // sized for `cpu_total_blocks`. Empty when swap is disabled.
+    cpu_k_pools: Vec<PinnedBuffer>,
+    cpu_v_pools: Vec<PinnedBuffer>,
+    cpu_allocator: BlockAllocator,
+
+    // Bytes occupied by one block within a single layer pool:
+    //   num_kv_heads * BLOCK_SIZE * head_dim * elem_size.
+    block_bytes: usize,
+
+    allocator: BlockAllocator,
+    seq_states: Vec<Option<SeqState>>,
+
+    // GPU-resident per-sequence metadata. Uploaded each step via sync_to_gpu().
+    // block_table_gpu: i32 [max_seqs, max_blocks_per_seq]
+    // context_lens_gpu: i32 [max_seqs]
+    block_table_gpu: GpuBuffer,
+    context_lens_gpu: GpuBuffer,
+    // Host-side staging mirroring the GPU buffers above.
+    block_table_host: Vec<i32>,
+    context_lens_host: Vec<i32>,
+
+    // Config
+    num_layers: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    elem_size: usize,
+    dtype: DType,
+    device: u32,
+    max_seqs: usize,
+    max_blocks_per_seq: usize,
+}
+
+impl PagedKVCache {
+    /// Bytes occupied by all KV blocks for ONE physical block across the whole
+    /// model (both K and V, all layers). Use this to size pools against VRAM.
+    pub fn bytes_per_block(config: &ModelConfig, dtype: DType) -> usize {
+        2 * config.num_layers()
+            * config.num_kv_heads()
+            * BLOCK_SIZE
+            * config.head_dim()
+            * dtype.size_bytes()
+    }
+
+    /// Create a new paged cache.
+    /// - `total_blocks`: total number of physical GPU blocks across all sequences.
+    /// - `cpu_total_blocks`: physical blocks in the pinned-host swap pool (0 = swap off).
+    /// - `max_seqs`: max number of concurrent sequences (slots), incl. swapped.
+    /// - `max_blocks_per_seq`: capacity of the block table per slot
+    ///   (must be >= ceil(max_seq_len / BLOCK_SIZE)).
+    pub fn new(
+        config: &ModelConfig,
+        total_blocks: usize,
+        cpu_total_blocks: usize,
+        max_seqs: usize,
+        max_blocks_per_seq: usize,
+        dtype: DType,
+        device: u32,
+    ) -> Self {
+        Self::new_tp(
+            config, config.num_kv_heads(), total_blocks, cpu_total_blocks,
+            max_seqs, max_blocks_per_seq, dtype, device,
+        )
+    }
+
+    /// Like `new`, but with an explicit `num_kv_heads` — under tensor parallelism
+    /// each rank only stores its `num_kv_heads / world` heads, so the pool is
+    /// sized for the local head count, not the model's full count.
+    #[allow(clippy::too_many_arguments)]
+    pub fn new_tp(
+        config: &ModelConfig,
+        num_kv_heads: usize,
+        total_blocks: usize,
+        cpu_total_blocks: usize,
+        max_seqs: usize,
+        max_blocks_per_seq: usize,
+        dtype: DType,
+        device: u32,
+    ) -> Self {
+        assert!(total_blocks >= 2, "need at least 2 blocks (one is sentinel)");
+        let num_layers = config.num_layers();
+        let head_dim = config.head_dim();
+        let elem_size = dtype.size_bytes();
+        let block_bytes = num_kv_heads * BLOCK_SIZE * head_dim * elem_size;
+        let pool_bytes = total_blocks * block_bytes;
+
+        let mut k_pools = Vec::with_capacity(num_layers);
+        let mut v_pools = Vec::with_capacity(num_layers);
+        for _ in 0..num_layers {
+            let mut k = GpuBuffer::alloc(pool_bytes).expect("alloc paged K pool");
+            let mut v = GpuBuffer::alloc(pool_bytes).expect("alloc paged V pool");
+            k.zero().unwrap();
+            v.zero().unwrap();
+            k_pools.push(k);
+            v_pools.push(v);
+        }
+
+        // Pinned-host swap pools (one per layer, mirroring the GPU layout).
+        let mut cpu_k_pools = Vec::new();
+        let mut cpu_v_pools = Vec::new();
+        if cpu_total_blocks >= 2 {
+            let cpu_pool_bytes = cpu_total_blocks * block_bytes;
+            for _ in 0..num_layers {
+                cpu_k_pools.push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU K swap pool"));
+                cpu_v_pools.push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU V swap pool"));
+            }
+        }
+        let cpu_allocator = BlockAllocator::new(if cpu_total_blocks >= 2 { cpu_total_blocks } else { 0 });
+
+        let block_table_gpu =
+            GpuBuffer::alloc(max_seqs * max_blocks_per_seq * std::mem::size_of::<i32>())
+                .expect("alloc block table");
+        let context_lens_gpu =
+            GpuBuffer::alloc(max_seqs * std::mem::size_of::<i32>()).expect("alloc context lens");
+
+        let block_table_host = vec![0i32; max_seqs * max_blocks_per_seq];
+        let context_lens_host = vec![0i32; max_seqs];
+
+        let seq_states = (0..max_seqs).map(|_| None).collect();
+
+        Self {
+            k_pools,
+            v_pools,
+            cpu_k_pools,
+            cpu_v_pools,
+            cpu_allocator,
+            block_bytes,
+            allocator: BlockAllocator::new(total_blocks),
+            seq_states,
+            block_table_gpu,
+            context_lens_gpu,
+            block_table_host,
+            context_lens_host,
+            num_layers,
+            num_kv_heads,
+            head_dim,
+            elem_size,
+            dtype,
+            device,
+            max_seqs,
+            max_blocks_per_seq,
+        }
+    }
+
+    pub fn num_layers(&self) -> usize { self.num_layers }
+    pub fn num_kv_heads(&self) -> usize { self.num_kv_heads }
+    pub fn head_dim(&self) -> usize { self.head_dim }
+    pub fn dtype(&self) -> DType { self.dtype }
+    pub fn max_seqs(&self) -> usize { self.max_seqs }
+    pub fn max_blocks_per_seq(&self) -> usize { self.max_blocks_per_seq }
+    pub fn free_blocks(&self) -> usize { self.allocator.free_count() }
+    pub fn total_blocks(&self) -> usize { self.allocator.total() }
+
+    pub fn k_pool(&self, layer: usize) -> &GpuBuffer { &self.k_pools[layer] }
+    pub fn v_pool(&self, layer: usize) -> &GpuBuffer { &self.v_pools[layer] }
+    pub fn block_table_gpu(&self) -> &GpuBuffer { &self.block_table_gpu }
+    pub fn context_lens_gpu(&self) -> &GpuBuffer { &self.context_lens_gpu }
+
+    pub fn seq_len(&self, slot: usize) -> usize {
+        self.seq_states[slot].as_ref().map(|s| s.seq_len).unwrap_or(0)
+    }
+
+    pub fn is_slot_free(&self, slot: usize) -> bool {
+        self.seq_states[slot].is_none()
+    }
+
+    /// Register a new sequence at `slot`. Allocates the first block.
+    /// Returns Err(()) if no slot or no blocks are available.
+    pub fn register_sequence(&mut self, slot: usize) -> Result<(), &'static str> {
+        if slot >= self.max_seqs {
+            return Err("slot out of range");
+        }
+        if self.seq_states[slot].is_some() {
+            return Err("slot already in use");
+        }
+        let block = self.allocator.alloc().ok_or("out of blocks")?;
+        self.seq_states[slot] = Some(SeqState {
+            block_ids: vec![block],
+            seq_len: 0,
+            location: Location::Gpu,
+        });
+        Ok(())
+    }
+
+    /// Free all blocks for `slot` and clear the slot. Frees from whichever pool
+    /// (GPU or CPU) the sequence currently lives in.
+    pub fn free_sequence(&mut self, slot: usize) {
+        if let Some(state) = self.seq_states[slot].take() {
+            let alloc = match state.location {
+                Location::Gpu => &mut self.allocator,
+                Location::Cpu => &mut self.cpu_allocator,
+            };
+            for b in state.block_ids {
+                alloc.free(b);
+            }
+        }
+    }
+
+    /// Number of blocks needed to hold `seq_len + new_tokens` tokens, beyond
+    /// what is currently allocated for `slot`.
+    pub fn additional_blocks_needed(&self, slot: usize, new_tokens: usize) -> usize {
+        let state = self.seq_states[slot].as_ref().expect("unregistered slot");
+        let cur = state.block_ids.len();
+        let needed_total = (state.seq_len + new_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE;
+        if needed_total > cur { needed_total - cur } else { 0 }
+    }
+
+    /// Pre-allocate enough physical blocks in `slot` to cover positions
+    /// `[0, end_pos)`. Call once before the per-layer append loop so that
+    /// every layer's append uses the same block table.
+    pub fn ensure_capacity(&mut self, slot: usize, end_pos: usize) {
+        let state = self.seq_states[slot].as_mut().expect("unregistered slot");
+        let needed_total = (end_pos + BLOCK_SIZE - 1) / BLOCK_SIZE;
+        while state.block_ids.len() < needed_total {
+            let b = self.allocator.alloc().expect("out of blocks (caller must check)");
+            assert!(state.block_ids.len() < self.max_blocks_per_seq, "block table overflow");
+            state.block_ids.push(b);
+        }
+    }
+
+    /// Append `num_tokens` of K/V into the paged pool for `slot` at logical
+    /// position `start_pos`. Caller must have called `ensure_capacity(slot, start_pos + num_tokens)`
+    /// first (or accept that this method may also extend block list).
+    /// Does NOT touch `seq_len`. Call `advance_seq_len(slot, num_tokens)` after
+    /// every layer has been written.
+    ///
+    /// `k_new`, `v_new`: GPU tensors with logical shape
+    ///   [1, num_kv_heads, num_tokens, head_dim]
+    /// stored contiguously (head-major, then tokens, then dim).
+    pub fn append_tokens(
+        &mut self,
+        slot: usize,
+        layer: usize,
+        k_new: &Tensor,
+        v_new: &Tensor,
+        num_tokens: usize,
+        start_pos: usize,
+    ) {
+        if num_tokens == 0 { return; }
+        // Make sure blocks exist for the target range.
+        self.ensure_capacity(slot, start_pos + num_tokens);
+
+        let block_ids = self.seq_states[slot].as_ref().unwrap().block_ids.clone();
+
+        let nkv = self.num_kv_heads;
+        let hd = self.head_dim;
+        let es = self.elem_size;
+        let bs = BLOCK_SIZE;
+
+        let k_src = k_new.storage().gpu_buffer();
+        let v_src = v_new.storage().gpu_buffer();
+
+        let k_pool = &mut self.k_pools[layer];
+        let v_pool = &mut self.v_pools[layer];
+
+        let mut t = 0usize;
+        while t < num_tokens {
+            let p = start_pos + t;
+            let logical_blk = p / bs;
+            let slot_in_blk = p % bs;
+            let chunk = (bs - slot_in_blk).min(num_tokens - t);
+            let phys = block_ids[logical_blk] as usize;
+
+            for h in 0..nkv {
+                let src_off = (h * num_tokens + t) * hd * es;
+                let dst_off = ((phys * nkv + h) * bs + slot_in_blk) * hd * es;
+                let count = chunk * hd * es;
+                k_pool.copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
+                v_pool.copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
+            }
+
+            t += chunk;
+        }
+    }
+
+    /// Advance the logical seq_len after append_tokens for ALL layers has completed.
+    pub fn advance_seq_len(&mut self, slot: usize, num_tokens: usize) {
+        let state = self.seq_states[slot].as_mut().expect("unregistered slot");
+        state.seq_len += num_tokens;
+    }
+
+    /// Refresh the host-side block table + context lens from `seq_states`,
+    /// then upload to GPU. Call once per decode step before the paged kernel.
+    pub fn sync_to_gpu(&mut self) {
+        let stride = self.max_blocks_per_seq;
+        for slot in 0..self.max_seqs {
+            let row = &mut self.block_table_host[slot * stride..(slot + 1) * stride];
+            row.fill(0);
+            let len = match &self.seq_states[slot] {
+                Some(s) => {
+                    for (i, b) in s.block_ids.iter().enumerate() {
+                        row[i] = *b as i32;
+                    }
+                    s.seq_len as i32
+                }
+                None => 0,
+            };
+            self.context_lens_host[slot] = len;
+        }
+
+        self.upload_metadata();
+    }
+
+    /// Pack the given active slots into rows 0..slots.len() of block_table_gpu
+    /// and context_lens_gpu, then upload. Used by paged decode where the kernel
+    /// iterates over `batch` active sequences in order.
+    pub fn sync_active_batch_to_gpu(&mut self, slots: &[usize]) {
+        let lens: Vec<i32> = slots
+            .iter()
+            .map(|&s| self.seq_states[s].as_ref().unwrap().seq_len as i32)
+            .collect();
+        self.sync_active_batch_with_lens(slots, &lens);
+    }
+
+    /// Like sync_active_batch_to_gpu but uses caller-supplied kv_lens (number
+    /// of valid K/V tokens to attend over per active row). Useful when the
+    /// kv_len for the current step differs from the cached seq_len (e.g.
+    /// before advance_seq_len has run).
+    pub fn sync_active_batch_with_lens(&mut self, slots: &[usize], kv_lens: &[i32]) {
+        assert_eq!(slots.len(), kv_lens.len());
+        assert!(slots.len() <= self.max_seqs, "active batch exceeds max_seqs");
+        let stride = self.max_blocks_per_seq;
+        for row in &mut self.block_table_host {
+            *row = 0;
+        }
+        for cl in &mut self.context_lens_host {
+            *cl = 0;
+        }
+        for (i, &slot) in slots.iter().enumerate() {
+            let s = self.seq_states[slot].as_ref().expect("unregistered slot in active batch");
+            let row = &mut self.block_table_host[i * stride..(i + 1) * stride];
+            for (j, b) in s.block_ids.iter().enumerate() {
+                row[j] = *b as i32;
+            }
+            self.context_lens_host[i] = kv_lens[i];
+        }
+        self.upload_metadata();
+    }
+
+    fn upload_metadata(&mut self) {
+        let bt_bytes = unsafe {
+            std::slice::from_raw_parts(
+                self.block_table_host.as_ptr() as *const u8,
+                self.block_table_host.len() * std::mem::size_of::<i32>(),
+            )
+        };
+        self.block_table_gpu.copy_from_host(bt_bytes).unwrap();
+
+        let cl_bytes = unsafe {
+            std::slice::from_raw_parts(
+                self.context_lens_host.as_ptr() as *const u8,
+                self.context_lens_host.len() * std::mem::size_of::<i32>(),
+            )
+        };
+        self.context_lens_gpu.copy_from_host(cl_bytes).unwrap();
+    }
+
+    /// Materialize a contiguous K/V tensor for a sequence at `layer`, shaped
+    /// [1, num_kv_heads, seq_len, head_dim]. Used for prefill, where Flash
+    /// Attention 2 expects contiguous K/V.
+    ///
+    /// Allocates from the cached allocator; the returned Tensors own their storage.
+    pub fn gather_kv_contiguous(&self, slot: usize, layer: usize) -> (Tensor, Tensor) {
+        let state = self.seq_states[slot].as_ref().expect("unregistered slot");
+        let sl = state.seq_len;
+        let nkv = self.num_kv_heads;
+        let hd = self.head_dim;
+        let es = self.elem_size;
+        let bs = BLOCK_SIZE;
+
+        let out_bytes = nkv * sl * hd * es;
+        let mut k_dst = xserv_cuda::allocator::cached_alloc(out_bytes).expect("alloc gather K");
+        let mut v_dst = xserv_cuda::allocator::cached_alloc(out_bytes).expect("alloc gather V");
+
+        let k_pool = &self.k_pools[layer];
+        let v_pool = &self.v_pools[layer];
+
+        let mut p = 0usize;
+        while p < sl {
+            let logical_blk = p / bs;
+            let slot_in_blk = p % bs;
+            let chunk = (bs - slot_in_blk).min(sl - p);
+            let phys = state.block_ids[logical_blk] as usize;
+
+            for h in 0..nkv {
+                let src_off = ((phys * nkv + h) * bs + slot_in_blk) * hd * es;
+                let dst_off = (h * sl + p) * hd * es;
+                let count = chunk * hd * es;
+                k_dst.copy_from_device_at(k_pool, src_off, dst_off, count).unwrap();
+                v_dst.copy_from_device_at(v_pool, src_off, dst_off, count).unwrap();
+            }
+            p += chunk;
+        }
+
+        let shape = &[1usize, nkv, sl, hd];
+        let k = unsafe { tensor_from_owned_buf(k_dst, shape, self.dtype, self.device) };
+        let v = unsafe { tensor_from_owned_buf(v_dst, shape, self.dtype, self.device) };
+        (k, v)
+    }
+
+    // ----- Swapping (vLLM-style preemption to pinned host memory) -----
+
+    pub fn free_cpu_blocks(&self) -> usize { self.cpu_allocator.free_count() }
+    pub fn swap_enabled(&self) -> bool { !self.cpu_k_pools.is_empty() }
+
+    pub fn is_swapped(&self, slot: usize) -> bool {
+        matches!(self.seq_states[slot].as_ref().map(|s| s.location), Some(Location::Cpu))
+    }
+
+    /// Number of physical blocks currently held by `slot` (in either pool).
+    pub fn block_count(&self, slot: usize) -> usize {
+        self.seq_states[slot].as_ref().map(|s| s.block_ids.len()).unwrap_or(0)
+    }
+
+    /// Whether a swapped sequence at `slot` can be brought back (enough free GPU blocks).
+    pub fn can_swap_in(&self, slot: usize) -> bool {
+        self.allocator.can_alloc(self.block_count(slot))
+    }
+
+    /// Whether the GPU sequence at `slot` can be evicted (enough free CPU blocks).
+    pub fn can_swap_out(&self, slot: usize) -> bool {
+        self.cpu_allocator.can_alloc(self.block_count(slot))
+    }
+
+    /// Evict `slot`'s KV from GPU to pinned host memory and free its GPU blocks.
+    /// The slot stays registered (location = Cpu); the sequence is paused.
+    pub fn swap_out(&mut self, slot: usize) -> Result<(), &'static str> {
+        let state = self.seq_states[slot].as_ref().ok_or("swap_out: empty slot")?;
+        if state.location == Location::Cpu { return Ok(()); }
+        let gpu_ids = state.block_ids.clone();
+        let n = gpu_ids.len();
+        if !self.cpu_allocator.can_alloc(n) { return Err("swap_out: CPU pool full"); }
+
+        let cpu_ids: Vec<u32> = (0..n)
+            .map(|_| self.cpu_allocator.alloc().expect("checked can_alloc"))
+            .collect();
+
+        let bb = self.block_bytes;
+        for layer in 0..self.num_layers {
+            for i in 0..n {
+                let g_off = gpu_ids[i] as usize * bb;
+                let c_off = cpu_ids[i] as usize * bb;
+                self.k_pools[layer]
+                    .copy_to_host_at(&mut self.cpu_k_pools[layer].as_mut_slice()[c_off..c_off + bb], g_off, bb)
+                    .unwrap();
+                self.v_pools[layer]
+                    .copy_to_host_at(&mut self.cpu_v_pools[layer].as_mut_slice()[c_off..c_off + bb], g_off, bb)
+                    .unwrap();
+            }
+        }
+
+        for b in gpu_ids {
+            self.allocator.free(b);
+        }
+        let state = self.seq_states[slot].as_mut().unwrap();
+        state.block_ids = cpu_ids;
+        state.location = Location::Cpu;
+        Ok(())
+    }
+
+    /// Bring `slot`'s KV back from host to GPU and free its CPU blocks.
+    pub fn swap_in(&mut self, slot: usize) -> Result<(), &'static str> {
+        let state = self.seq_states[slot].as_ref().ok_or("swap_in: empty slot")?;
+        if state.location == Location::Gpu { return Ok(()); }
+        let cpu_ids = state.block_ids.clone();
+        let n = cpu_ids.len();
+        if !self.allocator.can_alloc(n) { return Err("swap_in: GPU pool full"); }
+
+        let gpu_ids: Vec<u32> = (0..n)
+            .map(|_| self.allocator.alloc().expect("checked can_alloc"))
+            .collect();
+
+        let bb = self.block_bytes;
+        for layer in 0..self.num_layers {
+            for i in 0..n {
+                let g_off = gpu_ids[i] as usize * bb;
+                let c_off = cpu_ids[i] as usize * bb;
+                self.k_pools[layer]
+                    .copy_from_host_at(&self.cpu_k_pools[layer].as_slice()[c_off..c_off + bb], g_off, bb)
+                    .unwrap();
+                self.v_pools[layer]
+                    .copy_from_host_at(&self.cpu_v_pools[layer].as_slice()[c_off..c_off + bb], g_off, bb)
+                    .unwrap();
+            }
+        }
+
+        for b in cpu_ids {
+            self.cpu_allocator.free(b);
+        }
+        let state = self.seq_states[slot].as_mut().unwrap();
+        state.block_ids = gpu_ids;
+        state.location = Location::Gpu;
+        Ok(())
+    }
+}
+
+unsafe fn tensor_from_owned_buf(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
+    use smallvec::SmallVec;
+    use xserv_tensor::shape::contiguous_strides;
+    use xserv_tensor::storage::Storage;
+
+    let storage = Storage::cuda(buf, device);
+    Tensor::from_storage(
+        storage,
+        SmallVec::from_slice(shape),
+        contiguous_strides(shape),
+        0,
+        dtype,
+    )
+}

crates/xserv-model/src/qwen3.rs

@@ -6,6 +6,7 @@ use xserv_tensor::{DType, Device, Tensor};
 use crate::config::ModelConfig;
 use crate::gpt2::KVCache;
 use crate::kv_cache::GpuKVCache;
+use crate::paged_kv_cache::PagedKVCache;
 
 pub struct Qwen3 {
     pub config: ModelConfig,
@@ -14,6 +15,11 @@ pub struct Qwen3 {
     norm: Tensor,
     lm_head_t: Tensor,  // precomputed transpose
     rope_cache: RopeCache,
+    // Tensor parallelism. `tp` is None (or world==1) for single-GPU; otherwise
+    // this rank holds 1/world of the heads and AllReduces after o_proj/down_proj.
+    tp: Option<std::sync::Arc<xserv_distributed::TpContext>>,
+    local_num_heads: usize,    // = num_heads / world
+    local_num_kv_heads: usize, // = num_kv_heads / world
 }
 
 struct Qwen3Block {
@@ -31,48 +37,95 @@ struct Qwen3Block {
 }
 
 impl Qwen3 {
-    pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
+    /// Single-GPU load (weights already on the target GPU). Equivalent to
+    /// `from_weights_tp(.., rank=0, world=1, device=0, tp=None)`.
+    pub fn from_weights(config: ModelConfig, w: HashMap<String, Tensor>) -> Self {
+        Self::from_weights_tp(config, w, 0, 1, 0, None)
+    }
+
+    /// Tensor-parallel load. `w` may live on CPU or any device; each weight is
+    /// sharded for `rank`/`world`, uploaded to `device`, and transposed.
+    /// `world==1` shards are identity, so this is also the single-GPU path.
+    ///
+    /// Split scheme (Megatron-style):
+    /// - column-parallel (split output): q/k/v/gate/up  → shard rows of `[out,in]`
+    /// - row-parallel    (split input):  o/down         → shard cols of `[out,in]`
+    /// - replicated: norms, embed_tokens, lm_head
+    pub fn from_weights_tp(
+        config: ModelConfig,
+        mut w: HashMap<String, Tensor>,
+        rank: usize,
+        world: usize,
+        device: u32,
+        tp: Option<std::sync::Arc<xserv_distributed::TpContext>>,
+    ) -> Self {
+        crate::init_kernels();
+        let dev = Device::Cuda(device);
         let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
             w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
         };
+        // Replicated weight: upload whole to this rank's device.
+        let repl = |t: Tensor| -> Tensor { t.to_device(dev) };
+        // column-parallel: keep this rank's rows of [out, in], upload, transpose → [in, out/world].
+        let col = |t: Tensor| -> Tensor { shard_rows(&t, rank, world).to_device(dev).transpose(0, 1).contiguous() };
+        // row-parallel: keep this rank's cols of [out, in], upload, transpose → [in/world, out].
+        let row = |t: Tensor| -> Tensor { shard_cols(&t, rank, world).to_device(dev).transpose(0, 1).contiguous() };
 
-        let embed_tokens = take(&mut w, "model.embed_tokens.weight");
+        let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight"));
-        let norm = take(&mut w, "model.norm.weight");
+        let norm = repl(take(&mut w, "model.norm.weight"));
-        let lm_head_raw = take(&mut w, "lm_head.weight");
+        let lm_head_t = repl(take(&mut w, "lm_head.weight")).transpose(0, 1).contiguous();
 
         let rope_cache = RopeCache::new(
-            config.max_seq_len().min(8192), // limit for memory
+            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);
+        if rank == 0 {
+            eprintln!("Loading+sharding weights for {} layers (world={world})...", 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")),
+                input_norm: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))),
-                q_proj_wt: transpose_w(take(&mut w, &format!("{p}.self_attn.q_proj.weight"))),
+                q_proj_wt: col(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"))),
+                k_proj_wt: col(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"))),
+                v_proj_wt: col(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"))),
+                o_proj_wt: row(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
-                q_norm: take(&mut w, &format!("{p}.self_attn.q_norm.weight")),
+                q_norm: repl(take(&mut w, &format!("{p}.self_attn.q_norm.weight"))),
-                k_norm: take(&mut w, &format!("{p}.self_attn.k_norm.weight")),
+                k_norm: repl(take(&mut w, &format!("{p}.self_attn.k_norm.weight"))),
-                post_norm: take(&mut w, &format!("{p}.post_attention_layernorm.weight")),
+                post_norm: repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight"))),
-                gate_proj_wt: transpose_w(take(&mut w, &format!("{p}.mlp.gate_proj.weight"))),
+                gate_proj_wt: col(take(&mut w, &format!("{p}.mlp.gate_proj.weight"))),
-                up_proj_wt: transpose_w(take(&mut w, &format!("{p}.mlp.up_proj.weight"))),
+                up_proj_wt: col(take(&mut w, &format!("{p}.mlp.up_proj.weight"))),
-                down_proj_wt: transpose_w(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
+                down_proj_wt: row(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
             });
         }
 
-        let lm_head_t = transpose_w(lm_head_raw);
+        Self {
-        Self { config, embed_tokens, layers, norm, lm_head_t, rope_cache }
+            local_num_heads: config.num_heads() / world,
+            local_num_kv_heads: config.num_kv_heads() / world,
+            config,
+            embed_tokens,
+            layers,
+            norm,
+            lm_head_t,
+            rope_cache,
+            tp,
+        }
+    }
+
+    /// In-place AllReduce(sum) of a partial `[*, hidden]` BF16 activation across
+    /// TP ranks (no-op when not tensor-parallel). Used after o_proj and down_proj.
+    #[inline]
+    fn all_reduce(&self, t: &Tensor) {
+        if let Some(tp) = &self.tp {
+            if tp.world > 1 {
+                let ptr = t.storage().gpu_buffer().as_ptr() as *mut std::ffi::c_void;
+                tp.all_reduce_sum_bf16_ptr(ptr, t.numel());
+            }
+        }
     }
 
     pub fn forward_with_cache(&self, token_ids: &[u32], cache: &mut KVCache) -> Tensor {
@@ -147,6 +200,309 @@ impl Qwen3 {
         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]
+    }
+
+    /// Paged decode: process one token per sequence using a shared paged KV cache.
+    ///
+    /// tokens:    [B] one token per sequence
+    /// positions: [B] current logical position (BEFORE this step) per sequence
+    /// seq_slots: [B] slot ids in `paged_cache`
+    pub fn forward_decode_paged(
+        &self,
+        tokens: &[u32],
+        positions: &[usize],
+        seq_slots: &[usize],
+        paged_cache: &mut PagedKVCache,
+    ) -> Tensor {
+        let batch = tokens.len();
+        assert_eq!(positions.len(), batch);
+        assert_eq!(seq_slots.len(), batch);
+        assert!(batch > 0);
+
+        // TP: this rank owns a slice of the heads (local_* == full when world==1).
+        let num_heads = self.local_num_heads;
+        let num_kv_heads = self.local_num_kv_heads;
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        // Ensure all slots have enough physical blocks for this token, then
+        // upload block tables + context_lens once for the whole forward (the
+        // tables are identical across layers; only the layer's K/V pool changes).
+        let kv_lens: Vec<i32> = positions.iter().map(|&p| (p + 1) as i32).collect();
+        for (b, &slot) in seq_slots.iter().enumerate() {
+            paged_cache.ensure_capacity(slot, positions[b] + 1);
+        }
+        paged_cache.sync_active_batch_with_lens(seq_slots, &kv_lens);
+
+        let bt_ptr = paged_cache.block_table_gpu().as_ptr() as *const i32;
+        let cl_ptr = paged_cache.context_lens_gpu().as_ptr() as *const i32;
+        let max_blocks = paged_cache.max_blocks_per_seq();
+
+        // 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);
+
+            let q_all = matmul_2d(&normed, &layer.q_proj_wt);
+            let k_all = matmul_2d(&normed, &layer.k_proj_wt);
+            let v_all = matmul_2d(&normed, &layer.v_proj_wt);
+
+            let mut q_rows: Vec<Tensor> = Vec::with_capacity(batch);
+            for b in 0..batch {
+                let q_row = row_view(&q_all, b);
+                let k_row = row_view(&k_all, b);
+                let v_row = row_view(&v_all, b);
+
+                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);
+
+                let q = head_rmsnorm(&q, &layer.q_norm, eps);
+                let k = head_rmsnorm(&k, &layer.k_norm, eps);
+
+                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);
+
+                let pos = [positions[b] as u32];
+                rope_inplace(&q, &self.rope_cache, &pos);
+                rope_inplace(&k, &self.rope_cache, &pos);
+
+                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);
+
+                paged_cache.append_tokens(seq_slots[b], layer_idx, &k, &v, 1, positions[b]);
+
+                let q_flat = xserv_kernels::merge_heads_gpu(&q, 1, num_heads, head_dim);
+                q_rows.push(q_flat);
+            }
+
+            let q_batched_2d = concat_rows(&q_rows);
+            // q_batched_2d: [B, num_heads * head_dim]. Memory is [B, H, D] —
+            // a plain reshape view to [B, H, 1, D] is what the paged kernel expects.
+            let q_4d = q_batched_2d.reshape(&[batch, num_heads, 1, head_dim]);
+
+            let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
+            let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
+
+            let attn_out = xserv_kernels::paged_decode_attention(
+                &q_4d,
+                k_pool_ptr,
+                v_pool_ptr,
+                bt_ptr,
+                cl_ptr,
+                batch,
+                num_heads,
+                num_kv_heads,
+                head_dim,
+                max_blocks,
+            );
+
+            // attn_out shape [B, H, 1, D] is contiguous-equivalent to [B, H*D].
+            // Plain reshape is a view; merge_heads_gpu would incorrectly swap B<->H.
+            let attn_merged = attn_out.reshape(&[batch, num_heads * head_dim]);
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+            self.all_reduce(&attn_proj); // TP: sum partial attention outputs
+
+            let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            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);
+            self.all_reduce(&down); // TP: sum partial MLP outputs
+            x = add_any(&residual, &down);
+        }
+
+        // Advance logical seq_len now that all layers have been written.
+        for &slot in seq_slots {
+            paged_cache.advance_seq_len(slot, 1);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_2d(&x, &self.lm_head_t)
+    }
+
+    /// Paged prefill: write a sequence of `new_tokens` K/V into the paged
+    /// cache for `slot`, run flash attention via gathered contiguous K/V.
+    /// Returns logits [new_tokens, vocab_size].
+    pub fn forward_prefill_paged(
+        &self,
+        token_ids: &[u32],
+        slot: usize,
+        paged_cache: &mut PagedKVCache,
+    ) -> Tensor {
+        let new_tokens = token_ids.len();
+        let pos_offset = paged_cache.seq_len(slot);
+        // TP: this rank owns a slice of the heads (local_* == full when world==1).
+        let num_heads = self.local_num_heads;
+        let num_kv_heads = self.local_num_kv_heads;
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        // Pre-allocate enough blocks and bump seq_len up-front so per-layer
+        // gather_kv_contiguous returns the freshly written K/V range.
+        paged_cache.ensure_capacity(slot, pos_offset + new_tokens);
+        paged_cache.advance_seq_len(slot, new_tokens);
+
+        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);
+
+            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);
+
+            let q = head_rmsnorm(&q, &layer.q_norm, eps);
+            let k = head_rmsnorm(&k, &layer.k_norm, eps);
+
+            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);
+            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);
+
+            // Write into paged pool at the original (pre-advance) position.
+            paged_cache.append_tokens(slot, layer_idx, &k, &v, new_tokens, pos_offset);
+
+            // Gather contiguous K/V for the full sequence (seq_len already includes new_tokens).
+            let (k_full, v_full) = paged_cache.gather_kv_contiguous(slot, layer_idx);
+            let attn_out = flash_attention(&q, &k_full, &v_full, true);
+
+            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);
+            self.all_reduce(&attn_proj); // TP: sum partial attention outputs
+
+            let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            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);
+            self.all_reduce(&down); // TP: sum partial MLP outputs
+            x = add_any(&residual, &down);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_2d(&x, &self.lm_head_t)
+    }
+
     /// Forward with GPU-resident KV cache and GPU transpose/reshape kernels.
     pub fn forward_gpu_cache(&self, token_ids: &[u32], cache: &mut GpuKVCache) -> Tensor {
         let new_tokens = token_ids.len();
@@ -190,23 +546,20 @@ impl Qwen3 {
             cache.append(layer_idx, &k, &v, new_tokens, pos_offset);
             let (k_full, v_full) = cache.get_kv_len(layer_idx, pos_offset + new_tokens);
 
-            // GPU repeat KV for GQA
+            // Flash Attention with native GQA (no repeat_kv needed)
-            let n_rep = num_heads / num_kv_heads;
+            let attn_out = flash_attention(&q, &k_full, &v_full, true);
-            let k_full = xserv_kernels::repeat_kv_gpu(&k_full, n_rep);
-            let v_full = xserv_kernels::repeat_kv_gpu(&v_full, n_rep);
-
-            let attn_out = attention(&q, &k_full, &v_full, true);
             // GPU merge_heads: [1, H, S, D] → [S, H*D]
             let attn_merged = xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
             let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
-            x = add_any(&residual, &attn_proj);
 
-            let residual = x.clone();
+            // Fused add + rmsnorm: (normed, x) where x = residual + attn_proj
-            let normed = rmsnorm(&x, &layer.post_norm, eps);
+            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 gate_activated = silu(&gate);
+            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
-            let hidden_states = mul_any(&gate_activated, &up);
             let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
             x = add_any(&residual, &down);
         }
@@ -215,10 +568,81 @@ impl Qwen3 {
         let x = rmsnorm(&x, &self.norm, eps);
         matmul_2d(&x, &self.lm_head_t)
     }
+
+    /// Extract weight pointers for CUDA Graph capture.
+    pub fn layer_weight_ptrs(&self) -> Vec<crate::decode_graph::LayerWeightPtrs> {
+        self.layers.iter().map(|l| crate::decode_graph::LayerWeightPtrs {
+            input_norm: l.input_norm.data_ptr() as *const std::ffi::c_void,
+            q_proj_wt: l.q_proj_wt.data_ptr() as *const std::ffi::c_void,
+            k_proj_wt: l.k_proj_wt.data_ptr() as *const std::ffi::c_void,
+            v_proj_wt: l.v_proj_wt.data_ptr() as *const std::ffi::c_void,
+            o_proj_wt: l.o_proj_wt.data_ptr() as *const std::ffi::c_void,
+            q_norm: l.q_norm.data_ptr() as *const std::ffi::c_void,
+            k_norm: l.k_norm.data_ptr() as *const std::ffi::c_void,
+            post_norm: l.post_norm.data_ptr() as *const std::ffi::c_void,
+            gate_proj_wt: l.gate_proj_wt.data_ptr() as *const std::ffi::c_void,
+            up_proj_wt: l.up_proj_wt.data_ptr() as *const std::ffi::c_void,
+            down_proj_wt: l.down_proj_wt.data_ptr() as *const std::ffi::c_void,
+        }).collect()
+    }
+
+    /// Get pointers needed for CUDA Graph capture.
+    pub fn graph_capture_ptrs(&self) -> (
+        *const std::ffi::c_void,  // norm weight
+        *const std::ffi::c_void,  // lm_head_t
+        *const std::ffi::c_void,  // embed_tokens
+        *const std::ffi::c_void,  // rope cos
+        *const std::ffi::c_void,  // rope sin
+    ) {
+        (
+            self.norm.data_ptr() as *const std::ffi::c_void,
+            self.lm_head_t.data_ptr() as *const std::ffi::c_void,
+            self.embed_tokens.data_ptr() as *const std::ffi::c_void,
+            self.rope_cache.cos.as_ptr() as *const std::ffi::c_void,
+            self.rope_cache.sin.as_ptr() as *const std::ffi::c_void,
+        )
+    }
 }
 
 // --- Helpers ---
 
+/// Keep this rank's contiguous row-block of a 2D `[rows, cols]` BF16 tensor
+/// (column-parallel split: split the OUTPUT dim). `world==1` returns the whole.
+/// Input must be a contiguous CPU (or device) BF16 tensor.
+fn shard_rows(t: &Tensor, rank: usize, world: usize) -> Tensor {
+    if world == 1 { return t.clone(); }
+    let shape = t.shape();
+    assert_eq!(shape.len(), 2, "shard_rows expects 2D weight");
+    let (rows, cols) = (shape[0], shape[1]);
+    assert!(rows % world == 0, "rows {rows} not divisible by world {world}");
+    let local = rows / world;
+    let host = t.to_device(Device::Cpu);
+    let data = host.as_slice::<bf16>();
+    let start = rank * local * cols;
+    let shard = data[start..start + local * cols].to_vec();
+    Tensor::from_slice(&shard, &[local, cols])
+}
+
+/// Keep this rank's column-block of a 2D `[rows, cols]` BF16 tensor (row-parallel
+/// split: split the INPUT dim). Strided copy. `world==1` returns the whole.
+fn shard_cols(t: &Tensor, rank: usize, world: usize) -> Tensor {
+    if world == 1 { return t.clone(); }
+    let shape = t.shape();
+    assert_eq!(shape.len(), 2, "shard_cols expects 2D weight");
+    let (rows, cols) = (shape[0], shape[1]);
+    assert!(cols % world == 0, "cols {cols} not divisible by world {world}");
+    let local = cols / world;
+    let c0 = rank * local;
+    let host = t.to_device(Device::Cpu);
+    let data = host.as_slice::<bf16>();
+    let mut shard = Vec::with_capacity(rows * local);
+    for r in 0..rows {
+        let base = r * cols + c0;
+        shard.extend_from_slice(&data[base..base + local]);
+    }
+    Tensor::from_slice(&shard, &[rows, local])
+}
+
 fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
     assert_eq!(a.ndim(), 2);
     assert_eq!(b.ndim(), 2);
@@ -319,6 +743,53 @@ fn repeat_kv(x: &Tensor, n_rep: usize) -> Tensor {
     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)
 }

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

@@ -13,9 +13,11 @@ xserv-tensor = { path = "../xserv-tensor" }
 xserv-kernels = { path = "../xserv-kernels" }
 xserv-model = { path = "../xserv-model" }
 xserv-tokenizer = { path = "../xserv-tokenizer" }
+xserv-distributed = { path = "../xserv-distributed" }
 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

@@ -1,11 +1,18 @@
 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::engine::{GenerateEvent, GenerateRequest};
 use crate::AppState;
+use crate::engine::{GenerateEvent, GenerateRequest};
+use xserv_model::SamplingParams;
 
 #[derive(Deserialize)]
 pub struct ChatRequest {
@@ -14,6 +21,14 @@ pub struct ChatRequest {
     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)]
@@ -22,7 +37,9 @@ pub struct Message {
     pub content: String,
 }
 
-fn default_max_tokens() -> usize { 256 }
+fn default_max_tokens() -> usize {
+    256
+}
 
 #[derive(Serialize)]
 pub struct ModelsResponse {
@@ -37,7 +54,9 @@ pub struct ModelInfo {
     owned_by: &'static str,
 }
 
-pub async fn health() -> &'static str { "ok" }
+pub async fn health() -> &'static str {
+    "ok"
+}
 
 pub async fn list_models(Extension(state): Extension<Arc<AppState>>) -> Json<ModelsResponse> {
     Json(ModelsResponse {
@@ -53,34 +72,60 @@ pub async fn list_models(Extension(state): Extension<Arc<AppState>>) -> Json<Mod
 pub async fn chat_completions(
     Extension(state): Extension<Arc<AppState>>,
     Json(req): Json<ChatRequest>,
-) -> Json<serde_json::Value> {
+) -> 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 = std::time::SystemTime::now()
+    let created = unix_timestamp();
-        .duration_since(std::time::UNIX_EPOCH)
+
-        .unwrap()
+    if let Some(response) = validate_request(&req, &model_name) {
-        .as_secs();
+        return response;
+    }
 
-    // Prepare prompt tokens (MutexGuard scoped)
     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 bad_request(format!(
+            "prompt is {} tokens, exceeds max_seq_len {}",
+            prompt_token_count, max_seq_len
+        ));
+    }
+    let max_tokens = req.max_tokens.min(max_seq_len - prompt_token_count);
 
-    // Create channel and submit request (MutexGuard scoped)
     let (tx, mut rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
     let gen_req = GenerateRequest {
         prompt_tokens,
-        max_tokens: req.max_tokens,
+        max_tokens,
+        sampling: sampling_params(&req),
         sender: tx,
     };
-    state.engine_sender.lock().unwrap().send(gen_req).expect("engine channel closed");
+    if let Err(resp) = submit_to_engine(&state, gen_req) {
+        return resp;
+    }
 
-    // Now await — no MutexGuards held here
     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, .. } => content.push_str(&text),
+            GenerateEvent::Token { text, .. } => {
-            GenerateEvent::Done { finish_reason: fr } => { finish_reason = fr; break; }
+                completion_token_count += 1;
+                content.push_str(&text);
+            }
+            GenerateEvent::Done { finish_reason: fr } => {
+                finish_reason = fr;
+                break;
+            }
         }
     }
 
@@ -95,21 +140,206 @@ pub async fn chat_completions(
             "finish_reason": finish_reason,
         }],
         "usage": {
-            "prompt_tokens": 0,
+            "prompt_tokens": prompt_token_count,
-            "completion_tokens": 0,
+            "completion_tokens": completion_token_count,
-            "total_tokens": 0
+            "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();
+
+    if let Some(response) = validate_request(&req, &model_name) {
+        return response;
+    }
+
+    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 bad_request(format!(
+            "prompt is {} tokens, exceeds max_seq_len {}",
+            prompt_tokens.len(), max_seq_len
+        ));
+    }
+    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,
+    };
+    if let Err(resp) = submit_to_engine(&state, gen_req) {
+        return resp;
+    }
+
+    // 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 validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
+    if let Some(model) = &req.model {
+        if model != model_name {
+            return Some(bad_request(format!(
+                "model '{model}' is not loaded; available model is '{model_name}'"
+            )));
+        }
+    }
+
+    if req.max_tokens == 0 {
+        return Some(bad_request("max_tokens must be greater than 0"));
+    }
+
+    None
+}
+
+/// Hand a request to the engine thread. Poison-tolerant (recovers the lock if a
+/// prior handler panicked) and returns a clean 503 instead of panicking when the
+/// engine thread is gone, so one dead engine doesn't cascade into every request.
+fn submit_to_engine(state: &AppState, req: GenerateRequest) -> Result<(), Response> {
+    let sender = state.engine_sender.lock().unwrap_or_else(|e| e.into_inner());
+    sender.send(req).map_err(|_| service_unavailable("inference engine is not available"))
+}
+
+fn service_unavailable(message: impl Into<String>) -> Response {
+    (
+        StatusCode::SERVICE_UNAVAILABLE,
+        Json(serde_json::json!({
+            "error": { "message": message.into(), "type": "server_error" }
+        })),
+    )
+        .into_response()
+}
+
+fn bad_request(message: impl Into<String>) -> Response {
+    (
+        StatusCode::BAD_REQUEST,
+        Json(serde_json::json!({
+            "error": {
+                "message": message.into(),
+                "type": "invalid_request_error"
+            }
+        })),
+    )
+        .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" => { prompt.push_str(&msg.content); prompt.push('\n'); }
+            "system" | "user" | "assistant" => {
-            "user" | "assistant" => { prompt.push_str(&msg.content); }
+                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.push_str("<think>\n\n</think>\n\n");
     prompt
 }

crates/xserv-server/src/engine.rs

@@ -1,9 +1,10 @@
 use std::collections::VecDeque;
 use std::path::Path;
 use std::sync::mpsc;
-use xserv_model::{GpuKVCache, ModelConfig, Qwen3};
+use std::sync::Once;
+use std::time::Instant;
+use xserv_model::{ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample, BLOCK_SIZE};
 use xserv_model::loader;
-use xserv_model::qwen3::sample_greedy;
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
@@ -13,11 +14,13 @@ pub struct Engine {
     tokenizer: Tokenizer,
     max_batch_size: usize,
     max_seq_len: usize,
+    paged_cache: PagedKVCache,
 }
 
 pub struct GenerateRequest {
     pub prompt_tokens: Vec<u32>,
     pub max_tokens: usize,
+    pub sampling: SamplingParams,
     pub sender: tokio::sync::mpsc::Sender<GenerateEvent>,
 }
 
@@ -31,13 +34,26 @@ struct Sequence {
     prompt_tokens: Vec<u32>,
     generated_tokens: Vec<u32>,
     max_tokens: usize,
-    kv_cache: GpuKVCache,
+    sampling: SamplingParams,
+    seq_slot: Option<usize>,
     sender: tokio::sync::mpsc::Sender<GenerateEvent>,
     prefilled: bool,
+    eos_token_id: Option<u32>,
+    decode_buffer: Vec<u8>,
+    created_at: Instant,
 }
 
 impl Engine {
-    pub fn load(model_dir: &Path, max_batch_size: usize) -> Self {
+    pub fn load(model_dir: &Path, max_batch_size: usize, max_seq_len: usize) -> Self {
+        Self::load_with_swap(model_dir, max_batch_size, max_seq_len, 8)
+    }
+
+    pub fn load_with_swap(
+        model_dir: &Path,
+        max_batch_size: usize,
+        max_seq_len: usize,
+        swap_space_gb: usize,
+    ) -> Self {
         xserv_cuda::device::set_device(0).unwrap();
         let config = ModelConfig::from_file(&model_dir.join("config.json"));
         eprintln!("[engine] Loading weights...");
@@ -45,65 +61,226 @@ impl Engine {
         eprintln!("[engine] Loaded {} tensors", weights.len());
         let model = Qwen3::from_weights(config.clone(), weights);
         let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
-        let max_seq_len = 256;
+
-        eprintln!("[engine] Ready (max_batch_size={max_batch_size}, max_seq_len={max_seq_len})");
+        // Tier-1 sizing: size the GPU block pool to *available VRAM* after the
-        Self { model, config, tokenizer, max_batch_size, max_seq_len }
+        // weights are resident, not to worst-case max_batch * max_ctx. This is
+        // what makes paged attention elastic — sequences share the pool on
+        // demand, and overflow is swapped to host (Tier-2) rather than reserved.
+        let bytes_per_block = PagedKVCache::bytes_per_block(&config, DType::BF16);
+        let info = xserv_cuda::device::device_info(0).expect("device info");
+        // Reserve headroom for activations, cuBLAS workspace and the [B, vocab]
+        // logits buffer; the transpose peak during load is already behind us.
+        const ACTIVATION_RESERVE: usize = 3 * 1024 * 1024 * 1024; // 3 GiB
+        let util_num = 90; // use 90% of remaining free memory for KV
+        let usable = info.free_memory.saturating_sub(ACTIVATION_RESERVE);
+        let mut total_blocks = (usable * util_num / 100) / bytes_per_block;
+        // Cap at a sane upper bound and ensure a floor.
+        total_blocks = total_blocks.max(256);
+        // Test hook: force a small GPU pool to exercise the swap path. Must stay
+        // >= max_blocks_per_seq so a single max-length sequence still fits.
+        if let Ok(v) = std::env::var("XSERV_MAX_KV_BLOCKS") {
+            if let Ok(n) = v.parse::<usize>() {
+                total_blocks = total_blocks.min(n);
+                eprintln!("[engine] XSERV_MAX_KV_BLOCKS override: gpu_blocks={total_blocks}");
+            }
+        }
+
+        let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
+        // Slots must cover running + swapped sequences, so be generous (cheap:
+        // each slot is just a block-table row of i32s).
+        let max_seqs_slots = (max_batch_size * 8).max(32);
+        // CPU swap pool: swap_space_gb of pinned host memory.
+        let cpu_total_blocks = (swap_space_gb * 1024 * 1024 * 1024) / bytes_per_block;
+
+        let paged_cache = PagedKVCache::new(
+            &config,
+            total_blocks,
+            cpu_total_blocks,
+            max_seqs_slots,
+            max_blocks_per_seq,
+            DType::BF16,
+            0,
+        );
+
+        eprintln!(
+            "[engine] Ready (max_batch={max_batch_size}, max_seq_len={max_seq_len}, \
+             gpu_blocks={total_blocks} ({:.1} GiB), swap_blocks={cpu_total_blocks} ({swap_space_gb} GiB), \
+             free_vram={:.1} GiB)",
+            (total_blocks * bytes_per_block) as f64 / 1e9,
+            info.free_memory as f64 / 1e9,
+        );
+        Self { model, config, tokenizer, max_batch_size, max_seq_len, paged_cache }
     }
 
     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>) {
+    ///
+    /// Sequences move between three sets:
+    ///   waiting  — admitted to the queue, no GPU slot yet
+    ///   running  — KV resident on GPU, actively prefilling/decoding
+    ///   swapped  — KV evicted to pinned host memory (preempted), paused
+    /// When running sequences grow past the GPU block pool, the newest are
+    /// swapped out to host (vLLM-style) and swapped back in when blocks free up.
+    pub fn run(&mut self, rx: mpsc::Receiver<GenerateRequest>) {
         let mut waiting: VecDeque<Sequence> = VecDeque::new();
         let mut running: Vec<Sequence> = Vec::new();
+        let mut swapped: Vec<Sequence> = Vec::new();
         let mut next_id: u64 = 0;
 
         eprintln!("[scheduler] Listening for requests...");
 
         loop {
-            // Step 1: Remove finished sequences
+            // Step 1: Remove finished sequences and return their slots.
+            let finished_slots: Vec<usize> = running.iter()
+                .filter(|s| is_finished(s))
+                .filter_map(|s| s.seq_slot)
+                .collect();
+            for slot in finished_slots {
+                self.paged_cache.free_sequence(slot);
+            }
             running.retain(|seq| !is_finished(seq));
 
-            // Step 2: Admit new sequences from waiting queue
+            // Step 2: Swap previously-evicted sequences back in when there is
-            while running.len() < self.max_batch_size {
+            // room (oldest first). They resume decoding from where they paused.
-                if let Some(seq) = waiting.pop_front() {
+            while running.len() < self.max_batch_size && !swapped.is_empty() {
+                let slot = swapped[0].seq_slot.expect("swapped slot");
+                if !self.paged_cache.can_swap_in(slot) { break; }
+                self.paged_cache.swap_in(slot).expect("swap_in");
+                let seq = swapped.remove(0);
+                eprintln!("[scheduler] swapped in seq {} ({} blocks)", seq.id, self.paged_cache.block_count(slot));
                 running.push(seq);
-                } else {
+            }
-                    break;
+
+            // Step 3: Admit new sequences (block-aware). Only admit if the GPU
+            // pool can hold the prompt AND leave one block of decode headroom
+            // per already-running sequence, so admission never starves decode.
+            {
+                let mut avail = self.paged_cache.free_blocks();
+                let decode_reserve = running.len();
+                while running.len() < self.max_batch_size {
+                    let Some(front) = waiting.front() else { break; };
+                    let prompt_blocks = front.prompt_tokens.len().div_ceil(BLOCK_SIZE).max(1);
+                    if avail < prompt_blocks + decode_reserve { break; }
+                    let free_slot = (0..self.paged_cache.max_seqs())
+                        .find(|&s| self.paged_cache.is_slot_free(s));
+                    let Some(slot) = free_slot else { break; };
+                    let mut seq = waiting.pop_front().unwrap();
+                    self.paged_cache.register_sequence(slot).expect("register paged slot");
+                    seq.seq_slot = Some(slot);
+                    running.push(seq);
+                    avail -= prompt_blocks; // projected free after this seq prefills
                 }
             }
 
-            // Step 3: If nothing to do, blocking wait for new request
+            // Step 4: If nothing to do, blocking wait for new request.
-            if running.is_empty() {
+            if running.is_empty() && waiting.is_empty() && swapped.is_empty() {
                 match rx.recv() {
                     Ok(req) => {
                         let seq = self.make_sequence(req, &mut next_id);
-                        running.push(seq);
+                        waiting.push_back(seq);
+                        continue;
                     }
                     Err(_) => break, // channel closed
                 }
             }
+            // Nothing runnable this iteration (e.g. all swapped, waiting on
+            // blocks to free): loop to retry swap-in/admission next iteration.
+            if running.is_empty() {
+                continue;
+            }
 
-            // Step 4: Process one iteration for all running sequences
+            // Step 5a: Process prefills (one at a time — different prompt lengths).
+            // Admission guaranteed block headroom, so ensure_capacity won't starve.
+            let mut newly_prefilled = Vec::new();
             for seq in running.iter_mut() {
                 if !seq.prefilled {
-                    // Prefill
+                    let slot = seq.seq_slot.expect("slot");
-                    let logits = self.model.forward_gpu_cache(&seq.prompt_tokens, &mut seq.kv_cache);
+                    let logits = self.model.forward_prefill_paged(
-                    let next = sample_greedy(&logits);
+                        &seq.prompt_tokens, slot, &mut self.paged_cache,
+                    );
+                    let next = sample(&logits, &seq.sampling);
                     seq.generated_tokens.push(next);
                     seq.prefilled = true;
-                    self.emit_token(seq, next);
+                    emit_token(&self.tokenizer, seq, next);
-                } else {
+                    newly_prefilled.push(seq.id);
-                    // Decode one token
-                    let last = *seq.generated_tokens.last().unwrap();
-                    let logits = self.model.forward_gpu_cache(&[last], &mut seq.kv_cache);
-                    let next = sample_greedy(&logits);
-                    seq.generated_tokens.push(next);
-                    self.emit_token(seq, next);
                 }
             }
 
-            // Step 5: Check for newly arrived requests (non-blocking)
+            // Step 5b: Ensure block headroom for this decode step; preempt the
+            // newest running sequences to host if the pool can't cover it.
+            let mut needed = decode_block_need(&self.paged_cache, &running, &newly_prefilled);
+            while self.paged_cache.free_blocks() < needed {
+                // Victim: newest prefilled, decoding (not just-prefilled) sequence.
+                let victim = (0..running.len()).rev().find(|&p| {
+                    running[p].prefilled
+                        && !newly_prefilled.contains(&running[p].id)
+                        && running[p].seq_slot.is_some()
+                });
+                let Some(pos) = victim else { break; };
+                let seq = running.remove(pos);
+                let slot = seq.seq_slot.unwrap();
+                if self.paged_cache.can_swap_out(slot) {
+                    let nblocks = self.paged_cache.block_count(slot);
+                    self.paged_cache.swap_out(slot).expect("swap_out");
+                    eprintln!("[scheduler] preempt: swapped out seq {} ({nblocks} blocks) to host", seq.id);
+                    swapped.push(seq);
+                    needed = decode_block_need(&self.paged_cache, &running, &newly_prefilled);
+                } else {
+                    running.insert(pos, seq); // CPU pool full — can't evict further
+                    break;
+                }
+            }
+
+            // Step 5c: Batched paged decode for the surviving prefilled sequences.
+            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] paged decode active");
+                });
+
+                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| self.paged_cache.seq_len(running[i].seq_slot.unwrap()))
+                    .collect();
+                let slots: Vec<usize> = decode_indices.iter()
+                    .map(|&i| running[i].seq_slot.unwrap())
+                    .collect();
+
+                let logits = self.model.forward_decode_paged(
+                    &tokens, &positions, &slots, &mut self.paged_cache,
+                );
+
+                // 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 {
+                        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 {
+                        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);
+                    emit_token(&self.tokenizer, &mut running[i], next);
+                }
+            }
+
+            // Step 6: Check for newly arrived requests (non-blocking)
             loop {
                 match rx.try_recv() {
                     Ok(req) => {
@@ -117,45 +294,68 @@ impl Engine {
         }
     }
 
-    fn make_sequence(&self, req: GenerateRequest, next_id: &mut u64) -> Sequence {
+    fn make_sequence(&mut self, req: GenerateRequest, next_id: &mut u64) -> Sequence {
         let id = *next_id;
         *next_id += 1;
-        let kv_cache = GpuKVCache::new(&self.config, self.max_seq_len, DType::BF16);
         Sequence {
             id,
             prompt_tokens: req.prompt_tokens,
             generated_tokens: Vec::new(),
             max_tokens: req.max_tokens,
-            kv_cache,
+            sampling: req.sampling,
+            seq_slot: None,
             sender: req.sender,
             prefilled: false,
+            eos_token_id: self.tokenizer.eos_token_id(),
+            decode_buffer: Vec::new(),
+            created_at: Instant::now(),
+        }
     }
 }
 
-    fn emit_token(&self, seq: &Sequence, token_id: u32) {
+/// Total additional GPU blocks the next decode step needs across all
-        let text = self.tokenizer.decode(&[token_id]);
+/// currently-decoding (prefilled, not just-prefilled) sequences.
+fn decode_block_need(paged: &PagedKVCache, running: &[Sequence], newly_prefilled: &[u64]) -> usize {
+    running.iter()
+        .filter(|s| s.prefilled && !newly_prefilled.contains(&s.id))
+        .filter_map(|s| s.seq_slot)
+        .map(|slot| paged.additional_blocks_needed(slot, 1))
+        .sum()
+}
 
-        if self.tokenizer.eos_token_id() == Some(token_id) {
+fn emit_token(tokenizer: &Tokenizer, seq: &mut Sequence, token_id: u32) {
-            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
+    if tokenizer.eos_token_id() == Some(token_id) {
+        let tail = tokenizer.flush_decode_stream(&mut seq.decode_buffer);
+        send_token_if_nonempty(seq, tail);
         let _ = seq.sender.blocking_send(GenerateEvent::Done {
             finish_reason: "stop".to_string(),
         });
-        } else if seq.generated_tokens.len() >= seq.max_tokens {
+        return;
-            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
+    }
+
+    let text = tokenizer.decode_token_stream(token_id, &mut seq.decode_buffer);
+    if seq.generated_tokens.len() >= seq.max_tokens {
+        let tail = tokenizer.flush_decode_stream(&mut seq.decode_buffer);
+        send_token_if_nonempty(seq, text);
+        send_token_if_nonempty(seq, tail);
         let _ = seq.sender.blocking_send(GenerateEvent::Done {
             finish_reason: "length".to_string(),
         });
     } else {
-            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
+        send_token_if_nonempty(seq, text);
     }
 }
+
+fn send_token_if_nonempty(seq: &Sequence, text: String) {
+    if !text.is_empty() {
+        let id = *seq.generated_tokens.last().unwrap_or(&0);
+        let _ = seq.sender.blocking_send(GenerateEvent::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:
+    seq.sender.is_closed() || seq.eos_token_id == Some(last)
-    // If sender is closed (receiver dropped), also consider finished.
-    seq.sender.is_closed() || last == 151645 // Qwen3 EOS token ID (hardcoded for now)
 }

crates/xserv-server/src/main.rs

@@ -1,22 +1,25 @@
 mod api;
 mod engine;
+mod tp_engine;
 
 use axum::{routing::{get, post}, Extension, Router};
 use std::path::PathBuf;
 use std::sync::{mpsc, Arc, Mutex};
 use engine::GenerateRequest;
+use xserv_model::ModelConfig;
 
 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]");
+        eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N] [--tp N]");
         std::process::exit(1);
     }
 
@@ -30,7 +33,37 @@ async fn main() {
         .position(|a| a == "--max-batch")
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
-        .unwrap_or(4);
+        .unwrap_or(4)
+        .max(1);
+    let requested_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)
+        .max(1);
+    let swap_space_gb: usize = args.iter()
+        .position(|a| a == "--swap-space-gb")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(8);
+    let tp: usize = args.iter()
+        .position(|a| a == "--tp")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(1)
+        .max(1);
+    let model_config = ModelConfig::from_file(&model_dir.join("config.json"));
+    let model_max_seq_len = model_config.max_seq_len();
+    if model_max_seq_len == 0 {
+        eprintln!("model config has invalid max_seq_len=0");
+        std::process::exit(1);
+    }
+    let max_seq_len = requested_max_seq_len.min(model_max_seq_len);
+    if max_seq_len != requested_max_seq_len {
+        eprintln!(
+            "[server] --max-seq-len {requested_max_seq_len} exceeds model limit {model_max_seq_len}; using {max_seq_len}"
+        );
+    }
 
     let model_name = model_dir.file_name()
         .map(|n| n.to_string_lossy().to_string())
@@ -43,14 +76,20 @@ async fn main() {
 
     let model_dir_clone = model_dir.clone();
     std::thread::spawn(move || {
-        let engine = engine::Engine::load(&model_dir_clone, max_batch);
+        if tp <= 1 {
+            let mut engine = engine::Engine::load_with_swap(&model_dir_clone, max_batch, max_seq_len, swap_space_gb);
             engine.run(rx);
+        } else {
+            // Tensor-parallel path: rank-0 coordinator + worker rank threads.
+            tp_engine::run_tp(&model_dir_clone, tp, max_seq_len, 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()
@@ -60,7 +99,7 @@ async fn main() {
         .layer(Extension(state));
 
     let addr = format!("0.0.0.0:{port}");
-    eprintln!("[server] Listening on {addr} (max_batch={max_batch})");
+    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-server/src/tp_engine.rs

@@ -0,0 +1,195 @@
+//! Tensor-parallel inference engine for the HTTP server.
+//!
+//! Serial coordinator model: one rank-0 coordinator thread (the caller) drives
+//! generation and owns the scheduler; ranks 1..world are worker threads. For
+//! each step the coordinator broadcasts a command (Register/Prefill/Decode/Free)
+//! to the workers and runs the same op on its own shard; the per-layer NCCL
+//! AllReduces keep all ranks in lockstep. Only the coordinator samples — the
+//! chosen token is carried in the next Decode command, so this is correct for
+//! both greedy and stochastic sampling.
+//!
+//! Requests are processed one at a time (sufficient for the quality benchmark,
+//! which issues serial requests). Continuous batching across ranks is future
+//! work; the single-GPU `Engine` still handles TP=1.
+
+use std::path::{Path, PathBuf};
+use std::sync::mpsc;
+use std::sync::Arc;
+use std::thread;
+
+use xserv_distributed::{TpContext, UniqueId};
+use xserv_model::loader;
+use xserv_model::{sample, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
+use xserv_tensor::{DType, Device};
+use xserv_tokenizer::Tokenizer;
+
+use crate::engine::{GenerateEvent, GenerateRequest};
+
+#[derive(Clone)]
+enum TpCommand {
+    Register(usize),
+    Free(usize),
+    Prefill { tokens: Vec<u32>, slot: usize },
+    Decode { tokens: Vec<u32>, positions: Vec<usize>, slots: Vec<usize> },
+    Shutdown,
+}
+
+struct RankCtx {
+    model: Qwen3,
+    cache: PagedKVCache,
+}
+
+fn build_rank(
+    model_dir: &Path,
+    config: &ModelConfig,
+    rank: usize,
+    world: usize,
+    device: u32,
+    max_seq_len: usize,
+    tp: Option<Arc<TpContext>>,
+) -> RankCtx {
+    let weights = loader::load_model_dir(model_dir, Device::Cpu);
+    let model = Qwen3::from_weights_tp(config.clone(), weights, rank, world, device, tp);
+    let local_kv = config.num_kv_heads() / world;
+    let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE);
+    let total_blocks = max_blocks_per_seq + 8;
+    let cache = PagedKVCache::new_tp(
+        config, local_kv, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, device,
+    );
+    RankCtx { model, cache }
+}
+
+fn worker_loop(
+    rank: usize,
+    world: usize,
+    id: UniqueId,
+    model_dir: PathBuf,
+    config: ModelConfig,
+    max_seq_len: usize,
+    cmd_rx: mpsc::Receiver<TpCommand>,
+    ack_tx: mpsc::Sender<()>,
+) {
+    let tp = Arc::new(TpContext::init(rank, world, id, rank as u32));
+    let mut rc = build_rank(&model_dir, &config, rank, world, rank as u32, max_seq_len, Some(tp));
+    while let Ok(cmd) = cmd_rx.recv() {
+        match cmd {
+            TpCommand::Register(slot) => {
+                let _ = rc.cache.register_sequence(slot);
+            }
+            TpCommand::Free(slot) => rc.cache.free_sequence(slot),
+            TpCommand::Prefill { tokens, slot } => {
+                let _ = rc.model.forward_prefill_paged(&tokens, slot, &mut rc.cache);
+            }
+            TpCommand::Decode { tokens, positions, slots } => {
+                let _ = rc.model.forward_decode_paged(&tokens, &positions, &slots, &mut rc.cache);
+            }
+            TpCommand::Shutdown => {
+                let _ = ack_tx.send(());
+                break;
+            }
+        }
+        let _ = ack_tx.send(());
+    }
+}
+
+/// Run the TP coordinator (rank 0) on the calling thread. Spawns worker ranks
+/// internally and consumes generation requests from `rx`.
+pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Receiver<GenerateRequest>) {
+    assert!(world >= 2, "run_tp requires world >= 2");
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    assert!(
+        config.num_kv_heads() % world == 0,
+        "num_kv_heads {} not divisible by tp {world}",
+        config.num_kv_heads()
+    );
+    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+    let id = xserv_distributed::get_unique_id();
+
+    // Spawn worker ranks 1..world.
+    let (ack_tx, ack_rx) = mpsc::channel::<()>();
+    let mut cmd_txs: Vec<mpsc::Sender<TpCommand>> = Vec::new();
+    for rank in 1..world {
+        let (ctx_tx, ctx_rx) = mpsc::channel::<TpCommand>();
+        cmd_txs.push(ctx_tx);
+        let ack_tx = ack_tx.clone();
+        let model_dir = model_dir.to_path_buf();
+        let config = config.clone();
+        thread::spawn(move || {
+            worker_loop(rank, world, id, model_dir, config, max_seq_len, ctx_rx, ack_tx);
+        });
+    }
+
+    // Rank 0 (this thread).
+    let tp = Arc::new(TpContext::init(0, world, id, 0));
+    let mut rc = build_rank(model_dir, &config, 0, world, 0, max_seq_len, Some(tp));
+    eprintln!("[tp-engine] ready (tp={world}, max_seq_len={max_seq_len})");
+
+    let eos = tokenizer.eos_token_id();
+    let n_workers = world - 1;
+    let broadcast = |txs: &[mpsc::Sender<TpCommand>], cmd: TpCommand| {
+        for t in txs {
+            let _ = t.send(cmd.clone());
+        }
+    };
+    let wait_acks = |rx: &mpsc::Receiver<()>| {
+        for _ in 0..n_workers {
+            let _ = rx.recv();
+        }
+    };
+
+    let slot = 0usize;
+    while let Ok(req) = rx.recv() {
+        broadcast(&cmd_txs, TpCommand::Register(slot));
+        rc.cache.register_sequence(slot).expect("register slot");
+        wait_acks(&ack_rx);
+
+        // Prefill.
+        broadcast(&cmd_txs, TpCommand::Prefill { tokens: req.prompt_tokens.clone(), slot });
+        let logits = rc.model.forward_prefill_paged(&req.prompt_tokens, slot, &mut rc.cache);
+        wait_acks(&ack_rx);
+        let mut next = sample(&logits, &req.sampling);
+
+        let mut decode_buf: Vec<u8> = Vec::new();
+        let mut generated = 1usize;
+        emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
+
+        let finish = loop {
+            if eos == Some(next) {
+                break "stop";
+            }
+            if generated >= req.max_tokens {
+                break "length";
+            }
+            let pos = rc.cache.seq_len(slot);
+            broadcast(&cmd_txs, TpCommand::Decode { tokens: vec![next], positions: vec![pos], slots: vec![slot] });
+            let logits = rc.model.forward_decode_paged(&[next], &[pos], &[slot], &mut rc.cache);
+            wait_acks(&ack_rx);
+            next = sample(&logits, &req.sampling);
+            generated += 1;
+            emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
+        };
+
+        let tail = tokenizer.flush_decode_stream(&mut decode_buf);
+        if !tail.is_empty() {
+            let _ = req.sender.blocking_send(GenerateEvent::Token { id: next, text: tail });
+        }
+        let _ = req.sender.blocking_send(GenerateEvent::Done { finish_reason: finish.to_string() });
+
+        broadcast(&cmd_txs, TpCommand::Free(slot));
+        rc.cache.free_sequence(slot);
+        wait_acks(&ack_rx);
+    }
+
+    broadcast(&cmd_txs, TpCommand::Shutdown);
+}
+
+/// Stream a token's decoded text to the client (EOS contributes no text).
+fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, eos: Option<u32>, buf: &mut Vec<u8>) {
+    if eos == Some(token_id) {
+        return;
+    }
+    let text = tokenizer.decode_token_stream(token_id, buf);
+    if !text.is_empty() {
+        let _ = req.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
+    }
+}

crates/xserv-tensor/src/lib.rs

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

crates/xserv-tensor/src/storage.rs

@@ -3,7 +3,7 @@ use xserv_cuda::{GpuBuffer, Result as CudaResult};
 
 enum StorageInner {
     Cpu { data: Vec<u8> },
-    Cuda { buffer: GpuBuffer },
+    Cuda { buffer: GpuBuffer, device: u32 },
 }
 
 /// Reference-counted storage for tensor data. Multiple tensors can share
@@ -31,21 +31,21 @@ impl Storage {
         Self(Arc::new(StorageInner::Cpu { data }))
     }
 
-    pub fn cuda(buffer: GpuBuffer) -> Self {
+    pub fn cuda(buffer: GpuBuffer, device: u32) -> Self {
-        Self(Arc::new(StorageInner::Cuda { buffer }))
+        Self(Arc::new(StorageInner::Cuda { buffer, device }))
     }
 
     pub fn device(&self) -> Device {
         match self.0.as_ref() {
             StorageInner::Cpu { .. } => Device::Cpu,
-            StorageInner::Cuda { .. } => Device::Cuda(0),
+            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(),
+            StorageInner::Cuda { buffer, .. } => buffer.len(),
         }
     }
 
@@ -59,7 +59,7 @@ impl Storage {
 
     pub fn gpu_buffer(&self) -> &GpuBuffer {
         match self.0.as_ref() {
-            StorageInner::Cuda { buffer } => buffer,
+            StorageInner::Cuda { buffer, .. } => buffer,
             StorageInner::Cpu { .. } => panic!("cannot access CPU storage as GPU buffer"),
         }
     }
@@ -71,11 +71,11 @@ impl Storage {
             return Ok(self.clone());
         }
         match (current, target) {
-            (Device::Cpu, Device::Cuda(_dev)) => {
+            (Device::Cpu, Device::Cuda(dev)) => {
                 let cpu_data = self.as_cpu_bytes();
-                let mut buf = GpuBuffer::alloc(cpu_data.len())?;
+                let mut buf = xserv_cuda::allocator::cached_alloc(cpu_data.len())?;
                 buf.copy_from_host(cpu_data)?;
-                Ok(Storage::cuda(buf))
+                Ok(Storage::cuda(buf, dev))
             }
             (Device::Cuda(_), Device::Cpu) => {
                 let gpu_buf = self.gpu_buffer();
@@ -83,11 +83,11 @@ impl Storage {
                 gpu_buf.copy_to_host(&mut data)?;
                 Ok(Storage::cpu(data))
             }
-            (Device::Cuda(_), Device::Cuda(_)) => {
+            (Device::Cuda(_), Device::Cuda(dev)) => {
                 let src = self.gpu_buffer();
-                let mut dst = GpuBuffer::alloc(src.len())?;
+                let mut dst = xserv_cuda::allocator::cached_alloc(src.len())?;
                 dst.copy_from_device(src)?;
-                Ok(Storage::cuda(dst))
+                Ok(Storage::cuda(dst, dev))
             }
             _ => unreachable!(),
         }
@@ -97,10 +97,10 @@ impl Storage {
     pub fn deep_copy(&self) -> CudaResult<Self> {
         match self.0.as_ref() {
             StorageInner::Cpu { data } => Ok(Storage::cpu(data.clone())),
-            StorageInner::Cuda { buffer } => {
+            StorageInner::Cuda { buffer, device } => {
-                let mut dst = GpuBuffer::alloc(buffer.len())?;
+                let mut dst = xserv_cuda::allocator::cached_alloc(buffer.len())?;
                 dst.copy_from_device(buffer)?;
-                Ok(Storage::cuda(dst))
+                Ok(Storage::cuda(dst, *device))
             }
         }
     }
@@ -109,10 +109,24 @@ impl Storage {
     pub fn zeros(len_bytes: usize, device: Device) -> CudaResult<Self> {
         match device {
             Device::Cpu => Ok(Storage::cpu(vec![0u8; len_bytes])),
-            Device::Cuda(_) => {
+            Device::Cuda(dev) => {
-                let mut buf = GpuBuffer::alloc(len_bytes)?;
+                let mut buf = xserv_cuda::allocator::cached_alloc(len_bytes)?;
                 buf.zero()?;
-                Ok(Storage::cuda(buf))
+                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

@@ -1,7 +1,21 @@
+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
@@ -51,6 +65,22 @@ impl Tensor {
         }
     }
 
+    /// 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 {
@@ -123,10 +153,15 @@ impl Tensor {
     pub fn unsqueeze(&self, dim: usize) -> Self {
         assert!(dim <= self.ndim());
         let mut new_shape = self.shape.clone();
-        let mut new_strides = self.strides.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 };
-        new_strides.insert(dim, stride_val);
+            s.insert(dim, stride_val);
+            s
+        };
         Self {
             storage: self.storage.clone(),
             shape: new_shape,
@@ -142,9 +177,12 @@ impl Tensor {
         if self.is_contiguous() {
             return self.clone();
         }
-        // For GPU tensors: round-trip through CPU (correct but slow).
+        // For GPU tensors: use the registered GPU kernel if available,
-        // TODO: write a GPU contiguous-copy kernel for performance.
+        // 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());
@@ -237,3 +275,58 @@ impl std::fmt::Debug for Tensor {
         )
     }
 }
+
+#[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-tokenizer/src/bpe.rs

@@ -41,6 +41,7 @@ enum MergeEntry {
 struct AddedToken {
     id: u32,
     content: String,
+    #[allow(dead_code)]
     special: bool,
 }
 
@@ -51,8 +52,6 @@ impl Tokenizer {
         let tj: TokenizerJson = serde_json::from_str(&data)
             .unwrap_or_else(|e| panic!("failed to parse tokenizer.json: {e}"));
 
-        let byte_fallback = tj.model.byte_fallback;
-
         // Build encoder: token bytes → ID
         // All HF tokenizers use GPT-2 byte-to-unicode mapping for vocab keys.
         let mut encoder = HashMap::new();
@@ -92,21 +91,22 @@ impl Tokenizer {
             }
         }
 
-        // Special tokens
+        // Added tokens are matched as indivisible tokens by HF tokenizers,
+        // even when their `special` flag is false (for example Qwen3's
+        // <think> and </think> 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);
-                }
-            }
         }
+        let eos_token_id = special_tokens
+            .get("<|im_end|>")
+            .or_else(|| special_tokens.get("<|end_of_text|>"))
+            .or_else(|| special_tokens.get("<|endoftext|>"))
+            .copied();
 
         // Pre-tokenization regex
         let pre_tokenize_re = if byte_fallback {
@@ -170,10 +170,26 @@ impl Tokenizer {
             }
             // Fall back to per-byte encoding
             let word_bytes: Vec<u8> = word.bytes().collect();
-            let mut token_ids: Vec<u32> = word_bytes.iter().map(|&b| {
+            let mut token_ids: Vec<u32> = word_bytes.iter().filter_map(|&b| {
-                *self.encoder.get(&vec![b]).unwrap_or_else(|| {
+                if let Some(&id) = self.encoder.get(&vec![b]) {
-                    panic!("byte {b} (0x{b:02X}) not in vocab")
+                    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
@@ -216,6 +232,19 @@ impl Tokenizer {
         String::from_utf8_lossy(&bytes).into_owned()
     }
 
+    pub fn decode_token_stream(&self, token_id: u32, pending: &mut Vec<u8>) -> String {
+        if let Some(bytes) = self.decoder.get(token_id as usize) {
+            pending.extend_from_slice(bytes);
+        }
+        take_valid_utf8(pending)
+    }
+
+    pub fn flush_decode_stream(&self, pending: &mut Vec<u8>) -> String {
+        let text = String::from_utf8_lossy(pending).into_owned();
+        pending.clear();
+        text
+    }
+
     pub fn eos_token_id(&self) -> Option<u32> {
         self.eos_token_id
     }
@@ -236,6 +265,31 @@ fn token_str_to_bytes(s: &str) -> Vec<u8> {
     s.chars().map(|c| unicode_to_byte(c)).collect()
 }
 
+fn take_valid_utf8(pending: &mut Vec<u8>) -> String {
+    match std::str::from_utf8(pending) {
+        Ok(text) => {
+            let text = text.to_string();
+            pending.clear();
+            text
+        }
+        Err(err) => {
+            let valid_up_to = err.valid_up_to();
+            if valid_up_to == 0 {
+                if let Some(error_len) = err.error_len() {
+                    let invalid_len = error_len.min(pending.len());
+                    let text = String::from_utf8_lossy(&pending[..invalid_len]).into_owned();
+                    pending.drain(..invalid_len);
+                    return text;
+                }
+                return String::new();
+            }
+            let text = String::from_utf8_lossy(&pending[..valid_up_to]).into_owned();
+            pending.drain(..valid_up_to);
+            text
+        }
+    }
+}
+
 /// 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
@@ -265,3 +319,49 @@ fn unicode_to_byte(c: char) -> u8 {
         panic!("unmapped unicode char U+{:04X} in tokenizer", c as u32)
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{take_valid_utf8, Tokenizer};
+
+    #[test]
+    fn qwen_added_tokens_are_indivisible_and_im_end_is_eos() {
+        let path =
+            std::env::temp_dir().join(format!("xserv-tokenizer-test-{}.json", std::process::id()));
+        std::fs::write(
+            &path,
+            r#"{
+                "model": {
+                    "vocab": {},
+                    "merges": [],
+                    "byte_fallback": false
+                },
+                "added_tokens": [
+                    {"id":151643,"content":"<|endoftext|>","special":true},
+                    {"id":151644,"content":"<|im_start|>","special":true},
+                    {"id":151645,"content":"<|im_end|>","special":true},
+                    {"id":151667,"content":"<think>","special":false},
+                    {"id":151668,"content":"</think>","special":false}
+                ]
+            }"#,
+        )
+        .unwrap();
+
+        let tokenizer = Tokenizer::from_file(&path);
+        let _ = std::fs::remove_file(&path);
+
+        assert_eq!(tokenizer.eos_token_id(), Some(151645));
+        assert_eq!(tokenizer.encode("<think>"), vec![151667]);
+        assert_eq!(tokenizer.encode("</think>"), vec![151668]);
+        assert_eq!(tokenizer.decode(&[151645]), "<|im_end|>");
+    }
+
+    #[test]
+    fn stream_decode_buffers_incomplete_utf8() {
+        let mut pending = vec![0xF0, 0x9F];
+        assert_eq!(take_valid_utf8(&mut pending), "");
+        pending.extend_from_slice(&[0x98, 0x8A, b'!']);
+        assert_eq!(take_valid_utf8(&mut pending), "😊!");
+        assert!(pending.is_empty());
+    }
+}

csrc/activation/activations.cu

@@ -1,5 +1,6 @@
 #include <cuda_bf16.h>
 #include <math.h>
+#include "../common.cuh"
 
 // GELU (tanh approximation):
 // gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
@@ -45,6 +46,18 @@ __global__ void scale_bf16_kernel(const __nv_bfloat16* x, __nv_bfloat16* out, fl
     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;
@@ -71,6 +84,7 @@ 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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_gelu_bf16(const void* x, void* out, int n, void* stream) {
@@ -78,12 +92,14 @@ void launch_gelu_bf16(const void* x, void* out, int n, void* stream) {
     int grid = (n + block - 1) / block;
     gelu_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_silu_bf16(const void* x, void* out, int n, void* stream) {
@@ -91,6 +107,7 @@ void launch_silu_bf16(const void* x, void* out, int n, void* stream) {
     int grid = (n + block - 1) / block;
     silu_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_scale_f32(const void* x, void* out, float scale, int n, void* stream) {
@@ -98,6 +115,7 @@ void launch_scale_f32(const void* x, void* out, float scale, int n, void* stream
     int grid = (n + block - 1) / block;
     scale_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const float*)x, (float*)out, scale, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_scale_bf16(const void* x, void* out, float scale, int n, void* stream) {
@@ -105,6 +123,7 @@ void launch_scale_bf16(const void* x, void* out, float scale, int n, void* strea
     int grid = (n + block - 1) / block;
     scale_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)x, (__nv_bfloat16*)out, scale, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_add_f32(const void* a, const void* b, void* out, int n, void* stream) {
@@ -112,24 +131,36 @@ void launch_add_f32(const void* a, const void* b, void* out, int n, void* stream
     int grid = (n + block - 1) / block;
     add_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const float*)a, (const float*)b, (float*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 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);
+    CUDA_CHECK_LAST_ERROR();
 }
 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);
+    CUDA_CHECK_LAST_ERROR();
 }
 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);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/attention/causal_mask.cu

@@ -1,4 +1,5 @@
 #include <cuda_bf16.h>
+#include "../common.cuh"
 
 // Apply causal mask: set scores[row][col] = -inf where col > row + offset.
 // offset is used for KV cache: when query starts at position `offset`,
@@ -27,8 +28,7 @@ __global__ void causal_mask_bf16(
     int col = blockIdx.x * blockDim.x + threadIdx.x;
 
     if (col < cols && col > row + offset) {
-        // BF16 doesn't have proper -inf literal, use a very large negative
+        scores[batch_idx * rows * cols + row * cols + col] = __float2bfloat16(-INFINITY);
-        scores[batch_idx * rows * cols + row * cols + col] = __float2bfloat16(-1e9f);
     }
 }
 
@@ -40,6 +40,7 @@ void launch_causal_mask_f32(void* scores, int batch, int rows, int cols,
     dim3 grid((cols + block - 1) / block, rows, batch);
     causal_mask_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
         (float*)scores, rows, cols, offset);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_causal_mask_bf16(void* scores, int batch, int rows, int cols,
@@ -48,6 +49,7 @@ void launch_causal_mask_bf16(void* scores, int batch, int rows, int cols,
     dim3 grid((cols + block - 1) / block, rows, batch);
     causal_mask_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (__nv_bfloat16*)scores, rows, cols, offset);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/attention/flash_attention.cu

@@ -0,0 +1,419 @@
+#include <cuda_bf16.h>
+#include <float.h>
+#include "../common.cuh"
+
+// 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
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+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
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+}

csrc/attention/paged_attention.cu

@@ -0,0 +1,215 @@
+#include <cuda_bf16.h>
+#include <float.h>
+#include "../common.cuh"
+
+// Paged decode attention kernel for BF16 with FP32 accumulation.
+//
+// Reads K/V from a paged pool indexed by a per-sequence block table.
+// One CUDA block per (sequence, q_head). Each block streams over the
+// sequence's KV positions and accumulates attention output via online
+// softmax.
+//
+// Layouts:
+//   Q          [batch, num_q_heads, 1, head_dim]                    BF16
+//   K_cache    [num_blocks, num_kv_heads, BLOCK_SIZE, head_dim]     BF16
+//   V_cache    same
+//   block_tables  [max_seqs, max_blocks_per_seq]                    int32
+//                 — the i-th sequence in this launch reads row
+//                   block_tables[seq_slot[i] * stride + ...].
+//                 For simplicity the launch passes a packed row table
+//                 [batch, max_blocks_per_seq] (already gathered for the
+//                 active batch) so we just index by blockIdx.x_seq.
+//   context_lens  [batch] int32  — number of valid tokens per sequence.
+//
+// One CUDA block: 256 threads, head_dim <= 128.
+
+#define PAGED_BLOCK_SIZE 16
+#define PAGED_THREADS 256
+#define PAGED_HEAD_DIM_MAX 128
+
+__global__ void paged_decode_attention_bf16_kernel(
+    const __nv_bfloat16* __restrict__ Q,
+    const __nv_bfloat16* __restrict__ K_cache,
+    const __nv_bfloat16* __restrict__ V_cache,
+    __nv_bfloat16* __restrict__ O,
+    const int* __restrict__ block_tables,    // [batch, max_blocks_per_seq]
+    const int* __restrict__ context_lens,    // [batch]
+    int num_q_heads, int num_kv_heads,
+    int head_dim, int max_blocks_per_seq,
+    float scale
+) {
+    int seq_idx = blockIdx.y;          // batch dim
+    int q_head  = blockIdx.x;          // 0 .. num_q_heads-1
+    int tid     = threadIdx.x;
+
+    int kv_len = context_lens[seq_idx];
+    if (kv_len <= 0) {
+        // Nothing to attend over; zero output for safety.
+        if (tid < head_dim) {
+            O[((long long)seq_idx * num_q_heads + q_head) * head_dim + tid] =
+                __float2bfloat16(0.0f);
+        }
+        return;
+    }
+
+    // GQA mapping
+    int heads_per_group = num_q_heads / num_kv_heads;
+    int kv_head = q_head / heads_per_group;
+
+    // Pointers
+    const __nv_bfloat16* Q_ptr = Q +
+        ((long long)seq_idx * num_q_heads + q_head) * head_dim;
+    __nv_bfloat16* O_ptr = O +
+        ((long long)seq_idx * num_q_heads + q_head) * head_dim;
+    const int* bt = block_tables + (long long)seq_idx * max_blocks_per_seq;
+
+    // Load Q vector into registers.
+    float q_reg[PAGED_HEAD_DIM_MAX];
+    for (int d = 0; d < head_dim; d++) {
+        q_reg[d] = __bfloat162float(Q_ptr[d]);
+    }
+
+    // Per-thread online softmax state.
+    float local_max = -INFINITY;
+    float local_sum = 0.0f;
+    float local_O[PAGED_HEAD_DIM_MAX];
+    for (int d = 0; d < head_dim; d++) local_O[d] = 0.0f;
+
+    int kv_stride_block = num_kv_heads * PAGED_BLOCK_SIZE * head_dim;
+    int kv_stride_head  = PAGED_BLOCK_SIZE * head_dim;
+
+    // Each thread handles positions tid, tid+PAGED_THREADS, ...
+    for (int pos = tid; pos < kv_len; pos += PAGED_THREADS) {
+        int logical_blk = pos / PAGED_BLOCK_SIZE;
+        int slot_in_blk = pos % PAGED_BLOCK_SIZE;
+        int phys_blk = bt[logical_blk];
+
+        const __nv_bfloat16* K_pos = K_cache
+            + (long long)phys_blk * kv_stride_block
+            + kv_head * kv_stride_head
+            + slot_in_blk * head_dim;
+        const __nv_bfloat16* V_pos = V_cache
+            + (long long)phys_blk * kv_stride_block
+            + kv_head * kv_stride_head
+            + slot_in_blk * head_dim;
+
+        // dot(Q, K[pos]) * scale
+        float dot = 0.0f;
+        for (int d = 0; d < head_dim; d++) {
+            dot += q_reg[d] * __bfloat162float(K_pos[d]);
+        }
+        float s = dot * scale;
+
+        float new_max = fmaxf(local_max, s);
+        float correction = expf(local_max - new_max);
+        float p = expf(s - new_max);
+
+        local_sum = local_sum * correction + p;
+        for (int d = 0; d < head_dim; d++) local_O[d] *= correction;
+
+        // Accumulate weighted V.
+        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 ----
+    __shared__ float smem_max[32];
+    __shared__ float smem_sum[32];
+    __shared__ float smem_O[PAGED_HEAD_DIM_MAX];
+
+    int lane = tid & 31;
+    int warp_id = tid >> 5;
+    int num_warps = PAGED_THREADS >> 5;
+
+    // Step 1: block-wide max
+    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: rescale local state to 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
+    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, dim by dim
+    for (int d = tid; d < head_dim; d += PAGED_THREADS) smem_O[d] = 0.0f;
+    __syncthreads();
+
+    for (int d = 0; d < head_dim; d++) {
+        float val = local_O[d];
+        #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();
+
+    float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
+    for (int d = tid; d < head_dim; d += PAGED_THREADS) {
+        O_ptr[d] = __float2bfloat16(smem_O[d] * inv_sum);
+    }
+}
+
+extern "C" {
+
+void launch_paged_decode_attention_bf16(
+    const void* Q,
+    const void* K_cache,
+    const void* V_cache,
+    void* O,
+    const int* block_tables,
+    const int* context_lens,
+    int batch, int num_q_heads, int num_kv_heads,
+    int head_dim, int max_blocks_per_seq,
+    float scale, void* stream
+) {
+    dim3 grid(num_q_heads, batch);
+    int block = PAGED_THREADS;
+
+    paged_decode_attention_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)Q,
+        (const __nv_bfloat16*)K_cache,
+        (const __nv_bfloat16*)V_cache,
+        (__nv_bfloat16*)O,
+        block_tables, context_lens,
+        num_q_heads, num_kv_heads,
+        head_dim, max_blocks_per_seq,
+        scale
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+}

csrc/common.cuh

@@ -48,3 +48,17 @@ __device__ __forceinline__ float block_reduce_max(float val) {
     if (warp_id == 0) val = warp_reduce_max(val);
     return val;
 }
+
+// --- Launch error checking (debug builds only) ---
+#ifdef NDEBUG
+#define CUDA_CHECK_LAST_ERROR() ((void)0)
+#else
+#include <cstdio>
+#define CUDA_CHECK_LAST_ERROR() do { \
+    cudaError_t err = cudaGetLastError(); \
+    if (err != cudaSuccess) { \
+        fprintf(stderr, "CUDA kernel launch error at %s:%d: %s\n", \
+                __FILE__, __LINE__, cudaGetErrorString(err)); \
+    } \
+} while(0)
+#endif

csrc/embedding/embedding.cu

@@ -1,4 +1,5 @@
 #include <cuda_bf16.h>
+#include "../common.cuh"
 
 // Embedding lookup: out[seq_idx] = table[token_ids[seq_idx]]
 // Grid: num_tokens, Block: handles hidden_size elements per token.
@@ -7,10 +8,12 @@ __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 hidden_size,
+    int vocab_size
 ) {
     int token_idx = blockIdx.x;
     int tid = token_ids[token_idx];
+    if (tid < 0 || tid >= vocab_size) return;
     const float* row = table + tid * hidden_size;
     float* dst = out + token_idx * hidden_size;
 
@@ -23,10 +26,12 @@ __global__ void embedding_bf16(
     const __nv_bfloat16* __restrict__ table,
     const int* __restrict__ token_ids,
     __nv_bfloat16* __restrict__ out,
-    int hidden_size
+    int hidden_size,
+    int vocab_size
 ) {
     int token_idx = blockIdx.x;
     int tid = token_ids[token_idx];
+    if (tid < 0 || tid >= vocab_size) return;
     const __nv_bfloat16* row = table + tid * hidden_size;
     __nv_bfloat16* dst = out + token_idx * hidden_size;
 
@@ -38,18 +43,20 @@ __global__ void embedding_bf16(
 extern "C" {
 
 void launch_embedding_f32(const void* table, const void* token_ids, void* out,
-                          int num_tokens, int hidden_size, void* stream) {
+                          int num_tokens, int hidden_size, int vocab_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);
+        (const float*)table, (const int*)token_ids, (float*)out, hidden_size, vocab_size);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_embedding_bf16(const void* table, const void* token_ids, void* out,
-                           int num_tokens, int hidden_size, void* stream) {
+                           int num_tokens, int hidden_size, int vocab_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);
+        (__nv_bfloat16*)out, hidden_size, vocab_size);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/embedding/rope.cu

@@ -1,10 +1,11 @@
 #include <cuda_bf16.h>
 #include <math.h>
+#include "../common.cuh"
 
-// RoPE: Rotary Position Embedding
+// RoPE: Rotary Position Embedding, using the Qwen/Llama rotate_half layout.
-// For each pair (x[2i], x[2i+1]) at position `pos`:
+// For each dimension i in the first half at position `pos`:
-//   y[2i]   = x[2i] * cos - x[2i+1] * sin
+//   y[i]            = x[i] * cos - x[i + half_dim] * sin
-//   y[2i+1] = x[2i] * sin + x[2i+1] * cos
+//   y[i + half_dim] = x[i + half_dim] * cos + x[i] * sin
 // where cos/sin come from precomputed cos_cache/sin_cache.
 //
 // cos_cache[pos][i] = cos(pos * freq[i])
@@ -35,11 +36,11 @@ __global__ void rope_f32(
     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 x0 = x[base + pair_idx];
-    float x1 = x[base + 2 * pair_idx + 1];
+    float x1 = x[base + pair_idx + half_dim];
 
-    x[base + 2 * pair_idx]     = x0 * cos_val - x1 * sin_val;
+    x[base + pair_idx]            = x0 * cos_val - x1 * sin_val;
-    x[base + 2 * pair_idx + 1] = x0 * sin_val + x1 * cos_val;
+    x[base + pair_idx + half_dim] = x1 * cos_val + x0 * sin_val;
 }
 
 __global__ void rope_bf16(
@@ -61,11 +62,11 @@ __global__ void rope_bf16(
     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 x0 = __bfloat162float(x[base + pair_idx]);
-    float x1 = __bfloat162float(x[base + 2 * pair_idx + 1]);
+    float x1 = __bfloat162float(x[base + pair_idx + half_dim]);
 
-    x[base + 2 * pair_idx]     = __float2bfloat16(x0 * cos_val - x1 * sin_val);
+    x[base + pair_idx]            = __float2bfloat16(x0 * cos_val - x1 * sin_val);
-    x[base + 2 * pair_idx + 1] = __float2bfloat16(x0 * sin_val + x1 * cos_val);
+    x[base + pair_idx + half_dim] = __float2bfloat16(x1 * cos_val + x0 * sin_val);
 }
 
 // Precompute cos/sin cache on GPU
@@ -94,6 +95,7 @@ void launch_rope_f32(void* x, const void* cos_cache, const void* sin_cache,
     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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache,
@@ -104,6 +106,7 @@ void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache,
     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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_compute_rope_cache(void* cos_cache, void* sin_cache,
@@ -111,6 +114,7 @@ void launch_compute_rope_cache(void* cos_cache, void* sin_cache,
                                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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/embedding/transpose.cu

@@ -1,4 +1,5 @@
 #include <cuda_bf16.h>
+#include "../common.cuh"
 
 // 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).
@@ -111,6 +112,55 @@ __global__ void repeat_kv_bf16(
     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,
@@ -120,6 +170,7 @@ void launch_reshape_heads_bf16(const void* in, void* out,
     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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_merge_heads_bf16(const void* in, void* out,
@@ -129,6 +180,7 @@ void launch_merge_heads_bf16(const void* in, void* out,
     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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_transpose_hsd_to_shd_bf16(const void* in, void* out,
@@ -138,6 +190,7 @@ void launch_transpose_hsd_to_shd_bf16(const void* in, void* out,
     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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_transpose_shd_to_hsd_bf16(const void* in, void* out,
@@ -147,6 +200,7 @@ void launch_transpose_shd_to_hsd_bf16(const void* in, void* out,
     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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_repeat_kv_bf16(const void* in, void* out,
@@ -156,6 +210,33 @@ void launch_repeat_kv_bf16(const void* in, void* out,
     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);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+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);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/gemm/gemv.cu

@@ -0,0 +1,105 @@
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include "../common.cuh"
+
+// 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
+    );
+    CUDA_CHECK_LAST_ERROR();
+
+    // 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
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+} // extern "C"

csrc/gemm/naive.cu

@@ -1,4 +1,5 @@
 #include <cuda_bf16.h>
+#include "../common.cuh"
 
 // Naive GEMM: each thread computes one element of C.
 // C[i][j] = sum_k A[i][k] * B[k][j]
@@ -46,6 +47,7 @@ void launch_gemm_naive_bf16(
     gemm_naive_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)A, (const __nv_bfloat16*)B, (__nv_bfloat16*)C, M, N, K
     );
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_gemm_naive_f32(
@@ -57,6 +59,7 @@ void launch_gemm_naive_f32(
     gemm_naive_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const float*)A, (const float*)B, (float*)C, M, N, K
     );
+    CUDA_CHECK_LAST_ERROR();
 }
 
 } // extern "C"

csrc/gemm/tiled.cu

@@ -1,4 +1,5 @@
 #include <cuda_bf16.h>
+#include "../common.cuh"
 
 // Tiled GEMM using shared memory.
 // Each thread block loads TILE_SIZE x TILE_SIZE tiles of A and B
@@ -100,6 +101,7 @@ void launch_gemm_tiled_f32(
     gemm_tiled_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const float*)A, (const float*)B, (float*)C, M, N, K
     );
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_gemm_tiled_bf16(
@@ -111,6 +113,7 @@ void launch_gemm_tiled_bf16(
     gemm_tiled_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)A, (const __nv_bfloat16*)B, (__nv_bfloat16*)C, M, N, K
     );
+    CUDA_CHECK_LAST_ERROR();
 }
 
 } // extern "C"

csrc/normalization/layernorm.cu

@@ -14,27 +14,34 @@ __global__ void layernorm_f32(
     const float* x_row = x + row * hidden_size;
     float* out_row = out + row * hidden_size;
 
-    // Welford online: compute mean and variance in one pass
+    // Pass 1: compute mean
     float local_sum = 0.0f;
-    float local_sum_sq = 0.0f;
     for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
-        float v = x_row[i];
+        local_sum += x_row[i];
-        local_sum += v;
-        local_sum_sq += v * v;
     }
     local_sum = block_reduce_sum(local_sum);
-    local_sum_sq = block_reduce_sum(local_sum_sq);
 
     __shared__ float s_mean, s_inv_std;
     if (threadIdx.x == 0) {
-        float mean = local_sum / hidden_size;
+        s_mean = local_sum / hidden_size;
-        float var = local_sum_sq / hidden_size - mean * mean;
-        s_mean = mean;
-        s_inv_std = rsqrtf(var + eps);
     }
     __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];
@@ -52,26 +59,34 @@ __global__ void layernorm_bf16(
     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;
-    float local_sum_sq = 0.0f;
     for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
-        float v = __bfloat162float(x_row[i]);
+        local_sum += __bfloat162float(x_row[i]);
-        local_sum += v;
-        local_sum_sq += v * v;
     }
     local_sum = block_reduce_sum(local_sum);
-    local_sum_sq = block_reduce_sum(local_sum_sq);
 
     __shared__ float s_mean, s_inv_std;
     if (threadIdx.x == 0) {
-        float mean = local_sum / hidden_size;
+        s_mean = local_sum / hidden_size;
-        float var = local_sum_sq / hidden_size - mean * mean;
-        s_mean = mean;
-        s_inv_std = rsqrtf(var + eps);
     }
     __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]);
@@ -86,17 +101,21 @@ 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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/normalization/rmsnorm.cu

@@ -63,21 +63,78 @@ __global__ void rmsnorm_bf16(
     }
 }
 
+// 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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 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);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+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);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/reduce/softmax.cu

@@ -94,6 +94,7 @@ void launch_softmax_f32(const void* x, void* out, int rows, int cols, void* stre
     if (block < 32) block = 32;
     softmax_f32<<<rows, block, 0, (cudaStream_t)stream>>>(
         (const float*)x, (float*)out, cols);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_softmax_bf16(const void* x, void* out, int rows, int cols, void* stream) {
@@ -101,6 +102,7 @@ void launch_softmax_bf16(const void* x, void* out, int rows, int cols, void* str
     if (block < 32) block = 32;
     softmax_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)x, (__nv_bfloat16*)out, cols);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

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 都成功加载
+- [ ] 加载 Qwen3-8B sharded 权重，验证所有 tensor 都成功加载
-- [ ] 性能: 测量 7B 模型权重加载时间 (mmap → GPU 全流程)
+- [ ] 性能: 测量 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) |
+| Hidden Size | 768 | 4,096 (8B) |
-| Layers | 12 | 28 |
+| Layers | 12 | 36 |
 | Tied Embeddings | Yes | No |
 
 ### 需要新增/修改的组件
@@ -948,16 +948,16 @@ pub struct Qwen3DecoderLayer {
 ### 显存预算 (BF16, 单卡 5090 32GB)
 
 ```
-模型权重:  7B × 2B = ~14 GB
+模型权重:  8B × 2B = ~16 GB
-KV cache:  28 layers × 2(KV) × 8 heads × 4096 tokens × 128 dim × 2B ≈ 4.5 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
 
-    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 上限)
+   - 5090 SM120: shared memory per SM = 100 KB (需实测确认)
-   - 需要同时存 Q_tile, K_tile, V_tile, S_tile
+   - 需同时存 Q_tile, K_tile, V_tile, S_tile
-   - 典型值: Br=Bc=128 for D=128, BF16
+   - 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 的结合**:
+4. **GQA 支持**:
-   - Flash Attention 的 K/V tile 遍历逻辑需要适配间接寻址
+   - K/V heads 数量 < Q heads 时，kernel 中做 `kv_head = q_head / num_groups` 索引
-   - 每个 tile 查 block_table 得到物理地址
+   - 不需要 repeat_kv 操作，直接在 kernel 内部解决
-   - 这是 "Flash-Decoding" / "FlashInfer" 的核心
+
+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/10-qwen3.md

@@ -1,12 +1,12 @@
-# Phase 10: Qwen3-7B Support — Design Document (Milestone ②)
+# Phase 10: Qwen3-8B Support — Design Document (Milestone ②)
 
 ## Goal
 
-扩展模型定义支持 Qwen3-7B 架构，验证输出正确性。与 GPT-2 的关键差异：RMSNorm、RoPE、GQA、SwiGLU、不共享 embedding。
+扩展模型定义支持 Qwen3-8B 架构，验证输出正确性。与 GPT-2 的关键差异：RMSNorm、RoPE、GQA、SwiGLU、不共享 embedding。
 
 ## 架构差异 (GPT-2 → Qwen3)
 
-| 特性 | GPT-2 | Qwen3-7B |
+| 特性 | GPT-2 | Qwen3-8B |
 |------|-------|----------|
 | Norm | LayerNorm(gamma, beta) | RMSNorm(gamma only) |
 | Position | Learned absolute (wpe) | RoPE (no params) |
@@ -15,8 +15,8 @@
 | FFN | 2 Linear (fc, proj) + GELU | 3 Linear (gate, up, down) + SwiGLU |
 | Weight layout | [in, out] (Conv1D style) | [out, in] (standard Linear) |
 | Tied embeddings | Yes | No (separate lm_head) |
-| hidden_size | 768 | 3584 |
+| hidden_size | 768 | 4096 |
-| num_layers | 12 | 28 |
+| num_layers | 12 | 36 |
 | head_dim | 64 | 128 |
 
 ## Weight Names (HuggingFace)
@@ -67,17 +67,17 @@ out  = down_proj(out)    # [S, 18944] @ [18944, 3584]^T → [S, 3584]
 ## 显存预算 (BF16, 单卡 5090)
 
 ```
-权重: 7B × 2B = ~14 GB (BF16)
+权重: 8B × 2B = ~16 GB (BF16)
-        7B × 4B = ~28 GB (FP32) — 不够! 必须用 BF16
+        8B × 4B = ~32 GB (FP32) — 不够! 必须用 BF16
 KV cache (S=256, B=1): ~0.1 GB
-总计: ~14 GB (BF16), 单卡可运行
+总计: ~16 GB (BF16), 单卡可运行
 ```
 
-**关键**: Qwen3-7B 必须用 BF16 才能在单张 5090 (32GB) 上运行。当前 GPT-2 用 FP32，需要支持 BF16 forward pass。
+**关键**: Qwen3-8B 必须用 BF16 才能在单张 5090 (32GB) 上运行。当前 GPT-2 用 FP32，需要支持 BF16 forward pass。
 
 ## Implementation Plan
 
-1. 下载 Qwen3-7B 模型 (BF16, ~14GB)
+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)

docs/11-paged-attention.md

@@ -1,8 +1,10 @@
-# Phase 11: Paged Attention + KV Cache Manager — Design Document
+# 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，使用 block-based paging 管理显存。消除每步 decode 的 CPU round-trip（当前 KV cache 最大性能瓶颈之一）。
+将 KV cache 从 CPU Vec 迁移到 GPU，消除每步 decode 的 CPU round-trip（当前 KV cache 最大性能瓶颈之一）。
 
 ## 当前问题
 

docs/12-continuous-batching.md

@@ -150,4 +150,8 @@ HTTP Handler                              Engine Thread
 
 ## 当前状态
 
-**未实现**。当前是 FIFO 串行，一次只处理一个请求。本文档是实现的设计规格。
+**已实现: 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/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/16-llama-cpp-comparison.md

@@ -0,0 +1,201 @@
+# Phase 16: llama.cpp Comparison Baseline
+
+> **Goal.** Replace HF transformers with **llama.cpp** as the standing
+> performance baseline, and add a standard quality (response correctness)
+> benchmark suite (AIME 2025, GSM8K). Provide a one-click entrypoint that runs
+> both systems under identical workloads and emits a side-by-side report.
+
+## Motivation
+
+xserv has cleared 140% of HF transformers throughput on Qwen3-8B (Phase 15).
+HF is no longer a useful performance bar — it's a *correctness* baseline.
+
+**llama.cpp** is the right next bar because:
+- It's a serious C++/CUDA inference engine with active optimization
+- Same OpenAI-compatible API → black-box, fair comparison
+- Same GGUF↔safetensors weight source (we convert BF16, no quantization shortcuts)
+- Used widely as a reference point in the community
+
+We also need **quality benchmarks** so that performance improvements don't
+silently regress model quality (numerical precision, sampling, prompt
+formatting). AIME and GSM8K are the cheapest credible signals.
+
+## Architecture
+
+```
+xserv/
+├── third_party/llama.cpp/         # cloned by setup-llama-cpp.sh
+│   └── build/bin/llama-server     # CUDA build (SM120)
+├── tools/
+│   ├── setup-llama-cpp.sh         # clone + cmake build (idempotent)
+│   ├── convert-to-gguf.sh         # safetensors → BF16 GGUF (same weights)
+│   ├── sync-and-build.sh          # extended with `bench` subcommand
+│   └── bench/                     # Python benchmark driver
+│       ├── runner.py              # entrypoint
+│       ├── servers.py             # subprocess lifecycle (start/stop both)
+│       ├── client.py              # OpenAI streaming client + TTFT/TPOT
+│       ├── speed.py               # speed suite
+│       ├── quality.py             # quality suite
+│       ├── tasks/{aime,gsm8k}.py  # dataset loaders + scorers
+│       ├── report.py              # markdown + json output
+│       └── requirements.txt       # httpx, datasets
+└── bench-out/                     # report artifacts (gitignored)
+    ├── comparison-<stamp>.md
+    ├── comparison-<stamp>.json
+    └── logs/{xserv,llama_cpp}.log
+```
+
+Both systems are treated as **black-box HTTP servers** speaking the OpenAI
+streaming chat API. No in-process integration, no shared Python bindings. This
+keeps the comparison fair (same protocol, same prompt-template path) and
+isolates the test harness from internal API churn on either side.
+
+## Workflow
+
+The GPU host (dash5) has **no outbound network and no rsync**, so anything from
+the internet is fetched locally and shipped over via tar-over-ssh.
+
+```
+local repo (has network)              dash5 (GPU host, no network)
+────────────────────────              ────────────────────────────
+# one-time, on a networked machine:
+python3 -m tools.bench.fetch_datasets  →  tools/bench/data/{aime2025,gsm8k}.json
+git submodule update --init …          →  third_party/llama.cpp source
+
+tools/sync-and-build.sh bench   →  tar project   (excl. target, third_party, bench-out)
+                                →  tar llama.cpp source (excl. build, .git)
+                                →  setup-llama-cpp.sh   (build-only; no-op if built)
+                                →  convert-to-gguf.sh   (no-op if .gguf exists)
+                                →  cargo build --release
+                                →  python3 -m tools.bench.runner ...
+                                →  bench-out/comparison-<stamp>.md
+tools/sync-and-build.sh fetch-bench-out  ←  tar bench-out back
+```
+
+Behind a flaky proxy, fetch datasets through the HF mirror:
+`HF_ENDPOINT=https://hf-mirror.com python3 -m tools.bench.fetch_datasets`.
+
+`tools/__init__.py` exists so `python3 -m tools.bench.runner` resolves our
+package: some site-packages (e.g. nvfuser) ship a regular top-level `tools`
+package that would otherwise shadow a namespace `tools`.
+
+## What gets measured
+
+### Speed (TTFT / TPOT / throughput)
+
+- **Single-stream**, three prompt lengths (short / medium / long), `cfg.speed_prompts` repeats each
+  - `TTFT p50/p95`, `TPOT p50/p95`, per-request throughput
+- **Concurrent**, fixed medium prompt, sweep `concurrency ∈ {1, 2, 4, 8}`
+  - Aggregate `tok/s`, `TTFT p95`, error count
+- Both at `temperature=0`, `max_tokens=128` by default.
+
+### Quality (response correctness)
+
+| Task | N | Source | Scoring | Why |
+|---|---|---|---|---|
+| AIME 2025 | 30 | `MathArena/aime_2025`, fallback `yentinglin/aime_2025` (HF) | exact-match boxed integer (0..999) | reasoning + math, hard signal |
+| GSM8K | 1319 | `openai/gsm8k` (HF), `test` split | exact-match `\boxed{n}` or last number | broad sanity, decimals allowed |
+
+Same `temperature=0` sampling across both systems. Max tokens: 16384 for AIME
+(reasoning long), 2048 for GSM8K. Subsample with `--quality-limit N` for smoke.
+
+**Generation mode must match.** xserv's prompt builder hardcodes Qwen3 thinking
+OFF (it appends an empty `<think></think>` block). llama-server applies the
+GGUF's Qwen3 jinja template, which has thinking ON by default. The driver
+therefore sends `chat_template_kwargs={"enable_thinking": false}` to llama.cpp
+so both engines run the model in the same mode. Pass `--enable-thinking` to
+compare in thinking mode instead (xserv would need a matching change first).
+
+### Report
+
+`bench-out/comparison-<stamp>.md` contains:
+- Environment (GPU, driver, xserv commit, python)
+- Speed table per scenario (xserv | llama.cpp | xserv÷llama.cpp speedup)
+- Quality table per task (n, correct, accuracy, mean tokens, TTFT, TPOT, wall)
+
+A sibling `.json` holds all per-request raw rows and per-problem case detail
+(prediction, gold, response preview) so we can diff regressions in CI later.
+
+## Running it
+
+**One-time prerequisites (on a networked machine):**
+```bash
+git submodule update --init third_party/llama.cpp     # pinned to b9371
+HF_ENDPOINT=https://hf-mirror.com python3 -m tools.bench.fetch_datasets
+```
+
+**Full sweep on dash5 (recommended):**
+```bash
+./tools/sync-and-build.sh bench -- --max-seq-len 8192 --quality-limit 50
+./tools/sync-and-build.sh fetch-bench-out
+open bench-out/comparison-*.md
+```
+
+**Speed-only smoke (fast):**
+```bash
+./tools/sync-and-build.sh bench -- --suite speed --speed-prompts 2
+```
+
+**Quality smoke with 5 problems each:**
+```bash
+./tools/sync-and-build.sh bench -- --suite quality --quality-limit 5
+```
+
+**On a host that already has both servers running** (e.g. local dev with two
+shells open):
+```bash
+python3 -m tools.bench.runner \
+    --xserv-base-url http://127.0.0.1:8080 \
+    --llama-base-url http://127.0.0.1:8081 \
+    --suite all
+```
+
+## Design choices
+
+1. **Black-box HTTP, not FFI.** Both engines bind the same OpenAI surface and
+   real serving traffic uses HTTP. Anything that doesn't show up over the wire
+   doesn't matter for serving.
+2. **Same BF16 weights.** We convert the same safetensors with llama.cpp's
+   `convert_hf_to_gguf.py --outtype bf16`. No quantization at this stage; if we
+   want a quant comparison later we'll add a separate column, not replace this
+   one.
+3. **Streaming everywhere.** TTFT and TPOT only make sense with streaming. We
+   ask both servers for `stream=true` with `include_usage` so we can read
+   server-reported token counts when available.
+4. **Idempotent setup.** `setup-llama-cpp.sh` and `convert-to-gguf.sh` are
+   safe to re-run — they no-op when the build / file already exists. The
+   `bench` subcommand wires them so the first run does a full setup and
+   subsequent runs are fast.
+5. **Subprocess lifecycle owned by the driver.** We spawn each server in its
+   own process group and SIGTERM the group on exit so half-dead llama-server
+   children don't survive. If the user is already running a server somewhere,
+   pass `--xserv-base-url` / `--llama-base-url` to skip launch.
+6. **One server at a time.** The driver starts a system, runs every suite
+   against it, stops it, then moves to the next. Two BF16 8B models (~16GB each)
+   do not co-reside on a single 32GB GPU, and a resident idle engine would
+   distort the other's latency/throughput. This serialization is why the report
+   is assembled from per-system passes rather than a single interleaved run.
+
+## Known constraints / findings
+
+- **xserv OOM at `--max-seq-len 8192` — fixed.** xserv used to pre-allocate its
+  paged-KV pool (`total_blocks = blocks_per_seq · max_batch · 2`, ≈9GB at 8192)
+  on top of the 16GB weights, exceeding 32GB at startup (`paged_kv_cache.rs`
+  `alloc paged K pool: OutOfMemory`). Now the pool is sized to *available VRAM*
+  (`cudaMemGetInfo`) and overflow is swapped to pinned host memory (vLLM-style
+  preemption, `--swap-space-gb`). The 8192 comparison runs cleanly with 0 swap
+  events; swap is verified separately under a forced-small pool. The benchmark
+  surfaced the OOM — a good example of the baseline doing its job.
+- When the xserv engine thread dies, the API now returns a clean 503 (the
+  request handler uses a poison-tolerant lock instead of cascading
+  mutex-poison panics). The driver records any failure as a per-request error,
+  so a broken engine shows up as `errs=N` / `accuracy 0%` rather than a hung run.
+
+## Future extensions
+
+- Add quant runs (Q8_0, Q4_K_M) as separate "system" columns
+- Wire to GitHub Actions for nightly regression
+- Track results across commits to flag regressions (per-commit JSON in
+  `docs/benchmarks/history/`)
+- Add MMLU-Pro / HumanEval when budget allows
+- Long-context benchmark (8K, 32K prompts) to compare prefill scaling

docs/17-tensor-parallelism.md

@@ -0,0 +1,122 @@
+# Phase 17: Tensor Parallelism (TP)
+
+> 目标：在单机多卡上做 **张量并行**，把 Qwen3-8B 的权重、计算和 KV cache 按
+> head / 中间维切分到 TP 个 GPU 上，用 AllReduce 聚合，降低单卡显存压力并提升吞吐。
+> 先做 **TP=2 / 4（组内）**，跳过投机解码（原 Phase 16）。
+
+## 1. 硬件约束（dash5）
+
+- 8× RTX 5090（32GB，SM120），**无 NVLink**，纯 PCIe Gen5。
+- 拓扑：GPU 0–3 一组、4–7 一组，组内 `PHB`（同 host bridge，可 P2P），跨组 `NODE`。
+- **TP 必须在组内**（0–3 或 4–7），否则 AllReduce 走跨组 PCIe，延迟更高。
+- AllReduce 带宽受限于 PCIe（~单向 64GB/s），远低于 NVLink；通信会是 decode 的主要开销。
+
+## 2. 切分方案（Megatron-style）
+
+Qwen3-8B：`hidden=4096`、`num_heads=32`、`num_kv_heads=8`、`head_dim=128`、
+`intermediate=12288`、`layers=36`、`vocab=151936`。TP=2/4/8 都能整除 32/8/12288。
+
+每个 transformer block 的切分（设 world size = `T`，本 rank = `r`）：
+
+### Attention（column → row）
+| 权重 | 原 shape (已转置) | 切分 | 每 rank shape |
+|------|-------------------|------|---------------|
+| `q_proj_wt` | `[hidden, num_heads·head_dim]` | column（按 Q head） | `[hidden, (num_heads/T)·head_dim]` |
+| `k_proj_wt` | `[hidden, num_kv_heads·head_dim]` | column（按 KV head） | `[hidden, (num_kv_heads/T)·head_dim]` |
+| `v_proj_wt` | 同上 | column | 同上 |
+| `o_proj_wt` | `[num_heads·head_dim, hidden]` | **row** | `[(num_heads/T)·head_dim, hidden]` |
+
+- 每个 rank 只算自己的 `num_heads/T` 个 Q head 和对应的 `num_kv_heads/T` 个 KV head；
+  GQA 的 `n_rep = num_heads/num_kv_heads = 4` 在每个 rank 内保持不变。
+- `q_norm`/`k_norm`（`[head_dim]`）逐 head 应用，**复制**到每个 rank。
+- RoPE 逐 head、按 position 应用，每个 rank 独立做。
+- **KV cache 也切分**：每个 rank 的 paged KV 只存自己的 `num_kv_heads/T` 个 head
+  → 每卡 KV 显存降为 1/T（TP 的一大收益）。
+- `attn = merge_heads(...) @ o_proj_wt` 得到**部分** `[T_tok, hidden]` → **AllReduce(sum)** → 完整。
+
+### MLP / SwiGLU（column → row）
+| 权重 | 原 shape | 切分 | 每 rank shape |
+|------|----------|------|---------------|
+| `gate_proj_wt` | `[hidden, intermediate]` | column | `[hidden, intermediate/T]` |
+| `up_proj_wt` | 同上 | column | 同上 |
+| `down_proj_wt` | `[intermediate, hidden]` | **row** | `[intermediate/T, hidden]` |
+
+- `silu(gate)*up` 在切分后的 `[T_tok, intermediate/T]` 上逐元素做，无需通信。
+- `down = (...) @ down_proj_wt` 得到部分 `[T_tok, hidden]` → **AllReduce(sum)** → 完整。
+
+### 复制（不切分）
+- 所有 RMSNorm 权重（`input_norm`/`post_norm`/最终 `norm`）：每个 rank 在 AllReduce 后
+  拿到完整 hidden，本地用复制的权重归一化。
+- **第一版**：`embed_tokens` 和 `lm_head` 复制（各 ~1.2GB）。
+  后续优化：vocab-parallel embedding（local lookup + AllReduce）、column-parallel lm_head + AllGather。
+
+### 通信点
+每层 **2 次 AllReduce**（o_proj 后、down_proj 后）→ 每生成 1 token 共 `2·36 = 72` 次。
+decode 时每次 AllReduce 张量是 `[batch, 4096]` BF16（单 token batch=1 时 8KB），**延迟主导**。
+
+## 3. 进程 / 线程模型
+
+**单进程、多线程**：每个 TP rank 一个 OS 线程，线程启动时 `cudaSetDevice(rank_device)` 并绑定。
+
+选择理由：
+- xserv 的 caching allocator 是 `thread_local`，每线程独立池 → 天然契合「一线程一卡一池」。
+- CUDA context 隐式按 device/thread 管理；线程内只 set 一次 device、不再切换即可。
+- HTTP server / 调度器仍在主线程（rank 0 协调），无需多进程 IPC，改动最小。
+- 单机 8 卡足够；多进程（torchrun 式）留待真正跨节点时再说。
+
+执行流：主线程调度器准备一个 step 的输入（tokens/positions/slots），广播给 `T` 个 rank 线程；
+每个 rank 线程跑自己的分片 forward（含层内 AllReduce），rank 0 拿到完整 logits 后采样。
+用 barrier / channel 同步每个 step。
+
+## 4. 通信库：NCCL
+
+用 **NCCL**（dash5 已装：`/usr/lib/x86_64-linux-gnu/libnccl.so.2`，`/usr/include/nccl.h`）。
+
+- 新建 crate `xserv-distributed`：NCCL FFI（`ncclGetUniqueId`、`ncclCommInitRank`、
+  `ncclAllReduce`、`ncclGroupStart/End`）+ `TpContext`（rank/world/comm/stream）+
+  `all_reduce_sum(&mut GpuBuffer)` 原语。
+- NCCL 多线程模式：主线程生成 `ncclUniqueId`，各 rank 线程用 `ncclCommInitRank(comm, world, id, rank)`
+  初始化（需 `ncclGroupStart/End` 包裹并发 init）。
+- AllReduce 用 BF16（`ncclBfloat16`）+ `ncclSum`，在每个 rank 自己的 stream 上。
+
+> **决策点**：collective 用 NCCL 还是自己手写 P2P ring/tree AllReduce？
+> 本项目是「从零构建」，但 collective 属于基础设施（类比我们也用 cuBLAS 作为可切换后端）。
+> 推荐先用 NCCL 把 TP 跑通、拿到正确性与加速比，后续可选做手写 AllReduce 作为学习项。
+
+## 5. 权重分片加载
+
+每个 rank 只加载/保留自己的分片，省显存：
+- column-parallel 权重：按输出维切片取本 rank 的 `[*, dim/T]` 段。
+- row-parallel 权重：按输入维切片取本 rank 的 `[dim/T, *]` 段。
+- 复制权重（norm/embed/lm_head）：每个 rank 各留一份。
+
+实现：`loader` 读 safetensors（mmap）时按 rank 只搬运需要的切片到该 rank 的 GPU；
+`Qwen3::from_weights_tp(config, weights, rank, world)` 在转置/切分时按 rank 取段。
+
+## 6. 实现步骤（逐步可验证）
+
+1. **P17.1 — `xserv-distributed` 基础**：NCCL FFI + `TpContext` + `all_reduce_sum`。
+   验收：2 卡各放一个已知向量，AllReduce 后两卡结果都等于和。
+2. **P17.2 — 分片权重加载**：`from_weights_tp`，每 rank 只持有自己的分片。
+   验收：各 rank 权重 shape 正确、显存占用约为 1/T（+ 复制项）。
+3. **P17.3 — TP forward**：rank 内 attention/MLP + 层内 AllReduce。
+   验收：**TP=2 的 logits 与 TP=1 在 BF16 容差内一致**（top-1 一致，top-5 重合）。
+4. **P17.4 — 接入 engine/server**：`--tp N`，多线程 rank workers + rank0 调度。
+   验收：`--tp 2` 端到端可服务；用现有 llama.cpp bench 跑正确性；
+   测 TP=2 vs TP=1 的吞吐 / TTFT / 每卡显存。
+
+## 7. 预期与风险
+
+- **显存**：每卡权重 + KV 降到约 1/T（embed/lm_head 暂复制）。TP=2 时单卡 ~8GB 权重 + 更大 KV/并发空间。
+- **吞吐**：PCIe AllReduce 延迟会吃掉部分收益；decode 是延迟敏感的，72 次小 AllReduce/token
+  可能让 TP=2 的单流 tok/s **不一定线性提升**，但能跑更大 batch / 更长 context。先测实测数。
+- **风险**：NCCL 多线程初始化的同步、每 rank stream 与现有 kernel stream 的协调、
+  KV cache 按 rank 切 head 后 paged kernel 的 head 维参数要用 per-rank 值。
+- 正确性优先：先 TP=1 等价（logits 对齐），再谈性能。
+
+## 8. 不在本阶段范围
+
+- 跨组 TP=8、Pipeline Parallelism、多节点。
+- vocab-parallel embedding / lm_head（先复制）。
+- 手写 AllReduce（NCCL 跑通后可选）。
+- 与 CUDA Graph decode 的结合（先走非 graph 路径）。

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/llama-cpp-comparison.md

@@ -0,0 +1,89 @@
+# Benchmark: xserv vs llama.cpp (Qwen3-8B)
+
+**What this adds.** A standing baseline that compares xserv against **llama.cpp**
+on both **response quality (correctness)** and **performance (TTFT / TPOT /
+throughput)**, using the same model weights and standard public datasets. This
+replaces HF transformers as our reference point — xserv already beat HF, so it
+is no longer a useful performance bar.
+
+- **Baseline engine**: llama.cpp, vendored as a submodule pinned to `b9371`,
+  built with CUDA for SM120 (RTX 5090).
+- **Same weights**: the Qwen3-8B safetensors are converted to a **BF16 GGUF**
+  (`convert_hf_to_gguf.py --outtype bf16`) — no quantization, so the comparison
+  is apples-to-apples.
+- **Standard quality datasets**: **AIME 2025** (30 competition-math problems,
+  exact-match boxed integer) and **GSM8K** (grade-school math, exact-match).
+- **Black-box HTTP**: both engines are driven through the OpenAI-compatible
+  streaming API; the driver measures TTFT/TPOT/throughput and scores answers.
+
+See `docs/16-llama-cpp-comparison.md` for the design and `tools/bench/` for the
+driver. One-click: `tools/sync-and-build.sh bench`.
+
+## How it runs
+
+The GPU host (dash5) has no outbound network, so datasets are fetched locally
+(`tools/bench/fetch_datasets.py`) into JSON and the llama.cpp source is shipped
+over with the project; everything builds and runs on the GPU host. The driver
+runs **one engine at a time** (two BF16 8B models do not co-reside on a 32GB
+GPU, and a resident idle engine would distort the other's numbers).
+
+Generation mode is matched: xserv hardcodes Qwen3 **thinking off**, so the
+driver sends `chat_template_kwargs={enable_thinking:false}` to llama.cpp.
+
+## Results (RTX 5090, BF16, greedy, 8192 ctx, max_batch 4)
+
+### Performance — llama.cpp is the stronger baseline
+
+| scenario | metric | xserv | llama.cpp | xserv ÷ llama.cpp |
+|---|---|---|---|---|
+| single / medium | TTFT p50 (ms) | 28.0 | 17.7 | 0.63× |
+| single / medium | TPOT p50 (ms/tok) | 17.5 | 10.4 | 0.60× |
+| single / medium | throughput (tok/s) | 56.6 | 95.1 | 0.60× |
+| concurrent-4 | throughput (tok/s) | 135.2 | 317.1 | 0.43× |
+| concurrent-8 | throughput (tok/s) | 135.5 | 322.5 | 0.42× |
+
+xserv runs at **~0.42–0.60×** llama.cpp. It saturates at `max_batch` (~135 tok/s)
+while llama.cpp keeps scaling under load (~322 tok/s). This is the honest new bar.
+The ratio is the same at 4096 and 8192 — TPOT is bandwidth-bound, not
+context-bound at these sizes.
+
+### Quality — parity, confirming xserv's numerical fidelity
+
+| task | n | xserv | llama.cpp |
+|---|---|---|---|
+| GSM8K | 50 | 98.0% (49/50) | 96.0% (48/50) |
+| AIME 2025 | 30 | 20.0% (6/30) | 20.0% (6/30) |
+
+With equal context the two engines land at identical AIME accuracy and
+within one problem on GSM8K. At 8192 both generate full-length solutions
+(mean ~3.4k / ~4.2k tokens), so neither is truncated. Two independent engines
+agreeing at ~20% confirms that's genuine Qwen3-8B (thinking-off) capability and
+that xserv is numerically faithful. Response prefixes are byte-identical (same
+prompt templating); the only run-to-run wobble is greedy-decode divergence /
+nondeterminism on long (~3k-token) sequences (see finding 3).
+
+## Findings the benchmark surfaced
+
+1. **Context must be provisioned per-request, not total.** A first run showed
+   xserv 20.0% vs llama.cpp 3.3% on AIME — an artifact: llama.cpp divides total
+   `-c` across `--parallel` slots, so `-c 4096 --parallel 4` gave each request
+   only **1024 tokens**, truncating long AIME solutions before the boxed answer
+   (capped at ~940 generated tokens). GSM8K (~280 tokens) was unaffected, which
+   is how we caught it. Fixed: per-slot context = `max_seq_len` (total
+   `-c = max_seq_len × parallel`). After the fix, AIME is at parity (above).
+2. **xserv OOM'd at `--max-seq-len 8192` — now fixed.** xserv used to eagerly
+   pre-allocate its paged-KV pool (`blocks_per_seq × max_batch × 2`, ~9GB at
+   8192) on top of the 16GB weights, exceeding 32GB at startup. Fixed by sizing
+   the pool to *available VRAM* (`cudaMemGetInfo`) instead of worst-case demand,
+   plus vLLM-style **swap to pinned host memory**: when running sequences grow
+   past the GPU pool, the newest are evicted to host and swapped back when blocks
+   free up (`--swap-space-gb`, default 8). The results above run at 8192 with **0
+   swap events** — the VRAM-sized pool alone covers this load; swap is the
+   overload safety net (verified lossless under a forced-small pool).
+3. **xserv decode is not run-to-run deterministic.** The same greedy (temp 0)
+   AIME config produced 6/30 / 7/30 / 6/30 across runs — non-deterministic CUDA
+   reductions flip an argmax over long (~3k-token) generations. Harmless for
+   serving, but it explains why long-sequence accuracy wobbles by a problem.
+
+Raw artifacts (per-request timings, per-problem prediction/gold) are written to
+`bench-out/` as `comparison-<stamp>.{md,json}` (gitignored).

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/tensor-parallelism.md

@@ -0,0 +1,73 @@
+# Benchmark: Tensor Parallelism (TP=1/2/4) — xserv vs llama.cpp
+
+**Setup.** Qwen3-8B BF16 on 8× RTX 5090 (PCIe Gen5, **no NVLink**; GPUs grouped
+0-3 / 4-7 by PHB). Both engines driven over the same OpenAI HTTP harness, same
+scorers, thinking-off, greedy (temp 0), `max_tokens` 2048. Datasets: **AIME
+2025** (30) + **GSM8K** (30). The two engines run **concurrently on disjoint
+groups** — xserv on GPU 0..N-1, llama.cpp (`--split-mode row`) on GPU 4..4+N-1
+(`tools/bench/run_tp_parallel.sh`).
+
+## Correctness — on par across engines and TP
+
+| TP | task | xserv | llama.cpp |
+|----|------|-------|-----------|
+| 1 | AIME 2025 | 16.7% (5/30) | 13.3% (4/30) |
+| 1 | GSM8K | 96.7% (29/30) | 96.7% (29/30) |
+| 2 | AIME 2025 | 13.3% (4/30) | 13.3% (4/30) |
+| 2 | GSM8K | 93.3% (28/30) | 96.7% (29/30) |
+| 4 | AIME 2025 | 16.7% (5/30) | 13.3% (4/30) |
+| 4 | GSM8K | 96.7% (29/30) | 96.7% (29/30) |
+
+Within ±1 problem everywhere — TP changes nothing about quality on either
+engine, and the two engines agree. (AIME is low for both: Qwen3-8B thinking-off,
+capped at 2048 tokens.)
+
+## Performance — TPOT (ms/token, lower is better)
+
+| TP | xserv AIME / GSM8K | llama.cpp AIME / GSM8K |
+|----|--------------------|------------------------|
+| 1 | 21.0 / 17.8 | **10.4 / 10.3** |
+| 2 | 17.2 / 13.9 | 19.0 / 18.9 |
+| 4 | **15.2 / 12.1** | 20.2 / 20.2 |
+
+**Opposite TP scaling, with a crossover:**
+
+- **xserv TP scales positively**: TPOT 21.0 → 17.2 → 15.2 ms (AIME),
+  17.8 → 13.9 → 12.1 ms (GSM8K) — TP=4 is ~1.4–1.5× faster than TP=1. GPU 0-3
+  all ~82% utilized. (Sublinear because of the 72 PCIe AllReduces/token.)
+- **llama.cpp row-split regresses**: TPOT 10.4 → 19.0 → 20.2 ms — TP=1 is its
+  best; TP=2/4 nearly double the latency. GPU 4-7 only ~24% utilized
+  (communication-bound). Row-split's per-layer cross-GPU traffic over PCIe
+  without NVLink dominates.
+- **Crossover**: llama.cpp is ~2× faster at TP=1, still ahead at TP=2, and xserv
+  is ~1.3× faster at TP=4 (15.2 vs 20.2 ms AIME). AIME-30 wall clock: xserv
+  1046 → 846 → 730 s (falling), llama.cpp 520 → 952 → 1012 s (rising).
+
+On a NVLink-less PCIe box, **xserv's TP is a genuine win and llama.cpp's
+tensor-split is counterproductive** — exactly what the topology predicts.
+
+### Clean same-path xserv scaling (bench-tp)
+
+The HTTP numbers above mix engines (xserv TP=1 uses the production continuous-
+batching engine; TP≥2 uses the serial TP coordinator). The single-stream,
+same-code-path scaling from `bench-tp` (greedy, 8 prompts × 64 tokens):
+
+| TP | xserv decode tok/s | speedup | TTFT |
+|----|--------------------|---------|------|
+| 1 | 58.5 | 1.00× | 18.0 ms |
+| 2 | 75.7 | 1.29× | 13.4 ms |
+| 4 | 86.1 | 1.47× | 11.5 ms |
+
+## Caveats
+
+- xserv **TP=1 uses the production `Engine`**, TP≥2 the serial `tp_engine`
+  coordinator — different per-token paths, so the HTTP TP=1→2 step has an engine
+  confound. The clean same-path scaling (bench-tp, above) confirms the trend.
+- xserv **TTFT is weaker** on long AIME prompts (~460–500 ms vs llama ~100–190 ms)
+  — prefill is a known optimization target.
+- llama.cpp uses `--split-mode row` (its tensor-parallel mode); the default
+  `layer` split only memory-splits, without parallel compute.
+- The TP HTTP server processes requests **serially** (sufficient for this serial
+  quality benchmark); continuous-batching TP is future work.
+
+Raw artifacts: `bench-out/tp{1,2,4}-{xserv,llama}/comparison-*.{md,json}`.

tools/bench/client.py

@@ -0,0 +1,158 @@
+"""HTTP client for OpenAI-compatible /v1/chat/completions.
+
+Records per-request: TTFT (time to first content token), TPOT (mean
+inter-token latency over the decode phase), and end-to-end throughput.
+
+We don't care about parsing exact OpenAI envelope semantics, just enough to
+get the deltas + finish_reason + token counts.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+import httpx
+
+
+@dataclass
+class StreamResult:
+    text: str = ""
+    completion_tokens: int = 0
+    prompt_tokens: int = 0
+    finish_reason: str | None = None
+    # Timings (seconds; -1 means not measured)
+    ttft_s: float = -1.0
+    e2e_s: float = -1.0
+    chunk_times: list[float] = field(default_factory=list)   # absolute monotonic times of content chunks
+    error: str | None = None
+
+    @property
+    def tpot_s(self) -> float:
+        """Mean inter-content-chunk latency after the first chunk (seconds/token)."""
+        if len(self.chunk_times) < 2:
+            return -1.0
+        deltas = [self.chunk_times[i] - self.chunk_times[i - 1] for i in range(1, len(self.chunk_times))]
+        return sum(deltas) / len(deltas)
+
+    @property
+    def throughput_tok_s(self) -> float:
+        if self.e2e_s <= 0 or self.completion_tokens <= 0:
+            return -1.0
+        return self.completion_tokens / self.e2e_s
+
+
+async def chat_stream(
+    client: httpx.AsyncClient,
+    base_url: str,
+    model: str,
+    messages: list[dict[str, str]],
+    *,
+    max_tokens: int,
+    temperature: float = 0.0,
+    api_key: str | None = None,
+    timeout: float = 1800.0,
+    extra_body: dict | None = None,
+) -> StreamResult:
+    payload: dict[str, Any] = {
+        "model": model,
+        "messages": messages,
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "stream": True,
+    }
+    # llama-server returns usage in the final stream chunk when this is set;
+    # xserv ignores unknown fields, so this is harmless there.
+    payload["stream_options"] = {"include_usage": True}
+    if extra_body:
+        payload.update(extra_body)
+
+    headers = {"Content-Type": "application/json"}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    url = base_url.rstrip("/") + "/v1/chat/completions"
+    res = StreamResult()
+    t_start = time.perf_counter()
+
+    try:
+        async with client.stream(
+            "POST", url, json=payload, headers=headers, timeout=timeout,
+        ) as resp:
+            if resp.status_code != 200:
+                body = await resp.aread()
+                res.error = f"HTTP {resp.status_code}: {body.decode(errors='replace')[:400]}"
+                res.e2e_s = time.perf_counter() - t_start
+                return res
+
+            async for line in resp.aiter_lines():
+                if not line or not line.startswith("data:"):
+                    continue
+                data = line[len("data:"):].strip()
+                if data == "[DONE]":
+                    break
+                try:
+                    chunk = json.loads(data)
+                except json.JSONDecodeError:
+                    continue
+
+                if "usage" in chunk and chunk["usage"]:
+                    usage = chunk["usage"]
+                    res.prompt_tokens = usage.get("prompt_tokens", res.prompt_tokens)
+                    res.completion_tokens = usage.get("completion_tokens", res.completion_tokens)
+
+                choices = chunk.get("choices") or []
+                if not choices:
+                    continue
+                choice = choices[0]
+                delta = choice.get("delta") or {}
+                content = delta.get("content")
+                if content:
+                    now = time.perf_counter()
+                    if res.ttft_s < 0:
+                        res.ttft_s = now - t_start
+                    res.chunk_times.append(now)
+                    res.text += content
+                if choice.get("finish_reason"):
+                    res.finish_reason = choice["finish_reason"]
+    except Exception as e:  # noqa: BLE001 — surface any failure to the report
+        res.error = f"{type(e).__name__}: {e}"
+
+    res.e2e_s = time.perf_counter() - t_start
+    # Fall back to chunk count when server doesn't report usage (xserv stream path).
+    if res.completion_tokens == 0:
+        res.completion_tokens = len(res.chunk_times)
+    return res
+
+
+async def chat_concurrent(
+    base_url: str,
+    model: str,
+    prompts: list[list[dict[str, str]]],
+    *,
+    max_tokens: int,
+    temperature: float = 0.0,
+    api_key: str | None = None,
+    timeout: float = 1800.0,
+    concurrency: int,
+    extra_body: dict | None = None,
+) -> tuple[list[StreamResult], float]:
+    """Fire `concurrency` requests in parallel waves. Returns per-request results
+    plus wall-clock elapsed time of the entire batch."""
+    sem = asyncio.Semaphore(concurrency)
+    limits = httpx.Limits(max_connections=concurrency * 2, max_keepalive_connections=concurrency)
+    async with httpx.AsyncClient(timeout=timeout, limits=limits) as client:
+        async def one(messages: list[dict[str, str]]) -> StreamResult:
+            async with sem:
+                return await chat_stream(
+                    client, base_url, model, messages,
+                    max_tokens=max_tokens, temperature=temperature,
+                    api_key=api_key, timeout=timeout, extra_body=extra_body,
+                )
+        t0 = time.perf_counter()
+        results = await asyncio.gather(*(one(p) for p in prompts))
+        wall = time.perf_counter() - t0
+    return results, wall

tools/bench/config.py

@@ -0,0 +1,57 @@
+"""Defaults + CLI argument shapes for the benchmark driver.
+
+All paths default to the dash5 layout (/opt/wjh/...) because that's where the
+GPU lives — see docs/16-llama-cpp-comparison.md.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+
+
+# Names used in reports and as logical keys throughout the driver.
+SYSTEM_XSERV = "xserv"
+SYSTEM_LLAMA_CPP = "llama.cpp"
+DEFAULT_SYSTEMS = (SYSTEM_XSERV, SYSTEM_LLAMA_CPP)
+
+
+@dataclass
+class SystemEndpoint:
+    """How to reach (or how to start) one of the systems under test."""
+
+    name: str
+    base_url: str                  # http://host:port  (OpenAI-compatible root, no /v1)
+    model_id: str                  # what to put in the request body's "model" field
+    api_key: str | None = None     # llama-server doesn't need one; xserv ignores it
+    # Extra fields merged into every request body for this system. Used to keep
+    # the two engines in the SAME generation mode — xserv hardcodes Qwen3
+    # thinking OFF (empty <think></think> in its prompt builder), so we disable
+    # thinking on llama-server via chat_template_kwargs to match. Both engines
+    # ignore unknown fields, so this is safe.
+    extra_body: dict | None = None
+    # Process supervision is optional — if base_url is already serving, we skip launch.
+    launch_cmd: list[str] | None = None
+    launch_env: dict[str, str] = field(default_factory=dict)
+    launch_cwd: str | None = None
+    health_path: str = "/health"
+    ready_timeout_s: float = 600.0   # cold loads of 8B BF16 take a while
+
+
+@dataclass
+class BenchConfig:
+    out_dir: str = "bench-out"
+    # Speed suite
+    speed_prompts: int = 8           # synthetic prompts per length bucket
+    speed_max_tokens: int = 128
+    speed_concurrency: tuple[int, ...] = (1, 2, 4, 8)
+    # Quality suite
+    quality_max_tokens_aime: int = 16384
+    quality_max_tokens_gsm8k: int = 2048
+    quality_limit: int | None = None   # subsample for smoke tests; None = all
+    quality_temperature: float = 0.0
+    request_timeout_s: float = 1800.0
+
+
+def env_default(key: str, fallback: str) -> str:
+    return os.environ.get(key, fallback)

tools/bench/fetch_datasets.py

@@ -0,0 +1,40 @@
+"""Pre-fetch quality-benchmark datasets into local JSON.
+
+Run this on a machine WITH network (e.g. your laptop). The resulting
+tools/bench/data/*.json files are then shipped to the GPU host (which has no
+network) by the bench sync step.
+
+Usage:
+  python3 -m tools.bench.fetch_datasets               # all tasks
+  python3 -m tools.bench.fetch_datasets aime2025      # one task
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+if __package__ in (None, ""):
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+from tools.bench.tasks import aime, gsm8k, save_local
+
+FETCHERS = {
+    "aime2025": aime.load_remote,
+    "gsm8k": gsm8k.load_remote,
+}
+
+
+def main() -> None:
+    wanted = sys.argv[1:] or list(FETCHERS)
+    for name in wanted:
+        if name not in FETCHERS:
+            raise SystemExit(f"unknown task: {name} (have: {', '.join(FETCHERS)})")
+        print(f"[fetch] {name} ...")
+        records = FETCHERS[name]()
+        path = save_local(name, records)
+        print(f"[fetch] {name}: {len(records)} records -> {path}")
+
+
+if __name__ == "__main__":
+    main()

tools/bench/quality.py

@@ -0,0 +1,147 @@
+"""Quality suite — run dataset tasks against each system, score, report.
+
+Each task module exposes the same surface:
+    load() -> list[{id, problem, answer, source}]
+    make_messages(problem) -> list[dict]
+    extract_answer(text) -> str | None
+    score(pred, gold) -> bool
+
+Concurrency is fixed at 1 per system for quality runs. Mixing concurrent
+requests with quality scoring is fine (deterministic temperature=0) but the
+extra moving parts aren't worth it for the first iteration.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import statistics
+import time
+from dataclasses import asdict, dataclass
+from typing import Any
+
+import httpx
+
+from .client import chat_stream
+from .config import BenchConfig, SystemEndpoint
+from .tasks import aime, gsm8k
+
+TASKS = {
+    "aime2025": (aime, "quality_max_tokens_aime"),
+    "gsm8k":    (gsm8k, "quality_max_tokens_gsm8k"),
+}
+
+
+@dataclass
+class QualityRow:
+    system: str
+    task: str
+    n_total: int
+    n_correct: int
+    n_errors: int
+    accuracy: float
+    mean_completion_tokens: float
+    mean_ttft_ms: float
+    mean_tpot_ms: float
+    wall_s: float
+
+
+@dataclass
+class QualityCase:
+    system: str
+    task: str
+    problem_id: str
+    gold: str
+    pred: str | None
+    correct: bool
+    completion_tokens: int
+    ttft_ms: float
+    tpot_ms: float
+    e2e_s: float
+    error: str | None
+    response_preview: str
+
+
+async def _run_one_task(
+    ep: SystemEndpoint, task_name: str, task_mod, max_tokens: int, cfg: BenchConfig,
+) -> tuple[QualityRow, list[QualityCase]]:
+    problems = task_mod.load()
+    if cfg.quality_limit is not None:
+        problems = problems[: cfg.quality_limit]
+    print(f"[quality] {ep.name} / {task_name}: {len(problems)} problems "
+          f"(max_tokens={max_tokens})")
+
+    cases: list[QualityCase] = []
+    t_wall = time.perf_counter()
+    async with httpx.AsyncClient(timeout=cfg.request_timeout_s) as client:
+        for prob in problems:
+            messages = task_mod.make_messages(prob["problem"])
+            r = await chat_stream(
+                client, ep.base_url, ep.model_id, messages,
+                max_tokens=max_tokens,
+                temperature=cfg.quality_temperature,
+                api_key=ep.api_key,
+                timeout=cfg.request_timeout_s,
+                extra_body=ep.extra_body,
+            )
+            pred = task_mod.extract_answer(r.text) if r.error is None else None
+            correct = task_mod.score(pred, prob["answer"]) if r.error is None else False
+            cases.append(QualityCase(
+                system=ep.name, task=task_name,
+                problem_id=prob["id"], gold=prob["answer"], pred=pred,
+                correct=correct, completion_tokens=r.completion_tokens,
+                ttft_ms=r.ttft_s * 1000 if r.ttft_s > 0 else -1.0,
+                tpot_ms=r.tpot_s * 1000 if r.tpot_s > 0 else -1.0,
+                e2e_s=r.e2e_s, error=r.error,
+                response_preview=(r.text or "")[:240].replace("\n", " "),
+            ))
+            mark = "✓" if correct else ("E" if r.error else "✗")
+            print(f"  [{mark}] {prob['id']:>4s}  gold={prob['answer']:>6s}  "
+                  f"pred={str(pred):>6s}  tok={r.completion_tokens:5d}  "
+                  f"{r.e2e_s:6.1f}s")
+    wall = time.perf_counter() - t_wall
+
+    ok = [c for c in cases if c.error is None]
+    correct = sum(1 for c in cases if c.correct)
+    errors = sum(1 for c in cases if c.error)
+    row = QualityRow(
+        system=ep.name,
+        task=task_name,
+        n_total=len(cases),
+        n_correct=correct,
+        n_errors=errors,
+        accuracy=correct / max(len(cases) - errors, 1),
+        mean_completion_tokens=statistics.mean(c.completion_tokens for c in ok) if ok else 0.0,
+        mean_ttft_ms=statistics.mean(c.ttft_ms for c in ok if c.ttft_ms > 0) if ok else -1.0,
+        mean_tpot_ms=statistics.mean(c.tpot_ms for c in ok if c.tpot_ms > 0) if ok else -1.0,
+        wall_s=wall,
+    )
+    return row, cases
+
+
+def run_quality(
+    endpoints: list[SystemEndpoint], cfg: BenchConfig, tasks: list[str],
+) -> tuple[list[QualityRow], list[QualityCase]]:
+    all_rows: list[QualityRow] = []
+    all_cases: list[QualityCase] = []
+    for ep in endpoints:
+        print(f"[quality] === {ep.name} ===")
+        for task_name in tasks:
+            if task_name not in TASKS:
+                raise ValueError(f"unknown task: {task_name}")
+            task_mod, max_tok_attr = TASKS[task_name]
+            row, cases = asyncio.run(_run_one_task(
+                ep, task_name, task_mod, getattr(cfg, max_tok_attr), cfg,
+            ))
+            all_rows.append(row)
+            all_cases.extend(cases)
+            print(f"  -> {row.task}: {row.n_correct}/{row.n_total} = "
+                  f"{row.accuracy * 100:.1f}%  ({row.wall_s:.1f}s wall)")
+    return all_rows, all_cases
+
+
+def rows_to_dicts(rows: list[QualityRow]) -> list[dict[str, Any]]:
+    return [asdict(r) for r in rows]
+
+
+def cases_to_dicts(cases: list[QualityCase]) -> list[dict[str, Any]]:
+    return [asdict(c) for c in cases]

tools/bench/report.py

@@ -0,0 +1,122 @@
+"""Combined speed + quality report (markdown + json side-cars)."""
+
+from __future__ import annotations
+
+import datetime as dt
+import json
+import os
+from typing import Any
+
+from .config import DEFAULT_SYSTEMS
+
+
+def _fmt(x: float, nd: int = 1) -> str:
+    if x is None or x < 0:
+        return "—"
+    return f"{x:.{nd}f}"
+
+
+def _speed_table(rows: list[dict[str, Any]]) -> str:
+    if not rows:
+        return "_(no speed results)_\n"
+
+    # scenarios in stable order
+    scenarios: list[str] = []
+    for r in rows:
+        if r["scenario"] not in scenarios:
+            scenarios.append(r["scenario"])
+    systems: list[str] = []
+    for r in rows:
+        if r["system"] not in systems:
+            systems.append(r["system"])
+
+    by = {(r["system"], r["scenario"]): r for r in rows}
+    out = []
+    out.append("| scenario | metric | " + " | ".join(systems) + " | speedup (xserv ÷ llama.cpp) |")
+    out.append("|---|---|" + "|".join(["---"] * (len(systems) + 1)) + "|")
+
+    metrics = [
+        ("ttft_ms_p50", "TTFT p50 (ms)", "lower"),
+        ("ttft_ms_p95", "TTFT p95 (ms)", "lower"),
+        ("tpot_ms_p50", "TPOT p50 (ms/tok)", "lower"),
+        ("throughput_tok_s", "Throughput (tok/s)", "higher"),
+    ]
+    for sc in scenarios:
+        for key, label, direction in metrics:
+            cells = []
+            vals = {}
+            for s in systems:
+                row = by.get((s, sc))
+                v = row[key] if row else -1.0
+                vals[s] = v
+                cells.append(_fmt(v, 2 if "tpot" in key else 1))
+            x = vals.get("xserv", -1.0)
+            l = vals.get("llama.cpp", -1.0)
+            if x > 0 and l > 0:
+                ratio = (x / l) if direction == "higher" else (l / x)
+                cells.append(f"{ratio:.2f}×")
+            else:
+                cells.append("—")
+            out.append(f"| {sc} | {label} | " + " | ".join(cells) + " |")
+    return "\n".join(out) + "\n"
+
+
+def _quality_table(rows: list[dict[str, Any]]) -> str:
+    if not rows:
+        return "_(no quality results)_\n"
+    by_task: dict[str, list[dict[str, Any]]] = {}
+    for r in rows:
+        by_task.setdefault(r["task"], []).append(r)
+    out: list[str] = []
+    out.append("| task | system | n | correct | accuracy | mean tokens | TTFT (ms) | TPOT (ms/tok) | wall (s) |")
+    out.append("|---|---|---|---|---|---|---|---|---|")
+    for task, task_rows in by_task.items():
+        for r in task_rows:
+            out.append(
+                f"| {task} | {r['system']} | {r['n_total']} | {r['n_correct']} | "
+                f"{r['accuracy'] * 100:.1f}% | {r['mean_completion_tokens']:.0f} | "
+                f"{_fmt(r['mean_ttft_ms'])} | {_fmt(r['mean_tpot_ms'], 2)} | {r['wall_s']:.1f} |"
+            )
+    return "\n".join(out) + "\n"
+
+
+def write_report(
+    out_dir: str,
+    speed_rows: list[dict[str, Any]],
+    speed_raw: list[dict[str, Any]],
+    quality_rows: list[dict[str, Any]],
+    quality_cases: list[dict[str, Any]],
+    env: dict[str, Any],
+) -> str:
+    os.makedirs(out_dir, exist_ok=True)
+    stamp = dt.datetime.now().strftime("%Y%m%d-%H%M%S")
+    md_path = os.path.join(out_dir, f"comparison-{stamp}.md")
+    json_path = os.path.join(out_dir, f"comparison-{stamp}.json")
+
+    with open(json_path, "w") as f:
+        json.dump({
+            "stamp": stamp,
+            "env": env,
+            "speed": {"summary": speed_rows, "raw": speed_raw},
+            "quality": {"summary": quality_rows, "cases": quality_cases},
+        }, f, indent=2)
+
+    lines: list[str] = []
+    lines.append(f"# xserv vs llama.cpp — comparison\n")
+    lines.append(f"_Generated: {stamp}_\n")
+    lines.append("## Environment\n")
+    for k, v in env.items():
+        lines.append(f"- **{k}**: {v}")
+    lines.append("")
+    lines.append("## Speed\n")
+    lines.append(_speed_table(speed_rows))
+    lines.append("\n## Quality\n")
+    lines.append(_quality_table(quality_rows))
+    lines.append(f"\n_Raw results: `{os.path.basename(json_path)}`_\n")
+
+    with open(md_path, "w") as f:
+        f.write("\n".join(lines))
+
+    print(f"\n[report] wrote {md_path}")
+    print(f"[report] wrote {json_path}")
+    return md_path

tools/bench/requirements.txt

@@ -0,0 +1,2 @@
+httpx>=0.27
+datasets>=2.20

tools/bench/run_tp_parallel.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+# Run the TP=1/2/4 quality sweep with xserv and llama.cpp CONCURRENTLY on
+# disjoint GPU groups: xserv on GPUs 0..N-1, llama.cpp on GPUs 4..4+N-1.
+# The 8x5090 box is grouped 0-3 / 4-7 (PHB intra-group), so each engine's TP
+# comm stays intra-group and the two engines never contend for a GPU.
+#
+# Run from the repo root on the GPU host. Produces bench-out/tp{1,2,4}-{xserv,llama}.
+
+set -u
+MODEL="${MODEL:-/opt/wjh/models/qwen3-8b}"
+GGUF="${GGUF:-/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf}"
+LIMIT="${LIMIT:-30}"
+MAXSEQ="${MAXSEQ:-2048}"
+TPS="${TPS:-1 2 4}"
+
+for TP in $TPS; do
+    LD=$(seq -s, 4 $((3 + TP)))   # llama GPUs: 4 / 4,5 / 4,5,6,7
+    echo "##### TP=$TP  (xserv GPU 0..$((TP-1))  ||  llama GPU $LD) #####"
+    rm -rf "bench-out/tp$TP-xserv" "bench-out/tp$TP-llama"
+
+    python3 -u -m tools.bench.runner --systems xserv --tp "$TP" \
+        --xserv-bin ./target/release/xserv-server --xserv-model "$MODEL" \
+        --suite quality --quality-tasks aime2025,gsm8k --quality-limit "$LIMIT" \
+        --max-batch 1 --max-seq-len "$MAXSEQ" \
+        --out-dir "bench-out/tp$TP-xserv" > "/tmp/tp$TP-xserv.log" 2>&1 &
+    XP=$!
+
+    python3 -u -m tools.bench.runner --systems llama.cpp --tp "$TP" --llama-devices "$LD" \
+        --llama-bin third_party/llama.cpp/build/bin/llama-server --llama-gguf "$GGUF" \
+        --suite quality --quality-tasks aime2025,gsm8k --quality-limit "$LIMIT" \
+        --max-batch 1 --max-seq-len "$MAXSEQ" \
+        --out-dir "bench-out/tp$TP-llama" > "/tmp/tp$TP-llama.log" 2>&1 &
+    LP=$!
+
+    wait "$XP" "$LP"
+    echo "TP=$TP done (xserv exit=$? )"
+done
+echo ALL_DONE

tools/bench/runner.py

@@ -0,0 +1,234 @@
+"""One-click entrypoint: spin up both servers, run suites, write report.
+
+Usage examples:
+
+  # Full sweep against both systems
+  python3 -m tools.bench.runner \
+      --xserv-bin ./target/release/xserv-server \
+      --xserv-model /opt/wjh/models/qwen3-8b \
+      --llama-bin third_party/llama.cpp/build/bin/llama-server \
+      --llama-gguf /opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf \
+      --suite all
+
+  # Speed-only smoke test
+  python3 -m tools.bench.runner ... --suite speed
+
+  # Quality with 5-problem subsample
+  python3 -m tools.bench.runner ... --suite quality --quality-limit 5
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import platform
+import subprocess
+import sys
+from typing import Any
+
+# Allow running as `python3 tools/bench/runner.py` from repo root.
+if __package__ in (None, ""):
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+from tools.bench.config import (
+    BenchConfig, SystemEndpoint, SYSTEM_XSERV, SYSTEM_LLAMA_CPP,
+)
+from tools.bench.servers import (
+    start_server, stop_server,
+    xserv_launch_cmd, llama_cpp_launch_cmd,
+)
+from tools.bench.speed import run_speed, rows_to_dicts as speed_rows_to_dicts
+from tools.bench.quality import (
+    run_quality, rows_to_dicts as q_rows_to_dicts, cases_to_dicts,
+)
+from tools.bench.report import write_report
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="xserv vs llama.cpp benchmark suite")
+    # Targets
+    p.add_argument("--xserv-bin", default="./target/release/xserv-server")
+    p.add_argument("--xserv-model", required=False,
+                   help="HF model directory for xserv-server (defaults to $XSERV_MODEL_DIR)")
+    p.add_argument("--xserv-port", type=int, default=18080)
+    p.add_argument("--xserv-base-url", default=None,
+                   help="If set, skip launching xserv and target this URL.")
+    p.add_argument("--xserv-model-id", default="qwen3-8b")
+
+    p.add_argument("--llama-bin", default="third_party/llama.cpp/build/bin/llama-server")
+    p.add_argument("--llama-gguf", required=False,
+                   help="GGUF model for llama-server (defaults to $LLAMA_GGUF)")
+    p.add_argument("--llama-port", type=int, default=18081)
+    p.add_argument("--llama-base-url", default=None,
+                   help="If set, skip launching llama-server and target this URL.")
+    p.add_argument("--llama-model-id", default="qwen3-8b",
+                   help="String to send in OpenAI 'model' field; llama-server is permissive.")
+
+    # Shared
+    p.add_argument("--max-batch", type=int, default=4)
+    p.add_argument("--max-seq-len", type=int, default=8192)
+    p.add_argument("--systems", default="xserv,llama.cpp",
+                   help="Comma-separated subset to run, e.g. 'xserv' to skip llama.cpp")
+    p.add_argument("--tp", type=int, default=1,
+                   help="Tensor-parallel degree for BOTH engines (xserv --tp N; "
+                        "llama.cpp --split-mode row over the first N GPUs).")
+    p.add_argument("--llama-devices", default=None,
+                   help="Comma list of GPU ordinals for llama.cpp (first --tp used). "
+                        "Lets llama run on a disjoint GPU group (e.g. 4,5,6,7) so it "
+                        "can run concurrently with xserv on 0..N-1.")
+    p.add_argument("--enable-thinking", action="store_true",
+                   help="Enable Qwen3 thinking on llama.cpp. Default OFF to match "
+                        "xserv, which hardcodes thinking off in its prompt builder.")
+
+    # Suites
+    p.add_argument("--suite", choices=["speed", "quality", "all"], default="all")
+    p.add_argument("--quality-tasks", default="aime2025,gsm8k")
+    p.add_argument("--quality-limit", type=int, default=None,
+                   help="Cap problems per task (smoke test). None = all problems.")
+    p.add_argument("--speed-prompts", type=int, default=8)
+    p.add_argument("--speed-max-tokens", type=int, default=128)
+    p.add_argument("--speed-concurrency", default="1,2,4,8")
+
+    p.add_argument("--out-dir", default="bench-out")
+    return p.parse_args()
+
+
+def build_endpoints(args) -> list[SystemEndpoint]:
+    wanted = set(s.strip() for s in args.systems.split(",") if s.strip())
+    eps: list[SystemEndpoint] = []
+
+    if SYSTEM_XSERV in wanted:
+        if args.xserv_base_url:
+            eps.append(SystemEndpoint(
+                name=SYSTEM_XSERV, base_url=args.xserv_base_url,
+                model_id=args.xserv_model_id, launch_cmd=None,
+            ))
+        else:
+            model_dir = args.xserv_model or os.environ.get("XSERV_MODEL_DIR")
+            if not model_dir:
+                raise SystemExit("--xserv-model or XSERV_MODEL_DIR required (or pass --xserv-base-url)")
+            eps.append(SystemEndpoint(
+                name=SYSTEM_XSERV,
+                base_url=f"http://127.0.0.1:{args.xserv_port}",
+                model_id=args.xserv_model_id,
+                launch_cmd=xserv_launch_cmd(
+                    args.xserv_bin, model_dir, args.xserv_port,
+                    max_batch=args.max_batch, max_seq_len=args.max_seq_len, tp=args.tp,
+                ),
+                health_path="/health",
+                ready_timeout_s=1200.0,
+            ))
+
+    # Match xserv's hardcoded thinking-OFF mode unless explicitly overridden.
+    llama_extra_body = None if args.enable_thinking else {
+        "chat_template_kwargs": {"enable_thinking": False}
+    }
+
+    if SYSTEM_LLAMA_CPP in wanted:
+        if args.llama_base_url:
+            eps.append(SystemEndpoint(
+                name=SYSTEM_LLAMA_CPP, base_url=args.llama_base_url,
+                model_id=args.llama_model_id, launch_cmd=None,
+                extra_body=llama_extra_body,
+            ))
+        else:
+            gguf = args.llama_gguf or os.environ.get("LLAMA_GGUF")
+            if not gguf:
+                raise SystemExit("--llama-gguf or LLAMA_GGUF required (or pass --llama-base-url)")
+            # Pick the GPUs llama.cpp runs on. Default is the first `tp` GPUs;
+            # pass --llama-devices to place it on a disjoint group (e.g. 4,5,6,7)
+            # so it can run concurrently with xserv on 0..N-1. --split-mode row
+            # then tensor-parallel-splits across exactly these devices.
+            if args.llama_devices:
+                devs = [d.strip() for d in args.llama_devices.split(",") if d.strip()][: max(args.tp, 1)]
+                llama_env = {"CUDA_VISIBLE_DEVICES": ",".join(devs)}
+            elif args.tp > 1:
+                llama_env = {"CUDA_VISIBLE_DEVICES": ",".join(str(d) for d in range(args.tp))}
+            else:
+                llama_env = {}
+            eps.append(SystemEndpoint(
+                name=SYSTEM_LLAMA_CPP,
+                base_url=f"http://127.0.0.1:{args.llama_port}",
+                model_id=args.llama_model_id,
+                launch_cmd=llama_cpp_launch_cmd(
+                    args.llama_bin, gguf, args.llama_port,
+                    n_parallel=args.max_batch, ctx_per_slot=args.max_seq_len, tp=args.tp,
+                ),
+                launch_env=llama_env,
+                # llama-server's health endpoint also returns 200 only when model is loaded.
+                health_path="/health",
+                ready_timeout_s=1200.0,
+                extra_body=llama_extra_body,
+            ))
+    return eps
+
+
+def collect_env() -> dict[str, Any]:
+    env: dict[str, Any] = {
+        "platform": platform.platform(),
+        "python": sys.version.split()[0],
+    }
+    for cmd, key in [
+        (["nvidia-smi", "--query-gpu=name,driver_version,memory.total", "--format=csv,noheader"], "gpu"),
+        (["git", "rev-parse", "HEAD"], "xserv_commit"),
+    ]:
+        try:
+            out = subprocess.check_output(cmd, text=True, stderr=subprocess.DEVNULL, timeout=5).strip()
+            env[key] = out.splitlines()[0] if out else "?"
+        except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
+            env[key] = "?"
+    return env
+
+
+def main() -> None:
+    args = parse_args()
+    endpoints = build_endpoints(args)
+    if not endpoints:
+        raise SystemExit("no systems selected (check --systems)")
+
+    cfg = BenchConfig(
+        out_dir=args.out_dir,
+        speed_prompts=args.speed_prompts,
+        speed_max_tokens=args.speed_max_tokens,
+        speed_concurrency=tuple(int(c) for c in args.speed_concurrency.split(",") if c.strip()),
+        quality_limit=args.quality_limit,
+    )
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    log_dir = os.path.join(args.out_dir, "logs")
+
+    speed_rows: list[Any] = []
+    speed_raw: list[dict[str, Any]] = []
+    quality_rows: list[Any] = []
+    quality_cases: list[Any] = []
+    tasks = [t.strip() for t in args.quality_tasks.split(",") if t.strip()]
+
+    # One server at a time. Two BF16 8B models (~16GB each) do not co-reside on a
+    # single 32GB GPU, and even if they did, a resident idle engine would distort
+    # the other's measurements. Start → run all suites → stop, then next system.
+    for ep in endpoints:
+        h = start_server(ep, log_dir)
+        try:
+            if args.suite in ("speed", "all"):
+                rows, raw = run_speed([ep], cfg)
+                speed_rows.extend(rows)
+                speed_raw.extend(raw)
+            if args.suite in ("quality", "all"):
+                rows, cases = run_quality([ep], cfg, tasks)
+                quality_rows.extend(rows)
+                quality_cases.extend(cases)
+        finally:
+            stop_server(h)
+
+    write_report(
+        out_dir=args.out_dir,
+        speed_rows=speed_rows_to_dicts(speed_rows) if speed_rows else [],
+        speed_raw=speed_raw,
+        quality_rows=q_rows_to_dicts(quality_rows) if quality_rows else [],
+        quality_cases=cases_to_dicts(quality_cases) if quality_cases else [],
+        env=collect_env(),
+    )
+
+
+if __name__ == "__main__":
+    main()

tools/bench/servers.py

@@ -0,0 +1,161 @@
+"""Start/stop xserv-server and llama-server as subprocesses.
+
+The benchmark driver treats both systems as black-box HTTP servers — it does
+not import their Rust/C++ code. This keeps the comparison fair (same wire
+protocol, no in-process shortcut) and avoids coupling the bench harness to
+internal APIs.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import signal
+import subprocess
+import sys
+import time
+import urllib.error
+import urllib.request
+from dataclasses import dataclass
+
+from .config import SystemEndpoint
+
+
+@dataclass
+class ServerHandle:
+    endpoint: SystemEndpoint
+    proc: subprocess.Popen[bytes] | None
+    log_path: str | None
+
+
+def _wait_ready(base_url: str, health_path: str, timeout_s: float) -> bool:
+    url = base_url.rstrip("/") + health_path
+    deadline = time.monotonic() + timeout_s
+    last_err = ""
+    while time.monotonic() < deadline:
+        try:
+            with urllib.request.urlopen(url, timeout=5) as r:
+                if r.status == 200:
+                    return True
+        except (urllib.error.URLError, ConnectionError, TimeoutError) as e:
+            last_err = repr(e)
+            time.sleep(1.0)
+    print(f"[servers] not ready after {timeout_s}s ({url}): {last_err}", file=sys.stderr)
+    return False
+
+
+def start_server(ep: SystemEndpoint, log_dir: str) -> ServerHandle:
+    """Launch `ep.launch_cmd` if set; otherwise assume it's already running."""
+    if ep.launch_cmd is None:
+        if _wait_ready(ep.base_url, ep.health_path, timeout_s=10.0):
+            print(f"[servers] reusing already-running {ep.name} at {ep.base_url}")
+            return ServerHandle(endpoint=ep, proc=None, log_path=None)
+        raise RuntimeError(f"{ep.name}: no launch_cmd and not reachable at {ep.base_url}")
+
+    os.makedirs(log_dir, exist_ok=True)
+    log_path = os.path.join(log_dir, f"{ep.name.replace('.', '_')}.log")
+    log_f = open(log_path, "wb")
+    env = os.environ.copy()
+    env.update(ep.launch_env)
+
+    print(f"[servers] launching {ep.name}: {' '.join(ep.launch_cmd)}")
+    print(f"[servers]   log: {log_path}")
+    proc = subprocess.Popen(
+        ep.launch_cmd,
+        cwd=ep.launch_cwd,
+        env=env,
+        stdout=log_f,
+        stderr=subprocess.STDOUT,
+        # Own process group so SIGTERM kills children (llama-server in particular).
+        preexec_fn=os.setsid,
+    )
+
+    ok = _wait_ready(ep.base_url, ep.health_path, timeout_s=ep.ready_timeout_s)
+    if not ok:
+        # Hand back enough info so caller can drain logs before dying.
+        log_f.flush()
+        try:
+            os.killpg(proc.pid, signal.SIGTERM)
+        except ProcessLookupError:
+            pass
+        raise RuntimeError(
+            f"{ep.name} failed to become ready (see {log_path}). "
+            "Common causes: model path wrong, port already in use, OOM."
+        )
+
+    return ServerHandle(endpoint=ep, proc=proc, log_path=log_path)
+
+
+def stop_server(h: ServerHandle, *, grace_s: float = 10.0) -> None:
+    if h.proc is None:
+        return
+    print(f"[servers] stopping {h.endpoint.name} (pid {h.proc.pid})")
+    try:
+        os.killpg(h.proc.pid, signal.SIGTERM)
+    except ProcessLookupError:
+        return
+    try:
+        h.proc.wait(timeout=grace_s)
+    except subprocess.TimeoutExpired:
+        print(f"[servers] {h.endpoint.name} did not exit, sending SIGKILL")
+        with contextlib.suppress(ProcessLookupError):
+            os.killpg(h.proc.pid, signal.SIGKILL)
+        h.proc.wait(timeout=5)
+
+
+# ---------- launch-command builders ----------
+
+
+def xserv_launch_cmd(
+    bin_path: str,
+    model_dir: str,
+    port: int,
+    *,
+    max_batch: int,
+    max_seq_len: int,
+    tp: int = 1,
+) -> list[str]:
+    cmd = [
+        bin_path,
+        model_dir,
+        "--port", str(port),
+        "--max-batch", str(max_batch),
+        "--max-seq-len", str(max_seq_len),
+    ]
+    if tp > 1:
+        cmd += ["--tp", str(tp)]  # xserv binds rank r -> GPU r internally
+    return cmd
+
+
+def llama_cpp_launch_cmd(
+    bin_path: str,
+    gguf_path: str,
+    port: int,
+    *,
+    n_parallel: int,
+    ctx_per_slot: int,
+    n_gpu_layers: int = 99,
+    tp: int = 1,
+) -> list[str]:
+    # llama.cpp DIVIDES total -c across --parallel slots: per-slot context is
+    # n_ctx / n_parallel. xserv gives each sequence the full max_seq_len, so to
+    # match we must set total -c = ctx_per_slot * n_parallel. Getting this wrong
+    # silently truncates long generations (e.g. AIME) on llama.cpp's side.
+    total_ctx = ctx_per_slot * n_parallel
+    cmd = [
+        bin_path,
+        "-m", gguf_path,
+        "--port", str(port),
+        "--host", "0.0.0.0",
+        "-c", str(total_ctx),
+        "-ngl", str(n_gpu_layers),
+        "--parallel", str(n_parallel),
+        # NOTE: do NOT pass --log-disable; its startup log reports per-slot
+        # n_ctx, which is exactly the diagnostic that catches ctx misconfig.
+    ]
+    if tp > 1:
+        # Tensor-parallel split across the visible GPUs (caller restricts the
+        # set via CUDA_VISIBLE_DEVICES in launch_env). Row-split is llama.cpp's
+        # tensor-parallel mode (vs the default layer/pipeline split).
+        cmd += ["--split-mode", "row"]
+    return cmd

tools/bench/speed.py

@@ -0,0 +1,171 @@
+"""Speed suite: TTFT, TPOT, throughput; serial and concurrent.
+
+Single-stream and concurrent throughput are reported separately because they
+stress different things — TTFT/TPOT are kernel/latency bound (single stream),
+throughput at high concurrency is scheduler/batching bound.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import statistics
+from dataclasses import asdict, dataclass
+from typing import Any
+
+from .client import StreamResult, chat_concurrent
+from .config import BenchConfig, SystemEndpoint
+
+
+# Three prompt-length buckets cover the common interesting points:
+# short = greeting-style; medium = QA; long = summarize-ish (prefill-heavy).
+SPEED_PROMPTS = {
+    "short":  "What is 2 + 2?",
+    "medium": "Explain the difference between TCP and UDP, briefly.",
+    "long": (
+        "Write a detailed comparison of Python and Rust for systems programming. "
+        "Cover memory management, performance, ergonomics, ecosystem, and typical "
+        "use cases. Be specific."
+    ),
+}
+
+
+@dataclass
+class SpeedRow:
+    system: str
+    scenario: str          # e.g. "single/short", "concurrent-4"
+    requests: int
+    completion_tokens_total: int
+    wall_s: float
+    ttft_ms_p50: float
+    ttft_ms_p95: float
+    tpot_ms_p50: float
+    tpot_ms_p95: float
+    throughput_tok_s: float    # aggregate completion_tokens / wall
+    per_req_throughput_tok_s_mean: float
+    errors: int
+
+
+def _percentile(values: list[float], p: float) -> float:
+    if not values:
+        return -1.0
+    s = sorted(values)
+    idx = max(0, min(len(s) - 1, int(round((p / 100.0) * (len(s) - 1)))))
+    return s[idx]
+
+
+def _summarize(system: str, scenario: str, results: list[StreamResult], wall_s: float) -> SpeedRow:
+    ok = [r for r in results if r.error is None]
+    ttft_ms = [r.ttft_s * 1000 for r in ok if r.ttft_s >= 0]
+    tpot_ms = [r.tpot_s * 1000 for r in ok if r.tpot_s >= 0]
+    per_req_tps = [r.throughput_tok_s for r in ok if r.throughput_tok_s > 0]
+    total_tokens = sum(r.completion_tokens for r in ok)
+    return SpeedRow(
+        system=system,
+        scenario=scenario,
+        requests=len(results),
+        completion_tokens_total=total_tokens,
+        wall_s=wall_s,
+        ttft_ms_p50=_percentile(ttft_ms, 50),
+        ttft_ms_p95=_percentile(ttft_ms, 95),
+        tpot_ms_p50=_percentile(tpot_ms, 50),
+        tpot_ms_p95=_percentile(tpot_ms, 95),
+        throughput_tok_s=total_tokens / wall_s if wall_s > 0 else -1.0,
+        per_req_throughput_tok_s_mean=statistics.mean(per_req_tps) if per_req_tps else -1.0,
+        errors=len(results) - len(ok),
+    )
+
+
+async def run_single_stream(
+    ep: SystemEndpoint, cfg: BenchConfig,
+) -> tuple[list[SpeedRow], list[dict[str, Any]]]:
+    """One request at a time, three prompt lengths. Repeat each `cfg.speed_prompts` times."""
+    rows: list[SpeedRow] = []
+    raw: list[dict[str, Any]] = []
+    for bucket, prompt in SPEED_PROMPTS.items():
+        messages = [[{"role": "user", "content": prompt}]] * cfg.speed_prompts
+        results, wall = await chat_concurrent(
+            ep.base_url, ep.model_id, messages,
+            max_tokens=cfg.speed_max_tokens,
+            temperature=0.0,
+            api_key=ep.api_key,
+            timeout=cfg.request_timeout_s,
+            concurrency=1,
+            extra_body=ep.extra_body,
+        )
+        rows.append(_summarize(ep.name, f"single/{bucket}", results, wall))
+        for i, r in enumerate(results):
+            raw.append({
+                "system": ep.name, "scenario": f"single/{bucket}", "i": i,
+                "ttft_s": r.ttft_s, "tpot_s": r.tpot_s,
+                "completion_tokens": r.completion_tokens,
+                "e2e_s": r.e2e_s, "error": r.error,
+                "finish_reason": r.finish_reason,
+            })
+    return rows, raw
+
+
+async def run_concurrent(
+    ep: SystemEndpoint, cfg: BenchConfig,
+) -> tuple[list[SpeedRow], list[dict[str, Any]]]:
+    """Fixed medium-length prompt, sweep concurrency."""
+    rows: list[SpeedRow] = []
+    raw: list[dict[str, Any]] = []
+    prompt = SPEED_PROMPTS["medium"]
+    for c in cfg.speed_concurrency:
+        # Send 4x concurrency requests so the scheduler sees sustained load,
+        # not just one wave.
+        n = max(c * 4, 8)
+        messages = [[{"role": "user", "content": prompt}]] * n
+        results, wall = await chat_concurrent(
+            ep.base_url, ep.model_id, messages,
+            max_tokens=cfg.speed_max_tokens,
+            temperature=0.0,
+            api_key=ep.api_key,
+            timeout=cfg.request_timeout_s,
+            concurrency=c,
+            extra_body=ep.extra_body,
+        )
+        rows.append(_summarize(ep.name, f"concurrent-{c}", results, wall))
+        for i, r in enumerate(results):
+            raw.append({
+                "system": ep.name, "scenario": f"concurrent-{c}", "i": i,
+                "ttft_s": r.ttft_s, "tpot_s": r.tpot_s,
+                "completion_tokens": r.completion_tokens,
+                "e2e_s": r.e2e_s, "error": r.error,
+                "finish_reason": r.finish_reason,
+            })
+    return rows, raw
+
+
+def run_speed(
+    endpoints: list[SystemEndpoint], cfg: BenchConfig,
+) -> tuple[list[SpeedRow], list[dict[str, Any]]]:
+    all_rows: list[SpeedRow] = []
+    all_raw: list[dict[str, Any]] = []
+    for ep in endpoints:
+        print(f"[speed] === {ep.name} ===")
+        # Tiny warmup so the first row isn't penalized by lazy cache allocation.
+        warm_messages = [[{"role": "user", "content": "Hello"}]]
+        asyncio.run(chat_concurrent(
+            ep.base_url, ep.model_id, warm_messages,
+            max_tokens=8, temperature=0.0, api_key=ep.api_key,
+            timeout=120, concurrency=1, extra_body=ep.extra_body,
+        ))
+
+        rows1, raw1 = asyncio.run(run_single_stream(ep, cfg))
+        all_rows.extend(rows1); all_raw.extend(raw1)
+        for r in rows1:
+            print(f"  {r.scenario:18s} ttft p50={r.ttft_ms_p50:7.1f}ms  "
+                  f"tpot p50={r.tpot_ms_p50:6.2f}ms  thpt={r.throughput_tok_s:6.1f} tok/s")
+
+        rows2, raw2 = asyncio.run(run_concurrent(ep, cfg))
+        all_rows.extend(rows2); all_raw.extend(raw2)
+        for r in rows2:
+            print(f"  {r.scenario:18s} reqs={r.requests:3d}  thpt={r.throughput_tok_s:6.1f} tok/s  "
+                  f"ttft p95={r.ttft_ms_p95:7.1f}ms  errs={r.errors}")
+
+    return all_rows, all_raw
+
+
+def rows_to_dicts(rows: list[SpeedRow]) -> list[dict[str, Any]]:
+    return [asdict(r) for r in rows]

tools/bench/summarize_tp.py

@@ -0,0 +1,24 @@
+"""Summarize the concurrent TP sweep: bench-out/tp{1,2,4}-{xserv,llama}."""
+import glob
+import json
+import os
+import sys
+
+base = sys.argv[1] if len(sys.argv) > 1 else "bench-out"
+rows = []
+for tp in (1, 2, 4):
+    for sysname in ("xserv", "llama"):
+        files = sorted(glob.glob(os.path.join(base, f"tp{tp}-{sysname}", "comparison-*.json")))
+        if not files:
+            continue
+        d = json.load(open(files[-1]))
+        for r in d["quality"]["summary"]:
+            rows.append((tp, sysname, r["task"], r["n_correct"], r["n_total"],
+                         r["accuracy"] * 100, r["mean_completion_tokens"],
+                         r["mean_ttft_ms"], r["mean_tpot_ms"], r["wall_s"]))
+
+print("%-3s %-7s %-9s %-9s %7s %9s %9s %10s %9s" %
+      ("TP", "engine", "task", "correct", "acc%", "mean_tok", "TTFT_ms", "TPOT_ms", "wall_s"))
+for (tp, s, task, nc, nt, acc, tok, ttft, tpot, wall) in rows:
+    print("%-3d %-7s %-9s %-9s %6.1f%% %9.0f %9.1f %10.2f %9.0f" %
+          (tp, s, task, f"{nc}/{nt}", acc, tok, ttft, tpot, wall))

tools/bench/tasks/__init__.py

@@ -0,0 +1,46 @@
+"""Shared helpers for quality tasks.
+
+Each task can be backed by a pre-fetched local JSON file (so the GPU host
+doesn't need network). The JSON is a list of records:
+    [{"id": str, "problem": str, "answer": str, "source": str}, ...]
+
+Use tools/bench/fetch_datasets.py on a networked machine to produce these
+files, then ship them to the GPU host (the bench sync does this automatically).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+
+def data_dir() -> str:
+    """Directory holding pre-fetched dataset JSON. Override via BENCH_DATA_DIR."""
+    return os.environ.get(
+        "BENCH_DATA_DIR",
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data"),
+    )
+
+
+def local_json_path(task_name: str) -> str:
+    return os.path.normpath(os.path.join(data_dir(), f"{task_name}.json"))
+
+
+def load_local(task_name: str) -> list[dict[str, Any]] | None:
+    """Return records from the local JSON file if present, else None."""
+    path = local_json_path(task_name)
+    if not os.path.isfile(path):
+        return None
+    with open(path) as f:
+        records = json.load(f)
+    print(f"[tasks] loaded {len(records)} records from {path}")
+    return records
+
+
+def save_local(task_name: str, records: list[dict[str, Any]]) -> str:
+    path = local_json_path(task_name)
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "w") as f:
+        json.dump(records, f, ensure_ascii=False, indent=1)
+    return path

tools/bench/tasks/aime.py

@@ -0,0 +1,114 @@
+"""AIME 2025 — 30 problems, integer answers 0..999.
+
+Scoring: exact-match of the integer in the last `\\boxed{...}` in the response,
+falling back to the last standalone integer in the response. Matches the
+convention used by most reasoning-model leaderboards.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from . import load_local
+
+TASK_NAME = "aime2025"
+
+# Tried in order; first one to load wins. These are the most-cited HF datasets
+# for AIME 2025 at time of writing; we don't depend on any one being present.
+DATASET_CANDIDATES = [
+    ("MathArena/aime_2025", None, "test"),
+    ("yentinglin/aime_2025", None, "train"),
+    ("opencompass/AIME2025", "AIME2025-I", "test"),
+]
+
+
+def load() -> list[dict[str, Any]]:
+    # Prefer the pre-fetched local JSON (GPU host has no network).
+    local = load_local(TASK_NAME)
+    if local is not None:
+        return local
+    return load_remote()
+
+
+def load_remote() -> list[dict[str, Any]]:
+    """Fetch from HuggingFace. Requires network — used by fetch_datasets.py."""
+    from datasets import load_dataset  # noqa: PLC0415 — optional dep, see requirements.txt
+
+    last_err: Exception | None = None
+    for repo, config, split in DATASET_CANDIDATES:
+        try:
+            ds = load_dataset(repo, config, split=split) if config else load_dataset(repo, split=split)
+        except Exception as e:  # noqa: BLE001 — try the next candidate
+            last_err = e
+            continue
+
+        problems: list[dict[str, Any]] = []
+        for i, row in enumerate(ds):
+            problem = row.get("problem") or row.get("question") or row.get("Problem")
+            answer = row.get("answer") or row.get("Answer") or row.get("solution_int")
+            if problem is None or answer is None:
+                continue
+            problems.append({
+                "id": str(row.get("id") or row.get("ID") or i),
+                "problem": problem,
+                "answer": str(answer).strip(),
+                "source": repo,
+            })
+        if problems:
+            return problems
+
+    raise RuntimeError(
+        f"Could not load AIME 2025 from any of {[c[0] for c in DATASET_CANDIDATES]} "
+        f"(last error: {last_err!r}). Set HF_HOME / HF_TOKEN if needed."
+    )
+
+
+SYSTEM_PROMPT = (
+    "You are a careful math problem solver. Solve the problem step by step. "
+    "Put your final integer answer (an integer from 0 to 999) inside \\boxed{}."
+)
+
+
+def make_messages(problem: str) -> list[dict[str, str]]:
+    return [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": problem},
+    ]
+
+
+_BOXED_RE = re.compile(r"\\boxed\s*\{([^{}]*)\}")
+_INT_RE = re.compile(r"-?\d+")
+
+
+def extract_answer(text: str) -> str | None:
+    """Return canonical integer string, or None if nothing parseable."""
+    if not text:
+        return None
+    boxed = _BOXED_RE.findall(text)
+    candidates: list[str] = []
+    if boxed:
+        # Inside the \boxed{} there may be extra latex; grab the last integer.
+        ints = _INT_RE.findall(boxed[-1])
+        if ints:
+            candidates.append(ints[-1])
+    # Fallback: the last integer anywhere in the response.
+    if not candidates:
+        ints = _INT_RE.findall(text)
+        if ints:
+            candidates.append(ints[-1])
+    if not candidates:
+        return None
+    try:
+        return str(int(candidates[-1]))
+    except ValueError:
+        return None
+
+
+def score(pred: str | None, gold: str) -> bool:
+    if pred is None:
+        return False
+    try:
+        return int(pred) == int(gold)
+    except ValueError:
+        return False

tools/bench/tasks/gsm8k.py

@@ -0,0 +1,90 @@
+"""GSM8K — 1319 grade-school math problems with integer/decimal answers.
+
+Gold answers in the dataset are in the form `... #### 42`. We score by
+exact-match of the final number, with the same `\\boxed{}` / last-number
+extraction used for AIME, since for instruction-tuned models the response
+follows the prompt instructions, not the dataset's `####` convention.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from . import load_local
+
+TASK_NAME = "gsm8k"
+
+
+def load() -> list[dict[str, Any]]:
+    local = load_local(TASK_NAME)
+    if local is not None:
+        return local
+    return load_remote()
+
+
+def load_remote() -> list[dict[str, Any]]:
+    """Fetch from HuggingFace. Requires network — used by fetch_datasets.py."""
+    from datasets import load_dataset  # noqa: PLC0415
+
+    ds = load_dataset("openai/gsm8k", "main", split="test")
+    out: list[dict[str, Any]] = []
+    for i, row in enumerate(ds):
+        ans_full: str = row["answer"]
+        # gold format: "<chain of thought>\n#### 42"
+        gold = ans_full.split("####")[-1].strip().replace(",", "")
+        out.append({
+            "id": str(i),
+            "problem": row["question"],
+            "answer": gold,
+            "source": "openai/gsm8k",
+        })
+    return out
+
+
+SYSTEM_PROMPT = (
+    "You are a careful math problem solver. Solve the problem step by step. "
+    "Put your final numeric answer inside \\boxed{}."
+)
+
+
+def make_messages(problem: str) -> list[dict[str, str]]:
+    return [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": problem},
+    ]
+
+
+_BOXED_RE = re.compile(r"\\boxed\s*\{([^{}]*)\}")
+# Allow comma-grouped thousands (e.g. "3,500"); _normalize_num strips them.
+_NUM_RE = re.compile(r"-?\d+(?:,\d{3})*(?:\.\d+)?")
+
+
+def _normalize_num(s: str) -> str | None:
+    s = s.replace(",", "").strip()
+    try:
+        f = float(s)
+    except ValueError:
+        return None
+    return str(int(f)) if f.is_integer() else f"{f:g}"
+
+
+def extract_answer(text: str) -> str | None:
+    if not text:
+        return None
+    boxed = _BOXED_RE.findall(text)
+    if boxed:
+        nums = _NUM_RE.findall(boxed[-1])
+        if nums:
+            return _normalize_num(nums[-1])
+    nums = _NUM_RE.findall(text)
+    if nums:
+        return _normalize_num(nums[-1])
+    return None
+
+
+def score(pred: str | None, gold: str) -> bool:
+    if pred is None:
+        return False
+    gold_norm = _normalize_num(gold)
+    return gold_norm is not None and pred == gold_norm

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/convert-to-gguf.sh

@@ -0,0 +1,55 @@
+#!/usr/bin/env bash
+# Convert a HuggingFace safetensors model dir into a BF16 GGUF for llama.cpp.
+#
+# Why BF16: we run xserv in BF16, so the baseline must run BF16 too. If we
+# compared xserv-BF16 against llama.cpp-Q4_K_M the speed delta would be
+# dominated by quantization, not by our kernels — that's not an apples-to-
+# apples comparison.
+#
+# Usage:
+#   tools/convert-to-gguf.sh <hf-model-dir> [out.gguf]
+#
+# Example:
+#   tools/convert-to-gguf.sh /opt/wjh/models/qwen3-8b
+#   # → /opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf
+
+set -euo pipefail
+
+if [ "$#" -lt 1 ]; then
+    echo "Usage: $0 <hf-model-dir> [out.gguf]" >&2
+    exit 1
+fi
+
+SRC="$(realpath "$1")"
+ROOT_DIR="$(cd "$(dirname "$0")/.." && pwd)"
+CONVERT_PY="$ROOT_DIR/third_party/llama.cpp/convert_hf_to_gguf.py"
+
+if [ ! -f "$CONVERT_PY" ]; then
+    echo "convert script not found: $CONVERT_PY" >&2
+    echo "Run tools/setup-llama-cpp.sh first." >&2
+    exit 1
+fi
+
+if [ ! -d "$SRC" ]; then
+    echo "source model dir not found: $SRC" >&2
+    exit 1
+fi
+
+if [ "$#" -ge 2 ]; then
+    OUT="$2"
+else
+    BASENAME="$(basename "$SRC")"
+    OUT="$SRC/${BASENAME}-bf16.gguf"
+fi
+
+if [ -f "$OUT" ]; then
+    echo "==> already exists: $OUT (skipping; remove to force re-convert)"
+    echo "$OUT"
+    exit 0
+fi
+
+echo "==> converting $SRC -> $OUT (BF16)"
+python3 "$CONVERT_PY" "$SRC" --outfile "$OUT" --outtype bf16
+
+echo "=== done ==="
+echo "$OUT"

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/setup-llama-cpp.sh

@@ -0,0 +1,94 @@
+#!/usr/bin/env bash
+# Build the llama.cpp baseline (third_party/llama.cpp) with CUDA.
+#
+# Source is vendored as a git submodule pinned to a fixed tag (see .gitmodules
+# and the recorded gitlink commit). This script does NOT fetch from the network
+# by default — it expects the source to already be present, either via:
+#   - `git submodule update --init` (on a host with network), or
+#   - rsync/tar transfer (how it reaches dash5, which has no network).
+#
+# It only fetches as a convenience fallback when the source is missing AND
+# network is reachable.
+#
+# Idempotent. Safe to re-run.
+#
+# Usage:
+#   tools/setup-llama-cpp.sh            # build (configure if needed)
+#   tools/setup-llama-cpp.sh --rebuild  # wipe build dir, reconfigure, rebuild
+#
+# Env:
+#   CUDA_ARCH   CUDA architectures for cmake (default 120-real = RTX 5090 SM120)
+#   CUDA_HOME   CUDA toolkit root (auto-detected: /usr/local/cuda-12.9 then cuda)
+
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "$0")/.." && pwd)"
+VENDOR_DIR="$ROOT_DIR/third_party/llama.cpp"
+CUDA_ARCH="${CUDA_ARCH:-120-real}"
+REBUILD=0
+for arg in "$@"; do
+    case "$arg" in
+        --rebuild) REBUILD=1 ;;
+        --help|-h) grep -E '^#' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
+    esac
+done
+
+if [ -d /usr/local/cuda-12.9 ]; then
+    export CUDA_HOME="${CUDA_HOME:-/usr/local/cuda-12.9}"
+elif [ -d /usr/local/cuda ]; then
+    export CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
+fi
+[ -n "${CUDA_HOME:-}" ] && export PATH="$CUDA_HOME/bin:$PATH"
+
+echo "=== llama.cpp build ==="
+echo "  vendor dir : $VENDOR_DIR"
+echo "  CUDA arch  : $CUDA_ARCH"
+echo "  CUDA_HOME  : ${CUDA_HOME:-<not set>}"
+
+# --- Ensure source is present ---
+if [ ! -f "$VENDOR_DIR/CMakeLists.txt" ]; then
+    echo "==> source missing at $VENDOR_DIR"
+    if git -C "$ROOT_DIR" rev-parse --git-dir >/dev/null 2>&1 \
+        && timeout 8 git ls-remote https://github.com/ggerganov/llama.cpp HEAD >/dev/null 2>&1; then
+        echo "==> network OK, initializing submodule"
+        git -C "$ROOT_DIR" submodule update --init --recursive third_party/llama.cpp
+    else
+        echo "ERROR: llama.cpp source not present and network unavailable." >&2
+        echo "  On a networked host run: git submodule update --init third_party/llama.cpp" >&2
+        echo "  Then transfer the source here (the bench tooling does this via rsync)." >&2
+        exit 1
+    fi
+fi
+
+BUILD_DIR="$VENDOR_DIR/build"
+if [ "$REBUILD" -eq 1 ] && [ -d "$BUILD_DIR" ]; then
+    echo "==> --rebuild: removing $BUILD_DIR"
+    rm -rf "$BUILD_DIR"
+fi
+
+SERVER_BIN="$BUILD_DIR/bin/llama-server"
+if [ -x "$SERVER_BIN" ] && [ "$REBUILD" -eq 0 ]; then
+    echo "==> already built: $SERVER_BIN (use --rebuild to force)"
+    echo "$SERVER_BIN"
+    exit 0
+fi
+
+echo "==> cmake configure"
+cmake -S "$VENDOR_DIR" -B "$BUILD_DIR" \
+    -DGGML_CUDA=ON \
+    -DLLAMA_CURL=OFF \
+    -DLLAMA_BUILD_TESTS=OFF \
+    -DLLAMA_BUILD_EXAMPLES=OFF \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_CUDA_ARCHITECTURES="$CUDA_ARCH"
+
+echo "==> build llama-server llama-cli (jobs: $(nproc))"
+cmake --build "$BUILD_DIR" --target llama-server llama-cli -j "$(nproc)"
+
+if [ ! -x "$SERVER_BIN" ]; then
+    echo "ERROR: llama-server did not build at $SERVER_BIN" >&2
+    exit 1
+fi
+
+echo "=== done ==="
+echo "$SERVER_BIN"