Prime Intellect has released Prime Mega MoE, a training-capable extension of DeepSeek's fused Mega MoE kernel from the DeepGEMM library, adding a fast, overlapping backward pass to enable full training of mixture-of-experts (MoE) models. The kernel fuses expert-parallel dispatch, two expert linear layers, SwiGLU activation, and combine operations, overlapping communication with computation while using GPU-only scratch buffers to avoid intermediate tensors. It supports BF16, blockwise FP8 on Hopper, and MXFP8 on Blackwell GPUs, with an fp8_bwd mode for BF16 forward and FP8 backward passes. Integrated as a "mega_moe" dispatch backend in the prime-rl trainer via a prebuilt DeepGEMM wheel, it was tested on Qwen3-30B-A3B and Qwen3-235B-A22B models (128 experts, top-8 routing) using 8×H200 nodes (32–128 GPUs) with FP8 precision, 16,384-token sequences, and full activation checkpointing. Full training steps run 32–34% faster than the prime-rl baseline (2.78s to 2.08s on 30B; 12.72s to 9.65s on 235B) and 48–55% faster than TorchTitan. Single-layer forward pass speedups reach 3.11x (12.0ms to 3.9ms on 30B), while backward gains are smaller at 1.87x (30B) and 1.24x (235B). Throughput rises from 5,337 to 7,886 tokens/s/GPU (30B) and 1,094 to 1,697 (235B) versus TorchTitan. Licensed Apache 2.0, it outperforms TorchTitan+DeepEP's separate kernels and AMD's Primus on supported hardware.
Read original
marktechpost