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

@@ -2,6 +2,12 @@
 resolver = "2"
 members = [
     "crates/xserv-cuda",
+    "crates/xserv-tensor",
+    "crates/xserv-kernels",
+    "crates/xserv-model",
+    "crates/xserv-tokenizer",
+    "crates/xserv-server",
+    "crates/xserv-distributed",
 ]
 
 [workspace.package]
@@ -12,3 +18,12 @@ license = "MIT"
 [workspace.dependencies]
 half = "2"
 smallvec = "1"
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+safetensors = "0.5"
+regex = "1"
+tokio = { version = "1", features = ["full"] }
+axum = "0.8"
+uuid = { version = "1", features = ["v4"] }
+tokio-stream = "0.1"
+rand = "0.8"

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/error.rs

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

crates/xserv-cuda/src/ffi.rs

@@ -3,6 +3,8 @@ use std::os::raw::c_char;
 
 pub type CudaStream = *mut c_void;
 pub type CudaEvent = *mut c_void;
+pub type CudaGraph = *mut c_void;
+pub type CudaGraphExec = *mut c_void;
 
 pub const CUDA_MEMCPY_H2D: i32 = 1;
 pub const CUDA_MEMCPY_D2H: i32 = 2;
@@ -11,31 +13,16 @@ pub const CUDA_MEMCPY_D2D: i32 = 3;
 pub const CUDA_SUCCESS: i32 = 0;
 pub const CUDA_ERROR_OUT_OF_MEMORY: i32 = 2;
 
-#[repr(C)]
+/// 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 {
@@ -87,6 +102,70 @@ impl GpuBuffer {
         error::check(unsafe { ffi::cudaMemset(self.ptr, 0, self.len) })
     }
 
+    /// Copy `count` bytes from `src` buffer at `src_offset` to this buffer at `dst_offset`.
+    pub fn copy_from_device_at(&mut self, src: &GpuBuffer, src_offset: usize, dst_offset: usize, count: usize) -> Result<()> {
+        assert!(src_offset + count <= src.len);
+        assert!(dst_offset + count <= self.len);
+        error::check(unsafe {
+            ffi::cudaMemcpy(
+                self.ptr.add(dst_offset),
+                src.ptr.add(src_offset),
+                count,
+                ffi::CUDA_MEMCPY_D2D,
+            )
+        })
+    }
+
+    /// 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) {
@@ -99,14 +178,29 @@ impl GpuBuffer {
     /// Reconstruct a GpuBuffer from a raw pointer + length.
     /// Safety: ptr must have been allocated with cudaMalloc, len must be correct.
     pub unsafe fn from_raw(ptr: *mut u8, len: usize) -> Self {
-        Self { ptr, len }
+        Self { ptr, len, owned: true, pooled: false }
+    }
+
+    /// Create a non-owning view of GPU memory. Dropping this buffer does NOT
+    /// call `cudaFree`. The caller must ensure the underlying allocation
+    /// outlives this borrow.
+    ///
+    /// # Safety
+    /// `ptr` must point to a valid GPU allocation of at least `len` bytes that
+    /// will remain live for the lifetime of the returned `GpuBuffer`.
+    pub unsafe fn borrow_raw(ptr: *mut u8, len: usize) -> Self {
+        Self { ptr, len, owned: false, pooled: false }
     }
 }
 
 impl Drop for GpuBuffer {
     fn drop(&mut self) {
-        if !self.ptr.is_null() {
+        if self.owned && !self.ptr.is_null() {
-            unsafe { ffi::cudaFree(self.ptr) };
+            if self.pooled {
+                crate::allocator::return_to_pool(self.ptr, self.len);
+            } else {
+                unsafe { ffi::cudaFree(self.ptr) };
+            }
         }
     }
 }

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/Cargo.toml

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

crates/xserv-kernels/build.rs

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

crates/xserv-kernels/src/activation.rs

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

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

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

crates/xserv-kernels/src/gemm.rs

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

crates/xserv-kernels/src/layernorm.rs

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

crates/xserv-kernels/src/lib.rs

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

crates/xserv-kernels/src/rmsnorm.rs

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

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

crates/xserv-kernels/src/softmax.rs

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

crates/xserv-kernels/src/transpose.rs

@@ -0,0 +1,142 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_reshape_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_merge_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_repeat_kv_bf16(inp: *const c_void, out: *mut c_void, kv_heads: i32, n_rep: i32, seq_len: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_strided_copy_bf16(inp: *const c_void, out: *mut c_void, numel: i32, ndim: i32,
+        shape0: i32, shape1: i32, shape2: i32, shape3: i32,
+        in_stride0: i32, in_stride1: i32, in_stride2: i32, in_stride3: i32,
+        in_offset: i32, stream: *mut c_void);
+    fn launch_strided_copy_f32(inp: *const c_void, out: *mut c_void, numel: i32, ndim: i32,
+        shape0: i32, shape1: i32, shape2: i32, shape3: i32,
+        in_stride0: i32, in_stride1: i32, in_stride2: i32, in_stride3: i32,
+        in_offset: i32, stream: *mut c_void);
+}
+
+/// [S, H*D] → [1, H, S, D] on GPU (BF16)
+pub fn reshape_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
+    let out = Tensor::empty(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
+    unsafe {
+        launch_reshape_heads_bf16(
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
+        );
+    }
+    out
+}
+
+/// [1, H, S, D] → [S, H*D] on GPU (BF16)
+pub fn merge_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
+    let hidden = num_heads * head_dim;
+    let out = Tensor::empty(&[seq_len, hidden], DType::BF16, x.device());
+    unsafe {
+        launch_merge_heads_bf16(
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
+        );
+    }
+    out
+}
+
+/// [1, H, S, D] → [S, H, D] for RoPE on GPU (BF16)
+pub fn transpose_for_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
+    let out = Tensor::empty(&[seq_len, num_heads, head_dim], DType::BF16, x.device());
+    unsafe {
+        launch_transpose_hsd_to_shd_bf16(
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
+        );
+    }
+    out
+}
+
+/// [S, H, D] → [1, H, S, D] after RoPE on GPU (BF16)
+pub fn transpose_from_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
+    let out = Tensor::empty(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
+    unsafe {
+        launch_transpose_shd_to_hsd_bf16(
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
+        );
+    }
+    out
+}
+
+/// [1, KV_H, S, D] → [1, KV_H*n_rep, S, D] on GPU (BF16)
+pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
+    if n_rep == 1 { return x.clone(); }
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
+    let kv_heads = x.shape()[1];
+    let seq_len = x.shape()[2];
+    let head_dim = x.shape()[3];
+    let new_heads = kv_heads * n_rep;
+    let out = Tensor::empty(&[1, new_heads, seq_len, head_dim], DType::BF16, x.device());
+    unsafe {
+        launch_repeat_kv_bf16(
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            kv_heads as i32, n_rep as i32, seq_len as i32, head_dim as i32, std::ptr::null_mut(),
+        );
+    }
+    out
+}
+
+/// Make a non-contiguous GPU tensor contiguous via a strided copy kernel.
+/// Supports BF16 and F32, up to 4D tensors (padded to 4D internally).
+pub fn strided_to_contiguous_gpu(x: &Tensor) -> Tensor {
+    assert!(matches!(x.device(), Device::Cuda(_)), "expected GPU tensor");
+    assert!(!x.is_contiguous(), "tensor is already contiguous");
+    assert!(x.ndim() <= 4, "strided_to_contiguous_gpu supports up to 4D");
+
+    let ndim = x.ndim();
+    let numel = x.numel();
+
+    // Pad shape and strides to 4D (prepend 1s for shape, 0s for strides)
+    let mut shape4 = [1i32; 4];
+    let mut strides4 = [0i32; 4];
+    let pad = 4 - ndim;
+    for i in 0..ndim {
+        shape4[pad + i] = x.shape()[i] as i32;
+        strides4[pad + i] = x.strides()[i] as i32;
+    }
+
+    let out = Tensor::empty(x.shape(), x.dtype(), x.device());
+
+    // Use storage base pointer + element offset, because strides are relative to
+    // element 0 of the storage, not the data_ptr() (which already adds byte offset).
+    let storage_ptr = x.storage().gpu_buffer().as_ptr();
+    let in_offset = x.offset() as i32;
+
+    unsafe {
+        match x.dtype() {
+            DType::BF16 => launch_strided_copy_bf16(
+                storage_ptr as _, out.data_ptr() as *mut c_void,
+                numel as i32, ndim as i32,
+                shape4[0], shape4[1], shape4[2], shape4[3],
+                strides4[0], strides4[1], strides4[2], strides4[3],
+                in_offset, std::ptr::null_mut(),
+            ),
+            DType::F32 => launch_strided_copy_f32(
+                storage_ptr as _, out.data_ptr() as *mut c_void,
+                numel as i32, ndim as i32,
+                shape4[0], shape4[1], shape4[2], shape4[3],
+                strides4[0], strides4[1], strides4[2], strides4[3],
+                in_offset, std::ptr::null_mut(),
+            ),
+            _ => panic!("strided_to_contiguous_gpu: unsupported dtype {:?}", x.dtype()),
+        }
+    }
+    out
+}

crates/xserv-kernels/tests/attention_test.rs

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

crates/xserv-kernels/tests/gemm_test.rs

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

crates/xserv-kernels/tests/ops_test.rs

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

crates/xserv-model/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "xserv-model"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+xserv-cuda = { path = "../xserv-cuda" }
+xserv-tensor = { path = "../xserv-tensor" }
+xserv-kernels = { path = "../xserv-kernels" }
+xserv-tokenizer = { path = "../xserv-tokenizer" }
+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-gpt2.rs

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

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

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

crates/xserv-model/src/bin/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/dump-logits.rs

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

crates/xserv-model/src/bin/xserv-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/bin/xserv-cli.rs

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

crates/xserv-model/src/config.rs

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

crates/xserv-model/src/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

