Add Qwen3.6 MoE inference support

crates/xserv-kernels/build.rs

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

crates/xserv-kernels/src/activation.rs

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

crates/xserv-kernels/src/attention.rs

@@ -412,8 +412,8 @@ pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tens
         "num_q_heads must be divisible by num_kv_heads"
     );
     assert!(
-        head_dim <= 128,
-        "flash_attention supports head_dim up to 128"
+        head_dim <= 256,
+        "flash_attention supports head_dim up to 256"
     );
 
     // Dispatch to specialized decode kernel for single-token generation
@@ -479,7 +479,7 @@ pub fn flash_attention_sinks(
     assert_eq!(k.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
     assert_eq!(v.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
     assert!(num_q_heads % num_kv_heads == 0);
-    assert!(head_dim <= 128);
+    assert!(head_dim <= 256);
     assert_eq!(
         sinks.shape()[0],
         num_q_heads,
@@ -548,7 +548,7 @@ pub fn paged_decode_attention(
         num_q_heads % num_kv_heads == 0,
         "GQA: num_q_heads must be divisible by num_kv_heads"
     );
-    assert!(head_dim <= 128);
+    assert!(head_dim <= 256);
 
     let scale = 1.0 / (head_dim as f32).sqrt();
     let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
@@ -601,7 +601,7 @@ pub fn paged_decode_attention_tree(
     assert_eq!(q.shape()[2], 1);
     assert_eq!(q.dtype(), DType::BF16);
     assert!(num_q_heads % num_kv_heads == 0);
-    assert!(head_dim <= 128);
+    assert!(head_dim <= 256);
 
     let scale = 1.0 / (head_dim as f32).sqrt();
     let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
@@ -653,7 +653,7 @@ pub fn paged_decode_attention_sinks(
     assert_eq!(q.shape()[2], 1);
     assert_eq!(q.dtype(), DType::BF16);
     assert!(num_q_heads % num_kv_heads == 0);
-    assert!(head_dim <= 128);
+    assert!(head_dim <= 256);
 
     let scale = 1.0 / (head_dim as f32).sqrt();
     let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());

crates/xserv-kernels/src/deltanet.rs

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

crates/xserv-kernels/src/lib.rs

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

crates/xserv-kernels/src/moe.rs

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

crates/xserv-kernels/src/rope.rs

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