Prime Intellect's Prime Flash MoE Runs 2.4x Faster on Blackwell GPUs
Prime Intellect releases Blackwell-native CUDA kernels for MoE inference that fuse routing, SwiGLU, and quantization into a single pass — up to 2.4x faster than PyTorch grouped GEMM on B200s.
- Up to 2.4x faster than PyTorch grouped GEMM for MoE feed-forward inference on NVIDIA B200 GPUs.
- Fuses routing-aware GEMMs, SwiGLU activation, and quantization into a single kernel — intermediate activations never hit HBM.
- Supports both BF16 and MXFP8 data paths; MXFP8 path quantizes the intermediate activation on-chip between the two matrix multiplications.
- Two pipelines: a fully fused single-launch kernel (best at small token counts) and a split pipeline (default, scales better at large token counts).
- Benchmarked on 8× B200 GPUs; uses Blackwell-exclusive hardware features (tcgen05 tensor cores, TMA gather mode) — not portable to H100.
- Open-source at
PrimeIntellect-ai/prime-flash-moeand integrated into the prime-rl training framework.
Prime Flash MoE is a set of open-source, Blackwell-optimized CUDA kernels that accelerate the feed-forward pass in Mixture-of-Experts (MoE) models. The core idea is borrowed from FlashAttention: never write intermediate results to GPU memory if you can keep them on-chip. On benchmarks run on NVIDIA B200 GPUs, the kernels reach up to 2.4x faster than PyTorch's grouped GEMM baseline across the 4k–128k token range.
The kernels are integrated into Prime Intellect's prime-rl framework and are available as a standalone open-source repo at PrimeIntellect-ai/prime-flash-moe. There is no cost to use them.
The problem with naive MoE inference
In a standard MoE feed-forward layer, each token is routed to a small subset of experts (typically top-k out of E total). Each expert runs a two-stage projection with a SwiGLU activation in between. The naive PyTorch implementation looks like this:
for expert in experts:
gate_up = x[expert] @ w1[expert].T
gate, up = gate_up.chunk(2, dim=-1)
act = F.silu(gate) * up
expert_out = act @ w2[expert].T
out[expert] += routing_weight[expert] * expert_outThis launches separate kernels for the two matrix multiplications and the SwiGLU activation, and materializes the intermediate activation tensor in HBM (high-bandwidth memory), only for that tensor to be read back immediately by the down projection. That round-trip through HBM is pure waste , the activation is written and immediately consumed.
Even the improved version using grouped_mm (which batches all experts into one kernel call) still writes the intermediate activation to memory between the two GEMMs.
Prime Flash MoE is a set of Blackwell-optimized CUDA kernels which never materialize some intermediate tensors at all, and in the fused configuration never materialize the activation either, thus saving a lot of memory traffic.