server: VRAM-sized KV pool + vLLM-style swap scheduler

Fixes the paged-KV OOM at large --max-seq-len and adds elastic memory:

- Size the GPU block pool to available VRAM (cudaMemGetInfo) instead of the
  worst-case blocks_per_seq * max_batch * 2 reservation, which OOM'd at 8192.
- Scheduler tracks waiting/running/swapped sets: block-aware admission,
  swap-in of resumable sequences when blocks free, and preemption of the
  newest running sequence to host when the pool can't cover a decode step.
- --swap-space-gb (default 8) sizes the pinned host swap pool;
  XSERV_MAX_KV_BLOCKS forces a small pool to exercise swapping.
- api: poison-tolerant lock + clean 503 when the engine thread is gone,
  instead of cascading mutex-poison panics.

Verified on RTX 5090: serves at --max-seq-len 8192 (previously OOM), and a
forced 40-block pool drives 48 lossless swap-out/swap-in cycles under
concurrency with coherent output.

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

crates/xserv-cuda/src/ffi.rs

@@ -39,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;

crates/xserv-cuda/src/memory.rs

@@ -116,6 +116,56 @@ impl GpuBuffer {
         })
     }
 
+    /// Async copy `count` bytes from `src` at `src_offset` to `self` at `dst_offset` on `stream`.
+    pub fn copy_from_device_at_async(&mut self, src: &GpuBuffer, src_offset: usize, dst_offset: usize, count: usize, stream: &CudaStream) -> Result<()> {
+        assert!(src_offset + count <= src.len);
+        assert!(dst_offset + count <= self.len);
+        error::check(unsafe {
+            ffi::cudaMemcpyAsync(
+                self.ptr.add(dst_offset),
+                src.ptr.add(src_offset),
+                count,
+                ffi::CUDA_MEMCPY_D2D,
+                stream.as_raw(),
+            )
+        })
+    }
+
+    /// Copy `count` bytes from this GPU buffer at `src_offset` to a host slice (D2H).
+    pub fn copy_to_host_at(&self, dst: &mut [u8], src_offset: usize, count: usize) -> Result<()> {
+        assert!(src_offset + count <= self.len, "src range out of bounds");
+        assert!(count <= dst.len(), "host dst too small");
+        error::check(unsafe {
+            ffi::cudaMemcpy(
+                dst.as_mut_ptr(),
+                self.ptr.add(src_offset),
+                count,
+                ffi::CUDA_MEMCPY_D2H,
+            )
+        })
+    }
+
+    /// Copy `count` bytes from a host slice to this GPU buffer at `dst_offset` (H2D).
+    pub fn copy_from_host_at(&mut self, src: &[u8], dst_offset: usize, count: usize) -> Result<()> {
+        assert!(dst_offset + count <= self.len, "dst range out of bounds");
+        assert!(count <= src.len(), "host src too small");
+        error::check(unsafe {
+            ffi::cudaMemcpy(
+                self.ptr.add(dst_offset),
+                src.as_ptr(),
+                count,
+                ffi::CUDA_MEMCPY_H2D,
+            )
+        })
+    }
+
+    /// Async zero fill on stream.
+    pub fn zero_async(&mut self, stream: &CudaStream) -> Result<()> {
+        error::check(unsafe {
+            ffi::cudaMemsetAsync(self.ptr, 0, self.len, stream.as_raw())
+        })
+    }
+
     /// Consume the buffer without freeing GPU memory. Returns the raw pointer and length.
     /// Caller is responsible for eventually calling cudaFree.
     pub fn into_raw(self) -> (*mut u8, usize) {

crates/xserv-kernels/build.rs

@@ -26,6 +26,7 @@ fn main() {
         .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

@@ -19,7 +19,9 @@ fn dispatch_unary(x: &Tensor, f32_fn: unsafe extern "C" fn(*const c_void, *mut c
                   bf16_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void)) -> Tensor {
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
     let out = Tensor::empty(x.shape(), x.dtype(), x.device());
-    let n = x.numel() as i32;
+    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()),
@@ -38,7 +40,9 @@ fn dispatch_binary(a: &Tensor, b: &Tensor,
     assert!(matches!(a.device(), Device::Cuda(_)));
     assert_eq!(a.dtype(), b.dtype());
     let out = Tensor::empty(a.shape(), a.dtype(), a.device());
-    let n = a.numel() as i32;
+    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()),
@@ -55,7 +59,9 @@ pub fn silu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_silu_f32, launch_si
 pub fn scale(x: &Tensor, scale_val: f32) -> Tensor {
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
     let out = Tensor::empty(x.shape(), x.dtype(), x.device());
-    let n = x.numel() as i32;
+    let n = x.numel();
+    assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
+    let n = n as i32;
     unsafe {
         match x.dtype() {
             DType::F32 => launch_scale_f32(x.data_ptr() as _, out.data_ptr() as *mut c_void, scale_val, n, std::ptr::null_mut()),
@@ -77,7 +83,9 @@ pub fn silu_mul(gate: &Tensor, up: &Tensor) -> Tensor {
     assert!(matches!(gate.device(), Device::Cuda(_)));
     assert_eq!(gate.dtype(), DType::BF16, "silu_mul requires BF16");
     let out = Tensor::empty(gate.shape(), gate.dtype(), gate.device());
-    let n = gate.numel() as i32;
+    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,

crates/xserv-kernels/src/attention.rs

@@ -22,6 +22,17 @@ unsafe extern "C" {
         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) {
@@ -192,3 +203,58 @@ pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tens
 
     output
 }
+
+/// Paged decode attention.
+///
+/// q: [batch, num_q_heads, 1, head_dim] BF16, contiguous, GPU
+/// k_cache_ptr / v_cache_ptr: pointers to [num_blocks, num_kv_heads, BLOCK_SIZE, head_dim] BF16 pools
+/// block_tables_ptr: i32 [batch, max_blocks_per_seq] (rows already arranged for this batch)
+/// context_lens_ptr: i32 [batch]
+///
+/// Returns: [batch, num_q_heads, 1, head_dim] BF16
+#[allow(clippy::too_many_arguments)]
+pub fn paged_decode_attention(
+    q: &Tensor,
+    k_cache_ptr: *const c_void,
+    v_cache_ptr: *const c_void,
+    block_tables_ptr: *const i32,
+    context_lens_ptr: *const i32,
+    batch: usize,
+    num_q_heads: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    max_blocks_per_seq: usize,
+) -> Tensor {
+    assert_eq!(q.ndim(), 4);
+    assert_eq!(q.shape()[2], 1, "paged_decode_attention requires q_len == 1");
+    assert_eq!(q.dtype(), DType::BF16);
+    assert!(num_q_heads % num_kv_heads == 0, "GQA: num_q_heads must be divisible by num_kv_heads");
+    assert!(head_dim <= 128);
+
+    let scale = 1.0 / (head_dim as f32).sqrt();
+    let output = Tensor::empty(
+        &[batch, num_q_heads, 1, head_dim],
+        DType::BF16,
+        q.device(),
+    );
+
+    unsafe {
+        launch_paged_decode_attention_bf16(
+            q.data_ptr() as *const c_void,
+            k_cache_ptr,
+            v_cache_ptr,
+            output.data_ptr() as *mut c_void,
+            block_tables_ptr,
+            context_lens_ptr,
+            batch as i32,
+            num_q_heads as i32,
+            num_kv_heads as i32,
+            head_dim as i32,
+            max_blocks_per_seq as i32,
+            scale,
+            std::ptr::null_mut(),
+        );
+    }
+
+    output
+}

crates/xserv-kernels/src/dispatch.rs

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

crates/xserv-kernels/src/embedding.rs

@@ -4,9 +4,9 @@ use xserv_tensor::{DType, Device, Tensor};
 
 unsafe extern "C" {
     fn launch_embedding_f32(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
-                            num_tokens: i32, hidden_size: i32, stream: *mut c_void);
+                            num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void);
     fn launch_embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
-                             num_tokens: i32, hidden_size: i32, stream: *mut c_void);
+                             num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void);
 }
 
 /// Embedding lookup: table[token_ids[i]] for each i.
@@ -18,6 +18,9 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
 
     let hidden_size = table.shape()[1];
     let num_tokens = token_ids.len();
+    let vocab_size = table.shape()[0];
+    assert!(num_tokens <= i32::MAX as usize, "too many tokens for i32 kernel param");
+    assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
 
     // Upload token_ids to GPU
     let ids_bytes = unsafe {
@@ -29,6 +32,10 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
     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 {
@@ -36,12 +43,12 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
             DType::F32 => launch_embedding_f32(
                 table.data_ptr() as _, ids_gpu.as_ptr() as _,
                 out.data_ptr() as *mut c_void,
-                num_tokens as i32, hidden_size as i32, std::ptr::null_mut(),
+                num_tokens as i32, hidden_size as i32, vocab_size as i32, std::ptr::null_mut(),
             ),
             DType::BF16 => launch_embedding_bf16(
                 table.data_ptr() as _, ids_gpu.as_ptr() as _,
                 out.data_ptr() as *mut c_void,
-                num_tokens as i32, hidden_size as i32, std::ptr::null_mut(),
+                num_tokens as i32, hidden_size as i32, vocab_size as i32, std::ptr::null_mut(),
             ),
             _ => panic!("unsupported dtype for embedding"),
         }

crates/xserv-kernels/src/gemm.rs

@@ -20,7 +20,7 @@ unsafe extern "C" {
 }
 
 // --- FFI: cuBLAS ---
-type CublasHandle = *mut c_void;
+pub type CublasHandle = *mut c_void;
 
 #[allow(non_upper_case_globals)]
 const CUBLAS_OP_N: i32 = 0;
@@ -100,6 +100,13 @@ where
     })
 }
 
+/// 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.

crates/xserv-kernels/src/layernorm.rs

@@ -17,6 +17,8 @@ pub fn layernorm(x: &Tensor, gamma: &Tensor, beta: &Tensor, eps: f32) -> Tensor
     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 {

crates/xserv-kernels/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod activation;
 pub mod attention;
+pub mod dispatch;
 pub mod embedding;
 pub mod gemm;
 pub mod layernorm;
@@ -10,7 +11,7 @@ pub mod transpose;
 
 pub use activation::{add, gelu, mul, scale, silu, silu_mul};
 pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
-pub use attention::{attention, decode_attention, flash_attention};
+pub use attention::{attention, decode_attention, flash_attention, paged_decode_attention};
 pub use embedding::embedding;
 pub use gemm::{batched_matmul, matmul, GemmBackend};
 pub use layernorm::layernorm;

crates/xserv-kernels/src/rmsnorm.rs

@@ -20,6 +20,8 @@ pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
     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 {
@@ -54,6 +56,8 @@ pub fn add_rmsnorm(x: &Tensor, residual: &Tensor, gamma: &Tensor, eps: f32) -> (
     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());
 

crates/xserv-kernels/src/softmax.rs

@@ -14,6 +14,8 @@ pub fn softmax(x: &Tensor) -> Tensor {
 
     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 {

crates/xserv-kernels/tests/ops_test.rs

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

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

@@ -1,14 +1,14 @@
 use std::path::PathBuf;
 use std::time::Instant;
 use xserv_model::qwen3::sample_greedy;
-use xserv_model::{loader, GpuKVCache, ModelConfig, Qwen3};
+use xserv_model::{loader, DecodeGraphState, GpuKVCache, ModelConfig, Qwen3};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
 fn main() {
     let args: Vec<String> = std::env::args().collect();
     if args.len() < 2 {
-        eprintln!("Usage: bench-qwen3 <model-dir> [--gen-tokens N]");
+        eprintln!("Usage: bench-qwen3 <model-dir> [--gen-tokens N] [--cuda-graph]");
         std::process::exit(1);
     }
     let model_dir = PathBuf::from(&args[1]);
@@ -18,6 +18,7 @@ fn main() {
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
         .unwrap_or(20);
+    let use_cuda_graph = args.iter().any(|a| a == "--cuda-graph");
 
     xserv_cuda::device::set_device(0).unwrap();
 
@@ -34,6 +35,18 @@ fn main() {
         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![
@@ -96,6 +109,12 @@ fn main() {
 
         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);
@@ -109,8 +128,35 @@ fn main() {
         for _ in 1..gen_tokens {
             let last = *generated.last().unwrap();
             let t_start = Instant::now();
-            let logits = model.forward_gpu_cache(&[last], &mut cache);
-            let next = sample_greedy(&logits);
+
+            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; }

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

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

crates/xserv-model/src/decode_graph.rs

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

crates/xserv-model/src/gpt2.rs

@@ -280,45 +280,88 @@ fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
 fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) -> (Tensor, Tensor, Tensor) {
     let hidden = num_heads * head_dim;
     let qkv_cpu = qkv.to_device(Device::Cpu);
-    let data = qkv_cpu.as_slice::<f32>();
-
-    let mut q_data = vec![0.0f32; num_heads * seq_len * head_dim];
-    let mut k_data = vec![0.0f32; num_heads * seq_len * head_dim];
-    let mut v_data = vec![0.0f32; num_heads * seq_len * head_dim];
-
-    for s in 0..seq_len {
-        let row = &data[s * 3 * hidden..(s + 1) * 3 * hidden];
-        for h in 0..num_heads {
-            let src_off = h * head_dim;
-            let dst_off = (h * seq_len + s) * head_dim;
-            q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
-            k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
-            v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
-        }
-    }
-
     let device = qkv.device();
-    let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
-    let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
-    let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
-    (q, k, v)
+    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 src = x_cpu.as_slice::<f32>();
+    let device = x.device();
+    let dtype = x.dtype();
 
-    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]);
+    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),
     }
-    Tensor::from_slice(&out, &[seq_len, hidden]).to_device(x.device())
 }
 
 /// Greedy sampling: return the argmax token ID from the last position's logits.

crates/xserv-model/src/kv_cache.rs

@@ -76,6 +76,7 @@ impl GpuKVCache {
 
     pub fn advance_seq_len(&mut self, new_tokens: usize) {
         self.seq_len += new_tokens;
+        assert!(self.seq_len <= self.max_seq_len, "KV cache seq_len ({}) exceeds max_seq_len ({})", self.seq_len, self.max_seq_len);
     }
 
     /// Get K/V cache tensors for a layer up to `seq_len` tokens: [1, num_kv_heads, seq_len, head_dim]
@@ -85,6 +86,7 @@ impl GpuKVCache {
     }
 
     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;

crates/xserv-model/src/lib.rs

@@ -1,13 +1,17 @@
 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};
 

crates/xserv-model/src/paged_kv_cache.rs

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

crates/xserv-model/src/qwen3.rs

@@ -6,6 +6,7 @@ use xserv_tensor::{DType, Device, Tensor};
 use crate::config::ModelConfig;
 use crate::gpt2::KVCache;
 use crate::kv_cache::GpuKVCache;
+use crate::paged_kv_cache::PagedKVCache;
 
 pub struct Qwen3 {
     pub config: ModelConfig,
@@ -255,6 +256,196 @@ impl Qwen3 {
         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);
+
+        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;
+
+        // 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);
+
+            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);
+            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);
+        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;
+
+        // 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);
+
+            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);
+            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();
@@ -320,6 +511,40 @@ impl Qwen3 {
         let x = rmsnorm(&x, &self.norm, eps);
         matmul_2d(&x, &self.lm_head_t)
     }
+
+    /// Extract weight pointers for CUDA Graph capture.
+    pub fn layer_weight_ptrs(&self) -> Vec<crate::decode_graph::LayerWeightPtrs> {
+        self.layers.iter().map(|l| crate::decode_graph::LayerWeightPtrs {
+            input_norm: l.input_norm.data_ptr() as *const std::ffi::c_void,
+            q_proj_wt: l.q_proj_wt.data_ptr() as *const std::ffi::c_void,
+            k_proj_wt: l.k_proj_wt.data_ptr() as *const std::ffi::c_void,
+            v_proj_wt: l.v_proj_wt.data_ptr() as *const std::ffi::c_void,
+            o_proj_wt: l.o_proj_wt.data_ptr() as *const std::ffi::c_void,
+            q_norm: l.q_norm.data_ptr() as *const std::ffi::c_void,
+            k_norm: l.k_norm.data_ptr() as *const std::ffi::c_void,
+            post_norm: l.post_norm.data_ptr() as *const std::ffi::c_void,
+            gate_proj_wt: l.gate_proj_wt.data_ptr() as *const std::ffi::c_void,
+            up_proj_wt: l.up_proj_wt.data_ptr() as *const std::ffi::c_void,
+            down_proj_wt: l.down_proj_wt.data_ptr() as *const std::ffi::c_void,
+        }).collect()
+    }
+
+    /// Get pointers needed for CUDA Graph capture.
+    pub fn graph_capture_ptrs(&self) -> (
+        *const std::ffi::c_void,  // norm weight
+        *const std::ffi::c_void,  // lm_head_t
+        *const std::ffi::c_void,  // embed_tokens
+        *const std::ffi::c_void,  // rope cos
+        *const std::ffi::c_void,  // rope sin
+    ) {
+        (
+            self.norm.data_ptr() as *const std::ffi::c_void,
+            self.lm_head_t.data_ptr() as *const std::ffi::c_void,
+            self.embed_tokens.data_ptr() as *const std::ffi::c_void,
+            self.rope_cache.cos.as_ptr() as *const std::ffi::c_void,
+            self.rope_cache.sin.as_ptr() as *const std::ffi::c_void,
+        )
+    }
 }
 
 // --- Helpers ---

crates/xserv-server/src/api.rs

@@ -85,18 +85,20 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
     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 (StatusCode::BAD_REQUEST, Json(serde_json::json!({
-            "error": {
-                "message": format!("prompt is {} tokens, exceeds max_seq_len {}", prompt_token_count, max_seq_len),
-                "type": "invalid_request_error"
-            }
-        }))).into_response();
+        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);
 
@@ -107,12 +109,9 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
         sampling: sampling_params(&req),
         sender: tx,
     };