@@ -0,0 +1,380 @@
+use std::collections::HashMap;
+use xserv_kernels::*;
+use xserv_tensor::{DType, Device, Tensor};
+
+use crate::config::ModelConfig;
+
+pub struct GPT2 {
+    pub config: ModelConfig,
+    wte: Tensor,
+    wpe: Tensor,
+    layers: Vec<GPT2Block>,
+    ln_f_g: Tensor,
+    ln_f_b: Tensor,
+    lm_head: Tensor, // precomputed wte^T
+}
+
+struct GPT2Block {
+    ln_1_g: Tensor,
+    ln_1_b: Tensor,
+    attn_qkv_w: Tensor,
+    attn_qkv_b: Tensor,
+    attn_out_w: Tensor,
+    attn_out_b: Tensor,
+    ln_2_g: Tensor,
+    ln_2_b: Tensor,
+    mlp_fc_w: Tensor,
+    mlp_fc_b: Tensor,
+    mlp_proj_w: Tensor,
+    mlp_proj_b: Tensor,
+}
+
+pub struct KVCache {
+    // Per layer, per head: raw bytes (works for both f32 and bf16)
+    k: Vec<Vec<Vec<u8>>>,  // [num_layers][num_heads][seq_len * head_dim * elem_size]
+    v: Vec<Vec<Vec<u8>>>,
+    len: usize,
+    num_heads: usize,
+    head_dim: usize,
+    elem_size: usize,
+    dtype: DType,
+    device: Device,
+}
+
+impl KVCache {
+    pub fn new(num_layers: usize, num_heads: usize, head_dim: usize, dtype: DType, device: Device) -> Self {
+        Self {
+            k: (0..num_layers).map(|_| vec![vec![]; num_heads]).collect(),
+            v: (0..num_layers).map(|_| vec![vec![]; num_heads]).collect(),
+            len: 0,
+            num_heads,
+            head_dim,
+            elem_size: dtype.size_bytes(),
+            dtype,
+            device,
+        }
+    }
+
+    pub fn seq_len(&self) -> usize { self.len }
+
+    /// Append from a CPU tensor with shape [1, H, new_tokens, D].
+    pub fn append_kv_tensor(&mut self, layer: usize, k_cpu: &Tensor, v_cpu: &Tensor, new_tokens: usize) {
+        let hd = self.head_dim;
+        let es = self.elem_size;
+        let k_bytes = k_cpu.storage().as_cpu_bytes();
+        let v_bytes = v_cpu.storage().as_cpu_bytes();
+        let chunk = new_tokens * hd * es;
+        for h in 0..self.num_heads {
+            let off = h * chunk;
+            self.k[layer][h].extend_from_slice(&k_bytes[off..off + chunk]);
+            self.v[layer][h].extend_from_slice(&v_bytes[off..off + chunk]);
+        }
+        if layer == 0 {
+            self.len += new_tokens;
+        }
+    }
+
+    /// Reconstruct [1, H, seq_len, D] tensors.
+    pub fn get_kv_tensors(&self, layer: usize) -> (Tensor, Tensor) {
+        let sl = self.len;
+        let hd = self.head_dim;
+        let nh = self.num_heads;
+        let es = self.elem_size;
+        let head_bytes = sl * hd * es;
+        let total = nh * head_bytes;
+        let mut k_data = vec![0u8; total];
+        let mut v_data = vec![0u8; total];
+        for h in 0..nh {
+            let off = h * head_bytes;
+            k_data[off..off + head_bytes].copy_from_slice(&self.k[layer][h]);
+            v_data[off..off + head_bytes].copy_from_slice(&self.v[layer][h]);
+        }
+        let shape = &[1, nh, sl, hd];
+        let k = tensor_from_raw_bytes(&k_data, shape, self.dtype).to_device(self.device);
+        let v = tensor_from_raw_bytes(&v_data, shape, self.dtype).to_device(self.device);
+        (k, v)
+    }
+}
+
+fn tensor_from_raw_bytes(bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
+    match dtype {
+        DType::F32 => {
+            let data: &[f32] = unsafe {
+                std::slice::from_raw_parts(bytes.as_ptr() as *const f32, bytes.len() / 4)
+            };
+            Tensor::from_slice(data, shape)
+        }
+        DType::BF16 => {
+            let data: &[half::bf16] = unsafe {
+                std::slice::from_raw_parts(bytes.as_ptr() as *const half::bf16, bytes.len() / 2)
+            };
+            Tensor::from_slice(data, shape)
+        }
+        _ => panic!("unsupported dtype for KV cache"),
+    }
+}
+
+impl GPT2 {
+    pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
+        crate::init_kernels();
+        let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
+            w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
+        };
+
+        let wte = take(&mut w, "wte.weight");
+        let wpe = take(&mut w, "wpe.weight");
+        let ln_f_g = take(&mut w, "ln_f.weight");
+        let ln_f_b = take(&mut w, "ln_f.bias");
+        let lm_head = wte.transpose(0, 1).contiguous();
+
+        let num_layers = config.num_layers();
+        let mut layers = Vec::with_capacity(num_layers);
+        for i in 0..num_layers {
+            let p = format!("h.{i}");
+            layers.push(GPT2Block {
+                ln_1_g: take(&mut w, &format!("{p}.ln_1.weight")),
+                ln_1_b: take(&mut w, &format!("{p}.ln_1.bias")),
+                attn_qkv_w: take(&mut w, &format!("{p}.attn.c_attn.weight")),
+                attn_qkv_b: take(&mut w, &format!("{p}.attn.c_attn.bias")),
+                attn_out_w: take(&mut w, &format!("{p}.attn.c_proj.weight")),
+                attn_out_b: take(&mut w, &format!("{p}.attn.c_proj.bias")),
+                ln_2_g: take(&mut w, &format!("{p}.ln_2.weight")),
+                ln_2_b: take(&mut w, &format!("{p}.ln_2.bias")),
+                mlp_fc_w: take(&mut w, &format!("{p}.mlp.c_fc.weight")),
+                mlp_fc_b: take(&mut w, &format!("{p}.mlp.c_fc.bias")),
+                mlp_proj_w: take(&mut w, &format!("{p}.mlp.c_proj.weight")),
+                mlp_proj_b: take(&mut w, &format!("{p}.mlp.c_proj.bias")),
+            });
+        }
+
+        Self { config, wte, wpe, layers, ln_f_g, ln_f_b, lm_head }
+    }
+
+    /// Full forward pass without KV cache (for testing / correctness comparison).
+    pub fn forward(&self, token_ids: &[u32]) -> Tensor {
+        let seq_len = token_ids.len();
+        let hidden = self.config.hidden();
+        let num_heads = self.config.num_heads();
+        let head_dim = self.config.head_dim();
+
+        let tok_emb = embedding(&self.wte, token_ids);
+        let pos_ids: Vec<u32> = (0..seq_len as u32).collect();
+        let pos_emb = embedding(&self.wpe, &pos_ids);
+        let mut x = add_tensors(&tok_emb, &pos_emb);
+
+        for layer in &self.layers {
+            x = self.transformer_block(layer, &x, None, 0, seq_len, num_heads, head_dim, hidden);
+        }
+
+        let x = layernorm(&x, &self.ln_f_g, &self.ln_f_b, self.config.ln_eps());
+        matmul_2d(&x, &self.lm_head)
+    }
+
+    /// Forward pass with KV cache. First call = prefill, subsequent = decode.
+    pub fn forward_with_cache(&self, token_ids: &[u32], cache: &mut KVCache) -> Tensor {
+        let new_tokens = token_ids.len();
+        let pos_offset = cache.seq_len();
+        let hidden = self.config.hidden();
+        let num_heads = self.config.num_heads();
+        let head_dim = self.config.head_dim();
+
+        let tok_emb = embedding(&self.wte, token_ids);
+        let pos_ids: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
+        let pos_emb = embedding(&self.wpe, &pos_ids);
+        let mut x = add_tensors(&tok_emb, &pos_emb);
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            x = self.transformer_block(
+                layer, &x, Some((cache, layer_idx)),
+                pos_offset, new_tokens, num_heads, head_dim, hidden,
+            );
+        }
+
+        let x = layernorm(&x, &self.ln_f_g, &self.ln_f_b, self.config.ln_eps());
+        matmul_2d(&x, &self.lm_head)
+    }
+
+    fn transformer_block(
+        &self,
+        layer: &GPT2Block,
+        x: &Tensor,
+        cache: Option<(&mut KVCache, usize)>,
+        pos_offset: usize,
+        new_tokens: usize,
+        num_heads: usize,
+        head_dim: usize,
+        hidden: usize,
+    ) -> Tensor {
+        let residual = x.clone();
+        let normed = layernorm(x, &layer.ln_1_g, &layer.ln_1_b, self.config.ln_eps());
+
+        let qkv = linear(&normed, &layer.attn_qkv_w, Some(&layer.attn_qkv_b));
+        let (q, k_new, v_new) = split_qkv(&qkv, num_heads, head_dim, new_tokens);
+
+        let (k_full, v_full) = if let Some((cache, layer_idx)) = cache {
+            let k_cpu = k_new.to_device(Device::Cpu);
+            let v_cpu = v_new.to_device(Device::Cpu);
+            cache.append_kv_tensor(layer_idx, &k_cpu, &v_cpu, new_tokens);
+            cache.get_kv_tensors(layer_idx)
+        } else {
+            (k_new, v_new)
+        };
+
+        let attn_out = attention(&q, &k_full, &v_full, true);
+        let attn_out = merge_heads(&attn_out, new_tokens, hidden);
+        let attn_out = linear(&attn_out, &layer.attn_out_w, Some(&layer.attn_out_b));
+        let x = add_tensors(&residual, &attn_out);
+
+        let residual = x.clone();
+        let normed = layernorm(&x, &layer.ln_2_g, &layer.ln_2_b, self.config.ln_eps());
+        let fc = linear(&normed, &layer.mlp_fc_w, Some(&layer.mlp_fc_b));
+        let activated = gelu(&fc);
+        let proj = linear(&activated, &layer.mlp_proj_w, Some(&layer.mlp_proj_b));
+        add_tensors(&residual, &proj)
+    }
+}
+
+// --- Helper ops (unchanged) ---
+
+fn linear(x: &Tensor, weight: &Tensor, bias: Option<&Tensor>) -> Tensor {
+    let out = matmul_2d(x, weight);
+    if let Some(b) = bias { add_bias(&out, b) } else { out }
+}
+
+fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
+    assert_eq!(a.ndim(), 2);
+    assert_eq!(b.ndim(), 2);
+    matmul(a, b, GemmBackend::CuBlas)
+}
+
+fn add_tensors(a: &Tensor, b: &Tensor) -> Tensor {
+    xserv_kernels::add(a, b)
+}
+
+fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
+    // bias: [N], x: [S, N] — broadcast add via reshape
+    assert_eq!(x.ndim(), 2);
+    assert_eq!(bias.ndim(), 1);
+    let n = bias.shape()[0];
+    assert_eq!(x.shape()[1], n);
+    let rows = x.shape()[0];
+    // Broadcast: tile bias to [S, N] on CPU, then GPU add
+    let b_cpu = bias.to_device(Device::Cpu);
+    match x.dtype() {
+        DType::F32 => {
+            let bd = b_cpu.as_slice::<f32>();
+            let tiled: Vec<f32> = (0..rows).flat_map(|_| bd.iter().copied()).collect();
+            let b_full = Tensor::from_slice(&tiled, x.shape()).to_device(x.device());
+            xserv_kernels::add(x, &b_full)
+        }
+        DType::BF16 => {
+            let bd = b_cpu.as_slice::<half::bf16>();
+            let tiled: Vec<half::bf16> = (0..rows).flat_map(|_| bd.iter().copied()).collect();
+            let b_full = Tensor::from_slice(&tiled, x.shape()).to_device(x.device());
+            xserv_kernels::add(x, &b_full)
+        }
+        _ => panic!("unsupported dtype"),
+    }
+}
+
+fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) -> (Tensor, Tensor, Tensor) {
+    let hidden = num_heads * head_dim;
+    let qkv_cpu = qkv.to_device(Device::Cpu);
+    let 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 {
+                    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)
+        }
+        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 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 {
+                    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)
+        }
+        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.
+pub fn sample_greedy(logits: &Tensor) -> u32 {
+    assert_eq!(logits.ndim(), 2);
+    let logits_cpu = logits.to_device(Device::Cpu);
+    let data = logits_cpu.as_slice::<f32>();
+    let vocab_size = logits.shape()[1];
+    let seq_len = logits.shape()[0];
+    let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
+    last_row.iter()
+        .enumerate()
+        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
+        .map(|(idx, _)| idx as u32)
+        .unwrap()
+}

