Prime Intellect Boosts MoE Speed 2.4x on Blackwell GPUs

AI development startup Prime Intellect has released Prime Flash MoE, a set of open-source CUDA kernels designed to accelerate the feed-forward pass in Mixture-of-Experts (MoE) models. Benchmarked on an array of eight Nvidia B200 GPUs, the new kernels achieve up to a 2.4x speedup compared to PyTorch's grouped GEMM baseline across a token range of 4k to 128k. The software is free to use and has been integrated directly into the company's prime-rl training framework.

The performance leap comes from a design philosophy similar to FlashAttention, which avoids writing intermediate results to high-bandwidth memory (HBM). In standard MoE layers, separate kernels are launched for matrix multiplications and SwiGLU activations, forcing intermediate tensors to make a slow round-trip to HBM. Prime Flash MoE fuses routing-aware GEMMs, SwiGLU activation, and quantization into a single kernel. This ensures that intermediate activations remain entirely on-chip, eliminating unnecessary memory traffic.

Practitioners can utilize two distinct pipelines depending on their workload. A fully fused single-launch kernel delivers the best performance at small token counts, while a split pipeline serves as the default configuration to scale more effectively at larger token counts. The kernels support both BF16 and MXFP8 data paths. In the MXFP8 path, the system quantizes the intermediate activation on-chip between the two matrix multiplications.

For developers, this release offers a significant optimization path for running large-scale MoE models, but it comes with strict hardware requirements. Prime Flash MoE relies on Blackwell-exclusive hardware features, specifically tcgen05 tensor cores and Tensor Memory Accelerator (TMA) gather mode. As a result, these kernels are not portable to older Nvidia H100 GPUs. The code is publicly hosted in the PrimeIntellect-ai/prime-flash-moe repository.