Add Qwen3.6 MoE inference support

2026-07-13 20:24:41 +08:00
parent 588bfd9df3
commit a2de146fb6
27 changed files with 3153 additions and 149 deletions
--- a/README.md
+++ b/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
--- a/crates/xserv-kernels/build.rs
+++ b/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")
--- a/crates/xserv-kernels/src/activation.rs
+++ b/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);
--- a/crates/xserv-kernels/src/attention.rs
+++ b/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());
--- a/crates/xserv-kernels/src/deltanet.rs
+++ b/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
+}
--- a/crates/xserv-kernels/src/lib.rs
+++ b/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,
--- a/crates/xserv-kernels/src/moe.rs
+++ b/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
--- a/crates/xserv-kernels/src/rope.rs
+++ b/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.
--- a/crates/xserv-model/src/bin/xserv-chat.rs
+++ b/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
--- a/crates/xserv-model/src/bin/xserv-cli.rs
+++ b/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 {
--- a/crates/xserv-model/src/config.rs
+++ b/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
    }
 }
--- a/crates/xserv-model/src/decode_graph.rs
+++ b/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");
--- a/crates/xserv-model/src/lib.rs
+++ b/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).
--- a/crates/xserv-model/src/loader.rs
+++ b/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);
    }

--- a/crates/xserv-model/src/qwen35_moe.rs
+++ b/crates/xserv-model/src/qwen35_moe.rs
--- a/crates/xserv-model/src/sampling.rs
+++ b/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);
+    }
+}
--- a/crates/xserv-server/src/api.rs
+++ b/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.
--- a/crates/xserv-server/src/engine.rs
+++ b/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);
--- a/crates/xserv-server/src/main.rs
+++ b/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()
--- a/crates/xserv-server/src/pp_engine.rs
+++ b/crates/xserv-server/src/pp_engine.rs
@@ -16,14 +16,18 @@
 use std::ffi::c_void;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::mpsc;
 use std::thread;

 use half::bf16;
 use xserv_distributed::{PpContext, UniqueId};
 use xserv_model::loader;
-use xserv_model::sampling::SamplingParams;
-use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, sample};
+use xserv_model::sampling::{SamplingParams, sample_with_history};
+use xserv_model::{
+    BLOCK_SIZE, ModelConfig, ModelFamily, PagedKVCache, Qwen3, Qwen35Moe,
+    Qwen35RecurrentCache,
+};
 use xserv_tensor::{DType, Device, Tensor};
 use xserv_tokenizer::Tokenizer;