crates/xserv-model/src/kv_cache.rs

@@ -0,0 +1,151 @@
+use xserv_cuda::GpuBuffer;
+use xserv_tensor::{DType, Device, Tensor};
+use crate::config::ModelConfig;
+
+/// GPU-resident KV cache. Pre-allocates max_seq_len on GPU,
+/// appends new K/V via D2D copy at offset (no CPU round-trip).
+pub struct GpuKVCache {
+    // Per layer: contiguous GPU buffer for K and V
+    // Layout: [num_kv_heads, max_seq_len, head_dim] — contiguous per head
+    k_bufs: Vec<GpuBuffer>,
+    v_bufs: Vec<GpuBuffer>,
+    // Per layer: pre-allocated staging buffers for get_kv_len output.
+    // Size: num_kv_heads * max_seq_len * head_dim * elem_size (max possible output).
+    // Avoids cudaMalloc/cudaFree on every get_kv_len call.
+    k_staging: Vec<GpuBuffer>,
+    v_staging: Vec<GpuBuffer>,
+    seq_len: usize,
+    max_seq_len: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    elem_size: usize,
+    dtype: DType,
+    device: u32,
+}
+
+impl GpuKVCache {
+    pub fn new(config: &ModelConfig, max_seq_len: usize, dtype: DType, device: u32) -> Self {
+        let num_layers = config.num_layers();
+        let num_kv_heads = config.num_kv_heads();
+        let head_dim = config.head_dim();
+        let elem_size = dtype.size_bytes();
+        let buf_size = num_kv_heads * max_seq_len * head_dim * elem_size;
+
+        let mut k_bufs = Vec::with_capacity(num_layers);
+        let mut v_bufs = Vec::with_capacity(num_layers);
+        let mut k_staging = Vec::with_capacity(num_layers);
+        let mut v_staging = Vec::with_capacity(num_layers);
+        for _ in 0..num_layers {
+            let mut k = GpuBuffer::alloc(buf_size).expect("alloc KV cache K");
+            let mut v = GpuBuffer::alloc(buf_size).expect("alloc KV cache V");
+            k.zero().unwrap();
+            v.zero().unwrap();
+            k_bufs.push(k);
+            v_bufs.push(v);
+            k_staging.push(GpuBuffer::alloc(buf_size).expect("alloc KV staging K"));
+            v_staging.push(GpuBuffer::alloc(buf_size).expect("alloc KV staging V"));
+        }
+
+        Self { k_bufs, v_bufs, k_staging, v_staging, seq_len: 0, max_seq_len, num_kv_heads, head_dim, elem_size, dtype, device }
+    }
+
+    pub fn seq_len(&self) -> usize { self.seq_len }
+    pub fn max_seq_len(&self) -> usize { self.max_seq_len }
+
+    /// Append new K/V tensors for a given layer.
+    /// k_new, v_new: [1, num_kv_heads, new_tokens, head_dim] on GPU, contiguous.
+    /// `write_pos` is the sequence position to write at (caller manages this).
+    pub fn append(&mut self, layer: usize, k_new: &Tensor, v_new: &Tensor, new_tokens: usize, write_pos: usize) {
+        assert!(write_pos + new_tokens <= self.max_seq_len, "KV cache overflow");
+        let es = self.elem_size;
+        let hd = self.head_dim;
+        let max_s = self.max_seq_len;
+        let nh = self.num_kv_heads;
+
+        let k_src = k_new.storage().gpu_buffer();
+        let v_src = v_new.storage().gpu_buffer();
+
+        for h in 0..nh {
+            let src_off = h * new_tokens * hd * es;
+            let dst_off = (h * max_s + write_pos) * hd * es;
+            let count = new_tokens * hd * es;
+            self.k_bufs[layer].copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
+            self.v_bufs[layer].copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
+        }
+    }
+
+    pub fn advance_seq_len(&mut self, new_tokens: usize) {
+        self.seq_len += new_tokens;
+        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(&mut self, layer: usize) -> (Tensor, Tensor) {
+        let sl = self.seq_len;
+        self.get_kv_len(layer, sl)
+    }
+
+    pub fn get_kv_len(&mut self, layer: usize, sl: usize) -> (Tensor, Tensor) {
+        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;
+
+        // Copy each head's valid portion into pre-allocated staging buffers.
+        // Split borrows: staging (mut) vs cache (shared) are separate struct fields,
+        // so the borrow checker allows simultaneous &mut staging + &cache.
+        let out_size = nh * sl * hd * es;
+        let k_stg = &mut self.k_staging[layer];
+        let k_buf = &self.k_bufs[layer];
+        let v_stg = &mut self.v_staging[layer];
+        let v_buf = &self.v_bufs[layer];
+        for h in 0..nh {
+            let src_off = (h * max_s) * hd * es;
+            let dst_off = (h * sl) * hd * es;
+            let count = sl * hd * es;
+            k_stg.copy_from_device_at(k_buf, src_off, dst_off, count).unwrap();
+            v_stg.copy_from_device_at(v_buf, src_off, dst_off, count).unwrap();
+        }
+        // Grab raw pointers before dropping the mutable borrows
+        let k_ptr = k_stg.as_mut_ptr();
+        let v_ptr = v_stg.as_mut_ptr();
+
+        // Create Tensors that borrow from the staging buffers (no cudaMalloc/cudaFree).
+        // Safety: staging buffers are owned by GpuKVCache and outlive the returned Tensors
+        // in practice (Tensors are consumed within the same forward pass before the next
+        // get_kv_len call overwrites the staging buffer).
+        let shape = &[1usize, nh, sl, hd];
+        let k = unsafe {
+            tensor_from_gpu_buffer(GpuBuffer::borrow_raw(k_ptr, out_size), shape, self.dtype, self.device)
+        };
+        let v = unsafe {
+            tensor_from_gpu_buffer(GpuBuffer::borrow_raw(v_ptr, out_size), shape, self.dtype, self.device)
+        };
+        (k, v)
+    }
+}
+
+/// Create a Tensor from a GpuBuffer (takes ownership).
+unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
+    use xserv_tensor::storage::Storage;
+    use xserv_tensor::shape::contiguous_strides;
+    use smallvec::SmallVec;
+
+    let storage = Storage::cuda(buf, device);
+    Tensor::from_storage(
+        storage,
+        SmallVec::from_slice(shape),
+        contiguous_strides(shape),
+        0,
+        dtype,
+    )
+}
+
+/// Public version for use by other modules (e.g., batched decode concat).
+///
+/// # Safety
+/// `buf` must be a valid GPU allocation with at least `product(shape) * dtype.size_bytes()` bytes.
+pub unsafe fn tensor_from_gpu_buffer_pub(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
+    tensor_from_gpu_buffer(buf, shape, dtype, device)
+}

crates/xserv-model/src/lib.rs

@@ -0,0 +1,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/loader.rs

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

crates/xserv-model/src/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

