Accelerating MoE Inference with Warp Decode

Existing MoE (Mixture of Experts) inference systems organize computation around experts. This is efficient for large batches or the Prefill stage, but in the Autoregressive Decode stage, which generates only one token at a time, it creates unnecessary overhead for managing data layout.

Warp Decode completely flips the axis of parallelization from experts to outputs (Neurons). By assigning each Warp to a single output computation, it streams weights directly from memory and aggregates results without data staging, inter-warp synchronization, or intermediate buffers.

This approach compresses the computation process into two core kernels:

• Gate/Up Kernel: Each warp owns an intermediate neuron, converting MXFP8 weights to FP32 on the fly for computation. Efficiency is improved by reading the activation vector only once and reusing it for both projections.
• Down Kernel: Using Butterfly Reduction with __shfl_xor_sync, results are aggregated at the register level without going through shared memory.

On Blackwell GPUs, Warp Decode achieved a 1.84x improvement in throughput. It also achieved the feat of simultaneously improving computational accuracy, with outputs becoming 1.4x closer to the FP32 reference values.