Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine

NVIDIA published a technical post describing optimizations to dropless Mixture of Experts (MoE) training in JAX using NVIDIA Transformer Engine. The company reported that an unoptimized baseline for DeepSeek-V3 training on NVIDIA GB300 achieved 103 TFLOPS/GPU, with inter-GPU communication consuming 84% of accumulated kernel time. With JAX and Transformer Engine targeted kernel optimizations, throughput rose to 1,068 TFLOPS/GPU, a 10.4x improvement.

The post frames MoE as a common architectural trend in large-scale AI training, citing DeepSeek, Qwen, and Mixtral as examples of MoE models that match or exceed dense counterparts at a fraction of training compute. It describes dropless MoE as processing every token through its selected expert regardless of load imbalance, in contrast to capacity-based MoE, which assigns fixed token budgets and trims or pads overflow, trading model quality against regularity. Dropless routing produces variable, data-dependent expert token counts and ragged tensors.

Transformer Engine building blocks cited for JAX include group-aware MXFP8 quantization, MXFP8 grouped GEMM on expert matmuls, and optimized expert parallelism (EP) operations for dispatch and combine. Grouped GEMM handles all expert matmuls in a single kernel call using each expert's actual token count, backed by cuBLAS and cuBLASLt; on NVIDIA Blackwell GPUs it also enables MXFP8 block scaling. The EP path is powered by NCCL EP, which fuses dispatch and combine and includes token deduplication so a token sent to multiple experts on the same rank or to multiple ranks on a remote InfiniBand node traverses the network once.

Additional optimizations listed are JAX host offloading, which offloads query and value projection results, and XLA multistreaming collectives, which schedule independent collectives concurrently across CUDA streams to overlap cross-node InfiniBand and intra-node NVLink communication. The post states the stack sustains 97% scaling efficiency at 1,024 GPUs on NVIDIA GB300 NVL72 hardware when training DeepSeek-V3 671B. It reports a 10x end-to-end throughput gain on DeepSeek-V3 671B through MoE in JAX with the optimizations. Planned additions include NVFP4, quantization fused with GEMM, and A2A overlap.

The optimizations ship in the NVIDIA NGC MaxText container with Transformer Engine built in. The post directs readers to the MaxText MoE Configuration guide and Transformer Engine documentation, and notes using the container from September 10, 2026 (ghcr.io/nvidia/jax:maxtext-2026-09-10) or newer.