NVIDIA reports 10x dropless-MoE speedup in JAX
NVIDIA says new Transformer Engine kernels for dropless mixture-of-experts training in JAX raise DeepSeek-V3 throughput on GB200 from 103 to 1,068 TFLOPS per GPU, a self-reported 10.4x gain.
NVIDIA's developer blog documents a set of Transformer Engine kernels for dropless mixture-of-experts training in JAX, and reports that they lift DeepSeek-V3 (671B) training on GB200 from 103 TFLOPS/GPU to 1,068 TFLOPS/GPU — a 10.4x gain. In the unoptimized baseline, inter-GPU communication consumed 84% of accumulated kernel time. NVIDIA also cites 97% scaling efficiency at 1,024 GPUs.
The target problem is ragged tensors. Because the router is learned and its expert distribution skews during training, each expert receives a different, data-dependent number of tokens, so there is no clean rectangular GEMM to dispatch. Capacity-based MoE frameworks avoid this by fixing a per-expert token budget and either dropping overflow tokens or padding — trading model quality for hardware-friendly shapes. Dropless MoE processes every token, which means every kernel touching expert compute must handle variable shapes, and must do so without reading token counts back to the CPU, which would break CUDA graphs and force recompilation.
Transformer Engine's contribution is three pieces: group-aware MXFP8 quantization, an MXFP8 grouped GEMM for expert matmuls, and optimized expert-parallel dispatch and combine. The grouped GEMM (grouped_gemm/ragged_dot) runs all expert matmuls in one kernel call at actual token counts, backed by cuBLAS and cuBLASLt; on Blackwell it enables MXFP8 block scaling for expert matmuls.
The numbers are self-reported and measured against NVIDIA's own unoptimized baseline. Much of the post is MaxText configuration detail for reproducing the run.