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

@@ -412,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
@@ -479,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,
@@ -548,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());
@@ -601,7 +601,7 @@ pub fn paged_decode_attention_tree(
     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());
@@ -653,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/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,18 +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, 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, 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-eagle3.rs

@@ -83,18 +83,35 @@ fn main() {
     if args.len() < 3 {
         eprintln!(
             "Usage: bench-eagle3 <target-dir> <eagle3-dir> \
-             [--gen-tokens N] [--prompts N] [--max-seq-len N] [--device N]"
+             [--gen-tokens N] [--prompts N] [--max-seq-len N] [--device N] \
+             [--tree] [--gamma N] [--gsm8k PATH]"
         );
         std::process::exit(1);
     }
     let target_dir = PathBuf::from(&args[1]);
     let eagle_dir = PathBuf::from(&args[2]);
-    let gen_tokens = arg_usize(&args, "--gen-tokens", DEFAULT_GEN_TOKENS);
-    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 gamma = arg_usize(&args, "--gamma", 2).max(1);
     let use_tree = args.iter().any(|a| a == "--tree");
+    let gsm8k_path = arg_str(&args, "--gsm8k");
+    let (prompts_source, default_gen): (Vec<GsmProblem>, usize) = if let Some(p) = &gsm8k_path {
+        (load_gsm8k(p), 512)
+    } else {
+        (
+            PROMPTS
+                .iter()
+                .map(|s| GsmProblem {
+                    id: String::new(),
+                    problem: s.to_string(),
+                    answer: String::new(),
+                })
+                .collect(),
+            DEFAULT_GEN_TOKENS,
+        )
+    };
+    let gen_tokens = arg_usize(&args, "--gen-tokens", default_gen);
+    let prompt_count = arg_usize(&args, "--prompts", prompts_source.len()).min(prompts_source.len());
 
     xserv_cuda::device::set_device(device).unwrap();
     let info = xserv_cuda::device::device_info(device).unwrap();
@@ -137,8 +154,21 @@ fn main() {
     let mut spec_target_steps = 0usize;
     let mut mismatches = 0usize;
 
-    for (i, prompt) in PROMPTS.iter().take(prompt_count).enumerate() {
-        let ids = tokenizer.encode(prompt);
+    let is_gsm = gsm8k_path.is_some();
+    let mut baseline_correct = 0usize;
+    let mut spec_correct = 0usize;
+    let mut answer_agree = 0usize;
+    let mut both_scored = 0usize;
+    let im_end_id = tokenizer.special_token_id("<|im_end|>");
+
+    for (i, item) in prompts_source.iter().take(prompt_count).enumerate() {
+        let prompt = &item.problem;
+        let ids = if is_gsm {
+            let templated = build_chat_prompt(prompt);
+            tokenizer.encode(&templated)
+        } else {
+            tokenizer.encode(prompt)
+        };
         if ids.len() + gen_tokens >= max_seq_len {
             eprintln!("prompt {i} too long, skipping");
             continue;
@@ -195,6 +225,7 @@ fn main() {
         let ok = baseline.ids == spec.ids;
         if !ok {
             mismatches += 1;
+            if !is_gsm {
                 let common = baseline
                     .ids
                     .iter()
@@ -206,7 +237,42 @@ fn main() {
                     common, baseline.ids, spec.ids
                 );
             }
+        }
 
+        if is_gsm {
+            let gold = normalize_num(&item.answer);
+            let baseline_text = decode_until_im_end(&tokenizer, &baseline.ids, im_end_id);
+            let spec_text = decode_until_im_end(&tokenizer, &spec.ids, im_end_id);
+            let baseline_pred = extract_answer(&baseline_text);
+            let spec_pred = extract_answer(&spec_text);
+            let b_ok = gold.is_some() && baseline_pred == gold;
+            let s_ok = gold.is_some() && spec_pred == gold;
+            let agree = baseline_pred.is_some() && baseline_pred == spec_pred;
+            if b_ok {
+                baseline_correct += 1;
+            }
+            if s_ok {
+                spec_correct += 1;
+            }
+            if agree {
+                answer_agree += 1;
+            }
+            both_scored += 1;
+            println!(
+                "q={:03} id={} tok_match={} gold={} base={} spec={} b_ok={} s_ok={} agree={} base_tpot={:.2} spec_tpot={:.2}",
+                i,
+                item.id,
+                ok,
+                gold.as_deref().unwrap_or("?"),
+                baseline_pred.as_deref().unwrap_or("?"),
+                spec_pred.as_deref().unwrap_or("?"),
+                b_ok,
+                s_ok,
+                agree,
+                baseline.total_s * 1000.0 / baseline.ids.len() as f64,
+                spec.total_s * 1000.0 / spec.ids.len() as f64,
+            );
+        } else {
             println!(
                 "prompt={:02} match={} gen={} accept={}/{} target_steps={} baseline_tpot_ms={:.3} spec_tpot_ms={:.3}",
                 i,
@@ -219,6 +285,7 @@ fn main() {
                 spec.total_s * 1000.0 / spec.ids.len() as f64,
             );
         }
+    }
 
     let baseline_tpot = baseline_total_s * 1000.0 / baseline_tokens as f64;
     let spec_tpot = spec_total_s * 1000.0 / spec_tokens as f64;
@@ -240,6 +307,20 @@ fn main() {
         1000.0 / spec_tpot,
         baseline_tpot / spec_tpot
     );
+    if is_gsm && both_scored > 0 {
+        println!(
+            "gsm8k: baseline_acc={:.4} ({}/{}) spec_acc={:.4} ({}/{}) agreement={:.4} ({}/{})",
+            baseline_correct as f64 / both_scored as f64,
+            baseline_correct,
+            both_scored,
+            spec_correct as f64 / both_scored as f64,
+            spec_correct,
+            both_scored,
+            answer_agree as f64 / both_scored as f64,
+            answer_agree,
+            both_scored,
+        );
+    }
 }
 
 #[derive(Default)]
@@ -924,7 +1005,122 @@ fn arg_usize(args: &[String], flag: &str, default: usize) -> usize {
         .unwrap_or(default)
 }
 
+fn arg_str(args: &[String], flag: &str) -> Option<String> {
+    args.iter()
+        .position(|a| a == flag)
+        .and_then(|i| args.get(i + 1))
+        .cloned()
+}
+
 fn new_cache(config: &ModelConfig, max_seq_len: usize, device: u32) -> PagedKVCache {
     let num_blocks = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE + 2;
     PagedKVCache::new(config, num_blocks, 0, 16, num_blocks, DType::BF16, device)
 }
+
+struct GsmProblem {
+    id: String,
+    problem: String,
+    answer: String,
+}
+
+fn load_gsm8k(path: &str) -> Vec<GsmProblem> {
+    let text = std::fs::read_to_string(path)
+        .unwrap_or_else(|e| panic!("failed to read {path}: {e}"));
+    let v: serde_json::Value =
+        serde_json::from_str(&text).unwrap_or_else(|e| panic!("failed to parse {path}: {e}"));
+    let arr = v.as_array().expect("gsm8k json must be a top-level array");
+    arr.iter()
+        .map(|it| GsmProblem {
+            id: it.get("id").and_then(|x| x.as_str()).unwrap_or("").to_string(),
+            problem: it
+                .get("problem")
+                .and_then(|x| x.as_str())
+                .unwrap_or("")
+                .to_string(),
+            answer: it
+                .get("answer")
+                .and_then(|x| x.as_str())
+                .unwrap_or("")
+                .to_string(),
+        })
+        .collect()
+}
+
+const GSM_SYSTEM: &str = "You are a careful math problem solver. Solve the problem step by step. \
+Put your final numeric answer inside \\boxed{}.";
+
+fn build_chat_prompt(user: &str) -> String {
+    let mut s = String::new();
+    s.push_str("<|im_start|>system\n");
+    s.push_str(GSM_SYSTEM);
+    s.push_str("<|im_end|>\n");
+    s.push_str("<|im_start|>user\n");
+    s.push_str(user);
+    s.push_str("<|im_end|>\n");
+    s.push_str("<|im_start|>assistant\n");
+    // enable_thinking=false: skip the model's chain-of-thought preamble to
+    // save tokens and get straight to the answer.
+    s.push_str("<think>\n\n</think>\n\n");
+    s
+}
+
+fn decode_until_im_end(tokenizer: &Tokenizer, ids: &[u32], im_end: Option<u32>) -> String {
+    let cut = if let Some(e) = im_end {
+        ids.iter().position(|&t| t == e).unwrap_or(ids.len())
+    } else {
+        ids.len()
+    };
+    tokenizer.decode(&ids[..cut])
+}
+
+fn normalize_num(s: &str) -> Option<String> {
+    let cleaned: String = s.trim().replace(',', "");
+    let f: f64 = cleaned.parse().ok()?;
+    if f.fract() == 0.0 && f.abs() < 1e15 {
+        Some(format!("{}", f as i64))
+    } else {
+        Some(format!("{}", f))
+    }
+}
+
+fn extract_answer(text: &str) -> Option<String> {
+    // Prefer \boxed{...}
+    if let Some(idx) = text.rfind("\\boxed{") {
+        let start = idx + "\\boxed{".len();
+        let rest = &text[start..];
+        if let Some(end) = rest.find('}') {
+            let inner = &rest[..end];
+            if let Some(n) = last_number_in(inner) {
+                return normalize_num(&n);
+            }
+        }
+    }
+    // Fallback: last number anywhere in the text.
+    last_number_in(text).and_then(|n| normalize_num(&n))
+}
+
+fn last_number_in(text: &str) -> Option<String> {
+    let bytes = text.as_bytes();
+    let mut end: Option<usize> = None;
+    let mut start: Option<usize> = None;
+    let mut i = bytes.len();
+    while i > 0 {
+        i -= 1;
+        let c = bytes[i] as char;
+        let is_num = c.is_ascii_digit() || c == '.' || c == ',' || c == '-';
+        if end.is_none() {
+            if c.is_ascii_digit() {
+                end = Some(i + 1);
+                start = Some(i);
+            }
+        } else if is_num {
+            start = Some(i);
+        } else {
+            break;
+        }
+    }
+    match (start, end) {
+        (Some(s), Some(e)) => Some(text[s..e].to_string()),
+        _ => None,
+    }
+}

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,7 +1,7 @@
 use std::io::{self, Write};
 use std::path::PathBuf;
 use xserv_model::{
-    BLOCK_SIZE, KVCache, ModelConfig, PagedKVCache, SamplingParams, loader, sample,
+    BLOCK_SIZE, KVCache, ModelConfig, ModelFamily, PagedKVCache, SamplingParams, loader, sample,
     sample_greedy_penalized,
 };
 use xserv_tensor::{DType, Device};
@@ -43,6 +43,7 @@ fn main() {
         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);
@@ -56,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 {

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

@@ -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");

crates/xserv-model/src/lib.rs

@@ -8,10 +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};
@@ -19,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);
     }
 

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()
@@ -55,11 +69,6 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
         _ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
     };
 
-    // Greedy
-    if params.temperature == 0.0 {
-        return argmax(&last_row);
-    }
-
     // 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.
@@ -74,6 +83,29 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
         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);
+    }
+
     // Apply temperature
     let mut logits_f32: Vec<f32> = last_row.iter().map(|v| v / params.temperature).collect();
 
