[KongchangAI]
· 2 min read· 1,337 words

Prime Intellect Releases Prime Mega MoE: Trainable Fused MoE Kernel with 3.11x Forward Speedup

Prime Intellect Releases Prime Mega MoE: Trainable Fused MoE Kernel with 3.11x Forward Speedup

Prime Intellect adds backpropagation to DeepSeek's fused MoE kernel, boosting training step speed by up to 34%.

Prime Intellect releases Prime Mega MoE, implementing backpropagation on DeepSeek's open-source fused Mega MoE kernel to make it trainable. Official data shows up to 3.11x forward speedup and up to 34% faster Qwen3 MoE training steps in prime-rl. The article distinguishes kernel-level forward speedup from end-to-end training gains, and explains why MoE's routing structure amplifies fused kernel benefits in training.

Prime Intellect has released Prime Mega MoE, adding backpropagation to DeepSeek's fused Mega MoE kernel and making it truly trainable. In prime-rl, Qwen3 MoE training step speed improves by up to 34%, and forward propagation reaches 3.11x speedup. The significance of this work isn't about topping some benchmark—it's about pulling a high-performance kernel that was "inference-only" into the training loop, giving the efficiency bottleneck in MoE training a new solution.

Prime Intellect Releases Prime Mega MoE

From "Inference-Only" to "Trainable": Backpropagation Is the Key Piece

DeepSeekAd's previously open-sourced Mega MoE kernel is known for its fused operators, merging scattered operations in MoE such as expert routing and grouped matrix multiplication into a single execution, delivering strong performance in inference scenarios. But it lacked a backpropagation implementation and couldn't be used directly for training—like being handed a sports car without a steering wheel.

What Prime Intellect did was precisely add the steering wheel: implementing the backward pass for the fused kernel so gradients can propagate correctly. According to MarkTechPost, this change upgrades the kernel from an inference acceleration tool to a complete training component.

For training frameworks, this means there's no need to fall back to the slow, unfused path during training—the fusion gains from the forward pass can extend across the entire training pipeline. MoE models have large parameter counts and many experts, making operator scheduling overhead significant in both forward and backward passes. The value of fused kernels on the training side may be even greater than on the inference side.

Forward propagation and backpropagation are the two stages of neural network training. Forward propagation computes the output and loss layer by layer from input data; backpropagation starts from the loss and computes gradients layer by layer via the chain rule for the optimizer to update parameters. Inference only needs forward propagation, while training requires both. If a fused kernel only has a forward implementation, training must either fall back to unfused slow operators or hand-write backward logic—the former loses performance, the latter incurs high engineering costs. Completing the backward pass means every custom operator in the fused kernel must define its corresponding gradient computation while ensuring numerical correctness and memory efficiency. This is the threshold a kernel must cross to go from an inference tool to a training component.

Performance Numbers: 3.11x Forward and 34% Step Speed Improvement

The official disclosure includes two key figures: up to 3.11x speedup in forward propagation, and up to 34% improvement in Qwen3 MoE training step speed within prime-rl.

It's important to distinguish the scope of these two numbers. The 3.11x is a kernel-level forward speedup, measuring the throughput advantage of a single fused kernel over a baseline implementation; the 34% is an end-to-end training step speed improvement that already amortizes backpropagation, optimizer updates, communication, and other stages. The latter is closer to what users actually experience, while the former shows the kernel's own ceiling.

The gap between the two numbers also suggests: no matter how fast the kernel is, overall training remains constrained by other stages. Converting forward speedup into end-to-end gains requires the backward implementation to be efficient enough—otherwise the backward pass becomes the new bottleneck. Prime Intellect delivering the backward pass alongside is precisely to make both ends match.

Prime Mega MoE Performance Illustration

Why MoE Training Especially Needs Fused Kernels

The core structure of MoE is "routing + multiple experts." Each token is assigned to several experts, computation is split into many small matrix multiplications, and results are aggregated by routing weights. This structure naturally comes with irregular workloads and non-contiguous memory access in distributed training.

The fused kernel approach merges routing, grouped GEMM, activation, and aggregation steps to reduce kernel launch counts and intermediate tensor reads/writes. For MoE, the cost of moving intermediate results is often on the same order as computation itself, so the benefits of fusion are amplified.