-    state
-        .engine_sender
-        .lock()
-        .unwrap()
-        .send(gen_req)
-        .expect("engine channel closed");
+    if let Err(resp) = submit_to_engine(&state, gen_req) {
+        return resp;
+    }
 
     let mut content = String::new();
     let mut completion_token_count: usize = 0;
@@ -156,17 +155,19 @@ fn chat_stream(
     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 (StatusCode::BAD_REQUEST, Json(serde_json::json!({
-            "error": {
-                "message": format!("prompt is {} tokens, exceeds max_seq_len {}", prompt_tokens.len(), max_seq_len),
-                "type": "invalid_request_error"
-            }
-        }))).into_response();
+        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());
 
@@ -177,12 +178,9 @@ fn chat_stream(
         sampling: sampling_params(&req),
         sender: engine_tx,
     };
-    state
-        .engine_sender
-        .lock()
-        .unwrap()
-        .send(gen_req)
-        .expect("engine channel closed");
+    if let Err(resp) = submit_to_engine(&state, gen_req) {
+        return resp;
+    }
 
     // SSE event channel: engine events -> SSE events
     let (sse_tx, sse_rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
@@ -228,6 +226,53 @@ fn chat_stream(
     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,
@@ -295,5 +340,6 @@ fn build_prompt(messages: &[Message]) -> String {
         }
     }
     prompt.push_str("<|im_start|>assistant\n");
+    prompt.push_str("<think>\n\n</think>\n\n");
     prompt
 }

