Prime Intellect extends Mega MoE for distributed training
Prime Intellect says it extended DeepSeek's forward-only Mega MoE kernel with an overlapping backward pass, making the fused MoE implementation usable for distributed training. The vendor reports that its Prime Mega MoE kernels run up to 3.11 times faster than the default PyTorch implementation, while end-to-end training reaches 1.34 times the Prime-RL baseline throughput. Benchmarks used H200 GPUs with Qwen3-30B-A3B and Qwen3-235B-A22B in FP8 configurations.
Mega MoE’s fused computation and communication directly targets GPU idle time and inter-GPU transfers, which Prime Intellect identifies as major bottlenecks in distributed MoE training.
Prime Mega MoE: Fast Fused MoE Kernels for Training
Mixture-of-experts (MoE) models scale more efficiently by activating only a subset of experts per token, but communication between GPUs can become a major bottleneck during training.
We extended DeepSeek's Mega MoE kernel with a fast, overlapping backward pass, making it usable for training and delivering a significant speedup:
- Up to 3.11× faster MoE kernels compared to the default PyTorch-based implementation.
- Up to 1.34× faster end-to-end training compared to the Prime-RL baseline (34% higher throughput).
Technical details
In an MoE model, each token is routed to only a small subset of the model's experts. Those experts live on different GPUs, so training requires a lot of communication between them. In every layer, tokens are sent to the GPU that holds their expert, the expert computation runs, and the results are sent back. This happens twice per layer on every step: once in the forward pass and once in the backward pass.
#Prime Intellect#DeepSeek#Prime Mega MoE#Mega MoE#DeepGEMM#Prime-RL#PyTorch#Qwen3-30B-A3B#Qwen3-235B-A22B#H200