@@ -0,0 +1,811 @@
+use std::collections::HashMap;
+use half::bf16;
+use xserv_kernels::*;
+use xserv_tensor::{DType, Device, Tensor};
+
+use crate::config::ModelConfig;
+use crate::gpt2::KVCache;
+use crate::kv_cache::GpuKVCache;
+use crate::paged_kv_cache::PagedKVCache;
+
+pub struct Qwen3 {
+    pub config: ModelConfig,
+    embed_tokens: Tensor,
+    layers: Vec<Qwen3Block>,
+    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 {
+    input_norm: Tensor,       // [hidden]
+    q_proj_wt: Tensor,        // TRANSPOSED: [hidden, num_heads*head_dim]
+    k_proj_wt: Tensor,        // TRANSPOSED: [hidden, num_kv_heads*head_dim]
+    v_proj_wt: Tensor,
+    o_proj_wt: Tensor,        // TRANSPOSED: [num_heads*head_dim, hidden]
+    q_norm: Tensor,           // [head_dim]
+    k_norm: Tensor,           // [head_dim]
+    post_norm: Tensor,        // [hidden]
+    gate_proj_wt: Tensor,     // TRANSPOSED: [hidden, intermediate]
+    up_proj_wt: Tensor,
+    down_proj_wt: Tensor,     // TRANSPOSED: [intermediate, hidden]
+}
+
+impl Qwen3 {
+    /// 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 = repl(take(&mut w, "model.embed_tokens.weight"));
+        let norm = repl(take(&mut w, "model.norm.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(),
+            config.head_dim(),
+            config.rope_theta.unwrap_or(1_000_000.0) as f32,
+        );
+
+        let num_layers = config.num_layers();
+        let mut layers = Vec::with_capacity(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: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))),
+                q_proj_wt: col(take(&mut w, &format!("{p}.self_attn.q_proj.weight"))),
+                k_proj_wt: col(take(&mut w, &format!("{p}.self_attn.k_proj.weight"))),
+                v_proj_wt: col(take(&mut w, &format!("{p}.self_attn.v_proj.weight"))),
+                o_proj_wt: row(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
+                q_norm: repl(take(&mut w, &format!("{p}.self_attn.q_norm.weight"))),
+                k_norm: repl(take(&mut w, &format!("{p}.self_attn.k_norm.weight"))),
+                post_norm: repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight"))),
+                gate_proj_wt: col(take(&mut w, &format!("{p}.mlp.gate_proj.weight"))),
+                up_proj_wt: col(take(&mut w, &format!("{p}.mlp.up_proj.weight"))),
+                down_proj_wt: row(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
+            });
+        }
+
+        Self {
+            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 {
+        let new_tokens = token_ids.len();
+        let pos_offset = cache.seq_len();
+        let hidden = self.config.hidden();
+        let num_heads = self.config.num_heads();
+        let num_kv_heads = self.config.num_kv_heads();
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        let mut x = embedding(&self.embed_tokens, token_ids);
+        let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps);
+
+            // Q/K/V projections (pre-transposed weights, x @ wt)
+            let q = matmul_2d(&normed, &layer.q_proj_wt);
+            let k = matmul_2d(&normed, &layer.k_proj_wt);
+            let v = matmul_2d(&normed, &layer.v_proj_wt);
+
+            // Reshape to [1, heads, seq, head_dim]
+            let q = reshape_heads(&q, new_tokens, num_heads, head_dim);
+            let k = reshape_heads(&k, new_tokens, num_kv_heads, head_dim);
+            let v = reshape_heads(&v, new_tokens, num_kv_heads, head_dim);
+
+            // QK normalization (per-head RMSNorm)
+            let q = head_rmsnorm(&q, &layer.q_norm, eps);
+            let k = head_rmsnorm(&k, &layer.k_norm, eps);
+
+            // RoPE — kernel expects [S, H, D], our tensors are [1, H, S, D]
+            // Transpose to [1, S, H, D] → reshape to [S, H, D] for RoPE
+            let q = transpose_for_rope(&q, new_tokens, num_heads, head_dim);
+            let k = transpose_for_rope(&k, new_tokens, num_kv_heads, head_dim);
+            rope_inplace(&q, &self.rope_cache, &positions);
+            rope_inplace(&k, &self.rope_cache, &positions);
+            // Transpose back to [1, H, S, D]
+            let q = transpose_from_rope(&q, new_tokens, num_heads, head_dim);
+            let k = transpose_from_rope(&k, new_tokens, num_kv_heads, head_dim);
+
+            // KV cache
+            let k_cpu = k.to_device(Device::Cpu);
+            let v_cpu = v.to_device(Device::Cpu);
+            cache.append_kv_tensor(layer_idx, &k_cpu, &v_cpu, new_tokens);
+            let (k_full, v_full) = cache.get_kv_tensors(layer_idx);
+
+            // GQA: repeat K/V
+            let n_rep = num_heads / num_kv_heads;
+            let k_full = repeat_kv(&k_full, n_rep);
+            let v_full = repeat_kv(&v_full, n_rep);
+
+            // Attention
+            let attn_out = attention(&q, &k_full, &v_full, true);
+            let attn_merged = merge_heads_any(&attn_out, new_tokens, hidden);
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+            x = add_any(&residual, &attn_proj);
+
+            // SwiGLU FFN
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.post_norm, eps);
+            let gate = matmul_2d(&normed, &layer.gate_proj_wt);
+            let up = matmul_2d(&normed, &layer.up_proj_wt);
+            let gate_activated = silu(&gate);
+            let hidden_states = mul_any(&gate_activated, &up);
+            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            x = add_any(&residual, &down);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_2d(&x, &self.lm_head_t)
+    }
+
+    /// Batched decode: process one token per sequence simultaneously.
+    /// All compute-heavy ops (projections, FFN) operate on [B, hidden] tensors.
+    /// Per-sequence ops (RoPE, KV cache, attention) are handled individually.
+    ///
+    /// tokens: one token per sequence (len = batch_size)
+    /// positions: position offset for each sequence (len = batch_size)
+    /// caches: one mutable KV cache per sequence (len = batch_size)
+    ///
+    /// Returns logits: [batch_size, vocab_size]
+    pub fn forward_decode_batch(
+        &self,
+        tokens: &[u32],
+        positions: &[usize],
+        caches: &mut [&mut GpuKVCache],
+    ) -> Tensor {
+        let batch = tokens.len();
+        assert_eq!(positions.len(), batch);
+        assert_eq!(caches.len(), batch);
+        assert!(batch > 0);
+
+        let num_heads = self.config.num_heads();
+        let num_kv_heads = self.config.num_kv_heads();
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        // Batched embedding: [B, hidden]
+        let mut x = embedding(&self.embed_tokens, tokens);
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps); // [B, hidden]
+
+            // Batched projections: [B, hidden] × [hidden, X] = [B, X]
+            let q_all = matmul_2d(&normed, &layer.q_proj_wt); // [B, num_heads*head_dim]
+            let k_all = matmul_2d(&normed, &layer.k_proj_wt); // [B, num_kv_heads*head_dim]
+            let v_all = matmul_2d(&normed, &layer.v_proj_wt); // [B, num_kv_heads*head_dim]
+
+            // Per-sequence: reshape, qk-norm, RoPE, KV cache, attention, merge
+            let mut attn_outputs: Vec<Tensor> = Vec::with_capacity(batch);
+            for b in 0..batch {
+                // Extract row b: [1, X] — view into contiguous [B, X]
+                let q_row = row_view(&q_all, b); // [1, num_heads*head_dim]
+                let k_row = row_view(&k_all, b); // [1, num_kv_heads*head_dim]
+                let v_row = row_view(&v_all, b); // [1, num_kv_heads*head_dim]
+
+                // GPU reshape: [1, H*D] → [1, H, 1, D]
+                let q = xserv_kernels::reshape_heads_gpu(&q_row, 1, num_heads, head_dim);
+                let k = xserv_kernels::reshape_heads_gpu(&k_row, 1, num_kv_heads, head_dim);
+                let v = xserv_kernels::reshape_heads_gpu(&v_row, 1, num_kv_heads, head_dim);
+
+                // QK norm
+                let q = head_rmsnorm(&q, &layer.q_norm, eps);
+                let k = head_rmsnorm(&k, &layer.k_norm, eps);
+
+                // GPU transpose for RoPE: [1, H, 1, D] → [1, H, D]
+                let q = xserv_kernels::transpose_for_rope_gpu(&q, 1, num_heads, head_dim);
+                let k = xserv_kernels::transpose_for_rope_gpu(&k, 1, num_kv_heads, head_dim);
+
+                // RoPE with per-sequence position
+                let pos = [positions[b] as u32];
+                rope_inplace(&q, &self.rope_cache, &pos);
+                rope_inplace(&k, &self.rope_cache, &pos);
+
+                // Transpose back: [1, H, D] → [1, H, 1, D]
+                let q = xserv_kernels::transpose_from_rope_gpu(&q, 1, num_heads, head_dim);
+                let k = xserv_kernels::transpose_from_rope_gpu(&k, 1, num_kv_heads, head_dim);
+
+                // KV cache: append and get full cache
+                let pos_b = positions[b];
+                caches[b].append(layer_idx, &k, &v, 1, pos_b);
+                let (k_full, v_full) = caches[b].get_kv_len(layer_idx, pos_b + 1);
+
+                // Decode attention (uses native GQA, no repeat_kv needed)
+                let attn_out = flash_attention(&q, &k_full, &v_full, true);
+
+                // Merge heads: [1, H, 1, D] → [1, hidden]
+                let merged = xserv_kernels::merge_heads_gpu(&attn_out, 1, num_heads, head_dim);
+                attn_outputs.push(merged);
+            }
+
+            // Concat attention outputs: [B, hidden]
+            let attn_merged = concat_rows(&attn_outputs);
+
+            // Batched O projection: [B, hidden] × [hidden, hidden] = [B, hidden]
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+
+            // Fused add + rmsnorm
+            let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            // Batched FFN: all projections on [B, hidden]
+            let gate = matmul_2d(&normed, &layer.gate_proj_wt);
+            let up = matmul_2d(&normed, &layer.up_proj_wt);
+            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
+            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            x = add_any(&residual, &down);
+        }
+
+        // Advance KV cache seq_len for each sequence
+        for b in 0..batch {
+            caches[b].advance_seq_len(1);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_2d(&x, &self.lm_head_t) // [B, vocab_size]
+    }
+
+    /// 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();
+        let pos_offset = cache.seq_len();
+        let hidden = self.config.hidden();
+        let num_heads = self.config.num_heads();
+        let num_kv_heads = self.config.num_kv_heads();
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        let mut x = embedding(&self.embed_tokens, token_ids);
+        let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps);
+
+            let q = matmul_2d(&normed, &layer.q_proj_wt);
+            let k = matmul_2d(&normed, &layer.k_proj_wt);
+            let v = matmul_2d(&normed, &layer.v_proj_wt);
+
+            // GPU reshape: [S, H*D] → [1, H, S, D]
+            let q = xserv_kernels::reshape_heads_gpu(&q, new_tokens, num_heads, head_dim);
+            let k = xserv_kernels::reshape_heads_gpu(&k, new_tokens, num_kv_heads, head_dim);
+            let v = xserv_kernels::reshape_heads_gpu(&v, new_tokens, num_kv_heads, head_dim);
+
+            // QK norm (reshape to [H*S, D], rmsnorm, reshape back — stays on GPU)
+            let q = head_rmsnorm(&q, &layer.q_norm, eps);
+            let k = head_rmsnorm(&k, &layer.k_norm, eps);
+
+            // GPU transpose for RoPE: [1, H, S, D] → [S, H, D]
+            let q = xserv_kernels::transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim);
+            let k = xserv_kernels::transpose_for_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
+            rope_inplace(&q, &self.rope_cache, &positions);
+            rope_inplace(&k, &self.rope_cache, &positions);
+            // GPU transpose back: [S, H, D] → [1, H, S, D]
+            let q = xserv_kernels::transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim);
+            let k = xserv_kernels::transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
+
+            // GPU KV cache
+            cache.append(layer_idx, &k, &v, new_tokens, pos_offset);
+            let (k_full, v_full) = cache.get_kv_len(layer_idx, pos_offset + new_tokens);
+
+            // Flash Attention with native GQA (no repeat_kv needed)
+            let attn_out = flash_attention(&q, &k_full, &v_full, true);
+            // GPU merge_heads: [1, H, S, D] → [S, H*D]
+            let attn_merged = xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+
+            // Fused add + rmsnorm: (normed, x) where x = residual + attn_proj
+            let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            // Fused SiLU×Mul
+            let gate = matmul_2d(&normed, &layer.gate_proj_wt);
+            let up = matmul_2d(&normed, &layer.up_proj_wt);
+            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
+            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            x = add_any(&residual, &down);
+        }
+
+        cache.advance_seq_len(new_tokens);
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_2d(&x, &self.lm_head_t)
+    }
+
+    /// 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);
+    matmul(a, b, GemmBackend::CuBlas)
+}
+
+fn reshape_heads(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    let x_cpu = x.to_device(Device::Cpu);
+    let hidden = num_heads * head_dim;
+    let src = x_cpu.as_slice::<bf16>();
+    let mut out = vec![bf16::ZERO; num_heads * seq_len * head_dim];
+    for s in 0..seq_len {
+        for h in 0..num_heads {
+            let si = s * hidden + h * head_dim;
+            let di = (h * seq_len + s) * head_dim;
+            out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
+        }
+    }
+    Tensor::from_slice(&out, &[1, num_heads, seq_len, head_dim]).to_device(x.device())
+}
+
+fn merge_heads_any(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
+    let num_heads = x.shape()[1];
+    let head_dim = x.shape()[3];
+    let x_cpu = x.to_device(Device::Cpu);
+    let src = x_cpu.as_slice::<bf16>();
+    let mut out = vec![bf16::ZERO; seq_len * hidden];
+    for s in 0..seq_len {
+        for h in 0..num_heads {
+            let si = (h * seq_len + s) * head_dim;
+            let di = s * hidden + h * head_dim;
+            out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
+        }
+    }
+    Tensor::from_slice(&out, &[seq_len, hidden]).to_device(x.device())
+}
+
+/// Per-head RMSNorm: apply RMSNorm to each [head_dim] slice independently.
+/// x: [1, H, S, D], norm_weight: [D]
+fn head_rmsnorm(x: &Tensor, norm_weight: &Tensor, eps: f32) -> Tensor {
+    let num_heads = x.shape()[1];
+    let seq_len = x.shape()[2];
+    let head_dim = x.shape()[3];
+    // Reshape to [H*S, D], apply rmsnorm, reshape back
+    let total_rows = num_heads * seq_len;
+    let flat = x.reshape(&[total_rows, head_dim]);
+    let normed = rmsnorm(&flat, norm_weight, eps);
+    normed.reshape(&[1, num_heads, seq_len, head_dim])
+}
+
+/// [1, H, S, D] → [S, H, D] for RoPE kernel
+fn transpose_for_rope(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    let x_cpu = x.to_device(Device::Cpu);
+    let src = x_cpu.as_slice::<bf16>();
+    let mut out = vec![bf16::ZERO; seq_len * num_heads * head_dim];
+    for h in 0..num_heads {
+        for s in 0..seq_len {
+            let si = (h * seq_len + s) * head_dim;
+            let di = (s * num_heads + h) * head_dim;
+            out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
+        }
+    }
+    Tensor::from_slice(&out, &[seq_len, num_heads, head_dim]).to_device(x.device())
+}
+
+/// [S, H, D] → [1, H, S, D] after RoPE
+fn transpose_from_rope(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
+    let x_cpu = x.to_device(Device::Cpu);
+    let src = x_cpu.as_slice::<bf16>();
+    let mut out = vec![bf16::ZERO; num_heads * seq_len * head_dim];
+    for s in 0..seq_len {
+        for h in 0..num_heads {
+            let si = (s * num_heads + h) * head_dim;
+            let di = (h * seq_len + s) * head_dim;
+            out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
+        }
+    }
+    Tensor::from_slice(&out, &[1, num_heads, seq_len, head_dim]).to_device(x.device())
+}
+
+fn repeat_kv(x: &Tensor, n_rep: usize) -> Tensor {
+    if n_rep == 1 { return x.clone(); }
+    let kv_heads = x.shape()[1];
+    let seq_len = x.shape()[2];
+    let head_dim = x.shape()[3];
+    let x_cpu = x.to_device(Device::Cpu);
+    let src = x_cpu.as_slice::<bf16>();
+    let new_heads = kv_heads * n_rep;
+    let mut out = vec![bf16::ZERO; new_heads * seq_len * head_dim];
+    let chunk = seq_len * head_dim;
+    for kv_h in 0..kv_heads {
+        for r in 0..n_rep {
+            let dst_h = kv_h * n_rep + r;
+            out[dst_h * chunk..(dst_h + 1) * chunk]
+                .copy_from_slice(&src[kv_h * chunk..(kv_h + 1) * chunk]);
+        }
+    }
+    Tensor::from_slice(&out, &[1, new_heads, seq_len, head_dim]).to_device(x.device())
+}
+
+/// Extract row `b` from a contiguous 2D tensor [B, cols] as a [1, cols] view.
+/// Zero-copy: shares storage with the original tensor.
+fn row_view(t: &Tensor, row: usize) -> Tensor {
+    assert_eq!(t.ndim(), 2);
+    assert!(t.is_contiguous());
+    let cols = t.shape()[1];
+    assert!(row < t.shape()[0]);
+    let new_offset = t.offset() + row * cols;
+    Tensor::from_storage(
+        t.storage().clone(),
+        smallvec::SmallVec::from_slice(&[1, cols]),
+        xserv_tensor::shape::contiguous_strides(&[1, cols]),
+        new_offset,
+        t.dtype(),
+    )
+}
+
+/// Concatenate row tensors [1, cols] into a single [B, cols] tensor via D2D memcpy.
+fn concat_rows(rows: &[Tensor]) -> Tensor {
+    assert!(!rows.is_empty());
+    let batch = rows.len();
+    let cols = rows[0].shape()[1];
+    let dtype = rows[0].dtype();
+    let device = rows[0].device();
+    let elem_size = dtype.size_bytes();
+    let row_bytes = cols * elem_size;
+
+    // Allocate output [B, cols] and copy each row into it
+    let total_bytes = batch * row_bytes;
+    let mut out_buf = xserv_cuda::allocator::cached_alloc(total_bytes).expect("alloc concat_rows");
+
+    for (b, row) in rows.iter().enumerate() {
+        assert_eq!(row.shape(), &[1, cols]);
+        assert!(row.is_contiguous());
+        let src_buf = row.storage().gpu_buffer();
+        let src_offset = row.offset() * elem_size;
+        let dst_offset = b * row_bytes;
+        out_buf.copy_from_device_at(src_buf, src_offset, dst_offset, row_bytes).unwrap();
+    }
+
+    // Wrap in a Tensor
+    let device_id = match device { Device::Cuda(id) => id, _ => panic!("expected CUDA device") };
+    unsafe {
+        crate::kv_cache::tensor_from_gpu_buffer_pub(out_buf, &[batch, cols], dtype, device_id)
+    }
+}
+
+fn add_any(a: &Tensor, b: &Tensor) -> Tensor {
+    xserv_kernels::add(a, b)
+}
+
+fn mul_any(a: &Tensor, b: &Tensor) -> Tensor {
+    xserv_kernels::mul(a, b)
+}
+
+pub fn sample_greedy(logits: &Tensor) -> u32 {
+    assert_eq!(logits.ndim(), 2);
+    let logits_cpu = logits.to_device(Device::Cpu);
+    let vocab_size = logits.shape()[1];
+    let seq_len = logits.shape()[0];
+    let data = logits_cpu.as_slice::<bf16>();
+    let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
+    last.iter().enumerate()
+        .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
+        .map(|(i, _)| i as u32).unwrap()
+}