crates/xserv-server/src/engine.rs

@@ -3,7 +3,7 @@ use std::path::Path;
 use std::sync::mpsc;
 use std::sync::Once;
 use std::time::Instant;
-use xserv_model::{GpuKVCache, ModelConfig, Qwen3, SamplingParams, sample};
+use xserv_model::{ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample, BLOCK_SIZE};
 use xserv_model::loader;
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
@@ -14,6 +14,7 @@ pub struct Engine {
     tokenizer: Tokenizer,
     max_batch_size: usize,
     max_seq_len: usize,
+    paged_cache: PagedKVCache,
 }
 
 pub struct GenerateRequest {
@@ -34,15 +35,25 @@ struct Sequence {
     generated_tokens: Vec<u32>,
     max_tokens: usize,
     sampling: SamplingParams,
-    kv_cache: Option<GpuKVCache>,
+    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...");
@@ -50,8 +61,55 @@ impl Engine {
         eprintln!("[engine] Loaded {} tensors", weights.len());
         let model = Qwen3::from_weights(config.clone(), weights);
         let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
-        eprintln!("[engine] Ready (max_batch_size={max_batch_size}, max_seq_len={max_seq_len})");
-        Self { model, config, tokenizer, max_batch_size, max_seq_len }
+
+        // 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 }
@@ -59,54 +117,124 @@ impl Engine {
     pub fn max_seq_len(&self) -> usize { self.max_seq_len }
 
     /// Main scheduler loop. Receives requests from channel, manages concurrent sequences.
-    pub fn run(&self, rx: mpsc::Receiver<GenerateRequest>) {
+    ///
+    /// Sequences move between three sets:
+    ///   waiting  — admitted to the queue, no GPU slot yet
+    ///   running  — KV resident on GPU, actively prefilling/decoding
+    ///   swapped  — KV evicted to pinned host memory (preempted), paused
+    /// When running sequences grow past the GPU block pool, the newest are
+    /// swapped out to host (vLLM-style) and swapped back in when blocks free up.
+    pub fn run(&mut self, rx: mpsc::Receiver<GenerateRequest>) {
         let mut waiting: VecDeque<Sequence> = VecDeque::new();
         let mut running: Vec<Sequence> = Vec::new();
+        let mut swapped: Vec<Sequence> = Vec::new();
         let mut next_id: u64 = 0;
 
         eprintln!("[scheduler] Listening for requests...");
 
         loop {
-            // Step 1: Remove finished sequences
+            // Step 1: Remove finished sequences and return their slots.
+            let finished_slots: Vec<usize> = running.iter()
+                .filter(|s| is_finished(s))
+                .filter_map(|s| s.seq_slot)
+                .collect();
+            for slot in finished_slots {
+                self.paged_cache.free_sequence(slot);
+            }
             running.retain(|seq| !is_finished(seq));
 
-            // Step 2: Admit new sequences from waiting queue
-            while running.len() < self.max_batch_size {
-                if let Some(seq) = waiting.pop_front() {
+            // 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);
-                } else {
-                    break;
+                    avail -= prompt_blocks; // projected free after this seq prefills
                 }
             }
 
-            // Step 3: If nothing to do, blocking wait for new request
-            if running.is_empty() {
+            // 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);
-                        running.push(seq);
+                        waiting.push_back(seq);
+                        continue;
                     }
                     Err(_) => break, // channel closed
                 }
             }
+            // Nothing runnable this iteration (e.g. all swapped, waiting on
+            // blocks to free): loop to retry swap-in/admission next iteration.
+            if running.is_empty() {
+                continue;
+            }
 
-            // Step 4a: Process prefills (one at a time — different prompt lengths)
-            // Prefill sequences must be processed individually because they have
-            // different prompt lengths and each needs a full forward pass.
+            // 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 logits = self.model.forward_gpu_cache(&seq.prompt_tokens, seq.kv_cache.as_mut().unwrap());
+                    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;
-                    self.emit_token(seq, next);
+                    emit_token(&self.tokenizer, seq, next);
                     newly_prefilled.push(seq.id);
                 }
             }
 
-            // Step 4b: Batched decode — batch all decode-ready sequences into one forward pass.
-            // Projections and FFN run as [B, hidden] matmuls; attention remains per-seq.
+            // 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)
@@ -115,65 +243,44 @@ impl Engine {
             if !decode_indices.is_empty() {
                 static LOG_ONCE: Once = Once::new();
                 LOG_ONCE.call_once(|| {
-                    eprintln!("[scheduler] batched decode active");
+                    eprintln!("[scheduler] paged decode active");
                 });
-                eprintln!("[scheduler] decode batch_size={}", decode_indices.len());
 
-                if decode_indices.len() == 1 {
-                    // Single sequence: use per-seq path (no batching overhead)
-                    let i = decode_indices[0];
-                    let last = *running[i].generated_tokens.last().unwrap();
-                    let logits = self.model.forward_gpu_cache(&[last], running[i].kv_cache.as_mut().unwrap());
-                    let next = sample(&logits, &running[i].sampling);
+                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);
-                    self.emit_token(&running[i], next);
-                } else {
-                    // Batched decode: extract tokens and positions
-                    let tokens: Vec<u32> = decode_indices.iter()
-                        .map(|&i| *running[i].generated_tokens.last().unwrap())
-                        .collect();
-                    let positions: Vec<usize> = decode_indices.iter()
-                        .map(|&i| running[i].kv_cache.as_ref().unwrap().seq_len())
-                        .collect();
-
-                    // Take caches out of sequences via Option::take (no dummy allocation).
-                    let mut caches: Vec<GpuKVCache> = decode_indices.iter()
-                        .map(|&i| running[i].kv_cache.take().unwrap())
-                        .collect();
-                    let mut cache_refs: Vec<&mut GpuKVCache> = caches.iter_mut().collect();
-
-                    let logits = self.model.forward_decode_batch(&tokens, &positions, &mut cache_refs);
-
-                    // Put caches back: pop from end while iterating in reverse
-                    drop(cache_refs);
-                    for &i in decode_indices.iter().rev() {
-                        running[i].kv_cache = Some(caches.pop().unwrap());
-                    }
-
-                    // Sample per-sequence from batched logits [B, vocab_size]
-                    let vocab_size = logits.shape()[1];
-                    let logits_cpu = logits.to_device(xserv_tensor::Device::Cpu);
-                    let data = logits_cpu.as_slice::<half::bf16>();
-                    for (j, &i) in decode_indices.iter().enumerate() {
-                        let row_start = j * vocab_size;
-                        let row_logits = &data[row_start..row_start + vocab_size];
-                        let next = if running[i].sampling.temperature == 0.0 {
-                            // Greedy: argmax
-                            row_logits.iter().enumerate()
-                                .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
-                                .map(|(idx, _)| idx as u32).unwrap()
-                        } else {
-                            // Use the row as a single-row tensor for full sampling
-                            let row_tensor = xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]);
-                            sample(&row_tensor, &running[i].sampling)
-                        };
-                        running[i].generated_tokens.push(next);
-                        self.emit_token(&running[i], next);
-                    }
+                    emit_token(&self.tokenizer, &mut running[i], next);
                 }
             }
 
-            // Step 5: Check for newly arrived requests (non-blocking)
+            // Step 6: Check for newly arrived requests (non-blocking)
             loop {
                 match rx.try_recv() {
                     Ok(req) => {
@@ -187,39 +294,62 @@ impl Engine {
         }
     }
 
-    fn make_sequence(&self, req: GenerateRequest, next_id: &mut u64) -> Sequence {
+    fn make_sequence(&mut self, req: GenerateRequest, next_id: &mut u64) -> Sequence {
         let id = *next_id;
         *next_id += 1;
-        let kv_cache = GpuKVCache::new(&self.config, self.max_seq_len, DType::BF16, 0);
         Sequence {
             id,
             prompt_tokens: req.prompt_tokens,
             generated_tokens: Vec::new(),
             max_tokens: req.max_tokens,
             sampling: req.sampling,
-            kv_cache: Some(kv_cache),
+            seq_slot: None,
             sender: req.sender,
             prefilled: false,
             eos_token_id: self.tokenizer.eos_token_id(),
+            decode_buffer: Vec::new(),
             created_at: Instant::now(),
         }
     }
+}
 
-    fn emit_token(&self, seq: &Sequence, token_id: u32) {
-        let text = self.tokenizer.decode(&[token_id]);
+/// 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()
+}
 
-        if self.tokenizer.eos_token_id() == Some(token_id) {
-            let _ = seq.sender.blocking_send(GenerateEvent::Done {
-                finish_reason: "stop".to_string(),
-            });
-        } else if seq.generated_tokens.len() >= seq.max_tokens {
-            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
-            let _ = seq.sender.blocking_send(GenerateEvent::Done {
-                finish_reason: "length".to_string(),
-            });
-        } else {
-            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
-        }
+fn 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 });
     }
 }
 
