Prime Intellect's Prime Flash MoE Runs 2.4x Faster on Blackwell GPUs

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.

Two tricks that make fusion possible

Fusing the up-projection, SwiGLU, and down-projection into a single kernel sounds straightforward, but there are two structural blockers that make it genuinely hard.

Problem 1: SwiGLU couples distant columns. SwiGLU (a gating activation used in most modern LLMs) computes SiLU(gate) * up, where gate and up are two halves of the same projection output stored H columns apart. A GPU thread block (CTA) tiling the output dimension gets one half or the other, never both , making on-chip fusion impossible without cross-CTA communication.

The fix is a folded rank-4 tensor map that places every SwiGLU pair in the same accumulator. The cross-CTA dependency disappears, and the kernel still consumes the standard [E, N, K] weight layout. This is done by encoding the gate/up interleave directly into a TMA (Tensor Memory Accelerator) descriptor , a Blackwell hardware feature that translates multidimensional coordinates into memory addresses entirely in hardware.

Problem 2: The down projection needs the full intermediate row. After SwiGLU, each CTA owns only a slice of the intermediate dimension, not the complete row needed to compute the final output. The solution is to make the down projection a split-K GEMM, where each CTA computes a partial output from its own intermediate slice, and all partial results are accumulated into the output tensor. The reductions , top-k expert weighting, split-K accumulation, and the scatter back to token order , are all handled in a single hardware instruction (cp.reduce.async.bulk.add.bf16), performing the entire reduction in the memory system rather than on the SM.

Fused vs. split: an adaptive choice

The kernel ships with two pipelines, selectable via a split flag. The choice matters and depends on token count.

• Fused pipeline (split=False): Everything is one kernel launch. The intermediate activation never leaves the SM. At small token counts, the split-K reduction stays cache-resident in L2, making the fused path a clear win.
• Split pipeline (split=True, default): Three kernels , a gather, the up projection with SwiGLU, and a separate down projection. The activation is written to HBM exactly once. As token count grows, the split-K reduction traffic (32 KB per routed token-expert pair at the default shape) outweighs the cost of that single HBM round-trip, so the split pipeline pulls ahead and stays ahead.

The split flag is the one that moves with token count: --no-split is the single fused kernel, --split is the split pipeline and is the default because it holds up across the whole sweep.

BF16 and MXFP8, both on Blackwell

There are two data paths, one for BF16 and one for MXFP8. They share the same structure and differ in how they feed the tensor cores. MXFP8 (Microscaling FP8, a format standardized by the OCP MX spec) uses 8-bit floats with one shared exponent per block of 32 elements, roughly halving memory bandwidth versus BF16.

The MXFP8 path has one additional trick: the SwiGLU epilogue additionally computes a per-block amax, derives the e8m0 exponent and stores e4m3, emitting a fresh scale tile for the down projection as it goes. So the intermediate is quantized on the fly, inside the SM , it is never materialized in HBM and the down projection consumes it with its scales already in the expected layout.

Benchmarks on B200

The benchmark shape is E=32, top_k=4, K=2048, N=2048, H=1024, measured as median of 50 runs with CUDA graph replays on 8× NVIDIA B200 GPUs. Routing, sorting, and scatter are excluded from the timed region for both sides, making the comparison a clean kernel-vs-kernel measurement.

The primary baseline is two grouped GEMMs with an explicit SwiGLU in between , the realistic "no fused kernel" option. Key results:

• BF16: up to 2.4x speedup over grouped GEMM, ~2.3x sustained across the 4k–128k token range
• MXFP8: consistent speedup over scaled_grouped_mm, with a second baseline (no intermediate requantization) isolating the cost of the HBM round-trip that the fused kernel eliminates

Why this matters now

Prime Intellect recently trained INTELLECT-3, a 100B+ parameter MoE model using their RL stack , so these kernels are directly motivated by real production workloads. When NVIDIA announced Blackwell's native hardware capabilities, the promise was clear , but hardware capabilities are only half the story. The kernel engineering gap between what the hardware can theoretically do and what production inference frameworks actually achieve is significant: benchmarks across leading MoE backends on B200 show a 142 TFLOPS gap between the best and worst implementations, with the fastest being 1.84x faster at batch size 1.

Prime Flash MoE targets that gap directly, using Blackwell-specific hardware features , tcgen05 5th-gen tensor cores, TMA gather mode, and cp.reduce.async.bulk , that have no equivalent on Hopper (H100). The kernels also integrate into prime-rl, a framework for large-scale reinforcement learning designed to scale to 1000+ GPUs , accelerating the forward pass during RL training of MoE models.

If you are running MoE inference or RL training on B200s, this is a drop-in acceleration worth benchmarking. If you are on H100s, the kernels won't apply , they are explicitly Blackwell-native and rely on hardware instructions that do not exist on earlier architectures.