crates/xserv-model/src/sampling.rs

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

crates/xserv-server/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "xserv-server"
+version.workspace = true
+edition.workspace = true
+
+[[bin]]
+name = "xserv-server"
+path = "src/main.rs"
+
+[dependencies]
+xserv-cuda = { path = "../xserv-cuda" }
+xserv-tensor = { path = "../xserv-tensor" }
+xserv-kernels = { path = "../xserv-kernels" }
+xserv-model = { path = "../xserv-model" }
+xserv-tokenizer = { path = "../xserv-tokenizer" }
+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

@@ -0,0 +1,345 @@
+use axum::Extension;
+use axum::Json;
+use axum::http::StatusCode;
+use axum::response::sse::{Event, KeepAlive, Sse};
+use axum::response::{IntoResponse, Response};
+use serde::{Deserialize, Serialize};
+use std::convert::Infallible;
+use std::sync::Arc;
+use tokio_stream::StreamExt;
+use tokio_stream::wrappers::ReceiverStream;
+use uuid::Uuid;
+
+use crate::AppState;
+use crate::engine::{GenerateEvent, GenerateRequest};
+use xserv_model::SamplingParams;
+
+#[derive(Deserialize)]
+pub struct ChatRequest {
+    #[serde(default)]
+    pub model: Option<String>,
+    pub messages: Vec<Message>,
+    #[serde(default = "default_max_tokens")]
+    pub max_tokens: usize,
+    #[serde(default)]
+    pub stream: Option<bool>,
+    #[serde(default)]
+    pub temperature: Option<f32>,
+    #[serde(default)]
+    pub top_k: Option<usize>,
+    #[serde(default)]
+    pub top_p: Option<f32>,
+}
+
+#[derive(Deserialize)]
+pub struct Message {
+    pub role: String,
+    pub content: String,
+}
+
+fn default_max_tokens() -> usize {
+    256
+}
+
+#[derive(Serialize)]
+pub struct ModelsResponse {
+    object: &'static str,
+    data: Vec<ModelInfo>,
+}
+
+#[derive(Serialize)]
+pub struct ModelInfo {
+    id: String,
+    object: &'static str,
+    owned_by: &'static str,
+}
+
+pub async fn health() -> &'static str {
+    "ok"
+}
+
+pub async fn list_models(Extension(state): Extension<Arc<AppState>>) -> Json<ModelsResponse> {
+    Json(ModelsResponse {
+        object: "list",
+        data: vec![ModelInfo {
+            id: state.model_name.clone(),
+            object: "model",
+            owned_by: "xserv",
+        }],
+    })
+}
+
+pub async fn chat_completions(
+    Extension(state): Extension<Arc<AppState>>,
+    Json(req): Json<ChatRequest>,
+) -> Response {
+    if req.stream == Some(true) {
+        chat_stream(state, req)
+    } else {
+        chat_non_stream(state, req).await
+    }
+}
+
+async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
+    let id = format!("chatcmpl-{}", Uuid::new_v4());
+    let model_name = state.model_name.clone();
+    let created = unix_timestamp();
+
+    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 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);
+
+    let (tx, mut rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
+    let gen_req = GenerateRequest {
+        prompt_tokens,
+        max_tokens,
+        sampling: sampling_params(&req),
+        sender: tx,
+    };
+    if let Err(resp) = submit_to_engine(&state, gen_req) {
+        return resp;
+    }
+
+    let mut content = String::new();
+    let mut completion_token_count: usize = 0;
+    let mut finish_reason = "length".to_string();
+    while let Some(event) = rx.recv().await {
+        match event {
+            GenerateEvent::Token { text, .. } => {
+                completion_token_count += 1;
+                content.push_str(&text);
+            }
+            GenerateEvent::Done { finish_reason: fr } => {
+                finish_reason = fr;
+                break;
+            }
+        }
+    }
+
+    Json(serde_json::json!({
+        "id": id,
+        "object": "chat.completion",
+        "created": created,
+        "model": model_name,
+        "choices": [{
+            "index": 0,
+            "message": { "role": "assistant", "content": content },
+            "finish_reason": finish_reason,
+        }],
+        "usage": {
+            "prompt_tokens": prompt_token_count,
+            "completion_tokens": completion_token_count,
+            "total_tokens": prompt_token_count + completion_token_count
+        }
+    })).into_response()
+}
+
+fn chat_stream(
+    state: Arc<AppState>,
+    req: ChatRequest,
+) -> Response {
+    let id = format!("chatcmpl-{}", Uuid::new_v4());
+    let model_name = state.model_name.clone();
+    let created = unix_timestamp();
+
+    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" | "user" | "assistant" => {
+                prompt.push_str("<|im_start|>");
+                prompt.push_str(&msg.role);
+                prompt.push('\n');
+                prompt.push_str(&msg.content);
+                prompt.push_str("<|im_end|>\n");
+            }
+            _ => {}
+        }
+    }
+    prompt.push_str("<|im_start|>assistant\n");
+    prompt.push_str("<think>\n\n</think>\n\n");
+    prompt
+}

