model: fused GPU MoE kernel — eliminate CPU roundtrip

Replace the per-token CPU-routed MoE forward with an all-GPU path:

  1. moe_topk_softmax: GPU top-k + softmax (was CPU sort + softmax)
  2. moe_replicate: broadcast input to all local experts
  3. cublasGemmStridedBatchedEx: batched expert matmul (was per-expert cuBLAS)
  4. moe_weighted_sum: FP32-accumulated weighted sum on GPU (was GPU→CPU→F32→BF16→GPU)

Expert weights stored as contiguous 3D tensors for strided batched GEMM.
Zero CPU↔GPU transfers per MoE layer (was ~40 per token per layer).

Also: configurable geglu_alpha, LayerNorm bias auto-detect, unused-weight
diagnostic at load time.

GSM8K 30-problem: 11/30 → 23/30 (76.7%) vs llama.cpp 30/30 (100%).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

crates/xserv-kernels/src/lib.rs

@@ -5,6 +5,7 @@ pub mod dispatch;
 pub mod embedding;
 pub mod gemm;
 pub mod layernorm;
+pub mod moe;
 pub mod rmsnorm;
 pub mod rope;
 pub mod softmax;

crates/xserv-kernels/src/moe.rs

@@ -0,0 +1,223 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+use crate::gemm::{cublas_handle, CublasHandle};
+
+unsafe extern "C" {
+    fn launch_moe_topk_softmax_bf16(
+        router_logits: *const c_void,
+        topk_ids: *mut c_void, topk_weights: *mut c_void,
+        num_tokens: i32, num_experts: i32, top_k: i32,
+        stream: *mut c_void,
+    );
+    fn launch_moe_replicate_bf16(
+        x: *const c_void, x_rep: *mut c_void,
+        num_tokens: i32, hidden: i32, local_experts: i32,
+        stream: *mut c_void,
+    );
+    fn launch_moe_bias_add_3d_bf16(
+        x: *mut c_void, bias: *const c_void,
+        batch: i32, num_tokens: i32, dim: i32,
+        stream: *mut c_void,
+    );
+    fn launch_moe_weighted_sum_bf16(
+        expert_out: *const c_void,
+        topk_ids: *const c_void, topk_weights: *const c_void,
+        out: *mut c_void,
+        num_tokens: i32, hidden: i32, top_k: i32,
+        expert_start: i32, local_experts: i32,
+        stream: *mut c_void,
+    );
+
+    fn cublasGemmStridedBatchedEx(
+        handle: CublasHandle,
+        transa: i32, transb: i32,
+        m: i32, n: i32, k: i32,
+        alpha: *const c_void,
+        a: *const c_void, a_type: i32, lda: i32, stride_a: i64,
+        b: *const c_void, b_type: i32, ldb: i32, stride_b: i64,
+        beta: *const c_void,
+        c: *mut c_void, c_type: i32, ldc: i32, stride_c: i64,
+        batch_count: i32,
+        compute_type: i32,
+        algo: i32,
+    ) -> i32;
+
+    fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
+}
+
+const CUDA_R_16BF: i32 = 14;
+const CUBLAS_COMPUTE_32F: i32 = 68;
+const CUBLAS_GEMM_DEFAULT: i32 = -1;
+
+/// GPU top-k selection + softmax over router logits.
+///
+/// Input:  router_logits [num_tokens, num_experts] BF16 on GPU
+/// Output: (topk_ids [num_tokens, top_k] i32, topk_weights [num_tokens, top_k] f32)
+pub fn moe_topk_softmax(
+    router_logits: &Tensor,
+    num_experts: usize,
+    top_k: usize,
+) -> (Tensor, Tensor) {
+    assert_eq!(router_logits.ndim(), 2);
+    assert_eq!(router_logits.dtype(), DType::BF16);
+    assert!(router_logits.is_contiguous());
+    let num_tokens = router_logits.shape()[0];
+    assert_eq!(router_logits.shape()[1], num_experts);
+
+    let topk_ids = Tensor::empty(&[num_tokens, top_k], DType::F32, router_logits.device());
+    let topk_weights = Tensor::empty(&[num_tokens, top_k], DType::F32, router_logits.device());
+
+    unsafe {
+        launch_moe_topk_softmax_bf16(
+            router_logits.data_ptr() as *const c_void,
+            topk_ids.data_ptr() as *mut c_void,
+            topk_weights.data_ptr() as *mut c_void,
+            num_tokens as i32, num_experts as i32, top_k as i32,
+            std::ptr::null_mut(),
+        );
+    }
+
+    (topk_ids, topk_weights)
+}
+
+/// Replicate x [num_tokens, hidden] → [local_experts, num_tokens, hidden].
+pub fn moe_replicate(x: &Tensor, local_experts: usize) -> Tensor {
+    assert_eq!(x.ndim(), 2);
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous());
+    let num_tokens = x.shape()[0];
+    let hidden = x.shape()[1];
+    let out = Tensor::empty(&[local_experts, num_tokens, hidden], DType::BF16, x.device());
+
+    unsafe {
+        launch_moe_replicate_bf16(
+            x.data_ptr() as *const c_void,
+            out.data_ptr() as *mut c_void,
+            num_tokens as i32, hidden as i32, local_experts as i32,
+            std::ptr::null_mut(),
+        );
+    }
+
+    out
+}
+
+/// In-place 3D bias add: x [batch, num_tokens, dim] += bias [batch, dim].
+pub fn moe_bias_add_3d(x: &Tensor, bias: &Tensor) {
+    assert_eq!(x.ndim(), 3);
+    assert_eq!(bias.ndim(), 2);
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous());
+    let batch = x.shape()[0];
+    let num_tokens = x.shape()[1];
+    let dim = x.shape()[2];
+    assert_eq!(bias.shape(), &[batch, dim]);
+
+    unsafe {
+        launch_moe_bias_add_3d_bf16(
+            x.data_ptr() as *mut c_void,
+            bias.data_ptr() as *const c_void,
+            batch as i32, num_tokens as i32, dim as i32,
+            std::ptr::null_mut(),
+        );
+    }
+}
+
+/// Weighted sum of expert outputs → [num_tokens, hidden].
+///
+/// expert_out:   [local_experts, num_tokens, hidden] BF16
+/// topk_ids:     [num_tokens, top_k] i32 (global expert indices)
+/// topk_weights: [num_tokens, top_k] f32
+pub fn moe_weighted_sum(
+    expert_out: &Tensor,
+    topk_ids: &Tensor,
+    topk_weights: &Tensor,
+    expert_start: usize,
+    local_experts: usize,
+    top_k: usize,
+) -> Tensor {
+    assert_eq!(expert_out.ndim(), 3);
+    assert_eq!(expert_out.dtype(), DType::BF16);
+    let num_tokens = expert_out.shape()[1];
+    let hidden = expert_out.shape()[2];
+
+    let out = Tensor::empty(&[num_tokens, hidden], DType::BF16, expert_out.device());
+
+    unsafe {
+        launch_moe_weighted_sum_bf16(
+            expert_out.data_ptr() as *const c_void,
+            topk_ids.data_ptr() as *const c_void,
+            topk_weights.data_ptr() as *const c_void,
+            out.data_ptr() as *mut c_void,
+            num_tokens as i32, hidden as i32, top_k as i32,
+            expert_start as i32, local_experts as i32,
+            std::ptr::null_mut(),
+        );
+    }
+
+    out
+}
+
+/// Strided batched GEMM for MoE expert forward.
+/// C[b] = A[b] @ B[b]  for b in 0..batch
+///
+/// A: [batch, M, K] BF16 contiguous
+/// B: [batch, K, N] BF16 contiguous
+/// Returns C: [batch, M, N] BF16
+#[allow(clippy::too_many_arguments)]
+pub fn batched_gemm_strided(a: &Tensor, b: &Tensor) -> Tensor {
+    assert_eq!(a.ndim(), 3);
+    assert_eq!(b.ndim(), 3);
+    assert_eq!(a.dtype(), DType::BF16);
+    assert_eq!(b.dtype(), DType::BF16);
+    assert!(a.is_contiguous() && b.is_contiguous());
+    assert_eq!(a.shape()[0], b.shape()[0]);
+    assert_eq!(a.shape()[2], b.shape()[1]);
+
+    let batch = a.shape()[0];
+    let m = a.shape()[1];
+    let k = a.shape()[2];
+    let n = b.shape()[2];
+
+    let c = Tensor::empty(&[batch, m, n], DType::BF16, a.device());
+
+    let alpha: f32 = 1.0;
+    let beta: f32 = 0.0;
+
+    // cuBLAS column-major: we compute C^T = B^T @ A^T
+    // A is [batch, M, K] row-major → A^T is [K, M] col-major, lda=K
+    // B is [batch, K, N] row-major → B^T is [N, K] col-major, ldb=N? No...
+    //
+    // Actually for row-major: A[M,K] in memory = col-major A^T[K,M] with lda=K.
+    // So we call cublasGemmStridedBatchedEx with:
+    //   transa=N, transb=N
+    //   m=N, n=M, k=K  (because cuBLAS sees col-major)
+    //   A_cublas = B_row (pointer), lda=N
+    //   B_cublas = A_row (pointer), ldb=K
+    //   C_cublas = C_row (pointer), ldc=N
+
+    let stride_a = (m * k) as i64;
+    let stride_b = (k * n) as i64;
+    let stride_c = (m * n) as i64;
+
+    let handle = cublas_handle();
+    unsafe {
+        cublasSetStream_v2(handle, std::ptr::null_mut());
+        let status = cublasGemmStridedBatchedEx(
+            handle,
+            0, 0, // CUBLAS_OP_N, CUBLAS_OP_N
+            n as i32, m as i32, k as i32,
+            &alpha as *const f32 as *const c_void,
+            b.data_ptr() as *const c_void, CUDA_R_16BF, n as i32, stride_b,
+            a.data_ptr() as *const c_void, CUDA_R_16BF, k as i32, stride_a,
+            &beta as *const f32 as *const c_void,
+            c.data_ptr() as *mut c_void, CUDA_R_16BF, n as i32, stride_c,
+            batch as i32,
+            CUBLAS_COMPUTE_32F,
+            CUBLAS_GEMM_DEFAULT,
+        );
+        assert_eq!(status, 0, "cublasGemmStridedBatchedEx failed: {status}");
+    }
+
+    c
+}