Prime Intellect has released Prime Mega MoE, a training extension for DeepSeek’s fused GPU kernel Mega MoE. While DeepSeek’s kernel only covered the forward pass, Prime Intellect adds a fast backward pass that overlaps with other computations, making the kernel usable for training and not just inference. According to the company, in its prime-rl trainer, full training steps on Qwen3 MoE models run up to 34% faster, and up to 55% faster compared to TorchTitan.
In a MoE model (mixture of experts), each token is only sent to a small subset of experts, distributed across different GPUs. Each layer must therefore send the tokens, have the experts compute, then return the results, during both the forward pass and the backward pass, the backpropagation phase used for learning. In a standard configuration, these operations are separate: the GPU sorts tokens, waits for transport, computes, then waits again. These waiting times are all idle compute time.
A fused kernel, completed for learning
DeepSeek’s Mega MoE, added to the DeepGEMM library in April 2026, replaces this chain with a single kernel call. A kernel (kernel) is an optimized program that runs directly on the GPU. It fuses the expert-parallel dispatch, the two linear layers of the experts, SwiGLU, and the combination step. While one group of tokens is in transit, the kernel computes on a group that has already arrived, and it avoids intermediate tensors thanks to work buffers located solely on the GPU. However, this kernel only included a forward pass implementation.
Prime Intellect wrote the missing backward pass, also overlapping. The supported data types are BF16, MXFP8, reserved for Blackwell GPUs, and block FP8, reserved for Hopper GPUs. The prime-rl configuration also offers an fp8_bwd mode: forward pass in BF16, backward pass in FP8.
Self-reported figures
All results are published by Prime Intellect itself. The training steps use nodes of 8 H200s, in FP8, with sequences of 16,384 tokens and full activation checkpointing. On Qwen3-30B-A3B, distributed across 32 H200s, a full step takes 2.78 seconds with the base prime-rl, 3.07 seconds with TorchTitan (FP8) and 2.08 seconds with Prime Mega MoE, representing gains of 34% and 48%. On Qwen3-235B-A22B, distributed across 128 H200s, the durations are 12.72 seconds, 14.98 seconds and 9.65 seconds, for gains of 32% and 55%. Throughput increases, on the 30B model, from 5,337 to 7,886 tokens per second per GPU compared to TorchTitan, and on the 235B model, from 1,094 to 1,697.
Measured layer by layer on a single node, the gain is more pronounced on the forward than the backward pass. The forward pass of the Qwen3-30B-A3B model goes from 12.0 ms to 3.9 ms, a factor of 3.11. The backward pass of the Qwen3-235B-A22B model goes from 30.5 ms to 24.6 ms, a factor of 1.24. The figure “up to 3.11x” therefore only concerns the forward pass. TorchTitan claims a gain of up to 41%, but in a different configuration (B200, 671 billion parameter model, BF16 baseline), which the source considers not directly comparable.
Prime Mega MoE is distributed under the Apache 2.0 license. In prime-rl, it is activated as a dispatch backend: simply set type = “mega_moe” under [trainer.model.moe.dispatch], then choose the precision among bf16, mxfp8, fp8 or fp8_bwd. The training kernels are provided in a precompiled DeepGEMM wheel, available on the prime-kernels repository, as the standard version of DeepGEMM does not contain them.