crates/xserv-server/src/engine.rs

@@ -0,0 +1,361 @@
+use std::collections::VecDeque;
+use std::path::Path;
+use std::sync::mpsc;
+use std::sync::Once;
+use std::time::Instant;
+use xserv_model::{ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample, BLOCK_SIZE};
+use xserv_model::loader;
+use xserv_tensor::{DType, Device};
+use xserv_tokenizer::Tokenizer;
+
+pub struct Engine {
+    model: Qwen3,
+    config: ModelConfig,
+    tokenizer: Tokenizer,
+    max_batch_size: usize,
+    max_seq_len: usize,
+    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>,
+}
+
+pub enum GenerateEvent {
+    Token { id: u32, text: String },
+    Done { finish_reason: String },
+}
+
+struct Sequence {
+    id: u64,
+    prompt_tokens: Vec<u32>,
+    generated_tokens: Vec<u32>,
+    max_tokens: usize,
+    sampling: SamplingParams,
+    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, 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...");
+        let weights = loader::load_model_dir(model_dir, Device::Cuda(0));
+        eprintln!("[engine] Loaded {} tensors", weights.len());
+        let model = Qwen3::from_weights(config.clone(), weights);
+        let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+
+        // Tier-1 sizing: size the GPU block pool to *available VRAM* after the
+        // 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.
+    ///
+    /// 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 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: Swap previously-evicted sequences back in when there is
+            // room (oldest first). They resume decoding from where they paused.
+            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);
+            }
+
+            // 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 4: If nothing to do, blocking wait for new request.
+            if running.is_empty() && waiting.is_empty() && swapped.is_empty() {
+                match rx.recv() {
+                    Ok(req) => {
+                        let seq = self.make_sequence(req, &mut next_id);
+                        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 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 {
+                    let slot = seq.seq_slot.expect("slot");
+                    let logits = self.model.forward_prefill_paged(
+                        &seq.prompt_tokens, slot, &mut self.paged_cache,
+                    );
+                    let next = sample(&logits, &seq.sampling);
+                    seq.generated_tokens.push(next);
+                    seq.prefilled = true;
+                    emit_token(&self.tokenizer, seq, next);
+                    newly_prefilled.push(seq.id);
+                }
+            }
+
+            // 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) => {
+                        let seq = self.make_sequence(req, &mut next_id);
+                        waiting.push_back(seq);
+                    }
+                    Err(mpsc::TryRecvError::Empty) => break,
+                    Err(mpsc::TryRecvError::Disconnected) => return,
+                }
+            }
+        }
+    }
+
+    fn make_sequence(&mut self, req: GenerateRequest, next_id: &mut u64) -> Sequence {
+        let id = *next_id;
+        *next_id += 1;
+        Sequence {
+            id,
+            prompt_tokens: req.prompt_tokens,
+            generated_tokens: Vec::new(),
+            max_tokens: req.max_tokens,
+            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(),
+        }
+    }
+}
+
+/// Total additional GPU blocks the next decode step needs across all
+/// 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()
+}
+
+fn emit_token(tokenizer: &Tokenizer, seq: &mut Sequence, token_id: u32) {
+    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(),
+        });
+        return;
+    }
+
+    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 {
+        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; }
+    seq.sender.is_closed() || seq.eos_token_id == Some(last)
+}

crates/xserv-server/src/main.rs