@@ -39,7 +43,7 @@ enum PpCommand {
    /// Receive `[n_tokens, hidden]` from the previous stage, run this stage's
    /// layers; if last stage, sample with `sampling` and return the token.
    Prefill {
-        n_tokens: usize,
+        tokens: Vec<u32>,
        slot: usize,
        sampling: SamplingParams,
    },
@@ -52,13 +56,69 @@ enum PpCommand {
 }

 struct StageCtx {
-    model: Qwen3,
+    model: PpModel,
+    qwen35_recurrent: Option<Qwen35RecurrentCache>,
    cache: PagedKVCache,
    pp: Arc<PpContext>,
    hidden: usize,
    device: u32,
 }

+enum PpModel {
+    Qwen3(Qwen3),
+    Qwen35(Qwen35Moe),
+}
+
+impl StageCtx {
+    fn reset_slot(&mut self, slot: usize) {
+        if let Some(cache) = &mut self.qwen35_recurrent {
+            cache.reset_slot(slot);
+        }
+    }
+
+    fn embed(&self, tokens: &[u32]) -> Tensor {
+        match &self.model {
+            PpModel::Qwen3(model) => model.embed(tokens),
+            PpModel::Qwen35(model) => model.embed(tokens),
+        }
+    }
+
+    fn head(&self, x: &Tensor) -> Tensor {
+        match &self.model {
+            PpModel::Qwen3(model) => model.head(x),
+            PpModel::Qwen35(model) => model.head(x),
+        }
+    }
+
+    fn forward_prefill(&mut self, x: Tensor, slot: usize) -> Tensor {
+        match &self.model {
+            PpModel::Qwen3(model) => model.forward_layers_prefill(x, slot, &mut self.cache),
+            PpModel::Qwen35(model) => model.forward_layers_prefill(
+                x,
+                slot,
+                &mut self.cache,
+                self.qwen35_recurrent
+                    .as_mut()
+                    .expect("Qwen3.6 recurrent cache"),
+            ),
+        }
+    }
+
+    fn forward_decode(&mut self, x: Tensor, slot: usize) -> Tensor {
+        match &self.model {
+            PpModel::Qwen3(model) => model.forward_layers_decode(x, &[slot], &mut self.cache),
+            PpModel::Qwen35(model) => model.forward_layers_decode(
+                x,
+                slot,
+                &mut self.cache,
+                self.qwen35_recurrent
+                    .as_mut()
+                    .expect("Qwen3.6 recurrent cache"),
+            ),
+        }
+    }
+}
+
 /// Build this stage: NCCL init, load + slice weights, size a per-stage KV pool
 /// for THIS stage's layers only (so per-GPU KV is ~1/P).
 fn build_stage(
@@ -71,14 +131,40 @@ fn build_stage(
    id: UniqueId,
 ) -> StageCtx {
    let pp = Arc::new(PpContext::init(stage, world, id, device));
-    let weights = loader::load_model_dir(model_dir, Device::Cpu);
-    let model = Qwen3::from_weights_pp(config.clone(), weights, stage, world, device);
+    let model = match config.family() {
+        ModelFamily::Qwen35Moe => {
+            let weights = loader::load_model_dir_filtered(model_dir, Device::Cpu, |name| {
+                Qwen35Moe::weight_belongs_to_pp_stage(config, name, stage, world)
+            });
+            PpModel::Qwen35(Qwen35Moe::from_weights_pp(
+                config.clone(),
+                weights,
+                stage,
+                world,
+                device,
+            ))
+        }
+        _ => {
+            let weights = loader::load_model_dir(model_dir, Device::Cpu);
+            PpModel::Qwen3(Qwen3::from_weights_pp(
+                config.clone(),
+                weights,
+                stage,
+                world,
+                device,
+            ))
+        }
+    };

    // The KV cache only needs this stage's layers; build it from a config clone
    // whose layer count is the per-stage count (heads are NOT split under PP).
    let per_stage = config.num_layers() / world;
    let mut stage_config = config.clone();
-    stage_config.num_hidden_layers = Some(per_stage);
+    if let Some(text) = stage_config.text_config.as_mut() {
+        text.num_hidden_layers = Some(per_stage);
+    } else {
+        stage_config.num_hidden_layers = Some(per_stage);
+    }

    let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE);
    let total_blocks = max_blocks_per_seq + 8; // v1 serial: one active sequence
@@ -91,8 +177,13 @@ fn build_stage(
        DType::BF16,
        device,
    );
+    let qwen35_recurrent = match &model {
+        PpModel::Qwen35(model) => Some(model.new_recurrent_cache(4)),
+        PpModel::Qwen3(_) => None,
+    };
    StageCtx {
        model,
+        qwen35_recurrent,
        cache,
        pp,
        hidden: config.hidden(),
@@ -138,40 +229,51 @@ fn worker_loop(
        max_seq_len,
        id,
    );
+    let _ = ack_tx.send(());
    let is_last = stage == world - 1;
+    let mut token_history = vec![Vec::<u32>::new(); 4];
    let prev = stage - 1;
    let next = stage + 1;

    while let Ok(cmd) = cmd_rx.recv() {
        match cmd {
            PpCommand::Register(slot) => {
+                token_history[slot].clear();
+                sc.reset_slot(slot);
                let _ = sc.cache.register_sequence(slot);
                let _ = ack_tx.send(());
            }
            PpCommand::Free(slot) => {
+                token_history[slot].clear();
                sc.cache.free_sequence(slot);
                let _ = ack_tx.send(());
            }
            PpCommand::Prefill {
-                n_tokens,
+                tokens,
                slot,
                sampling,
            } => {
+                let n_tokens = tokens.len();
+                token_history[slot] = tokens;
                let x = recv_hidden(&sc, n_tokens, prev);
-                let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache);
+                let x = sc.forward_prefill(x, slot);
                if is_last {
-                    let logits = sc.model.head(&x);
-                    let _ = token_tx.send(sample(&logits, &sampling));
+                    let logits = sc.head(&x);
+                    let token = sample_with_history(&logits, &sampling, &token_history[slot]);
+                    token_history[slot].push(token);
+                    let _ = token_tx.send(token);
                } else {
                    send_hidden(&sc, &x, next);
                }
            }
            PpCommand::Decode { slot, sampling } => {
                let x = recv_hidden(&sc, 1, prev);
-                let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache);
+                let x = sc.forward_decode(x, slot);
                if is_last {
-                    let logits = sc.model.head(&x);
-                    let _ = token_tx.send(sample(&logits, &sampling));
+                    let logits = sc.head(&x);
+                    let token = sample_with_history(&logits, &sampling, &token_history[slot]);
+                    token_history[slot].push(token);
+                    let _ = token_tx.send(token);
                } else {
                    send_hidden(&sc, &x, next);
                }
@@ -191,6 +293,7 @@ pub fn run_pp(
    world: usize,
    max_seq_len: usize,
    rx: mpsc::Receiver<GenerateRequest>,
+    ready: Arc<AtomicBool>,
 ) {
    assert!(world >= 2, "run_pp requires world >= 2");
    let config = ModelConfig::from_file(&model_dir.join("config.json"));
@@ -231,6 +334,10 @@ pub fn run_pp(

    // Stage 0 (this thread): coordinator + embedding + first layers.
    let mut sc = build_stage(model_dir, &config, 0, world, 0, max_seq_len, id);
+    for _ in 1..world {
+        ack_rx.recv().expect("PP worker exited during model load");
+    }
+    ready.store(true, Ordering::Release);
    eprintln!("[pp-engine] ready (pp={world}, max_seq_len={max_seq_len})");

    let n_workers = world - 1;
@@ -249,6 +356,7 @@ pub fn run_pp(
    let slot = 0usize;
    while let Ok(req) = rx.recv() {
        broadcast(&cmd_txs, PpCommand::Register(slot));
+        sc.reset_slot(slot);
        sc.cache.register_sequence(slot).expect("register slot");
        wait_acks(&ack_rx);

@@ -256,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;
--- a/crates/xserv-server/src/tp_engine.rs
+++ b/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);
        };
--- a/csrc/activation/activations.cu
+++ b/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;
--- a/csrc/attention/deltanet.cu
+++ b/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();
+}
+
+}
--- a/csrc/attention/flash_attention.cu
+++ b/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]);
--- a/csrc/attention/paged_attention.cu
+++ b/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,
--- a/csrc/embedding/rope.cu
+++ b/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) {
--- a/csrc/moe/moe_sparse.cu
+++ b/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,