NVIDIA says its JAX stack makes dropless MoE training 10x faster

NVIDIA says it has made dropless Mixture-of-Experts training dramatically faster in JAX, taking DeepSeek-V3 from 103 to 1,068 TFLOPS per GPU on GB200. The company also reports 97% scaling efficiency across 1,024 GPUs on GB300 NVL72 hardware, turning a specialist optimisation story into a useful glimpse of where large-model training infrastructure is heading.

NVIDIA Watch analysis

What happened

In a technical blog post, NVIDIA describes a stack built around Transformer Engine, JAX and NCCL expert parallelism. Its aim is to keep every token in a dropless MoE model while making the resulting irregular expert workloads efficient enough for large clusters.

The reported improvement comes from several layers working together: grouped GEMM handles variable token counts without padding, fused dispatch and combine operations reduce communication overhead, and MXFP8 quantisation, host offloading and XLA multistreaming help keep memory and networking from becoming the bottleneck. NVIDIA says the optimised path is available through its NGC MaxText container, with configuration instructions for reproducing the DeepSeek-V3 setup.

Key findings

• 10.4x reported throughput gain
  NVIDIA says performance rose from 103 to 1,068 TFLOPS per GPU on GB200.
• 97% scaling efficiency
  The company reports this result at 1,024 GPUs on GB300 NVL72 hardware.
• Dropless routing
  Every token is processed instead of being dropped or padded to fit fixed expert capacity.
• Fused expert parallelism
  NVIDIA says dispatch and combine operations reduce duplicated network traffic and idle time.
• Reproducible software path
  The optimisations ship in an NGC MaxText container with reference configuration guidance.

Why it matters

MoE models save compute by routing each token to a subset of experts, but that flexibility creates a systems headache: experts receive uneven token loads, and communication between GPUs can dominate the useful work. NVIDIA's result suggests that the software stack, not just the accelerator, is becoming a decisive part of large-model training economics.

For practitioners, the useful detail is that this is more than a headline number. NVIDIA supplies a container and configuration flags, so teams with suitable Blackwell hardware can test the path and track step time, TFLOPS per GPU, grouped-GEMM latency and dispatch or combine latency. The catch is equally important: these are NVIDIA's measurements on NVIDIA hardware and a specific DeepSeek-V3 configuration, not a universal benchmark verdict.

Our read

This is a substantial engineering result, even with the usual vendor-benchmark asterisk. The interesting advance is the combination of irregular-token kernels and communication scheduling that lets dropless MoE preserve its quality advantage without surrendering the cluster to waiting time.

If you run large MoE training, the MaxText recipe is worth examining. If you do not, watch the broader trend: model architecture and distributed-systems plumbing are increasingly inseparable, which is inconvenient for neat product comparisons and excellent news for anyone who enjoys performance graphs with plot twists.

What to watch

• Independent reproduction of the 10.4x gain on comparable hardware.
• Performance and memory behaviour on models other than DeepSeek-V3.
• NVIDIA's planned NVFP4 and all-to-all overlap additions.
• Whether competing JAX and PyTorch stacks close the gap with similar kernels.

Discussion spark: Do NVIDIA's reported gains make dropless MoE practical enough to justify the added training-stack complexity?

Independent WittyWires tracker for public updates about NVIDIA. Not affiliated with or endorsed by NVIDIA; this is not an official account.