@@ -0,0 +1,105 @@
+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] [--max-seq-len N] [--swap-space-gb N] [--tp N]");
+        std::process::exit(1);
+    }
+
+    let model_dir = PathBuf::from(&args[1]);
+    let port: u16 = args.iter()
+        .position(|a| a == "--port")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(8080);
+    let max_batch: usize = args.iter()
+        .position(|a| a == "--max-batch")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(4)
+        .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())
+        .unwrap_or_else(|| "unknown".to_string());
+
+    let tokenizer = xserv_tokenizer::Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+
+    // Unbounded channel: allows multiple requests to queue up
+    let (tx, rx) = mpsc::channel::<GenerateRequest>();
+
+    let model_dir_clone = model_dir.clone();
+    std::thread::spawn(move || {
+        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()
+        .route("/health", get(api::health))
+        .route("/v1/models", get(api::list_models))
+        .route("/v1/chat/completions", post(api::chat_completions))
+        .layer(Extension(state));
+
+    let addr = format!("0.0.0.0:{port}");
+    eprintln!("[server] Listening on {addr} (max_batch={max_batch}, max_seq_len={max_seq_len})");
+    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
+    axum::serve(listener, app).await.unwrap();
+}

crates/xserv-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/Cargo.toml

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

crates/xserv-tensor/src/dtype.rs

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

crates/xserv-tensor/src/lib.rs

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

crates/xserv-tensor/src/shape.rs

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

crates/xserv-tensor/src/storage.rs

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

crates/xserv-tensor/src/tensor.rs

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

crates/xserv-tensor/tests/integration.rs

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

crates/xserv-tokenizer/Cargo.toml

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

crates/xserv-tokenizer/src/bpe.rs

@@ -0,0 +1,367 @@
+use regex::Regex;
+use serde::Deserialize;
+use std::collections::HashMap;
+use std::path::Path;
+
+pub struct Tokenizer {
+    encoder: HashMap<Vec<u8>, u32>,
+    decoder: Vec<Vec<u8>>,
+    merge_ranks: HashMap<(u32, u32), usize>,
+    special_tokens: HashMap<String, u32>,
+    #[allow(dead_code)]
+    special_token_ids: HashMap<u32, String>,
+    pre_tokenize_re: Regex,
+    eos_token_id: Option<u32>,
+    byte_fallback: bool,
+}
+
+#[derive(Deserialize)]
+struct TokenizerJson {
+    model: ModelSection,
+    #[serde(default)]
+    added_tokens: Vec<AddedToken>,
+}
+
+#[derive(Deserialize)]
+struct ModelSection {
+    vocab: HashMap<String, u32>,
+    merges: Vec<MergeEntry>,
+    #[serde(default)]
+    byte_fallback: bool,
+}
+
+#[derive(Deserialize)]
+#[serde(untagged)]
+enum MergeEntry {
+    Str(String),
+    Pair(Vec<String>),
+}
+
+#[derive(Deserialize)]
+struct AddedToken {
+    id: u32,
+    content: String,
+    #[allow(dead_code)]
+    special: bool,
+}
+
+impl Tokenizer {
+    pub fn from_file(path: &Path) -> Self {
+        let data = std::fs::read_to_string(path)
+            .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
+        let tj: TokenizerJson = serde_json::from_str(&data)
+            .unwrap_or_else(|e| panic!("failed to parse tokenizer.json: {e}"));
+
+        // Build encoder: token bytes → ID
+        // All HF tokenizers use GPT-2 byte-to-unicode mapping for vocab keys.
+        let mut encoder = HashMap::new();
+        for (token_str, &id) in &tj.model.vocab {
+            let bytes = token_str_to_bytes(token_str);
+            encoder.insert(bytes, id);
+        }
+
+        // Build decoder: ID → token bytes
+        let max_id = tj.model.vocab.values().copied().max().unwrap_or(0);
+        let added_max = tj.added_tokens.iter().map(|t| t.id).max().unwrap_or(0);
+        let vocab_size = (max_id.max(added_max) + 1) as usize;
+        let mut decoder = vec![vec![]; vocab_size];
+        for (token_str, &id) in &tj.model.vocab {
+            decoder[id as usize] = token_str_to_bytes(token_str);
+        }
+
+        // Parse merges (supports both "a b" string format and ["a", "b"] array format)
+        let byte_fallback = tj.model.byte_fallback;
+        let mut merge_ranks = HashMap::new();
+        for (rank, entry) in tj.model.merges.iter().enumerate() {
+            let (a_str, b_str) = match entry {
+                MergeEntry::Str(s) => {
+                    let parts: Vec<&str> = s.splitn(2, ' ').collect();
+                    if parts.len() != 2 { continue; }
+                    (parts[0].to_string(), parts[1].to_string())
+                }
+                MergeEntry::Pair(v) => {
+                    if v.len() != 2 { continue; }
+                    (v[0].clone(), v[1].clone())
+                }
+            };
+            let a_bytes = token_str_to_bytes(&a_str);
+            let b_bytes = token_str_to_bytes(&b_str);
+            if let (Some(&a_id), Some(&b_id)) = (encoder.get(&a_bytes), encoder.get(&b_bytes)) {
+                merge_ranks.insert((a_id, b_id), rank);
+            }
+        }
+
+        // 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();
+        for at in &tj.added_tokens {
+            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();
+        }
+        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 {
+            // Qwen-style: split on whitespace boundaries, keep Unicode words/numbers
+            Regex::new(r"[\p{L}\p{N}]+|[^\s\p{L}\p{N}]|\s+").unwrap()
+        } else {
+            // GPT-2 style
+            Regex::new(r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+").unwrap()
+        };
+
+        Self {
+            encoder,
+            decoder,
+            merge_ranks,
+            special_tokens,
+            special_token_ids,
+            pre_tokenize_re,
+            eos_token_id,
+            byte_fallback,
+        }
+    }
+
+    pub fn encode(&self, text: &str) -> Vec<u32> {
+        let mut tokens = Vec::new();
+
+        // Check for special tokens first (split around them)
+        let mut remaining = text;
+        while !remaining.is_empty() {
+            // Find earliest special token
+            let mut earliest: Option<(usize, &str, u32)> = None;
+            for (st, &id) in &self.special_tokens {
+                if let Some(pos) = remaining.find(st.as_str()) {
+                    if earliest.is_none() || pos < earliest.unwrap().0 {
+                        earliest = Some((pos, st, id));
+                    }
+                }
+            }
+
+            if let Some((pos, st, id)) = earliest {
+                if pos > 0 {
+                    self.encode_ordinary(&remaining[..pos], &mut tokens);
+                }
+                tokens.push(id);
+                remaining = &remaining[pos + st.len()..];
+            } else {
+                self.encode_ordinary(remaining, &mut tokens);
+                break;
+            }
+        }
+
+        tokens
+    }
+
+    fn encode_ordinary(&self, text: &str, out: &mut Vec<u32>) {
+        for mat in self.pre_tokenize_re.find_iter(text) {
+            let word = mat.as_str();
+            // Try to encode the whole word first
+            if let Some(&id) = self.encoder.get(word.as_bytes()) {
+                out.push(id);
+                continue;
+            }
+            // Fall back to per-byte encoding
+            let word_bytes: Vec<u8> = word.bytes().collect();
+            let mut token_ids: Vec<u32> = word_bytes.iter().filter_map(|&b| {
+                if let Some(&id) = self.encoder.get(&vec![b]) {
+                    Some(id)
+                } else if self.byte_fallback {
+                    let hex_token = format!("<0x{:02X}>", b);
+                    if let Some(&id) = self.special_tokens.get(&hex_token) {
+                        Some(id)
+                    } else if let Some(&id) = self.encoder.get(hex_token.as_bytes()) {
+                        Some(id)
+                    } else if let Some(&unk_id) = self.special_tokens.get("<unk>") {
+                        eprintln!("warning: byte 0x{b:02X} not in vocab, using <unk> token");
+                        Some(unk_id)
+                    } else {
+                        eprintln!("warning: byte 0x{b:02X} not in vocab and no fallback token, using token 0");
+                        Some(0)
+                    }
+                } else {
+                    eprintln!("warning: byte {b} (0x{b:02X}) not in vocab, skipping");
+                    None
+                }
+            }).collect();
+
+            // BPE merges
+            loop {
+                if token_ids.len() < 2 { break; }
+                let mut best_rank = usize::MAX;
+                let mut best_idx = 0;
+                for i in 0..token_ids.len() - 1 {
+                    if let Some(&rank) = self.merge_ranks.get(&(token_ids[i], token_ids[i + 1])) {
+                        if rank < best_rank {
+                            best_rank = rank;
+                            best_idx = i;
+                        }
+                    }
+                }
+                if best_rank == usize::MAX { break; }
+
+                let merged_bytes = [
+                    self.decoder[token_ids[best_idx] as usize].as_slice(),
+                    self.decoder[token_ids[best_idx + 1] as usize].as_slice(),
+                ].concat();
+                let merged_id = *self.encoder.get(&merged_bytes).unwrap_or_else(|| {
+                    panic!("merged token not in vocab");
+                });
+                token_ids[best_idx] = merged_id;
+                token_ids.remove(best_idx + 1);
+            }
+
+            out.extend_from_slice(&token_ids);
+        }
+    }
+
+    pub fn decode(&self, token_ids: &[u32]) -> String {
+        let mut bytes = Vec::new();
+        for &id in token_ids {
+            if let Some(b) = self.decoder.get(id as usize) {
+                bytes.extend_from_slice(b);
+            }
+        }
+        String::from_utf8_lossy(&bytes).into_owned()
+    }
+
+    pub fn 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
+    }
+
+    pub fn vocab_size(&self) -> usize {
+        self.decoder.len()
+    }
+
+    pub fn special_token_id(&self, name: &str) -> Option<u32> {
+        self.special_tokens.get(name).copied()
+    }
+}
+
+/// Convert a token string from HF vocab (which uses Unicode replacements for bytes)
+/// back to raw bytes. GPT-2 uses a byte-to-unicode mapping where e.g. byte 0x20 (space)
+/// is represented as 'Ġ' (U+0120).
+fn token_str_to_bytes(s: &str) -> Vec<u8> {
+    s.chars().map(|c| unicode_to_byte(c)).collect()
+}
+
+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
+    use std::sync::OnceLock;
+    static INV_MAP: OnceLock<HashMap<u32, u8>> = OnceLock::new();
+
+    let map = INV_MAP.get_or_init(|| {
+        let mut m = HashMap::new();
+        // Build GPT-2's bytes_to_unicode forward map, then invert
+        let mut n = 0u32;
+        for b in 0..=255u16 {
+            let byte = b as u8;
+            let unicode = match byte {
+                0x21..=0x7E | 0xA1..=0xAC | 0xAE..=0xFF => byte as u32,
+                _ => {
+                    let u = 256 + n;
+                    n += 1;
+                    u
+                }
+            };
+            m.insert(unicode, byte);
+        }
+        m
+    });
+
+    *map.get(&(c as u32)).unwrap_or_else(|| {
+        panic!("unmapped unicode char U+{:04X} in tokenizer", c as u32)
+    })
+}
+
+#[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());
+    }
+}

crates/xserv-tokenizer/src/lib.rs

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

csrc/activation/activations.cu

@@ -0,0 +1,166 @@
+#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)))
+__device__ __forceinline__ float gelu_f(float x) {
+    const float SQRT_2_OVER_PI = 0.7978845608f;
+    float cube = x * x * x;
+    float inner = SQRT_2_OVER_PI * (x + 0.044715f * cube);
+    return 0.5f * x * (1.0f + tanhf(inner));
+}
+
+// SiLU (Swish): silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
+__device__ __forceinline__ float silu_f(float x) {
+    return x / (1.0f + expf(-x));
+}
+
+__global__ void gelu_f32(const float* x, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = gelu_f(x[idx]);
+}
+
+__global__ void gelu_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(gelu_f(__bfloat162float(x[idx])));
+}
+
+__global__ void silu_f32(const float* x, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = silu_f(x[idx]);
+}
+
+__global__ void silu_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(silu_f(__bfloat162float(x[idx])));
+}
+
+__global__ void scale_f32_kernel(const float* x, float* out, float scale, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = x[idx] * scale;
+}
+
+__global__ void scale_bf16_kernel(const __nv_bfloat16* x, __nv_bfloat16* out, float scale, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(x[idx]) * scale);
+}
+
+// Fused SiLU×Mul: out = silu(gate) * up
+__global__ void silu_mul_bf16_kernel(const __nv_bfloat16* gate, const __nv_bfloat16* up,
+                                     __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        float g = __bfloat162float(gate[idx]);
+        float u = __bfloat162float(up[idx]);
+        float silu_g = g / (1.0f + expf(-g));
+        out[idx] = __float2bfloat16(silu_g * u);
+    }
+}
+
+// Element-wise add: out = a + b
+__global__ void add_f32_kernel(const float* a, const float* b, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = a[idx] + b[idx];
+}
+__global__ void add_bf16_kernel(const __nv_bfloat16* a, const __nv_bfloat16* b, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(a[idx]) + __bfloat162float(b[idx]));
+}
+
+// Element-wise mul: out = a * b
+__global__ void mul_f32_kernel(const float* a, const float* b, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = a[idx] * b[idx];
+}
+__global__ void mul_bf16_kernel(const __nv_bfloat16* a, const __nv_bfloat16* b, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(a[idx]) * __bfloat162float(b[idx]));
+}
+
+extern "C" {
+
+void launch_gelu_f32(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    gelu_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_gelu_bf16(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    gelu_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+    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) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    silu_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_scale_f32(const void* x, void* out, float scale, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    scale_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)x, (float*)out, scale, n);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_scale_bf16(const void* x, void* out, float scale, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    scale_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, scale, n);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_add_f32(const void* a, const void* b, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    add_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)a, (const float*)b, (float*)out, n);
+    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

