moe: sparse top-k decode — compute only routed experts (1.8x, beats llama TP=2)

Dense MoE replicated x across all 16 local experts and ran the full
batched GEMM, reading every expert's weights per token; the weighted
sum then discarded 12 of 16 results. Decode is memory-bound, so this
was ~8x wasted expert bytes — the entire decode gap vs llama.cpp.

New fused expert-indexed GEMVs (csrc/moe/moe_sparse.cu) read
topk_ids on-device (no host sync) and early-return block-uniformly
for experts other ranks own. FP8 runs W8A16 (activations stay BF16 —
tensor cores are irrelevant at M=1, and activation quantization error
disappears); MXFP4 runs W4A16. Per-expert bias + scale fused into the
GEMV epilogue; slot-indexed weighted sum skips (never multiplies)
unwritten non-local slots. Dense path retained for num_tokens > 8
(prefill) and via XSERV_DENSE_MOE=1 for A/B.

dash5 (RTX 5090), gpt-oss-20b FP8, TP=2: decode TPOT 13.9 -> 7.6 ms.
Warm-server vs llama.cpp MXFP4 TP=2: TPOT 7.19-7.32 vs 7.54-8.42 ms —
first config where xserv wins decode outright. GSM8K-100: 96% (dense
FP8: 91%). llama TP=1 (2.9 ms) remains ahead: next levers are decode
CUDA graphs, non-expert quantization, sparse prefill (docs/20).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

crates/xserv-kernels/src/moe.rs

@@ -29,6 +29,29 @@ unsafe extern "C" {
         stream: *mut c_void,
     );
 
+    fn launch_moe_sparse_gemv_fp8_bf16(
+        x: *const c_void, w: *const c_void, w_scales: *const c_void,
+        bias: *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_mxfp4_bf16(
+        x: *const c_void, w_packed: *const c_void, w_scales: *const c_void,
+        bias: *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_weighted_sum_sparse_bf16(
+        down: *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,
@@ -158,6 +181,110 @@ pub fn moe_weighted_sum(
     out
 }
 
+/// Sparse MoE GEMV (FP8 W8A16): compute only the routed experts.
+///
+/// x:        [num_tokens, K] BF16 (x_per_slot=false, gate_up) or
+///           [num_tokens * top_k, K] BF16 (x_per_slot=true, down)
+/// w_fp8_t:  [local_experts, N, K] FP8E4M3 (transposed weight layout)
+/// w_scales: [local_experts] F32 per-expert scalar scales
+/// bias:     [local_experts, N] BF16 (fused into the epilogue)
+/// topk_ids: [num_tokens, top_k] i32 global expert ids (GPU)
+///
+/// Returns y [num_tokens, top_k, N] BF16. Slots routed to experts NOT
+/// owned by this rank are left UNWRITTEN (uninitialized memory) — the
+/// consumer must skip them (see moe_weighted_sum_sparse).
+#[allow(clippy::too_many_arguments)]
+pub fn moe_sparse_gemv_fp8(
+    x: &Tensor, w_fp8_t: &Tensor, w_scales: &Tensor, bias: &Tensor,
+    topk_ids: &Tensor, num_tokens: usize, top_k: usize,
+    expert_start: usize, local_experts: usize, x_per_slot: bool,
+) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous());
+    let n = w_fp8_t.shape()[1];
+    let k = w_fp8_t.shape()[2];
+    assert_eq!(x.shape()[x.ndim() - 1], k);
+    assert_eq!(x.shape()[0], if x_per_slot { num_tokens * top_k } else { num_tokens });
+
+    let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
+    unsafe {
+        launch_moe_sparse_gemv_fp8_bf16(
+            x.data_ptr() as *const c_void,
+            w_fp8_t.data_ptr() as *const c_void,
+            w_scales.data_ptr() as *const c_void,
+            bias.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, x_per_slot as i32,
+            std::ptr::null_mut(),
+        );
+    }
+    y
+}
+
+/// Sparse MoE GEMV (MXFP4 W4A16): same contract as moe_sparse_gemv_fp8,
+/// with packed 4-bit weights [E, N, K/2] + UE8M0 block scales [E, N, K/32].
+#[allow(clippy::too_many_arguments)]
+pub fn moe_sparse_gemv_mxfp4(
+    x: &Tensor, w_packed: &Tensor, w_scales: &Tensor, bias: &Tensor,
+    topk_ids: &Tensor, num_tokens: usize, top_k: usize, n: usize, k: usize,
+    expert_start: usize, local_experts: usize, x_per_slot: bool,
+) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous());
+    assert_eq!(x.shape()[x.ndim() - 1], k);
+    assert_eq!(x.shape()[0], if x_per_slot { num_tokens * top_k } else { num_tokens });
+
+    let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
+    unsafe {
+        launch_moe_sparse_gemv_mxfp4_bf16(
+            x.data_ptr() as *const c_void,
+            w_packed.data_ptr() as *const c_void,
+            w_scales.data_ptr() as *const c_void,
+            bias.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, x_per_slot as i32,
+            std::ptr::null_mut(),
+        );
+    }
+    y
+}
+
+/// Weighted sum over the slot axis of the sparse GEMV output.
+///
+/// down: [num_tokens, top_k, hidden] BF16 (non-local slots uninitialized
+/// and skipped, never multiplied by zero — NaN * 0 = NaN).
+pub fn moe_weighted_sum_sparse(
+    down: &Tensor,
+    topk_ids: &Tensor,
+    topk_weights: &Tensor,
+    expert_start: usize,
+    local_experts: usize,
+) -> Tensor {
+    assert_eq!(down.ndim(), 3);
+    assert_eq!(down.dtype(), DType::BF16);
+    let num_tokens = down.shape()[0];
+    let top_k = down.shape()[1];
+    let hidden = down.shape()[2];
+
+    let out = Tensor::empty(&[num_tokens, hidden], DType::BF16, down.device());
+    unsafe {
+        launch_moe_weighted_sum_sparse_bf16(
+            down.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
 ///