Prime Flash MoE: Blackwell-Optimized CUDA Kernels for MoE Inference

sloppenheimer · x · 2026-08-14

PrimeIntellect introduced Prime Flash MoE, a set of Blackwell-optimized CUDA kernels designed for Mixture-of-Experts (MoE) inference.

Core Optimizations & Performance:

Avoids materializing intermediate tensors and activations (in fused mode), significantly saving memory traffic.
Delivers up to 2.4× speedup over PyTorch grouped GEMM and 2.3× speedup across the 4k-128k token range.
Integrated into their prime-rl framework to accelerate the forward pass of MoE models.

Dual Data Paths:

Features paths for both BF16 and MXFP8, sharing the same structure but differing in how they feed tensor cores.
Offers fused and split pipelines. Fusion avoids activation round-trips for small problem sizes, while the split pipeline becomes more cost-effective as the working set grows by materializing smaller activations.

More from Infra