@@ -0,0 +1,55 @@
+#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`,
+// we allow attending to positions [0, offset + row].
+// scores: [batch, rows, cols]  (flattened batch×heads)
+
+__global__ void causal_mask_f32(
+    float* __restrict__ scores,
+    int rows, int cols, int offset
+) {
+    int batch_idx = blockIdx.z;
+    int row = blockIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (col < cols && col > row + offset) {
+        scores[batch_idx * rows * cols + row * cols + col] = -INFINITY;
+    }
+}
+
+__global__ void causal_mask_bf16(
+    __nv_bfloat16* __restrict__ scores,
+    int rows, int cols, int offset
+) {
+    int batch_idx = blockIdx.z;
+    int row = blockIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (col < cols && col > row + offset) {
+        scores[batch_idx * rows * cols + row * cols + col] = __float2bfloat16(-INFINITY);
+    }
+}
+
+extern "C" {
+
+void launch_causal_mask_f32(void* scores, int batch, int rows, int cols,
+                            int offset, void* stream) {
+    int block = 256;
+    dim3 grid((cols + block - 1) / block, rows, batch);
+    causal_mask_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (float*)scores, rows, cols, offset);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_causal_mask_bf16(void* scores, int batch, int rows, int cols,
+                             int offset, void* stream) {
+    int block = 256;
+    dim3 grid((cols + block - 1) / block, rows, batch);
+    causal_mask_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (__nv_bfloat16*)scores, rows, cols, offset);
+    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

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

@@ -0,0 +1,62 @@
+#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.
+
+__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 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;
+
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        dst[i] = row[i];
+    }
+}
+
+__global__ void embedding_bf16(
+    const __nv_bfloat16* __restrict__ table,
+    const int* __restrict__ token_ids,
+    __nv_bfloat16* __restrict__ out,
+    int hidden_size,
+    int 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;
+
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        dst[i] = row[i];
+    }
+}
+
+extern "C" {
+
+void launch_embedding_f32(const void* table, const void* token_ids, void* out,
+                          int num_tokens, int hidden_size, 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, 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, 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, vocab_size);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+}

csrc/embedding/rope.cu

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

csrc/embedding/transpose.cu

@@ -0,0 +1,242 @@
+#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).
+
+// reshape_heads: [S, H*D] → [1, H, S, D]
+// Input layout:  element at [s, h*D + d] = flat[s * H*D + h*D + d]
+// Output layout: element at [0, h, s, d] = flat[h * S*D + s*D + d]
+__global__ void reshape_heads_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len, int num_heads, int head_dim
+) {
+    int hidden = num_heads * head_dim;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = seq_len * hidden;
+    if (idx >= total) return;
+
+    int s = idx / hidden;
+    int rem = idx % hidden;
+    int h = rem / head_dim;
+    int d = rem % head_dim;
+
+    int out_idx = h * seq_len * head_dim + s * head_dim + d;
+    out[out_idx] = in[idx];
+}
+
+// merge_heads: [1, H, S, D] → [S, H*D]
+// Input layout:  element at [0, h, s, d] = flat[h * S*D + s*D + d]
+// Output layout: element at [s, h*D + d] = flat[s * H*D + h*D + d]
+__global__ void merge_heads_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len, int num_heads, int head_dim
+) {
+    int hidden = num_heads * head_dim;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = seq_len * hidden;
+    if (idx >= total) return;
+
+    // idx is output index: [s, h*D + d]
+    int s = idx / hidden;
+    int rem = idx % hidden;
+    int h = rem / head_dim;
+    int d = rem % head_dim;
+
+    int in_idx = h * seq_len * head_dim + s * head_dim + d;
+    out[idx] = in[in_idx];
+}
+
+// transpose_for_rope: [1, H, S, D] → [S, H, D]
+// Input:  [h, s, d] at h*S*D + s*D + d
+// Output: [s, h, d] at s*H*D + h*D + d
+__global__ void transpose_hsd_to_shd_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len, int num_heads, int head_dim
+) {
+    int total = seq_len * num_heads * head_dim;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total) return;
+
+    // idx = output flat index: s*H*D + h*D + d
+    int s = idx / (num_heads * head_dim);
+    int rem = idx % (num_heads * head_dim);
+    int h = rem / head_dim;
+    int d = rem % head_dim;
+
+    int in_idx = h * seq_len * head_dim + s * head_dim + d;
+    out[idx] = in[in_idx];
+}
+
+// transpose_from_rope: [S, H, D] → [1, H, S, D]
+// Input:  [s, h, d] at s*H*D + h*D + d
+// Output: [h, s, d] at h*S*D + s*D + d
+__global__ void transpose_shd_to_hsd_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len, int num_heads, int head_dim
+) {
+    int total = seq_len * num_heads * head_dim;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total) return;
+
+    // idx = output flat index: h*S*D + s*D + d
+    int h = idx / (seq_len * head_dim);
+    int rem = idx % (seq_len * head_dim);
+    int s = rem / head_dim;
+    int d = rem % head_dim;
+
+    int in_idx = s * num_heads * head_dim + h * head_dim + d;
+    out[idx] = in[in_idx];
+}
+
+// repeat_kv: [1, KV_H, S, D] → [1, KV_H * n_rep, S, D]
+__global__ void repeat_kv_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int kv_heads, int n_rep, int seq_len, int head_dim
+) {
+    int total_heads = kv_heads * n_rep;
+    int total = total_heads * seq_len * head_dim;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total) return;
+
+    int out_h = idx / (seq_len * head_dim);
+    int rem = idx % (seq_len * head_dim);
+    int kv_h = out_h / n_rep;
+
+    int in_idx = kv_h * seq_len * head_dim + rem;
+    out[idx] = in[in_idx];
+}
+
+// ---- Generic strided copy (up to 4D) ----
+// Each thread copies one element. Maps flat contiguous output index to strided input index.
+// Unused dimensions are padded with shape=1, stride=0.
+
+__global__ void strided_copy_bf16(
+    const __nv_bfloat16* __restrict__ in,
+    __nv_bfloat16* __restrict__ out,
+    int numel,
+    int ndim,
+    int shape0, int shape1, int shape2, int shape3,
+    int in_stride0, int in_stride1, int in_stride2, int in_stride3,
+    int in_offset
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= numel) return;
+
+    // Decompose flat output index into multi-dim indices (rightmost = fastest)
+    int remaining = idx;
+    int i3 = remaining % shape3; remaining /= shape3;
+    int i2 = remaining % shape2; remaining /= shape2;
+    int i1 = remaining % shape1; remaining /= shape1;
+    int i0 = remaining;
+
+    int in_idx = in_offset + i0 * in_stride0 + i1 * in_stride1 + i2 * in_stride2 + i3 * in_stride3;
+    out[idx] = in[in_idx];
+}
+
+__global__ void strided_copy_f32(
+    const float* __restrict__ in,
+    float* __restrict__ out,
+    int numel,
+    int ndim,
+    int shape0, int shape1, int shape2, int shape3,
+    int in_stride0, int in_stride1, int in_stride2, int in_stride3,
+    int in_offset
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= numel) return;
+
+    int remaining = idx;
+    int i3 = remaining % shape3; remaining /= shape3;
+    int i2 = remaining % shape2; remaining /= shape2;
+    int i1 = remaining % shape1; remaining /= shape1;
+    int i0 = remaining;
+
+    int in_idx = in_offset + i0 * in_stride0 + i1 * in_stride1 + i2 * in_stride2 + i3 * in_stride3;
+    out[idx] = in[in_idx];
+}
+
+extern "C" {
+
+void launch_reshape_heads_bf16(const void* in, void* out,
+                               int seq_len, int num_heads, int head_dim, void* stream) {
+    int total = seq_len * num_heads * head_dim;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    reshape_heads_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_merge_heads_bf16(const void* in, void* out,
+                             int seq_len, int num_heads, int head_dim, void* stream) {
+    int total = seq_len * num_heads * head_dim;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    merge_heads_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_transpose_hsd_to_shd_bf16(const void* in, void* out,
+                                       int seq_len, int num_heads, int head_dim, void* stream) {
+    int total = seq_len * num_heads * head_dim;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    transpose_hsd_to_shd_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_transpose_shd_to_hsd_bf16(const void* in, void* out,
+                                       int seq_len, int num_heads, int head_dim, void* stream) {
+    int total = seq_len * num_heads * head_dim;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    transpose_shd_to_hsd_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_repeat_kv_bf16(const void* in, void* out,
+                           int kv_heads, int n_rep, int seq_len, int head_dim, void* stream) {
+    int total = kv_heads * n_rep * seq_len * head_dim;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    repeat_kv_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, kv_heads, n_rep, seq_len, head_dim);
+    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

@@ -0,0 +1,65 @@
+#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]
+// All matrices are row-major.
+__global__ void gemm_naive_bf16(
+    const __nv_bfloat16* A, const __nv_bfloat16* B, __nv_bfloat16* C,
+    int M, int N, int K
+) {
+    int row = blockIdx.y * blockDim.y + threadIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (row < M && col < N) {
+        float sum = 0.0f;
+        for (int k = 0; k < K; k++) {
+            sum += __bfloat162float(A[row * K + k]) * __bfloat162float(B[k * N + col]);
+        }
+        C[row * N + col] = __float2bfloat16(sum);
+    }
+}
+
+__global__ void gemm_naive_f32(
+    const float* A, const float* B, float* C,
+    int M, int N, int K
+) {
+    int row = blockIdx.y * blockDim.y + threadIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (row < M && col < N) {
+        float sum = 0.0f;
+        for (int k = 0; k < K; k++) {
+            sum += A[row * K + k] * B[k * N + col];
+        }
+        C[row * N + col] = sum;
+    }
+}
+
+extern "C" {
+
+void launch_gemm_naive_bf16(
+    const void* A, const void* B, void* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(16, 16);
+    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
+    gemm_naive_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)A, (const __nv_bfloat16*)B, (__nv_bfloat16*)C, M, N, K
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_gemm_naive_f32(
+    const void* A, const void* B, void* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(16, 16);
+    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
+    gemm_naive_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)A, (const float*)B, (float*)C, M, N, K
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+} // extern "C"

csrc/gemm/tiled.cu

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

csrc/normalization/layernorm.cu

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

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

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

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

docs/02-tensor.md

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

docs/03-gemm.md

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

docs/04-transformer-kernels.md

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

docs/05-attention.md

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

docs/06-model-loading.md

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

docs/07-tokenizer.md

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

docs/08-gpt2.md

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

docs/09-kv-cache.md

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

docs/10-qwen3.md

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

docs/11-paged-attention.md

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

docs/12-continuous-batching.md

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

docs/13-http-api.md

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

docs/14-flash-attention.md

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

docs/15-performance.md

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

docs/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/phase10-qwen3.md

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

docs/benchmarks/phase14-flash-attention.md

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