Better MoE model inference with warp decode

Summary

MoE inference commonly organizes decode around experts, but the post describes “warp decode,” an output-centric kernel for small-batch generation on NVIDIA Blackwell GPUs. Instead of staging tokens into expert-major buffers, each warp owns one output scalar, streams the required weight rows, accumulates contributions from all eight routed experts in FP32, and uses warp-level reductions; the gate/up and down projections are compressed into two kernels. This removes padding, scatter, combine, activation gather, and per-expert output buffers, while keeping BF16 activations and avoiding intermediate quantization. On B200s, the approach delivered 1.84x higher decode throughput, sustained 3.95 TB/s at batch size 32, and produced outputs 1.4x closer to a full FP32 reference.

Takeaways

1. 1

   Warp decode assigns each warp a single output scalar, allowing expert contributions and routing weights to be accumulated without materializing eight intermediate results in memory.

2. 2

   On NVIDIA B200 GPUs, warp decode sustained 3.95 TB/s at batch size 32, equal to 58% of the measured 6.8 TB/s peak for contiguous memory reads.

3. 3

   Outputs were 1.4x closer to full 32-bit ground truth because BF16 activations remain unquantized and reductions use FP32 accumulators.