Prime Mega MoE chooses to build training support on top of DeepSeek's Mega MoE, effectively reusing the existing fusion design and then completing the gradient path needed for training. For teams using prime-rl for MoE post-training, this is an upgrade path that requires no model rewriting.

Grouped GEMM is one of the core operators in MoE fused kernels. Standard GEMM performs multiply-add on a batch of matrices with identical shapes; in MoE, different experts receive different numbers of tokens, with each expert corresponding to a small matrix multiplication of different shape. Grouped GEMM packs multiple small matrix multiplications of varying shapes into a single kernel call, indexing each input/output by group offsets. This reduces kernel launch counts and scheduling overhead, and avoids GPU idle time from per-expert loops. Fused kernels further chain routing indices, token reordering, grouped GEMM, activation, and weighted aggregation into a pipeline, keeping intermediate results on-chip as much as possible or reducing memory round-trips—this is precisely the main source of optimization for memory-intensive stages in MoE training.

Impact on the Open-Source MoE Training Ecosystem

Prime Intellect is known for open training infrastructure, and prime-rl is its reinforcement learning training framework. Putting a trainable fused MoE kernel into this system points toward post-training and RL fine-tuning scenarios for MoE models—scenarios that are sensitive to training step speed and frequently alternate between forward and backward passes.

From an ecosystem perspective, high-performance MoE kernels have long been concentrated inside a few frameworks, with even fewer reusable implementations on the training side. Prime Mega MoE is released as a standalone component, lowering the barrier for other frameworks to integrate. If more frameworks adopt it going forward, the efficiency baseline for MoE training could rise across the board.

Of course, the 34% step speed improvement was measured on a specific model (Qwen3 MoE) and a specific framework (prime-rl); gains will vary when transferred to other model architectures or parallelism strategies. The number of experts, routing strategy, and sequence length in MoE all affect the actual performance of fused kernels, so teams should evaluate based on their own workloads.

prime-rl is Prime Intellect's open-source reinforcement learning training framework, aimed at post-training and RL fine-tuning of large models. RL training is characterized by high-frequency alternation between forward and backward passes: the sampling stage performs large amounts of forward inference to generate trajectories, and the update stage then performs backpropagation. This alternating pattern is sensitive to per-step latency, and also allows the forward gains of fused kernels to directly show up in sampling throughput. Connecting the trainable Mega MoE kernel to prime-rl means using the same fusion implementation on both the forward and backward sides of the RL loop, avoiding path-switching overhead between sampling and updating—especially important for RL fine-tuning of MoE models.

Summary

The value of Prime Mega MoE can be summed up in one sentence: turning a fused MoE kernel from an inference accelerator into a training accelerator. The 3.11x forward speedup and 34% training step speed improvement are metrics at two different levels—the former shows the kernel's ceiling, the latter reflects end-to-end gains. For teams doing MoE post-training and RL fine-tuning, this is an open component worth watching; for the broader training framework ecosystem, it provides a ready-made path to reuse high-performance kernels.

One-line Summary

Prime Intellect adds backpropagation to DeepSeek's fused MoE kernel, improving training step speed by up to 34%.

Paragraph Summary

Prime Intellect has released Prime Mega MoE, implementing backpropagation on top of DeepSeek's open-source fused Mega MoE kernel, upgrading it from an inference-only accelerator to a trainable component. Official data shows up to 3.11x speedup in forward propagation and up to 34% improvement in Qwen3 MoE training step speed within the prime-rl framework. The article distinguishes the scope of the two metrics: 3.11x is kernel-level forward speedup reflecting the ceiling; 34% is end-to-end gain including backward pass, optimizer, and communication. MoE's routing structure and irregular memory access amplify the benefits of fused kernels on the training side. The component is released openly, lowering the barrier for other frameworks to reuse high-performance MoE kernels, though gains will vary with model architecture and parallelism strategy.

#Prime Intellect#Prime Mega MoE#DeepSeek#MoE training#fused kernel#backpropagation#prime-rl#Qwen3 MoE
Share:

Related articles