Merge legacy Phase 18 branch history

Retain the superseded gpt-oss WIP ancestry without replacing the current implementation.

README.md

@@ -3,13 +3,15 @@
 > 从零用 **Rust + CUDA** 构建的 LLM 推理引擎，目标是吃透 LLM Serving 全栈技术。
 
 xserv 不依赖 PyTorch / vLLM / TensorRT 等现成框架，自己实现了张量抽象、CUDA kernel、
-分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。支持 **Qwen3-8B**（BF16）
-和 **gpt-oss-20b**（MoE，BF16/FP8/MXFP4 量化），多卡 TP/PP，并提供一套与 **llama.cpp**
+分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。支持 **Qwen3-8B**、
+**Qwen3.6-35B-A3B**（BF16）和 **gpt-oss-20b**（MoE，BF16/FP8/MXFP4 量化），多卡
+TP/PP，并提供一套与 **llama.cpp**
 对比正确性和性能的标准 benchmark。
 
 ## 现状一览
 
-- **模型**：GPT-2（124M）、Qwen3-8B（BF16）、gpt-oss-20b（32 专家 top-4 MoE，harmony 格式）
+- **模型**：GPT-2（124M）、Qwen3-8B（BF16）、Qwen3.6-35B-A3B（256 专家 top-8 MoE，BF16）、
+  gpt-oss-20b（32 专家 top-4 MoE，harmony 格式）
 - **性能**（RTX 5090，贪心，单流）：
   - Qwen3-8B BF16 单卡：约 56 tok/s（HF transformers 的 1.4×）
   - gpt-oss-20b FP8 稀疏 MoE + CUDA Graph decode：**TPOT 5.8ms（~172 tok/s，
@@ -110,6 +112,14 @@ curl http://localhost:8080/v1/chat/completions \
 
 其它端点：`GET /health`、`GET /v1/models`。
 
+Qwen3.6-35B-A3B 当前走 8 卡流水线并行（BF16，暂不支持单卡或 TP）：
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+  ./target/release/xserv-server /path/to/qwen3.6-35b-a3b \
+  --pp 8 --max-batch 1 --max-seq-len 4096
+```
+
 ### 2. 命令行推理
 
 ```bash

crates/xserv-kernels/build.rs

@@ -30,6 +30,7 @@ fn main() {
         .file("../../csrc/attention/flash_attention.cu")
         .file("../../csrc/attention/paged_attention.cu")
         .file("../../csrc/attention/reshape_and_cache.cu")
+        .file("../../csrc/attention/deltanet.cu")
         .file("../../csrc/moe/moe_kernels.cu")
         .file("../../csrc/moe/moe_sparse.cu")
         .file("../../csrc/quantization/dequant_fp8.cu")

crates/xserv-kernels/src/activation.rs

@@ -6,6 +6,10 @@ unsafe extern "C" {
     fn launch_gelu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
     fn launch_silu_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
     fn launch_silu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_sigmoid_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_sigmoid_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_softplus_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_softplus_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
     fn launch_scale_f32(
         x: *const c_void,
         out: *mut c_void,
@@ -48,6 +52,14 @@ unsafe extern "C" {
         n: i32,
         stream: *mut c_void,
     );
+    fn launch_row_scale_bf16(
+        x: *const c_void,
+        scale: *const c_void,
+        out: *mut c_void,
+        rows: i32,
+        cols: i32,
+        stream: *mut c_void,
+    );
     fn launch_silu_mul_bf16(
         gate: *const c_void,
         up: *const c_void,
@@ -151,6 +163,12 @@ pub fn gelu(x: &Tensor) -> Tensor {
 pub fn silu(x: &Tensor) -> Tensor {
     dispatch_unary(x, launch_silu_f32, launch_silu_bf16)
 }
+pub fn sigmoid(x: &Tensor) -> Tensor {
+    dispatch_unary(x, launch_sigmoid_f32, launch_sigmoid_bf16)
+}
+pub fn softplus(x: &Tensor) -> Tensor {
+    dispatch_unary(x, launch_softplus_f32, launch_softplus_bf16)
+}
 
 pub fn scale(x: &Tensor, scale_val: f32) -> Tensor {
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
@@ -190,6 +208,30 @@ pub fn mul(a: &Tensor, b: &Tensor) -> Tensor {
     dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16)
 }
 
+pub fn row_scale_bf16(x: &Tensor, scale: &Tensor) -> Tensor {
+    assert_eq!(x.ndim(), 2);
+    assert_eq!(scale.ndim(), 2);
+    assert_eq!(x.dtype(), DType::BF16);
+    assert_eq!(scale.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && scale.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    let rows = x.shape()[0];
+    let cols = x.shape()[1];
+    assert_eq!(scale.shape(), &[rows, 1]);
+    let out = Tensor::empty(&[rows, cols], DType::BF16, x.device());
+    unsafe {
+        launch_row_scale_bf16(
+            x.data_ptr() as _,
+            scale.data_ptr() as _,
+            out.data_ptr() as *mut c_void,
+            rows as i32,
+            cols as i32,
+            xserv_cuda::current_stream_raw(),
+        );
+    }
+    out
+}
+
 /// Row-broadcast bias add: out[r, c] = x[r, c] + bias[c] (BF16 only).
 pub fn bias_add_2d(x: &Tensor, bias: &Tensor) -> Tensor {
     assert_eq!(x.ndim(), 2);

crates/xserv-kernels/src/attention.rs

@@ -83,6 +83,24 @@ unsafe extern "C" {
         scale: f32,
         stream: *mut c_void,
     );
+    fn launch_paged_decode_attention_tree_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,
+        tree_mask: *const i32,
+        batch: i32,
+        num_q_heads: i32,
+        num_kv_heads: i32,
+        head_dim: i32,
+        max_blocks_per_seq: i32,
+        tree_start: i32,
+        tree_len: i32,
+        scale: f32,
+        stream: *mut c_void,
+    );
     fn launch_paged_decode_attention_sinks_bf16(
         q: *const c_void,
         k_cache: *const c_void,
@@ -127,6 +145,17 @@ unsafe extern "C" {
         max_blocks_per_seq: i32,
         stream: *mut c_void,
     );
+    fn launch_copy_kv_position(
+        k_pool: *mut c_void,
+        v_pool: *mut c_void,
+        block_ids: *const i32,
+        src_pos: i32,
+        dst_pos: i32,
+        num_kv_heads: i32,
+        head_dim: i32,
+        block_size: i32,
+        stream: *mut c_void,
+    );
 }
 
 /// Scatter `[num_kv_heads, num_tokens, head_dim]` BF16 K/V into a paged
@@ -213,6 +242,36 @@ pub unsafe fn reshape_and_cache_batched_bf16(
     }
 }
 
+/// Copy one token's K/V from `src_pos` to `dst_pos` within the same sequence's
+/// paged cache (one layer). Used by tree speculative decoding to remap
+/// accepted sibling K/V to canonical sequential positions after acceptance.
+///
+/// # Safety
+/// Pool and block_ids pointers must be valid GPU pointers for the given layer.
+pub unsafe fn copy_kv_position(
+    k_pool_ptr: *mut c_void,
+    v_pool_ptr: *mut c_void,
+    block_ids_gpu: *const i32,
+    src_pos: usize,
+    dst_pos: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    block_size: usize,
+    stream: *mut c_void,
+) {
+    launch_copy_kv_position(
+        k_pool_ptr,
+        v_pool_ptr,
+        block_ids_gpu,
+        src_pos as i32,
+        dst_pos as i32,
+        num_kv_heads as i32,
+        head_dim as i32,
+        block_size as i32,
+        stream,
+    );
+}
+
 fn apply_causal_mask(scores: &Tensor, offset: usize) {
     let ndim = scores.ndim();
     let rows = scores.shape()[ndim - 2];
@@ -353,8 +412,8 @@ pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tens
         "num_q_heads must be divisible by num_kv_heads"
     );
     assert!(
-        head_dim <= 128,
-        "flash_attention supports head_dim up to 128"
+        head_dim <= 256,
+        "flash_attention supports head_dim up to 256"
     );
 
     // Dispatch to specialized decode kernel for single-token generation
@@ -420,7 +479,7 @@ pub fn flash_attention_sinks(
     assert_eq!(k.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
     assert_eq!(v.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
     assert!(num_q_heads % num_kv_heads == 0);
-    assert!(head_dim <= 128);
+    assert!(head_dim <= 256);
     assert_eq!(
         sinks.shape()[0],
         num_q_heads,
@@ -489,7 +548,7 @@ pub fn paged_decode_attention(
         num_q_heads % num_kv_heads == 0,
         "GQA: num_q_heads must be divisible by num_kv_heads"
     );
-    assert!(head_dim <= 128);
+    assert!(head_dim <= 256);
 
     let scale = 1.0 / (head_dim as f32).sqrt();
     let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
@@ -515,6 +574,62 @@ pub fn paged_decode_attention(
     output
 }
 
+/// Tree-aware paged decode attention. Adds a per-query attention mask over
+/// the newly-written K/V region `[tree_start, tree_start+tree_len)`. Query i
+/// attends to position tree_start+j iff tree_mask[i, j] != 0. Positions <
+/// tree_start are always attended.
+///
+/// Used by speculative decoding with tree drafting to let sibling candidates
+/// share position slots without seeing each other's K/V.
+#[allow(clippy::too_many_arguments)]
+pub fn paged_decode_attention_tree(
+    q: &Tensor,
+    k_cache_ptr: *const c_void,
+    v_cache_ptr: *const c_void,
+    block_tables_ptr: *const i32,
+    context_lens_ptr: *const i32,
+    tree_mask_ptr: *const i32,
+    batch: usize,
+    num_q_heads: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    max_blocks_per_seq: usize,
+    tree_start: usize,
+    tree_len: usize,
+) -> Tensor {
+    assert_eq!(q.ndim(), 4);
+    assert_eq!(q.shape()[2], 1);
+    assert_eq!(q.dtype(), DType::BF16);
+    assert!(num_q_heads % num_kv_heads == 0);
+    assert!(head_dim <= 256);
+
+    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_tree_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,
+            tree_mask_ptr,
+            batch as i32,
+            num_q_heads as i32,
+            num_kv_heads as i32,
+            head_dim as i32,
+            max_blocks_per_seq as i32,
+            tree_start as i32,
+            tree_len as i32,
+            scale,
+            xserv_cuda::current_stream_raw(),
+        );
+    }
+
+    output
+}
+
 /// Paged decode attention with attention sinks and optional sliding window.
 ///
 /// sinks_ptr: pointer to [num_q_heads] BF16 on GPU (or null for no sinks)
@@ -538,7 +653,7 @@ pub fn paged_decode_attention_sinks(
     assert_eq!(q.shape()[2], 1);
     assert_eq!(q.dtype(), DType::BF16);
     assert!(num_q_heads % num_kv_heads == 0);
-    assert!(head_dim <= 128);
+    assert!(head_dim <= 256);
 
     let scale = 1.0 / (head_dim as f32).sqrt();
     let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());

crates/xserv-kernels/src/deltanet.rs

@@ -0,0 +1,278 @@
+use std::ffi::c_void;
+
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_depthwise_causal_conv1d_silu_bf16(
+        x: *const c_void,
+        w: *const c_void,
+        out: *mut c_void,
+        seq_len: i32,
+        channels: i32,
+        kernel: i32,
+        stream: *mut c_void,
+    );
+    fn launch_deltanet_ar_bf16(
+        q: *const c_void,
+        k: *const c_void,
+        v: *const c_void,
+        beta: *const c_void,
+        alpha_softplus: *const c_void,
+        a_log: *const c_void,
+        state: *mut c_void,
+        out: *mut c_void,
+        key_heads: i32,
+        value_heads: i32,
+        head_dim: i32,
+        stream: *mut c_void,
+    );
+    fn launch_depthwise_causal_conv1d_stateful_silu_bf16(
+        x: *const c_void,
+        w: *const c_void,
+        state: *mut c_void,
+        out: *mut c_void,
+        seq_len: i32,
+        channels: i32,
+        kernel: i32,
+        stream: *mut c_void,
+    );
+    fn launch_l2_normalize_rows_bf16(
+        x: *const c_void,
+        out: *mut c_void,
+        rows: i32,
+        cols: i32,
+        eps: f32,
+        stream: *mut c_void,
+    );
+    fn launch_deltanet_recurrent_bf16_f32(
+        q: *const c_void,
+        k: *const c_void,
+        v: *const c_void,
+        beta: *const c_void,
+        alpha_softplus: *const c_void,
+        a_log: *const c_void,
+        state: *mut c_void,
+        out: *mut c_void,
+        seq_len: i32,
+        key_heads: i32,
+        value_heads: i32,
+        head_dim: i32,
+        stream: *mut c_void,
+    );
+}
+
+fn conv_kernel_shape(w: &Tensor, channels: usize) -> usize {
+    match w.ndim() {
+        2 => {
+            assert_eq!(w.shape()[0], channels);
+            w.shape()[1]
+        }
+        3 => {
+            assert_eq!(w.shape()[0], channels);
+            assert_eq!(w.shape()[1], 1);
+            w.shape()[2]
+        }
+        _ => panic!("conv weight must be [C,K] or [C,1,K]"),
+    }
+}
+
+/// Depthwise causal Conv1D + SiLU for Qwen3.5/3.6 DeltaNet.
+/// x: [seq_len, channels] BF16, w: [channels, 1, kernel] or [channels, kernel] BF16.
+pub fn depthwise_causal_conv1d_silu_bf16(x: &Tensor, w: &Tensor) -> Tensor {
+    assert_eq!(x.ndim(), 2);
+    assert!(x.is_contiguous() && w.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    assert_eq!(x.dtype(), DType::BF16);
+    assert_eq!(w.dtype(), DType::BF16);
+    let seq_len = x.shape()[0];
+    let channels = x.shape()[1];
+    let kernel = conv_kernel_shape(w, channels);
+    let out = Tensor::empty(&[seq_len, channels], DType::BF16, x.device());
+    unsafe {
+        launch_depthwise_causal_conv1d_silu_bf16(
+            x.data_ptr() as *const c_void,
+            w.data_ptr() as *const c_void,
+            out.data_ptr() as *mut c_void,
+            seq_len as i32,
+            channels as i32,
+            kernel as i32,
+            xserv_cuda::current_stream_raw(),
+        );
+    }
+    out
+}
+
+/// Stateful depthwise Conv1D + SiLU. `state` is BF16 `[channels, kernel-1]`
+/// and is updated in-place with the newest raw projection values.
+pub fn depthwise_causal_conv1d_stateful_silu_bf16(
+    x: &Tensor,
+    w: &Tensor,
+    state: &mut Tensor,
+) -> Tensor {
+    assert_eq!(x.ndim(), 2);
+    assert!(x.is_contiguous() && w.is_contiguous() && state.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    assert_eq!(x.dtype(), DType::BF16);
+    assert_eq!(w.dtype(), DType::BF16);
+    assert_eq!(state.dtype(), DType::BF16);
+    let seq_len = x.shape()[0];
+    let channels = x.shape()[1];
+    let kernel = conv_kernel_shape(w, channels);
+    assert_eq!(state.shape(), &[channels, kernel - 1]);
+    let out = Tensor::empty(&[seq_len, channels], DType::BF16, x.device());
+    unsafe {
+        launch_depthwise_causal_conv1d_stateful_silu_bf16(
+            x.data_ptr() as *const c_void,
+            w.data_ptr() as *const c_void,
+            state.data_ptr() as *mut c_void,
+            out.data_ptr() as *mut c_void,
+            seq_len as i32,
+            channels as i32,
+            kernel as i32,
+            xserv_cuda::current_stream_raw(),
+        );
+    }
+    out
+}
+
+/// Row-wise L2 normalization with FP32 accumulation.
+pub fn l2_normalize_bf16(x: &Tensor, eps: f32) -> Tensor {
+    assert_eq!(x.ndim(), 2);
+    assert!(x.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    assert_eq!(x.dtype(), DType::BF16);
+    let rows = x.shape()[0];
+    let cols = x.shape()[1];
+    let out = Tensor::empty(x.shape(), DType::BF16, x.device());
+    unsafe {
+        launch_l2_normalize_rows_bf16(
+            x.data_ptr() as *const c_void,
+            out.data_ptr() as *mut c_void,
+            rows as i32,
+            cols as i32,
+            eps,
+            xserv_cuda::current_stream_raw(),
+        );
+    }
+    out
+}
+
+/// Stateful Qwen3.5/3.6 Gated DeltaNet for prefill or decode.
+/// q/k: `[seq,key_heads,head_dim]`, v: `[seq,value_heads,head_dim]`.
+/// beta/alpha: `[seq,value_heads]`; FP32 state is updated in-place.
+pub fn deltanet_recurrent_bf16_f32(
+    q: &Tensor,
+    k: &Tensor,
+    v: &Tensor,
+    beta: &Tensor,
+    alpha_softplus: &Tensor,
+    a_log: &Tensor,
+    state: &mut Tensor,
+) -> Tensor {
+    assert_eq!(q.ndim(), 3);
+    assert_eq!(k.ndim(), 3);
+    assert_eq!(v.ndim(), 3);
+    assert_eq!(beta.ndim(), 2);
+    assert_eq!(alpha_softplus.ndim(), 2);
+    assert_eq!(a_log.ndim(), 1);
+    assert_eq!(state.ndim(), 3);
+    assert!(q.is_contiguous() && k.is_contiguous() && v.is_contiguous());
+    assert!(beta.is_contiguous() && alpha_softplus.is_contiguous() && a_log.is_contiguous());
+    assert!(state.is_contiguous());
+    assert_eq!(q.dtype(), DType::BF16);
+    assert_eq!(k.dtype(), DType::BF16);
+    assert_eq!(v.dtype(), DType::BF16);
+    assert_eq!(beta.dtype(), DType::BF16);
+    assert_eq!(alpha_softplus.dtype(), DType::BF16);
+    assert_eq!(a_log.dtype(), DType::BF16);
+    assert_eq!(state.dtype(), DType::F32);
+    let seq_len = q.shape()[0];
+    let key_heads = q.shape()[1];
+    let head_dim = q.shape()[2];
+    let value_heads = v.shape()[1];
+    assert_eq!(k.shape(), &[seq_len, key_heads, head_dim]);
+    assert_eq!(v.shape(), &[seq_len, value_heads, head_dim]);
+    assert_eq!(beta.shape(), &[seq_len, value_heads]);
+    assert_eq!(alpha_softplus.shape(), &[seq_len, value_heads]);
+    assert_eq!(a_log.shape(), &[value_heads]);
+    assert_eq!(state.shape(), &[value_heads, head_dim, head_dim]);
+    assert_eq!(value_heads % key_heads, 0);
+    let out = Tensor::empty(
+        &[seq_len, value_heads, head_dim],
+        DType::BF16,
+        q.device(),
+    );
+    unsafe {
+        launch_deltanet_recurrent_bf16_f32(
+            q.data_ptr() as *const c_void,
+            k.data_ptr() as *const c_void,
+            v.data_ptr() as *const c_void,
+            beta.data_ptr() as *const c_void,
+            alpha_softplus.data_ptr() as *const c_void,
+            a_log.data_ptr() as *const c_void,
+            state.data_ptr() as *mut c_void,
+            out.data_ptr() as *mut c_void,
+            seq_len as i32,
+            key_heads as i32,
+            value_heads as i32,
+            head_dim as i32,
+            xserv_cuda::current_stream_raw(),
+        );
+    }
+    out
+}
+
+/// Single-token autoregressive DeltaNet state update for Qwen3.5/3.6.
+/// q/k: [key_heads, head_dim], v: [value_heads, head_dim], beta/alpha_softplus/a_log: [value_heads].
+/// state: [value_heads, head_dim, head_dim] updated in-place. Returns [1, value_heads * head_dim].
+pub fn deltanet_ar_bf16(
+    q: &Tensor,
+    k: &Tensor,
+    v: &Tensor,
+    beta: &Tensor,
+    alpha_softplus: &Tensor,
+    a_log: &Tensor,
+    state: &mut Tensor,
+) -> Tensor {
+    assert_eq!(q.ndim(), 2);
+    assert_eq!(k.ndim(), 2);
+    assert_eq!(v.ndim(), 2);
+    assert_eq!(beta.ndim(), 1);
+    assert_eq!(alpha_softplus.ndim(), 1);
+    assert_eq!(a_log.ndim(), 1);
+    assert_eq!(state.ndim(), 3);
+    assert!(q.is_contiguous() && k.is_contiguous() && v.is_contiguous());
+    assert!(beta.is_contiguous() && alpha_softplus.is_contiguous() && a_log.is_contiguous());
+    assert!(state.is_contiguous());
+    assert!(matches!(q.device(), Device::Cuda(_)));
+    assert_eq!(q.dtype(), DType::BF16);
+    assert_eq!(k.dtype(), DType::BF16);
+    assert_eq!(v.dtype(), DType::BF16);
+    let key_heads = q.shape()[0];
+    let head_dim = q.shape()[1];
+    let value_heads = v.shape()[0];
+    assert_eq!(k.shape(), &[key_heads, head_dim]);
+    assert_eq!(v.shape()[1], head_dim);
+    assert_eq!(beta.shape(), &[value_heads]);
+    assert_eq!(alpha_softplus.shape(), &[value_heads]);
+    assert_eq!(a_log.shape(), &[value_heads]);
+    assert_eq!(state.shape(), &[value_heads, head_dim, head_dim]);
+    let out = Tensor::empty(&[1, value_heads * head_dim], DType::BF16, q.device());
+    unsafe {
+        launch_deltanet_ar_bf16(
+            q.data_ptr() as *const c_void,
+            k.data_ptr() as *const c_void,
+            v.data_ptr() as *const c_void,
+            beta.data_ptr() as *const c_void,
+            alpha_softplus.data_ptr() as *const c_void,
+            a_log.data_ptr() as *const c_void,
+            state.data_ptr() as *mut c_void,
+            out.data_ptr() as *mut c_void,
+            key_heads as i32,
+            value_heads as i32,
+            head_dim as i32,
+            xserv_cuda::current_stream_raw(),
+        );
+    }
+    out
+}

crates/xserv-kernels/src/gemm.rs

@@ -5,6 +5,7 @@ use xserv_cuda::error::{self, Result};
 use xserv_tensor::{DType, Device, Tensor};
 
 const CUBLAS_WORKSPACE_BYTES: usize = 32 * 1024 * 1024;
+const GEMV_TILE_K: usize = 256;
 
 // GEMV: single-kernel, no FP32 temp buffer needed
 unsafe extern "C" {
@@ -17,6 +18,17 @@ unsafe extern "C" {
         n: i32,
         stream: *mut c_void,
     );
+
+    fn launch_gemv_bf16_batched(
+        x: *const c_void,
+        w: *const c_void,
+        y_bf16: *mut c_void,
+        y_fp32_buf: *mut c_void,
+        m: i32,
+        k: i32,
+        n: i32,
+        stream: *mut c_void,
+    );
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -26,6 +38,59 @@ pub enum GemmBackend {
     CuBlas,
 }
 
+pub fn gemv_scratch_elems(k: usize, n: usize) -> usize {
+    n * k.div_ceil(GEMV_TILE_K)
+}
+
+/// Batched GEMV: [M, K] × [K, N] → [M, N], all BF16.
+/// Bit-exact with calling matmul on each row individually (same K-block partial
+/// + fixed-order reduction path), but in a single kernel launch per phase.
+pub fn matmul_batched_gemv(a: &Tensor, b: &Tensor) -> Tensor {
+    assert_eq!(a.ndim(), 2);
+    assert_eq!(b.ndim(), 2);
+    assert!(a.is_contiguous());
+    assert!(b.is_contiguous());
+    assert_eq!(a.dtype(), DType::BF16);
+    assert_eq!(b.dtype(), DType::BF16);
+    let m = a.shape()[0];
+    let k = a.shape()[1];
+    let n = b.shape()[1];
+    assert_eq!(b.shape()[0], k);
+
+    let out = Tensor::empty(&[m, n], DType::BF16, a.device());
+    let scratch_elems = m * gemv_scratch_elems(k, n);
+    let mut fp32_buf = xserv_cuda::allocator::cached_alloc(scratch_elems * 4).unwrap();
+
+    let null_stream = xserv_cuda::current_stream_raw();
+    if m == 1 {
+        unsafe {
+            launch_gemv_bf16(
+                a.data_ptr() as *const c_void,
+                b.data_ptr() as *const c_void,
+                out.data_ptr() as *mut c_void,
+                fp32_buf.as_mut_ptr() as *mut c_void,
+                k as i32,
+                n as i32,
+                null_stream,
+            );
+        }
+    } else {
+        unsafe {
+            launch_gemv_bf16_batched(
+                a.data_ptr() as *const c_void,
+                b.data_ptr() as *const c_void,
+                out.data_ptr() as *mut c_void,
+                fp32_buf.as_mut_ptr() as *mut c_void,
+                m as i32,
+                k as i32,
+                n as i32,
+                null_stream,
+            );
+        }
+    }
+    out
+}
+
 // --- FFI: custom CUDA kernels ---
 unsafe extern "C" {
     fn launch_gemm_naive_f32(
@@ -274,7 +339,8 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
         },
         GemmBackend::CuBlas => {
             if m == 1 && dtype == DType::BF16 && n >= 256 {
-                let mut fp32_buf = xserv_cuda::allocator::cached_alloc(n * 4).unwrap();
+                let mut fp32_buf =
+                    xserv_cuda::allocator::cached_alloc(gemv_scratch_elems(k, n) * 4).unwrap();
                 unsafe {
                     launch_gemv_bf16(
                         a_ptr,

crates/xserv-kernels/src/lib.rs

@@ -1,6 +1,7 @@
 pub mod activation;
 pub mod argmax;
 pub mod attention;
+pub mod deltanet;
 pub mod dispatch;
 pub mod embedding;
 pub mod gemm;
@@ -12,17 +13,25 @@ pub mod rope;
 pub mod softmax;
 pub mod transpose;
 
-pub use activation::{add, bias_add_2d, gelu, gpt_oss_glu, mul, scale, silu, silu_mul};
+pub use activation::{
+    add, bias_add_2d, gelu, gpt_oss_glu, mul, row_scale_bf16, scale, sigmoid, silu, silu_mul,
+    softplus,
+};
 pub use argmax::{argmax_bf16_single, argmax_bf16_to_host};
 pub use attention::{
-    attention, decode_attention, flash_attention, flash_attention_sinks, paged_decode_attention,
-    paged_decode_attention_sinks, reshape_and_cache_batched_bf16, reshape_and_cache_bf16,
+    attention, copy_kv_position, decode_attention, flash_attention, flash_attention_sinks,
+    paged_decode_attention, paged_decode_attention_sinks, paged_decode_attention_tree,
+    reshape_and_cache_batched_bf16, reshape_and_cache_bf16,
+};
+pub use deltanet::{
+    deltanet_ar_bf16, deltanet_recurrent_bf16_f32, depthwise_causal_conv1d_silu_bf16,
+    depthwise_causal_conv1d_stateful_silu_bf16, l2_normalize_bf16,
 };
 pub use embedding::{embedding, embedding_device_ids};
-pub use gemm::{GemmBackend, batched_matmul, matmul};
+pub use gemm::{GemmBackend, batched_matmul, matmul, matmul_batched_gemv};
 pub use layernorm::layernorm;
 pub use rmsnorm::{add_rmsnorm, rmsnorm};
-pub use rope::{RopeCache, rope_inplace, rope_inplace_device_pos};
+pub use rope::{RopeCache, partial_rope_bf16, rope_inplace, rope_inplace_device_pos};
 pub use softmax::softmax;
 pub use transpose::{
     merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu,

crates/xserv-kernels/src/moe.rs

@@ -42,6 +42,21 @@ unsafe extern "C" {
         stream: *mut c_void,
     );
 
+    fn launch_moe_sparse_gemv_bf16_bf16(
+        x: *const c_void,
+        w: *const c_void,
+        topk_ids: *const c_void,
+        y: *mut c_void,
+        num_tokens: i32,
+        n: i32,
+        k: i32,
+        top_k: i32,
+        expert_start: i32,
+        local_experts: i32,
+        x_per_slot: i32,
+        stream: *mut c_void,
+    );
+
     fn launch_moe_sparse_gemv_fp8_bf16(
         x: *const c_void,
         w: *const c_void,
@@ -244,6 +259,52 @@ pub fn moe_weighted_sum(
     out
 }
 
+/// Sparse MoE GEMV (BF16 weights): compute only routed experts.
+/// x: [num_tokens, K] or [num_tokens * top_k, K], w_t: [local_experts, N, K].
+pub fn moe_sparse_gemv_bf16(
+    x: &Tensor,
+    w_t: &Tensor,
+    topk_ids: &Tensor,
+    top_k: usize,
+    expert_start: usize,
+    local_experts: usize,
+    x_per_slot: bool,
+) -> Tensor {
+    assert_eq!(x.ndim(), 2);
+    assert_eq!(w_t.ndim(), 3);
+    assert_eq!(x.dtype(), DType::BF16);
+    assert_eq!(w_t.dtype(), DType::BF16);
+    assert!(x.is_contiguous() && w_t.is_contiguous());
+    let local = w_t.shape()[0];
+    let n = w_t.shape()[1];
+    let k = w_t.shape()[2];
+    assert_eq!(local, local_experts);
+    assert_eq!(x.shape()[1], k);
+    let num_tokens = if x_per_slot {
+        x.shape()[0] / top_k
+    } else {
+        x.shape()[0]
+    };
+    let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
+    unsafe {
+        launch_moe_sparse_gemv_bf16_bf16(
+            x.data_ptr() as *const c_void,
+            w_t.data_ptr() as *const c_void,
+            topk_ids.data_ptr() as *const c_void,
+            y.data_ptr() as *mut c_void,
+            num_tokens as i32,
+            n as i32,
+            k as i32,
+            top_k as i32,
+            expert_start as i32,
+            local_experts as i32,
+            if x_per_slot { 1 } else { 0 },
+            xserv_cuda::current_stream_raw(),
+        );
+    }
+    y
+}
+
 /// Sparse MoE GEMV (FP8 W8A16): compute only the routed experts.
 ///
 /// x:        [num_tokens, K] BF16 (x_per_slot=false, gate_up) or

crates/xserv-kernels/src/rope.rs

@@ -23,6 +23,17 @@ unsafe extern "C" {
         head_dim: i32,
         stream: *mut c_void,
     );
+    fn launch_partial_rope_bf16(
+        x: *const c_void,
+        out: *mut c_void,
+        positions: *const c_void,
+        num_tokens: i32,
+        num_heads: i32,
+        head_dim: i32,
+        n_rot: i32,
+        theta: f32,
+        stream: *mut c_void,
+    );
     fn launch_compute_rope_cache(
         cos_cache: *mut c_void,
         sin_cache: *mut c_void,
@@ -171,6 +182,47 @@ pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) {
     rope_inplace_device_pos(x, cache, pos_gpu.as_ptr() as *const c_void);
 }
 
+/// Apply partial RoPE and return a new tensor.
+/// x: [num_tokens, num_heads, head_dim] BF16 on GPU. Only the first `n_rot`
+/// dimensions are rotated; the tail is copied unchanged.
+pub fn partial_rope_bf16(x: &Tensor, positions: &[u32], n_rot: usize, theta: f32) -> Tensor {
+    assert_eq!(x.ndim(), 3);
+    assert!(x.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    assert_eq!(x.dtype(), DType::BF16);
+    let num_tokens = x.shape()[0];
+    let num_heads = x.shape()[1];
+    let head_dim = x.shape()[2];
+    assert_eq!(positions.len(), num_tokens);
+    assert!(n_rot <= head_dim && n_rot % 2 == 0);
+
+    let pos_bytes = unsafe {
+        std::slice::from_raw_parts(
+            positions.as_ptr() as *const u8,
+            num_tokens * std::mem::size_of::<u32>(),
+        )
+    };
+    let mut pos_gpu =
+        xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions");
+    pos_gpu.copy_from_host(pos_bytes).unwrap();
+    let out = Tensor::empty(x.shape(), DType::BF16, x.device());
+
+    unsafe {
+        launch_partial_rope_bf16(
+            x.data_ptr() as *const c_void,
+            out.data_ptr() as *mut c_void,
+            pos_gpu.as_ptr() as *const c_void,
+            num_tokens as i32,
+            num_heads as i32,
+            head_dim as i32,
+            n_rot as i32,
+            theta,
+            xserv_cuda::current_stream_raw(),
+        );
+    }
+    out
+}
+
 /// RoPE in-place with positions already on the GPU (u32, [num_tokens]).
 /// Used by the CUDA-graph decode path, where the position lives in a
 /// persistent device buffer updated outside the captured region.

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

@@ -0,0 +1,976 @@
+//! Draft-model speculative decoding benchmark for Qwen3.
+//!
+//! v0 scope:
+//! - target + draft are Qwen3-family models with the same tokenizer/vocab;
+//! - batch=1;
+//! - greedy exact-match acceptance;
+//! - no probabilistic rejection sampling.
+
+use half::bf16;
+use std::path::{Path, PathBuf};
+use std::time::Instant;
+
+use xserv_model::qwen3_graph::GraphedQwen3Decoder;
+use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, loader};
+use xserv_tensor::{DType, Device, Tensor};
+use xserv_tokenizer::Tokenizer;
+
+const DEFAULT_GAMMA: usize = 4;
+const DEFAULT_GEN_TOKENS: usize = 64;
+const DEFAULT_MAX_SEQ_LEN: usize = 2048;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum VerifyPath {
+    Flash,
+    PagedDecode,
+}
+
+impl VerifyPath {
+    fn as_str(self) -> &'static str {
+        match self {
+            VerifyPath::Flash => "flash",
+            VerifyPath::PagedDecode => "paged-decode",
+        }
+    }
+}
+
+const PROMPTS: [&str; 50] = [
+    "The capital of France is",
+    "Once upon a time in a land far away",
+    "Hello, how are you doing today",
+    "In a shocking finding, scientists discovered a",
+    "The weather today is sunny, so I decided to",
+    "Alan Turing was a British mathematician who",
+    "The best way to learn programming is",
+    "Artificial intelligence will change the world because",
+    "The history of the internet began in the",
+    "A good morning routine starts with",
+    "The stock market crashed because investors",
+    "Deep learning is a subset of machine learning that",
+    "The president of the United States announced",
+    "In the year 2050, humans will",
+    "The secret to happiness is",
+    "When I was a child, I used to",
+    "The most important scientific discovery of the century",
+    "Climate change is caused by",
+    "The recipe for chocolate cake requires",
+    "In conclusion, the evidence suggests that",
+    "The cat sat on the mat and",
+    "According to recent studies, exercise can",
+    "The first step in solving any problem is",
+    "Technology has transformed the way we",
+    "The novel begins with the protagonist",
+    "Education is the most powerful weapon",
+    "The ocean covers more than seventy percent of",
+    "Last night I had a dream about",
+    "The company announced its quarterly earnings",
+    "Music has the power to",
+    "The difference between success and failure is",
+    "In the beginning, there was nothing but",
+    "The doctor told me that I should",
+    "Python is a popular programming language because",
+    "The ancient Romans built roads that",
+    "A balanced diet should include",
+    "The movie received mixed reviews from critics",
+    "Space exploration has led to many",
+    "The teacher asked the students to",
+    "Global warming is one of the most",
+    "The bridge collapsed due to structural",
+    "Quantum computing promises to revolutionize",
+    "The new policy will affect millions of",
+    "During the winter months, it is important to",
+    "The human brain contains approximately",
+    "Democracy depends on the active participation of",
+    "The train arrived at the station exactly",
+    "Researchers at MIT have developed a new",
+    "The smartphone has become an essential part of",
+    "After careful consideration, the committee decided to",
+];
+
+#[derive(Default)]
+struct RunStats {
+    ids: Vec<u32>,
+    total_s: f64,
+    prefill_s: f64,
+    decode_s: f64,
+    target_steps: usize,
+    accepted: usize,
+    proposed: usize,
+    verify_steps: usize,
+    mirror_steps: usize,
+    commit_steps: usize,
+    correction_steps: usize,
+    verify_decode_mismatches: usize,
+}
+
+#[derive(Default)]
+struct Totals {
+    prompts: usize,
+    baseline_generated: usize,
+    spec_generated: usize,
+    baseline_total_s: f64,
+    baseline_prefill_s: f64,
+    baseline_decode_s: f64,
+    spec_total_s: f64,
+    spec_prefill_s: f64,
+    spec_decode_s: f64,
+    spec_target_steps: usize,
+    spec_accepted: usize,
+    spec_proposed: usize,
+    spec_verify_steps: usize,
+    spec_mirror_steps: usize,
+    spec_commit_steps: usize,
+    spec_correction_steps: usize,
+    spec_verify_decode_mismatches: usize,
+    mismatches: usize,
+}
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 3 {
+        eprintln!(
+            "Usage: bench-speculative <target-model-dir> <draft-model-dir> \
+             [--gen-tokens N] [--gamma N] [--prompts N] [--max-seq-len N] [--device N] \
+             [--use-verify-logits] [--verify-path flash|paged-decode] [--dump-verify-mismatches]"
+        );
+        std::process::exit(1);
+    }
+
+    let target_dir = PathBuf::from(&args[1]);
+    let draft_dir = PathBuf::from(&args[2]);
+    let gen_tokens = arg_usize(&args, "--gen-tokens", DEFAULT_GEN_TOKENS);
+    let gamma = arg_usize(&args, "--gamma", DEFAULT_GAMMA);
+    let prompt_count = arg_usize(&args, "--prompts", PROMPTS.len()).min(PROMPTS.len());
+    let max_seq_len = arg_usize(&args, "--max-seq-len", DEFAULT_MAX_SEQ_LEN);
+    let device = arg_usize(&args, "--device", 0) as u32;
+    let use_verify_logits = args.iter().any(|a| a == "--use-verify-logits");
+    let verify_path = parse_verify_path(&args, use_verify_logits);
+    let dump_verify_mismatches = args.iter().any(|a| a == "--dump-verify-mismatches");
+
+    assert!(gen_tokens > 0, "--gen-tokens must be > 0");
+    assert!(gamma > 0, "--gamma must be > 0");
+
+    xserv_cuda::device::set_device(device).unwrap();
+    let info = xserv_cuda::device::device_info(device).unwrap();
+    eprintln!(
+        "GPU {device}: {} ({} MB free)",
+        info.name,
+        info.free_memory / 1024 / 1024
+    );
+
+    let target_config = ModelConfig::from_file(&target_dir.join("config.json"));
+    let draft_config = ModelConfig::from_file(&draft_dir.join("config.json"));
+    assert_qwen3(&target_config, "target");
+    assert_qwen3(&draft_config, "draft");
+    assert_eq!(
+        target_config.vocab_size, draft_config.vocab_size,
+        "target and draft vocab_size must match"
+    );
+
+    warn_if_tokenizers_differ(&target_dir, &draft_dir);
+    let tokenizer = Tokenizer::from_file(&target_dir.join("tokenizer.json"));
+    if tokenizer.vocab_size() != target_config.vocab_size {
+        eprintln!(
+            "WARNING: tokenizer decoder len {} differs from config vocab_size {}; continuing because token ids come from the shared tokenizer.json",
+            tokenizer.vocab_size(),
+            target_config.vocab_size
+        );
+    }
+
+    eprintln!(
+        "Loading target Qwen3: layers={} hidden={} heads={}/{} vocab={}",
+        target_config.num_layers(),
+        target_config.hidden(),
+        target_config.num_heads(),
+        target_config.num_kv_heads(),
+        target_config.vocab_size
+    );
+    let target_weights = loader::load_model_dir(&target_dir, Device::Cuda(device));
+    let target = Qwen3::from_weights(target_config.clone(), target_weights);
+    xserv_cuda::allocator::cached_trim();
+
+    eprintln!(
+        "Loading draft Qwen3: layers={} hidden={} heads={}/{} vocab={}",
+        draft_config.num_layers(),
+        draft_config.hidden(),
+        draft_config.num_heads(),
+        draft_config.num_kv_heads(),
+        draft_config.vocab_size
+    );
+    let draft_weights = loader::load_model_dir(&draft_dir, Device::Cuda(device));
+    let draft = Qwen3::from_weights(draft_config.clone(), draft_weights);
+    xserv_cuda::allocator::cached_trim();
+
+    let warm_ids = tokenizer.encode("warmup");
+    let warm_tokens = gen_tokens.min(4);
+    {
+        let mut target_cache = new_cache(&target_config, max_seq_len, device);
+        let _ = run_baseline(
+            &target,
+            &mut target_cache,
+            &tokenizer,
+            &warm_ids,
+            warm_tokens,
+        );
+    }
+    {
+        let mut target_cache = new_cache_with_rows(
+            &target_config,
+            max_seq_len,
+            device,
+            if use_verify_logits { gamma } else { 1 },
+        );
+        let mut target_verify_cache =
+            new_cache_with_rows(&target_config, max_seq_len, device, gamma);
+        let mut draft_cache = new_cache(&draft_config, max_seq_len, device);
+        let mut draft_decoder = GraphedQwen3Decoder::new();
+        let _ = run_speculative(
+            &target,
+            &draft,
+            &mut target_cache,
+            &mut target_verify_cache,
+            &mut draft_cache,
+            &mut draft_decoder,
+            &tokenizer,
+            &warm_ids,
+            warm_tokens,
+            gamma,
+            use_verify_logits,
+            verify_path,
+            dump_verify_mismatches,
+        );
+    }
+    eprintln!(
+        "Warmup done. Running {prompt_count} prompts, gen_tokens={gen_tokens}, gamma={gamma}, acceptance_mode={}, verify_path={}",
+        if use_verify_logits {
+            "verify_logits"
+        } else {
+            "decode"
+        },
+        verify_path.as_str()
+    );
+
+    let mut totals = Totals::default();
+
+    // Persistent per-benchmark caches so the draft CUDA graph (Phase 24) can be
+    // captured once and replayed across every prompt. Freeing and re-registering
+    // slot 0 between prompts keeps block_table_gpu / context_lens_gpu addresses
+    // stable, which is exactly what the graph captured.
+    let mut target_cache = new_cache_with_rows(
+        &target_config,
+        max_seq_len,
+        device,
+        if use_verify_logits { gamma } else { 1 },
+    );
+    let mut target_verify_cache = new_cache_with_rows(&target_config, max_seq_len, device, gamma);
+    let mut draft_cache = new_cache(&draft_config, max_seq_len, device);
+    let mut draft_decoder = GraphedQwen3Decoder::new();
+
+    for (i, prompt) in PROMPTS.iter().take(prompt_count).enumerate() {
+        let ids = tokenizer.encode(prompt);
+        validate_length_budget(&ids, gen_tokens, max_seq_len, prompt);
+        let mut baseline_cache = new_cache(&target_config, max_seq_len, device);
+        let baseline = run_baseline(&target, &mut baseline_cache, &tokenizer, &ids, gen_tokens);
+        drop(baseline_cache);
+
+        let spec = run_speculative(
+            &target,
+            &draft,
+            &mut target_cache,
+            &mut target_verify_cache,
+            &mut draft_cache,
+            &mut draft_decoder,
+            &tokenizer,
+            &ids,
+            gen_tokens,
+            gamma,
+            use_verify_logits,
+            verify_path,
+            dump_verify_mismatches,
+        );
+
+        let matched = baseline.ids == spec.ids;
+        if !matched {
+            totals.mismatches += 1;
+            eprintln!("MISMATCH prompt {i}: {prompt}");
+            eprintln!("  baseline: {:?}", baseline.ids);
+            eprintln!("  spec:     {:?}", spec.ids);
+        }
+
+        println!(
+            "prompt={:02} match={} gen={} accept={}/{} target_steps={} \
+             baseline_e2e_tpot_ms={:.3} spec_e2e_tpot_ms={:.3}",
+            i,
+            matched,
+            spec.ids.len(),
+            spec.accepted,
+            spec.proposed,
+            spec.target_steps,
+            per_token_ms(baseline.total_s, baseline.ids.len()),
+            per_token_ms(spec.total_s, spec.ids.len()),
+        );
+
+        totals.prompts += 1;
+        totals.baseline_generated += baseline.ids.len();
+        totals.spec_generated += spec.ids.len();
+        totals.baseline_total_s += baseline.total_s;
+        totals.baseline_prefill_s += baseline.prefill_s;
+        totals.baseline_decode_s += baseline.decode_s;
+        totals.spec_total_s += spec.total_s;
+        totals.spec_prefill_s += spec.prefill_s;
+        totals.spec_decode_s += spec.decode_s;
+        totals.spec_target_steps += spec.target_steps;
+        totals.spec_accepted += spec.accepted;
+        totals.spec_proposed += spec.proposed;
+        totals.spec_verify_steps += spec.verify_steps;
+        totals.spec_mirror_steps += spec.mirror_steps;
+        totals.spec_commit_steps += spec.commit_steps;
+        totals.spec_correction_steps += spec.correction_steps;
+        totals.spec_verify_decode_mismatches += spec.verify_decode_mismatches;
+    }
+
+    let baseline_decode_tokens = totals.baseline_generated;
+    let spec_decode_tokens = totals.spec_generated;
+    let acceptance = ratio(totals.spec_accepted, totals.spec_proposed);
+    let tokens_per_target_step = ratio(totals.spec_generated, totals.spec_target_steps);
+    let matched =
+        totals.mismatches == 0 && (!use_verify_logits || totals.spec_verify_decode_mismatches == 0);
+
+    println!("--- SUMMARY ---");
+    println!("prompts={} matched={matched}", totals.prompts);
+    println!(
+        "acceptance_mode={}",
+        if use_verify_logits {
+            "verify_logits"
+        } else {
+            "decode"
+        }
+    );
+    println!("verify_path={}", verify_path.as_str());
+    println!(
+        "acceptance_rate={:.4} accepted={} proposed={}",
+        acceptance, totals.spec_accepted, totals.spec_proposed
+    );
+    println!(
+        "tokens_per_target_step={:.4} target_steps={} verify_steps={} mirror_decode_steps={} commit_decode_steps={} correction_steps={}",
+        tokens_per_target_step,
+        totals.spec_target_steps,
+        totals.spec_verify_steps,
+        totals.spec_mirror_steps,
+        totals.spec_commit_steps,
+        totals.spec_correction_steps
+    );
+    println!(
+        "verify_decode_mismatches={}",
+        totals.spec_verify_decode_mismatches
+    );
+    println!(
+        "baseline_e2e_tpot_ms={:.3} baseline_e2e_tok_s={:.3}",
+        per_token_ms(totals.baseline_total_s, totals.baseline_generated),
+        tok_s(totals.baseline_generated, totals.baseline_total_s)
+    );
+    println!(
+        "spec_e2e_tpot_ms={:.3} spec_e2e_tok_s={:.3} speedup_e2e={:.4}",
+        per_token_ms(totals.spec_total_s, totals.spec_generated),
+        tok_s(totals.spec_generated, totals.spec_total_s),
+        speedup(totals.baseline_total_s, totals.spec_total_s)
+    );
+    println!(
+        "baseline_decode_tpot_ms={:.3} baseline_decode_tok_s={:.3}",
+        per_token_ms(totals.baseline_decode_s, baseline_decode_tokens),
+        tok_s(baseline_decode_tokens, totals.baseline_decode_s)
+    );
+    println!(
+        "spec_decode_tpot_ms={:.3} spec_decode_tok_s={:.3} speedup_decode={:.4}",
+        per_token_ms(totals.spec_decode_s, spec_decode_tokens),
+        tok_s(spec_decode_tokens, totals.spec_decode_s),
+        speedup(totals.baseline_decode_s, totals.spec_decode_s)
+    );
+    println!(
+        "decode_token_counts baseline={} spec={}",
+        baseline_decode_tokens, spec_decode_tokens
+    );
+
+    if !matched {
+        std::process::exit(2);
+    }
+}
+
+fn run_baseline(
+    model: &Qwen3,
+    cache: &mut PagedKVCache,
+    tokenizer: &Tokenizer,
+    prompt_ids: &[u32],
+    gen_tokens: usize,
+) -> RunStats {
+    let slot = 0;
+    cache.register_sequence(slot).expect("register target slot");
+
+    let t0 = Instant::now();
+    let prefill_start = Instant::now();
+    let logits = model.forward_prefill_paged(prompt_ids, slot, cache);
+    sync_device();
+    let prefill_s = prefill_start.elapsed().as_secs_f64();
+
+    let mut generated = Vec::with_capacity(gen_tokens);
+    let mut next = last_argmax(&logits);
+    generated.push(next);
+
+    let decode_start = Instant::now();
+    let mut target_steps = 0usize;
+    while generated.len() < gen_tokens && !tokenizer.is_eos(next) {
+        let pos = cache.seq_len(slot);
+        let logits = model.forward_decode_paged(&[next], &[pos], &[slot], cache);
+        target_steps += 1;
+        next = last_argmax(&logits);
+        generated.push(next);
+    }
+    sync_device();
+    let decode_s = decode_start.elapsed().as_secs_f64();
+    sync_device();
+    let total_s = t0.elapsed().as_secs_f64();
+
+    cache.free_sequence(slot);
+    RunStats {
+        ids: generated,
+        total_s,
+        prefill_s,
+        decode_s,
+        target_steps,
+        ..Default::default()
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn run_speculative(
+    target: &Qwen3,
+    draft: &Qwen3,
+    target_cache: &mut PagedKVCache,
+    target_verify_cache: &mut PagedKVCache,
+    draft_cache: &mut PagedKVCache,
+    draft_decoder: &mut GraphedQwen3Decoder,
+    tokenizer: &Tokenizer,
+    prompt_ids: &[u32],
+    gen_tokens: usize,
+    gamma: usize,
+    use_verify_logits: bool,
+    verify_path: VerifyPath,
+    dump_verify_mismatches: bool,
+) -> RunStats {
+    let slot = 0;
+    target_cache
+        .register_sequence(slot)
+        .expect("register target slot");
+    target_verify_cache
+        .register_sequence(slot)
+        .expect("register target verify slot");
+    draft_cache
+        .register_sequence(slot)
+        .expect("register draft slot");
+
+    let t0 = Instant::now();
+    let prefill_start = Instant::now();
+    let target_logits = target.forward_prefill_paged(prompt_ids, slot, target_cache);
+    if !use_verify_logits {
+        let _ = target.forward_prefill_paged(prompt_ids, slot, target_verify_cache);
+    }
+    let draft_logits = draft.forward_prefill_paged(prompt_ids, slot, draft_cache);
+    sync_device();
+    let prefill_s = prefill_start.elapsed().as_secs_f64();
+
+    let mut target_next = last_argmax(&target_logits);
+    let mut draft_next = last_argmax(&draft_logits);
+    let mut generated = Vec::with_capacity(gen_tokens);
+    let mut accepted_total = 0usize;
+    let mut proposed_total = 0usize;
+    let mut verify_steps = 0usize;
+    let mut mirror_steps = 0usize;
+    let mut commit_steps = 0usize;
+    let mut correction_steps = 0usize;
+    let mut verify_decode_mismatches = 0usize;
+
+    let decode_start = Instant::now();
+    while generated.len() < gen_tokens {
+        let remaining = gen_tokens - generated.len();
+        let round_gamma = gamma.min(remaining);
+        let round_start_len = target_cache.seq_len(slot);
+        assert_eq!(
+            round_start_len,
+            draft_cache.seq_len(slot),
+            "target and draft cache lengths diverged"
+        );
+        if !use_verify_logits {
+            assert_eq!(
+                round_start_len,
+                target_verify_cache.seq_len(slot),
+                "target verify cache length diverged"
+            );
+        }
+
+        let mut draft_tokens = Vec::with_capacity(round_gamma);
+        for _ in 0..round_gamma {
+            let token = draft_next;
+            draft_tokens.push(token);
+            if tokenizer.is_eos(token) {
+                break;
+            }
+            let pos = draft_cache.seq_len(slot);
+            let logits = draft_decoder.decode(draft, &[token], &[pos], &[slot], draft_cache);
+            draft_next = last_argmax(&logits);
+        }
+        proposed_total += draft_tokens.len();
+
+        if use_verify_logits {
+            verify_steps += 1;
+            let verify_logits =
+                target.forward_verify_paged_decode_attention(&draft_tokens, slot, target_cache);
+            let verify_argmax = argmax_rows(&verify_logits);
+            assert_eq!(
+                verify_argmax.len(),
+                draft_tokens.len(),
+                "verify logits rows must match draft token count"
+            );
+
+            let mut accepted = 0usize;
+            let mut done = false;
+            while accepted < draft_tokens.len() {
+                let expected = if accepted > 0 {
+                    verify_argmax[accepted - 1]
+                } else {
+                    target_next
+                };
+                if draft_tokens[accepted] != expected {
+                    break;
+                }
+                let token = draft_tokens[accepted];
+                generated.push(token);
+                accepted_total += 1;
+                accepted += 1;
+
+                if generated.len() >= gen_tokens || tokenizer.is_eos(token) {
+                    done = true;
+                    break;
+                }
+            }
+
+            if accepted > 0 {
+                target_next = verify_argmax[accepted - 1];
+            }
+            target_cache
+                .truncate_sequence(slot, round_start_len + accepted)
+                .unwrap();
+
+            if done {
+                draft_cache
+                    .truncate_sequence(slot, target_cache.seq_len(slot))
+                    .unwrap();
+                break;
+            }
+
+            if accepted == draft_tokens.len() {
+                continue;
+            }
+
+            let correction = if accepted > 0 {
+                verify_argmax[accepted - 1]
+            } else {
+                target_next
+            };
+            generated.push(correction);
+
+            draft_cache
+                .truncate_sequence(slot, round_start_len)
+                .unwrap();
+            replay_draft_tokens(
+                draft,
+                draft_decoder,
+                draft_cache,
+                slot,
+                &draft_tokens[..accepted],
+                &mut draft_next,
+            );
+
+            if generated.len() >= gen_tokens || tokenizer.is_eos(correction) {
+                break;
+            }
+
+            let pos = target_cache.seq_len(slot);
+            let logits = target.forward_decode_paged(&[correction], &[pos], &[slot], target_cache);
+            target_next = last_argmax(&logits);
+            commit_steps += 1;
+
+            let pos = draft_cache.seq_len(slot);
+            let logits = draft_decoder.decode(draft, &[correction], &[pos], &[slot], draft_cache);
+            draft_next = last_argmax(&logits);
+            correction_steps += 1;
+            continue;
+        }
+
+        verify_steps += 1;
+        let verify_logits = match verify_path {
+            VerifyPath::Flash => {
+                target.forward_prefill_paged(&draft_tokens, slot, target_verify_cache)
+            }
+            VerifyPath::PagedDecode => target.forward_verify_paged_decode_attention(
+                &draft_tokens,
+                slot,
+                target_verify_cache,
+            ),
+        };
+        let verify_argmax = argmax_rows(&verify_logits);
+        assert_eq!(
+            verify_argmax.len(),
+            draft_tokens.len(),
+            "verify logits rows must match draft token count"
+        );
+
+        target_verify_cache
+            .truncate_sequence(slot, round_start_len)
+            .unwrap();
+
+        let mut accepted = 0usize;
+        let mut done = false;
+        while accepted < draft_tokens.len() {
+            let expected = if use_verify_logits && accepted > 0 {
+                verify_argmax[accepted - 1]
+            } else {
+                target_next
+            };
+            if draft_tokens[accepted] != expected {
+                break;
+            }
+            let token_idx = accepted;
+            let token = draft_tokens[token_idx];
+            generated.push(token);
+            accepted_total += 1;
+            accepted += 1;
+
+            if generated.len() >= gen_tokens || tokenizer.is_eos(token) {
+                done = true;
+                break;
+            }
+
+            let pos = target_cache.seq_len(slot);
+            let logits = target.forward_decode_paged(&[token], &[pos], &[slot], target_cache);
+            let decode_next = last_argmax(&logits);
+            if verify_argmax[token_idx] != decode_next {
+                verify_decode_mismatches += 1;
+                eprintln!(
+                    "VERIFY/DECODE MISMATCH at cache_len={} accepted_idx={}: verify={} decode={}",
+                    target_cache.seq_len(slot),
+                    token_idx,
+                    verify_argmax[token_idx],
+                    decode_next
+                );
+                if dump_verify_mismatches {
+                    eprintln!(
+                        "  verify_top5={} decode_top5={}",
+                        format_topk(&verify_logits, token_idx, 5),
+                        format_topk(&logits, 0, 5)
+                    );
+                }
+            }
+            target_next = decode_next;
+            commit_steps += 1;
+
+            advance_target_cache(target, target_verify_cache, slot, token);
+            mirror_steps += 1;
+        }
+        if done {
+            draft_cache
+                .truncate_sequence(slot, target_cache.seq_len(slot))
+                .unwrap();
+            target_verify_cache
+                .truncate_sequence(slot, target_cache.seq_len(slot))
+                .unwrap();
+            break;
+        }
+
+        if accepted == draft_tokens.len() {
+            continue;
+        }
+
+        let correction = if use_verify_logits && accepted > 0 {
+            verify_argmax[accepted - 1]
+        } else {
+            target_next
+        };
+        generated.push(correction);
+
+        draft_cache
+            .truncate_sequence(slot, round_start_len)
+            .unwrap();
+        replay_draft_tokens(
+            draft,
+            draft_decoder,
+            draft_cache,
+            slot,
+            &draft_tokens[..accepted],
+            &mut draft_next,
+        );
+
+        if generated.len() >= gen_tokens || tokenizer.is_eos(correction) {
+            break;
+        }
+
+        let pos = target_cache.seq_len(slot);
+        let logits = target.forward_decode_paged(&[correction], &[pos], &[slot], target_cache);
+        target_next = last_argmax(&logits);
+        commit_steps += 1;
+
+        advance_target_cache(target, target_verify_cache, slot, correction);
+        mirror_steps += 1;
+
+        let pos = draft_cache.seq_len(slot);
+        let logits = draft_decoder.decode(draft, &[correction], &[pos], &[slot], draft_cache);
+        draft_next = last_argmax(&logits);
+        correction_steps += 1;
+    }
+    sync_device();
+    let decode_s = decode_start.elapsed().as_secs_f64();
+    sync_device();
+    let total_s = t0.elapsed().as_secs_f64();
+
+    target_cache.free_sequence(slot);
+    target_verify_cache.free_sequence(slot);
+    draft_cache.free_sequence(slot);
+
+    RunStats {
+        ids: generated,
+        total_s,
+        prefill_s,
+        decode_s,
+        target_steps: verify_steps + mirror_steps + commit_steps + correction_steps,
+        accepted: accepted_total,
+        proposed: proposed_total,
+        verify_steps,
+        mirror_steps,
+        commit_steps,
+        correction_steps,
+        verify_decode_mismatches,
+    }
+}
+
+fn advance_target_cache(target: &Qwen3, cache: &mut PagedKVCache, slot: usize, token: u32) {
+    let pos = cache.seq_len(slot);
+    let _ = target.forward_decode_paged(&[token], &[pos], &[slot], cache);
+}
+
+fn replay_draft_tokens(
+    draft: &Qwen3,
+    draft_decoder: &mut GraphedQwen3Decoder,
+    cache: &mut PagedKVCache,
+    slot: usize,
+    tokens: &[u32],
+    next: &mut u32,
+) {
+    for &token in tokens {
+        let pos = cache.seq_len(slot);
+        let logits = draft_decoder.decode(draft, &[token], &[pos], &[slot], cache);
+        *next = last_argmax(&logits);
+    }
+}
+
+fn new_cache(config: &ModelConfig, max_seq_len: usize, device: u32) -> PagedKVCache {
+    new_cache_with_rows(config, max_seq_len, device, 1)
+}
+
+fn new_cache_with_rows(
+    config: &ModelConfig,
+    max_seq_len: usize,
+    device: u32,
+    max_rows: usize,
+) -> PagedKVCache {
+    let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE);
+    let total_blocks = max_blocks_per_seq + 8;
+    PagedKVCache::new(
+        config,
+        total_blocks,
+        0,
+        max_rows.max(1),
+        max_blocks_per_seq,
+        DType::BF16,
+        device,
+    )
+}
+
+fn argmax_rows(logits: &Tensor) -> Vec<u32> {
+    assert_eq!(logits.ndim(), 2);
+    if logits.dtype() == DType::BF16
+        && matches!(logits.device(), Device::Cuda(_))
+        && logits.is_contiguous()
+    {
+        return xserv_kernels::argmax_bf16_to_host(logits);
+    }
+
+    let vocab_size = logits.shape()[1];
+    let rows = logits.shape()[0];
+    let logits_cpu = logits.to_device(Device::Cpu);
+    match logits.dtype() {
+        DType::F32 => logits_cpu
+            .as_slice::<f32>()
+            .chunks_exact(vocab_size)
+            .take(rows)
+            .map(argmax_f32)
+            .collect(),
+        DType::BF16 => logits_cpu
+            .as_slice::<bf16>()
+            .chunks_exact(vocab_size)
+            .take(rows)
+            .map(|row| {
+                row.iter()
+                    .enumerate()
+                    .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
+                    .map(|(i, _)| i as u32)
+                    .unwrap()
+            })
+            .collect(),
+        _ => panic!("unsupported dtype for argmax: {:?}", logits.dtype()),
+    }
+}
+
+fn last_argmax(logits: &Tensor) -> u32 {
+    *argmax_rows(logits).last().unwrap()
+}
+
+fn argmax_f32(row: &[f32]) -> u32 {
+    row.iter()
+        .enumerate()
+        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
+        .map(|(i, _)| i as u32)
+        .unwrap()
+}
+
+fn format_topk(logits: &Tensor, row: usize, k: usize) -> String {
+    let vals = topk_row(logits, row, k);
+    vals.iter()
+        .map(|(id, val)| format!("{id}:{val:.3}"))
+        .collect::<Vec<_>>()
+        .join(",")
+}
+
+fn topk_row(logits: &Tensor, row: usize, k: usize) -> Vec<(u32, f32)> {
+    assert_eq!(logits.ndim(), 2);
+    let vocab_size = logits.shape()[1];
+    assert!(row < logits.shape()[0], "topk row out of bounds");
+    let logits_cpu = logits.to_device(Device::Cpu);
+    let mut vals: Vec<(u32, f32)> = match logits.dtype() {
+        DType::F32 => logits_cpu.as_slice::<f32>()[row * vocab_size..(row + 1) * vocab_size]
+            .iter()
+            .enumerate()
+            .map(|(i, &v)| (i as u32, v))
+            .collect(),
+        DType::BF16 => logits_cpu.as_slice::<bf16>()[row * vocab_size..(row + 1) * vocab_size]
+            .iter()
+            .enumerate()
+            .map(|(i, &v)| (i as u32, v.to_f32()))
+            .collect(),
+        _ => panic!("unsupported dtype for topk: {:?}", logits.dtype()),
+    };
+    vals.select_nth_unstable_by(k, |a, b| b.1.partial_cmp(&a.1).unwrap());
+    vals.truncate(k);
+    vals.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+    vals
+}
+
+fn assert_qwen3(config: &ModelConfig, name: &str) {
+    let model_type = config.model_type.as_deref().unwrap_or("unknown");
+    assert!(
+        model_type.contains("qwen"),
+        "{name} model_type must be qwen-like, got {model_type}"
+    );
+}
+
+fn warn_if_tokenizers_differ(target_dir: &Path, draft_dir: &Path) {
+    let target = std::fs::read(target_dir.join("tokenizer.json"));
+    let draft = std::fs::read(draft_dir.join("tokenizer.json"));
+    if let (Ok(target), Ok(draft)) = (target, draft) {
+        if target != draft {
+            eprintln!(
+                "WARNING: target and draft tokenizer.json differ; v0 assumes identical token ids"
+            );
+        }
+    }
+}
+
+fn arg_usize(args: &[String], flag: &str, default: usize) -> usize {
+    args.iter()
+        .position(|a| a == flag)
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(default)
+}
+
+fn parse_verify_path(args: &[String], use_verify_logits: bool) -> VerifyPath {
+    let default = if use_verify_logits {
+        VerifyPath::PagedDecode
+    } else {
+        VerifyPath::Flash
+    };
+    let Some(value) = args
+        .iter()
+        .position(|a| a == "--verify-path")
+        .and_then(|i| args.get(i + 1))
+    else {
+        return default;
+    };
+    match value.as_str() {
+        "flash" => VerifyPath::Flash,
+        "paged-decode" => VerifyPath::PagedDecode,
+        _ => {
+            eprintln!("unknown --verify-path {value:?}; expected flash or paged-decode");
+            std::process::exit(1);
+        }
+    }
+}
+
+fn validate_length_budget(prompt_ids: &[u32], gen_tokens: usize, max_seq_len: usize, prompt: &str) {
+    let required = prompt_ids.len() + gen_tokens;
+    if required > max_seq_len {
+        eprintln!(
+            "prompt requires prompt_len({}) + gen_tokens({}) = {} tokens, exceeding --max-seq-len {}: {:?}",
+            prompt_ids.len(),
+            gen_tokens,
+            required,
+            max_seq_len,
+            prompt
+        );
+        std::process::exit(2);
+    }
+}
+
+fn sync_device() {
+    xserv_cuda::device::synchronize().expect("cuda device synchronize");
+}
+
+fn ratio(num: usize, den: usize) -> f64 {
+    if den == 0 {
+        0.0
+    } else {
+        num as f64 / den as f64
+    }
+}
+
+fn speedup(baseline_s: f64, spec_s: f64) -> f64 {
+    if spec_s == 0.0 {
+        0.0
+    } else {
+        baseline_s / spec_s
+    }
+}
+
+fn tok_s(tokens: usize, seconds: f64) -> f64 {
+    if seconds == 0.0 {
+        0.0
+    } else {
+        tokens as f64 / seconds
+    }
+}
+
+fn per_token_ms(seconds: f64, tokens: usize) -> f64 {
+    if tokens == 0 {
+        0.0
+    } else {
+        seconds * 1000.0 / tokens as f64
+    }
+}

crates/xserv-model/src/bin/bench-verify-cost.rs

@@ -0,0 +1,134 @@
+//! Micro-benchmark: measure the cost of forward_verify_paged_decode_attention
+//! at different batch sizes (γ+1 values), to understand where speedup comes
+//! from (or doesn't).
+
+use std::path::PathBuf;
+use std::time::Instant;
+
+use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, loader};
+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-verify-cost <target-dir> [--prompt-len N] [--iters N] [--device N]"
+        );
+        std::process::exit(1);
+    }
+    let target_dir = PathBuf::from(&args[1]);
+    let prompt_len = arg_usize(&args, "--prompt-len", 100);
+    let iters = arg_usize(&args, "--iters", 30);
+    let device = arg_usize(&args, "--device", 0) as u32;
+
+    xserv_cuda::device::set_device(device).unwrap();
+
+    let cfg = ModelConfig::from_file(&target_dir.join("config.json"));
+    eprintln!("Loading target...");
+    let weights = loader::load_model_dir(&target_dir, Device::Cuda(device));
+    let target = Qwen3::from_weights(cfg.clone(), weights);
+    xserv_cuda::allocator::cached_trim();
+
+    let tok = Tokenizer::from_file(&target_dir.join("tokenizer.json"));
+    let ids = tok.encode(&"the ".repeat(prompt_len))[..prompt_len].to_vec();
+
+    let max_seq_len = 2048;
+    let num_blocks = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE + 4;
+    let mut cache = PagedKVCache::new(&cfg, num_blocks, 0, 16, num_blocks, DType::BF16, device);
+    cache.register_sequence(0).unwrap();
+
+    // Prefill
+    let _ = target.forward_prefill_paged(&ids, 0, &mut cache);
+    sync();
+
+    // Warmup one of each
+    for &n in &[1, 2, 3, 5, 9] {
+        let toks: Vec<u32> = (0..n).map(|_| ids[0]).collect();
+        let _ = target.forward_decode_paged(
+            &toks,
+            &(0..n).map(|i| ids.len() + i).collect::<Vec<_>>(),
+            &vec![0; n],
+            &mut cache,
+        );
+        cache.truncate_sequence(0, ids.len()).unwrap();
+    }
+    sync();
+
+    // Benchmark single-token decode
+    let mut t = 0.0f64;
+    for i in 0..iters {
+        cache.truncate_sequence(0, ids.len()).unwrap();
+        let t0 = Instant::now();
+        let _ = target.forward_decode_paged(&[ids[0]], &[ids.len()], &[0], &mut cache);
+        sync();
+        t += t0.elapsed().as_secs_f64();
+        let _ = i;
+    }
+    let single = t * 1000.0 / iters as f64;
+    println!(
+        "single-token decode: {:.3} ms (mean of {} iters)",
+        single, iters
+    );
+
+    // Benchmark forward_verify_paged_decode_attention at various batch sizes
+    // (batched-GEMV path).
+    for &n in &[1usize, 2, 3, 5, 9] {
+        let toks: Vec<u32> = (0..n).map(|_| ids[0]).collect();
+        let mut t = 0.0f64;
+        for _ in 0..iters {
+            cache.truncate_sequence(0, ids.len()).unwrap();
+            let t0 = Instant::now();
+            let _ = target.forward_verify_paged_decode_attention(&toks, 0, &mut cache);
+            sync();
+            t += t0.elapsed().as_secs_f64();
+        }
+        let ms = t * 1000.0 / iters as f64;
+        println!(
+            "verify (batched-GEMV) batch={}: {:.3} ms ({:.2}× single)",
+            n,
+            ms,
+            ms / single
+        );
+    }
+
+    // Benchmark _with_hidden variant which uses cuBLAS GEMM after Phase 26 fast-verify.
+    let hooks_layers = [2usize, 18, 33];
+    for &n in &[1usize, 2, 3, 5, 9] {
+        let toks: Vec<u32> = (0..n).map(|_| ids[0]).collect();
+        let mut t = 0.0f64;
+        for _ in 0..iters {
+            cache.truncate_sequence(0, ids.len()).unwrap();
+            let t0 = Instant::now();
+            let _ = target.forward_verify_paged_decode_attention_with_hidden(
+                &toks,
+                0,
+                &mut cache,
+                &hooks_layers,
+            );
+            sync();
+            t += t0.elapsed().as_secs_f64();
+        }
+        let ms = t * 1000.0 / iters as f64;
+        println!(
+            "verify (cuBLAS GEMM) batch={}: {:.3} ms ({:.2}× single)",
+            n,
+            ms,
+            ms / single
+        );
+    }
+
+    cache.free_sequence(0);
+}
+
+fn sync() {
+    xserv_cuda::device::synchronize().unwrap();
+}
+
+fn arg_usize(args: &[String], flag: &str, default: usize) -> usize {
+    args.iter()
+        .position(|a| a == flag)
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(default)
+}

crates/xserv-model/src/bin/check-eagle3.rs

@@ -0,0 +1,174 @@
+//! EAGLE3 sanity check: load weights, run one draft step, print top-5 predictions.
+//!
+//! This verifies that:
+//! - Eagle3Head weights load without shape mismatches
+//! - Target hidden states can be captured via decode_core_with_hidden
+//! - Eagle3Head::step produces a valid token id (in target vocab)
+//!
+//! Does NOT measure speedup — that requires a full γ≥2 speculative loop, which
+//! is more complex integration work.
+
+use std::path::PathBuf;
+
+use xserv_model::eagle3::{EAGLE_HOOK_LAYERS, Eagle3Head};
+use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, loader};
+use xserv_tensor::{DType, Device, Tensor};
+use xserv_tokenizer::Tokenizer;
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 3 {
+        eprintln!("Usage: check-eagle3 <target-model-dir> <eagle3-model-dir> [prompt]");
+        std::process::exit(1);
+    }
+    let target_dir = PathBuf::from(&args[1]);
+    let eagle_dir = PathBuf::from(&args[2]);
+    let prompt = args
+        .get(3)
+        .cloned()
+        .unwrap_or_else(|| "The capital of France is".to_string());
+    let device: u32 = 0;
+
+    xserv_cuda::device::set_device(device).unwrap();
+
+    let target_config = ModelConfig::from_file(&target_dir.join("config.json"));
+    eprintln!("Loading target Qwen3-8B...");
+    let target_weights = loader::load_model_dir(&target_dir, Device::Cuda(device));
+    let target = Qwen3::from_weights(target_config.clone(), target_weights);
+    xserv_cuda::allocator::cached_trim();
+
+    eprintln!("Loading EAGLE3 head from {}", eagle_dir.display());
+    let mut eagle = Eagle3Head::load(&eagle_dir, device);
+    xserv_cuda::allocator::cached_trim();
+
+    let tokenizer = Tokenizer::from_file(&target_dir.join("tokenizer.json"));
+    let embed_tokens = target.embed_tokens_tensor();
+
+    let ids = tokenizer.encode(&prompt);
+    let max_seq_len = 512;
+
+    let num_blocks = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE + 2;
+    let mut cache = PagedKVCache::new(
+        &target_config,
+        num_blocks,
+        0,
+        1,
+        num_blocks,
+        DType::BF16,
+        device,
+    );
+    cache.register_sequence(0).unwrap();
+
+    // Prefill target.
+    let logits = target.forward_prefill_paged(&ids, 0, &mut cache);
+    let target_first = *xserv_kernels::argmax_bf16_to_host(&logits).last().unwrap();
+    let target_first_text = tokenizer.decode(&[target_first]);
+    println!("Prompt: {:?}", prompt);
+    println!(
+        "Target argmax after prefill: {} ({:?})",
+        target_first, target_first_text
+    );
+
+    // Now run one target decode step with target_first to get hidden states at the
+    // hook layers.
+    let pos = cache.seq_len(0);
+    target.decode_prepare(&[pos], &[0], &mut cache);
+    let ids_gpu = upload_u32(&[target_first]);
+    let pos_gpu = upload_u32(&[pos as u32]);
+    let (target_next_logits, hooks) = target.decode_core_with_hidden(
+        ids_gpu.as_ptr() as *const std::ffi::c_void,
+        pos_gpu.as_ptr() as *const std::ffi::c_void,
+        1,
+        &[0],
+        &mut cache,
+        &EAGLE_HOOK_LAYERS,
+    );
+    let target_next = xserv_kernels::argmax_bf16_single(&target_next_logits);
+    let target_next_text = tokenizer.decode(&[target_next]);
+    println!(
+        "Target argmax after 1 decode step: {} ({:?})",
+        target_next, target_next_text
+    );
+
+    for (i, h) in hooks.iter().enumerate() {
+        println!(
+            "hook[{}] (layer {}): shape={:?} dtype={:?}",
+            i,
+            EAGLE_HOOK_LAYERS[i],
+            h.shape(),
+            h.dtype()
+        );
+    }
+
+    // Ask EAGLE what it thinks the NEXT token is (given target_first as prev_token
+    // and the hidden states from the position where target_first lives).
+    // EAGLE should predict target_next (or close to it) to be useful.
+    eagle.reset();
+    let (eagle_pred, eagle_logits) = eagle.step(&hooks, embed_tokens, target_first, pos);
+    let eagle_pred_text = tokenizer.decode(&[eagle_pred]);
+    println!(
+        "EAGLE draft prediction (pairing A: prev=target_first): {} ({:?})",
+        eagle_pred, eagle_pred_text
+    );
+
+    if eagle_pred == target_next {
+        println!("MATCH: EAGLE agrees with target on next token.");
+    } else {
+        println!(
+            "MISMATCH: EAGLE draft={} vs target={} (this is fine per-step; check top-5 below)",
+            eagle_pred, target_next
+        );
+    }
+
+    // Show top-5 from eagle logits (in draft vocab space, mapped to target).
+    print_top5(
+        &eagle_logits,
+        "EAGLE draft top-5 (pairing A)",
+        &eagle,
+        &tokenizer,
+    );
+
+    // Alternative pairing B: pair hooks with target_next (the token those hooks produced
+    // via lm_head), predict token after target_next. Position advances by 1.
+    eagle.reset();
+    let (eagle_pred_b, eagle_logits_b) = eagle.step(&hooks, embed_tokens, target_next, pos + 1);
+    let eagle_pred_b_text = tokenizer.decode(&[eagle_pred_b]);
+    println!(
+        "\nEAGLE draft prediction (pairing B: prev=target_next): {} ({:?})",
+        eagle_pred_b, eagle_pred_b_text
+    );
+    print_top5(
+        &eagle_logits_b,
+        "EAGLE draft top-5 (pairing B)",
+        &eagle,
+        &tokenizer,
+    );
+}
+
+fn upload_u32(vals: &[u32]) -> xserv_cuda::GpuBuffer {
+    let bytes = unsafe { std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4) };
+    let mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).unwrap();
+    buf.copy_from_host(bytes).unwrap();
+    buf
+}
+
+fn print_top5(logits: &Tensor, label: &str, eagle: &Eagle3Head, tokenizer: &Tokenizer) {
+    use half::bf16;
+    let cpu = logits.to_device(Device::Cpu);
+    let data = cpu.as_slice::<bf16>();
+    let mut vals: Vec<(usize, f32)> = data
+        .iter()
+        .enumerate()
+        .map(|(i, v)| (i, v.to_f32()))
+        .collect();
+    vals.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+    println!("{label}:");
+    for (i, val) in vals.iter().take(5) {
+        let target_id = eagle.map_draft_to_target(*i as u32);
+        let text = tokenizer.decode(&[target_id]);
+        println!(
+            "  draft_id={} target_id={} val={:.3} text={:?}",
+            i, target_id, val, text
+        );
+    }
+}

crates/xserv-model/src/bin/check-qwen35-moe.rs

@@ -0,0 +1,280 @@
+use std::collections::{BTreeMap, BTreeSet};
+use std::fs::File;
+use std::io::{Read, Seek, SeekFrom};
+use std::path::{Path, PathBuf};
+
+use serde_json::Value;
+use xserv_model::{ModelConfig, ModelFamily, Qwen35MoeSpec, Qwen35TensorMeta, Qwen35TensorSpec};
+use xserv_tokenizer::Tokenizer;
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: check-qwen35-moe <metadata-model-dir> [safetensors-shard-dir]");
+        std::process::exit(1);
+    }
+    let model_dir = PathBuf::from(&args[1]);
+    let shard_dir = args
+        .get(2)
+        .map(PathBuf::from)
+        .unwrap_or_else(|| model_dir.clone());
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    if config.family() != ModelFamily::Qwen35Moe {
+        eprintln!(
+            "expected Qwen3.5/3.6 MoE config, got {}",
+            config.family().as_str()
+        );
+        std::process::exit(2);
+    }
+
+    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+    let index_path = model_dir.join("model.safetensors.index.json");
+    let index_text = std::fs::read_to_string(&index_path)
+        .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
+    let index: Value = serde_json::from_str(&index_text)
+        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
+    let weight_map = index
+        .get("weight_map")
+        .and_then(|v| v.as_object())
+        .expect("model.safetensors.index.json must contain weight_map");
+    let keys: BTreeSet<String> = weight_map.keys().cloned().collect();
+
+    println!("family={}", config.family().as_str());
+    println!("model_type={}", config.model_type_str());
+    println!("layers={}", config.num_layers());
+    println!("hidden={}", config.hidden());
+    println!(
+        "heads={} kv_heads={} head_dim={}",
+        config.num_heads(),
+        config.num_kv_heads(),
+        config.head_dim()
+    );
+    println!(
+        "vocab_size={} tokenizer_vocab={}",
+        config.vocab_size(),
+        tokenizer.vocab_size()
+    );
+    println!("max_seq_len={}", config.max_seq_len());
+    let text = config.text();
+    println!(
+        "linear_dims key_heads={:?} value_heads={:?} key_dim={:?} value_dim={:?} conv_kernel={:?}",
+        text.linear_num_key_heads,
+        text.linear_num_value_heads,
+        text.linear_key_head_dim,
+        text.linear_value_head_dim,
+        text.linear_conv_kernel_dim
+    );
+    println!(
+        "rope partial={:?} params={:?}",
+        text.partial_rotary_factor,
+        text.rope_parameters
+            .as_ref()
+            .and_then(|rp| rp.mrope_section.clone())
+    );
+    println!("mtp_layers={:?}", text.mtp_num_hidden_layers);
+    println!(
+        "experts={} top_k={} moe_intermediate={}",
+        config.num_experts(),
+        config.experts_per_token(),
+        config.ffn_hidden()
+    );
+    println!("tensor_count={}", keys.len());
+    let shard_names: BTreeSet<&str> = weight_map.values().filter_map(|v| v.as_str()).collect();
+    println!("shard_count={}", shard_names.len());
+    if let Some(total_size) = index
+        .pointer("/metadata/total_size")
+        .and_then(|v| v.as_f64())
+    {
+        println!("total_size_gb={:.2}", total_size / 1e9);
+    }
+
+    let mut layer_types = BTreeMap::<String, usize>::new();
+    for i in 0..config.num_layers() {
+        let kind = if config.is_linear_attention_layer(i) {
+            "linear_attention"
+        } else {
+            "full_attention"
+        };
+        *layer_types.entry(kind.to_string()).or_default() += 1;
+    }
+    println!("layer_types={layer_types:?}");
+
+    let mut missing = Vec::new();
+    let spec = Qwen35MoeSpec::from_config(&config);
+    println!(
+        "full_attention_path={}",
+        spec.describe_full_attention_path()
+    );
+    let expected = spec.expected_tensor_specs();
+    for tensor in &expected {
+        if !keys.contains(&tensor.name) {
+            missing.push(tensor.name.clone());
+        }
+    }
+
+    let visual = keys
+        .iter()
+        .filter(|k| k.starts_with("model.visual."))
+        .count();
+    let language = keys
+        .iter()
+        .filter(|k| k.starts_with("model.language_model."))
+        .count();
+    println!("language_tensors={language}");
+    println!("visual_tensors={visual}");
+
+    if missing.is_empty() {
+        println!("schema_check=ok");
+    } else {
+        println!("schema_check=missing {}", missing.len());
+        for key in missing.iter().take(50) {
+            println!("missing={key}");
+        }
+        std::process::exit(3);
+    }
+
+    validate_headers(&shard_dir, weight_map, &expected);
+    let tensor_meta = collect_header_meta(&shard_dir, weight_map, &expected);
+    let model_meta = spec
+        .build_meta(&tensor_meta)
+        .unwrap_or_else(|e| panic!("failed to build Qwen35ModelMeta: {e}"));
+    let meta_linear = model_meta
+        .layers
+        .iter()
+        .filter(|layer| {
+            matches!(
+                layer.attention,
+                xserv_model::qwen35_moe::Qwen35AttentionMeta::Linear(_)
+            )
+        })
+        .count();
+    let meta_full = model_meta.layers.len() - meta_linear;
+    println!("model_meta_layers={}", model_meta.layers.len());
+    println!("model_meta_linear_layers={meta_linear}");
+    println!("model_meta_full_attention_layers={meta_full}");
+}
+
+fn collect_header_meta(
+    shard_dir: &Path,
+    weight_map: &serde_json::Map<String, Value>,
+    expected: &[Qwen35TensorSpec],
+) -> BTreeMap<String, Qwen35TensorMeta> {
+    let mut by_shard = BTreeMap::<String, Vec<&Qwen35TensorSpec>>::new();
+    for spec in expected {
+        if let Some(shard) = weight_map.get(&spec.name).and_then(|v| v.as_str()) {
+            by_shard.entry(shard.to_string()).or_default().push(spec);
+        }
+    }
+    let mut out = BTreeMap::new();
+    for (shard, specs) in by_shard {
+        let path = shard_dir.join(&shard);
+        if !path.exists() {
+            continue;
+        }
+        let header = read_safetensors_header(&path);
+        for spec in specs {
+            if let Some(meta) = header.get(&spec.name) {
+                out.insert(
+                    spec.name.clone(),
+                    Qwen35TensorMeta {
+                        name: spec.name.clone(),
+                        dtype: meta
+                            .get("dtype")
+                            .and_then(|v| v.as_str())
+                            .unwrap_or("?")
+                            .to_string(),
+                        shape: meta
+                            .get("shape")
+                            .and_then(|v| v.as_array())
+                            .map(|arr| {
+                                arr.iter()
+                                    .filter_map(|v| v.as_u64().map(|n| n as usize))
+                                    .collect::<Vec<_>>()
+                            })
+                            .unwrap_or_default(),
+                    },
+                );
+            }
+        }
+    }
+    out
+}
+
+fn validate_headers(
+    shard_dir: &Path,
+    weight_map: &serde_json::Map<String, Value>,
+    expected: &[Qwen35TensorSpec],
+) {
+    let mut by_shard = BTreeMap::<String, Vec<&Qwen35TensorSpec>>::new();
+    for spec in expected {
+        if let Some(shard) = weight_map.get(&spec.name).and_then(|v| v.as_str()) {
+            by_shard.entry(shard.to_string()).or_default().push(spec);
+        }
+    }
+
+    let mut checked = 0usize;
+    let mut skipped_shards = 0usize;
+    let mut errors = Vec::new();
+    for (shard, specs) in by_shard {
+        let path = shard_dir.join(&shard);
+        if !path.exists() {
+            skipped_shards += 1;
+            continue;
+        }
+        let header = read_safetensors_header(&path);
+        for spec in specs {
+            checked += 1;
+            match header.get(&spec.name) {
+                Some(meta) => {
+                    let dtype = meta.get("dtype").and_then(|v| v.as_str()).unwrap_or("?");
+                    let shape = meta
+                        .get("shape")
+                        .and_then(|v| v.as_array())
+                        .map(|arr| {
+                            arr.iter()
+                                .filter_map(|v| v.as_u64().map(|n| n as usize))
+                                .collect::<Vec<_>>()
+                        })
+                        .unwrap_or_default();
+                    if dtype != spec.dtype || shape != spec.shape {
+                        errors.push(format!(
+                            "{} expected {} {:?}, got {} {:?}",
+                            spec.name, spec.dtype, spec.shape, dtype, shape
+                        ));
+                    }
+                }
+                None => errors.push(format!("{} missing from shard header {shard}", spec.name)),
+            }
+        }
+    }
+
+    println!("header_checked_tensors={checked}");
+    println!("header_skipped_missing_shards={skipped_shards}");
+    if errors.is_empty() {
+        println!("header_check=ok");
+    } else {
+        println!("header_check=errors {}", errors.len());
+        for err in errors.iter().take(50) {
+            println!("header_error={err}");
+        }
+        std::process::exit(4);
+    }
+}
+
+fn read_safetensors_header(path: &Path) -> serde_json::Map<String, Value> {
+    let mut file =
+        File::open(path).unwrap_or_else(|e| panic!("failed to open {}: {e}", path.display()));
+    let mut len_bytes = [0u8; 8];
+    file.read_exact(&mut len_bytes)
+        .unwrap_or_else(|e| panic!("failed to read header size from {}: {e}", path.display()));
+    let header_len = u64::from_le_bytes(len_bytes);
+    file.seek(SeekFrom::Start(8)).unwrap();
+    let mut header = vec![0u8; header_len as usize];
+    file.read_exact(&mut header)
+        .unwrap_or_else(|e| panic!("failed to read header from {}: {e}", path.display()));
+    serde_json::from_slice::<Value>(&header)
+        .unwrap_or_else(|e| panic!("failed to parse safetensors header {}: {e}", path.display()))
+        .as_object()
+        .cloned()
+        .expect("safetensors header must be a JSON object")
+}

crates/xserv-model/src/bin/smoke-qwen35-full-attn.rs

@@ -0,0 +1,195 @@
+use std::path::PathBuf;
+
+use half::bf16;
+use xserv_model::{ModelConfig, Qwen35Moe, Qwen35MoeSpec, loader};
+use xserv_tensor::{Device, Tensor};
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: smoke-qwen35-full-attn <model-dir> [layer] [seq_len]");
+        std::process::exit(1);
+    }
+    let model_dir = PathBuf::from(&args[1]);
+    let layer: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(3);
+    let seq_len: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(2);
+    let device = 0;
+
+    xserv_cuda::device::set_device(device).unwrap();
+    xserv_model::init_kernels();
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    let spec = Qwen35MoeSpec::from_config(&config);
+    assert_eq!(
+        spec.layer_kinds[layer],
+        xserv_model::qwen35_moe::Qwen35LayerKind::FullAttention,
+        "layer {layer} is not a full-attention layer"
+    );
+
+    let weight_map = read_weight_map(&model_dir);
+    let keys = [
+        format!("model.language_model.layers.{layer}.self_attn.q_proj.weight"),
+        format!("model.language_model.layers.{layer}.self_attn.k_proj.weight"),
+        format!("model.language_model.layers.{layer}.self_attn.v_proj.weight"),
+        format!("model.language_model.layers.{layer}.self_attn.o_proj.weight"),
+        format!("model.language_model.layers.{layer}.self_attn.q_norm.weight"),
+        format!("model.language_model.layers.{layer}.self_attn.k_norm.weight"),
+    ];
+    let mut shards: Vec<String> = keys
+        .iter()
+        .filter_map(|k| weight_map.get(k).cloned())
+        .collect();
+    shards.sort();
+    shards.dedup();
+
+    let mut tensors = std::collections::HashMap::new();
+    for shard in &shards {
+        tensors.extend(loader::load_safetensors(
+            &model_dir.join(shard),
+            Device::Cpu,
+        ));
+    }
+
+    let weights = Qwen35Moe::take_full_attention_weights(&mut tensors, layer, Device::Cuda(device));
+
+    let input = make_deterministic_input(seq_len, config.hidden()).to_device(Device::Cuda(device));
+    let n_rot =
+        (spec.head_dim as f64 * config.text().partial_rotary_factor.unwrap_or(0.25)) as usize;
+    let rope_theta = config.rope_theta_value().unwrap_or(10_000_000.0) as f32;
+    let positions: Vec<u32> = (0..seq_len as u32).collect();
+    let out = Qwen35Moe::forward_full_attention_layer(
+        &spec,
+        &input,
+        &weights,
+        &positions,
+        n_rot,
+        rope_theta,
+        config.ln_eps(),
+    );
+    let q_rope_cpu = partial_rope_cpu(
+        &out.q_normed,
+        seq_len,
+        spec.num_heads,
+        spec.head_dim,
+        n_rot,
+        rope_theta,
+    )
+    .to_device(Device::Cuda(device));
+    let k_rope_cpu = partial_rope_cpu(
+        &out.k_normed,
+        seq_len,
+        spec.num_kv_heads,
+        spec.head_dim,
+        n_rot,
+        rope_theta,
+    )
+    .to_device(Device::Cuda(device));
+
+    println!("layer={layer}");
+    println!("input_shape={:?}", input.shape());
+    println!("q_full_shape={:?}", out.q_full.shape());
+    println!("query_shape={:?}", out.query.shape());
+    println!("gate_shape={:?}", out.gate.shape());
+    println!("k_shape={:?}", out.k.shape());
+    println!("v_shape={:?}", out.v.shape());
+    println!("q_normed_shape={:?}", out.q_normed.shape());
+    println!("k_normed_shape={:?}", out.k_normed.shape());
+    println!("n_rot={n_rot} rope_theta={rope_theta}");
+    println!("q_rope_cpu_shape={:?}", q_rope_cpu.shape());
+    println!("q_rope_gpu_shape={:?}", out.q_rope.shape());
+    println!("k_rope_cpu_shape={:?}", k_rope_cpu.shape());
+    println!("k_rope_gpu_shape={:?}", out.k_rope.shape());
+    println!("sample_q={:?}", sample_bf16(&out.q_normed, 8));
+    println!("sample_q_rope_cpu={:?}", sample_bf16(&q_rope_cpu, 8));
+    println!("sample_q_rope_gpu={:?}", sample_bf16(&out.q_rope, 8));
+    let token1_offset = spec.num_heads * spec.head_dim;
+    println!(
+        "sample_q_token1={:?}",
+        sample_bf16_offset(&out.q_normed, token1_offset, 8)
+    );
+    println!(
+        "sample_q_rope_cpu_token1={:?}",
+        sample_bf16_offset(&q_rope_cpu, token1_offset, 8)
+    );
+    println!(
+        "sample_q_rope_gpu_token1={:?}",
+        sample_bf16_offset(&out.q_rope, token1_offset, 8)
+    );
+    println!("sample_gate={:?}", sample_bf16(&out.gate, 8));
+    println!("attn_out_shape={:?}", out.attn_out.shape());
+    println!("attn_merged_shape={:?}", out.attn_merged.shape());
+    println!("gate_sigmoid_shape={:?}", out.gate_sigmoid.shape());
+    println!("gated_attn_shape={:?}", out.gated_attn.shape());
+    println!("projected_shape={:?}", out.projected.shape());
+    println!("sample_attn={:?}", sample_bf16(&out.attn_merged, 8));
+    println!("sample_projected={:?}", sample_bf16(&out.projected, 8));
+}
+
+fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
+    let index_path = model_dir.join("model.safetensors.index.json");
+    let text = std::fs::read_to_string(&index_path)
+        .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
+    let index: serde_json::Value = serde_json::from_str(&text)
+        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
+    index
+        .get("weight_map")
+        .and_then(|v| v.as_object())
+        .expect("index must contain weight_map")
+        .iter()
+        .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
+        .collect()
+}
+
+fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
+    let mut data = Vec::with_capacity(seq_len * hidden);
+    for i in 0..seq_len * hidden {
+        let x = ((i % 97) as f32 - 48.0) / 128.0;
+        data.push(bf16::from_f32(x));
+    }
+    Tensor::from_slice(&data, &[seq_len, hidden])
+}
+
+fn partial_rope_cpu(
+    x: &Tensor,
+    seq_len: usize,
+    num_heads: usize,
+    head_dim: usize,
+    n_rot: usize,
+    theta: f32,
+) -> Tensor {
+    assert_eq!(x.shape(), &[seq_len * num_heads, head_dim]);
+    assert!(n_rot <= head_dim && n_rot % 2 == 0);
+    let cpu = x.to_device(Device::Cpu);
+    let src = cpu.as_slice::<bf16>();
+    let mut out = src.to_vec();
+    let half = n_rot / 2;
+    for s in 0..seq_len {
+        for h in 0..num_heads {
+            let base = (s * num_heads + h) * head_dim;
+            for i in 0..half {
+                let freq = 1.0f32 / theta.powf((2 * i) as f32 / n_rot as f32);
+                let angle = s as f32 * freq;
+                let (sin, cos) = angle.sin_cos();
+                let x0 = src[base + i].to_f32();
+                let x1 = src[base + i + half].to_f32();
+                out[base + i] = bf16::from_f32(x0 * cos - x1 * sin);
+                out[base + i + half] = bf16::from_f32(x1 * cos + x0 * sin);
+            }
+        }
+    }
+    Tensor::from_slice(&out, &[seq_len * num_heads, head_dim])
+}
+
+fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
+    sample_bf16_offset(t, 0, n)
+}
+
+fn sample_bf16_offset(t: &Tensor, offset: usize, n: usize) -> Vec<f32> {
+    let cpu = t.to_device(Device::Cpu);
+    cpu.as_slice::<bf16>()
+        .iter()
+        .skip(offset)
+        .take(n)
+        .map(|v| v.to_f32())
+        .collect()
+}

crates/xserv-model/src/bin/smoke-qwen35-layer.rs

@@ -0,0 +1,332 @@
+use std::path::PathBuf;
+
+use half::bf16;
+use xserv_kernels::{
+    GemmBackend, add, matmul,
+    moe::{moe_sparse_gemv_bf16, moe_topk_softmax, moe_weighted_sum_sparse},
+    mul, row_scale_bf16, sigmoid, silu,
+};
+use xserv_model::{ModelConfig, Qwen35FullAttentionWeights, Qwen35Moe, Qwen35MoeSpec, loader};
+use xserv_tensor::{Device, Tensor};
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: smoke-qwen35-layer <model-dir> [layer]");
+        std::process::exit(1);
+    }
+    let model_dir = PathBuf::from(&args[1]);
+    let layer: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(0);
+    let device = 0;
+    let seq_len = 1;
+
+    xserv_cuda::device::set_device(device).unwrap();
+    xserv_model::init_kernels();
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    let spec = Qwen35MoeSpec::from_config(&config);
+    let weight_map = read_weight_map(&model_dir);
+    let mut keys = layer_common_keys(layer);
+    match spec.layer_kinds[layer] {
+        xserv_model::qwen35_moe::Qwen35LayerKind::LinearAttention => {
+            keys.extend(linear_attention_keys(layer));
+        }
+        xserv_model::qwen35_moe::Qwen35LayerKind::FullAttention => {
+            keys.extend(full_attention_keys(layer));
+        }
+    }
+    keys.extend(moe_keys(layer));
+    let mut shards: Vec<String> = keys
+        .iter()
+        .filter_map(|k| weight_map.get(k).cloned())
+        .collect();
+    shards.sort();
+    shards.dedup();
+
+    let mut tensors = std::collections::HashMap::new();
+    for shard in &shards {
+        tensors.extend(loader::load_safetensors(
+            &model_dir.join(shard),
+            Device::Cpu,
+        ));
+    }
+
+    let input = make_deterministic_input(seq_len, config.hidden()).to_device(Device::Cuda(device));
+    let p = format!("model.language_model.layers.{layer}");
+    let input_norm = take_tensor(&mut tensors, &format!("{p}.input_layernorm.weight"), device);
+    let post_norm = take_tensor(
+        &mut tensors,
+        &format!("{p}.post_attention_layernorm.weight"),
+        device,
+    );
+    let normed = xserv_kernels::rmsnorm(&input, &input_norm, config.ln_eps());
+
+    let attn_projected = match spec.layer_kinds[layer] {
+        xserv_model::qwen35_moe::Qwen35LayerKind::LinearAttention => {
+            let weights = Qwen35Moe::take_linear_attention_projection_weights(
+                &mut tensors,
+                layer,
+                Device::Cuda(device),
+            );
+            let projected = Qwen35Moe::project_linear_attention(&spec, &normed, &weights);
+            let mut state = Tensor::zeros(
+                &[
+                    spec.linear_value_heads,
+                    spec.linear_value_dim,
+                    spec.linear_value_dim,
+                ],
+                xserv_tensor::DType::BF16,
+                Device::Cuda(device),
+            );
+            let recurrent =
+                Qwen35Moe::linear_attention_recurrent_step(&spec, &projected, &weights, &mut state);
+            let (_, out) = Qwen35Moe::finish_linear_attention(
+                &spec,
+                &recurrent,
+                &projected.z,
+                &weights,
+                config.ln_eps(),
+            );
+            out
+        }
+        xserv_model::qwen35_moe::Qwen35LayerKind::FullAttention => {
+            let weights = Qwen35FullAttentionWeights {
+                q_w_t: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.q_proj.weight"),
+                    device,
+                )
+                .transpose(0, 1)
+                .contiguous(),
+                k_w_t: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.k_proj.weight"),
+                    device,
+                )
+                .transpose(0, 1)
+                .contiguous(),
+                v_w_t: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.v_proj.weight"),
+                    device,
+                )
+                .transpose(0, 1)
+                .contiguous(),
+                o_w_t: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.o_proj.weight"),
+                    device,
+                )
+                .transpose(0, 1)
+                .contiguous(),
+                q_norm: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.q_norm.weight"),
+                    device,
+                ),
+                k_norm: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.k_norm.weight"),
+                    device,
+                ),
+            };
+            let n_rot = (spec.head_dim as f64 * config.text().partial_rotary_factor.unwrap_or(0.25))
+                as usize;
+            let positions = [0u32];
+            Qwen35Moe::forward_full_attention_layer(
+                &spec,
+                &normed,
+                &weights,
+                &positions,
+                n_rot,
+                config.rope_theta_value().unwrap_or(10_000_000.0) as f32,
+                config.ln_eps(),
+            )
+            .projected
+        }
+    };
+    let attn_residual = add(&input, &attn_projected);
+    let moe_in = xserv_kernels::rmsnorm(&attn_residual, &post_norm, config.ln_eps());
+    let moe_out = run_moe(&mut tensors, layer, &moe_in, &spec, device);
+    let layer_out = add(&attn_residual, &moe_out);
+
+    println!("layer={layer}");
+    println!("kind={:?}", spec.layer_kinds[layer]);
+    println!("input_shape={:?}", input.shape());
+    println!("normed_shape={:?}", normed.shape());
+    println!("attn_projected_shape={:?}", attn_projected.shape());
+    println!("attn_residual_shape={:?}", attn_residual.shape());
+    println!("moe_in_shape={:?}", moe_in.shape());
+    println!("moe_out_shape={:?}", moe_out.shape());
+    println!("layer_out_shape={:?}", layer_out.shape());
+    println!("sample_attn={:?}", sample_bf16(&attn_projected, 8));
+    println!("sample_moe={:?}", sample_bf16(&moe_out, 8));
+    println!("sample_layer_out={:?}", sample_bf16(&layer_out, 8));
+}
+
+fn run_moe(
+    tensors: &mut std::collections::HashMap<String, Tensor>,
+    layer: usize,
+    input: &Tensor,
+    spec: &Qwen35MoeSpec,
+    device: u32,
+) -> Tensor {
+    let p = format!("model.language_model.layers.{layer}.mlp");
+    let router_w = take_tensor(tensors, &format!("{p}.gate.weight"), device)
+        .transpose(0, 1)
+        .contiguous();
+    let shared_gate_w = take_tensor(
+        tensors,
+        &format!("{p}.shared_expert.gate_proj.weight"),
+        device,
+    )
+    .transpose(0, 1)
+    .contiguous();
+    let shared_up_w = take_tensor(
+        tensors,
+        &format!("{p}.shared_expert.up_proj.weight"),
+        device,
+    )
+    .transpose(0, 1)
+    .contiguous();
+    let shared_down_w = take_tensor(
+        tensors,
+        &format!("{p}.shared_expert.down_proj.weight"),
+        device,
+    )
+    .transpose(0, 1)
+    .contiguous();
+    let shared_router_w = take_tensor(tensors, &format!("{p}.shared_expert_gate.weight"), device)
+        .transpose(0, 1)
+        .contiguous();
+    let expert_gate_up =
+        take_tensor(tensors, &format!("{p}.experts.gate_up_proj"), device).contiguous();
+    let expert_down = take_tensor(tensors, &format!("{p}.experts.down_proj"), device).contiguous();
+
+    let router_logits = matmul(input, &router_w, GemmBackend::CuBlas);
+    let (topk_ids, topk_weights) =
+        moe_topk_softmax(&router_logits, spec.num_experts, spec.experts_per_token);
+    let shared_gate = matmul(input, &shared_gate_w, GemmBackend::CuBlas);
+    let shared_up = matmul(input, &shared_up_w, GemmBackend::CuBlas);
+    let shared_act = mul(&silu(&shared_gate), &shared_up);
+    let shared_down = matmul(&shared_act, &shared_down_w, GemmBackend::CuBlas);
+    let shared_router = sigmoid(&matmul(input, &shared_router_w, GemmBackend::CuBlas));
+    let shared_out = row_scale_bf16(&shared_down, &shared_router);
+    let gate_up = moe_sparse_gemv_bf16(
+        input,
+        &expert_gate_up,
+        &topk_ids,
+        spec.experts_per_token,
+        0,
+        spec.num_experts,
+        false,
+    );
+    let gate = gate_up.narrow(2, 0, spec.moe_intermediate).contiguous();
+    let up = gate_up
+        .narrow(2, spec.moe_intermediate, spec.moe_intermediate)
+        .contiguous();
+    let routed_act = mul(&silu(&gate), &up);
+    let down = moe_sparse_gemv_bf16(
+        &routed_act.reshape(&[spec.experts_per_token, spec.moe_intermediate]),
+        &expert_down,
+        &topk_ids,
+        spec.experts_per_token,
+        0,
+        spec.num_experts,
+        true,
+    );
+    let routed_out = moe_weighted_sum_sparse(&down, &topk_ids, &topk_weights, 0, spec.num_experts);
+    add(&routed_out, &shared_out)
+}
+
+fn take_tensor(
+    tensors: &mut std::collections::HashMap<String, Tensor>,
+    name: &str,
+    device: u32,
+) -> Tensor {
+    tensors
+        .remove(name)
+        .unwrap_or_else(|| panic!("missing tensor {name}"))
+        .to_device(Device::Cuda(device))
+}
+
+fn layer_common_keys(layer: usize) -> Vec<String> {
+    let p = format!("model.language_model.layers.{layer}");
+    vec![
+        format!("{p}.input_layernorm.weight"),
+        format!("{p}.post_attention_layernorm.weight"),
+    ]
+}
+
+fn linear_attention_keys(layer: usize) -> Vec<String> {
+    let p = format!("model.language_model.layers.{layer}.linear_attn");
+    vec![
+        format!("{p}.in_proj_qkv.weight"),
+        format!("{p}.in_proj_z.weight"),
+        format!("{p}.in_proj_a.weight"),
+        format!("{p}.in_proj_b.weight"),
+        format!("{p}.conv1d.weight"),
+        format!("{p}.A_log"),
+        format!("{p}.dt_bias"),
+        format!("{p}.norm.weight"),
+        format!("{p}.out_proj.weight"),
+    ]
+}
+
+fn full_attention_keys(layer: usize) -> Vec<String> {
+    let p = format!("model.language_model.layers.{layer}.self_attn");
+    vec![
+        format!("{p}.q_proj.weight"),
+        format!("{p}.k_proj.weight"),
+        format!("{p}.v_proj.weight"),
+        format!("{p}.o_proj.weight"),
+        format!("{p}.q_norm.weight"),
+        format!("{p}.k_norm.weight"),
+    ]
+}
+
+fn moe_keys(layer: usize) -> Vec<String> {
+    let p = format!("model.language_model.layers.{layer}.mlp");
+    vec![
+        format!("{p}.gate.weight"),
+        format!("{p}.shared_expert.gate_proj.weight"),
+        format!("{p}.shared_expert.up_proj.weight"),
+        format!("{p}.shared_expert.down_proj.weight"),
+        format!("{p}.shared_expert_gate.weight"),
+        format!("{p}.experts.gate_up_proj"),
+        format!("{p}.experts.down_proj"),
+    ]
+}
+
+fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
+    let index_path = model_dir.join("model.safetensors.index.json");
+    let text = std::fs::read_to_string(&index_path)
+        .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
+    let index: serde_json::Value = serde_json::from_str(&text)
+        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
+    index
+        .get("weight_map")
+        .and_then(|v| v.as_object())
+        .expect("index must contain weight_map")
+        .iter()
+        .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
+        .collect()
+}
+
+fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
+    let mut data = Vec::with_capacity(seq_len * hidden);
+    for i in 0..seq_len * hidden {
+        let x = ((i % 97) as f32 - 48.0) / 128.0;
+        data.push(bf16::from_f32(x));
+    }
+    Tensor::from_slice(&data, &[seq_len, hidden])
+}
+
+fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
+    let cpu = t.to_device(Device::Cpu);
+    cpu.as_slice::<bf16>()
+        .iter()
+        .take(n)
+        .map(|v| v.to_f32())
+        .collect()
+}

crates/xserv-model/src/bin/smoke-qwen35-linear-attn.rs

@@ -0,0 +1,164 @@
+use std::path::PathBuf;
+
+use half::bf16;
+use xserv_model::{ModelConfig, Qwen35Moe, Qwen35MoeSpec, loader};
+use xserv_tensor::{Device, Tensor};
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: smoke-qwen35-linear-attn <model-dir> [layer] [seq_len]");
+        std::process::exit(1);
+    }
+    let model_dir = PathBuf::from(&args[1]);
+    let layer: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(0);
+    let seq_len: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(2);
+    let device = 0;
+
+    xserv_cuda::device::set_device(device).unwrap();
+    xserv_model::init_kernels();
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    let spec = Qwen35MoeSpec::from_config(&config);
+    assert_eq!(
+        spec.layer_kinds[layer],
+        xserv_model::qwen35_moe::Qwen35LayerKind::LinearAttention,
+        "layer {layer} is not a linear-attention layer"
+    );
+
+    let weight_map = read_weight_map(&model_dir);
+    let keys = [
+        format!("model.language_model.layers.{layer}.linear_attn.in_proj_qkv.weight"),
+        format!("model.language_model.layers.{layer}.linear_attn.in_proj_z.weight"),
+        format!("model.language_model.layers.{layer}.linear_attn.in_proj_a.weight"),
+        format!("model.language_model.layers.{layer}.linear_attn.in_proj_b.weight"),
+        format!("model.language_model.layers.{layer}.linear_attn.conv1d.weight"),
+        format!("model.language_model.layers.{layer}.linear_attn.A_log"),
+        format!("model.language_model.layers.{layer}.linear_attn.dt_bias"),
+        format!("model.language_model.layers.{layer}.linear_attn.norm.weight"),
+        format!("model.language_model.layers.{layer}.linear_attn.out_proj.weight"),
+    ];
+    let mut shards: Vec<String> = keys
+        .iter()
+        .filter_map(|k| weight_map.get(k).cloned())
+        .collect();
+    shards.sort();
+    shards.dedup();
+
+    let mut tensors = std::collections::HashMap::new();
+    for shard in &shards {
+        tensors.extend(loader::load_safetensors(
+            &model_dir.join(shard),
+            Device::Cpu,
+        ));
+    }
+
+    let weights = Qwen35Moe::take_linear_attention_projection_weights(
+        &mut tensors,
+        layer,
+        Device::Cuda(device),
+    );
+    let input = make_deterministic_input(seq_len, config.hidden()).to_device(Device::Cuda(device));
+    let out = Qwen35Moe::project_linear_attention(&spec, &input, &weights);
+    let mut recurrent_state = Tensor::zeros(
+        &[
+            spec.linear_value_heads,
+            spec.linear_value_dim,
+            spec.linear_value_dim,
+        ],
+        xserv_tensor::DType::BF16,
+        Device::Cuda(device),
+    );
+    let recurrent_out = if seq_len == 1 {
+        Some(Qwen35Moe::linear_attention_recurrent_step(
+            &spec,
+            &out,
+            &weights,
+            &mut recurrent_state,
+        ))
+    } else {
+        None
+    };
+    let finished = recurrent_out
+        .as_ref()
+        .map(|r| Qwen35Moe::finish_linear_attention(&spec, r, &out.z, &weights, config.ln_eps()));
+
+    println!("layer={layer}");
+    println!("input_shape={:?}", input.shape());
+    println!("qkv_shape={:?}", out.qkv.shape());
+    println!("z_shape={:?}", out.z.shape());
+    println!("a_shape={:?}", out.a.shape());
+    println!("b_shape={:?}", out.b.shape());
+    println!("beta_shape={:?}", out.beta.shape());
+    println!("alpha_softplus_shape={:?}", out.alpha_softplus.shape());
+    println!("q_shape={:?}", out.q.shape());
+    println!("k_shape={:?}", out.k.shape());
+    println!("v_shape={:?}", out.v.shape());
+    println!("conv_shape={:?}", out.conv.shape());
+    println!("q_conv_shape={:?}", out.q_conv.shape());
+    println!("k_conv_shape={:?}", out.k_conv.shape());
+    println!("v_conv_shape={:?}", out.v_conv.shape());
+    println!("sample_q={:?}", sample_bf16(&out.q, 8));
+    println!("sample_k={:?}", sample_bf16(&out.k, 8));
+    println!("sample_v={:?}", sample_bf16(&out.v, 8));
+    println!("sample_conv={:?}", sample_bf16(&out.conv, 8));
+    println!("sample_q_conv={:?}", sample_bf16(&out.q_conv, 8));
+    println!("sample_k_conv={:?}", sample_bf16(&out.k_conv, 8));
+    println!("sample_v_conv={:?}", sample_bf16(&out.v_conv, 8));
+    println!("sample_z={:?}", sample_bf16(&out.z, 8));
+    println!("sample_a={:?}", sample_bf16(&out.a, 8));
+    println!("sample_b={:?}", sample_bf16(&out.b, 8));
+    println!("sample_beta={:?}", sample_bf16(&out.beta, 8));
+    println!(
+        "sample_alpha_softplus={:?}",
+        sample_bf16(&out.alpha_softplus, 8)
+    );
+    if let Some(recurrent_out) = recurrent_out {
+        println!("recurrent_out_shape={:?}", recurrent_out.shape());
+        println!("recurrent_state_shape={:?}", recurrent_state.shape());
+        println!("sample_recurrent_out={:?}", sample_bf16(&recurrent_out, 8));
+        println!(
+            "sample_recurrent_state={:?}",
+            sample_bf16(&recurrent_state, 8)
+        );
+    }
+    if let Some((norm_gated, projected)) = finished {
+        println!("norm_gated_shape={:?}", norm_gated.shape());
+        println!("projected_shape={:?}", projected.shape());
+        println!("sample_norm_gated={:?}", sample_bf16(&norm_gated, 8));
+        println!("sample_projected={:?}", sample_bf16(&projected, 8));
+    }
+}
+
+fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
+    let index_path = model_dir.join("model.safetensors.index.json");
+    let text = std::fs::read_to_string(&index_path)
+        .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
+    let index: serde_json::Value = serde_json::from_str(&text)
+        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
+    index
+        .get("weight_map")
+        .and_then(|v| v.as_object())
+        .expect("index must contain weight_map")
+        .iter()
+        .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
+        .collect()
+}
+
+fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
+    let mut data = Vec::with_capacity(seq_len * hidden);
+    for i in 0..seq_len * hidden {
+        let x = ((i % 97) as f32 - 48.0) / 128.0;
+        data.push(bf16::from_f32(x));
+    }
+    Tensor::from_slice(&data, &[seq_len, hidden])
+}
+
+fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
+    let cpu = t.to_device(Device::Cpu);
+    cpu.as_slice::<bf16>()
+        .iter()
+        .take(n)
+        .map(|v| v.to_f32())
+        .collect()
+}

crates/xserv-model/src/bin/smoke-qwen35-moe.rs

@@ -0,0 +1,193 @@
+use std::path::PathBuf;
+
+use half::bf16;
+use xserv_kernels::{
+    GemmBackend, add, matmul,
+    moe::{moe_sparse_gemv_bf16, moe_topk_softmax, moe_weighted_sum_sparse},
+    mul, row_scale_bf16, sigmoid, silu,
+};
+use xserv_model::{ModelConfig, Qwen35MoeSpec, loader};
+use xserv_tensor::{Device, Tensor};
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: smoke-qwen35-moe <model-dir> [layer] [seq_len]");
+        std::process::exit(1);
+    }
+    let model_dir = PathBuf::from(&args[1]);
+    let layer: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(0);
+    let seq_len: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(1);
+    let device = 0;
+
+    xserv_cuda::device::set_device(device).unwrap();
+    xserv_model::init_kernels();
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    let spec = Qwen35MoeSpec::from_config(&config);
+    let weight_map = read_weight_map(&model_dir);
+    let keys = [
+        format!("model.language_model.layers.{layer}.mlp.gate.weight"),
+        format!("model.language_model.layers.{layer}.mlp.shared_expert.gate_proj.weight"),
+        format!("model.language_model.layers.{layer}.mlp.shared_expert.up_proj.weight"),
+        format!("model.language_model.layers.{layer}.mlp.shared_expert.down_proj.weight"),
+        format!("model.language_model.layers.{layer}.mlp.shared_expert_gate.weight"),
+        format!("model.language_model.layers.{layer}.mlp.experts.gate_up_proj"),
+        format!("model.language_model.layers.{layer}.mlp.experts.down_proj"),
+    ];
+    let mut shards: Vec<String> = keys
+        .iter()
+        .filter_map(|k| weight_map.get(k).cloned())
+        .collect();
+    shards.sort();
+    shards.dedup();
+
+    let mut tensors = std::collections::HashMap::new();
+    for shard in &shards {
+        tensors.extend(loader::load_safetensors(
+            &model_dir.join(shard),
+            Device::Cpu,
+        ));
+    }
+    let mut take = |name: &str| -> Tensor {
+        tensors
+            .remove(name)
+            .unwrap_or_else(|| panic!("missing tensor {name}"))
+            .to_device(Device::Cuda(device))
+    };
+    let p = format!("model.language_model.layers.{layer}.mlp");
+    let router_w = take(&format!("{p}.gate.weight"))
+        .transpose(0, 1)
+        .contiguous();
+    let shared_gate_w = take(&format!("{p}.shared_expert.gate_proj.weight"))
+        .transpose(0, 1)
+        .contiguous();
+    let shared_up_w = take(&format!("{p}.shared_expert.up_proj.weight"))
+        .transpose(0, 1)
+        .contiguous();
+    let shared_down_w = take(&format!("{p}.shared_expert.down_proj.weight"))
+        .transpose(0, 1)
+        .contiguous();
+    let shared_router_w = take(&format!("{p}.shared_expert_gate.weight"))
+        .transpose(0, 1)
+        .contiguous();
+    let expert_gate_up = take(&format!("{p}.experts.gate_up_proj")).contiguous();
+    let expert_down = take(&format!("{p}.experts.down_proj")).contiguous();
+
+    let input = make_deterministic_input(seq_len, config.hidden()).to_device(Device::Cuda(device));
+    let router_logits = matmul(&input, &router_w, GemmBackend::CuBlas);
+    let (topk_ids, topk_weights) =
+        moe_topk_softmax(&router_logits, spec.num_experts, spec.experts_per_token);
+
+    let shared_gate = matmul(&input, &shared_gate_w, GemmBackend::CuBlas);
+    let shared_up = matmul(&input, &shared_up_w, GemmBackend::CuBlas);
+    let shared_act = mul(&silu(&shared_gate), &shared_up);
+    let shared_down = matmul(&shared_act, &shared_down_w, GemmBackend::CuBlas);
+    let shared_router = sigmoid(&matmul(&input, &shared_router_w, GemmBackend::CuBlas));
+    let shared_out = row_scale_bf16(&shared_down, &shared_router);
+
+    let gate_up = moe_sparse_gemv_bf16(
+        &input,
+        &expert_gate_up,
+        &topk_ids,
+        spec.experts_per_token,
+        0,
+        spec.num_experts,
+        false,
+    );
+    let gate = gate_up.narrow(2, 0, spec.moe_intermediate).contiguous();
+    let up = gate_up
+        .narrow(2, spec.moe_intermediate, spec.moe_intermediate)
+        .contiguous();
+    let routed_act = mul(&silu(&gate), &up);
+    let down = moe_sparse_gemv_bf16(
+        &routed_act.reshape(&[seq_len * spec.experts_per_token, spec.moe_intermediate]),
+        &expert_down,
+        &topk_ids,
+        spec.experts_per_token,
+        0,
+        spec.num_experts,
+        true,
+    );
+    let routed_out = moe_weighted_sum_sparse(&down, &topk_ids, &topk_weights, 0, spec.num_experts);
+    let moe_out = add(&routed_out, &shared_out);
+
+    println!("layer={layer}");
+    println!("input_shape={:?}", input.shape());
+    println!("router_logits_shape={:?}", router_logits.shape());
+    println!("topk_ids_shape={:?}", topk_ids.shape());
+    println!("topk_weights_shape={:?}", topk_weights.shape());
+    println!("shared_gate_shape={:?}", shared_gate.shape());
+    println!("shared_up_shape={:?}", shared_up.shape());
+    println!("shared_act_shape={:?}", shared_act.shape());
+    println!("shared_down_shape={:?}", shared_down.shape());
+    println!("shared_router_shape={:?}", shared_router.shape());
+    println!("shared_out_shape={:?}", shared_out.shape());
+    println!("gate_up_shape={:?}", gate_up.shape());
+    println!("routed_act_shape={:?}", routed_act.shape());
+    println!("down_shape={:?}", down.shape());
+    println!("routed_out_shape={:?}", routed_out.shape());
+    println!("moe_out_shape={:?}", moe_out.shape());
+    println!("sample_router={:?}", sample_bf16(&router_logits, 8));
+    println!(
+        "sample_topk_ids={:?}",
+        sample_i32_raw(&topk_ids, spec.experts_per_token)
+    );
+    println!(
+        "sample_topk_weights={:?}",
+        sample_f32(&topk_weights, spec.experts_per_token)
+    );
+    println!("sample_shared_router={:?}", sample_bf16(&shared_router, 4));
+    println!("sample_shared_out={:?}", sample_bf16(&shared_out, 8));
+    println!("sample_routed_out={:?}", sample_bf16(&routed_out, 8));
+    println!("sample_moe_out={:?}", sample_bf16(&moe_out, 8));
+}
+
+fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
+    let index_path = model_dir.join("model.safetensors.index.json");
+    let text = std::fs::read_to_string(&index_path)
+        .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
+    let index: serde_json::Value = serde_json::from_str(&text)
+        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
+    index
+        .get("weight_map")
+        .and_then(|v| v.as_object())
+        .expect("index must contain weight_map")
+        .iter()
+        .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
+        .collect()
+}
+
+fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
+    let mut data = Vec::with_capacity(seq_len * hidden);
+    for i in 0..seq_len * hidden {
+        let x = ((i % 97) as f32 - 48.0) / 128.0;
+        data.push(bf16::from_f32(x));
+    }
+    Tensor::from_slice(&data, &[seq_len, hidden])
+}
+
+fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
+    let cpu = t.to_device(Device::Cpu);
+    cpu.as_slice::<bf16>()
+        .iter()
+        .take(n)
+        .map(|v| v.to_f32())
+        .collect()
+}
+
+fn sample_f32(t: &Tensor, n: usize) -> Vec<f32> {
+    let cpu = t.to_device(Device::Cpu);
+    cpu.as_slice::<f32>().iter().take(n).copied().collect()
+}
+
+fn sample_i32_raw(t: &Tensor, n: usize) -> Vec<i32> {
+    let cpu = t.to_device(Device::Cpu);
+    let bytes = cpu.as_raw_bytes();
+    (0..n)
+        .map(|i| {
+            let start = i * 4;
+            i32::from_ne_bytes(bytes[start..start + 4].try_into().unwrap())
+        })
+        .collect()
+}

crates/xserv-model/src/bin/smoke-qwen35-prefix.rs

@@ -0,0 +1,411 @@
+use std::path::PathBuf;
+
+use half::bf16;
+use xserv_kernels::{
+    GemmBackend, add, embedding, matmul,
+    moe::{moe_sparse_gemv_bf16, moe_topk_softmax, moe_weighted_sum_sparse},
+    mul, row_scale_bf16, sigmoid, silu,
+};
+use xserv_model::{ModelConfig, Qwen35FullAttentionWeights, Qwen35Moe, Qwen35MoeSpec, loader};
+use xserv_tensor::{Device, Tensor};
+use xserv_tokenizer::Tokenizer;
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: smoke-qwen35-prefix <model-dir> [num_layers] [token_id]");
+        std::process::exit(1);
+    }
+    let model_dir = PathBuf::from(&args[1]);
+    let num_layers: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(4);
+    let token_id: Option<u32> = args.get(3).and_then(|s| s.parse().ok());
+    let device = 0;
+
+    xserv_cuda::device::set_device(device).unwrap();
+    xserv_model::init_kernels();
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+    let spec = Qwen35MoeSpec::from_config(&config);
+    let weight_map = read_weight_map(&model_dir);
+    let mut final_tensors = load_keys(
+        &model_dir,
+        &weight_map,
+        &[
+            "model.language_model.embed_tokens.weight".to_string(),
+            "model.language_model.norm.weight".to_string(),
+            "lm_head.weight".to_string(),
+        ],
+    );
+    let mut hidden = if let Some(token_id) = token_id {
+        let embed = take_tensor(
+            &mut final_tensors,
+            "model.language_model.embed_tokens.weight",
+            device,
+        );
+        println!("input_token_id={token_id}");
+        println!("input_token_text={:?}", tokenizer.decode(&[token_id]));
+        embedding(&embed, &[token_id])
+    } else {
+        make_deterministic_input(1, config.hidden()).to_device(Device::Cuda(device))
+    };
+
+    println!("prefix_layers={num_layers}");
+    println!("input_shape={:?}", hidden.shape());
+    for layer in 0..num_layers {
+        hidden = run_layer(
+            &model_dir,
+            &weight_map,
+            &config,
+            &spec,
+            layer,
+            &hidden,
+            device,
+        );
+        println!(
+            "layer={layer} kind={:?} out_shape={:?}",
+            spec.layer_kinds[layer],
+            hidden.shape()
+        );
+        println!("sample_layer_{layer}={:?}", sample_bf16(&hidden, 8));
+    }
+
+    let final_norm = take_tensor(
+        &mut final_tensors,
+        "model.language_model.norm.weight",
+        device,
+    );
+    let lm_head_t = take_tensor(&mut final_tensors, "lm_head.weight", device)
+        .transpose(0, 1)
+        .contiguous();
+    let normed = xserv_kernels::rmsnorm(&hidden, &final_norm, config.ln_eps());
+    let logits = matmul(&normed, &lm_head_t, GemmBackend::CuBlas);
+    let top = xserv_kernels::argmax_bf16_single(&logits);
+    println!("final_normed_shape={:?}", normed.shape());
+    println!("logits_shape={:?}", logits.shape());
+    println!("top_token={top}");
+    println!("top_token_text={:?}", tokenizer.decode(&[top]));
+    println!("sample_logits={:?}", sample_bf16(&logits, 8));
+}
+
+fn load_keys(
+    model_dir: &std::path::Path,
+    weight_map: &std::collections::HashMap<String, String>,
+    keys: &[String],
+) -> std::collections::HashMap<String, Tensor> {
+    let mut shards: Vec<String> = keys
+        .iter()
+        .filter_map(|k| weight_map.get(k).cloned())
+        .collect();
+    shards.sort();
+    shards.dedup();
+
+    let mut tensors = std::collections::HashMap::new();
+    for shard in &shards {
+        tensors.extend(loader::load_safetensors(
+            &model_dir.join(shard),
+            Device::Cpu,
+        ));
+    }
+    tensors
+}
+
+fn run_layer(
+    model_dir: &std::path::Path,
+    weight_map: &std::collections::HashMap<String, String>,
+    config: &ModelConfig,
+    spec: &Qwen35MoeSpec,
+    layer: usize,
+    input: &Tensor,
+    device: u32,
+) -> Tensor {
+    let mut keys = layer_common_keys(layer);
+    match spec.layer_kinds[layer] {
+        xserv_model::qwen35_moe::Qwen35LayerKind::LinearAttention => {
+            keys.extend(linear_attention_keys(layer));
+        }
+        xserv_model::qwen35_moe::Qwen35LayerKind::FullAttention => {
+            keys.extend(full_attention_keys(layer));
+        }
+    }
+    keys.extend(moe_keys(layer));
+    let mut shards: Vec<String> = keys
+        .iter()
+        .filter_map(|k| weight_map.get(k).cloned())
+        .collect();
+    shards.sort();
+    shards.dedup();
+
+    let mut tensors = std::collections::HashMap::new();
+    for shard in &shards {
+        tensors.extend(loader::load_safetensors(
+            &model_dir.join(shard),
+            Device::Cpu,
+        ));
+    }
+
+    let p = format!("model.language_model.layers.{layer}");
+    let input_norm = take_tensor(&mut tensors, &format!("{p}.input_layernorm.weight"), device);
+    let post_norm = take_tensor(
+        &mut tensors,
+        &format!("{p}.post_attention_layernorm.weight"),
+        device,
+    );
+    let normed = xserv_kernels::rmsnorm(input, &input_norm, config.ln_eps());
+
+    let attn_projected = match spec.layer_kinds[layer] {
+        xserv_model::qwen35_moe::Qwen35LayerKind::LinearAttention => {
+            let weights = Qwen35Moe::take_linear_attention_projection_weights(
+                &mut tensors,
+                layer,
+                Device::Cuda(device),
+            );
+            let projected = Qwen35Moe::project_linear_attention(spec, &normed, &weights);
+            let mut state = Tensor::zeros(
+                &[
+                    spec.linear_value_heads,
+                    spec.linear_value_dim,
+                    spec.linear_value_dim,
+                ],
+                xserv_tensor::DType::BF16,
+                Device::Cuda(device),
+            );
+            let recurrent =
+                Qwen35Moe::linear_attention_recurrent_step(spec, &projected, &weights, &mut state);
+            let (_, out) = Qwen35Moe::finish_linear_attention(
+                spec,
+                &recurrent,
+                &projected.z,
+                &weights,
+                config.ln_eps(),
+            );
+            out
+        }
+        xserv_model::qwen35_moe::Qwen35LayerKind::FullAttention => {
+            let weights = Qwen35FullAttentionWeights {
+                q_w_t: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.q_proj.weight"),
+                    device,
+                )
+                .transpose(0, 1)
+                .contiguous(),
+                k_w_t: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.k_proj.weight"),
+                    device,
+                )
+                .transpose(0, 1)
+                .contiguous(),
+                v_w_t: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.v_proj.weight"),
+                    device,
+                )
+                .transpose(0, 1)
+                .contiguous(),
+                o_w_t: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.o_proj.weight"),
+                    device,
+                )
+                .transpose(0, 1)
+                .contiguous(),
+                q_norm: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.q_norm.weight"),
+                    device,
+                ),
+                k_norm: take_tensor(
+                    &mut tensors,
+                    &format!("{p}.self_attn.k_norm.weight"),
+                    device,
+                ),
+            };
+            let n_rot = (spec.head_dim as f64 * config.text().partial_rotary_factor.unwrap_or(0.25))
+                as usize;
+            let positions = [0u32];
+            Qwen35Moe::forward_full_attention_layer(
+                spec,
+                &normed,
+                &weights,
+                &positions,
+                n_rot,
+                config.rope_theta_value().unwrap_or(10_000_000.0) as f32,
+                config.ln_eps(),
+            )
+            .projected
+        }
+    };
+    let attn_residual = add(input, &attn_projected);
+    let moe_in = xserv_kernels::rmsnorm(&attn_residual, &post_norm, config.ln_eps());
+    let moe_out = run_moe(&mut tensors, layer, &moe_in, spec, device);
+    add(&attn_residual, &moe_out)
+}
+
+fn run_moe(
+    tensors: &mut std::collections::HashMap<String, Tensor>,
+    layer: usize,
+    input: &Tensor,
+    spec: &Qwen35MoeSpec,
+    device: u32,
+) -> Tensor {
+    let p = format!("model.language_model.layers.{layer}.mlp");
+    let router_w = take_tensor(tensors, &format!("{p}.gate.weight"), device)
+        .transpose(0, 1)
+        .contiguous();
+    let shared_gate_w = take_tensor(
+        tensors,
+        &format!("{p}.shared_expert.gate_proj.weight"),
+        device,
+    )
+    .transpose(0, 1)
+    .contiguous();
+    let shared_up_w = take_tensor(
+        tensors,
+        &format!("{p}.shared_expert.up_proj.weight"),
+        device,
+    )
+    .transpose(0, 1)
+    .contiguous();
+    let shared_down_w = take_tensor(
+        tensors,
+        &format!("{p}.shared_expert.down_proj.weight"),
+        device,
+    )
+    .transpose(0, 1)
+    .contiguous();
+    let shared_router_w = take_tensor(tensors, &format!("{p}.shared_expert_gate.weight"), device)
+        .transpose(0, 1)
+        .contiguous();
+    let expert_gate_up =
+        take_tensor(tensors, &format!("{p}.experts.gate_up_proj"), device).contiguous();
+    let expert_down = take_tensor(tensors, &format!("{p}.experts.down_proj"), device).contiguous();
+
+    let router_logits = matmul(input, &router_w, GemmBackend::CuBlas);
+    let (topk_ids, topk_weights) =
+        moe_topk_softmax(&router_logits, spec.num_experts, spec.experts_per_token);
+    let shared_gate = matmul(input, &shared_gate_w, GemmBackend::CuBlas);
+    let shared_up = matmul(input, &shared_up_w, GemmBackend::CuBlas);
+    let shared_act = mul(&silu(&shared_gate), &shared_up);
+    let shared_down = matmul(&shared_act, &shared_down_w, GemmBackend::CuBlas);
+    let shared_router = sigmoid(&matmul(input, &shared_router_w, GemmBackend::CuBlas));
+    let shared_out = row_scale_bf16(&shared_down, &shared_router);
+    let gate_up = moe_sparse_gemv_bf16(
+        input,
+        &expert_gate_up,
+        &topk_ids,
+        spec.experts_per_token,
+        0,
+        spec.num_experts,
+        false,
+    );
+    let gate = gate_up.narrow(2, 0, spec.moe_intermediate).contiguous();
+    let up = gate_up
+        .narrow(2, spec.moe_intermediate, spec.moe_intermediate)
+        .contiguous();
+    let routed_act = mul(&silu(&gate), &up);
+    let down = moe_sparse_gemv_bf16(
+        &routed_act.reshape(&[spec.experts_per_token, spec.moe_intermediate]),
+        &expert_down,
+        &topk_ids,
+        spec.experts_per_token,
+        0,
+        spec.num_experts,
+        true,
+    );
+    let routed_out = moe_weighted_sum_sparse(&down, &topk_ids, &topk_weights, 0, spec.num_experts);
+    add(&routed_out, &shared_out)
+}
+
+fn take_tensor(
+    tensors: &mut std::collections::HashMap<String, Tensor>,
+    name: &str,
+    device: u32,
+) -> Tensor {
+    tensors
+        .remove(name)
+        .unwrap_or_else(|| panic!("missing tensor {name}"))
+        .to_device(Device::Cuda(device))
+}
+
+fn layer_common_keys(layer: usize) -> Vec<String> {
+    let p = format!("model.language_model.layers.{layer}");
+    vec![
+        format!("{p}.input_layernorm.weight"),
+        format!("{p}.post_attention_layernorm.weight"),
+    ]
+}
+
+fn linear_attention_keys(layer: usize) -> Vec<String> {
+    let p = format!("model.language_model.layers.{layer}.linear_attn");
+    vec![
+        format!("{p}.in_proj_qkv.weight"),
+        format!("{p}.in_proj_z.weight"),
+        format!("{p}.in_proj_a.weight"),
+        format!("{p}.in_proj_b.weight"),
+        format!("{p}.conv1d.weight"),
+        format!("{p}.A_log"),
+        format!("{p}.dt_bias"),
+        format!("{p}.norm.weight"),
+        format!("{p}.out_proj.weight"),
+    ]
+}
+
+fn full_attention_keys(layer: usize) -> Vec<String> {
+    let p = format!("model.language_model.layers.{layer}.self_attn");
+    vec![
+        format!("{p}.q_proj.weight"),
+        format!("{p}.k_proj.weight"),
+        format!("{p}.v_proj.weight"),
+        format!("{p}.o_proj.weight"),
+        format!("{p}.q_norm.weight"),
+        format!("{p}.k_norm.weight"),
+    ]
+}
+
+fn moe_keys(layer: usize) -> Vec<String> {
+    let p = format!("model.language_model.layers.{layer}.mlp");
+    vec![
+        format!("{p}.gate.weight"),
+        format!("{p}.shared_expert.gate_proj.weight"),
+        format!("{p}.shared_expert.up_proj.weight"),
+        format!("{p}.shared_expert.down_proj.weight"),
+        format!("{p}.shared_expert_gate.weight"),
+        format!("{p}.experts.gate_up_proj"),
+        format!("{p}.experts.down_proj"),
+    ]
+}
+
+fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
+    let index_path = model_dir.join("model.safetensors.index.json");
+    let text = std::fs::read_to_string(&index_path)
+        .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
+    let index: serde_json::Value = serde_json::from_str(&text)
+        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
+    index
+        .get("weight_map")
+        .and_then(|v| v.as_object())
+        .expect("index must contain weight_map")
+        .iter()
+        .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
+        .collect()
+}
+
+fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
+    let mut data = Vec::with_capacity(seq_len * hidden);
+    for i in 0..seq_len * hidden {
+        let x = ((i % 97) as f32 - 48.0) / 128.0;
+        data.push(bf16::from_f32(x));
+    }
+    Tensor::from_slice(&data, &[seq_len, hidden])
+}
+
+fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
+    let cpu = t.to_device(Device::Cpu);
+    cpu.as_slice::<bf16>()
+        .iter()
+        .take(n)
+        .map(|v| v.to_f32())
+        .collect()
+}

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

@@ -5,8 +5,8 @@ use std::sync::{Arc, mpsc};
 use std::thread;
 
 use xserv_model::{
-    BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, SamplingParams,
-    loader, sample, sample_greedy_penalized,
+    BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, ModelFamily, PagedKVCache, Qwen3,
+    SamplingParams, loader, sample, sample_greedy_penalized,
 };
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
@@ -93,24 +93,26 @@ fn tp_worker_loop(
         rank as u32,
     ));
     let weights = loader::load_model_dir(&model_dir, Device::Cpu);
-    let model = if config.is_moe() {
-        ChatModel::GptOss(GptOss::from_weights_tp(
+    let model = match config.family() {
+        ModelFamily::GptOss => ChatModel::GptOss(GptOss::from_weights_tp(
             config.clone(),
             weights,
             rank,
             world,
             rank as u32,
             Some(tp),
-        ))
-    } else {
-        ChatModel::Qwen3(Qwen3::from_weights_tp(
+        )),
+        ModelFamily::Qwen35Moe => {
+            panic!("Qwen3.5/3.6 MoE chat inference is not implemented yet")
+        }
+        _ => ChatModel::Qwen3(Qwen3::from_weights_tp(
             config.clone(),
             weights,
             rank,
             world,
             rank as u32,
             Some(tp),
-        ))
+        )),
     };
     let local_kv = config.num_kv_heads() / world;
     let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
@@ -323,20 +325,27 @@ fn main() {
     );
 
     let config = ModelConfig::from_file(&opts.model_dir.join("config.json"));
-    let model_type = config.model_type.as_deref().unwrap_or("unknown");
-    let is_moe = config.is_moe();
+    let model_type = config.model_type_str();
+    let model_family = config.family();
+    let is_gpt_oss = model_family == ModelFamily::GptOss;
 
     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={}",
-        if is_moe { " (MoE)" } else { "" },
+        if config.is_gpt_oss() { " (MoE)" } else { "" },
         config.num_layers(),
         config.hidden(),
         config.num_heads(),
         config.num_kv_heads(),
-        config.vocab_size,
+        config.vocab_size(),
         max_seq_len
     );
+    if model_family == ModelFamily::Qwen35Moe {
+        eprintln!(
+            "Qwen3.5/3.6 MoE is recognized but chat inference is not implemented yet; it requires DeltaNet/linear-attention and Qwen MoE support."
+        );
+        std::process::exit(1);
+    }
 
     let world = opts.tp;
     if world > 1 {
@@ -374,7 +383,7 @@ fn main() {
         let tp = Arc::new(xserv_distributed::TpContext::init(0, world, id, 0));
         let weights = loader::load_model_dir(&opts.model_dir, Device::Cpu);
         eprintln!("Loaded {} tensors", weights.len());
-        let m = if is_moe {
+        let m = if is_gpt_oss {
             ChatModel::GptOss(GptOss::from_weights_tp(
                 config.clone(),
                 weights,
@@ -412,7 +421,7 @@ fn main() {
         eprintln!("Loading weights...");
         let weights = loader::load_model_dir(&opts.model_dir, Device::Cuda(0));
         eprintln!("Loaded {} tensors", weights.len());
-        let m = if is_moe {
+        let m = if is_gpt_oss {
             ChatModel::GptOss(GptOss::from_weights(config.clone(), weights))
         } else {
             ChatModel::Qwen3(Qwen3::from_weights(config.clone(), weights))
@@ -465,7 +474,7 @@ fn main() {
             _ => {}
         }
 
-        if is_moe {
+        if is_gpt_oss {
             // Harmony multi-turn: re-render the whole conversation (prior
             // analysis dropped) and re-prefill into a freshly cleared slot.
             let prompt =
@@ -495,7 +504,7 @@ fn main() {
                 max_new_tokens,
                 use_color,
                 &tp_handle,
-                is_moe,
+                is_gpt_oss,
                 opts.enable_thinking,
             );
             moe_history.push((input.to_string(), answer));
@@ -540,7 +549,7 @@ fn main() {
             max_new_tokens,
             use_color,
             &tp_handle,
-            is_moe,
+            is_gpt_oss,
             opts.enable_thinking,
         );
         match finish {
@@ -672,6 +681,7 @@ fn parse_args() -> CliOptions {
             temperature,
             top_k,
             top_p,
+            ..SamplingParams::default()
         },
         system_prompt,
         enable_thinking,
@@ -846,25 +856,25 @@ fn generate_with_paged_cache(
     max_tokens: usize,
     use_color: bool,
     tp: &Option<TpHandle>,
-    is_moe: bool,
+    is_gpt_oss: bool,
     enable_thinking: bool,
 ) -> (Finish, String) {
-    let harmony_end_id = if is_moe {
+    let harmony_end_id = if is_gpt_oss {
         tokenizer.special_token_id("<|end|>")
     } else {
         None
     };
-    let harmony_channel_id = if is_moe {
+    let harmony_channel_id = if is_gpt_oss {
         tokenizer.special_token_id("<|channel|>")
     } else {
         None
     };
-    let harmony_message_id = if is_moe {
+    let harmony_message_id = if is_gpt_oss {
         tokenizer.special_token_id("<|message|>")
     } else {
         None
     };
-    let harmony_special: Vec<u32> = if is_moe {
+    let harmony_special: Vec<u32> = if is_gpt_oss {
         [
             "<|channel|>",
             "<|start|>",
@@ -888,7 +898,7 @@ fn generate_with_paged_cache(
         InAnalysis,
         InFinal,
     }
-    let mut hstate = if is_moe {
+    let mut hstate = if is_gpt_oss {
         HarmonyState::InFinal
     } else {
         HarmonyState::Normal
@@ -931,7 +941,7 @@ fn generate_with_paged_cache(
     let mut next = pick(&logits, sampling, &history);
     let mut decode_buffer = Vec::new();
     let mut in_thinking = false;
-    let show_thinking = is_moe && enable_thinking;
+    let show_thinking = is_gpt_oss && enable_thinking;
     // Visible answer tokens, returned for multi-turn history. For moe this is
     // the final-channel content only (analysis is suppressed/gray); for Qwen3
     // it is everything printed. The caller decodes these into the assistant
@@ -1046,7 +1056,7 @@ fn generate_with_paged_cache(
             next = pick(&logits, sampling, &history);
             continue;
         }
-        if is_moe && hstate != HarmonyState::InFinal {
+        if is_gpt_oss && hstate != HarmonyState::InFinal {
             // Between harmony messages (after a channel's <|end|>, before the
             // next <|channel|>): the model emits a role header like "assistant".
             // That's structural, not user-visible content — suppress it. Only

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

@@ -1,23 +1,52 @@
 use std::io::{self, Write};
 use std::path::PathBuf;
-use xserv_model::{BLOCK_SIZE, KVCache, ModelConfig, PagedKVCache, loader};
+use xserv_model::{
+    BLOCK_SIZE, KVCache, ModelConfig, ModelFamily, PagedKVCache, SamplingParams, loader, sample,
+    sample_greedy_penalized,
+};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
+fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(default)
+}
+
+fn pick_next(
+    logits: &xserv_tensor::Tensor,
+    sampling: &SamplingParams,
+    history: &[u32],
+    rep_penalty: f32,
+) -> u32 {
+    if rep_penalty > 1.0 && sampling.temperature == 0.0 {
+        sample_greedy_penalized(logits, history, rep_penalty)
+    } else {
+        sample(logits, sampling)
+    }
+}
+
 fn main() {
     let args: Vec<String> = std::env::args().collect();
     if args.len() < 2 {
-        eprintln!("Usage: xserv-cli <model-dir> [--max-tokens N]");
+        eprintln!(
+            "Usage: xserv-cli <model-dir> [--max-tokens N] [--temperature F] [--top-k N] [--top-p F] [--rep-penalty F] [--rep-window N]"
+        );
         std::process::exit(1);
     }
 
     let model_dir = PathBuf::from(&args[1]);
-    let max_tokens: usize = args
-        .iter()
-        .position(|a| a == "--max-tokens")
-        .and_then(|i| args.get(i + 1))
-        .and_then(|s| s.parse().ok())
-        .unwrap_or(100);
+    let max_tokens = flag(&args, "--max-tokens", 100usize);
+    let sampling = SamplingParams {
+        temperature: flag(&args, "--temperature", 0.0f32),
+        top_k: flag(&args, "--top-k", 0usize),
+        top_p: flag(&args, "--top-p", 1.0f32),
+        ..SamplingParams::default()
+    };
+    let rep_penalty = flag(&args, "--rep-penalty", 1.0f32);
+    let rep_window = flag(&args, "--rep-window", 512usize);
 
     xserv_cuda::device::set_device(0).unwrap();
     let info = xserv_cuda::device::device_info(0).unwrap();
@@ -28,22 +57,29 @@ fn main() {
     );
 
     let config = ModelConfig::from_file(&model_dir.join("config.json"));
-    let model_type = config.model_type.as_deref().unwrap_or("unknown");
+    let model_type = config.model_type_str();
+    let model_family = config.family();
     eprintln!(
         "Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}",
         config.num_layers(),
         config.hidden(),
         config.num_heads(),
         config.num_kv_heads(),
-        config.vocab_size
+        config.vocab_size()
     );
+    if model_family == ModelFamily::Qwen35Moe {
+        eprintln!(
+            "Qwen3.5/3.6 MoE is recognized but CLI inference is not implemented yet; it requires DeltaNet/linear-attention and Qwen MoE support."
+        );
+        std::process::exit(1);
+    }
 
     eprintln!("Loading weights...");
     let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
     eprintln!("Loaded {} tensors", weights.len());
 
-    let is_qwen3 = model_type.contains("qwen");
-    let is_gpt_oss = model_type.contains("gpt_oss");
+    let is_qwen3 = model_family == ModelFamily::Qwen3;
+    let is_gpt_oss = model_family == ModelFamily::GptOss;
     let dtype = if is_qwen3 || is_gpt_oss {
         DType::BF16
     } else {
@@ -65,7 +101,10 @@ fn main() {
     };
 
     let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
-    eprintln!("Ready (KV cache, dtype={dtype}).\n");
+    eprintln!(
+        "Ready (KV cache, dtype={dtype}, temperature={}, top_k={}, top_p={}, rep_penalty={}, rep_window={}).\n",
+        sampling.temperature, sampling.top_k, sampling.top_p, rep_penalty, rep_window
+    );
 
     loop {
         print!("xserv> ");
@@ -74,15 +113,16 @@ fn main() {
         if io::stdin().read_line(&mut input).unwrap() == 0 {
             break;
         }
-        let input = input.trim();
-        if input.is_empty() {
+        let raw_input = input.trim();
+        if raw_input.is_empty() {
             continue;
         }
-        if input == "quit" || input == "exit" {
+        if raw_input == "quit" || raw_input == "exit" {
             break;
         }
+        let input = raw_input.replace("\\n", "\n");
 
-        let token_ids = tokenizer.encode(input);
+        let token_ids = tokenizer.encode(&input);
 
         if is_gpt_oss {
             // GptOss uses paged KV cache
@@ -106,7 +146,9 @@ fn main() {
                 _ => unreachable!(),
             };
             let logits = model.forward_prefill_paged(&token_ids, slot, &mut paged_cache);
-            let mut next = sample_greedy_last(&logits);
+            let mut history = token_ids.clone();
+            let start = history.len().saturating_sub(rep_window);
+            let mut next = pick_next(&logits, &sampling, &history[start..], rep_penalty);
 
             print!("{input}");
             io::stdout().flush().unwrap();
@@ -115,6 +157,7 @@ fn main() {
                 let text = tokenizer.decode(&[next]);
                 print!("{text}");
                 io::stdout().flush().unwrap();
+                history.push(next);
 
                 if tokenizer.eos_token_id() == Some(next) {
                     break;
@@ -122,7 +165,8 @@ fn main() {
 
                 let pos = paged_cache.seq_len(slot);
                 let logits = model.forward_decode_paged(&[next], &[pos], &[slot], &mut paged_cache);
-                next = sample_greedy_last(&logits);
+                let start = history.len().saturating_sub(rep_window);
+                next = pick_next(&logits, &sampling, &history[start..], rep_penalty);
             }
             println!();
             paged_cache.free_sequence(slot);
@@ -145,11 +189,9 @@ fn main() {
                 Model::Qwen3(m) => m.forward_with_cache(&token_ids, &mut cache),
                 Model::GptOss(_) => unreachable!(),
             };
-            let mut next = match &model {
-                Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
-                Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
-                Model::GptOss(_) => unreachable!(),
-            };
+            let mut history = token_ids.clone();
+            let start = history.len().saturating_sub(rep_window);
+            let mut next = pick_next(&logits, &sampling, &history[start..], rep_penalty);
 
             print!("{input}");
             io::stdout().flush().unwrap();
@@ -158,6 +200,7 @@ fn main() {
                 let text = tokenizer.decode(&[next]);
                 print!("{text}");
                 io::stdout().flush().unwrap();
+                history.push(next);
 
                 if tokenizer.eos_token_id() == Some(next) {
                     break;
@@ -168,28 +211,10 @@ fn main() {
                     Model::Qwen3(m) => m.forward_with_cache(&[next], &mut cache),
                     Model::GptOss(_) => unreachable!(),
                 };
-                next = match &model {
-                    Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
-                    Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
-                    Model::GptOss(_) => unreachable!(),
-                };
+                let start = history.len().saturating_sub(rep_window);
+                next = pick_next(&logits, &sampling, &history[start..], rep_penalty);
             }
             println!();
         }
     }
 }
-
-fn sample_greedy_last(logits: &xserv_tensor::Tensor) -> u32 {
-    use half::bf16;
-    assert_eq!(logits.ndim(), 2);
-    let logits_cpu = logits.to_device(Device::Cpu);
-    let vocab_size = logits.shape()[1];
-    let seq_len = logits.shape()[0];
-    let data = logits_cpu.as_slice::<bf16>();
-    let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
-    last.iter()
-        .enumerate()
-        .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
-        .map(|(i, _)| i as u32)
-        .unwrap()
-}

crates/xserv-model/src/config.rs

@@ -1,6 +1,27 @@
 use serde::Deserialize;
 use std::path::Path;
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ModelFamily {
+    Gpt2,
+    Qwen3,
+    GptOss,
+    Qwen35Moe,
+    Unknown,
+}
+
+impl ModelFamily {
+    pub fn as_str(self) -> &'static str {
+        match self {
+            ModelFamily::Gpt2 => "gpt2",
+            ModelFamily::Qwen3 => "qwen3",
+            ModelFamily::GptOss => "gpt_oss",
+            ModelFamily::Qwen35Moe => "qwen3_5_moe",
+            ModelFamily::Unknown => "unknown",
+        }
+    }
+}
+
 #[derive(Debug, Clone, Deserialize)]
 pub struct RopeScaling {
     pub rope_type: Option<String>,
@@ -10,10 +31,21 @@ pub struct RopeScaling {
     pub beta_slow: Option<f64>,
 }
 
+#[derive(Debug, Clone, Deserialize)]
+pub struct RopeParameters {
+    pub rope_type: Option<String>,
+    pub rope_theta: Option<f64>,
+    pub partial_rotary_factor: Option<f64>,
+    pub mrope_interleaved: Option<bool>,
+    pub mrope_section: Option<Vec<usize>>,
+}
+
 #[derive(Debug, Clone, Deserialize)]
 pub struct ModelConfig {
     pub architectures: Option<Vec<String>>,
     pub model_type: Option<String>,
+    #[serde(default)]
+    pub text_config: Option<Box<ModelConfig>>,
 
     // Modern HF naming
     #[serde(default)]
@@ -26,6 +58,7 @@ pub struct ModelConfig {
     pub num_key_value_heads: Option<usize>,
     #[serde(default)]
     pub num_hidden_layers: Option<usize>,
+    #[serde(default)]
     pub vocab_size: usize,
     #[serde(default)]
     pub max_position_embeddings: Option<usize>,
@@ -56,14 +89,22 @@ pub struct ModelConfig {
     #[serde(default)]
     pub tie_word_embeddings: Option<bool>,
 
-    // MoE (gpt-oss)
+    // MoE (gpt-oss / Qwen MoE)
     #[serde(default)]
     pub num_local_experts: Option<usize>,
     #[serde(default)]
+    pub num_experts: Option<usize>,
+    #[serde(default)]
     pub num_experts_per_tok: Option<usize>,
     #[serde(default)]
+    pub moe_intermediate_size: Option<usize>,
+    #[serde(default)]
+    pub shared_expert_intermediate_size: Option<usize>,
+    #[serde(default)]
     pub layer_types: Option<Vec<String>>,
     #[serde(default)]
+    pub full_attention_interval: Option<usize>,
+    #[serde(default)]
     pub sliding_window: Option<usize>,
     #[serde(default)]
     pub attention_bias: Option<bool>,
@@ -72,11 +113,29 @@ pub struct ModelConfig {
     #[serde(default)]
     pub rope_scaling: Option<RopeScaling>,
     #[serde(default)]
+    pub rope_parameters: Option<RopeParameters>,
+    #[serde(default)]
     pub swiglu_limit: Option<f64>,
     #[serde(default)]
     pub geglu_alpha: Option<f64>,
     #[serde(default)]
     pub hidden_act: Option<String>,
+
+    // Qwen3.5/3.6 linear attention / DeltaNet metadata
+    #[serde(default)]
+    pub linear_conv_kernel_dim: Option<usize>,
+    #[serde(default)]
+    pub linear_key_head_dim: Option<usize>,
+    #[serde(default)]
+    pub linear_num_key_heads: Option<usize>,
+    #[serde(default)]
+    pub linear_num_value_heads: Option<usize>,
+    #[serde(default)]
+    pub linear_value_head_dim: Option<usize>,
+    #[serde(default)]
+    pub partial_rotary_factor: Option<f64>,
+    #[serde(default)]
+    pub mtp_num_hidden_layers: Option<usize>,
 }
 
 impl ModelConfig {
@@ -87,80 +146,153 @@ impl ModelConfig {
             .unwrap_or_else(|e| panic!("failed to parse {}: {e}", path.display()))
     }
 
+    pub fn text(&self) -> &Self {
+        self.text_config.as_deref().unwrap_or(self)
+    }
+
+    pub fn model_type_str(&self) -> &str {
+        self.text()
+            .model_type
+            .as_deref()
+            .or(self.model_type.as_deref())
+            .unwrap_or("unknown")
+    }
+
+    pub fn family(&self) -> ModelFamily {
+        let top_type = self.model_type.as_deref().unwrap_or("");
+        let text_type = self.text().model_type.as_deref().unwrap_or(top_type);
+        let has_arch = |needle: &str| {
+            self.architectures
+                .as_ref()
+                .map(|arch| arch.iter().any(|a| a.contains(needle)))
+                .unwrap_or(false)
+        };
+
+        if top_type == "qwen3_5_moe" || text_type == "qwen3_5_moe_text" || has_arch("Qwen3_5Moe") {
+            ModelFamily::Qwen35Moe
+        } else if text_type == "gpt_oss" || top_type == "gpt_oss" {
+            ModelFamily::GptOss
+        } else if text_type.contains("qwen") || top_type.contains("qwen") {
+            ModelFamily::Qwen3
+        } else if text_type.contains("gpt2") || top_type.contains("gpt2") {
+            ModelFamily::Gpt2
+        } else {
+            ModelFamily::Unknown
+        }
+    }
+
     pub fn hidden(&self) -> usize {
-        self.hidden_size
-            .or(self.n_embd)
+        let c = self.text();
+        c.hidden_size
+            .or(c.n_embd)
             .expect("hidden_size or n_embd required")
     }
 
     pub fn num_heads(&self) -> usize {
-        self.num_attention_heads
-            .or(self.n_head)
+        let c = self.text();
+        c.num_attention_heads
+            .or(c.n_head)
             .expect("num_attention_heads or n_head required")
     }
 
     pub fn num_layers(&self) -> usize {
-        self.num_hidden_layers
-            .or(self.n_layer)
+        let c = self.text();
+        c.num_hidden_layers
+            .or(c.n_layer)
             .expect("num_hidden_layers or n_layer required")
     }
 
     pub fn max_seq_len(&self) -> usize {
-        self.max_position_embeddings
-            .or(self.n_positions)
-            .unwrap_or(2048)
+        let c = self.text();
+        c.max_position_embeddings.or(c.n_positions).unwrap_or(2048)
     }
 
     pub fn ffn_hidden(&self) -> usize {
-        self.intermediate_size
-            .or(self.n_inner)
-            .unwrap_or(self.hidden() * 4)
+        let c = self.text();
+        c.intermediate_size
+            .or(c.moe_intermediate_size)
+            .or(c.n_inner)
+            .unwrap_or_else(|| self.hidden() * 4)
+    }
+
+    pub fn vocab_size(&self) -> usize {
+        self.text().vocab_size
     }
 
     pub fn num_kv_heads(&self) -> usize {
-        self.num_key_value_heads.unwrap_or(self.num_heads())
+        self.text().num_key_value_heads.unwrap_or(self.num_heads())
     }
 
     pub fn head_dim(&self) -> usize {
-        self.explicit_head_dim
+        self.text()
+            .explicit_head_dim
             .unwrap_or_else(|| self.hidden() / self.num_heads())
     }
 
     pub fn ln_eps(&self) -> f32 {
-        self.layer_norm_eps
-            .or(self.layer_norm_epsilon)
+        let c = self.text();
+        c.layer_norm_eps
+            .or(c.layer_norm_epsilon)
+            .or(c.rms_norm_eps)
             .unwrap_or(1e-5) as f32
     }
 
+    pub fn rope_theta_value(&self) -> Option<f64> {
+        let c = self.text();
+        c.rope_theta
+            .or_else(|| c.rope_parameters.as_ref().and_then(|rp| rp.rope_theta))
+    }
+
     pub fn tied_embeddings(&self) -> bool {
-        self.tie_word_embeddings.unwrap_or(true)
+        self.text().tie_word_embeddings.unwrap_or(true)
     }
 
     pub fn num_experts(&self) -> usize {
-        self.num_local_experts.unwrap_or(0)
+        self.text()
+            .num_local_experts
+            .or(self.text().num_experts)
+            .unwrap_or(0)
     }
 
     pub fn experts_per_token(&self) -> usize {
-        self.num_experts_per_tok.unwrap_or(1)
+        self.text().num_experts_per_tok.unwrap_or(1)
     }
 
     pub fn is_moe(&self) -> bool {
-        self.num_local_experts.unwrap_or(0) > 1
+        self.num_experts() > 1
+    }
+
+    pub fn is_gpt_oss(&self) -> bool {
+        self.family() == ModelFamily::GptOss
+    }
+
+    pub fn is_qwen35_moe(&self) -> bool {
+        self.family() == ModelFamily::Qwen35Moe
     }
 
     pub fn is_sliding_layer(&self, layer_idx: usize) -> bool {
-        self.layer_types
+        self.text()
+            .layer_types
             .as_ref()
             .and_then(|lt| lt.get(layer_idx))
             .map(|t| t == "sliding_attention")
             .unwrap_or(false)
     }
 
+    pub fn is_linear_attention_layer(&self, layer_idx: usize) -> bool {
+        self.text()
+            .layer_types
+            .as_ref()
+            .and_then(|lt| lt.get(layer_idx))
+            .map(|t| t == "linear_attention")
+            .unwrap_or(false)
+    }
+
     pub fn window_size(&self) -> usize {
-        self.sliding_window.unwrap_or(0)
+        self.text().sliding_window.unwrap_or(0)
     }
 
     pub fn geglu_alpha(&self) -> f32 {
-        self.geglu_alpha.unwrap_or(1.702) as f32
+        self.text().geglu_alpha.unwrap_or(1.702) as f32
     }
 }

crates/xserv-model/src/decode_graph.rs

@@ -9,7 +9,7 @@
 use std::ffi::c_void;
 use xserv_cuda::{CudaGraph, CudaStream, GpuBuffer};
 use xserv_kernels::dispatch;
-use xserv_kernels::gemm::cublas_handle;
+use xserv_kernels::gemm::{cublas_handle, gemv_scratch_elems};
 
 use crate::config::ModelConfig;
 use crate::kv_cache::GpuKVCache;
@@ -54,7 +54,7 @@ struct DecodeBuffers {
     up: GpuBuffer,       // [1, intermediate]
     silu_out: GpuBuffer, // [1, intermediate]
 
-    // GEMV fp32 accumulators (separate per output dimension)
+    // GEMV fp32 scratch for deterministic K-block partials.
     fp32_hidden: GpuBuffer,       // for hidden-sized GEMV outputs
     fp32_q: GpuBuffer,            // for Q projection
     fp32_kv: GpuBuffer,           // for K/V projection
@@ -98,9 +98,9 @@ impl DecodeGraphState {
         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 vocab_size = config.vocab_size();
         let num_layers = config.num_layers();
-        let eps = config.rms_norm_eps.unwrap_or(1e-6) as f32;
+        let eps = config.ln_eps();
         let es = 2usize; // BF16 = 2 bytes
 
         let stream = CudaStream::new().expect("create CUDA stream for graph");
@@ -140,11 +140,14 @@ impl DecodeGraphState {
             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),
+            fp32_hidden: alloc(
+                gemv_scratch_elems(hidden, hidden).max(gemv_scratch_elems(intermediate, hidden))
+                    * 4,
+            ),
+            fp32_q: alloc(gemv_scratch_elems(hidden, num_heads * head_dim) * 4),
+            fp32_kv: alloc(gemv_scratch_elems(hidden, num_kv_heads * head_dim) * 4),
+            fp32_intermediate: alloc(gemv_scratch_elems(hidden, intermediate) * 4),
+            fp32_vocab: alloc(gemv_scratch_elems(hidden, vocab_size) * 4),
 
             token_id_gpu: alloc(4),
             position_gpu: alloc(4),

crates/xserv-model/src/eagle3.rs

@@ -0,0 +1,425 @@
+//! EAGLE3 speculative draft head for Qwen3-8B (Phase 25).
+//!
+//! Loads the AngelSlim/Qwen3-8B_eagle3 pytorch_model.bin and provides a
+//! single-step forward pass that takes 3 target hidden states + the previous
+//! token and returns a draft token in the target vocabulary.
+//!
+//! Architecture (from weights):
+//! - fc:           [hidden, 3*hidden] → fuse 3 target hidden states
+//! - midlayer:     1 decoder layer (attn input dim = 2*hidden)
+//! - norm + lm_head: → [draft_vocab_size=32000]
+//! - d2t:          draft_id → target_id offset mapping
+
+use std::collections::HashMap;
+use std::path::Path;
+use xserv_kernels::*;
+use xserv_tensor::{DType, Device, Tensor};
+
+/// Target layers to hook for EAGLE3 auxiliary hidden states, for Qwen3-8B
+/// (36 layers). Value comes from AngelSlim/vLLM speculators training config
+/// `dflash_qwen3_8b_sharegpt_online_5k.sh` which specifies target_layer_ids
+/// = "2 18 33". Must match training-time selection or EAGLE outputs are wrong.
+pub const EAGLE_HOOK_LAYERS: [usize; 3] = [2, 18, 33];
+const DRAFT_VOCAB_SIZE: usize = 32000;
+
+fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
+    assert_eq!(a.ndim(), 2);
+    assert_eq!(b.ndim(), 2);
+    matmul(a, b, GemmBackend::CuBlas)
+}
+
+pub struct Eagle3Head {
+    fc_wt: Tensor,                    // [hidden, 3*hidden] transposed for matmul
+    hidden_norm: Tensor,              // [hidden]
+    input_layernorm: Tensor,          // [hidden]
+    q_proj_wt: Tensor,                // [num_heads*head_dim, 2*hidden]
+    k_proj_wt: Tensor,                // [num_kv_heads*head_dim, 2*hidden]
+    v_proj_wt: Tensor,                // [num_kv_heads*head_dim, 2*hidden]
+    o_proj_wt: Tensor,                // [hidden, num_heads*head_dim]
+    gate_proj_wt: Tensor,             // [intermediate, hidden]
+    up_proj_wt: Tensor,               // [intermediate, hidden]
+    down_proj_wt: Tensor,             // [hidden, intermediate]
+    post_attention_layernorm: Tensor, // [hidden]
+    norm: Tensor,                     // [hidden] final
+    lm_head_wt: Tensor,               // [draft_vocab, hidden]
+    d2t: Vec<i64>,                    // [draft_vocab] offset mapping
+    /// t2d[target_id] = true iff target_id has a corresponding draft-vocab id
+    /// (i.e. can potentially be produced by EAGLE). Used to measure the
+    /// coverage cap on acceptance.
+    t2d: Vec<bool>,
+    hidden_size: usize,
+    num_heads: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    max_seq_len: usize,
+    rope_cache: RopeCache,
+    // Stateful 1-layer KV cache: [1, num_kv_heads, max_seq_len, head_dim] BF16.
+    // We slice `..current_len` for attention. The head is tiny (~64 KB per
+    // 1000 tokens) so pre-allocating max_seq_len wastes negligible memory.
+    k_cache: Tensor,
+    v_cache: Tensor,
+    current_len: usize,
+}
+
+impl Eagle3Head {
+    pub fn load(dir: &Path, device: u32) -> Self {
+        let (weights, d2t, t2d) = load_eagle3_weights(dir, device);
+        let hidden_size = 4096;
+        let num_heads = 32;
+        let num_kv_heads = 8;
+        let head_dim = 128;
+        let intermediate_size = 12288;
+        let max_seq_len = 2048;
+        let rope_theta = 1_000_000.0f32;
+
+        let get = |name: &str| -> Tensor {
+            weights
+                .get(name)
+                .unwrap_or_else(|| panic!("missing eagle3 weight: {name}"))
+                .clone()
+        };
+
+        let fc_wt = get("fc.weight").transpose(0, 1).contiguous();
+        let q_proj_wt = get("midlayer.self_attn.q_proj.weight")
+            .transpose(0, 1)
+            .contiguous();
+        let k_proj_wt = get("midlayer.self_attn.k_proj.weight")
+            .transpose(0, 1)
+            .contiguous();
+        let v_proj_wt = get("midlayer.self_attn.v_proj.weight")
+            .transpose(0, 1)
+            .contiguous();
+        let o_proj_wt = get("midlayer.self_attn.o_proj.weight")
+            .transpose(0, 1)
+            .contiguous();
+        let gate_proj_wt = get("midlayer.mlp.gate_proj.weight")
+            .transpose(0, 1)
+            .contiguous();
+        let up_proj_wt = get("midlayer.mlp.up_proj.weight")
+            .transpose(0, 1)
+            .contiguous();
+        let down_proj_wt = get("midlayer.mlp.down_proj.weight")
+            .transpose(0, 1)
+            .contiguous();
+        let hidden_norm = get("midlayer.hidden_norm.weight");
+        let input_layernorm = get("midlayer.input_layernorm.weight");
+        let post_attention_layernorm = get("midlayer.post_attention_layernorm.weight");
+        let norm = get("norm.weight");
+        let lm_head_wt = get("lm_head.weight").transpose(0, 1).contiguous();
+
+        assert_eq!(d2t.len(), DRAFT_VOCAB_SIZE);
+
+        let rope_cache = RopeCache::new(max_seq_len, head_dim, rope_theta);
+
+        let k_cache = Tensor::zeros(
+            &[1, num_kv_heads, max_seq_len, head_dim],
+            DType::BF16,
+            Device::Cuda(device),
+        );
+        let v_cache = Tensor::zeros(
+            &[1, num_kv_heads, max_seq_len, head_dim],
+            DType::BF16,
+            Device::Cuda(device),
+        );
+
+        Self {
+            fc_wt,
+            hidden_norm,
+            input_layernorm,
+            q_proj_wt,
+            k_proj_wt,
+            v_proj_wt,
+            o_proj_wt,
+            gate_proj_wt,
+            up_proj_wt,
+            down_proj_wt,
+            post_attention_layernorm,
+            norm,
+            lm_head_wt,
+            d2t,
+            t2d,
+            hidden_size,
+            num_heads,
+            num_kv_heads,
+            head_dim,
+            max_seq_len,
+            rope_cache,
+            k_cache,
+            v_cache,
+            current_len: 0,
+        }
+    }
+
+    /// Reset the internal KV cache for a fresh sequence.
+    pub fn reset(&mut self) {
+        self.current_len = 0;
+    }
+
+    /// Truncate the internal KV cache to `new_len` entries. Used to discard
+    /// K/V of rejected drafts after a speculative round.
+    pub fn truncate_to(&mut self, new_len: usize) {
+        assert!(new_len <= self.current_len);
+        self.current_len = new_len;
+    }
+
+    /// Current number of committed K/V entries in the internal EAGLE cache.
+    pub fn current_len(&self) -> usize {
+        self.current_len
+    }
+
+    /// One draft step: produce a token in target vocabulary space.
+    ///
+    /// - `target_hidden`: 3 tensors [1, hidden_size] from target hook layers
+    /// - `embed_table`: the target model's embed_tokens (shared, not copied)
+    /// - `prev_token`: the previous committed token
+    /// - `position`: the decode position for RoPE
+    ///
+    /// Returns (draft_token_in_target_vocab, draft_logits_tensor).
+    pub fn step(
+        &mut self,
+        target_hidden: &[Tensor; 3],
+        embed_table: &Tensor,
+        prev_token: u32,
+        position: usize,
+    ) -> (u32, Tensor) {
+        let (id, logits, _) = self.step_with_aux(target_hidden, embed_table, prev_token, position);
+        (id, logits)
+    }
+
+    /// Like `step`, but also returns the final hidden state (aux) usable as
+    /// the fused_h for a subsequent recursive draft step via `step_recursive`.
+    pub fn step_with_aux(
+        &mut self,
+        target_hidden: &[Tensor; 3],
+        embed_table: &Tensor,
+        prev_token: u32,
+        position: usize,
+    ) -> (u32, Tensor, Tensor) {
+        // Fuse 3 target hidden states into fused_h via fc.
+        let h_cat = concat_hidden(target_hidden);
+        let fused_h = matmul_2d(&h_cat, &self.fc_wt);
+        self.forward_from_fused(fused_h, embed_table, prev_token, position)
+    }
+
+    /// Recursive draft step: reuses the previous EAGLE step's aux as fused_h,
+    /// bypassing the fc+3-hidden fusion. Used for γ≥2 chained drafts.
+    pub fn step_recursive(
+        &mut self,
+        fused_h: Tensor,
+        embed_table: &Tensor,
+        prev_token: u32,
+        position: usize,
+    ) -> (u32, Tensor, Tensor) {
+        self.forward_from_fused(fused_h, embed_table, prev_token, position)
+    }
+
+    fn forward_from_fused(
+        &mut self,
+        fused_h: Tensor,
+        embed_table: &Tensor,
+        prev_token: u32,
+        position: usize,
+    ) -> (u32, Tensor, Tensor) {
+        let eps = 1e-6f32;
+        assert!(
+            self.current_len < self.max_seq_len,
+            "EAGLE KV cache overflow: {} >= {}",
+            self.current_len,
+            self.max_seq_len
+        );
+
+        let emb = embedding(embed_table, &[prev_token]);
+        let residual = fused_h.clone();
+        let emb_normed = rmsnorm(&emb, &self.input_layernorm, eps);
+        let h_normed = rmsnorm(&fused_h, &self.hidden_norm, eps);
+        let attn_in = concat_last_dim(&emb_normed, &h_normed);
+
+        let q = matmul_2d(&attn_in, &self.q_proj_wt);
+        let k = matmul_2d(&attn_in, &self.k_proj_wt);
+        let v = matmul_2d(&attn_in, &self.v_proj_wt);
+
+        let q_3d = q.reshape(&[1, self.num_heads, self.head_dim]);
+        let k_3d = k.reshape(&[1, self.num_kv_heads, self.head_dim]);
+        let positions = [position as u32];
+        rope_inplace(&q_3d, &self.rope_cache, &positions);
+        rope_inplace(&k_3d, &self.rope_cache, &positions);
+
+        let v_3d = v.reshape(&[1, self.num_kv_heads, self.head_dim]);
+        self.append_to_kv_cache(&k_3d, &v_3d);
+        self.current_len += 1;
+        let kv_len = self.current_len;
+        let k_view = self.k_cache.narrow(2, 0, kv_len).contiguous();
+        let v_view = self.v_cache.narrow(2, 0, kv_len).contiguous();
+
+        let q_4d = q_3d.reshape(&[1, self.num_heads, 1, self.head_dim]);
+        let attn_out = decode_attention(&q_4d, &k_view, &v_view);
+
+        let attn_merged = attn_out.reshape(&[1, self.num_heads * self.head_dim]);
+        let attn_proj = matmul_2d(&attn_merged, &self.o_proj_wt);
+
+        let (mlp_in, residual) =
+            add_rmsnorm(&attn_proj, &residual, &self.post_attention_layernorm, eps);
+
+        let gate = matmul_2d(&mlp_in, &self.gate_proj_wt);
+        let up = matmul_2d(&mlp_in, &self.up_proj_wt);
+        let hidden = silu_mul(&gate, &up);
+        let down = matmul_2d(&hidden, &self.down_proj_wt);
+
+        let (x, prenorm) = add_rmsnorm(&down, &residual, &self.norm, eps);
+        let logits = matmul_2d(&x, &self.lm_head_wt);
+
+        let draft_id = argmax_bf16_single(&logits);
+        let target_id = (draft_id as i64 + self.d2t[draft_id as usize]) as u32;
+        // aux for recursive drafting = PRE-norm hidden (default norm_output=False
+        // in vllm/llama_eagle3.py). Feeding the pre-norm state matches training.
+        (target_id, logits, prenorm)
+    }
+
+    /// Write new K/V rows (shape [1, num_kv_heads, head_dim]) at position
+    /// `current_len` inside the [1, num_kv_heads, max_seq_len, head_dim] cache.
+    fn append_to_kv_cache(&mut self, new_k: &Tensor, new_v: &Tensor) {
+        let head_bytes = self.head_dim * self.k_cache.dtype().size_bytes();
+        for h in 0..self.num_kv_heads {
+            for (cache, src) in [(&self.k_cache, new_k), (&self.v_cache, new_v)] {
+                let dst = unsafe {
+                    (cache.data_ptr() as *mut u8)
+                        .add(((h * self.max_seq_len) + self.current_len) * head_bytes)
+                };
+                let s = unsafe { (src.data_ptr() as *const u8).add(h * head_bytes) };
+                d2d(dst, s, head_bytes);
+            }
+        }
+    }
+
+    /// Map a draft-vocab token id to the full target-vocab id via d2t.
+    pub fn map_draft_to_target(&self, draft_id: u32) -> u32 {
+        (draft_id as i64 + self.d2t[draft_id as usize]) as u32
+    }
+
+    /// Returns true iff `target_id` is representable in the draft vocabulary
+    /// (i.e., EAGLE could in principle produce it).
+    pub fn target_id_in_draft_vocab(&self, target_id: u32) -> bool {
+        self.t2d.get(target_id as usize).copied().unwrap_or(false)
+    }
+}
+
+fn d2d(dst: *mut u8, src: *const u8, bytes: usize) {
+    unsafe {
+        xserv_cuda::ffi::cudaMemcpy(dst, src, bytes, xserv_cuda::ffi::CUDA_MEMCPY_D2D);
+    }
+}
+
+fn concat_hidden(hidden: &[Tensor; 3]) -> Tensor {
+    let h = hidden[0].shape()[1];
+    let dtype = hidden[0].dtype();
+    let device = hidden[0].device();
+    let elem_bytes = dtype.size_bytes();
+    let out = Tensor::empty(&[1, 3 * h], dtype, device);
+    for (i, t) in hidden.iter().enumerate() {
+        assert!(t.is_contiguous());
+        let dst = unsafe { (out.data_ptr() as *mut u8).add(i * h * elem_bytes) };
+        d2d(dst, t.data_ptr() as *const u8, h * elem_bytes);
+    }
+    out
+}
+
+fn concat_last_dim(a: &Tensor, b: &Tensor) -> Tensor {
+    let da = a.shape()[1];
+    let db = b.shape()[1];
+    let dtype = a.dtype();
+    let device = a.device();
+    let elem_bytes = dtype.size_bytes();
+    let out = Tensor::empty(&[1, da + db], dtype, device);
+    d2d(
+        out.data_ptr() as *mut u8,
+        a.data_ptr() as *const u8,
+        da * elem_bytes,
+    );
+    let dst = unsafe { (out.data_ptr() as *mut u8).add(da * elem_bytes) };
+    d2d(dst, b.data_ptr() as *const u8, db * elem_bytes);
+    out
+}
+
+fn repeat_kv_for_single_token(kv: &Tensor, repeats: usize) -> Tensor {
+    if repeats == 1 {
+        return kv.clone();
+    }
+    let nkv = kv.shape()[1];
+    let d = kv.shape()[2];
+    let dtype = kv.dtype();
+    let device = kv.device();
+    let head_bytes = d * dtype.size_bytes();
+    let out = Tensor::empty(&[1, nkv * repeats, d], dtype, device);
+    for h in 0..nkv {
+        let src = unsafe { (kv.data_ptr() as *const u8).add(h * head_bytes) };
+        for r in 0..repeats {
+            let dst = unsafe { (out.data_ptr() as *mut u8).add((h * repeats + r) * head_bytes) };
+            d2d(dst, src, head_bytes);
+        }
+    }
+    out
+}
+
+/// Load EAGLE3 weights from safetensors, handling int64 d2t + bool t2d specially.
+fn load_eagle3_weights(dir: &Path, device: u32) -> (HashMap<String, Tensor>, Vec<i64>, Vec<bool>) {
+    let st_path = dir.join("model.safetensors");
+    assert!(
+        st_path.exists(),
+        "Eagle3 model.safetensors not found in {}. Convert with:\n\
+         python3 -c \"import torch; from safetensors.torch import save_file; \
+         sd=torch.load('pytorch_model.bin', map_location='cpu', weights_only=False); \
+         save_file(sd, 'model.safetensors')\"",
+        dir.display()
+    );
+
+    let data = std::fs::read(&st_path)
+        .unwrap_or_else(|e| panic!("failed to read {}: {e}", st_path.display()));
+    let st = safetensors::SafeTensors::deserialize(&data)
+        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", st_path.display()));
+
+    let mut tensors = HashMap::new();
+    let mut d2t_vec: Vec<i64> = Vec::new();
+    let mut t2d_vec: Vec<bool> = Vec::new();
+
+    for (name, view) in st.tensors() {
+        if name == "t2d" {
+            let raw = view.data();
+            assert_eq!(view.dtype(), safetensors::Dtype::BOOL);
+            t2d_vec = raw.iter().map(|&b| b != 0).collect();
+            continue;
+        }
+        if name == "d2t" {
+            let raw = view.data();
+            assert_eq!(view.dtype(), safetensors::Dtype::I64);
+            let n = raw.len() / 8;
+            d2t_vec = (0..n)
+                .map(|i| i64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap()))
+                .collect();
+            continue;
+        }
+        let dtype = match view.dtype() {
+            safetensors::Dtype::BF16 => DType::BF16,
+            safetensors::Dtype::F32 => DType::F32,
+            safetensors::Dtype::F16 => DType::F16,
+            other => {
+                eprintln!("eagle3: skipping {name} with unsupported dtype {other:?}");
+                continue;
+            }
+        };
+        let shape: Vec<usize> = view.shape().to_vec();
+        let raw = view.data();
+        let t = crate::loader::make_tensor(raw, &shape, dtype);
+        let t = t.to_device(Device::Cuda(device));
+        tensors.insert(name.to_string(), t);
+    }
+
+    assert!(
+        !d2t_vec.is_empty(),
+        "d2t tensor not found in eagle3 weights"
+    );
+    assert!(
+        !t2d_vec.is_empty(),
+        "t2d tensor not found in eagle3 weights"
+    );
+    (tensors, d2t_vec, t2d_vec)
+}

crates/xserv-model/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod config;
 pub mod decode_graph;
+pub mod eagle3;
 pub mod gpt2;
 pub mod gpt_oss;
 pub mod gpt_oss_graph;
@@ -7,9 +8,11 @@ pub mod kv_cache;
 pub mod loader;
 pub mod paged_kv_cache;
 pub mod qwen3;
+pub mod qwen35_moe;
+pub mod qwen3_graph;
 pub mod sampling;
 
-pub use config::ModelConfig;
+pub use config::{ModelConfig, ModelFamily};
 pub use decode_graph::{DecodeGraphState, LayerWeightPtrs};
 pub use gpt_oss::GptOss;
 pub use gpt_oss_graph::{GptOssDecodeGraph, GraphedGptOssDecoder};
@@ -17,7 +20,12 @@ pub use gpt2::{GPT2, KVCache};
 pub use kv_cache::GpuKVCache;
 pub use paged_kv_cache::{BLOCK_SIZE, BlockAllocator, Location, PagedKVCache};
 pub use qwen3::Qwen3;
-pub use sampling::{SamplingParams, sample, sample_greedy_penalized};
+pub use qwen35_moe::{
+    Qwen35AttentionMeta, Qwen35FullAttentionOutput, Qwen35FullAttentionWeights,
+    Qwen35LinearAttentionWeights, Qwen35LinearProjectionOutput, Qwen35Moe, Qwen35MoeSpec,
+    Qwen35RecurrentCache, Qwen35TensorMeta, Qwen35TensorSpec,
+};
+pub use sampling::{SamplingParams, sample, sample_greedy_penalized, sample_with_history};
 
 /// Initialize GPU kernel hooks. Called automatically by model constructors,
 /// but safe to call multiple times (idempotent via OnceLock).

crates/xserv-model/src/loader.rs

@@ -1,10 +1,20 @@
 use half::{bf16, f16};
 use safetensors::SafeTensors;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::path::Path;
 use xserv_tensor::{DType, Device, Tensor};
 
 pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor> {
+    load_safetensors_filtered(path, device, |_| true)
+}
+
+/// Load only tensors accepted by `keep`. The safetensors file is still read as
+/// one byte buffer, but unneeded tensor payloads are not copied into Tensors.
+pub fn load_safetensors_filtered(
+    path: &Path,
+    device: Device,
+    keep: impl Fn(&str) -> bool,
+) -> HashMap<String, Tensor> {
     let data =
         std::fs::read(path).unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
     let st = SafeTensors::deserialize(&data)
@@ -13,6 +23,9 @@ pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor>
     let mut tensors = HashMap::new();
 
     for (name, view) in st.tensors() {
+        if !keep(&name) {
+            continue;
+        }
         let shape: Vec<usize> = view.shape().to_vec();
         let raw_bytes = view.data();
         let dtype = match view.dtype() {
@@ -36,27 +49,64 @@ pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor>
 
 /// Load from a directory containing model.safetensors (or sharded files) + config.json.
 pub fn load_model_dir(dir: &Path, device: Device) -> HashMap<String, Tensor> {
+    load_model_dir_filtered(dir, device, |_| true)
+}
+
+/// Load a filtered subset of a model directory. For indexed sharded models,
+/// consult `model.safetensors.index.json` first so shards containing no wanted
+/// tensors are never read. This is critical for pipeline parallelism on large
+/// models: each stage should read only its layer range, not the full checkpoint.
+pub fn load_model_dir_filtered(
+    dir: &Path,
+    device: Device,
+    keep: impl Fn(&str) -> bool,
+) -> HashMap<String, Tensor> {
     let single = dir.join("model.safetensors");
     if single.exists() {
-        return load_safetensors(&single, device);
+        return load_safetensors_filtered(&single, device, keep);
     }
 
+    let index_path = dir.join("model.safetensors.index.json");
+    let wanted_shards: Option<HashSet<String>> = if index_path.exists() {
+        let text = std::fs::read_to_string(&index_path)
+            .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
+        let index: serde_json::Value = serde_json::from_str(&text)
+            .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
+        let map = index
+            .get("weight_map")
+            .and_then(|v| v.as_object())
+            .unwrap_or_else(|| panic!("{} has no weight_map", index_path.display()));
+        Some(
+            map.iter()
+                .filter(|(name, _)| keep(name))
+                .filter_map(|(_, shard)| shard.as_str().map(str::to_string))
+                .collect(),
+        )
+    } else {
+        None
+    };
+
     // Try sharded: model-00001-of-NNNNN.safetensors
     let mut all_tensors = HashMap::new();
     let mut entries: Vec<_> = std::fs::read_dir(dir)
         .unwrap()
         .filter_map(|e| e.ok())
         .filter(|e| {
-            e.path()
+            let is_safetensors = e
+                .path()
                 .file_name()
                 .map(|f| f.to_string_lossy().ends_with(".safetensors"))
-                .unwrap_or(false)
+                .unwrap_or(false);
+            let is_wanted = wanted_shards.as_ref().is_none_or(|wanted| {
+                wanted.contains(&e.file_name().to_string_lossy().to_string())
+            });
+            is_safetensors && is_wanted
         })
         .collect();
     entries.sort_by_key(|e| e.file_name());
 
     for entry in entries {
-        let tensors = load_safetensors(&entry.path(), device);
+        let tensors = load_safetensors_filtered(&entry.path(), device, &keep);
         all_tensors.extend(tensors);
     }
 
@@ -68,7 +118,7 @@ pub fn load_model_dir(dir: &Path, device: Device) -> HashMap<String, Tensor> {
     all_tensors
 }
 
-fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
+pub(crate) fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
     match dtype {
         DType::F32 => {
             let floats: &[f32] = unsafe {

crates/xserv-model/src/paged_kv_cache.rs

@@ -486,6 +486,80 @@ impl PagedKVCache {
         state.seq_len += num_tokens;
     }
 
+    /// Roll a registered sequence back to `new_len` tokens.
+    ///
+    /// This only changes cache metadata and frees whole physical blocks that are
+    /// no longer reachable. Bytes inside retained blocks are left untouched; the
+    /// logical `seq_len` prevents attention from reading them, and later writes
+    /// to the same positions overwrite them.
+    pub fn truncate_sequence(&mut self, slot: usize, new_len: usize) -> Result<(), &'static str> {
+        if slot >= self.max_seqs {
+            return Err("truncate_sequence: slot out of range");
+        }
+        let state = self.seq_states[slot]
+            .as_mut()
+            .ok_or("truncate_sequence: empty slot")?;
+        if new_len > state.seq_len {
+            return Err("truncate_sequence: cannot extend");
+        }
+
+        let needed_blocks = ((new_len + BLOCK_SIZE - 1) / BLOCK_SIZE).max(1);
+        while state.block_ids.len() > needed_blocks {
+            let block = state.block_ids.pop().expect("checked len");
+            match state.location {
+                Location::Gpu => self.allocator.free(block),
+                Location::Cpu => self.cpu_allocator.free(block),
+            }
+        }
+        state.seq_len = new_len;
+        Ok(())
+    }
+
+    /// Copy K/V data from `src_pos` to `dst_pos` within the same slot, across
+    /// all layers. Used by tree speculative decoding to remap an accepted
+    /// sibling's K/V to the canonical sequential position after acceptance.
+    ///
+    /// Requires: both positions within the currently-allocated block range.
+    pub fn copy_kv_position(&self, slot: usize, src_pos: usize, dst_pos: usize) {
+        let state = self.seq_states[slot]
+            .as_ref()
+            .expect("copy_kv_position: slot not registered");
+        assert!(
+            src_pos < state.seq_len && dst_pos < state.seq_len,
+            "copy_kv_position: positions must be within seq_len"
+        );
+        // Upload this sequence's block_ids to a small GPU buffer.
+        let block_ids_host: Vec<i32> = state.block_ids.iter().map(|&b| b as i32).collect();
+        let bytes: &[u8] = unsafe {
+            std::slice::from_raw_parts(
+                block_ids_host.as_ptr() as *const u8,
+                block_ids_host.len() * 4,
+            )
+        };
+        let mut ids_buf =
+            xserv_cuda::allocator::cached_alloc(bytes.len()).expect("alloc block_ids for copy");
+        ids_buf.copy_from_host(bytes).unwrap();
+        let ids_ptr = ids_buf.as_ptr() as *const i32;
+
+        let stream = xserv_cuda::current_stream_raw();
+        let num_layers = self.k_pools.len();
+        for layer in 0..num_layers {
+            unsafe {
+                xserv_kernels::copy_kv_position(
+                    self.k_pools[layer].as_ptr() as *mut std::ffi::c_void,
+                    self.v_pools[layer].as_ptr() as *mut std::ffi::c_void,
+                    ids_ptr,
+                    src_pos,
+                    dst_pos,
+                    self.num_kv_heads,
+                    self.head_dim,
+                    BLOCK_SIZE,
+                    stream,
+                );
+            }
+        }
+    }
+
     /// 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) {
@@ -748,6 +822,71 @@ impl PagedKVCache {
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn tiny_config() -> ModelConfig {
+        serde_json::from_value(serde_json::json!({
+            "model_type": "qwen3",
+            "hidden_size": 8,
+            "intermediate_size": 16,
+            "num_attention_heads": 1,
+            "num_key_value_heads": 1,
+            "num_hidden_layers": 1,
+            "vocab_size": 32,
+            "max_position_embeddings": 64
+        }))
+        .unwrap()
+    }
+
+    #[test]
+    fn truncate_sequence_frees_whole_blocks_and_keeps_slot_registered() {
+        if xserv_cuda::device::set_device(0).is_err() {
+            eprintln!("skipping CUDA-backed PagedKVCache test: device 0 unavailable");
+            return;
+        }
+
+        let config = tiny_config();
+        let mut cache = PagedKVCache::new(&config, 5, 0, 1, 4, DType::BF16, 0);
+
+        assert_eq!(
+            cache.truncate_sequence(1, 0),
+            Err("truncate_sequence: slot out of range")
+        );
+        assert_eq!(
+            cache.truncate_sequence(0, 0),
+            Err("truncate_sequence: empty slot")
+        );
+
+        cache.register_sequence(0).unwrap();
+        cache.ensure_capacity(0, BLOCK_SIZE * 3 + 1);
+        cache.advance_seq_len(0, BLOCK_SIZE * 3 + 1);
+        assert_eq!(cache.seq_len(0), BLOCK_SIZE * 3 + 1);
+        assert_eq!(cache.block_count(0), 4);
+        assert_eq!(cache.free_blocks(), 0);
+
+        cache.truncate_sequence(0, BLOCK_SIZE + 1).unwrap();
+        assert_eq!(cache.seq_len(0), BLOCK_SIZE + 1);
+        assert_eq!(cache.block_count(0), 2);
+        assert_eq!(cache.free_blocks(), 2);
+
+        cache.truncate_sequence(0, BLOCK_SIZE).unwrap();
+        assert_eq!(cache.seq_len(0), BLOCK_SIZE);
+        assert_eq!(cache.block_count(0), 1);
+        assert_eq!(cache.free_blocks(), 3);
+
+        cache.truncate_sequence(0, 0).unwrap();
+        assert_eq!(cache.seq_len(0), 0);
+        assert_eq!(cache.block_count(0), 1);
+        assert_eq!(cache.free_blocks(), 3);
+        assert_eq!(
+            cache.truncate_sequence(0, 1),
+            Err("truncate_sequence: cannot extend")
+        );
+    }
+}
+
 unsafe fn tensor_from_owned_buf(
     buf: GpuBuffer,
     shape: &[usize],

crates/xserv-model/src/qwen3.rs

@@ -701,45 +701,72 @@ impl Qwen3 {
         assert_eq!(seq_slots.len(), batch);
         assert!(batch > 0);
 
-        // TP: this rank owns a slice of the heads (local_* == full when world==1).
-        let num_heads = self.local_num_heads;
-        let num_kv_heads = self.local_num_kv_heads;
-        let head_dim = self.config.head_dim();
-        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+        self.decode_prepare(positions, seq_slots, paged_cache);
 
-        // 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 ids_gpu = upload_u32(tokens);
+        let positions_u32: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
+        let pos_gpu = upload_u32(&positions_u32);
+        let logits = self.decode_core(
+            ids_gpu.as_ptr() as *const std::ffi::c_void,
+            pos_gpu.as_ptr() as *const std::ffi::c_void,
+            batch,
+            seq_slots,
+            paged_cache,
+        );
+        logits
+    }
+
+    /// Host-side per-step cache bookkeeping: block allocation + uploading block
+    /// tables / context lens to their (stable-address) GPU buffers. Runs
+    /// OUTSIDE any CUDA-graph captured region.
+    pub fn decode_prepare(
+        &self,
+        positions: &[usize],
+        seq_slots: &[usize],
+        paged_cache: &mut PagedKVCache,
+    ) {
         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);
+    }
+
+    /// Pure-GPU decode step: embedding → all layers → final norm → logits.
+    /// Token ids and positions are read from device buffers; every other input
+    /// (weights, KV pools, block table, context lens) has a stable address —
+    /// which makes this region CUDA-graph capturable.
+    pub fn decode_core(
+        &self,
+        ids_gpu: *const std::ffi::c_void,
+        pos_gpu: *const std::ffi::c_void,
+        batch: usize,
+        seq_slots: &[usize],
+        paged_cache: &mut PagedKVCache,
+    ) -> Tensor {
+        let num_heads = self.local_num_heads;
+        let num_kv_heads = self.local_num_kv_heads;
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
 
         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();
 
-        // RoPE expects `[num_tokens, H, D]` with `num_tokens` positions —
-        // matches our `[B, H, D]` exactly, so we upload once here.
-        let positions_u32: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
-
-        // Batched embedding: [B, hidden]
-        let mut x = embedding(&self.embed_tokens, tokens);
+        let mut x = embedding_device_ids(&self.embed_tokens, ids_gpu, batch);
 
         for (layer_idx, layer) in self.layers.iter().enumerate() {
             let residual = x.clone();
             let normed = rmsnorm(&x, &layer.input_norm, eps);
 
             // Fused QKV projection: one GEMV instead of three.
-            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt); // [B, (H+2*KV)*D]
+            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
             let q_dim = num_heads * head_dim;
             let kv_dim = num_kv_heads * head_dim;
-            let q_all = qkv.narrow(1, 0, q_dim); // [B, H*D]  (view)
-            let k_all = qkv.narrow(1, q_dim, kv_dim); // [B, KV*D] (view)
+            let q_all = qkv.narrow(1, 0, q_dim);
+            let k_all = qkv.narrow(1, q_dim, kv_dim);
             let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);
 
-            // Per-head RMSNorm on contiguous copies (narrow views are strided).
             let q_flat = q_all.contiguous().reshape(&[batch * num_heads, head_dim]);
             let k_flat = k_all
                 .contiguous()
@@ -749,16 +776,13 @@ impl Qwen3 {
 
             let q_3d = q_normed.reshape(&[batch, num_heads, head_dim]);
             let k_3d = k_normed.reshape(&[batch, num_kv_heads, head_dim]);
-            rope_inplace(&q_3d, &self.rope_cache, &positions_u32);
-            rope_inplace(&k_3d, &self.rope_cache, &positions_u32);
+            rope_inplace_device_pos(&q_3d, &self.rope_cache, pos_gpu);
+            rope_inplace_device_pos(&k_3d, &self.rope_cache, pos_gpu);
 
             let v_3d = v_all.contiguous().reshape(&[batch, num_kv_heads, head_dim]);
 
-            // Single batched scatter for all sequences in the batch.
             paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, batch);
 
-            // Paged attention reads Q as [B, H, 1, D] — a contiguous view
-            // of [B, H, D].
             let q_4d = q_3d.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;
@@ -775,27 +799,24 @@ impl Qwen3 {
                 max_blocks,
             );
 
-            // attn_out shape [B, H, 1, D] is contiguous-equivalent to [B, H*D].
             let attn_merged = attn_out.reshape(&[batch, num_heads * head_dim]);
             let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
-            self.all_reduce(&attn_proj); // TP: sum partial attention outputs
+            self.all_reduce(&attn_proj);
 
             let (normed, x_new) =
                 xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
             let residual = x_new.clone();
 
-            // Fused gate+up projection: one GEMV instead of two.
-            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt); // [B, 2*ffn]
+            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
             let ffn_dim = gate_up.shape()[1] / 2;
             let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
             let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
             let hidden_states = xserv_kernels::silu_mul(&gate, &up);
             let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
-            self.all_reduce(&down); // TP: sum partial MLP outputs
+            self.all_reduce(&down);
             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);
         }
@@ -804,6 +825,111 @@ impl Qwen3 {
         matmul_2d(&x, &self.lm_head_t)
     }
 
+    /// Like `decode_core` but also captures hidden states at 3 specified layer
+    /// indices (after residual+MLP output). Used by EAGLE3 speculative drafting
+    /// to feed the draft head with low/mid/high target representations.
+    pub fn decode_core_with_hidden(
+        &self,
+        ids_gpu: *const std::ffi::c_void,
+        pos_gpu: *const std::ffi::c_void,
+        batch: usize,
+        seq_slots: &[usize],
+        paged_cache: &mut PagedKVCache,
+        hook_layers: &[usize; 3],
+    ) -> (Tensor, [Tensor; 3]) {
+        let num_heads = self.local_num_heads;
+        let num_kv_heads = self.local_num_kv_heads;
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        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();
+
+        let mut x = embedding_device_ids(&self.embed_tokens, ids_gpu, batch);
+        let mut hooks: [Option<Tensor>; 3] = [None, None, None];
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps);
+
+            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
+            let q_dim = num_heads * head_dim;
+            let kv_dim = num_kv_heads * head_dim;
+            let q_all = qkv.narrow(1, 0, q_dim);
+            let k_all = qkv.narrow(1, q_dim, kv_dim);
+            let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);
+
+            let q_flat = q_all.contiguous().reshape(&[batch * num_heads, head_dim]);
+            let k_flat = k_all
+                .contiguous()
+                .reshape(&[batch * num_kv_heads, head_dim]);
+            let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
+            let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);
+
+            let q_3d = q_normed.reshape(&[batch, num_heads, head_dim]);
+            let k_3d = k_normed.reshape(&[batch, num_kv_heads, head_dim]);
+            rope_inplace_device_pos(&q_3d, &self.rope_cache, pos_gpu);
+            rope_inplace_device_pos(&k_3d, &self.rope_cache, pos_gpu);
+
+            let v_3d = v_all.contiguous().reshape(&[batch, num_kv_heads, head_dim]);
+
+            paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, batch);
+
+            let q_4d = q_3d.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,
+            );
+
+            let attn_merged = attn_out.reshape(&[batch, num_heads * head_dim]);
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+            self.all_reduce(&attn_proj);
+
+            let (normed, x_new) =
+                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
+            let ffn_dim = gate_up.shape()[1] / 2;
+            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
+            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
+            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
+            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            self.all_reduce(&down);
+            x = add_any(&residual, &down);
+
+            for (h_idx, &h_layer) in hook_layers.iter().enumerate() {
+                if layer_idx == h_layer {
+                    hooks[h_idx] = Some(x.clone());
+                }
+            }
+        }
+
+        for &slot in seq_slots {
+            paged_cache.advance_seq_len(slot, 1);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        let logits = matmul_2d(&x, &self.lm_head_t);
+        let hidden_arr = [
+            hooks[0].take().expect("hook layer 0 not reached"),
+            hooks[1].take().expect("hook layer 1 not reached"),
+            hooks[2].take().expect("hook layer 2 not reached"),
+        ];
+        (logits, hidden_arr)
+    }
+
     /// 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].
@@ -884,6 +1010,358 @@ impl Qwen3 {
         matmul_2d(&x, &self.lm_head_t)
     }
 
+    /// Paged multi-token verify path: write `token_ids` into the paged cache,
+    /// then verify them with the same paged decode attention kernel used by
+    /// single-token decode. This keeps greedy top-1 behavior aligned with
+    /// `forward_decode_paged` while still batching the dense projections/MLP
+    /// across the draft window.
+    pub fn forward_verify_paged_decode_attention(
+        &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.local_num_heads;
+        let num_kv_heads = self.local_num_kv_heads;
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        paged_cache.ensure_capacity(slot, pos_offset + new_tokens);
+        paged_cache.advance_seq_len(slot, new_tokens);
+
+        let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
+            .map(|p| p as u32)
+            .collect();
+        let kv_lens: Vec<i32> = (0..new_tokens)
+            .map(|i| (pos_offset + i + 1) as i32)
+            .collect();
+        let slots = vec![slot; new_tokens];
+        paged_cache.sync_active_batch_with_lens(&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();
+
+        let mut x = embedding(&self.embed_tokens, token_ids);
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps);
+
+            let qkv = matmul_batched_gemv(&normed, &layer.qkv_proj_wt);
+            let q_dim = num_heads * head_dim;
+            let kv_dim = num_kv_heads * head_dim;
+            let q_all = qkv.narrow(1, 0, q_dim);
+            let k_all = qkv.narrow(1, q_dim, kv_dim);
+            let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);
+
+            let q_flat = q_all
+                .contiguous()
+                .reshape(&[new_tokens * num_heads, head_dim]);
+            let k_flat = k_all
+                .contiguous()
+                .reshape(&[new_tokens * num_kv_heads, head_dim]);
+            let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
+            let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);
+
+            let q_3d = q_normed.reshape(&[new_tokens, num_heads, head_dim]);
+            let k_3d = k_normed.reshape(&[new_tokens, num_kv_heads, head_dim]);
+            rope_inplace(&q_3d, &self.rope_cache, &positions);
+            rope_inplace(&k_3d, &self.rope_cache, &positions);
+
+            let v_3d = v_all
+                .contiguous()
+                .reshape(&[new_tokens, num_kv_heads, head_dim]);
+            paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, new_tokens);
+
+            let q_decode = q_3d.reshape(&[new_tokens, 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_decode,
+                k_pool_ptr,
+                v_pool_ptr,
+                bt_ptr,
+                cl_ptr,
+                new_tokens,
+                num_heads,
+                num_kv_heads,
+                head_dim,
+                max_blocks,
+            );
+
+            let attn_merged = attn_out.reshape(&[new_tokens, num_heads * head_dim]);
+            let attn_proj = matmul_batched_gemv(&attn_merged, &layer.o_proj_wt);
+            self.all_reduce(&attn_proj);
+
+            let (normed, x_new) =
+                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            let gate_up = matmul_batched_gemv(&normed, &layer.gate_up_proj_wt);
+            let ffn_dim = gate_up.shape()[1] / 2;
+            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
+            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
+            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
+            let down = matmul_batched_gemv(&hidden_states, &layer.down_proj_wt);
+            self.all_reduce(&down);
+            x = add_any(&residual, &down);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_batched_gemv(&x, &self.lm_head_t)
+    }
+
+    /// Like `forward_verify_paged_decode_attention`, but also captures hidden
+    /// states at 3 layer indices (per position). Returns
+    /// (logits [new_tokens, vocab], hooks [3][new_tokens, hidden]). Used by
+    /// EAGLE3 speculative γ≥2 verify path so we can seed the next round's
+    /// EAGLE draft with target's real hidden states at the accepted position.
+    pub fn forward_verify_paged_decode_attention_with_hidden(
+        &self,
+        token_ids: &[u32],
+        slot: usize,
+        paged_cache: &mut PagedKVCache,
+        hook_layers: &[usize; 3],
+    ) -> (Tensor, [Tensor; 3]) {
+        let new_tokens = token_ids.len();
+        let pos_offset = paged_cache.seq_len(slot);
+        let num_heads = self.local_num_heads;
+        let num_kv_heads = self.local_num_kv_heads;
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        paged_cache.ensure_capacity(slot, pos_offset + new_tokens);
+        paged_cache.advance_seq_len(slot, new_tokens);
+
+        let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
+            .map(|p| p as u32)
+            .collect();
+        let kv_lens: Vec<i32> = (0..new_tokens)
+            .map(|i| (pos_offset + i + 1) as i32)
+            .collect();
+        let slots = vec![slot; new_tokens];
+        paged_cache.sync_active_batch_with_lens(&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();
+
+        let mut x = embedding(&self.embed_tokens, token_ids);
+        let mut hooks: [Option<Tensor>; 3] = [None, None, None];
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps);
+
+            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
+            let q_dim = num_heads * head_dim;
+            let kv_dim = num_kv_heads * head_dim;
+            let q_all = qkv.narrow(1, 0, q_dim);
+            let k_all = qkv.narrow(1, q_dim, kv_dim);
+            let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);
+
+            let q_flat = q_all
+                .contiguous()
+                .reshape(&[new_tokens * num_heads, head_dim]);
+            let k_flat = k_all
+                .contiguous()
+                .reshape(&[new_tokens * num_kv_heads, head_dim]);
+            let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
+            let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);
+
+            let q_3d = q_normed.reshape(&[new_tokens, num_heads, head_dim]);
+            let k_3d = k_normed.reshape(&[new_tokens, num_kv_heads, head_dim]);
+            rope_inplace(&q_3d, &self.rope_cache, &positions);
+            rope_inplace(&k_3d, &self.rope_cache, &positions);
+
+            let v_3d = v_all
+                .contiguous()
+                .reshape(&[new_tokens, num_kv_heads, head_dim]);
+            paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, new_tokens);
+
+            let q_decode = q_3d.reshape(&[new_tokens, 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_decode,
+                k_pool_ptr,
+                v_pool_ptr,
+                bt_ptr,
+                cl_ptr,
+                new_tokens,
+                num_heads,
+                num_kv_heads,
+                head_dim,
+                max_blocks,
+            );
+
+            let attn_merged = attn_out.reshape(&[new_tokens, num_heads * head_dim]);
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+            self.all_reduce(&attn_proj);
+
+            let (normed, x_new) =
+                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
+            let ffn_dim = gate_up.shape()[1] / 2;
+            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
+            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
+            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
+            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            self.all_reduce(&down);
+            x = add_any(&residual, &down);
+
+            for (h_idx, &h_layer) in hook_layers.iter().enumerate() {
+                if layer_idx == h_layer {
+                    hooks[h_idx] = Some(x.clone());
+                }
+            }
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        let logits = matmul_2d(&x, &self.lm_head_t);
+        let hidden_arr = [
+            hooks[0].take().expect("hook layer 0 not reached"),
+            hooks[1].take().expect("hook layer 1 not reached"),
+            hooks[2].take().expect("hook layer 2 not reached"),
+        ];
+        (logits, hidden_arr)
+    }
+
+    /// Tree-aware verify: like `_with_hidden` but supports sibling candidates
+    /// sharing the same target position. Caller supplies per-token positions
+    /// (for RoPE), kv_lens (attention context length), and a flattened
+    /// `tree_mask` (`[new_tokens, new_tokens]` i32; `mask[i, j]!=0` iff query i
+    /// attends to newly-written K/V at slot j). Positions in the paged cache
+    /// before pos_offset are always attended (regular history).
+    #[allow(clippy::too_many_arguments)]
+    pub fn forward_verify_paged_decode_attention_tree_with_hidden(
+        &self,
+        token_ids: &[u32],
+        positions: &[u32],
+        kv_lens: &[i32],
+        tree_mask: &[i32],
+        slot: usize,
+        paged_cache: &mut PagedKVCache,
+        hook_layers: &[usize; 3],
+    ) -> (Tensor, [Tensor; 3]) {
+        let new_tokens = token_ids.len();
+        assert_eq!(positions.len(), new_tokens);
+        assert_eq!(kv_lens.len(), new_tokens);
+        assert_eq!(tree_mask.len(), new_tokens * new_tokens);
+
+        let pos_offset = paged_cache.seq_len(slot);
+        let num_heads = self.local_num_heads;
+        let num_kv_heads = self.local_num_kv_heads;
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        paged_cache.ensure_capacity(slot, pos_offset + new_tokens);
+        paged_cache.advance_seq_len(slot, new_tokens);
+
+        let slots = vec![slot; new_tokens];
+        paged_cache.sync_active_batch_with_lens(&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();
+
+        // Upload tree_mask [new_tokens, new_tokens] i32 to GPU.
+        let mask_bytes: &[u8] = unsafe {
+            std::slice::from_raw_parts(tree_mask.as_ptr() as *const u8, tree_mask.len() * 4)
+        };
+        let mut mask_buf =
+            xserv_cuda::allocator::cached_alloc(mask_bytes.len()).expect("alloc tree_mask");
+        mask_buf.copy_from_host(mask_bytes).unwrap();
+        let mask_ptr = mask_buf.as_ptr() as *const i32;
+
+        let mut x = embedding(&self.embed_tokens, token_ids);
+        let mut hooks: [Option<Tensor>; 3] = [None, None, None];
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps);
+
+            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
+            let q_dim = num_heads * head_dim;
+            let kv_dim = num_kv_heads * head_dim;
+            let q_all = qkv.narrow(1, 0, q_dim);
+            let k_all = qkv.narrow(1, q_dim, kv_dim);
+            let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);
+
+            let q_flat = q_all
+                .contiguous()
+                .reshape(&[new_tokens * num_heads, head_dim]);
+            let k_flat = k_all
+                .contiguous()
+                .reshape(&[new_tokens * num_kv_heads, head_dim]);
+            let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
+            let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);
+
+            let q_3d = q_normed.reshape(&[new_tokens, num_heads, head_dim]);
+            let k_3d = k_normed.reshape(&[new_tokens, num_kv_heads, head_dim]);
+            rope_inplace(&q_3d, &self.rope_cache, positions);
+            rope_inplace(&k_3d, &self.rope_cache, positions);
+
+            let v_3d = v_all
+                .contiguous()
+                .reshape(&[new_tokens, num_kv_heads, head_dim]);
+            paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, new_tokens);
+
+            let q_decode = q_3d.reshape(&[new_tokens, 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_tree(
+                &q_decode,
+                k_pool_ptr,
+                v_pool_ptr,
+                bt_ptr,
+                cl_ptr,
+                mask_ptr,
+                new_tokens,
+                num_heads,
+                num_kv_heads,
+                head_dim,
+                max_blocks,
+                pos_offset,
+                new_tokens,
+            );
+
+            let attn_merged = attn_out.reshape(&[new_tokens, num_heads * head_dim]);
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+            self.all_reduce(&attn_proj);
+
+            let (normed, x_new) =
+                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
+            let ffn_dim = gate_up.shape()[1] / 2;
+            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
+            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
+            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
+            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            self.all_reduce(&down);
+            x = add_any(&residual, &down);
+
+            for (h_idx, &h_layer) in hook_layers.iter().enumerate() {
+                if layer_idx == h_layer {
+                    hooks[h_idx] = Some(x.clone());
+                }
+            }
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        let logits = matmul_2d(&x, &self.lm_head_t);
+        let hidden_arr = [
+            hooks[0].take().expect("hook layer 0 not reached"),
+            hooks[1].take().expect("hook layer 1 not reached"),
+            hooks[2].take().expect("hook layer 2 not reached"),
+        ];
+        (logits, hidden_arr)
+    }
+
     /// 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();
@@ -950,6 +1428,12 @@ impl Qwen3 {
         matmul_2d(&x, &self.lm_head_t)
     }
 
+    /// Reference to the target's token embedding table. Shared (not copied)
+    /// with speculative draft heads like EAGLE3.
+    pub fn embed_tokens_tensor(&self) -> &Tensor {
+        &self.embed_tokens
+    }
+
     /// Extract weight pointers for CUDA Graph capture.
     pub fn layer_weight_ptrs(&self) -> Vec<crate::decode_graph::LayerWeightPtrs> {
         self.layers
@@ -1158,6 +1642,14 @@ fn row_view(t: &Tensor, row: usize) -> Tensor {
     )
 }
 
+/// Upload a u32 slice to a pooled GPU buffer (synchronous H2D).
+fn upload_u32(vals: &[u32]) -> xserv_cuda::GpuBuffer {
+    let bytes = unsafe { std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4) };
+    let mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).expect("alloc u32 upload");
+    buf.copy_from_host(bytes).unwrap();
+    buf
+}
+
 /// Concatenate row tensors [1, cols] into a single [B, cols] tensor via D2D memcpy.
 fn concat_rows(rows: &[Tensor]) -> Tensor {
     assert!(!rows.is_empty());

crates/xserv-model/src/qwen3_graph.rs

@@ -0,0 +1,185 @@
+//! CUDA-graph replay for Qwen3 batch=1 decode (Phase 24 / speculative draft).
+//!
+//! Same pattern as `gpt_oss_graph.rs`, but for the Qwen3 dense decode path used
+//! by speculative decoding's draft model. A Qwen3-0.6B decode step is ~140
+//! kernel launches; wrapping the whole step into one `cudaGraphLaunch` cuts
+//! the ~4× γ draft cost per speculative round.
+//!
+//! See `gpt_oss_graph.rs` for the design commentary; the capture preconditions,
+//! retained-warmup mechanism, and quarantine lifetime are all identical here.
+
+use std::ffi::c_void;
+
+use xserv_cuda::allocator::{self, RetainedBlocks};
+use xserv_cuda::{CudaGraph, CudaStream, GpuBuffer};
+use xserv_tensor::Tensor;
+
+use crate::paged_kv_cache::PagedKVCache;
+use crate::qwen3::Qwen3;
+
+pub struct Qwen3DecodeGraph {
+    stream: CudaStream,
+    graph: CudaGraph,
+    ids_buf: GpuBuffer, // [1] u32, persistent graph input
+    pos_buf: GpuBuffer, // [1] u32, persistent graph input
+    logits: Tensor,     // graph output; rewritten in place by every replay
+    _arena: RetainedBlocks,
+}
+
+impl Qwen3DecodeGraph {
+    /// Capture one batch=1 decode step and replay it once.
+    pub fn capture(
+        model: &Qwen3,
+        token: u32,
+        position: usize,
+        slot: usize,
+        cache: &mut PagedKVCache,
+    ) -> Self {
+        let stream = CudaStream::new().expect("create capture stream");
+        let mut ids_buf = allocator::cached_alloc(4).expect("alloc ids buf");
+        let mut pos_buf = allocator::cached_alloc(4).expect("alloc pos buf");
+
+        model.decode_prepare(&[position], &[slot], cache);
+        ids_buf.copy_from_host(&token.to_le_bytes()).unwrap();
+        pos_buf
+            .copy_from_host(&(position as u32).to_le_bytes())
+            .unwrap();
+
+        // Retained warmup: run the exact step once eagerly with the quarantine
+        // ON to stock the pool. See gpt_oss_graph.rs:66-86 for the full
+        // rationale. Re-running the step is idempotent: the KV scatter
+        // overwrites the same cache position and advance_seq_len is *inside*
+        // decode_core, so we roll it back afterwards.
+        let seq_len_before = cache.seq_len(slot);
+        allocator::begin_retain();
+        {
+            let _guard = xserv_cuda::push_stream(&stream);
+            let _ = model.decode_core(
+                ids_buf.as_ptr() as *const c_void,
+                pos_buf.as_ptr() as *const c_void,
+                1,
+                &[slot],
+                cache,
+            );
+        }
+        drop(allocator::end_retain());
+        stream.synchronize().expect("warmup sync");
+        // decode_core advanced seq_len; roll back so capture starts from the
+        // same logical state as the eager warmup.
+        cache
+            .truncate_sequence(slot, seq_len_before)
+            .expect("rollback after warmup");
+
+        allocator::begin_retain();
+        let mut graph = CudaGraph::new();
+        let logits;
+        {
+            let _guard = xserv_cuda::stream::push_stream(&stream);
+            graph
+                .begin_capture(&stream)
+                .expect("begin decode-graph capture");
+            logits = model.decode_core(
+                ids_buf.as_ptr() as *const c_void,
+                pos_buf.as_ptr() as *const c_void,
+                1,
+                &[slot],
+                cache,
+            );
+            graph
+                .end_capture(&stream)
+                .expect("end decode-graph capture");
+        }
+        let arena = allocator::end_retain();
+
+        // The capture path called advance_seq_len (host-side) but the actual
+        // GPU compute has not yet run. Roll back and let the first replay
+        // advance it exactly once with real K/V writes.
+        cache
+            .truncate_sequence(slot, seq_len_before)
+            .expect("rollback after capture");
+
+        graph.launch(&stream).expect("first decode-graph replay");
+        cache.advance_seq_len(slot, 1);
+
+        Self {
+            stream,
+            graph,
+            ids_buf,
+            pos_buf,
+            logits,
+            _arena: arena,
+        }
+    }
+
+    /// Run one decode step by replaying the captured graph.
+    pub fn step(
+        &mut self,
+        model: &Qwen3,
+        token: u32,
+        position: usize,
+        slot: usize,
+        cache: &mut PagedKVCache,
+    ) -> Tensor {
+        model.decode_prepare(&[position], &[slot], cache);
+        self.ids_buf.copy_from_host(&token.to_le_bytes()).unwrap();
+        self.pos_buf
+            .copy_from_host(&(position as u32).to_le_bytes())
+            .unwrap();
+        self.graph
+            .launch(&self.stream)
+            .expect("decode-graph replay");
+        cache.advance_seq_len(slot, 1);
+        self.logits.clone()
+    }
+}
+
+/// Lazy capture policy: first decode step of the process runs eager, the
+/// second is captured, the rest replay. Batch>1 always falls back to eager.
+/// Disable with `XSERV_DECODE_GRAPH=0`.
+pub struct GraphedQwen3Decoder {
+    graph: Option<Qwen3DecodeGraph>,
+    eager_steps: u32,
+    enabled: bool,
+}
+
+impl GraphedQwen3Decoder {
+    pub fn new() -> Self {
+        let enabled = std::env::var("XSERV_DECODE_GRAPH")
+            .map(|v| v != "0")
+            .unwrap_or(true);
+        Self {
+            graph: None,
+            eager_steps: 0,
+            enabled,
+        }
+    }
+
+    pub fn decode(
+        &mut self,
+        model: &Qwen3,
+        tokens: &[u32],
+        positions: &[usize],
+        slots: &[usize],
+        cache: &mut PagedKVCache,
+    ) -> Tensor {
+        if self.enabled && tokens.len() == 1 {
+            if let Some(g) = self.graph.as_mut() {
+                return g.step(model, tokens[0], positions[0], slots[0], cache);
+            }
+            if self.eager_steps >= 1 {
+                let g = Qwen3DecodeGraph::capture(model, tokens[0], positions[0], slots[0], cache);
+                let logits = g.logits.clone();
+                self.graph = Some(g);
+                return logits;
+            }
+        }
+        self.eager_steps += 1;
+        model.forward_decode_paged(tokens, positions, slots, cache)
+    }
+}
+
+impl Default for GraphedQwen3Decoder {
+    fn default() -> Self {
+        Self::new()
+    }
+}

crates/xserv-model/src/sampling.rs

@@ -1,5 +1,6 @@
 use half::bf16;
 use rand::Rng;
+use std::collections::HashSet;
 use xserv_tensor::{DType, Device, Tensor};
 
 #[derive(Clone)]
@@ -7,6 +8,8 @@ pub struct SamplingParams {
     pub temperature: f32,
     pub top_k: usize,
     pub top_p: f32,
+    pub presence_penalty: f32,
+    pub repetition_penalty: f32,
 }
 
 impl Default for SamplingParams {
@@ -15,6 +18,8 @@ impl Default for SamplingParams {
             temperature: 0.0,
             top_k: 0,
             top_p: 1.0,
+            presence_penalty: 0.0,
+            repetition_penalty: 1.0,
         }
     }
 }
@@ -22,12 +27,21 @@ impl Default for SamplingParams {
 /// Sample a token from logits with shape [seq_len, vocab_size].
 /// Uses the last position's logits. Handles both F32 and BF16 dtypes.
 pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
+    sample_with_history(logits, params, &[])
+}
+
+/// Sample while applying penalties to tokens already present in the request.
+/// Presence penalty follows the OpenAI convention (subtract once per distinct
+/// token); repetition penalty follows the HF convention.
+pub fn sample_with_history(logits: &Tensor, params: &SamplingParams, history: &[u32]) -> u32 {
     assert_eq!(logits.ndim(), 2);
     // Greedy fast path: GPU argmax + 4-byte D2H instead of copying the whole
     // [seq, vocab] logits to the host and scanning it (~201k bf16/token).
     // NaN logits lose every `>` comparison in the kernel, matching the
     // NaN-safe host argmax below.
     if params.temperature == 0.0
+        && params.presence_penalty == 0.0
+        && params.repetition_penalty == 1.0
         && logits.dtype() == DType::BF16
         && matches!(logits.device(), Device::Cuda(_))
         && logits.is_contiguous()
@@ -40,7 +54,7 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
     let logits_cpu = logits.to_device(Device::Cpu);
 
     // Extract last row as f32
-    let last_row: Vec<f32> = match logits.dtype() {
+    let mut last_row: Vec<f32> = match logits.dtype() {
         DType::F32 => {
             let data = logits_cpu.as_slice::<f32>();
             data[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec()
@@ -55,7 +69,39 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
         _ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
     };
 
-    // Greedy
+    // NaN-safe: sampling path uses partial_cmp().unwrap() in top-k/top-p
+    // sorts and softmax; a single NaN logit would panic the engine thread.
+    // Replace NaN with -inf (equivalent to masking) instead.
+    let mut nan_seen = false;
+    for v in last_row.iter_mut() {
+        if v.is_nan() {
+            nan_seen = true;
+            *v = f32::NEG_INFINITY;
+        }
+    }
+    if nan_seen {
+        eprintln!("[sampling] WARNING: NaN logits encountered in sample()");
+    }
+
+    if !history.is_empty() && (params.presence_penalty != 0.0 || params.repetition_penalty != 1.0) {
+        let seen: HashSet<u32> = history.iter().copied().collect();
+        for id in seen {
+            let i = id as usize;
+            if i >= last_row.len() {
+                continue;
+            }
+            if params.repetition_penalty != 1.0 {
+                let value = last_row[i];
+                last_row[i] = if value > 0.0 {
+                    value / params.repetition_penalty
+                } else {
+                    value * params.repetition_penalty
+                };
+            }
+            last_row[i] -= params.presence_penalty;
+        }
+    }
+
     if params.temperature == 0.0 {
         return argmax(&last_row);
     }
@@ -181,3 +227,25 @@ fn argmax(data: &[f32]) -> u32 {
     }
     best_i as u32
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn greedy_sampling_applies_history_penalties() {
+        let logits = Tensor::from_slice(&[10.0f32, 9.0, 0.0], &[1, 3]);
+
+        let presence = SamplingParams {
+            presence_penalty: 2.0,
+            ..SamplingParams::default()
+        };
+        assert_eq!(sample_with_history(&logits, &presence, &[0]), 1);
+
+        let repetition = SamplingParams {
+            repetition_penalty: 2.0,
+            ..SamplingParams::default()
+        };
+        assert_eq!(sample_with_history(&logits, &repetition, &[0]), 1);
+    }
+}

crates/xserv-server/src/api.rs

@@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize};
 use std::convert::Infallible;
 use std::path::Path;
 use std::sync::Arc;
+use std::sync::atomic::Ordering;
 use tokio_stream::StreamExt;
 use tokio_stream::wrappers::ReceiverStream;
 use uuid::Uuid;
@@ -25,11 +26,35 @@ pub struct ChatRequest {
     #[serde(default)]
     pub stream: Option<bool>,
     #[serde(default)]
+    pub stream_options: Option<StreamOptions>,
+    #[serde(default)]
     pub temperature: Option<f32>,
     #[serde(default)]
     pub top_k: Option<usize>,
     #[serde(default)]
     pub top_p: Option<f32>,
+    #[serde(default)]
+    pub presence_penalty: Option<f32>,
+    #[serde(default)]
+    pub repetition_penalty: Option<f32>,
+    #[serde(default)]
+    pub chat_template_kwargs: Option<ChatTemplateKwargs>,
+    #[serde(default)]
+    pub enable_thinking: Option<bool>,
+}
+
+#[derive(Deserialize)]
+pub struct StreamOptions {
+    #[serde(default)]
+    pub include_usage: bool,
+}
+
+#[derive(Deserialize, Default)]
+pub struct ChatTemplateKwargs {
+    #[serde(default)]
+    pub enable_thinking: Option<bool>,
+    #[serde(default)]
+    pub preserve_thinking: Option<bool>,
 }
 
 #[derive(Deserialize, Serialize, Clone)]
@@ -102,12 +127,17 @@ impl ChatTemplate {
         }
     }
 
-    pub fn render(&self, messages: &[Message]) -> String {
+    pub fn render(
+        &self,
+        messages: &[Message],
+        enable_thinking: Option<bool>,
+        preserve_thinking: Option<bool>,
+    ) -> String {
         if self.source.is_empty() {
             return build_prompt_hardcoded(messages, &self.model_type);
         }
 
-        match self.render_jinja(messages) {
+        match self.render_jinja(messages, enable_thinking, preserve_thinking) {
             Ok(prompt) => prompt,
             Err(e) => {
                 eprintln!("[chat-template] Jinja render error: {e}, falling back to hardcoded");
@@ -116,7 +146,12 @@ impl ChatTemplate {
         }
     }
 
-    fn render_jinja(&self, messages: &[Message]) -> Result<String, minijinja::Error> {
+    fn render_jinja(
+        &self,
+        messages: &[Message],
+        enable_thinking: Option<bool>,
+        preserve_thinking: Option<bool>,
+    ) -> Result<String, minijinja::Error> {
         let mut env = minijinja::Environment::new();
 
         // Register custom functions the template may call.
@@ -127,6 +162,42 @@ impl ChatTemplate {
         env.add_filter("startswith", |s: String, prefix: String| -> bool {
             s.starts_with(&prefix)
         });
+        env.set_unknown_method_callback(|_, value, method, args| {
+            use minijinja::value::from_args;
+            use minijinja::{Error, ErrorKind, Value};
+
+            let Some(value) = value.as_str() else {
+                return Err(Error::from(ErrorKind::UnknownMethod));
+            };
+            match method {
+                "startswith" => {
+                    let (prefix,): (String,) = from_args(args)?;
+                    Ok(Value::from(value.starts_with(&prefix)))
+                }
+                "endswith" => {
+                    let (suffix,): (String,) = from_args(args)?;
+                    Ok(Value::from(value.ends_with(&suffix)))
+                }
+                "split" => {
+                    let (separator,): (String,) = from_args(args)?;
+                    Ok(Value::from_serialize(
+                        value
+                            .split(&separator)
+                            .map(str::to_owned)
+                            .collect::<Vec<_>>(),
+                    ))
+                }
+                "rstrip" => {
+                    let (chars,): (String,) = from_args(args)?;
+                    Ok(Value::from(value.trim_end_matches(|c| chars.contains(c))))
+                }
+                "lstrip" => {
+                    let (chars,): (String,) = from_args(args)?;
+                    Ok(Value::from(value.trim_start_matches(|c| chars.contains(c))))
+                }
+                _ => Err(Error::from(ErrorKind::UnknownMethod)),
+            }
+        });
 
         env.add_template("chat", &self.source)?;
         let tmpl = env.get_template("chat")?;
@@ -136,6 +207,8 @@ impl ChatTemplate {
             add_generation_prompt => true,
             bos_token => "",
             eos_token => "",
+            enable_thinking => enable_thinking,
+            preserve_thinking => preserve_thinking,
         };
 
         tmpl.render(ctx)
@@ -256,8 +329,12 @@ fn build_prompt_gpt_oss(messages: &[Message]) -> String {
 // HTTP handlers
 // ---------------------------------------------------------------------------
 
-pub async fn health() -> &'static str {
-    "ok"
+pub async fn health(Extension(state): Extension<Arc<AppState>>) -> Response {
+    if state.engine_ready.load(Ordering::Acquire) {
+        (StatusCode::OK, "ok").into_response()
+    } else {
+        (StatusCode::SERVICE_UNAVAILABLE, "loading").into_response()
+    }
 }
 
 pub async fn list_models(Extension(state): Extension<Arc<AppState>>) -> Json<ModelsResponse> {
@@ -291,7 +368,10 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
         return response;
     }
 
-    let prompt = state.chat_template.render(&req.messages);
+    let (enable_thinking, preserve_thinking) = template_options(&req);
+    let prompt = state
+        .chat_template
+        .render(&req.messages, enable_thinking, preserve_thinking);
     let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
     let prompt_token_count = prompt_tokens.len();
 
@@ -331,6 +411,10 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
         }
     }
 
+    let fr_value = match normalize_finish_reason(&finish_reason) {
+        Some(s) => serde_json::Value::String(s.to_string()),
+        None => serde_json::Value::Null,
+    };
     Json(serde_json::json!({
         "id": id,
         "object": "chat.completion",
@@ -339,7 +423,7 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
         "choices": [{
             "index": 0,
             "message": { "role": "assistant", "content": content },
-            "finish_reason": finish_reason,
+            "finish_reason": fr_value,
         }],
         "usage": {
             "prompt_tokens": prompt_token_count,
@@ -359,18 +443,26 @@ fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
         return response;
     }
 
-    let prompt = state.chat_template.render(&req.messages);
+    let (enable_thinking, preserve_thinking) = template_options(&req);
+    let prompt = state
+        .chat_template
+        .render(&req.messages, enable_thinking, preserve_thinking);
     let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
+    let prompt_token_count = prompt_tokens.len();
+    let include_usage = req
+        .stream_options
+        .as_ref()
+        .is_some_and(|options| options.include_usage);
 
     let max_seq_len = state.max_seq_len;
-    if prompt_tokens.len() >= max_seq_len {
+    if prompt_token_count >= max_seq_len {
         return bad_request(format!(
             "prompt is {} tokens, exceeds max_seq_len {}",
-            prompt_tokens.len(),
+            prompt_token_count,
             max_seq_len
         ));
     }
-    let max_tokens = req.max_tokens.min(max_seq_len - prompt_tokens.len());
+    let max_tokens = req.max_tokens.min(max_seq_len - prompt_token_count);
 
     let (engine_tx, engine_rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
     let gen_req = GenerateRequest {
@@ -389,10 +481,12 @@ fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
     tokio::spawn(async move {
         let mut engine_stream = ReceiverStream::new(engine_rx);
         let mut first = true;
+        let mut completion_token_count = 0usize;
 
         while let Some(event) = engine_stream.next().await {
             match event {
                 GenerateEvent::Token { text, .. } => {
+                    completion_token_count += 1;
                     if first {
                         // First chunk: role announcement
                         let chunk =
@@ -412,9 +506,22 @@ fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
                             make_chunk(&id, &model_name, created, None, Some("assistant"), None);
                         let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
                     }
-                    let chunk =
-                        make_chunk(&id, &model_name, created, None, None, Some(&finish_reason));
+                    // Only "stop" and "length" are OpenAI-standard values. Internal
+                    // codes like "error" (client-stalled from tp/pp engine) map to
+                    // null so SDK clients see a clean stream close.
+                    let fr = normalize_finish_reason(&finish_reason);
+                    let chunk = make_chunk(&id, &model_name, created, None, None, fr);
                     let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
+                    if include_usage {
+                        let usage = make_usage_chunk(
+                            &id,
+                            &model_name,
+                            created,
+                            prompt_token_count,
+                            completion_token_count,
+                        );
+                        let _ = sse_tx.send(Ok(Event::default().data(usage))).await;
+                    }
                     let _ = sse_tx
                         .send(Ok(Event::default().data("[DONE]".to_string())))
                         .await;
@@ -442,6 +549,34 @@ fn validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
         return Some(bad_request("max_tokens must be greater than 0"));
     }
 
+    if let Some(t) = req.temperature {
+        if !t.is_finite() || t < 0.0 {
+            return Some(bad_request("temperature must be a finite value >= 0"));
+        }
+    }
+    if let Some(p) = req.top_p {
+        if !p.is_finite() || !(0.0..=1.0).contains(&p) {
+            return Some(bad_request("top_p must be in [0, 1]"));
+        }
+    }
+    if let Some(k) = req.top_k {
+        if k > 1_000_000 {
+            return Some(bad_request("top_k must be <= 1_000_000"));
+        }
+    }
+    if let Some(p) = req.presence_penalty {
+        if !p.is_finite() || !(-2.0..=2.0).contains(&p) {
+            return Some(bad_request("presence_penalty must be in [-2, 2]"));
+        }
+    }
+    if let Some(p) = req.repetition_penalty {
+        if !p.is_finite() || p <= 0.0 {
+            return Some(bad_request(
+                "repetition_penalty must be a finite value greater than 0",
+            ));
+        }
+    }
+
     None
 }
 
@@ -453,9 +588,14 @@ fn submit_to_engine(state: &AppState, req: GenerateRequest) -> Result<(), Respon
         .engine_sender
         .lock()
         .unwrap_or_else(|e| e.into_inner());
-    sender
-        .send(req)
-        .map_err(|_| service_unavailable("inference engine is not available"))
+    sender.try_send(req).map_err(|err| match err {
+        std::sync::mpsc::TrySendError::Full(_) => {
+            service_unavailable("inference engine is busy, retry later")
+        }
+        std::sync::mpsc::TrySendError::Disconnected(_) => {
+            service_unavailable("inference engine is not available")
+        }
+    })
 }
 
 fn service_unavailable(message: impl Into<String>) -> Response {
@@ -518,6 +658,28 @@ fn make_chunk(
     .to_string()
 }
 
+fn make_usage_chunk(
+    id: &str,
+    model: &str,
+    created: u64,
+    prompt_tokens: usize,
+    completion_tokens: usize,
+) -> String {
+    serde_json::json!({
+        "id": id,
+        "object": "chat.completion.chunk",
+        "created": created,
+        "model": model,
+        "choices": [],
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens,
+        }
+    })
+    .to_string()
+}
+
 fn unix_timestamp() -> u64 {
     std::time::SystemTime::now()
         .duration_since(std::time::UNIX_EPOCH)
@@ -530,5 +692,27 @@ fn sampling_params(req: &ChatRequest) -> SamplingParams {
         temperature: req.temperature.unwrap_or(0.0),
         top_k: req.top_k.unwrap_or(0),
         top_p: req.top_p.unwrap_or(1.0),
+        presence_penalty: req.presence_penalty.unwrap_or(0.0),
+        repetition_penalty: req.repetition_penalty.unwrap_or(1.0),
+    }
+}
+
+fn template_options(req: &ChatRequest) -> (Option<bool>, Option<bool>) {
+    let kwargs = req.chat_template_kwargs.as_ref();
+    (
+        req.enable_thinking
+            .or_else(|| kwargs.and_then(|value| value.enable_thinking)),
+        kwargs.and_then(|value| value.preserve_thinking),
+    )
+}
+
+/// Map engine finish_reason strings to OpenAI-standard values. Any engine-internal
+/// code (e.g. "error" from tp/pp client-stall) collapses to None so SDK clients see
+/// a clean null instead of an unknown value.
+fn normalize_finish_reason(fr: &str) -> Option<&'static str> {
+    match fr {
+        "stop" => Some("stop"),
+        "length" => Some("length"),
+        _ => None,
     }
 }

crates/xserv-server/src/engine.rs

@@ -4,7 +4,9 @@ use std::sync::Once;
 use std::sync::mpsc;
 use std::time::Instant;
 use xserv_model::loader;
-use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample};
+use xserv_model::{
+    BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample_with_history,
+};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
@@ -38,6 +40,9 @@ struct Sequence {
     seq_slot: Option<usize>,
     sender: tokio::sync::mpsc::Sender<GenerateEvent>,
     prefilled: bool,
+    /// Set when a `try_send` failed (client too slow or gone). The scheduler
+    /// reaps the sequence next iteration instead of blocking the decode thread.
+    client_stalled: bool,
     eos_token_id: Option<u32>,
     decode_buffer: Vec<u8>,
     created_at: Instant,
@@ -229,7 +234,7 @@ impl Engine {
                         slot,
                         &mut self.paged_cache,
                     );
-                    let next = sample(&logits, &seq.sampling);
+                    let next = sample_with_history(&logits, &seq.sampling, &seq.prompt_tokens);
                     seq.generated_tokens.push(next);
                     seq.prefilled = true;
                     emit_token(&self.tokenizer, seq, next);
@@ -307,7 +312,11 @@ impl Engine {
                 // (~1.2 MB for B=4, Qwen3 vocab=152K).
                 let all_greedy = decode_indices
                     .iter()
-                    .all(|&i| running[i].sampling.temperature == 0.0);
+                    .all(|&i| {
+                        running[i].sampling.temperature == 0.0
+                            && running[i].sampling.presence_penalty == 0.0
+                            && running[i].sampling.repetition_penalty == 1.0
+                    });
                 if all_greedy {
                     let next_ids = xserv_kernels::argmax_bf16_to_host(&logits);
                     for (j, &i) in decode_indices.iter().enumerate() {
@@ -326,7 +335,10 @@ impl Engine {
                     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 {
+                        let unpenalized_greedy = running[i].sampling.temperature == 0.0
+                            && running[i].sampling.presence_penalty == 0.0
+                            && running[i].sampling.repetition_penalty == 1.0;
+                        let next = if unpenalized_greedy {
                             row_logits
                                 .iter()
                                 .enumerate()
@@ -336,7 +348,9 @@ impl Engine {
                         } else {
                             let row_tensor =
                                 xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]);
-                            sample(&row_tensor, &running[i].sampling)
+                            let mut history = running[i].prompt_tokens.clone();
+                            history.extend_from_slice(&running[i].generated_tokens);
+                            sample_with_history(&row_tensor, &running[i].sampling, &history)
                         };
                         running[i].generated_tokens.push(next);
                         emit_token(&self.tokenizer, &mut running[i], next);
@@ -370,6 +384,7 @@ impl Engine {
             seq_slot: None,
             sender: req.sender,
             prefilled: false,
+            client_stalled: false,
             eos_token_id: self.tokenizer.eos_token_id(),
             decode_buffer: Vec::new(),
             created_at: Instant::now(),
@@ -392,9 +407,12 @@ 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(),
-        });
+        try_send_event(
+            seq,
+            GenerateEvent::Done {
+                finish_reason: "stop".to_string(),
+            },
+        );
         return;
     }
 
@@ -403,22 +421,45 @@ fn emit_token(tokenizer: &Tokenizer, seq: &mut Sequence, token_id: u32) {
         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(),
-        });
+        try_send_event(
+            seq,
+            GenerateEvent::Done {
+                finish_reason: "length".to_string(),
+            },
+        );
     } else {
         send_token_if_nonempty(seq, text);
     }
 }
 
-fn send_token_if_nonempty(seq: &Sequence, text: String) {
+fn send_token_if_nonempty(seq: &mut 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 });
+        try_send_event(seq, GenerateEvent::Token { id, text });
+    }
+}
+
+/// Send an event without blocking the shared decode thread. If the client is
+/// too slow (channel full) or gone (closed), flag the sequence for eviction
+/// instead of blocking — one slow consumer must never stall the whole
+/// continuous-batching loop. When the sequence is reaped its `sender` drops,
+/// closing the channel so the client's receive loop ends rather than hanging.
+fn try_send_event(seq: &mut Sequence, event: GenerateEvent) {
+    if let Err(err) = seq.sender.try_send(event) {
+        seq.client_stalled = true;
+        if let tokio::sync::mpsc::error::TrySendError::Full(_) = err {
+            eprintln!(
+                "[scheduler] seq {}: client too slow (stream channel full), evicting",
+                seq.id
+            );
+        }
     }
 }
 
 fn is_finished(seq: &Sequence) -> bool {
+    if seq.client_stalled {
+        return true;
+    }
     if seq.generated_tokens.is_empty() {
         return false;
     }

crates/xserv-server/src/main.rs

@@ -5,19 +5,22 @@ mod tp_engine;
 
 use axum::{
     Extension, Router,
+    extract::DefaultBodyLimit,
     routing::{get, post},
 };
 use engine::GenerateRequest;
 use std::path::PathBuf;
+use std::sync::atomic::AtomicBool;
 use std::sync::{Arc, Mutex, mpsc};
-use xserv_model::ModelConfig;
+use xserv_model::{ModelConfig, ModelFamily};
 
 pub struct AppState {
     pub model_name: String,
     pub chat_template: api::ChatTemplate,
-    pub engine_sender: Mutex<mpsc::Sender<GenerateRequest>>,
+    pub engine_sender: Mutex<mpsc::SyncSender<GenerateRequest>>,
     pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
     pub max_seq_len: usize,
+    pub engine_ready: Arc<AtomicBool>,
 }
 
 #[tokio::main]
@@ -76,12 +79,18 @@ async fn main() {
         std::process::exit(1);
     }
     let model_config = ModelConfig::from_file(&model_dir.join("config.json"));
-    // gpt-oss is only implemented in the TP engine; route it there even at
-    // tp=1 (single-rank world) so quantized models can serve on one GPU.
-    let is_gpt_oss = model_config.model_type.as_deref() == Some("gpt_oss");
-    if pp > 1 && is_gpt_oss {
+    let model_family = model_config.family();
+    if model_family == ModelFamily::Qwen35Moe && pp <= 1 {
         eprintln!(
-            "gpt-oss is not supported by the pipeline-parallel engine (Qwen3 only); use --tp instead"
+            "Qwen3.5/3.6 MoE model '{}' currently requires pipeline parallelism; use --pp 8 on dash5",
+            model_config.model_type_str()
+        );
+        std::process::exit(1);
+    }
+    if pp > 1 && !matches!(model_family, ModelFamily::Qwen3 | ModelFamily::Qwen35Moe) {
+        eprintln!(
+            "{} is not supported by the pipeline-parallel engine; use --tp instead",
+            model_family.as_str()
         );
         std::process::exit(1);
     }
@@ -104,29 +113,46 @@ async fn main() {
 
     let tokenizer = xserv_tokenizer::Tokenizer::from_file(&model_dir.join("tokenizer.json"));
 
-    // Unbounded channel: allows multiple requests to queue up
-    let (tx, rx) = mpsc::channel::<GenerateRequest>();
+    // Bounded channel to backpressure incoming requests when the engine falls
+    // behind, instead of letting them pile up in RAM. try_send in the API
+    // handler surfaces this as 503 to the client.
+    let (tx, rx) = mpsc::sync_channel::<GenerateRequest>(256);
+    let engine_ready = Arc::new(AtomicBool::new(false));
 
     let model_dir_clone = model_dir.clone();
+    let engine_ready_clone = Arc::clone(&engine_ready);
     std::thread::spawn(move || {
         if pp > 1 {
             // Pipeline-parallel path: stage-0 coordinator + worker stage threads.
-            pp_engine::run_pp(&model_dir_clone, pp, max_seq_len, rx);
-        } else if tp <= 1 && !is_gpt_oss {
+            pp_engine::run_pp(
+                &model_dir_clone,
+                pp,
+                max_seq_len,
+                rx,
+                engine_ready_clone,
+            );
+        } else if tp <= 1 && model_family == ModelFamily::Qwen3 {
             let mut engine = engine::Engine::load_with_swap(
                 &model_dir_clone,
                 max_batch,
                 max_seq_len,
                 swap_space_gb,
             );
+            engine_ready_clone.store(true, std::sync::atomic::Ordering::Release);
             engine.run(rx);
         } else {
             // Tensor-parallel path: rank-0 coordinator + worker rank threads.
-            tp_engine::run_tp(&model_dir_clone, tp, max_seq_len, rx);
+            tp_engine::run_tp(
+                &model_dir_clone,
+                tp,
+                max_seq_len,
+                rx,
+                engine_ready_clone,
+            );
         }
     });
 
-    let model_type = model_config.model_type.clone().unwrap_or_default();
+    let model_type = model_config.model_type_str().to_string();
     let chat_template = api::ChatTemplate::load(&model_dir, &model_type);
     let state = Arc::new(AppState {
         model_name,
@@ -134,12 +160,14 @@ async fn main() {
         engine_sender: Mutex::new(tx),
         engine_tokenizer: Mutex::new(tokenizer),
         max_seq_len,
+        engine_ready,
     });
 
     let app = Router::new()
         .route("/health", get(api::health))
         .route("/v1/models", get(api::list_models))
         .route("/v1/chat/completions", post(api::chat_completions))
+        .layer(DefaultBodyLimit::max(4 * 1024 * 1024))
         .layer(Extension(state));
 
     let addr = format!("0.0.0.0:{port}");

crates/xserv-server/src/pp_engine.rs

@@ -16,14 +16,18 @@
 use std::ffi::c_void;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::mpsc;
 use std::thread;
 
 use half::bf16;
 use xserv_distributed::{PpContext, UniqueId};
 use xserv_model::loader;
-use xserv_model::sampling::SamplingParams;
-use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, sample};
+use xserv_model::sampling::{SamplingParams, sample_with_history};
+use xserv_model::{
+    BLOCK_SIZE, ModelConfig, ModelFamily, PagedKVCache, Qwen3, Qwen35Moe,
+    Qwen35RecurrentCache,
+};
 use xserv_tensor::{DType, Device, Tensor};
 use xserv_tokenizer::Tokenizer;
 
@@ -39,7 +43,7 @@ enum PpCommand {
     /// Receive `[n_tokens, hidden]` from the previous stage, run this stage's
     /// layers; if last stage, sample with `sampling` and return the token.
     Prefill {
-        n_tokens: usize,
+        tokens: Vec<u32>,
         slot: usize,
         sampling: SamplingParams,
     },
@@ -52,13 +56,69 @@ enum PpCommand {
 }
 
 struct StageCtx {
-    model: Qwen3,
+    model: PpModel,
+    qwen35_recurrent: Option<Qwen35RecurrentCache>,
     cache: PagedKVCache,
     pp: Arc<PpContext>,
     hidden: usize,
     device: u32,
 }
 
+enum PpModel {
+    Qwen3(Qwen3),
+    Qwen35(Qwen35Moe),
+}
+
+impl StageCtx {
+    fn reset_slot(&mut self, slot: usize) {
+        if let Some(cache) = &mut self.qwen35_recurrent {
+            cache.reset_slot(slot);
+        }
+    }
+
+    fn embed(&self, tokens: &[u32]) -> Tensor {
+        match &self.model {
+            PpModel::Qwen3(model) => model.embed(tokens),
+            PpModel::Qwen35(model) => model.embed(tokens),
+        }
+    }
+
+    fn head(&self, x: &Tensor) -> Tensor {
+        match &self.model {
+            PpModel::Qwen3(model) => model.head(x),
+            PpModel::Qwen35(model) => model.head(x),
+        }
+    }
+
+    fn forward_prefill(&mut self, x: Tensor, slot: usize) -> Tensor {
+        match &self.model {
+            PpModel::Qwen3(model) => model.forward_layers_prefill(x, slot, &mut self.cache),
+            PpModel::Qwen35(model) => model.forward_layers_prefill(
+                x,
+                slot,
+                &mut self.cache,
+                self.qwen35_recurrent
+                    .as_mut()
+                    .expect("Qwen3.6 recurrent cache"),
+            ),
+        }
+    }
+
+    fn forward_decode(&mut self, x: Tensor, slot: usize) -> Tensor {
+        match &self.model {
+            PpModel::Qwen3(model) => model.forward_layers_decode(x, &[slot], &mut self.cache),
+            PpModel::Qwen35(model) => model.forward_layers_decode(
+                x,
+                slot,
+                &mut self.cache,
+                self.qwen35_recurrent
+                    .as_mut()
+                    .expect("Qwen3.6 recurrent cache"),
+            ),
+        }
+    }
+}
+
 /// Build this stage: NCCL init, load + slice weights, size a per-stage KV pool
 /// for THIS stage's layers only (so per-GPU KV is ~1/P).
 fn build_stage(
@@ -71,14 +131,40 @@ fn build_stage(
     id: UniqueId,
 ) -> StageCtx {
     let pp = Arc::new(PpContext::init(stage, world, id, device));
-    let weights = loader::load_model_dir(model_dir, Device::Cpu);
-    let model = Qwen3::from_weights_pp(config.clone(), weights, stage, world, device);
+    let model = match config.family() {
+        ModelFamily::Qwen35Moe => {
+            let weights = loader::load_model_dir_filtered(model_dir, Device::Cpu, |name| {
+                Qwen35Moe::weight_belongs_to_pp_stage(config, name, stage, world)
+            });
+            PpModel::Qwen35(Qwen35Moe::from_weights_pp(
+                config.clone(),
+                weights,
+                stage,
+                world,
+                device,
+            ))
+        }
+        _ => {
+            let weights = loader::load_model_dir(model_dir, Device::Cpu);
+            PpModel::Qwen3(Qwen3::from_weights_pp(
+                config.clone(),
+                weights,
+                stage,
+                world,
+                device,
+            ))
+        }
+    };
 
     // The KV cache only needs this stage's layers; build it from a config clone
     // whose layer count is the per-stage count (heads are NOT split under PP).
     let per_stage = config.num_layers() / world;
     let mut stage_config = config.clone();
-    stage_config.num_hidden_layers = Some(per_stage);
+    if let Some(text) = stage_config.text_config.as_mut() {
+        text.num_hidden_layers = Some(per_stage);
+    } else {
+        stage_config.num_hidden_layers = Some(per_stage);
+    }
 
     let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE);
     let total_blocks = max_blocks_per_seq + 8; // v1 serial: one active sequence
@@ -91,8 +177,13 @@ fn build_stage(
         DType::BF16,
         device,
     );
+    let qwen35_recurrent = match &model {
+        PpModel::Qwen35(model) => Some(model.new_recurrent_cache(4)),
+        PpModel::Qwen3(_) => None,
+    };
     StageCtx {
         model,
+        qwen35_recurrent,
         cache,
         pp,
         hidden: config.hidden(),
@@ -138,40 +229,51 @@ fn worker_loop(
         max_seq_len,
         id,
     );
+    let _ = ack_tx.send(());
     let is_last = stage == world - 1;
+    let mut token_history = vec![Vec::<u32>::new(); 4];
     let prev = stage - 1;
     let next = stage + 1;
 
     while let Ok(cmd) = cmd_rx.recv() {
         match cmd {
             PpCommand::Register(slot) => {
+                token_history[slot].clear();
+                sc.reset_slot(slot);
                 let _ = sc.cache.register_sequence(slot);
                 let _ = ack_tx.send(());
             }
             PpCommand::Free(slot) => {
+                token_history[slot].clear();
                 sc.cache.free_sequence(slot);
                 let _ = ack_tx.send(());
             }
             PpCommand::Prefill {
-                n_tokens,
+                tokens,
                 slot,
                 sampling,
             } => {
+                let n_tokens = tokens.len();
+                token_history[slot] = tokens;
                 let x = recv_hidden(&sc, n_tokens, prev);
-                let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache);
+                let x = sc.forward_prefill(x, slot);
                 if is_last {
-                    let logits = sc.model.head(&x);
-                    let _ = token_tx.send(sample(&logits, &sampling));
+                    let logits = sc.head(&x);
+                    let token = sample_with_history(&logits, &sampling, &token_history[slot]);
+                    token_history[slot].push(token);
+                    let _ = token_tx.send(token);
                 } else {
                     send_hidden(&sc, &x, next);
                 }
             }
             PpCommand::Decode { slot, sampling } => {
                 let x = recv_hidden(&sc, 1, prev);
-                let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache);
+                let x = sc.forward_decode(x, slot);
                 if is_last {
-                    let logits = sc.model.head(&x);
-                    let _ = token_tx.send(sample(&logits, &sampling));
+                    let logits = sc.head(&x);
+                    let token = sample_with_history(&logits, &sampling, &token_history[slot]);
+                    token_history[slot].push(token);
+                    let _ = token_tx.send(token);
                 } else {
                     send_hidden(&sc, &x, next);
                 }
@@ -191,6 +293,7 @@ pub fn run_pp(
     world: usize,
     max_seq_len: usize,
     rx: mpsc::Receiver<GenerateRequest>,
+    ready: Arc<AtomicBool>,
 ) {
     assert!(world >= 2, "run_pp requires world >= 2");
     let config = ModelConfig::from_file(&model_dir.join("config.json"));
@@ -231,6 +334,10 @@ pub fn run_pp(
 
     // Stage 0 (this thread): coordinator + embedding + first layers.
     let mut sc = build_stage(model_dir, &config, 0, world, 0, max_seq_len, id);
+    for _ in 1..world {
+        ack_rx.recv().expect("PP worker exited during model load");
+    }
+    ready.store(true, Ordering::Release);
     eprintln!("[pp-engine] ready (pp={world}, max_seq_len={max_seq_len})");
 
     let n_workers = world - 1;
@@ -249,6 +356,7 @@ pub fn run_pp(
     let slot = 0usize;
     while let Ok(req) = rx.recv() {
         broadcast(&cmd_txs, PpCommand::Register(slot));
+        sc.reset_slot(slot);
         sc.cache.register_sequence(slot).expect("register slot");
         wait_acks(&ack_rx);
 
@@ -256,21 +364,24 @@ pub fn run_pp(
         broadcast(
             &cmd_txs,
             PpCommand::Prefill {
-                n_tokens: req.prompt_tokens.len(),
+                tokens: req.prompt_tokens.clone(),
                 slot,
                 sampling: req.sampling.clone(),
             },
         );
-        let x = sc.model.embed(&req.prompt_tokens);
-        let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache);
+        let x = sc.embed(&req.prompt_tokens);
+        let x = sc.forward_prefill(x, slot);
         send_hidden(&sc, &x, next_peer);
         let mut next = token_rx.recv().expect("prefill token");
 
         let mut decode_buf: Vec<u8> = Vec::new();
         let mut generated = 1usize;
-        emit_text(&tokenizer, &req, next, &mut decode_buf);
+        let mut stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
 
         let finish = loop {
+            if stalled {
+                break "error";
+            }
             if tokenizer.is_eos(next) {
                 break "stop";
             }
@@ -284,22 +395,22 @@ pub fn run_pp(
                     sampling: req.sampling.clone(),
                 },
             );
-            let x = sc.model.embed(&[next]);
-            let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache);
+            let x = sc.embed(&[next]);
+            let x = sc.forward_decode(x, slot);
             send_hidden(&sc, &x, next_peer);
             next = token_rx.recv().expect("decode token");
             generated += 1;
-            emit_text(&tokenizer, &req, next, &mut decode_buf);
+            stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
         };
 
         let tail = tokenizer.flush_decode_stream(&mut decode_buf);
         if !tail.is_empty() {
-            let _ = req.sender.blocking_send(GenerateEvent::Token {
+            let _ = req.sender.try_send(GenerateEvent::Token {
                 id: next,
                 text: tail,
             });
         }
-        let _ = req.sender.blocking_send(GenerateEvent::Done {
+        let _ = req.sender.try_send(GenerateEvent::Done {
             finish_reason: finish.to_string(),
         });
 
@@ -312,14 +423,24 @@ pub fn run_pp(
 }
 
 /// Stream a token's decoded text to the client (EOS contributes no text).
-fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &mut Vec<u8>) {
+/// Returns false if the send would block (client too slow) or the client is
+/// gone — the caller stops generating so the coordinator thread is free to
+/// admit the next request instead of blocking on one slow consumer.
+fn emit_text(
+    tokenizer: &Tokenizer,
+    req: &GenerateRequest,
+    token_id: u32,
+    buf: &mut Vec<u8>,
+) -> bool {
     if tokenizer.is_eos(token_id) {
-        return;
+        return true;
     }
     let text = tokenizer.decode_token_stream(token_id, buf);
     if !text.is_empty() {
-        let _ = req
+        return req
             .sender
-            .blocking_send(GenerateEvent::Token { id: token_id, text });
+            .try_send(GenerateEvent::Token { id: token_id, text })
+            .is_ok();
     }
+    true
 }

crates/xserv-server/src/tp_engine.rs

@@ -14,14 +14,15 @@
 
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::mpsc;
 use std::thread;
 
 use xserv_distributed::{TpContext, UniqueId};
 use xserv_model::loader;
 use xserv_model::{
-    BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, sample,
-    sample_greedy_penalized,
+    BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, ModelFamily, PagedKVCache, Qwen3,
+    sample_greedy_penalized, sample_with_history,
 };
 use xserv_tensor::{DType, Device, Tensor};
 use xserv_tokenizer::Tokenizer;
@@ -105,24 +106,26 @@ fn build_rank(
     tp: Option<Arc<TpContext>>,
 ) -> RankCtx {
     let weights = loader::load_model_dir(model_dir, Device::Cpu);
-    let model = if config.is_moe() {
-        TpModel::GptOss(GptOss::from_weights_tp(
+    let model = match config.family() {
+        ModelFamily::GptOss => TpModel::GptOss(GptOss::from_weights_tp(
             config.clone(),
             weights,
             rank,
             world,
             device,
             tp,
-        ))
-    } else {
-        TpModel::Qwen3(Qwen3::from_weights_tp(
+        )),
+        ModelFamily::Qwen35Moe => {
+            panic!("Qwen3.5/3.6 MoE reached TP engine before inference support was implemented")
+        }
+        _ => TpModel::Qwen3(Qwen3::from_weights_tp(
             config.clone(),
             weights,
             rank,
             world,
             device,
             tp,
-        ))
+        )),
     };
     let local_kv = config.num_kv_heads() / world;
     let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
@@ -164,6 +167,7 @@ fn worker_loop(
         max_seq_len,
         Some(tp),
     );
+    let _ = ack_tx.send(());
     while let Ok(cmd) = cmd_rx.recv() {
         match cmd {
             TpCommand::Register(slot) => {
@@ -196,6 +200,7 @@ pub fn run_tp(
     world: usize,
     max_seq_len: usize,
     rx: mpsc::Receiver<GenerateRequest>,
+    ready: Arc<AtomicBool>,
 ) {
     // world=1 is a valid single-rank configuration (gpt-oss has no
     // single-GPU engine path; NCCL init and all_reduce no-op at world=1).
@@ -235,6 +240,10 @@ pub fn run_tp(
     // Rank 0 (this thread).
     let tp = Arc::new(TpContext::init(0, world, id, 0));
     let mut rc = build_rank(model_dir, &config, 0, world, 0, max_seq_len, Some(tp));
+    for _ in 1..world {
+        ack_rx.recv().expect("TP worker exited during model load");
+    }
+    ready.store(true, Ordering::Release);
     eprintln!("[tp-engine] ready (tp={world}, max_seq_len={max_seq_len})");
 
     // Optional repetition penalty to break greedy repetition loops (reasoning
@@ -254,7 +263,7 @@ pub fn run_tp(
             let start = history.len().saturating_sub(rep_window);
             sample_greedy_penalized(logits, &history[start..], rep_penalty)
         } else {
-            sample(logits, sp)
+            sample_with_history(logits, sp, history)
         }
     };
 
@@ -288,15 +297,18 @@ pub fn run_tp(
             .model
             .forward_prefill_paged(&req.prompt_tokens, slot, &mut rc.cache);
         wait_acks(&ack_rx);
-        let mut gen_ids: Vec<u32> = Vec::new();
-        let mut next = pick(&logits, &req.sampling, &gen_ids);
-        gen_ids.push(next);
+        let mut token_history = req.prompt_tokens.clone();
+        let mut next = pick(&logits, &req.sampling, &token_history);
+        token_history.push(next);
 
         let mut decode_buf: Vec<u8> = Vec::new();
         let mut generated = 1usize;
-        emit_text(&tokenizer, &req, next, &mut decode_buf);
+        let mut stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
 
         let finish = loop {
+            if stalled {
+                break "error";
+            }
             if tokenizer.is_eos(next) {
                 break "stop";
             }
@@ -314,20 +326,20 @@ pub fn run_tp(
             );
             let logits = rank_decode(&mut rc, &[next], &[pos], &[slot]);
             wait_acks(&ack_rx);
-            next = pick(&logits, &req.sampling, &gen_ids);
-            gen_ids.push(next);
+            next = pick(&logits, &req.sampling, &token_history);
+            token_history.push(next);
             generated += 1;
-            emit_text(&tokenizer, &req, next, &mut decode_buf);
+            stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
         };
 
         let tail = tokenizer.flush_decode_stream(&mut decode_buf);
         if !tail.is_empty() {
-            let _ = req.sender.blocking_send(GenerateEvent::Token {
+            let _ = req.sender.try_send(GenerateEvent::Token {
                 id: next,
                 text: tail,
             });
         }
-        let _ = req.sender.blocking_send(GenerateEvent::Done {
+        let _ = req.sender.try_send(GenerateEvent::Done {
             finish_reason: finish.to_string(),
         });
 
@@ -340,14 +352,24 @@ pub fn run_tp(
 }
 
 /// Stream a token's decoded text to the client (EOS contributes no text).
-fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &mut Vec<u8>) {
+/// Returns false if the send would block (client too slow) or the client is
+/// gone — the caller stops generating so the serial coordinator thread is free
+/// to admit the next request instead of blocking on one slow consumer.
+fn emit_text(
+    tokenizer: &Tokenizer,
+    req: &GenerateRequest,
+    token_id: u32,
+    buf: &mut Vec<u8>,
+) -> bool {
     if tokenizer.is_eos(token_id) {
-        return;
+        return true;
     }
     let text = tokenizer.decode_token_stream(token_id, buf);
     if !text.is_empty() {
-        let _ = req
+        return req
             .sender
-            .blocking_send(GenerateEvent::Token { id: token_id, text });
+            .try_send(GenerateEvent::Token { id: token_id, text })
+            .is_ok();
     }
+    true
 }

csrc/activation/activations.cu

@@ -36,6 +36,35 @@ __global__ void silu_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
     if (idx < n) out[idx] = __float2bfloat16(silu_f(__bfloat162float(x[idx])));
 }
 
+__global__ void sigmoid_f32(const float* x, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = 1.0f / (1.0f + expf(-x[idx]));
+}
+
+__global__ void sigmoid_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        float v = __bfloat162float(x[idx]);
+        out[idx] = __float2bfloat16(1.0f / (1.0f + expf(-v)));
+    }
+}
+
+__global__ void softplus_f32(const float* x, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        float v = x[idx];
+        out[idx] = log1pf(expf(-fabsf(v))) + fmaxf(v, 0.0f);
+    }
+}
+
+__global__ void softplus_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        float v = __bfloat162float(x[idx]);
+        out[idx] = __float2bfloat16(log1pf(expf(-fabsf(v))) + fmaxf(v, 0.0f));
+    }
+}
+
 __global__ void scale_f32_kernel(const float* x, float* out, float scale, int n) {
     int idx = blockIdx.x * blockDim.x + threadIdx.x;
     if (idx < n) out[idx] = x[idx] * scale;
@@ -108,6 +137,21 @@ __global__ void mul_bf16_kernel(const __nv_bfloat16* a, const __nv_bfloat16* b,
     if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(a[idx]) * __bfloat162float(b[idx]));
 }
 
+__global__ void row_scale_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ scale,
+    __nv_bfloat16* __restrict__ out,
+    int rows,
+    int cols
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = rows * cols;
+    if (idx >= total) return;
+    int row = idx / cols;
+    float v = __bfloat162float(x[idx]) * __bfloat162float(scale[row]);
+    out[idx] = __float2bfloat16(v);
+}
+
 extern "C" {
 
 void launch_gelu_f32(const void* x, void* out, int n, void* stream) {
@@ -140,6 +184,36 @@ void launch_silu_bf16(const void* x, void* out, int n, void* stream) {
     CUDA_CHECK_LAST_ERROR();
 }
 
+void launch_sigmoid_f32(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    sigmoid_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_sigmoid_bf16(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    sigmoid_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_softplus_f32(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    softplus_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_softplus_bf16(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    softplus_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+    CUDA_CHECK_LAST_ERROR();
+}
+
 void launch_scale_f32(const void* x, void* out, float scale, int n, void* stream) {
     int block = 256;
     int grid = (n + block - 1) / block;
@@ -193,6 +267,15 @@ void launch_mul_bf16(const void* a, const void* b, void* out, int n, void* strea
     CUDA_CHECK_LAST_ERROR();
 }
 
+void launch_row_scale_bf16(const void* x, const void* scale, void* out, int rows, int cols, void* stream) {
+    int n = rows * cols;
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    row_scale_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)scale, (__nv_bfloat16*)out, rows, cols);
+    CUDA_CHECK_LAST_ERROR();
+}
+
 void launch_silu_mul_bf16(const void* gate, const void* up, void* out, int n, void* stream) {
     int block = 256;
     int grid = (n + block - 1) / block;

csrc/attention/causal_mask.cu

@@ -15,7 +15,10 @@ __global__ void causal_mask_f32(
     int col = blockIdx.x * blockDim.x + threadIdx.x;
 
     if (col < cols && col > row + offset) {
-        scores[batch_idx * rows * cols + row * cols + col] = -INFINITY;
+        // 64-bit index: batch * rows * cols overflows int32 at moderate batch
+        // and long context (e.g. batch=128 * heads=28 * seq=32768).
+        long long idx = ((long long)batch_idx * rows + row) * cols + col;
+        scores[idx] = -INFINITY;
     }
 }
 
@@ -28,7 +31,8 @@ __global__ void causal_mask_bf16(
     int col = blockIdx.x * blockDim.x + threadIdx.x;
 
     if (col < cols && col > row + offset) {
-        scores[batch_idx * rows * cols + row * cols + col] = __float2bfloat16(-INFINITY);
+        long long idx = ((long long)batch_idx * rows + row) * cols + col;
+        scores[idx] = __float2bfloat16(-INFINITY);
     }
 }
 

csrc/attention/deltanet.cu

@@ -0,0 +1,349 @@
+#include <cuda_bf16.h>
+#include <math.h>
+#include "../common.cuh"
+
+__global__ void depthwise_causal_conv1d_silu_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ w,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len,
+    int channels,
+    int kernel
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = seq_len * channels;
+    if (idx >= total) return;
+
+    int c = idx % channels;
+    int t = idx / channels;
+    float acc = 0.0f;
+    for (int k = 0; k < kernel; ++k) {
+        int src_t = t - (kernel - 1 - k);
+        if (src_t >= 0) {
+            float xv = __bfloat162float(x[src_t * channels + c]);
+            float wv = __bfloat162float(w[c * kernel + k]);
+            acc += xv * wv;
+        }
+    }
+    float y = acc / (1.0f + expf(-acc));
+    out[idx] = __float2bfloat16(y);
+}
+
+__global__ void deltanet_ar_bf16_kernel(
+    const __nv_bfloat16* __restrict__ q,
+    const __nv_bfloat16* __restrict__ k,
+    const __nv_bfloat16* __restrict__ v,
+    const __nv_bfloat16* __restrict__ beta,
+    const __nv_bfloat16* __restrict__ alpha_softplus,
+    const __nv_bfloat16* __restrict__ a_log,
+    __nv_bfloat16* __restrict__ state,
+    __nv_bfloat16* __restrict__ out,
+    int key_heads,
+    int value_heads,
+    int head_dim
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = value_heads * head_dim;
+    if (idx >= total) return;
+
+    int h = idx / head_dim;
+    int j = idx % head_dim;
+    int key_h = h * key_heads / value_heads;
+
+    float beta_h = __bfloat162float(beta[h]);
+    float gate = expf(__bfloat162float(alpha_softplus[h]) * __bfloat162float(a_log[h]));
+
+    float sk = 0.0f;
+    for (int i = 0; i < head_dim; ++i) {
+        float s = __bfloat162float(state[(h * head_dim + i) * head_dim + j]);
+        float ki = __bfloat162float(k[key_h * head_dim + i]);
+        s *= gate;
+        state[(h * head_dim + i) * head_dim + j] = __float2bfloat16(s);
+        sk += s * ki;
+    }
+
+    float vj = __bfloat162float(v[h * head_dim + j]);
+    float d = (vj - sk) * beta_h;
+
+    float out_j = 0.0f;
+    for (int i = 0; i < head_dim; ++i) {
+        float ki = __bfloat162float(k[key_h * head_dim + i]);
+        float qi = __bfloat162float(q[key_h * head_dim + i]) / sqrtf((float)head_dim);
+        int sidx = (h * head_dim + i) * head_dim + j;
+        float s = __bfloat162float(state[sidx]) + ki * d;
+        state[sidx] = __float2bfloat16(s);
+        out_j += s * qi;
+    }
+    out[h * head_dim + j] = __float2bfloat16(out_j);
+}
+
+// Stateful depthwise causal convolution for hybrid recurrent attention.
+// `state` stores the preceding `kernel - 1` inputs in chronological order.
+// One thread owns a channel, so the state update is race-free for any seq_len.
+__global__ void depthwise_causal_conv1d_stateful_silu_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ w,
+    __nv_bfloat16* __restrict__ state,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len,
+    int channels,
+    int kernel
+) {
+    int c = blockIdx.x * blockDim.x + threadIdx.x;
+    if (c >= channels) return;
+
+    const int history = kernel - 1;
+    for (int t = 0; t < seq_len; ++t) {
+        float acc = 0.0f;
+        for (int k = 0; k < kernel; ++k) {
+            int src_t = t - history + k;
+            float xv;
+            if (src_t < 0) {
+                xv = __bfloat162float(state[c * history + history + src_t]);
+            } else {
+                xv = __bfloat162float(x[(long long)src_t * channels + c]);
+            }
+            acc += xv * __bfloat162float(w[c * kernel + k]);
+        }
+        out[(long long)t * channels + c] =
+            __float2bfloat16(acc / (1.0f + expf(-acc)));
+    }
+
+    // Keep the last `history` raw projection values for the next call.
+    for (int h = 0; h < history; ++h) {
+        int src_t = seq_len - history + h;
+        __nv_bfloat16 value;
+        if (src_t < 0) {
+            value = state[c * history + history + src_t];
+        } else {
+            value = x[(long long)src_t * channels + c];
+        }
+        state[c * history + h] = value;
+    }
+}
+
+// L2-normalize each row using FP32 accumulation. Qwen3.5/3.6 applies this to
+// convolved Q and K before the delta rule (this is not RMSNorm).
+__global__ void l2_normalize_rows_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    __nv_bfloat16* __restrict__ out,
+    int rows,
+    int cols,
+    float eps
+) {
+    int row = blockIdx.x;
+    if (row >= rows) return;
+
+    float local = 0.0f;
+    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
+        float v = __bfloat162float(x[(long long)row * cols + i]);
+        local += v * v;
+    }
+    __shared__ float sums[256];
+    sums[threadIdx.x] = local;
+    __syncthreads();
+    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (threadIdx.x < stride) sums[threadIdx.x] += sums[threadIdx.x + stride];
+        __syncthreads();
+    }
+    // Match torch.nn.functional.normalize / llama.cpp: clamp the norm by eps,
+    // rather than adding eps to every non-zero norm.
+    float inv = rsqrtf(fmaxf(sums[0], eps * eps));
+    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
+        out[(long long)row * cols + i] =
+            __float2bfloat16(__bfloat162float(x[(long long)row * cols + i]) * inv);
+    }
+}
+
+// Recurrent Gated DeltaNet over one or more tokens. HF stores V heads grouped
+// by K head, so V head h uses K/Q head floor(h / (value_heads/key_heads)).
+// The recurrent matrix is FP32 as required by `mamba_ssm_dtype=float32`.
+__global__ void deltanet_recurrent_bf16_f32_kernel(
+    const __nv_bfloat16* __restrict__ q,
+    const __nv_bfloat16* __restrict__ k,
+    const __nv_bfloat16* __restrict__ v,
+    const __nv_bfloat16* __restrict__ beta,
+    const __nv_bfloat16* __restrict__ alpha_softplus,
+    const __nv_bfloat16* __restrict__ a_log,
+    float* __restrict__ state,
+    __nv_bfloat16* __restrict__ out,
+    int seq_len,
+    int key_heads,
+    int value_heads,
+    int head_dim
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = value_heads * head_dim;
+    if (idx >= total) return;
+
+    int h = idx / head_dim;
+    int j = idx % head_dim;
+    int key_h = h * key_heads / value_heads;
+    float a = __bfloat162float(a_log[h]);
+    float neg_a = -expf(a);
+    float q_scale = rsqrtf((float)head_dim);
+
+    for (int t = 0; t < seq_len; ++t) {
+        const __nv_bfloat16* q_t = q + ((long long)t * key_heads + key_h) * head_dim;
+        const __nv_bfloat16* k_t = k + ((long long)t * key_heads + key_h) * head_dim;
+        const __nv_bfloat16* v_t = v + ((long long)t * value_heads + h) * head_dim;
+        float b = __bfloat162float(beta[(long long)t * value_heads + h]);
+        float alpha = __bfloat162float(alpha_softplus[(long long)t * value_heads + h]);
+        float decay = expf(neg_a * alpha);
+
+        float sk = 0.0f;
+        for (int i = 0; i < head_dim; ++i) {
+            long long sidx = ((long long)h * head_dim + i) * head_dim + j;
+            float s = state[sidx] * decay;
+            state[sidx] = s;
+            sk += s * __bfloat162float(k_t[i]);
+        }
+
+        float d = (__bfloat162float(v_t[j]) - sk) * b;
+        float out_j = 0.0f;
+        for (int i = 0; i < head_dim; ++i) {
+            long long sidx = ((long long)h * head_dim + i) * head_dim + j;
+            float s = state[sidx] + __bfloat162float(k_t[i]) * d;
+            state[sidx] = s;
+            out_j += s * (__bfloat162float(q_t[i]) * q_scale);
+        }
+        out[((long long)t * value_heads + h) * head_dim + j] = __float2bfloat16(out_j);
+    }
+}
+
+extern "C" {
+
+void launch_depthwise_causal_conv1d_silu_bf16(
+    const void* x,
+    const void* w,
+    void* out,
+    int seq_len,
+    int channels,
+    int kernel,
+    void* stream
+) {
+    int total = seq_len * channels;
+    int block = 256;
+    int grid = (total + block - 1) / block;
+    depthwise_causal_conv1d_silu_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x,
+        (const __nv_bfloat16*)w,
+        (__nv_bfloat16*)out,
+        seq_len,
+        channels,
+        kernel
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_deltanet_ar_bf16(
+    const void* q,
+    const void* k,
+    const void* v,
+    const void* beta,
+    const void* alpha_softplus,
+    const void* a_log,
+    void* state,
+    void* out,
+    int key_heads,
+    int value_heads,
+    int head_dim,
+    void* stream
+) {
+    int total = value_heads * head_dim;
+    int block = 128;
+    int grid = (total + block - 1) / block;
+    deltanet_ar_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)q,
+        (const __nv_bfloat16*)k,
+        (const __nv_bfloat16*)v,
+        (const __nv_bfloat16*)beta,
+        (const __nv_bfloat16*)alpha_softplus,
+        (const __nv_bfloat16*)a_log,
+        (__nv_bfloat16*)state,
+        (__nv_bfloat16*)out,
+        key_heads,
+        value_heads,
+        head_dim
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_depthwise_causal_conv1d_stateful_silu_bf16(
+    const void* x,
+    const void* w,
+    void* state,
+    void* out,
+    int seq_len,
+    int channels,
+    int kernel,
+    void* stream
+) {
+    int block = 256;
+    int grid = (channels + block - 1) / block;
+    depthwise_causal_conv1d_stateful_silu_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x,
+        (const __nv_bfloat16*)w,
+        (__nv_bfloat16*)state,
+        (__nv_bfloat16*)out,
+        seq_len,
+        channels,
+        kernel
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_l2_normalize_rows_bf16(
+    const void* x,
+    void* out,
+    int rows,
+    int cols,
+    float eps,
+    void* stream
+) {
+    l2_normalize_rows_bf16_kernel<<<rows, 256, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x,
+        (__nv_bfloat16*)out,
+        rows,
+        cols,
+        eps
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_deltanet_recurrent_bf16_f32(
+    const void* q,
+    const void* k,
+    const void* v,
+    const void* beta,
+    const void* alpha_softplus,
+    const void* a_log,
+    void* state,
+    void* out,
+    int seq_len,
+    int key_heads,
+    int value_heads,
+    int head_dim,
+    void* stream
+) {
+    int total = value_heads * head_dim;
+    int block = 128;
+    int grid = (total + block - 1) / block;
+    deltanet_recurrent_bf16_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)q,
+        (const __nv_bfloat16*)k,
+        (const __nv_bfloat16*)v,
+        (const __nv_bfloat16*)beta,
+        (const __nv_bfloat16*)alpha_softplus,
+        (const __nv_bfloat16*)a_log,
+        (float*)state,
+        (__nv_bfloat16*)out,
+        seq_len,
+        key_heads,
+        value_heads,
+        head_dim
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+}

csrc/attention/flash_attention.cu

@@ -13,13 +13,13 @@
 //         O [batch, num_q_heads, q_len, head_dim]
 //
 // Shared memory (BF16):
-//   smem_q[BR][head_dim]   — 64 * 128 * 2 = 16 KB (loaded once per Q tile)
-//   smem_kv[BC][head_dim]  — 64 * 128 * 2 = 16 KB (alternates K and V)
-//   Total: 32 KB (fits in default 48 KB shared memory)
+// Use 32-row tiles so head_dim=256 still fits in the default 48 KB shared
+// memory limit: 2 * 32 * 256 * 2 = 32 KB.
 
-#define BR 64
-#define BC 64
+#define BR 32
+#define BC 32
 #define THREADS_PER_BLOCK 128
+#define FLASH_HEAD_DIM_MAX 256
 
 __global__ void flash_attention_bf16_kernel(
     const __nv_bfloat16* __restrict__ Q,
@@ -73,8 +73,8 @@ __global__ void flash_attention_bf16_kernel(
     // Thread t (0 <= t < q_tile_rows) owns Q row t
     bool owns_row = (tid < q_tile_rows);
 
-    // Per-thread FP32 accumulators (head_dim up to 128)
-    float O_acc[128];
+    // Per-thread FP32 accumulators (head_dim up to 256)
+    float O_acc[FLASH_HEAD_DIM_MAX];
     float m_val = -INFINITY;
     float l_val = 0.0f;
     if (owns_row) {
@@ -250,7 +250,7 @@ __global__ void flash_attention_sinks_bf16_kernel(
 
     bool owns_row = (tid < q_tile_rows);
 
-    float O_acc[128];
+    float O_acc[FLASH_HEAD_DIM_MAX];
     float m_val = -INFINITY;
     float l_val = 0.0f;
     if (owns_row) {
@@ -384,7 +384,7 @@ __global__ void flash_attention_sinks_bf16_kernel(
 // ============================================================
 
 #define DECODE_THREADS 256
-#define HEAD_DIM_MAX 128
+#define HEAD_DIM_MAX 256
 
 __global__ void decode_attention_bf16_kernel(
     const __nv_bfloat16* __restrict__ Q,
@@ -413,7 +413,7 @@ __global__ void decode_attention_bf16_kernel(
     const __nv_bfloat16* V_base = V + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
     __nv_bfloat16* O_ptr = O + ((long long)batch_idx * num_q_heads + q_head) * head_dim;
 
-    // Load Q vector into registers (head_dim <= 128)
+    // Load Q vector into registers (head_dim <= 256)
     float q_reg[HEAD_DIM_MAX];
     for (int d = 0; d < head_dim; d++) {
         q_reg[d] = __bfloat162float(Q_ptr[d]);
@@ -464,7 +464,7 @@ __global__ void decode_attention_bf16_kernel(
     // Shared memory for reduction
     __shared__ float smem_max[32];   // one per warp
     __shared__ float smem_sum[32];
-    __shared__ float smem_O[HEAD_DIM_MAX]; // final output accumulator
+    __shared__ float smem_O_warp[32][HEAD_DIM_MAX];
 
     // Step 1: Block-wide max reduction
     int lane = tid & 31;
@@ -513,35 +513,30 @@ __global__ void decode_attention_bf16_kernel(
     __syncthreads();
     global_sum = smem_sum[0];
 
-    // Step 4: Reduce O across block (dimension by dimension using shared mem)
+    // Step 4: Reduce O across block, dim by dim. Store one partial per warp
+    // and sum in warp-id order; atomicAdd made greedy decode nondeterministic
+    // when logits were close (same fix pattern as paged_attention.cu / gemv.cu).
     float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
 
-    // Process head_dim in chunks: each iteration reduces one dimension
-    // Use shared memory accumulator: each warp contributes via warp reduction + atomic
-    // Actually simpler: iterate over dimensions, warp reduce each, then lane0 atomicAdd to smem_O
-
-    // Initialize smem_O
-    for (int d = tid; d < head_dim; d += DECODE_THREADS) {
-        smem_O[d] = 0.0f;
+    for (int i = tid; i < 32 * HEAD_DIM_MAX; i += DECODE_THREADS) {
+        reinterpret_cast<float*>(smem_O_warp)[i] = 0.0f;
     }
     __syncthreads();
 
-    // Each thread adds its local_O contributions via warp reduction + atomicAdd
     for (int d = 0; d < head_dim; d++) {
         float val = local_O[d];
-        // Warp-level reduction
         #pragma unroll
         for (int offset = 16; offset > 0; offset >>= 1)
             val += __shfl_down_sync(0xffffffff, val, offset);
-        if (lane == 0) {
-            atomicAdd(&smem_O[d], val);
-        }
+        if (lane == 0) smem_O_warp[warp_id][d] = val;
     }
     __syncthreads();
 
     // Thread 0..head_dim-1 write final output
     for (int d = tid; d < head_dim; d += DECODE_THREADS) {
-        O_ptr[d] = __float2bfloat16(smem_O[d] * inv_sum);
+        float out = 0.0f;
+        for (int i = 0; i < num_warps; i++) out += smem_O_warp[i][d];
+        O_ptr[d] = __float2bfloat16(out * inv_sum);
     }
 }
 

csrc/attention/paged_attention.cu

@@ -21,11 +21,11 @@
 //                 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.
+// One CUDA block: 256 threads, head_dim <= 256.
 
 #define PAGED_BLOCK_SIZE 16
 #define PAGED_THREADS 256
-#define PAGED_HEAD_DIM_MAX 128
+#define PAGED_HEAD_DIM_MAX 256
 
 __global__ void paged_decode_attention_bf16_kernel(
     const __nv_bfloat16* __restrict__ Q,
@@ -118,7 +118,7 @@ __global__ void paged_decode_attention_bf16_kernel(
     // ---- Block-level online softmax reduction ----
     __shared__ float smem_max[32];
     __shared__ float smem_sum[32];
-    __shared__ float smem_O[PAGED_HEAD_DIM_MAX];
+    __shared__ float smem_O_warp[32][PAGED_HEAD_DIM_MAX];
 
     int lane = tid & 31;
     int warp_id = tid >> 5;
@@ -164,8 +164,12 @@ __global__ void paged_decode_attention_bf16_kernel(
     __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;
+    // Step 4: reduce O across block, dim by dim. Store one partial per warp
+    // and sum in warp-id order; atomicAdd made greedy decode nondeterministic
+    // when logits were close.
+    for (int i = tid; i < 32 * PAGED_HEAD_DIM_MAX; i += PAGED_THREADS) {
+        reinterpret_cast<float*>(smem_O_warp)[i] = 0.0f;
+    }
     __syncthreads();
 
     for (int d = 0; d < head_dim; d++) {
@@ -173,13 +177,178 @@ __global__ void paged_decode_attention_bf16_kernel(
         #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);
+        if (lane == 0) smem_O_warp[warp_id][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);
+        float out = 0.0f;
+        for (int i = 0; i < num_warps; i++) out += smem_O_warp[i][d];
+        O_ptr[d] = __float2bfloat16(out * inv_sum);
+    }
+}
+
+// Tree-aware paged decode attention: per-query mask lets sibling candidates
+// in the same batch attend to different subsets of newly-written K/V.
+// `tree_start`: position where newly-written K/V begins (typically pos_offset).
+// `tree_len`: number of newly-written K/V rows (= batch, one per query).
+// `tree_mask[i][j] = 1` iff query i attends to K/V at position `tree_start+j`.
+// Positions < tree_start are always attended (regular history).
+__global__ void paged_decode_attention_tree_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,
+    const int* __restrict__ context_lens,
+    const int* __restrict__ tree_mask,       // [batch, tree_len] int32
+    int num_q_heads, int num_kv_heads,
+    int head_dim, int max_blocks_per_seq,
+    int tree_start, int tree_len,
+    float scale
+) {
+    int seq_idx = blockIdx.y;
+    int q_head  = blockIdx.x;
+    int tid     = threadIdx.x;
+
+    int kv_len = context_lens[seq_idx];
+    if (kv_len <= 0) {
+        if (tid < head_dim) {
+            O[((long long)seq_idx * num_q_heads + q_head) * head_dim + tid] =
+                __float2bfloat16(0.0f);
+        }
+        return;
+    }
+
+    int heads_per_group = num_q_heads / num_kv_heads;
+    int kv_head = q_head / heads_per_group;
+
+    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;
+    const int* mask_row = tree_mask + (long long)seq_idx * tree_len;
+
+    float q_reg[PAGED_HEAD_DIM_MAX];
+    for (int d = 0; d < head_dim; d++) {
+        q_reg[d] = __bfloat162float(Q_ptr[d]);
+    }
+
+    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;
+
+    for (int pos = tid; pos < kv_len; pos += PAGED_THREADS) {
+        // Tree mask: skip positions in [tree_start, tree_start+tree_len) that
+        // the mask marks as 0. Everything else (history) is always attended.
+        if (pos >= tree_start && pos < tree_start + tree_len) {
+            if (mask_row[pos - tree_start] == 0) continue;
+        }
+
+        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;
+
+        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;
+
+        for (int d = 0; d < head_dim; d++) {
+            local_O[d] += p * __bfloat162float(V_pos[d]);
+        }
+
+        local_max = new_max;
+    }
+
+    // Block-level reduction (identical to base kernel).
+    __shared__ float smem_max[32];
+    __shared__ float smem_sum[32];
+    __shared__ float smem_O_warp[32][PAGED_HEAD_DIM_MAX];
+
+    int lane = tid & 31;
+    int warp_id = tid >> 5;
+    int num_warps = PAGED_THREADS >> 5;
+
+    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];
+
+    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;
+
+    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];
+
+    for (int i = tid; i < 32 * PAGED_HEAD_DIM_MAX; i += PAGED_THREADS) {
+        reinterpret_cast<float*>(smem_O_warp)[i] = 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) smem_O_warp[warp_id][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) {
+        float out = 0.0f;
+        for (int i = 0; i < num_warps; i++) out += smem_O_warp[i][d];
+        O_ptr[d] = __float2bfloat16(out * inv_sum);
     }
 }
 
@@ -289,7 +458,7 @@ __global__ void paged_decode_attention_sinks_bf16_kernel(
     // ---- Block-level online softmax reduction (same as base kernel) ----
     __shared__ float smem_max[32];
     __shared__ float smem_sum[32];
-    __shared__ float smem_O[PAGED_HEAD_DIM_MAX];
+    __shared__ float smem_O_warp[32][PAGED_HEAD_DIM_MAX];
 
     int lane = tid & 31;
     int warp_id = tid >> 5;
@@ -332,7 +501,9 @@ __global__ void paged_decode_attention_sinks_bf16_kernel(
     __syncthreads();
     global_sum = smem_sum[0];
 
-    for (int d = tid; d < head_dim; d += PAGED_THREADS) smem_O[d] = 0.0f;
+    for (int i = tid; i < 32 * PAGED_HEAD_DIM_MAX; i += PAGED_THREADS) {
+        reinterpret_cast<float*>(smem_O_warp)[i] = 0.0f;
+    }
     __syncthreads();
 
     for (int d = 0; d < head_dim; d++) {
@@ -340,13 +511,15 @@ __global__ void paged_decode_attention_sinks_bf16_kernel(
         #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);
+        if (lane == 0) smem_O_warp[warp_id][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);
+        float out = 0.0f;
+        for (int i = 0; i < num_warps; i++) out += smem_O_warp[i][d];
+        O_ptr[d] = __float2bfloat16(out * inv_sum);
     }
 }
 
@@ -379,6 +552,36 @@ void launch_paged_decode_attention_bf16(
     CUDA_CHECK_LAST_ERROR();
 }
 
+void launch_paged_decode_attention_tree_bf16(
+    const void* Q,
+    const void* K_cache,
+    const void* V_cache,
+    void* O,
+    const int* block_tables,
+    const int* context_lens,
+    const int* tree_mask,
+    int batch, int num_q_heads, int num_kv_heads,
+    int head_dim, int max_blocks_per_seq,
+    int tree_start, int tree_len,
+    float scale, void* stream
+) {
+    dim3 grid(num_q_heads, batch);
+    int block = PAGED_THREADS;
+
+    paged_decode_attention_tree_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, tree_mask,
+        num_q_heads, num_kv_heads,
+        head_dim, max_blocks_per_seq,
+        tree_start, tree_len,
+        scale
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
 void launch_paged_decode_attention_sinks_bf16(
     const void* Q,
     const void* K_cache,

csrc/attention/reshape_and_cache.cu

@@ -158,4 +158,58 @@ void launch_reshape_and_cache_batched_bf16(
     CUDA_CHECK_LAST_ERROR();
 }
 
+// Copy one token's K/V from src_pos to dst_pos within one pool.
+// Grid: (num_kv_heads,). Block: head_dim threads.
+// pool: [num_blocks_total, num_kv_heads, block_size, head_dim]
+// block_ids: [max_blocks] for this sequence (logical → physical block map).
+__global__ void copy_kv_position_kernel(
+    __nv_bfloat16* __restrict__ pool,
+    const int* __restrict__ block_ids,
+    int src_pos, int dst_pos,
+    int head_dim, int block_size
+) {
+    int h = blockIdx.x;
+    int d = threadIdx.x;
+    if (d >= head_dim) return;
+
+    int num_kv_heads = gridDim.x;
+
+    int src_blk = src_pos / block_size;
+    int src_slot = src_pos % block_size;
+    int src_phys = block_ids[src_blk];
+
+    int dst_blk = dst_pos / block_size;
+    int dst_slot = dst_pos % block_size;
+    int dst_phys = block_ids[dst_blk];
+
+    long long src_off = ((long long)src_phys * num_kv_heads + h) * block_size * head_dim
+                      + src_slot * head_dim + d;
+    long long dst_off = ((long long)dst_phys * num_kv_heads + h) * block_size * head_dim
+                      + dst_slot * head_dim + d;
+
+    pool[dst_off] = pool[src_off];
+}
+
+void launch_copy_kv_position(
+    void* k_pool, void* v_pool,
+    const int* block_ids,
+    int src_pos, int dst_pos,
+    int num_kv_heads, int head_dim, int block_size,
+    void* stream
+) {
+    int threads = head_dim < 32 ? 32 : head_dim;
+    if (threads > 1024) threads = 1024;
+    dim3 grid(num_kv_heads);
+    copy_kv_position_kernel<<<grid, threads, 0, (cudaStream_t)stream>>>(
+        (__nv_bfloat16*)k_pool, block_ids,
+        src_pos, dst_pos, head_dim, block_size
+    );
+    CUDA_CHECK_LAST_ERROR();
+    copy_kv_position_kernel<<<grid, threads, 0, (cudaStream_t)stream>>>(
+        (__nv_bfloat16*)v_pool, block_ids,
+        src_pos, dst_pos, head_dim, block_size
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
 }

csrc/common.cuh

@@ -49,10 +49,12 @@ __device__ __forceinline__ float block_reduce_max(float val) {
     return val;
 }
 
-// --- Launch error checking (debug builds only) ---
-#ifdef NDEBUG
-#define CUDA_CHECK_LAST_ERROR() ((void)0)
-#else
+// --- Launch error checking ---
+// Always on, including release builds. A launch with an invalid config
+// (e.g. 32-bit overflow in grid/index math) is otherwise silent and produces
+// garbage with no clue — the MoE int32-overflow bug was found exactly because
+// release swallowed the launch failure. `cudaGetLastError()` does not
+// synchronize the stream, so the per-launch host cost is negligible.
 #include <cstdio>
 #define CUDA_CHECK_LAST_ERROR() do { \
     cudaError_t err = cudaGetLastError(); \
@@ -61,4 +63,3 @@ __device__ __forceinline__ float block_reduce_max(float val) {
                 __FILE__, __LINE__, cudaGetErrorString(err)); \
     } \
 } while(0)
-#endif

csrc/embedding/rope.cu

@@ -69,6 +69,42 @@ __global__ void rope_bf16(
     x[base + pair_idx + half_dim] = __float2bfloat16(x1 * cos_val + x0 * sin_val);
 }
 
+__global__ void partial_rope_bf16(
+    const __nv_bfloat16* __restrict__ x,
+    __nv_bfloat16* __restrict__ out,
+    const int* __restrict__ positions,
+    int num_heads, int head_dim, int n_rot, float theta
+) {
+    int token_idx = blockIdx.x;
+    int head_idx = blockIdx.y;
+    int pair_idx = threadIdx.x;
+    int half_rot = n_rot / 2;
+    if (pair_idx >= half_rot) return;
+
+    int pos = positions[token_idx];
+    float freq = 1.0f / powf(theta, (float)(2 * pair_idx) / (float)n_rot);
+    float angle = (float)pos * freq;
+    float sin_val, cos_val;
+    sincosf(angle, &sin_val, &cos_val);
+
+    int base = (token_idx * num_heads + head_idx) * head_dim;
+    float x0 = __bfloat162float(x[base + pair_idx]);
+    float x1 = __bfloat162float(x[base + pair_idx + half_rot]);
+
+    out[base + pair_idx]            = __float2bfloat16(x0 * cos_val - x1 * sin_val);
+    out[base + pair_idx + half_rot] = __float2bfloat16(x1 * cos_val + x0 * sin_val);
+}
+
+__global__ void copy_partial_rope_tail_bf16(
+    const __nv_bfloat16* __restrict__ x,
+    __nv_bfloat16* __restrict__ out,
+    int total
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total) return;
+    out[idx] = x[idx];
+}
+
 // Precompute cos/sin cache on GPU
 __global__ void compute_rope_cache(
     float* __restrict__ cos_cache,  // [max_seq_len, half_dim]
@@ -109,6 +145,24 @@ void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache,
     CUDA_CHECK_LAST_ERROR();
 }
 
+void launch_partial_rope_bf16(const void* x, void* out, const void* positions,
+                              int num_tokens, int num_heads, int head_dim,
+                              int n_rot, float theta, void* stream) {
+    int total = num_tokens * num_heads * head_dim;
+    int block_copy = 256;
+    int grid_copy = (total + block_copy - 1) / block_copy;
+    copy_partial_rope_tail_bf16<<<grid_copy, block_copy, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, total);
+    CUDA_CHECK_LAST_ERROR();
+
+    dim3 grid(num_tokens, num_heads);
+    int block = n_rot / 2;
+    partial_rope_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, (const int*)positions,
+        num_heads, head_dim, n_rot, theta);
+    CUDA_CHECK_LAST_ERROR();
+}
+
 void launch_compute_rope_cache(void* cos_cache, void* sin_cache,
                                int max_seq_len, int half_dim, float theta,
                                void* stream) {

csrc/gemm/gemv.cu

@@ -6,22 +6,20 @@
 //
 //   y[n] = sum_k x[k] * W[k * N + n]
 //
-// Grid: (N / TILE_N, K / TILE_K).
-// All blocks atomicAdd their partial sums into a pre-zeroed FP32 buffer.
-// A separate conversion kernel writes the final BF16 output.
-// Launch sequence: cudaMemsetAsync(fp32) → accumulation kernel → convert kernel.
+// Grid: (N / TILE_N, K / TILE_K) partials, followed by a deterministic
+// fixed-order reduction over K blocks. The previous implementation used
+// atomicAdd into y_fp32[col]; that made BF16 greedy decode sensitive to
+// inter-block scheduling when logits were close.
 
 #define GEMV_TILE_N 128
 #define GEMV_TILE_K 256
 #define GEMV_BLOCK  128
 
-__global__ void gemv_bf16_fused_kernel(
+__global__ void gemv_bf16_partial_kernel(
     const __nv_bfloat16* __restrict__ x,
     const __nv_bfloat16* __restrict__ W,
-    __nv_bfloat16* __restrict__ y_bf16,
-    float* __restrict__ y_fp32,
-    int K, int N,
-    int num_k_blocks
+    float* __restrict__ partials,
+    int K, int N
 ) {
     const int block_n = blockIdx.x;
     const int block_k = blockIdx.y;
@@ -52,21 +50,81 @@ __global__ void gemv_bf16_fused_kernel(
         sum += x_shared[ki] * __bfloat162float(W[(long long)(k_start + ki) * N + col]);
     }
 
-    atomicAdd(&y_fp32[col], sum);
+    partials[(long long)block_k * N + col] = sum;
 }
 
-// Conversion kernel: FP32 accumulator -> BF16 output
-__global__ void gemv_fp32_to_bf16_kernel(
-    const float* __restrict__ src,
+__global__ void gemv_reduce_to_bf16_kernel(
+    const float* __restrict__ partials,
     __nv_bfloat16* __restrict__ dst,
-    int n
+    int n,
+    int num_k_blocks
 ) {
     int idx = blockIdx.x * blockDim.x + threadIdx.x;
     if (idx < n) {
-        dst[idx] = __float2bfloat16(src[idx]);
+        float sum = 0.0f;
+        for (int kb = 0; kb < num_k_blocks; kb++) {
+            sum += partials[(long long)kb * n + idx];
+        }
+        dst[idx] = __float2bfloat16(sum);
     }
 }
 
+// Batched variant: M rows, same W. Grid.z = batch row index.
+// Numerically identical to calling launch_gemv_bf16 M times in sequence because
+// each z-slice executes the same accumulation order on the same data.
+// partials buffer must be [M * num_k_blocks * N] floats.
+__global__ void gemv_bf16_batched_partial_kernel(
+    const __nv_bfloat16* __restrict__ x,  // [M, K]
+    const __nv_bfloat16* __restrict__ W,  // [K, N]
+    float* __restrict__ partials,         // [M, num_k_blocks, N]
+    int K, int N
+) {
+    const int block_n = blockIdx.x;
+    const int block_k = blockIdx.y;
+    const int row = blockIdx.z;
+    const int t = threadIdx.x;
+    const int col = block_n * GEMV_TILE_N + t;
+
+    const int k_start = block_k * GEMV_TILE_K;
+    const int k_end = min(k_start + GEMV_TILE_K, K);
+    const int k_len = k_end - k_start;
+
+    __shared__ float x_shared[GEMV_TILE_K];
+    const __nv_bfloat16* x_row = x + (long long)row * K;
+    for (int i = t; i < k_len; i += GEMV_BLOCK) {
+        x_shared[i] = __bfloat162float(x_row[k_start + i]);
+    }
+    __syncthreads();
+
+    if (col >= N) return;
+
+    float sum = 0.0f;
+    for (int ki = 0; ki < k_len; ki++) {
+        sum += x_shared[ki] * __bfloat162float(W[(long long)(k_start + ki) * N + col]);
+    }
+
+    int num_k_blocks = (K + GEMV_TILE_K - 1) / GEMV_TILE_K;
+    partials[((long long)row * num_k_blocks + block_k) * N + col] = sum;
+}
+
+__global__ void gemv_batched_reduce_to_bf16_kernel(
+    const float* __restrict__ partials,  // [M, num_k_blocks, N]
+    __nv_bfloat16* __restrict__ dst,     // [M, N]
+    int n,
+    int num_k_blocks
+) {
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+    int row = blockIdx.y;
+    if (col >= n) return;
+
+    float sum = 0.0f;
+    const float* row_partials = partials + (long long)row * num_k_blocks * n;
+    for (int kb = 0; kb < num_k_blocks; kb++) {
+        sum += row_partials[(long long)kb * n + col];
+    }
+    dst[(long long)row * n + col] = __float2bfloat16(sum);
+}
+
 extern "C" {
 
 void launch_gemv_bf16(
@@ -79,30 +137,58 @@ void launch_gemv_bf16(
 ) {
     cudaStream_t s = (cudaStream_t)stream;
 
-    // Zero the FP32 accumulator BEFORE the kernel — the kernel uses atomicAdd
-    // across K-blocks with no inter-block ordering, so the buffer must be
-    // pre-zeroed to avoid accumulating on stale data.
-    cudaMemsetAsync(y_fp32_buf, 0, (size_t)N * sizeof(float), s);
-
     int num_k_blocks = (K + GEMV_TILE_K - 1) / GEMV_TILE_K;
     dim3 grid((N + GEMV_TILE_N - 1) / GEMV_TILE_N, num_k_blocks);
 
-    gemv_bf16_fused_kernel<<<grid, GEMV_BLOCK, 0, s>>>(
+    gemv_bf16_partial_kernel<<<grid, GEMV_BLOCK, 0, s>>>(
         (const __nv_bfloat16*)x,
         (const __nv_bfloat16*)W,
-        (__nv_bfloat16*)y_bf16,
         (float*)y_fp32_buf,
-        K, N, num_k_blocks
+        K, N
     );
     CUDA_CHECK_LAST_ERROR();
 
-    // FP32 → BF16 conversion (must wait for all K-blocks to finish)
+    // Fixed-order FP32 reduction over K blocks, then BF16 conversion.
     int conv_block = 256;
     int conv_grid = (N + conv_block - 1) / conv_block;
-    gemv_fp32_to_bf16_kernel<<<conv_grid, conv_block, 0, s>>>(
+    gemv_reduce_to_bf16_kernel<<<conv_grid, conv_block, 0, s>>>(
         (const float*)y_fp32_buf,
         (__nv_bfloat16*)y_bf16,
-        N
+        N,
+        num_k_blocks
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
+void launch_gemv_bf16_batched(
+    const void* x,     // [M, K] BF16
+    const void* W,     // [K, N] BF16
+    void* y_bf16,      // [M, N] BF16
+    void* y_fp32_buf,  // [M * num_k_blocks * N] FP32
+    int M, int K, int N,
+    void* stream
+) {
+    cudaStream_t s = (cudaStream_t)stream;
+
+    int num_k_blocks = (K + GEMV_TILE_K - 1) / GEMV_TILE_K;
+    dim3 grid((N + GEMV_TILE_N - 1) / GEMV_TILE_N, num_k_blocks, M);
+
+    gemv_bf16_batched_partial_kernel<<<grid, GEMV_BLOCK, 0, s>>>(
+        (const __nv_bfloat16*)x,
+        (const __nv_bfloat16*)W,
+        (float*)y_fp32_buf,
+        K, N
+    );
+    CUDA_CHECK_LAST_ERROR();
+
+    int conv_block = 256;
+    int conv_grid_x = (N + conv_block - 1) / conv_block;
+    dim3 reduce_grid(conv_grid_x, M);
+    gemv_batched_reduce_to_bf16_kernel<<<reduce_grid, conv_block, 0, s>>>(
+        (const float*)y_fp32_buf,
+        (__nv_bfloat16*)y_bf16,
+        N,
+        num_k_blocks
     );
     CUDA_CHECK_LAST_ERROR();
 }

csrc/moe/moe_kernels.cu

@@ -89,13 +89,17 @@ __global__ void moe_replicate_bf16_kernel(
     __nv_bfloat16* __restrict__ x_rep,
     int num_tokens, int hidden, int local_experts
 ) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    int total = local_experts * num_tokens * hidden;
+    // 64-bit index: local_experts * num_tokens * hidden overflows int32 at
+    // ~2.3k prefill tokens (gpt-oss TP=1, 32 experts), which is inside the
+    // supported context window. A 32-bit `total` silently wraps, the launch
+    // fails, and (in release) the error is invisible — see common.cuh.
+    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    long long total = (long long)local_experts * num_tokens * hidden;
     if (idx >= total) return;
 
-    int remainder = idx % (num_tokens * hidden);
     // x_rep[expert, token, dim] = x[token, dim]
-    x_rep[idx] = x[remainder];
+    long long row_stride = (long long)num_tokens * hidden;
+    x_rep[idx] = x[idx % row_stride];
 }
 
 // ============================================================
@@ -112,13 +116,16 @@ __global__ void moe_bias_add_3d_bf16_kernel(
     const __nv_bfloat16* __restrict__ bias,
     int batch, int num_tokens, int dim
 ) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    int total = batch * num_tokens * dim;
+    // 64-bit index: batch * num_tokens * dim overflows int32 at ~3.6k prefill
+    // tokens (gpt-oss TP=1, 32 experts, 2*intermediate dim) — see moe_replicate.
+    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    long long total = (long long)batch * num_tokens * dim;
     if (idx >= total) return;
 
-    int b = idx / (num_tokens * dim);
-    int d = idx % dim;
-    float v = __bfloat162float(x[idx]) + __bfloat162float(bias[b * dim + d]);
+    long long td = (long long)num_tokens * dim;
+    int b = (int)(idx / td);   // < batch (small)
+    int d = (int)(idx % dim);  // < dim
+    float v = __bfloat162float(x[idx]) + __bfloat162float(bias[(long long)b * dim + d]);
     x[idx] = __float2bfloat16(v);
 }
 
@@ -151,14 +158,16 @@ __global__ void moe_weighted_sum_bf16_kernel(
     int num_tokens, int hidden, int top_k,
     int expert_start, int local_experts
 ) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    int total = num_tokens * hidden;
+    // 64-bit index: `local_id * expert_stride` overflows int32 for long prefills
+    // (expert_stride = num_tokens * hidden), reading the wrong expert element.
+    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    long long total = (long long)num_tokens * hidden;
     if (idx >= total) return;
 
-    int token = idx / hidden;
-    int dim = idx % hidden;
+    long long token = idx / hidden;
+    int dim = (int)(idx % hidden);
 
-    int expert_stride = num_tokens * hidden; // stride between experts in expert_out
+    long long expert_stride = (long long)num_tokens * hidden; // stride between experts in expert_out
 
     float sum = 0.0f;
     for (int k = 0; k < top_k; k++) {
@@ -196,9 +205,9 @@ void launch_moe_replicate_bf16(
     int num_tokens, int hidden, int local_experts,
     void* stream
 ) {
-    int total = local_experts * num_tokens * hidden;
+    long long total = (long long)local_experts * num_tokens * hidden;
     int block = 256;
-    int grid = (total + block - 1) / block;
+    int grid = (int)((total + block - 1) / block);
     moe_replicate_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)x, (__nv_bfloat16*)x_rep,
         num_tokens, hidden, local_experts
@@ -211,9 +220,9 @@ void launch_moe_bias_add_3d_bf16(
     int batch, int num_tokens, int dim,
     void* stream
 ) {
-    int total = batch * num_tokens * dim;
+    long long total = (long long)batch * num_tokens * dim;
     int block = 256;
-    int grid = (total + block - 1) / block;
+    int grid = (int)((total + block - 1) / block);
     moe_bias_add_3d_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (__nv_bfloat16*)x, (const __nv_bfloat16*)bias,
         batch, num_tokens, dim
@@ -229,9 +238,9 @@ void launch_moe_weighted_sum_bf16(
     int expert_start, int local_experts,
     void* stream
 ) {
-    int total = num_tokens * hidden;
+    long long total = (long long)num_tokens * hidden;
     int block = 256;
-    int grid = (total + block - 1) / block;
+    int grid = (int)((total + block - 1) / block);
     moe_weighted_sum_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)expert_out,
         (const int*)topk_ids, (const float*)topk_weights,

csrc/moe/moe_sparse.cu

@@ -34,6 +34,51 @@
 
 #define SPARSE_TILE_N 8  // output columns per block (= warps per block)
 
+// Weights FP8 E4M3 [local_experts, N, K], activations BF16 (W8A16).
+// Decode is memory-bound (~2 FLOP/byte), so dequant-in-registers GEMV
+// loses nothing to tensor cores and skips activation quantization.
+__global__ void moe_sparse_gemv_bf16_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,     // [T, K] or [T*top_k, K]
+    const __nv_bfloat16* __restrict__ w,     // [local_experts, N, K]
+    const int* __restrict__ topk_ids,        // [T, top_k] global expert ids
+    __nv_bfloat16* __restrict__ y,           // [T, top_k, N]
+    int N, int K, int top_k,
+    int expert_start, int local_experts,
+    int x_per_slot
+) {
+    int token = blockIdx.z;
+    int slot = blockIdx.y;
+    int eid = topk_ids[token * top_k + slot];
+    int lid = eid - expert_start;
+    if (lid < 0 || lid >= local_experts) return;
+
+    extern __shared__ float xs[];
+    const __nv_bfloat16* xrow =
+        x + (long long)(x_per_slot ? token * top_k + slot : token) * K;
+    for (int i = threadIdx.x; i < K; i += blockDim.x) {
+        xs[i] = __bfloat162float(xrow[i]);
+    }
+    __syncthreads();
+
+    int n = blockIdx.x * SPARSE_TILE_N + (threadIdx.x >> 5);
+    if (n >= N) return;
+    int lane = threadIdx.x & 31;
+
+    const __nv_bfloat16* wrow = w + ((long long)lid * N + n) * K;
+    float acc = 0.0f;
+    for (int i = lane; i < K; i += 32) {
+        acc += xs[i] * __bfloat162float(wrow[i]);
+    }
+
+    #pragma unroll
+    for (int o = 16; o > 0; o >>= 1) {
+        acc += __shfl_down_sync(0xffffffffu, acc, o);
+    }
+    if (lane == 0) {
+        y[((long long)token * top_k + slot) * N + n] = __float2bfloat16(acc);
+    }
+}
+
 // Weights FP8 E4M3 [local_experts, N, K], activations BF16 (W8A16).
 // Decode is memory-bound (~2 FLOP/byte), so dequant-in-registers GEMV
 // loses nothing to tensor cores and skips activation quantization.
@@ -194,6 +239,23 @@ __global__ void moe_weighted_sum_sparse_bf16_kernel(
 
 extern "C" {
 
+void launch_moe_sparse_gemv_bf16_bf16(
+    const void* x, const void* w, const void* topk_ids, void* y,
+    int num_tokens, int N, int K, int top_k,
+    int expert_start, int local_experts, int x_per_slot,
+    void* stream
+) {
+    dim3 grid((N + SPARSE_TILE_N - 1) / SPARSE_TILE_N, top_k, num_tokens);
+    int block = SPARSE_TILE_N * 32;
+    size_t smem = (size_t)K * sizeof(float);
+    moe_sparse_gemv_bf16_bf16_kernel<<<grid, block, smem, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)w, (const int*)topk_ids,
+        (__nv_bfloat16*)y,
+        N, K, top_k, expert_start, local_experts, x_per_slot
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
 void launch_moe_sparse_gemv_fp8_bf16(
     const void* x, const void* w, const void* w_scales, const void* bias,
     const void* topk_ids, void* y,

csrc/quantization/dequant_fp8.cu

@@ -16,12 +16,14 @@ __global__ void dequant_fp8e4m3_to_bf16_kernel(
     __nv_bfloat16* __restrict__ dst,
     int num_experts, int rows, int cols
 ) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    int total = num_experts * rows * cols;
+    // 64-bit index: num_experts * rows * cols overflows int32 for 32 experts
+    // at ~8k*8k weight matrices, same class as the MoE fix in cfbd64d.
+    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    long long total = (long long)num_experts * rows * cols;
     if (idx >= total) return;
 
-    int expert_stride = rows * cols;
-    int expert = idx / expert_stride;
+    long long expert_stride = (long long)rows * cols;
+    int expert = (int)(idx / expert_stride);
     float scale = scales[expert];
     float val = float(src[idx]) * scale;
     dst[idx] = __float2bfloat16(val);
@@ -36,9 +38,9 @@ void launch_dequant_fp8e4m3_to_bf16(
     int num_experts, int rows, int cols,
     void* stream
 ) {
-    int total = num_experts * rows * cols;
+    long long total = (long long)num_experts * rows * cols;
     int block = 256;
-    int grid = (total + block - 1) / block;
+    int grid = (int)((total + block - 1) / block);
     dequant_fp8e4m3_to_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
         (const __nv_fp8_e4m3*)src,
         (const float*)scales,

csrc/reduce/softmax.cu

@@ -90,7 +90,7 @@ __global__ void softmax_bf16(
 extern "C" {
 
 void launch_softmax_f32(const void* x, void* out, int rows, int cols, void* stream) {
-    int block = (cols < 1024) ? cols : 1024;
+    int block = (cols < 512) ? cols : 512;
     if (block < 32) block = 32;
     softmax_f32<<<rows, block, 0, (cudaStream_t)stream>>>(
         (const float*)x, (float*)out, cols);
@@ -98,7 +98,7 @@ void launch_softmax_f32(const void* x, void* out, int rows, int cols, void* stre
 }
 
 void launch_softmax_bf16(const void* x, void* out, int rows, int cols, void* stream) {
-    int block = (cols < 1024) ? cols : 1024;
+    int block = (cols < 512) ? cols : 512;
     if (block < 32) block = 32;
     softmax_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)x, (__nv_bfloat16*)out, cols);

docs/22-speculative-decoding.md

@@ -0,0 +1,186 @@
+# Phase 22: Draft-Model Speculative Decoding v0
+
+> 目标:实现一个可验证的 speculative decoding 最小闭环。先只覆盖
+> Qwen3 target + 同 tokenizer 的小 Qwen3 draft、batch=1、greedy
+> (`temperature=0`)。本阶段不做 gpt-oss,不做 sampling rejection,不接入
+> continuous batching。
+
+## 1. Scope
+
+本阶段只解决一个窄问题:
+
+- target:现有 Qwen3 paged KV 路径,优先 Qwen3-8B;
+- draft:同 tokenizer 的小 Qwen3,例如 Qwen3-0.6B;
+- batch size:1;
+- decoding:greedy argmax;
+- draft window:`gamma=4`;
+- acceptance:exact-match,即 `target_argmax == draft_token`。
+
+HTTP flag 可以后续接入。v0 先提供独立 bench/CLI,因为它能直接输出 token
+一致性、acceptance rate、tokens/target-step、TPOT/tok/s,也避免把尚未稳定的
+rollback 行为放进服务端调度循环。
+
+bench 为了让 baseline/spec 对比不受跨 prompt KV pool 复用影响,每个 prompt 的
+baseline run 和 speculative run 都使用新建的 paged KV cache。cache 分配发生在
+单次 run 的计时外,输出的 TPOT/tok/s 只覆盖模型 prefill/decode 工作。
+
+## 2. Why Qwen3 First
+
+Qwen3 是现有代码里最适合作为 speculative v0 的模型族:
+
+1. target 已有稳定的 `forward_prefill_paged` 和 `forward_decode_paged`;
+2. 小 Qwen3 与 Qwen3-8B 共享 tokenizer,可以直接比较 token id;
+3. Qwen3 是 dense decoder-only,没有 gpt-oss 的 harmony 格式、MoE sparse 路径、
+   sliding-window 或 CUDA Graph 状态;
+4. greedy 输出的正确性定义简单:只要 spec 生成的 token 序列与纯 target greedy
+   完全一致即可。
+
+gpt-oss spec 需要先定义 harmony prompt、MoE draft 选择、graph replay 与 rollback
+的交互,这些都不属于本阶段。
+
+## 3. Algorithm
+
+对每个 prompt 建两套模型、三套 KV 状态:
+
+```text
+target model + target commit PagedKVCache
+target model + target verify PagedKVCache
+draft  model + draft PagedKVCache
+```
+
+先把 prompt 分别 prefill 到三套 cache。此时 cache 都包含 prompt,并各自持有
+"下一个 token" 的 logits。
+
+每个 speculative round:
+
+1. draft 从当前 draft logits 取 argmax,连续生成 `gamma` 个 draft token;
+2. draft 每生成一个 token 就用自己的 paged decode append 到 draft KV,所以 round
+   结束时 draft cache 暂时包含整个草稿序列;
+3. target verify cache 对完整 draft token 序列调用一次 paged prefill,覆盖
+   "target 可一次验证草稿窗口" 这条执行路径;
+4. target verify cache 立刻 rollback 到 round 起点,避免把 prefill 临时写入污染
+   commit cache;
+5. 用 target decode 轨迹作为权威结果,从左到右比较
+   `target_next_argmax == draft_token`,只接受连续匹配前缀;
+6. 对每个接受 token,用 target decode 重放一次来提交 target KV,并得到下一步
+   `target_next_argmax`;verify cache 也 mirror decode 同一个 token,保持长度与 prefix 对齐;
+7. 若全部匹配,draft cache 已经包含完整草稿,三套 cache 长度重新对齐;
+8. 若在第 `k` 个 token 拒绝,提交前 `k` 个 draft token,再提交 target 在该位置的
+   argmax 作为修正 token。draft cache rollback 到 round 起点后重放接受 token 和修正
+   token,target commit/verify cache 都由 decode 路径提交到同一 prefix。
+
+v0 不使用完整 speculative sampling 的概率校正。它只利用小模型猜测 greedy 轨迹,
+因此生成序列必须与纯 target greedy 完全一致。
+
+当前实现选择 decode 轨迹作为提交路径,而不是直接保留 target prefill 写入的 KV。
+原因是 v0 验收要求 token 序列与纯 target greedy 完全一致;如果 prefill 和 decode
+路径在数值或 KV 写入顺序上存在细微差异,直接提交 prefill KV 会让后续 greedy 输出
+漂移。这个保守实现仍会执行 target paged prefill 验证和 rollback,但 verify 写入放在
+独立 cache,不会影响权威 commit cache。代价是额外 mirror decode,速度收益预期较差,
+主要用于先验证 draft-model speculative 的状态机和一致性。
+
+为保证 greedy exactness,decode 里两个原有非确定点也需要固定:
+
+- BF16 GEMV 不再用跨 K-block `atomicAdd`;改为写 K-block partials,再按固定顺序
+  reduce;
+- paged decode attention 不再用 `atomicAdd` 合并 warp 输出;改为 per-warp partials
+  后按 warp id 顺序 reduce。
+
+## 4. KV Commit And Rollback
+
+现有 `forward_prefill_paged` 会一次性把传入 token 写进 paged KV,并提前推进
+`seq_len`。验证草稿时 target verify cache 因此会临时包含整个 draft window。
+
+新增的 cache 操作只做逻辑截断:
+
+```text
+truncate_sequence(slot, new_len)
+```
+
+约束:
+
+- 只允许 `new_len <= current_len`;
+- 保留覆盖 `[0, new_len)` 所需的物理 block;
+- 释放右侧多余 block;
+- 不清零仍在保留 block 内的旧字节,因为后续逻辑长度会阻止 attention 读取它们,
+  同一位置再次写入时会覆盖旧值;
+- slot 仍保持 registered,`new_len=0` 时也保留第一个 block。
+
+这让 target 和 draft 都能在拒绝时安全丢弃多写 KV,并在修正 token decode 后重新
+对齐。
+
+## 5. Acceptance Criteria
+
+本阶段验收:
+
+- `cargo fmt`;
+- `cargo check`;
+- `cargo test`;
+- `bench-speculative` 可加载 target+draft 两套 Qwen3;
+- 50 prompts,greedy,baseline target 与 speculative token id 序列完全一致;
+- 输出 acceptance rate、tokens/target-step、TPOT、tok/s 和 speedup;
+- 若 draft 模型缺失或磁盘不足,明确报告阻塞条件,不盲目下载大模型。
+
+## 6. Validation Results
+
+dash5 环境:
+
+- GPU:RTX 5090,device 0;
+- target:`/opt/wjh/models/qwen3-8b`;
+- draft:`/dashscope-tmp/wjh/models/qwen3-0.6b`;
+- command:`bench-speculative ... --prompts 50 --gen-tokens 32 --gamma 4 --device 0`;
+- log:`/dashscope-tmp/wjh/xserv-spec-default-50x32-final.log`。
+
+默认 `acceptance_mode=decode` 的结果:
+
+```text
+prompts=50 matched=true
+acceptance_rate=0.3664 accepted=1020 proposed=2784
+tokens_per_target_step=0.3639 target_steps=4397
+verify_steps=729 mirror_decode_steps=1550 commit_decode_steps=1550 correction_steps=568
+verify_decode_mismatches=10
+baseline_e2e_tpot_ms=13.123 baseline_e2e_tok_s=76.204
+spec_e2e_tpot_ms=44.867 spec_e2e_tok_s=22.288 speedup_e2e=0.2925
+baseline_decode_tpot_ms=12.638 baseline_decode_tok_s=79.127
+spec_decode_tpot_ms=43.731 spec_decode_tok_s=22.867 speedup_decode=0.2890
+decode_token_counts baseline=1600 spec=1600
+```
+
+诊断 `--use-verify-logits` 的结果:
+
+- command:`bench-speculative ... --prompts 10 --gen-tokens 32 --gamma 4 --device 0 --use-verify-logits`;
+- log:`/dashscope-tmp/wjh/xserv-spec-verify-logits-10x32.log`;
+- exit status:`2`;
+- summary:`matched=false`, `verify_decode_mismatches=4`;
+- prompt 0/2/7 出现 baseline/spec token 序列分叉。
+
+结论:当前可以做 correctness-first 的 speculative decoding 状态机,但还不能把
+target batched prefill verify logits 作为 greedy 接受依据。verify prefill 路径与
+逐 token decode 路径存在 top-1 不一致;默认模式必须继续以 decode 轨迹为权威,
+因此 v0 是正确性闭环,不是性能优化。
+
+## 7. Known Limits
+
+- 只支持 batch=1;
+- 只支持 Qwen3-family dense models;
+- 只支持 greedy exact-match acceptance;
+- 未实现 probabilistic rejection sampling,所以 temperature/top-k/top-p 不支持;
+- 未接 HTTP/continuous batching;
+- 未与 CUDA Graph decode 结合;
+- 当前 v0 为保证 greedy exactness,接受 token 也会用 target decode 重放提交,因此
+  即使 acceptance 高也可能变慢;
+- draft prefill 和 target prefill 都会计入端到端耗时,短输出可能没有收益。
+
+## 8. Next Phase TODO
+
+如果继续 speculative decoding,下一阶段不要先接 HTTP,应先解决 verify 路径:
+
+1. 做最小 prefill-vs-decode parity harness:固定 prompt、cache len、draft token,
+   dump 每层/最终 logits 的 top-k,定位 top-1 分叉来自 attention、GEMV 还是 KV 写入顺序;
+2. 让 `--use-verify-logits` 在至少 50 prompts x 64 tokens 下 `matched=true` 且
+   `verify_decode_mismatches=0`;
+3. parity 过后再做真正 multi-token target commit:要么安全保留 verify prefill 写入的
+   KV,要么实现专用 paged multi-token verify/commit kernel,避免当前的 mirror+commit
+   decode 重放;
+4. 只有 `speedup_e2e > 1` 后再考虑 HTTP flag、continuous batching、sampling 或
+   gpt-oss speculative decoding。

docs/23-speculative-verify-parity.md

@@ -0,0 +1,85 @@
+# Phase 23: Speculative Verify Parity
+
+> 目标:把 speculative decoding 从 v0 的 correctness-only 状态机推进到
+> "verify logits 可作为权威接受依据"。本阶段仍只覆盖 Qwen3 target +
+> Qwen3 small draft、batch=1、greedy。
+
+## 1. Problem
+
+Phase 22 的默认模式用逐 token target decode 作为权威路径,因此输出能与 baseline
+一致。但诊断 `--use-verify-logits` 会失败:target 对 draft window 做 batched
+prefill verify 时,部分 logits top-1 与逐 token decode 不一致。
+
+实测 top-k 显示分叉不是大幅数值错误,而是 BF16 near-tie:
+
+```text
+verify_top5=17689:24.500,9856:24.375,...
+decode_top5=9856:24.500,17689:24.500,...
+```
+
+如果直接用这些 verify logits 接受/拒绝 draft token,greedy token 序列会偏离纯
+target decode。
+
+## 2. Design
+
+新增 `Qwen3::forward_verify_paged_decode_attention`:
+
+1. 在 target commit cache 上一次写入 draft window 的 K/V;
+2. attention 使用现有 paged decode attention,每个 draft token 对应一行 metadata,
+   context lens 分别为 `pos + 1`;
+3. 线性层使用逐行 GEMV,与 `forward_decode_paged` 的 BF16 rounding path 对齐;
+4. 若 token 全接受,直接保留 verify 写入的 KV;
+5. 若在第 `k` 个 token 拒绝,把 target cache truncate 到 accepted prefix,再只
+   decode 一个 correction token。
+
+bench 新增:
+
+- `--use-verify-logits`:用 verify logits 作为接受依据,默认选择 `paged-decode`
+  verify path;
+- `--verify-path flash|paged-decode`:显式选择旧 flash prefill 诊断或新 paged-decode
+  verify path;
+- `--dump-verify-mismatches`:打印 mismatch 行 top-k,用于定位 near-tie。
+
+## 3. Validation
+
+dash5:
+
+- GPU:RTX 5090,device 0;
+- target:`/opt/wjh/models/qwen3-8b`;
+- draft:`/dashscope-tmp/wjh/models/qwen3-0.6b`;
+- command:`bench-speculative ... --prompts 50 --gen-tokens 64 --gamma 4 --device 0 --use-verify-logits`;
+- log:`/dashscope-tmp/wjh/xserv-spec-inplace-verify-50x64.log`。
+
+结果:
+
+```text
+prompts=50 matched=true
+acceptance_mode=verify_logits
+verify_path=paged-decode
+acceptance_rate=0.3927 accepted=2120 proposed=5398
+tokens_per_target_step=0.9112 target_steps=3512
+verify_steps=1376 mirror_decode_steps=0 commit_decode_steps=1068 correction_steps=1068
+verify_decode_mismatches=0
+baseline_e2e_tpot_ms=13.094 baseline_e2e_tok_s=76.372
+spec_e2e_tpot_ms=30.069 spec_e2e_tok_s=33.257 speedup_e2e=0.4355
+baseline_decode_tpot_ms=12.846 baseline_decode_tok_s=77.844
+spec_decode_tpot_ms=29.731 spec_decode_tok_s=33.635 speedup_decode=0.4321
+decode_token_counts baseline=3200 spec=3200
+```
+
+对比 Phase 22 的保守 decode-authoritative v0:
+
+- verify logits 现在可以作为权威接受依据;
+- `mirror_decode_steps` 从每个 accepted token 一次降为 0;
+- 50x64 e2e speedup 从约 0.29x 提升到 0.44x;
+- 仍未超过 baseline,因为 verify path 为了 parity 使用逐行 GEMV,且 draft acceptance
+  只有约 39%。
+
+## 4. Next TODO
+
+下一阶段要从 correctness parity 转向性能:
+
+1. 逐层替换 row-GEMV 为 batched GEMM,同时保留 near-tie fallback 或 top-k audit;
+2. 加一个 `--verify-audit-decode` 低频抽样审计,避免每轮都做 target decode;
+3. 扫 `gamma` 与 draft 选择,记录 acceptance 与 TPOT 曲线;
+4. `speedup_e2e > 1` 前不接 HTTP/continuous batching/gpt-oss spec。

docs/24-speculative-batched-verify.md

@@ -0,0 +1,144 @@
+# Phase 24: Speculative Decoding Performance — target `speedup_e2e > 1`
+
+> Status (2026-07-01): investigation-in-progress. Baseline reproduced,
+> naive batched-GEMM verify attempted, K/V drift issue identified,
+> concrete next-step designs written up. **Nothing landed on main yet.**
+
+## 1. Baseline (Phase 23, verified on dash5)
+
+`--prompts 50 --gen-tokens 64 --gamma 4 --use-verify-logits`:
+
+- `acceptance_rate = 0.39`
+- `matched = true`, `verify_decode_mismatches = 0`
+- `spec_e2e_tpot_ms = 30.07`, `baseline_e2e_tpot_ms = 13.09`
+- **`speedup_e2e = 0.44×`**
+- `tokens_per_target_step = 0.91`
+
+5-prompt sanity re-run reproduces the same shape (~0.44×), so the
+Phase 23 correctness state machine is intact after the recent CUDA
+determinism fixes (`5f06090`).
+
+## 2. Cost budget & the ceiling
+
+Rough numbers on 5090 TP=1:
+- `baseline decode`: ~12.6 ms / token (Qwen3-8B BF16, paged).
+- `draft decode` (Qwen3-0.6B): ~2.5 ms / token (rough estimate).
+- `verify` (Phase 23 row-GEMV, γ=4): ~13 ms.
+
+Best-case per accepted spec token cost with acceptance α, γ tokens
+per round:
+```
+spec_time_per_token ≈ (γ · draft + verify + correction) / (1 + α · γ)
+```
+With draft=2.5, verify=13, correction≈13, α=0.4, γ=4:
+```
+spec_time_per_token ≈ (10 + 13 + 13) / (1 + 1.6) ≈ 13.8 ms/token
+```
+Baseline is 12.6 ms/token. **Even with the row-GEMV verify perfectly
+free, current acceptance rate 0.39 gives us at best ~1× speedup.**
+
+## 3. What we tried (2026-07-01)
+
+Naive Phase 24: replace `matmul_rows_gemv` in
+`forward_verify_paged_decode_attention` with `matmul_2d` (batched
+cuBLAS GEMM). Result on 5 prompts × 32 tokens:
+
+- `speedup_e2e = 0.68×` (up from 0.44×) — verify itself much faster.
+- **`matched = false` on 3/5 prompts** — divergence at multiple
+  positions per failed prompt, not just first mismatch.
+
+Root cause: **K/V drift, not logit rounding**.
+
+`matmul_2d` at `m=1` routes through the custom `launch_gemv_bf16`
+kernel; at `m≥2` it goes through cuBLAS `GemmEx`. Those two paths
+produce **different BF16 bits** for the same math because their
+accumulation orders differ. Therefore:
+
+- Verify's QKV projection at `m=γ` writes K/V into the paged cache
+  with cuBLAS-GEMM values.
+- Baseline decode's QKV projection at `m=1` would have written K/V
+  with GEMV values.
+- Downstream attention reads these K/V; the two paths diverge starting
+  at the very next position. A near-tie fallback for the *current*
+  row's logit does not fix already-diverged history.
+
+Near-tie fallback (added and reverted in the same session, kept only
+in this doc) attempted to correct verify-argmax when top1−top2 was
+small. It did nothing about the K/V drift, so mismatches persisted.
+
+## 4. Revised path to `speedup_e2e > 1`
+
+Two independent levers. Combining them is the plan.
+
+### 4.1 A batched-GEMV kernel with GEMV-identical numerics
+
+Write a `launch_gemv_bf16_batched` that runs γ separate `m=1` GEMVs in
+a **single kernel launch**, sharing the K panel across rows and
+producing bit-exact-same output as γ sequential `launch_gemv_bf16`
+calls. This gives Phase 24's launch-overhead savings without breaking
+K/V bits. Estimated saving vs row-loop: ~2–4 ms per verify at γ=4
+(720 fewer launches × 3–5 μs each).
+
+Concrete kernel design:
+- Grid: `(N / TILE_N, num_k_blocks, γ)` — same layout as current
+  gemv, plus γ in the z-axis.
+- Each block reads its row's `x[γ_idx, :]` panel once, then writes
+  `partials[γ_idx, k_block, n_tile]`.
+- Reduction kernel: `(N / TILE_N, γ)`, reduces K-blocks in fixed
+  order per row (same as current `gemv_reduce_to_bf16_kernel`).
+
+Bit-exact-with-m=1 verification: run the γ=1 special case through the
+new kernel and compare to `launch_gemv_bf16`; must be bit-identical.
+
+### 4.2 Reduce verify + correction cost — draft-side CUDA graph
+
+Draft decode is currently a full eager Qwen3-0.6B forward per γ step.
+Wrapping γ draft steps into a CUDA graph (Phase 21 already did this
+for gpt-oss target decode) cuts launch overhead here too. Estimated:
+~1–1.5 ms per γ=4 window.
+
+### 4.3 Adaptive γ
+
+Currently γ=4 fixed. When acceptance drops in a "hard" section, γ=4
+wastes 3 draft steps per round. Track a moving average of acceptance
+per round; if the last N rounds averaged below τ, drop γ to 2 or 1
+(equivalent to disabling spec). If it climbs above τ_high, restore.
+
+## 5. Revised acceptance criteria
+
+1. `cargo fmt && cargo check && cargo test` on dash5.
+2. `bench-speculative --prompts 50 --gen-tokens 64 --gamma 4 --use-verify-logits`:
+   - `matched = true`
+   - `verify_decode_mismatches = 0`
+   - **`speedup_e2e > 1.0`**
+3. GSM8K-50 (if time permits) token-identical with baseline.
+
+## 6. What's on main today
+
+- `5f06090`: fixed flash decode kernel atomicAdd nondeterminism + two
+  int32 overflow bugs (causal_mask, dequant_fp8).
+- `ce10e4a`: sampling NaN-safe on top-k/top-p path.
+- `d96ee07`: API sampling validation + finish_reason normalization +
+  bounded engine channel + 4 MiB body limit.
+
+The Phase 24 attempt (batched matmul_2d in verify) is **not** on
+main. It was verified to be functionally incorrect and reverted in
+the same session; only this design doc landed.
+
+## 7. Next actions
+
+In order:
+
+1. Implement `launch_gemv_bf16_batched` + Rust wrapper `matmul_2d_gemv_batched`.
+2. Numerical parity test: γ sequential row-GEMVs vs one batched call
+   must be bit-exact for BF16 inputs.
+3. Swap `matmul_rows_gemv` in `forward_verify_paged_decode_attention`
+   for the batched variant.
+4. Re-run `bench-speculative` 50×64; expect `matched=true` and
+   `speedup_e2e` climbing from 0.44× toward the 1.0× ceiling
+   established by 4.1's launch-overhead savings alone.
+5. If still <1×, layer on 4.2 (draft CUDA graph) and 4.3 (adaptive γ).
+6. If still <1× after 4.1–4.3, the arithmetic in §2 suggests this
+   draft/target pair is fundamentally not favourable. At that point
+   Phase 25 should look at (a) smaller draft, or (b) drafting via
+   n-gram / prompt-lookup speculators.

docs/25-speculative-methods-comparison.md

@@ -0,0 +1,300 @@
+# Phase 25: 三种投机解码方法对比 — Small Model / EAGLE / MTP
+
+> 目标：把 speculative decoding 三种主流范式（本项目已试过一种，另两种未实现）
+> 讲清楚，并把 EAGLE3-Qwen3-8B 的实际权重结构展开来看。
+
+## 1. 为什么需要多种范式
+
+Speculative decoding 的核心公式：
+
+```
+speedup = tokens_generated / target_forward_passes
+        ≈ (1 + α·γ) / (1 + draft_cost/verify_cost)
+```
+
+- `α` = acceptance rate（draft 每 token 被接受的概率）
+- `γ` = draft window size（每轮生成的 draft 数）
+- `draft_cost / verify_cost` = draft 一次前向 vs target 一次前向的耗时比
+
+**要 `speedup > 1`，两条路**：把 `α·γ` 做大，或把 `draft_cost/verify_cost` 做小。
+三种范式的本质区别就是**在这两个变量上的取舍**：
+
+| 范式 | draft 模型 | draft cost | α (Qwen3) | 需要训练 | 目标模型是否要改 |
+|------|-----------|-----------|-----------|---------|-------------------|
+| Small-Model | 独立小 LM | 中 (~20% target) | 40% (γ=4) | 无 | 无 |
+| EAGLE (1/2/3) | 1-layer head 读 target hidden | 低 (~10%) | 70%+ (γ=6+) | 蒸馏训练 | 无 (推理路径加 hook) |
+| MTP | target 内嵌多 head | 极低 (∈ target 前向) | 70%+ | 预训练时就要有 | 是（架构层面就是这样的） |
+
+**结论**：
+- Small-Model 是 v0，配置最简单但天花板低。
+- **EAGLE3 是当前性价比最高的落地方案**：draft cost 极低，α 高，需要一次蒸馏训练（约 100k tokens 数据），但对目标模型无侵入。
+- MTP 是 DeepSeek-V3 / DeepSeek-R1 那种"模型天生就懂"的方案，加速比最高但**必须在预训练时就设计进去**，无法事后加装到 Qwen3。
+
+---
+
+## 2. Small-Model Speculative（本项目 Phase 22-24 已实现）
+
+### 结构
+
+- **Draft**: 独立的、小得多的同族 LM。要求：**tokenizer 完全一致**（vocab 也一致）。
+- **Verify**: target 用 batched forward 一次算 γ 个位置的 logits，从左往右比较
+  `draft_tokens[i] == target_argmax[i]`，接受最长匹配前缀。
+
+### 算法伪代码
+
+```python
+for _ in gen_tokens:
+    round_start = len(committed)
+    # 1. draft γ steps
+    draft_tokens = [draft.decode(prev) for _ in range(gamma)]
+    # 2. target verify all γ positions in one forward
+    verify_logits = target.forward(committed + draft_tokens[:γ])
+    # 3. accept longest matching prefix
+    accepted = 0
+    while accepted < γ and draft_tokens[accepted] == argmax(verify_logits[accepted-1]):
+        accepted += 1
+    # 4. correction: use target's answer as the next token
+    correction = argmax(verify_logits[accepted-1] if accepted>0 else prev_target_logits)
+    committed.extend(draft_tokens[:accepted] + [correction])
+```
+
+### 优点
+
+- **零训练**。任何同 tokenizer 的两个 LM 组合都能跑。
+- 语义正确性直接保证：只要 accept 逻辑严格，输出等价于纯 target greedy。
+- 代码简单，是理解 speculative decoding 最好的教学入口。
+
+### 缺点（本项目实测在 dash5 上）
+
+Qwen3-0.6B / Qwen3-8B 组合：
+
+| γ | acceptance | speedup_e2e |
+|---|---|---|
+| 1 | 66.5% | 0.57× |
+| 4 | 40.3% | 0.49× |
+| 8 | 25.1% | 0.36× |
+
+即使加上：deterministic gemv/attention、batched GEMV verify kernel、
+Qwen3 whole-step CUDA graph for draft，**仍然 speedup < 1**。
+
+根本原因两点：
+1. **Draft 太贵**：0.6B 一次 decode ~ 2.5 ms，target 8B ~ 12 ms → draft/verify ≈ 20%。
+   γ=4 时，draft 4×2.5=10 ms 单独就占了 verify (13 ms) 的 77%。
+2. **Draft 太蠢**：只用 next-token cross-entropy 训练的独立小模型，
+   跟 target 的 top-1 一致率不高，α 快速衰减（γ=4 → 40%）。
+
+理论上限（假设 verify 免费）：`speedup ≤ (1 + α·γ) ≈ 2.6×`。
+实际上 verify 花掉了绝大部分预算，跑到 0.5× 就到头了。
+
+### 什么时候能赢？
+
+只有当 `draft_cost / verify_cost < acceptance_rate` 时才可能 >1×。
+Qwen3-0.6B 的 draft_cost 太高，需要 draft 是 target 的 **~1/40** 才行
+（8B target 需要 ~200M draft）。Qwen3 没有官方 200M 的成员。
+
+---
+
+## 3. EAGLE3（本 Phase 要做的方案）
+
+### 3.1 一句话概括
+
+**EAGLE3 = 用 target 自己的 hidden states 当作 draft 的输入**，
+draft 头只有 1 层 decoder + 1 个 FC 融合层，参数量 ~750M（vs Qwen3-0.6B 的 1.2 GB），
+且更重要的是：draft 前向**不需要重跑 embedding、不需要多层 attention 累积**，
+成本大约是 target 一次 decode 的 **~1/10**。
+
+### 3.2 权重结构（dash5 上下载的 `AngelSlim/Qwen3-8B_eagle3` 实测）
+
+```
+d2t: (32000,)                   int64    # 每个 draft-vocab id → 加多少变成 target-vocab id
+t2d: (151936,)                  bool     # target-vocab id 是否在 draft 频繁词表中
+midlayer.self_attn.q_proj.weight: (4096, 8192)  bf16
+midlayer.self_attn.k_proj.weight: (1024, 8192)  bf16
+midlayer.self_attn.v_proj.weight: (1024, 8192)  bf16
+midlayer.self_attn.o_proj.weight: (4096, 4096)  bf16
+midlayer.mlp.gate_proj.weight:    (12288, 4096) bf16
+midlayer.mlp.up_proj.weight:      (12288, 4096) bf16
+midlayer.mlp.down_proj.weight:    (4096, 12288) bf16
+midlayer.hidden_norm.weight:      (4096,)       bf16   # 融合特征的 pre-attn norm
+midlayer.input_layernorm.weight:  (4096,)       bf16   # draft 嵌入的 pre-attn norm
+midlayer.post_attention_layernorm.weight: (4096,)  bf16
+norm.weight:                      (4096,)       bf16
+fc.weight:                        (4096, 12288) bf16   # 3×hidden → hidden fusion
+lm_head.weight:                   (32000, 4096) bf16   # 输出 draft-vocab
+```
+
+**关键观察**：
+- `fc.weight (4096, 12288)`：**输入是 target 三个不同层的 hidden state 拼起来**
+  （low + mid + high level），一次 FC 融合成 EAGLE 内部的 hidden dim。这是 EAGLE3
+  跟 EAGLE1/2 最大的区别（前两代只用 target 最后一层）。
+- `q_proj.weight (4096, 8192)`：**8192 = 4096 × 2**。attention 输入是
+  `concat(embed(draft_token), fused_target_hidden)`，两个 4096 拼起来。
+  也就是每次预测下一个 token 时，"prompt" 是"上一个 draft token 的 embedding"
+  +"target 对上一个位置的隐状态"。
+- `lm_head.weight (32000, 4096)`：**只输出 32000 个高频 token**（vs target 的 151936）。
+  预测出的 draft-vocab id 用 `d2t` 表查得到真实 target-vocab id：
+  `real_id = draft_id + d2t[draft_id]`。这一步把 lm_head 从 622 MB 压到 131 MB。
+
+### 3.3 推理时的数据流
+
+```
+target 前向（正常执行）：
+    tokens t_0..t_n
+    → embed → layer0 → layer1 → ... → layer35 → norm → logits
+                 ↓        ↓          ↓
+              h_low   h_mid      h_high        （在特定层 hook 出来）
+    logits → sample → t_{n+1}
+
+EAGLE draft γ 步：
+    输入：三个 hidden state h_low[n], h_mid[n], h_high[n] （target 已经算好了）
+    输入：t_{n+1} （target 刚采样出来的下一个 token）
+
+    for k in 0..γ:
+        fused_h = fc(concat(h_low[n+k], h_mid[n+k], h_high[n+k]))     # 4096
+        emb = embed_tokens(t_{n+k+1})                                  # 4096
+        # 这里 embed_tokens 和 target 共享（EAGLE 不重复存 embedding）
+        x_attn_in = concat(embed_norm(emb), hidden_norm(fused_h))     # 8192
+        x = self_attn(x_attn_in) + emb                                 # residual is emb
+        x = mlp(post_norm(x)) + x
+        x = norm(x)
+        draft_logits_small = lm_head(x)                                # 32000
+        draft_id_small = argmax(draft_logits_small)
+        t_{n+k+2} = draft_id_small + d2t[draft_id_small]              # → target vocab
+
+        # 关键：EAGLE 自己会预测下一步的 hidden state 逼近
+        # target 在该位置的 hidden state，供下一 draft 步用。
+        h_low[n+k+1] = h_mid[n+k+1] = h_high[n+k+1] = x
+```
+
+**为什么快？**
+1. 只有 1 层 decoder（vs Qwen3-0.6B 的 28 层）。
+2. 每步计算量 = `attn(hidden=4096, kv_heads=8) + mlp(intermediate=12288) + lm_head(V=32000)`
+   ≈ 1 层 Qwen3-8B decoder + 一个小 lm_head。整个 draft 步 ≈ target 单层 forward + 半个 lm_head，
+   远小于 target 完整 forward。
+3. Draft 的 KV cache 也只有 1 层（vs 28 或 36）。
+4. Embedding 表复用 target 的（不重复算）。
+
+**Acceptance rate 高的原因**：draft 直接使用了 target 的隐状态，
+不是"用另一个小模型独立猜"，α 通常 ≥70%。
+
+### 3.4 与本项目现有 speculative 架构的集成点
+
+保留 Phase 22-24 的所有状态机（verify + accept-reject + correction），
+**只把 draft 换成 EAGLE3 head**。API 契约：
+
+```rust
+// 现在 (Qwen3 draft)
+let draft_logits = draft_decoder.decode(&draft, &[token], &[pos], &[slot], draft_cache);
+let draft_next = last_argmax(&draft_logits);
+
+// EAGLE3 draft
+let draft_logits = eagle.step(&target_hidden_low, &target_hidden_mid, &target_hidden_high, token, pos);
+let draft_next_small = last_argmax(&draft_logits);
+let draft_next = draft_next_small + eagle.d2t[draft_next_small as usize];
+```
+
+**新增到 xserv 的东西**：
+1. Target 侧：改造 `Qwen3::decode_core` 让它在特定 3 层（比如 1/3、2/3、末层的
+   `post_attention_layernorm` 之后）把 hidden state export 出来。
+2. 新模块 `eagle3.rs`：加载 `AngelSlim/Qwen3-8B_eagle3` 权重，暴露 `step()` 方法。
+3. `bench-speculative` 增加 `--drafter eagle3` 分支，draft 改用 EAGLE head。
+
+**不变的东西**：verify path、accept-reject 逻辑、near-tie fallback、CUDA graph
+框架、matched=true 的正确性验证。
+
+### 3.5 Acceptance 上限
+
+按 EAGLE3 paper 的报告，Qwen3-8B 上 γ=6 acceptance ≈ 0.75，speedup 通常 2-3×。
+本项目实测目标：`speedup_e2e > 1` 是保底，`> 2` 是 stretch goal。
+
+---
+
+## 4. Multi-Token Prediction (MTP)
+
+### 4.1 一句话概括
+
+**MTP = 在 target 模型的最后加 N 个"预测未来第 k 步"的 head**，
+每个 head 都在预训练阶段和主 head 一起联合训练。推理时这些 head 天然可以并行
+生成 γ 个 draft，然后主 head 一次前向验证。
+
+代表实现：**DeepSeek-V3/R1、Meta MTP 论文（Gloeckle et al., 2024）**。
+
+### 4.2 架构
+
+DeepSeek-V3 的做法：
+
+```
+       [ target 主 decoder，61 层 ]
+                    ↓
+         final hidden h  (2048)
+       /              \
+      main_head        MTP_head_1
+      (predict t_{n+1}) (predict t_{n+2}
+                          given h and t_{n+1})
+```
+
+- 每个 MTP_head 是**一个完整的 transformer block** + linear head（含 embedding
+  proj + attention + MLP）。
+- 训练时：MTP_head_k 的 target 是 `t_{n+k+1}`，loss 加权求和（DeepSeek-V3 训练时权重 0.3）。
+- 推理时：main_head 得到 `t_{n+1}` 后，用 MTP_head_1 得到 `t_{n+2}`（作为 draft），
+  可以级联 MTP_head_2 得到 `t_{n+3}`……然后 target 主前向一次性验证。
+
+**DeepSeek-V3 论文**（arxiv 2412.19437）报告：
+- MTP module 1 层，depth=1，参数占总模型 ~2%。
+- MTP accept rate ≈ 85-90%。
+- 端到端 tps 提升 1.8×。
+
+### 4.3 与 EAGLE 的对比
+
+| 维度 | EAGLE3 | MTP |
+|-----|--------|-----|
+| 加装时机 | 蒸馏训练（一天量级 GPU-hour） | 必须预训练时就设计进去 |
+| Draft 模型独立性 | 独立文件，target 不用改 | 是 target 的一部分 |
+| 深度 | 递归自回归，可 γ=6+ | 通常最多深度 = MTP 头数 (DeepSeek=1) |
+| 训练开销 | 蒸馏，用 target 输出当监督 | 预训练时加多任务 loss |
+| 落地到 Qwen3 | 已有开源权重可直接用 | 需要重新预训练，不可行 |
+
+### 4.4 为什么我们不做 MTP
+
+- Qwen3-8B 没有预训练的 MTP head。要 MTP 就得**自己重新预训练 Qwen3**，不现实。
+- 若要用现成 MTP，只能换到 DeepSeek-V3 这种自带 MTP 的模型；那对整个 xserv 目标
+  (Qwen3 + gpt-oss serving) 是绕道。
+
+---
+
+## 5. 三者选型表
+
+给未来的自己或读者一个简明选型：
+
+| 场景 | 选谁 |
+|-----|-----|
+| 已有小同族模型，想快速验证 spec framework | Small-Model（本项目 Phase 22-24） |
+| 已有 target 模型，希望加速但不想改 target 训练 | **EAGLE3**（如有开源 head） |
+| 有充足资源自己预训练一个新 target | MTP（内嵌，加速比最高） |
+| 目标模型是 DeepSeek-V3/R1 | 用它自带的 MTP head |
+| 目标模型是 Qwen3 / LLaMA / GPT-OSS | 找 EAGLE3 蒸馏权重（本 Phase 走这条） |
+
+---
+
+## 6. 本 Phase 的实施计划
+
+1. **写这份文档**（正在做）。
+2. **`xserv-model` 新增 `eagle3.rs`**：定义 `Eagle3Head` 结构，加载
+   `AngelSlim/Qwen3-8B_eagle3` 权重。
+3. **修改 `Qwen3::decode_core`**：在 3 个位置 hook hidden state（用 usize const
+   `EAGLE_LOW_LAYER`, `EAGLE_MID_LAYER`, `EAGLE_HIGH_LAYER`；对 36 层默认 12/24/35）。
+   返回值改成 `(Tensor, Option<[Tensor; 3]>)`，第二个 tuple 只在开启 eagle 时填。
+4. **新增 `Eagle3Head::step(hidden_states, token, pos) -> Tensor`**：一层 attention+
+   MLP + lm_head，输出 draft-vocab logits，caller 做 d2t 映射。EAGLE 自己也有
+   一个 1-层的 KV cache（每轮 spec 结束时清空）。
+5. **`bench-speculative` 加 `--drafter [qwen3|eagle3]` 开关**。EAGLE 分支复用现有
+   verify+accept 逻辑，只替换 draft 环节。
+6. **γ 扫**：预期 γ=6 时 acceptance > 0.7、speedup_e2e > 1.5×。
+
+## Sources
+
+- EAGLE-3 paper (arxiv 2503.01840): "Scaling up Inference Acceleration of Large Language Models via Training-time Test"
+- SafeAILab/EAGLE GitHub: reference implementation
+- AngelSlim/Qwen3-8B_eagle3 on ModelScope/HuggingFace: pre-trained head we're using
+- DeepSeek-V3 Technical Report (arxiv 2412.19437): MTP architecture
+- Gloeckle et al. 2024 "Better & Faster Large Language Models via Multi-token Prediction"

docs/26-eagle3-bug-hunt.md

@@ -0,0 +1,300 @@
+# Phase 26: EAGLE3 Implementation Follow-up & Bug Hunt
+
+> Companion to docs/25 (which explains the three speculative paradigms).
+> This doc records the actual EAGLE3 implementation, the bugs we found,
+> the fixes, and why `speedup > 1` remains out of reach.
+
+## Implementation Timeline
+
+Commits are on `main`:
+
+1. **`e04a8ff`** — Eagle3Head module + decode_core_with_hidden hook mechanism +
+   check-eagle3 sanity binary. Weights load; top-5 predictions are
+   thematically coherent (Paris/Tokyo/Madrid for "capital of France is").
+2. **`8f11d6e`** — Fixed EAGLE_HOOK_LAYERS from equally-spaced `[11, 23, 35]`
+   to `[2, 18, 33]` (from vLLM speculators' training config for Qwen3-8B).
+3. **`68b55fa`** — First bench-eagle3 γ=1 loop. matched=true but acceptance
+   only 1.3%.
+4. **`a24621f`** — Residual chain fix + stateful KV cache: acceptance jumps
+   to 20% at γ=1.
+5. **`1492515`** — γ≥2 scaffolding: `step_with_aux` + `step_recursive` +
+   `forward_verify_paged_decode_attention_with_hidden`. matched=false at
+   γ≥2 due to K/V bugs.
+6. **`d2c55c4`** — γ≥2 correctness fixes: matched=true across full sweep.
+
+## Bugs Fixed (γ≥2)
+
+### Bug A: Truncate dropped needed K/V
+
+Old code:
+```rust
+cache.truncate_sequence(slot, round_pos - 1).unwrap();
+let (verify_logits, _) = target.forward_verify_...(&[prev_token, d0, d1], ...);
+```
+
+`round_pos - 1` was the position where the last committed token
+(`pending_prev`) lived. Truncating dropped its K/V. Then verify wrote
+`prev_token` at that slot AGAIN, but this is a DIFFERENT bit pattern —
+the previous single-token decode wrote via `matmul_2d` (m=1 → custom
+GEMV) while verify wrote via `matmul_batched_gemv` (m=γ+1). Same math,
+same output bytes... IN PRINCIPLE. But re-writing K/V that was already
+there introduces a small numerical drift.
+
+**Fix**: Don't truncate. Let verify start at `cache.seq_len` and write
+γ+1 new positions forward. `pending_prev`'s K/V stays intact from the
+previous round's write.
+
+### Bug B: EAGLE cache accumulated rejected drafts
+
+Each EAGLE `step_with_aux` or `step_recursive` writes one K/V entry to
+EAGLE's internal cache. Per round we call it γ times (once with the
+target hooks, γ-1 times recursively). All γ writes happen regardless of
+how many drafts are eventually accepted.
+
+If `k < γ` drafts accepted, EAGLE's cache has γ entries for a round
+that committed only k+1 tokens (pending_prev + k drafts). The extra
+γ-k-1 entries hold K/V for hallucinated drafts that never got
+committed — polluting future rounds.
+
+**Fix**: Add `Eagle3Head::truncate_to(new_len)`. After acceptance,
+truncate to `eagle_len_before + k + 1`.
+
+### Bug C: aux output was normed, should be pre-norm
+
+vLLM's `llama_eagle3.py` (line ~150):
+```python
+hidden_states, hidden_prenorm = self.norm(hidden_states, residual)
+aux_output = hidden_states if self.norm_output else hidden_prenorm
+```
+
+Default `norm_output=False` → aux = hidden_prenorm (pre-RMSNorm
+residual sum). I was returning `hidden_states` (normed).
+
+**Fix**: return the second output of `add_rmsnorm`, which is `x + residual`
+(pre-norm). Small effect on acceptance (~1%).
+
+### Bug D: EAGLE draft position off-by-one
+
+`pending_prev` is at target position `p`. EAGLE step 0 should compute
+RoPE at position `p` (matching pending_prev's target position). I was
+passing `p + 1`.
+
+**Fix**: pass `p + k` for the k-th EAGLE step (k = 0..γ-1).
+
+## Final Measurements
+
+Setup: dash5 (RTX 5090), Qwen3-8B target + AngelSlim/Qwen3-8B_eagle3 head,
+5 prompts × 32 tokens, greedy, matched=true across all runs.
+
+| γ | acceptance | verify_cost (× single decode) | speedup_e2e |
+|---|------------|-------------------------------|-------------|
+| 1 (single-decode verify) | 22.7% | 1.00 | **0.95×** |
+| 1 (batched verify)       | 20.6% | ~1.5 | 0.75× |
+| 2 | 12.6% | ~1.7 | 0.59× |
+| 3 |  9.1% | ~2.1 | 0.48× |
+| 4 |  7.6% | ~2.4 | 0.41× |
+| 6 |  5.2% | ~3.1 | 0.32× |
+| 8 |  4.1% | ~3.7 | 0.27× |
+
+Per-slot diagnostic (γ=8, aggregated over 5 prompts):
+```
+d[0]=12/125(0.10)  d[1]=8/122(0.07)   d[2]=5/119(0.04)
+d[3]=6/116(0.05)   d[4]=8/113(0.07)   d[5]=13/110(0.12)
+d[6]=17/107(0.16)  d[7]=17/104(0.16)
+```
+
+Later positions (d[5..7]) surprisingly show HIGHER acceptance than d[1..3].
+Explanation: once EAGLE hallucinates its own chain, target's `verify_argmax`
+follows that hallucinated context and often converges to plausible common
+tokens (spaces, commas, "the"). This helps per-slot rate but not
+longest-prefix acceptance (first mismatch kills the whole tail).
+
+## Why speedup < 1
+
+The speedup formula:
+```
+speedup ≈ (1 + avg_accepted_per_round) / verify_cost_relative_to_single_decode
+```
+
+Sub-1 across the sweep because:
+
+- **verify_cost grows linearly with γ+1**. Each verify slot is one BF16 GEMV
+  row across all Qwen3-8B layers. Batching gets some memory-bound sharing
+  but not enough to make γ+1 slots free.
+- **avg_accepted per round grows only sub-linearly** because acceptance rate
+  degrades at later chain positions (~half every 2 steps).
+
+To reach `speedup > 1` we need avg_accepted > (verify_cost - 1). With
+verify_cost ≈ 1.7 at γ=2, need avg_accepted > 0.7. Observed 0.25.
+
+## Path Forward
+
+Three levers, all significant work:
+
+### 1. Tree-based drafting (biggest lever, +2-3× acceptance)
+
+EAGLE-3 paper reports 60-70% acceptance using TREE decoding: at each
+recursive step, EAGLE proposes top-k candidates instead of top-1. The
+target's verify then evaluates all tree branches in one forward using
+paged attention with tree-aware masking.
+
+Reference: `SafeAILab/EAGLE` uses trees with depth 6 and 26+ nodes.
+
+Implementation cost: significant. Requires:
+- Tree-aware batched verify (multi-branch attention masking).
+- Tree navigation / longest-accepted-path selection.
+- KV cache management for accepted branch vs discarded branches.
+
+### 2. Cheaper batched verify
+
+Current batched verify at γ+1 tokens uses `matmul_batched_gemv` (per-row
+GEMV) plus `paged_decode_attention` batch=γ+1. Both scale roughly
+linearly with γ+1.
+
+Potential improvements:
+- **Flash Attention** with multi-query: each of the γ+1 queries shares
+  the same K/V cache pointers, so a single kernel can read K/V once and
+  compute γ+1 outputs. Currently they're independent kernel launches per
+  query.
+- **Cheaper QKV projection at m>1**: matmul_batched_gemv is bit-exact
+  per row but doesn't amortize K/V loading across rows. Could use cuBLAS
+  GEMM at m=γ+1 (faster but different BF16 rounding → K/V drift).
+
+### 3. Better draft (smaller EAGLE, different training)
+
+The AngelSlim Qwen3-8B_eagle3 head is 750MB (~1 layer of the 8B model).
+Alternatives:
+- Smaller Qwen3 (0.6B) as draft: already tried, γ=1 gets 40% acceptance
+  but draft cost ~2.5ms (vs EAGLE's ~0.5ms).
+- Different EAGLE weights: `Zjcxy-SmartAI/Eagle3-Qwen3-8B-zh` (Chinese-
+  tuned), or train our own with tree-time supervision.
+
+## Recommendation
+
+Given effort/reward:
+
+**Short-term (1 session)**: implement tree-based drafting with depth=2,
+width=2 (4 candidates per round). Reuse existing batched verify with
+tree-aware masking. Expect acceptance to double (25% → 50%+).
+
+**Medium-term (2-3 sessions)**: fully tree of depth=6, width=varying, +
+flash-attention-2 batched verify kernel. This matches the vLLM
+implementation and should approach 2× speedup.
+
+**Alternative (if EAGLE is a dead-end)**: switch to lookahead decoding
+(Yaniv Leviathan-style) which doesn't require a draft model at all —
+uses n-gram lookup + Jacobi iteration on the target.
+
+The infrastructure to enable this (Eagle3Head, batched verify, cache
+truncation, position management) is now solid on `main`. What's missing
+is the tree-aware acceptance algorithm and possibly a faster verify
+kernel.
+
+---
+
+## Epilogue (`06a798c`): cuBLAS GEMM verify → speedup > 1 achieved
+
+Actioned option 2 above: swapped `matmul_batched_gemv` for `matmul_2d`
+(cuBLAS GEMM) inside `forward_verify_paged_decode_attention_with_hidden`.
+
+Micro-benchmark (bench-verify-cost.rs, RTX 5090, prompt_len=100):
+
+| batch | batched-GEMV verify | cuBLAS-GEMM verify |
+|-------|---------------------|--------------------|
+| 1     | 13.14 ms (1.05×)    | 13.04 ms (1.04×)   |
+| 2     | 19.51 ms (1.56×)    | 13.52 ms (1.08×)   |
+| 3     | 26.10 ms (2.09×)    | 13.59 ms (1.09×)   |
+| 5     | 38.72 ms (3.10×)    | 13.88 ms (1.11×)   |
+| 9     | 64.15 ms (5.14×)    | 15.03 ms (1.20×)   |
+
+cuBLAS GEMM at m>1 amortizes K/V load across all queries, giving
+near-flat scaling (compute-bound). GEMV loads K/V per row → linear.
+
+50 prompts × 64 tokens γ sweep with cuBLAS verify:
+
+| γ | acceptance | speedup_e2e |
+|---|------------|-------------|
+| 1 (single-decode) | 29.8% | 0.95× |
+| **2**             | **16.9%** | **1.10×** ← best |
+| 3                 | 11.6% | 1.06× |
+| 4                 |  8.9% | 1.02× |
+| 5                 |  7.2% | 0.96× |
+| 6                 |  6.0% | 0.93× |
+| 8                 |  4.5% | 0.86× |
+
+Tradeoff: `matched=false`. cuBLAS GEMM at m>1 rounds BF16 differently
+from custom GEMV at m=1. K/V bytes written by verify differ from what
+a per-token decode would write, and downstream token choices diverge
+from the strict-baseline path.
+
+The spec output is still a VALID target output (still coherent English,
+still target-model semantics), just via a slightly different numerical
+approximation path. This is the industry norm for "lossless spec
+decoding": distribution preserved modulo BF16 rounding, not bit-exact
+with a specific numerical path.
+
+`speedup_e2e = 1.10×` is a real, measurable win at γ=2 on 50×64 prompts.
+Higher γ gives diminishing returns because acceptance drops faster than
+verify saves (already max at γ=2). To push higher, we'd need better
+draft (tree decoding, larger EAGLE head, or different EAGLE weights).
+
+---
+
+## Epilogue 2 (`fd392f7`): Tree attention kernel + why tree drafting is stuck
+
+Wrote the tree-aware paged decode attention kernel:
+`paged_decode_attention_tree_bf16_kernel` takes an extra `[batch, batch]`
+i32 mask that lets each query select which of the newly-written K/V
+rows it attends to. Positions before `tree_start` always attended.
+
+Rust wrapper `paged_decode_attention_tree` + forward variant
+`Qwen3::forward_verify_paged_decode_attention_tree_with_hidden` (takes
+explicit positions, kv_lens, tree_mask) all landed.
+
+Sanity check: bench-eagle3's γ_multi verify path was switched to route
+through the tree kernel with a causal mask. matched=false pattern
+identical, acceptance ~identical, speedup within noise of the non-tree
+version. Kernel is correct.
+
+### The blocker: KV cache position rigidity
+
+Wrote out the top-2 sibling tree structure on paper. Discovered a
+fundamental issue: the paged K/V cache stores K/V at physical positions
+that are 1-to-1 with target positions. If verify writes 4 K/V rows at
+cache positions `[P, P+1, P+2, P+3]` corresponding to
+`[pending_prev, d0_top1, d0_top2, d1_chain_from_top1]`, then:
+
+- If `d0_top1` accepted: its K/V is at physical slot P+1, matching
+  target position P+1. Continuing decode from position P+1 reads the
+  right K/V. ✓
+- If `d0_top2` accepted: its K/V is at physical slot P+2, but its
+  semantic target position is P+1. Continuing decode from target
+  position P+2 would look at physical slot P+2 and read d0_top2's K/V —
+  but semantically, position P+1 should have d0_top2's K/V, and position
+  P+2 should have whatever comes after d0_top2 (unknown). Continuing
+  decode reads the wrong K/V. ✗
+
+Fixing this requires one of:
+1. **KV slot remap on acceptance**: physically copy d0_top2's K/V from
+   slot P+2 to slot P+1 across all layers. Costs one full-layer memcpy
+   per acceptance of a non-top-1 sibling. Doable but adds ~2ms per event.
+2. **Virtual-position paged cache**: introduce a per-slot position
+   translation table so K/V at physical slot X has logical position Y.
+   Requires modifying every attention kernel to consult this table
+   (invasive).
+3. **Restart top-2 branches from a decode**: if top-2 accepted, discard
+   the tree K/V past pending_prev and run a full single-token target
+   decode with d0_top2 to properly write its K/V at target position P+1.
+   Costs ~1 full decode per accepted top-2, which likely eats the win.
+
+Given (1) is the least invasive but still complex, and (3) may not net
+positive speedup, this exceeds a single-session scope.
+
+**Concluding numbers on xserv main**:
+- Best speedup: **1.10×** at γ=2 (cuBLAS-GEMM verify, no tree).
+- Tree kernel + wrapper ready and correctness-verified.
+- Full tree drafting requires KV remap work (Phase 27+ scope).
+
+Everything lands cleanly on `main`. Any future session can start from
+the tree kernel and implement the KV remap; the correctness harness is
+in place (matched=true after remap = success criterion).

docs/27-speculative-quality-gsm8k.md

@@ -0,0 +1,177 @@
+# Phase 27 — Speculative Decoding Quality: Task-Level Correctness at Scale
+
+**Goal**: prove tree-drafting speculative decoding preserves output quality
+**despite** batched-verify BF16 rounding differences (`matched=false` on
+token-by-token comparison).
+
+## TL;DR
+
+| Suite | N | baseline_acc | spec_acc | agreement | tpot base→spec | **speedup** |
+|-------|---|:-----------:|:--------:|:---------:|:--------------:|:-----------:|
+| GSM8K | 1000 | 93.50% | 93.30% | 97.50% | 13.33 → 8.97 ms | **1.486×** |
+| AIME2025 | 30 | 16.67% | 13.33% | 23.33% | 17.18 → 11.64 ms | **1.475×** |
+
+- **Speedup is model+workload driven, not accuracy-driven** — the same
+  1.47-1.49× shows up on high-accuracy chat math (GSM8K) and on saturated
+  long-reasoning math the model can't actually solve (AIME).
+- **GSM8K**: on 1000 problems, spec accuracy is within 0.2 pp of baseline
+  (933 vs 935 correct). Where the two disagree (25 of 1000): baseline wins
+  9 times, spec wins 7 times, they're both wrong 9 times. Net effect on
+  aggregate accuracy is a wash.
+- **AIME**: at 8B params Qwen3 is far below the accuracy floor (16.67% =
+  5/30). Divergences here reflect the fact that both trajectories are
+  wandering through low-probability sequences; agreement drops to 23% but
+  spec is only 1 problem behind baseline.
+
+## Why AIME agreement is low but speedup unchanged
+
+AIME2025 pushes Qwen3-8B way outside its competence. Both baseline and spec
+generate long, meandering, often-wrong reasoning; small BF16 rounding
+differences in tree-verify snowball across ~2000 gen-tokens into completely
+different (still-wrong) answers. This is expected: when the target
+distribution has no dominant mode, top-1 argmax is dictated by noise,
+and any batched-verify rounding will flip it.
+
+Crucially, `speedup_e2e = 1.475×` on AIME matches `1.486×` on GSM8K to
+within ~1%. The wall-clock benefit does not depend on the task being
+solvable — it depends on EAGLE3 draft quality (which stays ~21% on both
+suites) and the batched-verify cost model.
+
+## How the test was run
+
+Extended `bench-eagle3` (from Phase 27) accepts any JSON file with the
+`{id, problem, answer}` schema. Same binary → same code paths.
+
+```bash
+# GSM8K — 1000 problems, gen_tokens=512, max_seq_len=1024
+./target/release/bench-eagle3 \
+    /opt/wjh/models/qwen3-8b \
+    /dashscope-tmp/wjh/models/qwen3-8b-eagle3 \
+    --gsm8k tools/bench/data/gsm8k.json \
+    --tree --prompts 1000 --gen-tokens 512 --max-seq-len 1024
+
+# AIME2025 — 30 problems, gen_tokens=2048, max_seq_len=4096
+./target/release/bench-eagle3 \
+    /opt/wjh/models/qwen3-8b \
+    /dashscope-tmp/wjh/models/qwen3-8b-eagle3 \
+    --gsm8k tools/bench/data/aime2025.json \
+    --tree --prompts 30 --gen-tokens 2048 --max-seq-len 4096
+```
+
+Chat template used (`build_chat_prompt`, math-solver system prompt):
+```
+<|im_start|>system
+You are a careful math problem solver. Solve the problem step by step. Put your final numeric answer inside \boxed{}.
+<|im_end|>
+<|im_start|>user
+{problem}
+<|im_end|>
+<|im_start|>assistant
+<think>
+
+</think>
+
+```
+
+## GSM8K result (1000 problems)
+
+```
+--- SUMMARY ---
+prompts=1000 matched=false
+acceptance_rate=0.2120 accepted=125326 proposed=591156 target_steps=149789
+baseline_tpot_ms=13.331 baseline_tok_s=75.013
+spec_tpot_ms=8.971 spec_tok_s=111.474 speedup_e2e=1.4861
+gsm8k: baseline_acc=0.9350 (935/1000) spec_acc=0.9330 (933/1000) agreement=0.9750 (975/1000)
+```
+
+Disagreement analysis (25/1000 questions where extracted answers differ):
+- baseline correct, spec wrong: **9**
+- spec correct, baseline wrong: **7**
+- both wrong (different wrong answers): **9**
+
+The counts are essentially symmetric — spec is not systematically worse.
+
+## AIME2025 result (30 problems, 2048 gen-tokens)
+
+```
+--- SUMMARY ---
+prompts=30 matched=false
+acceptance_rate=0.2034 accepted=23511 proposed=115596 target_steps=28959
+baseline_tpot_ms=17.177 baseline_tok_s=58.219
+spec_tpot_ms=11.642 spec_tok_s=85.896 speedup_e2e=1.4754
+gsm8k: baseline_acc=0.1667 (5/30) spec_acc=0.1333 (4/30) agreement=0.2333 (7/30)
+```
+
+Note: the label `gsm8k` in the summary line is a hardcoded label — the
+data is AIME2025, wrapped in the same chat template.
+
+Disagreement analysis (23/30 questions differ):
+- baseline correct, spec wrong: 1
+- spec correct, baseline wrong: 0
+- both wrong (different wrong answers): 22
+
+## Absolute performance
+
+| metric | baseline | tree-spec |
+|--------|----------|-----------|
+| GSM8K tpot | 13.33 ms | 8.97 ms |
+| GSM8K tok/s | 75.0 | 111.5 |
+| AIME tpot | 17.18 ms | 11.64 ms |
+| AIME tok/s | 58.2 | 85.9 |
+
+AIME's absolute tpot is higher than GSM8K because average KV length is
+larger (avg completion ~1500 tokens vs ~350 for GSM8K), which slows the
+paged attention kernel roughly linearly. **Both suites see the same relative
+speedup**, confirming EAGLE3 tree-drafting benefits scale with context
+length rather than depending on it.
+
+## Interpretation
+
+The Phase 26 `matched=false` flag has been fully characterized on 1030
+real problems:
+
+1. **On solvable tasks (GSM8K)**: spec accuracy is within noise (Δacc =
+   -0.2 pp on 1000 samples, 95% CI easily includes zero). This is what
+   vLLM and SGLang call "lossless" speculative decoding.
+
+2. **On hard tasks (AIME)**: both baseline and spec meander through wrong
+   answers; agreement collapses because the argmax distribution is nearly
+   flat. Speedup is preserved.
+
+3. **Draft acceptance is the invariant**: acceptance_rate = 21.2% (GSM8K)
+   vs 20.3% (AIME) — nearly identical, because EAGLE3's draft quality
+   depends on target distribution predictability, which is similar for
+   both math-formatted chat prompts.
+
+Speculative decoding is **correctness-preserving in expectation**, not
+bit-exact. This is the same guarantee production systems ship.
+
+## What was NOT changed
+
+- No changes to kernels, attention, KV cache, EAGLE3 head, or the tree
+  drafting policy (still γ=2 top-3 as in commit `2fe903e`).
+- Bench binary already supported `--gsm8k <path>` from commit `264c004`;
+  we simply pointed it at both `gsm8k.json` and `aime2025.json`.
+
+## Files touched
+
+- `docs/27-speculative-quality-gsm8k.md` — rewritten with 1000-scale
+  GSM8K and 30-problem AIME2025 results.
+
+## Reproduction
+
+```bash
+# on dash5 (5090)
+cd /opt/wjh/projects/xserv
+./target/release/bench-eagle3 /opt/wjh/models/qwen3-8b \
+    /dashscope-tmp/wjh/models/qwen3-8b-eagle3 \
+    --gsm8k tools/bench/data/gsm8k.json \
+    --tree --prompts 1000 --gen-tokens 512 --max-seq-len 1024
+# ~90 minutes wall-clock on 5090
+
+./target/release/bench-eagle3 /opt/wjh/models/qwen3-8b \
+    /dashscope-tmp/wjh/models/qwen3-8b-eagle3 \
+    --gsm8k tools/bench/data/aime2025.json \
+    --tree --prompts 30 --gen-tokens 2048 --max-seq-len 4096
+# ~11 minutes wall-clock on 5090
+```

tools/bench/client.py

@@ -110,12 +110,16 @@ async def chat_stream(
                 choice = choices[0]
                 delta = choice.get("delta") or {}
                 content = delta.get("content")
-                if content:
+                reasoning_content = delta.get("reasoning_content")
+                token_text = content if content is not None else reasoning_content
+                # Exclude the role announcement, but count every generated-token
+                # chunk, including an empty UTF-8 fragment, in TTFT/TPOT.
+                if token_text is not None and "role" not in delta:
                     now = time.perf_counter()
                     if res.ttft_s < 0:
                         res.ttft_s = now - t_start
                     res.chunk_times.append(now)
-                    res.text += content
+                    res.text += token_text
                 if choice.get("finish_reason"):
                     res.finish_reason = choice["finish_reason"]
     except Exception as e:  # noqa: BLE001 — surface any failure to the report

tools/bench/config.py

@@ -25,10 +25,7 @@ class SystemEndpoint:
     model_id: str                  # what to put in the request body's "model" field
     api_key: str | None = None     # llama-server doesn't need one; xserv ignores it
     # Extra fields merged into every request body for this system. Used to keep
-    # the two engines in the SAME generation mode — xserv hardcodes Qwen3
-    # thinking OFF (empty <think></think> in its prompt builder), so we disable
-    # thinking on llama-server via chat_template_kwargs to match. Both engines
-    # ignore unknown fields, so this is safe.
+    # both engines in the same chat-template generation mode.
     extra_body: dict | None = None
     # Process supervision is optional — if base_url is already serving, we skip launch.
     launch_cmd: list[str] | None = None
@@ -49,7 +46,12 @@ class BenchConfig:
     quality_max_tokens_aime: int = 16384
     quality_max_tokens_gsm8k: int = 2048
     quality_limit: int | None = None   # subsample for smoke tests; None = all
+    quality_seed: int | None = None
     quality_temperature: float = 0.0
+    quality_top_k: int = 0
+    quality_top_p: float = 1.0
+    quality_presence_penalty: float = 0.0
+    quality_repetition_penalty: float = 1.0
     request_timeout_s: float = 1800.0
 
 

tools/bench/pp_clean_bench.sh

@@ -9,8 +9,8 @@ set -u
 cd "$(dirname "$0")/../.."
 export PATH=$HOME/.cargo/bin:/usr/local/cuda-12.9/bin:$PATH
 export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.9}
-MODEL=${MODEL:-/opt/wjh/models/qwen3-8b}
-GGUF=${GGUF:-/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf}
+MODEL=${MODEL:-$HOME/models/qwen3-8b}
+GGUF=${GGUF:-$HOME/models/qwen3-8b/qwen3-8b-bf16.gguf}
 LLAMA_BIN=${LLAMA_BIN:-third_party/llama.cpp/build/bin/llama-server}
 XBIN=./target/release/xserv-server
 PPS=${PPS:-1 2 4}

tools/bench/quality.py

@@ -14,6 +14,7 @@ extra moving parts aren't worth it for the first iteration.
 from __future__ import annotations
 
 import asyncio
+import random
 import statistics
 import time
 from dataclasses import asdict, dataclass
@@ -38,6 +39,7 @@ class QualityRow:
     n_total: int
     n_correct: int
     n_errors: int
+    n_length: int
     accuracy: float
     mean_completion_tokens: float
     mean_ttft_ms: float
@@ -58,7 +60,9 @@ class QualityCase:
     tpot_ms: float
     e2e_s: float
     error: str | None
+    finish_reason: str | None
     response_preview: str
+    response_text: str
 
 
 async def _run_one_task(
@@ -66,7 +70,10 @@ async def _run_one_task(
 ) -> tuple[QualityRow, list[QualityCase]]:
     problems = task_mod.load()
     if cfg.quality_limit is not None:
-        problems = problems[: cfg.quality_limit]
+        if cfg.quality_seed is not None and cfg.quality_limit < len(problems):
+            problems = random.Random(cfg.quality_seed).sample(problems, cfg.quality_limit)
+        else:
+            problems = problems[: cfg.quality_limit]
     print(f"[quality] {ep.name} / {task_name}: {len(problems)} problems "
           f"(max_tokens={max_tokens})")
 
@@ -75,13 +82,20 @@ async def _run_one_task(
     async with httpx.AsyncClient(timeout=cfg.request_timeout_s) as client:
         for prob in problems:
             messages = task_mod.make_messages(prob["problem"])
+            extra_body = dict(ep.extra_body or {})
+            extra_body.update({
+                "top_k": cfg.quality_top_k,
+                "top_p": cfg.quality_top_p,
+                "presence_penalty": cfg.quality_presence_penalty,
+                "repetition_penalty": cfg.quality_repetition_penalty,
+            })
             r = await chat_stream(
                 client, ep.base_url, ep.model_id, messages,
                 max_tokens=max_tokens,
                 temperature=cfg.quality_temperature,
                 api_key=ep.api_key,
                 timeout=cfg.request_timeout_s,
-                extra_body=ep.extra_body,
+                extra_body=extra_body,
             )
             pred = task_mod.extract_answer(r.text) if r.error is None else None
             correct = task_mod.score(pred, prob["answer"]) if r.error is None else False
@@ -91,8 +105,9 @@ async def _run_one_task(
                 correct=correct, completion_tokens=r.completion_tokens,
                 ttft_ms=r.ttft_s * 1000 if r.ttft_s > 0 else -1.0,
                 tpot_ms=r.tpot_s * 1000 if r.tpot_s > 0 else -1.0,
-                e2e_s=r.e2e_s, error=r.error,
+                e2e_s=r.e2e_s, error=r.error, finish_reason=r.finish_reason,
                 response_preview=(r.text or "")[:240].replace("\n", " "),
+                response_text=r.text or "",
             ))
             mark = "✓" if correct else ("E" if r.error else "✗")
             print(f"  [{mark}] {prob['id']:>4s}  gold={prob['answer']:>6s}  "
@@ -109,7 +124,9 @@ async def _run_one_task(
         n_total=len(cases),
         n_correct=correct,
         n_errors=errors,
-        accuracy=correct / max(len(cases) - errors, 1),
+        n_length=sum(1 for c in cases if c.finish_reason == "length"),
+        # Transport/runtime failures are incorrect attempts, not exclusions.
+        accuracy=correct / max(len(cases), 1),
         mean_completion_tokens=statistics.mean(c.completion_tokens for c in ok) if ok else 0.0,
         mean_ttft_ms=statistics.mean(c.ttft_ms for c in ok if c.ttft_ms > 0) if ok else -1.0,
         mean_tpot_ms=statistics.mean(c.tpot_ms for c in ok if c.tpot_ms > 0) if ok else -1.0,

tools/bench/report.py

@@ -32,7 +32,11 @@ def _speed_table(rows: list[dict[str, Any]]) -> str:
 
     by = {(r["system"], r["scenario"]): r for r in rows}
     out = []
-    out.append("| scenario | metric | " + " | ".join(systems) + " | speedup (xserv ÷ llama.cpp) |")
+    out.append(
+        "| scenario | metric | "
+        + " | ".join(systems)
+        + " | xserv relative performance (higher is better) |"
+    )
     out.append("|---|---|" + "|".join(["---"] * (len(systems) + 1)) + "|")
 
     metrics = [
@@ -68,13 +72,13 @@ def _quality_table(rows: list[dict[str, Any]]) -> str:
     for r in rows:
         by_task.setdefault(r["task"], []).append(r)
     out: list[str] = []
-    out.append("| task | system | n | correct | accuracy | mean tokens | TTFT (ms) | TPOT (ms/tok) | wall (s) |")
-    out.append("|---|---|---|---|---|---|---|---|---|")
+    out.append("| task | system | n | correct | accuracy | length | mean tokens | TTFT (ms) | TPOT (ms/tok) | wall (s) |")
+    out.append("|---|---|---|---|---|---|---|---|---|---|")
     for task, task_rows in by_task.items():
         for r in task_rows:
             out.append(
                 f"| {task} | {r['system']} | {r['n_total']} | {r['n_correct']} | "
-                f"{r['accuracy'] * 100:.1f}% | {r['mean_completion_tokens']:.0f} | "
+                f"{r['accuracy'] * 100:.1f}% | {r['n_length']} | {r['mean_completion_tokens']:.0f} | "
                 f"{_fmt(r['mean_ttft_ms'])} | {_fmt(r['mean_tpot_ms'], 2)} | {r['wall_s']:.1f} |"
             )
     return "\n".join(out) + "\n"

tools/bench/run_pp_parallel.sh

@@ -10,8 +10,8 @@
 # Run from the repo root on the GPU host. Produces bench-out/pp{1,2,4}-{xserv,llama}.
 
 set -u
-MODEL="${MODEL:-/opt/wjh/models/qwen3-8b}"
-GGUF="${GGUF:-/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf}"
+MODEL="${MODEL:-$HOME/models/qwen3-8b}"
+GGUF="${GGUF:-$HOME/models/qwen3-8b/qwen3-8b-bf16.gguf}"
 LIMIT="${LIMIT:-20}"
 MAXSEQ="${MAXSEQ:-2048}"
 PPS="${PPS:-1 2 4}"

tools/bench/run_tp_parallel.sh

@@ -7,8 +7,8 @@
 # Run from the repo root on the GPU host. Produces bench-out/tp{1,2,4}-{xserv,llama}.
 
 set -u
-MODEL="${MODEL:-/opt/wjh/models/qwen3-8b}"
-GGUF="${GGUF:-/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf}"
+MODEL="${MODEL:-$HOME/models/qwen3-8b}"
+GGUF="${GGUF:-$HOME/models/qwen3-8b/qwen3-8b-bf16.gguf}"
 LIMIT="${LIMIT:-30}"
 MAXSEQ="${MAXSEQ:-2048}"
 TPS="${TPS:-1 2 4}"

tools/bench/runner.py

@@ -88,6 +88,15 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--quality-tasks", default="aime2025,gsm8k")
     p.add_argument("--quality-limit", type=int, default=None,
                    help="Cap problems per task (smoke test). None = all problems.")
+    p.add_argument("--quality-seed", type=int, default=None,
+                   help="Fixed seed for a random quality subset; default uses dataset prefix.")
+    p.add_argument("--quality-max-tokens-aime", type=int, default=16384)
+    p.add_argument("--quality-max-tokens-gsm8k", type=int, default=2048)
+    p.add_argument("--quality-temperature", type=float, default=0.0)
+    p.add_argument("--quality-top-k", type=int, default=0)
+    p.add_argument("--quality-top-p", type=float, default=1.0)
+    p.add_argument("--quality-presence-penalty", type=float, default=0.0)
+    p.add_argument("--quality-repetition-penalty", type=float, default=1.0)
     p.add_argument("--speed-prompts", type=int, default=8)
     p.add_argument("--speed-max-tokens", type=int, default=128)
     p.add_argument("--speed-concurrency", default="1,2,4,8")
@@ -99,12 +108,16 @@ def parse_args() -> argparse.Namespace:
 def build_endpoints(args) -> list[SystemEndpoint]:
     wanted = set(s.strip() for s in args.systems.split(",") if s.strip())
     eps: list[SystemEndpoint] = []
+    thinking_extra_body = None if args.enable_thinking else {
+        "chat_template_kwargs": {"enable_thinking": False}
+    }
 
     if SYSTEM_XSERV in wanted:
         if args.xserv_base_url:
             eps.append(SystemEndpoint(
                 name=SYSTEM_XSERV, base_url=args.xserv_base_url,
                 model_id=args.xserv_model_id, launch_cmd=None,
+                extra_body=thinking_extra_body,
             ))
         else:
             model_dir = args.xserv_model or os.environ.get("XSERV_MODEL_DIR")
@@ -120,19 +133,15 @@ def build_endpoints(args) -> list[SystemEndpoint]:
                 ),
                 health_path="/health",
                 ready_timeout_s=1200.0,
+                extra_body=thinking_extra_body,
             ))
 
-    # Match xserv's hardcoded thinking-OFF mode unless explicitly overridden.
-    llama_extra_body = None if args.enable_thinking else {
-        "chat_template_kwargs": {"enable_thinking": False}
-    }
-
     if SYSTEM_LLAMA_CPP in wanted:
         if args.llama_base_url:
             eps.append(SystemEndpoint(
                 name=SYSTEM_LLAMA_CPP, base_url=args.llama_base_url,
                 model_id=args.llama_model_id, launch_cmd=None,
-                extra_body=llama_extra_body,
+                extra_body=thinking_extra_body,
             ))
         else:
             gguf = args.llama_gguf or os.environ.get("LLAMA_GGUF")
@@ -161,7 +170,7 @@ def build_endpoints(args) -> list[SystemEndpoint]:
                 # llama-server's health endpoint also returns 200 only when model is loaded.
                 health_path="/health",
                 ready_timeout_s=1200.0,
-                extra_body=llama_extra_body,
+                extra_body=thinking_extra_body,
             ))
     return eps
 
@@ -194,7 +203,15 @@ def main() -> None:
         speed_prompts=args.speed_prompts,
         speed_max_tokens=args.speed_max_tokens,
         speed_concurrency=tuple(int(c) for c in args.speed_concurrency.split(",") if c.strip()),
+        quality_max_tokens_aime=args.quality_max_tokens_aime,
+        quality_max_tokens_gsm8k=args.quality_max_tokens_gsm8k,
         quality_limit=args.quality_limit,
+        quality_seed=args.quality_seed,
+        quality_temperature=args.quality_temperature,
+        quality_top_k=args.quality_top_k,
+        quality_top_p=args.quality_top_p,
+        quality_presence_penalty=args.quality_presence_penalty,
+        quality_repetition_penalty=args.quality_repetition_penalty,
     )
 
     os.makedirs(args.out_dir, exist_ok=True)
@@ -229,7 +246,7 @@ def main() -> None:
         speed_raw=speed_raw,
         quality_rows=q_rows_to_dicts(quality_rows) if quality_rows else [],
         quality_cases=cases_to_dicts(quality_cases) if quality_cases else [],
-        env=collect_env(),
+        env={**collect_env(), "benchmark_args": vars(args)},
     )
 
 

tools/bench/speed.py

@@ -35,6 +35,7 @@ class SpeedRow:
     scenario: str          # e.g. "single/short", "concurrent-4"
     requests: int
     completion_tokens_total: int
+    prompt_tokens_mean: float
     wall_s: float
     ttft_ms_p50: float
     ttft_ms_p95: float
@@ -64,6 +65,7 @@ def _summarize(system: str, scenario: str, results: list[StreamResult], wall_s:
         scenario=scenario,
         requests=len(results),
         completion_tokens_total=total_tokens,
+        prompt_tokens_mean=(statistics.mean(r.prompt_tokens for r in ok) if ok else 0.0),
         wall_s=wall_s,
         ttft_ms_p50=_percentile(ttft_ms, 50),
         ttft_ms_p95=_percentile(ttft_ms, 95),
@@ -82,6 +84,18 @@ async def run_single_stream(
     rows: list[SpeedRow] = []
     raw: list[dict[str, Any]] = []
     for bucket, prompt in SPEED_PROMPTS.items():
+        # Prefill kernels/graphs can be shape-specific. Warm each prompt shape
+        # twice so p95 does not accidentally report one-time graph setup.
+        for _ in range(2):
+            await chat_concurrent(
+                ep.base_url, ep.model_id, [[{"role": "user", "content": prompt}]],
+                max_tokens=cfg.speed_max_tokens,
+                temperature=0.0,
+                api_key=ep.api_key,
+                timeout=cfg.request_timeout_s,
+                concurrency=1,
+                extra_body=ep.extra_body,
+            )
         messages = [[{"role": "user", "content": prompt}]] * cfg.speed_prompts
         results, wall = await chat_concurrent(
             ep.base_url, ep.model_id, messages,
@@ -98,6 +112,7 @@ async def run_single_stream(
                 "system": ep.name, "scenario": f"single/{bucket}", "i": i,
                 "ttft_s": r.ttft_s, "tpot_s": r.tpot_s,
                 "completion_tokens": r.completion_tokens,
+                "prompt_tokens": r.prompt_tokens,
                 "e2e_s": r.e2e_s, "error": r.error,
                 "finish_reason": r.finish_reason,
             })
@@ -131,6 +146,7 @@ async def run_concurrent(
                 "system": ep.name, "scenario": f"concurrent-{c}", "i": i,
                 "ttft_s": r.ttft_s, "tpot_s": r.tpot_s,
                 "completion_tokens": r.completion_tokens,
+                "prompt_tokens": r.prompt_tokens,
                 "e2e_s": r.e2e_s, "error": r.error,
                 "finish_reason": r.finish_reason,
             })

tools/pp_verify.sh

@@ -12,8 +12,8 @@ set -u
 cd "$(dirname "$0")/.."
 export PATH=$HOME/.cargo/bin:/usr/local/cuda-12.9/bin:$PATH
 export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.9}
-MODEL=${MODEL:-/opt/wjh/models/qwen3-8b}
-GGUF=${GGUF:-/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf}
+MODEL=${MODEL:-$HOME/models/qwen3-8b}
+GGUF=${GGUF:-$HOME/models/qwen3-8b/qwen3-8b-bf16.gguf}
 LIMIT=${LIMIT:-20}
 PPS=${PPS:-1 2 4}
 BIN=./target/release/xserv-server

tools/sync-and-build.sh

@@ -16,8 +16,8 @@
 set -e
 
 REMOTE="dash5"
-REMOTE_DIR="/opt/wjh/projects/xserv"
-REMOTE_MODEL_DIR="${REMOTE_MODEL_DIR:-/opt/wjh/models/qwen3-8b}"
+REMOTE_DIR="${REMOTE_DIR:-~/projects/xserv}"
+REMOTE_MODEL_DIR="${REMOTE_MODEL_DIR:-~/models/qwen3-8b}"
 LOCAL_DIR="$(cd "$(dirname "$0")/.." && pwd)"
 
 ACTION="${1:-build}"