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_out

This 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.