Skip to content
Neural Daily

ResearchNews

Prime Mega MoE makes DeepSeek's MoE kernel trainable

🇭🇺 Olvasd el magyarul →

Prime Intellect added a backward pass to DeepSeek's fused Mega MoE kernel, turning a forward-only kernel into one that works for training. The result: MoE kernels up to 3.11x faster and up to 34% higher end-to-end training throughput.

By: Mikola Máté Patrik · · 2 min read

Prime Intellect title card reading Prime Mega MoE
Image: Prime Intellect · Press kit · Official social image from the Prime Intellect blog post · Source

The problem

In a mixture-of-experts model, each token is routed to only a few experts, but those experts live on different GPUs. During training, every layer moves data twice: tokens go to the GPU holding their expert, and the results come back. When these steps are separate operations, part of the GPU is always waiting, and that is wasted compute (Prime Intellect).

The Mega MoE kernel in DeepSeek's open-source DeepGEMM collection solves this with a single fused kernel: in one call it routes the tokens, runs the expert MLP math and sends the results back, overlapping communication with computation (GitHub). The catch was that it only had a forward implementation, so it couldn't be used for training.

What Prime Intellect did

The team wrote a fast backward kernel that also overlaps communication and compute, then replaced the standard MoE layers in its Prime-RL training stack with calls to these kernels. Supported data types are bfloat16, MXFP8 (Blackwell only) and blockwise FP8 (Hopper only).

  • Up to 3.11x faster MoE kernels than the default PyTorch-based implementation, measured on a single node with 8 H200s, one MoE layer, 128 experts, top-8 routing and FP8.
  • Up to 1.34x faster end-to-end training, or 34% higher throughput than the Prime-RL baseline.
  • End-to-end tests ran Qwen3-30B-A3B on 4 nodes and Qwen3-235B-A22B on 16 nodes, with 8 H200s per node, 16,384-token sequences, EP=8 plus FSDP and FP8.

What it means for you

If you train or RL-tune large MoE models, communication between GPUs quickly becomes the bottleneck. A 34% throughput gain means that much more training from the same GPU hours, which is real money on large runs. The project also shows that others can build meaningfully on DeepSeek's openly released infrastructure code. Watch for when and how Prime Intellect makes the kernel available in Prime-RL.

    Produced with the help of AI, reviewed by an editor. How we make our articles

    Reviewed and edited by the Neural Daily desk. About

    Related articles