@@ -227,7 +357,5 @@ fn is_finished(seq: &Sequence) -> bool {
     if seq.generated_tokens.is_empty() { return false; }
     let last = *seq.generated_tokens.last().unwrap();
     if seq.generated_tokens.len() >= seq.max_tokens { return true; }
-    // Check EOS — need tokenizer info. Use a simple heuristic:
-    // If sender is closed (receiver dropped), also consider finished.
     seq.sender.is_closed() || seq.eos_token_id == Some(last)
 }

crates/xserv-server/src/main.rs

@@ -5,6 +5,7 @@ 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,
@@ -17,7 +18,7 @@ pub struct AppState {
 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]");
+        eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N]");
         std::process::exit(1);
     }
 
@@ -31,12 +32,31 @@ async fn main() {
         .position(|a| a == "--max-batch")
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
-        .unwrap_or(4);
-    let max_seq_len: usize = args.iter()
+        .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);
+        .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 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())
@@ -49,7 +69,7 @@ async fn main() {
 
     let model_dir_clone = model_dir.clone();
     std::thread::spawn(move || {
-        let engine = engine::Engine::load(&model_dir_clone, max_batch, max_seq_len);
+        let mut engine = engine::Engine::load_with_swap(&model_dir_clone, max_batch, max_seq_len, swap_space_gb);
         engine.run(rx);
     });
 

crates/xserv-tokenizer/src/bpe.rs

@@ -41,6 +41,7 @@ enum MergeEntry {
 struct AddedToken {
     id: u32,
     content: String,
+    #[allow(dead_code)]
     special: bool,
 }
 
@@ -90,21 +91,22 @@ impl Tokenizer {
             }
         }
 
-        // Special tokens
+        // Added tokens are matched as indivisible tokens by HF tokenizers,
+        // even when their `special` flag is false (for example Qwen3's
+        // <think> and </think> tokens).
         let mut special_tokens = HashMap::new();
         let mut special_token_ids = HashMap::new();
-        let mut eos_token_id = None;
         for at in &tj.added_tokens {
-            if at.special {
-                special_tokens.insert(at.content.clone(), at.id);
-                special_token_ids.insert(at.id, at.content.clone());
-                decoder.resize(decoder.len().max(at.id as usize + 1), vec![]);
-                decoder[at.id as usize] = at.content.as_bytes().to_vec();
-                if at.content == "<|endoftext|>" || at.content == "<|end_of_text|>" {
-                    eos_token_id = Some(at.id);
-                }
-            }
+            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 {
@@ -230,6 +232,19 @@ impl Tokenizer {
         String::from_utf8_lossy(&bytes).into_owned()
     }
 
+    pub fn decode_token_stream(&self, token_id: u32, pending: &mut Vec<u8>) -> String {
+        if let Some(bytes) = self.decoder.get(token_id as usize) {
+            pending.extend_from_slice(bytes);
+        }
+        take_valid_utf8(pending)
+    }
+
+    pub fn flush_decode_stream(&self, pending: &mut Vec<u8>) -> String {
+        let text = String::from_utf8_lossy(pending).into_owned();
+        pending.clear();
+        text
+    }
+
     pub fn eos_token_id(&self) -> Option<u32> {
         self.eos_token_id
     }
@@ -250,6 +265,31 @@ fn token_str_to_bytes(s: &str) -> Vec<u8> {
     s.chars().map(|c| unicode_to_byte(c)).collect()
 }
 
+fn take_valid_utf8(pending: &mut Vec<u8>) -> String {
+    match std::str::from_utf8(pending) {
+        Ok(text) => {
+            let text = text.to_string();
+            pending.clear();
+            text
+        }
+        Err(err) => {
+            let valid_up_to = err.valid_up_to();
+            if valid_up_to == 0 {
+                if let Some(error_len) = err.error_len() {
+                    let invalid_len = error_len.min(pending.len());
+                    let text = String::from_utf8_lossy(&pending[..invalid_len]).into_owned();
+                    pending.drain(..invalid_len);
+                    return text;
+                }
+                return String::new();
+            }
+            let text = String::from_utf8_lossy(&pending[..valid_up_to]).into_owned();
+            pending.drain(..valid_up_to);
+            text
+        }
+    }
+}
+
 /// Convert a Unicode char back to the byte it represents in GPT-2 encoding.
 fn unicode_to_byte(c: char) -> u8 {
     // Build the inverse map on first use
@@ -279,3 +319,49 @@ fn unicode_to_byte(c: char) -> u8 {
         panic!("unmapped unicode char U+{:04X} in tokenizer", c as u32)
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{take_valid_utf8, Tokenizer};
+
+    #[test]
+    fn qwen_added_tokens_are_indivisible_and_im_end_is_eos() {
+        let path =
+            std::env::temp_dir().join(format!("xserv-tokenizer-test-{}.json", std::process::id()));
+        std::fs::write(
+            &path,
+            r#"{
+                "model": {
+                    "vocab": {},
+                    "merges": [],
+                    "byte_fallback": false
+                },
+                "added_tokens": [
+                    {"id":151643,"content":"<|endoftext|>","special":true},
+                    {"id":151644,"content":"<|im_start|>","special":true},
+                    {"id":151645,"content":"<|im_end|>","special":true},
+                    {"id":151667,"content":"<think>","special":false},
+                    {"id":151668,"content":"</think>","special":false}
+                ]
+            }"#,
+        )
+        .unwrap();
+
+        let tokenizer = Tokenizer::from_file(&path);
+        let _ = std::fs::remove_file(&path);
+
+        assert_eq!(tokenizer.eos_token_id(), Some(151645));
+        assert_eq!(tokenizer.encode("<think>"), vec![151667]);
+        assert_eq!(tokenizer.encode("</think>"), vec![151668]);
+        assert_eq!(tokenizer.decode(&[151645]), "<|im_end|>");
+    }
+
+    #[test]
+    fn stream_decode_buffers_incomplete_utf8() {
+        let mut pending = vec![0xF0, 0x9F];
+        assert_eq!(take_valid_utf8(&mut pending), "");
+        pending.extend_from_slice(&[0x98, 0x8A, b'!']);
+        assert_eq!(take_valid_utf8(&mut pending), "😊!");
+        assert!(pending.is_empty());
+    }
+}

csrc/activation/activations.cu

@@ -1,5 +1,6 @@
 #include <cuda_bf16.h>
 #include <math.h>
+#include "../common.cuh"
 
 // GELU (tanh approximation):
 // gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
@@ -83,6 +84,7 @@ void launch_gelu_f32(const void* x, void* out, int n, void* stream) {
     int block = 256;
     int grid = (n + block - 1) / block;
     gelu_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_gelu_bf16(const void* x, void* out, int n, void* stream) {
@@ -90,12 +92,14 @@ void launch_gelu_bf16(const void* x, void* out, int n, void* stream) {
     int grid = (n + block - 1) / block;
     gelu_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_silu_f32(const void* x, void* out, int n, void* stream) {
     int block = 256;
     int grid = (n + block - 1) / block;
     silu_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_silu_bf16(const void* x, void* out, int n, void* stream) {
@@ -103,6 +107,7 @@ void launch_silu_bf16(const void* x, void* out, int n, void* stream) {
     int grid = (n + block - 1) / block;
     silu_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_scale_f32(const void* x, void* out, float scale, int n, void* stream) {
@@ -110,6 +115,7 @@ void launch_scale_f32(const void* x, void* out, float scale, int n, void* stream
     int grid = (n + block - 1) / block;
     scale_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const float*)x, (float*)out, scale, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_scale_bf16(const void* x, void* out, float scale, int n, void* stream) {
@@ -117,6 +123,7 @@ void launch_scale_bf16(const void* x, void* out, float scale, int n, void* strea
     int grid = (n + block - 1) / block;
     scale_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)x, (__nv_bfloat16*)out, scale, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_add_f32(const void* a, const void* b, void* out, int n, void* stream) {
@@ -124,24 +131,28 @@ void launch_add_f32(const void* a, const void* b, void* out, int n, void* stream
     int grid = (n + block - 1) / block;
     add_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const float*)a, (const float*)b, (float*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 void launch_add_bf16(const void* a, const void* b, void* out, int n, void* stream) {
     int block = 256;
     int grid = (n + block - 1) / block;
     add_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 void launch_mul_f32(const void* a, const void* b, void* out, int n, void* stream) {
     int block = 256;
     int grid = (n + block - 1) / block;
     mul_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const float*)a, (const float*)b, (float*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 void launch_mul_bf16(const void* a, const void* b, void* out, int n, void* stream) {
     int block = 256;
     int grid = (n + block - 1) / block;
     mul_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_silu_mul_bf16(const void* gate, const void* up, void* out, int n, void* stream) {
@@ -149,6 +160,7 @@ void launch_silu_mul_bf16(const void* gate, const void* up, void* out, int n, vo
     int grid = (n + block - 1) / block;
     silu_mul_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)gate, (const __nv_bfloat16*)up, (__nv_bfloat16*)out, n);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/attention/causal_mask.cu

@@ -1,4 +1,5 @@
 #include <cuda_bf16.h>
+#include "../common.cuh"
 
 // Apply causal mask: set scores[row][col] = -inf where col > row + offset.
 // offset is used for KV cache: when query starts at position `offset`,
@@ -39,6 +40,7 @@ void launch_causal_mask_f32(void* scores, int batch, int rows, int cols,
     dim3 grid((cols + block - 1) / block, rows, batch);
     causal_mask_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
         (float*)scores, rows, cols, offset);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_causal_mask_bf16(void* scores, int batch, int rows, int cols,
@@ -47,6 +49,7 @@ void launch_causal_mask_bf16(void* scores, int batch, int rows, int cols,
     dim3 grid((cols + block - 1) / block, rows, batch);
     causal_mask_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (__nv_bfloat16*)scores, rows, cols, offset);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/attention/flash_attention.cu

@@ -1,5 +1,6 @@
 #include <cuda_bf16.h>
 #include <float.h>
+#include "../common.cuh"
 
 // Flash Attention 2 forward kernel for BF16 with FP32 accumulation.
 //
@@ -391,6 +392,7 @@ void launch_flash_attention_bf16(
         q_len, kv_len, head_dim,
         scale, causal
     );
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_decode_attention_bf16(
@@ -411,6 +413,7 @@ void launch_decode_attention_bf16(
         kv_len, head_dim,
         scale
     );
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/attention/paged_attention.cu

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

csrc/common.cuh

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

csrc/embedding/embedding.cu

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

csrc/embedding/rope.cu

@@ -1,10 +1,11 @@
 #include <cuda_bf16.h>
 #include <math.h>
+#include "../common.cuh"
 
-// RoPE: Rotary Position Embedding
-// For each pair (x[2i], x[2i+1]) at position `pos`:
-//   y[2i]   = x[2i] * cos - x[2i+1] * sin
-//   y[2i+1] = x[2i] * sin + x[2i+1] * cos
+// 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])
@@ -35,11 +36,11 @@ __global__ void rope_f32(
     float sin_val = sin_cache[pos * half_dim + pair_idx];
 
     int base = (token_idx * num_heads + head_idx) * head_dim;
-    float x0 = x[base + 2 * pair_idx];
-    float x1 = x[base + 2 * pair_idx + 1];
+    float x0 = x[base + pair_idx];
+    float x1 = x[base + pair_idx + half_dim];
 
-    x[base + 2 * pair_idx]     = x0 * cos_val - x1 * sin_val;
-    x[base + 2 * pair_idx + 1] = x0 * sin_val + x1 * cos_val;
+    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(
@@ -61,11 +62,11 @@ __global__ void rope_bf16(
     float sin_val = sin_cache[pos * half_dim + pair_idx];
 
     int base = (token_idx * num_heads + head_idx) * head_dim;
-    float x0 = __bfloat162float(x[base + 2 * pair_idx]);
-    float x1 = __bfloat162float(x[base + 2 * pair_idx + 1]);
+    float x0 = __bfloat162float(x[base + pair_idx]);
+    float x1 = __bfloat162float(x[base + pair_idx + half_dim]);
 
-    x[base + 2 * pair_idx]     = __float2bfloat16(x0 * cos_val - x1 * sin_val);
-    x[base + 2 * pair_idx + 1] = __float2bfloat16(x0 * sin_val + x1 * cos_val);
+    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
@@ -94,6 +95,7 @@ void launch_rope_f32(void* x, const void* cos_cache, const void* sin_cache,
     rope_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
         (float*)x, (const float*)cos_cache, (const float*)sin_cache,
         (const int*)positions, num_heads, head_dim);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache,
@@ -104,6 +106,7 @@ void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache,
     rope_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (__nv_bfloat16*)x, (const float*)cos_cache, (const float*)sin_cache,
         (const int*)positions, num_heads, head_dim);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_compute_rope_cache(void* cos_cache, void* sin_cache,
@@ -111,6 +114,7 @@ void launch_compute_rope_cache(void* cos_cache, void* sin_cache,
                                void* stream) {
     compute_rope_cache<<<max_seq_len, half_dim, 0, (cudaStream_t)stream>>>(
         (float*)cos_cache, (float*)sin_cache, max_seq_len, half_dim, theta);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/embedding/transpose.cu

@@ -1,4 +1,5 @@
 #include <cuda_bf16.h>
+#include "../common.cuh"
 
 // Transpose between [S, H, D] and [H, S, D] layouts (used for RoPE and attention).
 // Also handles [S, H*D] → [H, S, D] (reshape_heads) and reverse (merge_heads).
@@ -169,6 +170,7 @@ void launch_reshape_heads_bf16(const void* in, void* out,
     int grid = (total + block - 1) / block;
     reshape_heads_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_merge_heads_bf16(const void* in, void* out,
@@ -178,6 +180,7 @@ void launch_merge_heads_bf16(const void* in, void* out,
     int grid = (total + block - 1) / block;
     merge_heads_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_transpose_hsd_to_shd_bf16(const void* in, void* out,
@@ -187,6 +190,7 @@ void launch_transpose_hsd_to_shd_bf16(const void* in, void* out,
     int grid = (total + block - 1) / block;
     transpose_hsd_to_shd_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_transpose_shd_to_hsd_bf16(const void* in, void* out,
@@ -196,6 +200,7 @@ void launch_transpose_shd_to_hsd_bf16(const void* in, void* out,
     int grid = (total + block - 1) / block;
     transpose_shd_to_hsd_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_repeat_kv_bf16(const void* in, void* out,
@@ -205,6 +210,7 @@ void launch_repeat_kv_bf16(const void* in, void* out,
     int grid = (total + block - 1) / block;
     repeat_kv_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)in, (__nv_bfloat16*)out, kv_heads, n_rep, seq_len, head_dim);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 void launch_strided_copy_bf16(const void* in, void* out, int numel, int ndim,
@@ -217,6 +223,7 @@ void launch_strided_copy_bf16(const void* in, void* out, int numel, int ndim,
         (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,
@@ -229,6 +236,7 @@ void launch_strided_copy_f32(const void* in, void* out, int numel, int ndim,
         (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

@@ -1,5 +1,6 @@
 #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]
@@ -88,6 +89,7 @@ void launch_gemv_bf16(
         (float*)y_fp32_buf,
         K, N
     );
+    CUDA_CHECK_LAST_ERROR();
 
     // Convert FP32 -> BF16
     int conv_block = 256;
@@ -97,6 +99,7 @@ void launch_gemv_bf16(
         (__nv_bfloat16*)y_bf16,
         N
     );
+    CUDA_CHECK_LAST_ERROR();
 }
 
 } // extern "C"

csrc/gemm/naive.cu

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

csrc/gemm/tiled.cu

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

csrc/normalization/layernorm.cu

@@ -105,6 +105,7 @@ void launch_layernorm_f32(const void* x, const void* gamma, const void* beta,
     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,
@@ -114,6 +115,7 @@ void launch_layernorm_bf16(const void* x, const void* gamma, const void* beta,
     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

@@ -111,6 +111,7 @@ void launch_rmsnorm_f32(const void* x, const void* gamma, void* out,
     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,
@@ -120,6 +121,7 @@ void launch_rmsnorm_bf16(const void* x, const void* gamma, void* out,
     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,
@@ -132,6 +134,7 @@ void launch_add_rmsnorm_bf16(const void* x, const void* residual, const void* ga
         (const __nv_bfloat16*)gamma,
         (__nv_bfloat16*)normed_out, (__nv_bfloat16*)sum_out,
         hidden_size, eps);
+    CUDA_CHECK_LAST_ERROR();
 }
 
 }

csrc/reduce/softmax.cu

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