Add Qwen3.6 MoE inference support

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