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-model/src/gpt_oss.rs

@@ -549,6 +549,60 @@ impl GptOss {
             &router_logits, num_experts, top_k,
         );
 
+        // Sparse decode path: compute ONLY the routed experts. The dense path
+        // below reads every local expert's weights per forward; the sparse
+        // GEMVs read ~top_k/num_experts of the bytes, which dominates decode
+        // (memory-bound). Dense reads each weight once for ALL tokens, so it
+        // wins back at num_tokens ≈ local_experts / E[local hits] ≈ 8.
+        const SPARSE_MAX_TOKENS: usize = 8;
+        let quantized = layer.expert_gate_up_fp8.is_some() || layer.expert_gate_up_mxfp4.is_some();
+        if num_tokens <= SPARSE_MAX_TOKENS && quantized && !dense_moe_forced() {
+            let gate_up = if let Some((ref packed, ref scales)) = layer.expert_gate_up_mxfp4 {
+                let n = packed.shape()[1];
+                let k = packed.shape()[2] * 2;
+                xserv_kernels::moe::moe_sparse_gemv_mxfp4(
+                    x, packed, scales, &layer.expert_gate_up_bias, &topk_ids,
+                    num_tokens, top_k, n, k, expert_start, local_experts, false,
+                )
+            } else {
+                xserv_kernels::moe::moe_sparse_gemv_fp8(
+                    x, layer.expert_gate_up_fp8.as_ref().unwrap(),
+                    layer.expert_gate_up_scale.as_ref().unwrap(),
+                    &layer.expert_gate_up_bias, &topk_ids,
+                    num_tokens, top_k, expert_start, local_experts, false,
+                )
+            };
+
+            // GLU over all slots. Non-local slots hold unwritten memory; they
+            // are never consumed (the down GEMV and the weighted sum both skip
+            // slots whose expert this rank does not own).
+            let inter2 = gate_up.shape()[2];
+            let gate_up_flat = gate_up.reshape(&[num_tokens * top_k, inter2]);
+            let activated = gpt_oss_glu(&gate_up_flat, layer.glu_alpha, layer.glu_limit);
+
+            let down = if let Some((ref packed, ref scales)) = layer.expert_down_mxfp4 {
+                let n = packed.shape()[1];
+                let k = packed.shape()[2] * 2;
+                xserv_kernels::moe::moe_sparse_gemv_mxfp4(
+                    &activated, packed, scales, &layer.expert_down_bias, &topk_ids,
+                    num_tokens, top_k, n, k, expert_start, local_experts, true,
+                )
+            } else {
+                xserv_kernels::moe::moe_sparse_gemv_fp8(
+                    &activated, layer.expert_down_fp8.as_ref().unwrap(),
+                    layer.expert_down_scale.as_ref().unwrap(),
+                    &layer.expert_down_bias, &topk_ids,
+                    num_tokens, top_k, expert_start, local_experts, true,
+                )
+            };
+
+            let moe_out = xserv_kernels::moe::moe_weighted_sum_sparse(
+                &down, &topk_ids, &topk_weights, expert_start, local_experts,
+            );
+            self.all_reduce(&moe_out);
+            return moe_out;
+        }
+
         // 3. Replicate input: [tokens, hidden] → [local_experts, tokens, hidden]
         let x_rep = xserv_kernels::moe::moe_replicate(x, local_experts);
 
@@ -625,6 +679,12 @@ impl GptOss {
 
 // --- Helpers ---
 
+/// XSERV_DENSE_MOE=1 forces the dense all-expert path (A/B benchmarking).
+fn dense_moe_forced() -> bool {
+    static FORCED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+    *FORCED.get_or_init(|| std::env::var("XSERV_DENSE_MOE").is_ok_and(|v| v != "0"))
+}
+
 fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
     assert_eq!(a.ndim(), 2);
     assert_eq!(b.ndim(), 2);