Prime Intellect's Prime Mega MoE Trains Giant AI Models 55% Faster

Prime Intellect has released Prime Mega MoE, which adds a training-ready backward pass to DeepSeek’s fused mixture-of-experts kernel. Prime Intellect reports up to 3.11 times faster MoE operations than a standard PyTorch implementation and up to 34% higher end-to-end training throughput than its previous Prime-RL baseline.

The implementation targets a costly part of distributed MoE training: GPUs waiting while tokens and gradients move between experts on different devices. By overlapping that communication with expert computation, the kernel keeps more of each GPU’s capacity occupied.

The reported gains

Prime Intellect measured kernel-level and full-training performance, with each comparison using a different baseline.

These figures come from Prime Intellect’s own benchmarks and represent peak reported improvements. Actual gains will vary with model shape, token routing, network topology, batch size, precision, and the share of each training step spent inside MoE layers.

Where MoE training loses time

A mixture-of-experts model contains many feed-forward networks called experts, while its router sends each token to only a small subset of them. This sparse design increases parameter capacity without applying every parameter to every token.

Large models distribute those experts across GPUs through expert parallelism. Each MoE layer must dispatch tokens to the devices hosting their selected experts, run the expert multilayer perceptrons, and return the outputs to their original positions. The backward pass repeats the exchange for gradients.

Conventional implementations execute token sorting, communication, expert matrix multiplication, and output combination as separate operations. Gaps between those operations leave GPUs waiting for network transfers, creating pipeline bubbles that reduce utilization.

One kernel overlaps the pipeline

DeepSeek’s Mega MoE implementation comes from its open-source DeepGEMM collection. The fused kernel handles token routing, expert computation, and output combination in one invocation, scheduling groups of tokens so computation proceeds while another group moves between GPUs.

Fusion also reduces memory traffic because intermediate results can remain in on-device scratch buffers instead of being written to and read from larger materialized tensors. That saves memory bandwidth and removes kernel-launch overhead.

DeepSeek released a forward implementation optimized for Nvidia Hopper and Blackwell GPUs. Prime Intellect supplied the corresponding backward kernel required for gradient computation, then integrated both directions into Prime-RL as replacements for its previous MoE operations.

Precision and hardware support

Prime Mega MoE supports bfloat16 on Hopper and Blackwell, blockwise FP8 on Hopper, and MXFP8 on Blackwell.

FP8 formats store values with eight bits and use scaling factors to preserve useful numeric range. Blockwise FP8 shares scales across blocks of values. MXFP8, introduced with Blackwell, uses smaller microscaling groups. Training teams should validate loss curves and final model quality when changing precision formats.

How the benchmarks were run

The published end-to-end tests used Nvidia H200 GPUs, placing eight GPUs in each node. The Qwen3-30B-A3B run used four nodes, while Qwen3-235B-A22B used 16 nodes.

Activation checkpointing saves memory by discarding selected activations during the forward pass and recomputing them during backpropagation. The measured full-layer cost therefore includes two forward executions and one backward execution, which explains why its 2.47x speedup falls between the forward and backward results.

The H200 setup demonstrates Hopper performance. The published end-to-end results do not establish equivalent gains on Blackwell hardware, despite the kernel’s support for Blackwell-specific MXFP8.

Where the upgrade fits

Teams training large MoE models with expert parallelism on Hopper or Blackwell hardware are the primary audience. Existing Prime-RL users can access the optimization through the project’s updated MoE implementation, while users of other training stacks would need separate integration work.

Dense models receive no benefit because they lack routed experts. Small MoE workloads with little inter-GPU expert communication may also see limited improvement, and Ampere GPUs fall outside the supported architectures.

Developers evaluating Prime-RL should benchmark complete training steps on their own cluster, check convergence under the selected FP8 format, and compare the migration cost with the measured throughput gain. For communication-heavy MoE runs, Prime Intellect’s results indicate that overlapping token exchange with expert computation can remove a substantial portion of training time.