Better MoE model inference with warp decode

Warp decode is a novel approach to token generation in mixture-of-experts (MoE) inference systems, particularly on Blackwell GPUs, that enhances both performance and accuracy by reorienting the parallelism around outputs instead of experts. This method resulted from exploring the maximum achievable memory bandwidth for MoE decode on Blackwell, leading to the assignment of each warp to a single output value, which significantly improves throughput and accuracy, achieving a 1.84x throughput improvement and outputs 1.4x closer to the full FP32 reference. The approach eliminates the overhead of traditional expert-centric methods by removing unnecessary stages and intermediate buffers, allowing for more efficient scheduling and latency hiding. Warp decode is particularly advantageous in scenarios where there isn’t enough shared work per expert to justify traditional overheads, thus accelerating the research and training pipeline for Composer, a machine learning system, by allowing for faster model iteration and deployment. This strategy maintains hardware efficiency by sustaining substantial memory throughput and ensuring correctness across batch sizes, although it is not a universal replacement for expert-centric execution in scenarios like prefill and large-batch inference.

Use this post, company, and trend context to find content marketing opportunities, perform competitive analysis, or address product feature gaps via the Plushcap MCP server or the Plushcap API.