# Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine

DevFeed: [Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine](<https://devfeed.tech/articles/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engine-21079.md>)

Original publisher: [Read original article](<https://developer.nvidia.com/blog/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engine/>)

Author: Tanya Lenz

Published: 2026-09-14T16:39:15Z

Content type: article

Language: en

Sources: [NVIDIA Developer](<https://devfeed.tech/sources/nvidia-developer.md>), [NVIDIA Technical Blog](<https://devfeed.tech/sources/nvidia-technical-blog.md>)

Topics: [Mixture of Experts (MoE)](<https://devfeed.tech/topics/mixture-of-experts-moe.md>), [Artificial Intelligence](<https://devfeed.tech/topics/ai.md>), [GPU](<https://devfeed.tech/topics/gpu.md>), [Large Language Model](<https://devfeed.tech/topics/llm.md>), [Nvidia](<https://devfeed.tech/topics/nvidia.md>), [Routing (disambiguation)](<https://devfeed.tech/topics/routing.md>), [GB200](<https://devfeed.tech/topics/gb200.md>), [Network](<https://devfeed.tech/topics/network.md>), [Python](<https://devfeed.tech/topics/python.md>), [deepseek](<https://devfeed.tech/topics/deepseek.md>), [mixtral](<https://devfeed.tech/topics/mixtral.md>), [qwen](<https://devfeed.tech/topics/qwen.md>)

Tags: [agentic-ai-generative-ai](<https://devfeed.tech/tags/agentic-ai-generative-ai.md>), [ai](<https://devfeed.tech/tags/ai.md>), [communication](<https://devfeed.tech/tags/communication.md>), [deepseek](<https://devfeed.tech/tags/deepseek.md>), [developer-tools-techniques](<https://devfeed.tech/tags/developer-tools-techniques.md>), [gb200](<https://devfeed.tech/tags/gb200.md>), [gpu](<https://devfeed.tech/tags/gpu.md>), [mixtral](<https://devfeed.tech/tags/mixtral.md>), [mixture-of-experts-moe](<https://devfeed.tech/tags/mixture-of-experts-moe.md>), [mlops](<https://devfeed.tech/tags/mlops.md>), [moe](<https://devfeed.tech/tags/moe.md>), [networks](<https://devfeed.tech/tags/networks.md>), [nvfp4](<https://devfeed.tech/tags/nvfp4.md>), [nvidia](<https://devfeed.tech/tags/nvidia.md>), [performance](<https://devfeed.tech/tags/performance.md>), [python](<https://devfeed.tech/tags/python.md>), [qwen](<https://devfeed.tech/tags/qwen.md>), [router](<https://devfeed.tech/tags/router.md>), [routing](<https://devfeed.tech/tags/routing.md>), [tensors](<https://devfeed.tech/tags/tensors.md>), [token](<https://devfeed.tech/tags/token.md>), [tokens](<https://devfeed.tech/tags/tokens.md>), [training-ai-models](<https://devfeed.tech/tags/training-ai-models.md>)

## AI overview

This article explains how NVIDIA Transformer Engine and JAX kernel optimizations accelerate dropless Mixture of Experts (MoE) training. It describes bottlenecks from token routing, expert dispatch and gathering, all-to-all communication, and ragged expert matrix operations. In DeepSeek-V3 training on NVIDIA GB200, the optimized approach increased performance from 103 to 1,068 TFLOPS per GPU, a 10.4x improvement.

## Source excerpt

Mixture of experts (MoE) has become one of the defining architectural trends in large-scale AI model training. DeepSeek, Qwen, and Mixtral are examples of MoE...