@@ -195,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();
 
@@ -363,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 {
@@ -393,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 =
@@ -422,6 +512,16 @@ fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
                     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;
@@ -464,6 +564,18 @@ fn validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
             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
 }
@@ -546,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)
@@ -558,9 +692,20 @@ 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.

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;
 
@@ -232,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);
@@ -310,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() {
@@ -329,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()
@@ -339,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);

crates/xserv-server/src/main.rs

@@ -10,8 +10,9 @@ use axum::{
 };
 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,
@@ -19,6 +20,7 @@ pub struct AppState {
     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]
@@ -77,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);
     }
@@ -109,27 +117,42 @@ async fn main() {
     // 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,
@@ -137,6 +160,7 @@ async fn main() {
         engine_sender: Mutex::new(tx),
         engine_tokenizer: Mutex::new(tokenizer),
         max_seq_len,
+        engine_ready,
     });
 
     let app = Router::new()

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 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);
-    let model = Qwen3::from_weights_pp(config.clone(), weights, stage, world, device);
+            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();
+    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,13 +364,13 @@ 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");
 
@@ -287,8 +395,8 @@ 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;

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,9 +297,9 @@ 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;
@@ -317,8 +326,8 @@ 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;
             stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
         };

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/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]);

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,

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/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,

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,6 +70,9 @@ async def _run_one_task(
 ) -> tuple[QualityRow, list[QualityCase]]:
     problems = task_mod.load()
     if cfg.quality_limit is not None:
+        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}"