Prime Mega MoE：用於訓練的高速融合 MoE kernel

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.

Normally these steps are separate operations. The GPU sorts the tokens, hands them to the transport layer and waits; then it runs the expert math, hands the results back to the transport layer and waits again. All of that waiting wastes GPU time. The question is whether we can overlap compute with communication and close the bubbles in the pipeline.

Mega MoE to the rescue

Mega MoE is a fused, overlapping kernel from DeepSeek's open-source DeepGEMM kernel collection. It replaces the chain of separate operations above with a single kernel invocation that handles both computation and communication directly, without round trips from the GPU back to the CPU. In one pass, it routes the tokens, runs the expert MLP math and combines the results, overlapping compute with communication throughout.

While one group of tokens is in transit, the kernel is already computing on a group that has just arrived, so the GPU no longer sits idle waiting for data. This overlap of compute and communication is a key part of optimizing modern distributed MoE training.

A fused kernel has another benefit: intermediate tensors can be skipped, since the kernel uses GPU-only buffers as workspace and scratch memory.

Prime Mega MoE

DeepSeek's Mega MoE kernel provides only a forward implementation, optimized for Hopper and Blackwell GPUs. We built on it by adding a fast, overlapping backward kernel, which makes it usable for training. Supported data types are bfloat16, MXFP8 (Blackwell only) and blockwise FP8 (Hopper only).

Combining the forward and backward kernels and replacing the standard MoE layers with calls to these custom kernels produced large speedups.

Benchmarks

Full-step speedup

Mega MoE on vs. off in Prime-RL

Setup: 8×H200 per node; Qwen3-30B-A3B on 4 nodes and Qwen3-235B-A22B on 16 nodes; seq_len=16384 (one packed row per GPU, micro_batch_size=1, no gradient accumulation); EP=8 (intra-node) plus FSDP over all GPUs; FP8; full activation checkpointing; fp32 master weights with AdamW; torch.compile(fullgraph=True); no offloading; real data and real routing. 20 steps, reporting the mean of steps 6–20 (steps 1–5 are compile/JIT warmup).

Mega MoE on vs. TorchTitan

Per-layer speedup

Mega MoE on vs. off in Prime-RL

Setup: 1 node, 8×H200, EP=8; a single MoE layer (dispatch + experts + combine, no router); 16,384 tokens per GPU; 128 experts, top-8; random routing (identical for both sides); FP8; torch.compile(fullgraph=True); full activation checkpointing, so cost = 2× forward + backward.

Conclusion

By overlapping communication and computation, Prime Mega MoE significantly reduces MoE overhead and accelerates end-to-end training. Fused kernels like this can unlock large performance gains in distributed MoE training.