Dropless MoE training in JAX gets a tenfold lift: DeepSeek-V3 671B holds 97% scaling efficiency at 1,024 GPUs

Dropless MoE training can now use the hardware fully on the JAX stack, with every token processed by the expert its router chose instead of being trimmed or padded when loads are uneven. The earlier JAX training stack reached only 103 TFLOPS/GPU on DeepSeek-V3 on GB300, and inter-GPU communication consumed 84% of accumulated kernel time, leaving GPUs waiting for data. The same DeepSeek-V3 training reached 1,068 TFLOPS/GPU after targeted Transformer Engine kernel optimizations, a 10.4x gain over that baseline, with about a tenfold end-to-end throughput improvement.

The approach replaces expert computation, quantization, token exchange and communication scheduling with implementations that handle variable-length groups directly. A grouped GEMM performs all expert matmuls in a single kernel call, each with its actual token count, computing only regions with valid tokens; it is backed by cuBLAS and cuBLASLt and enables MXFP8 block scaling on Blackwell. Expert-parallel dispatch and combine are fused into one kernel path driven by the NCCL EP backend tuned for imbalanced traffic, so a token sent to several experts on one node crosses the network once; XLA multistream collectives overlap cross-node InfiniBand transfers with intra-node NVLink communication, and a latency-hiding scheduler decides which collectives are safe to overlap.

At scale, communication overhead usually grows faster than compute, and this stack keeps that degradation in check. Multirack scaling is where most systems struggle, because communication overhead tends to outpace compute. With the JAX MoE and Transformer Engine stack applied, the system sustained 97% efficiency at 1,024 GPUs; the optimizations ship in the NVIDIA NGC MaxText container with Transformer Engine built in.

The demonstrated scope is JAX training of DeepSeek-V3 671B on NVIDIA GB300; the throughput figures come from that model's training configuration, and scaling efficiency is reported up to 1,024 GPUs. The reference configuration shipped in the container is specific to DeepSeek-V3, and the text notes that different models require different tuning; it recommends validating correctness on a small MoE model first, then scaling up while tracking step time, TFLOPS/GPU, MFU, grouped GEMM latency and MoE dispatch/combine latency.

The reported 97% scaling efficiency stops at 1,024 GPUs, and behavior at larger scale is not reported; the throughput and scaling numbers are tied to DeepSeek-V3 671B and its training configuration, and other models need different tuning. NVFP4, quantization fused with GEMM, and A2A overlap are listed as planned additions and are not part of these results.