1.5x faster MoE training with custom MXFP8 kernels

Summary

An upgrade from Hopper H100s to Blackwell B200s exposed the Mixture-of-Experts (MoE) layer, implemented with MegaBlocks, as a major training bottleneck: it consumed nearly 53% of forward-pass time and 27% of backward-pass time. The team rewrote the layer at the GPU-kernel level using CUDA, PTX, and some ThunderKittens, while moving from BF16 to microscaled FP8 (MXFP8) with nearly zero loss in training quality. The implementation addresses Blackwell-specific tensor memory constraints and quantization overhead, including custom MXFP8 quantization that sustains more than 6.2 TB/s and emits scale matrices directly compatible with tcgen05.mma. On an internal model, MoE execution became 3.5x faster in both passes, producing a 1.5x end-to-end speedup on Blackwell and 2x versus the original Hopper setup, measured in tokens per second per GPU.

Takeaways

1. 1

   On Blackwell, tensor-core results accumulate in tensor memory rather than registers, so custom arithmetic requires transfers between tensor memory and registers; the asynchronous movement can create tensor-core pipeline bubbles and reduce occupancy.

2. 2

   The implementation fused MXFP8 quantization to produce both non-transposed and transposed outputs for forward and backward passes, and attached quantization or dequantization logic to the prologue and epilogue of the fused MXFP8 SwiGLU kernel.

3. 3

   For the internal model, Blackwell MXFP8 reduced MoE forward time from 25.96 ms to 9.45 ms and backward time from 59.17 ms to 17.04 ms; the resulting end-to-end throughput was 24k TPS/GPU, versus 16k for Blackwell BF16.