Prime Flash MoE - Faster MoE Kernels optimized for Blackwell

Introduction

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. The kernel is up to 2.4× faster than PyTorch grouped GEMM, and ~2.3× across the 4k-128k token range. We also integrate the kernel into our prime-rl framework, serving as an acceleration to the forward pass of MoE models.

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. Each runs in one of two pipelines, chosen with split: the fused one, where all three operations really are a single kernel launch, and the split one, which is three kernels - a gather, the up projection, and a separate down projection.

Why a fused and a split pipeline you may ask?

Fusion trades one activation materialization for C partial-output reductions. For our default shape that is roughly 2 KB of activation data versus 32 KB of partial output contributions per routed token-expert pair. At small problem sizes the output accumulator remains cache-resident, making those reductions cheap enough that avoiding the activation round trip wins. As the working set grows and accumulator locality becomes worse, materializing the much smaller activation once becomes cheaper.

The rest of this post focuses on the structural constraints that make this fusion nontrivial.

1. The computation

We want to compute the forward pass of a standard Mixture-of-Experts (MoE) feed-forward network with a SwiGLU activation.

A router has already selected the top-k experts for each token, together with normalized routing weights.

Let:

• E be the number of experts.
• D be the model hidden dimension.
• H be the intermediate FFN dimension.

The fused gate and up-projection weights are stored as:

The down-projection weights are stored as:

The gate and up projections are concatenated along the output dimension of W1.

After routing, each token is paired with one of its selected experts. The expert computation can therefore be viewed as a batched matrix multiplication in which each routed token uses the weight matrices of its assigned expert.

First: for each routed token-expert pair $(t, e)$, let $x(t) \in \mathbb{R}^D$. Because the expert weights are stored output-major, the first projection is:

We split the result into the gate and up-projection outputs:

The SwiGLU activation is then:

Finally the down projection computes:

The routing weight is applied before the outputs of the selected experts are accumulated:

The equivalent naive PyTorch implementation is:

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 naive implementation launches separate kernels for the two matrix multiplications and the SwiGLU activation. It also materializes the intermediate activation tensor in HBM, only for that tensor to be read back immediately by the down projection. Even when using vanilla PyTorch, we can already greatly improve performance by using grouped matrix multiplications instead a plain for loop for each expert:

gate_up = F.grouped_mm(xg, w1t, offs=offs, ...)
gate, up = gate_up.chunk(2, dim=-1)
act = F.silu(gate) * up
expert_out = F.grouped_mm(act, w2t, offs=offs, ...)
out = wg.unsqueeze(1) * expert_out

These are also the two implementations we benchmark against later. The grouped gemms are already much faster than the for loop, but we’ll hopefully beat them too!

2. Why you can't just fuse it

To keep act on-chip, the CTA which computes a part of act must also be the CTA which consumes it, which naturally stands in the way.

• Problem 1: SwiGLU couples two columns that live far apart. Computing act[t,j] needs gate[t,j] and up[t,j]. In the natural [gate, up] layout those are H columns apart, so a CTA tiling the N axis gets one or the other, but never both. The activation would then need a cross-CTA exchange.
• Problem 2: the down projection wants a whole row. y[t,:] = sum over j of act[t,j] * W2[j,:] sums over the entire intermediate dim H. A CTA holding a slice of j can't produce any finished output element.

Problem 1 turned out to be quite easy to fix while Problem 2 was much more complex, as you will read in the next chapters.

Trick 1: Interleave gate and up in the tensor map

We want CTA j to see gate[128j:128j+128] immediately followed by up[128j:128j+128]. Written out, the row we need for interleaved row p = 128b + r of expert e is:

That expression is affine in (r, b mod 2, floor(b/2), e), which means it can be encoded directly as a Tensor Memory Accelerator (TMA) tensor map. A TMA tensor map is a hardware descriptor that translates multidimensional coordinates into global memory addresses entirely in hardware, eliminating explicit address arithmetic in the kernel. We want to use as much of the hardware subsystems as possible in this case.

The fold that makes it fit: the expert stride is exactly twice the parity stride, so both collapse into a single dim of extent 2E and stride H rows, indexed by 2e + (b mod 2). A box of extent 2 based at 2e then picks up the gate and up halves of one expert in one shot. The kernel fixes its output tile to 256 columns (static_assert(WN*BN == 256)), so CTA j loads rows [256j, 256j+256) - which the descriptor gathers as gate[128j:128j+128] followed by up[128j:128j+128], straight out of the natural [E, N, K] layout.

The matched pair lands in one accumulator: gate in tmem columns 0-127, its partner up in 128-255. SwiGLU then requires to read two halves of a tile:

tcgen05::ld_32x32b_x32(gate, tcgen05::tmem_addr(tmem_acc+slot*ACC, tmem_row, (c32<<5)));
tcgen05::ld_32x32b_x32(up, tcgen05::tmem_addr(tmem_acc+slot*ACC, tmem_row, BK2+(c32<<5)));

The cross-CTA dependency disappears. Each CTA now has both operands in registers, with no additional runtime cost or preprocessing.

Trick 2: Cut the intermediate

After SwiGLU, each CTA owns 128 intermediate columns, but not the complete intermediate activation. It therefore cannot compute and write the final output row by itself. To see how we can still fuse the down projection, consider its decomposition again.

Let the intermediate dimension be partitioned across C CTAs. CTA c owns a contiguous set of intermediate indices:

where B is the number of intermediate columns assigned to each CTA. The full output row is:

Because the intermediate dimension is partitioned across CTAs, we can rewrite this as:

Each CTA therefore computes the contribution from its own 128-column slice and accumulates that partial result into the output tensor. The down projection becomes a split-K GEMM, where the K-split is exactly the intermediate slice owned by each CTA.

Because multiple CTAs contribute to the same output row, the output tensor acts as a global accumulator and must be initialized to zero before the kernel runs.

On the fused path that is the caller's job, since the kernel reduces into out from its very first down stage. The split path zeroes it itself, on a side stream, which section 6 gets to.

The split pipeline always keeps Trick 1. The difference is the down projection. It runs with C=1, so one CTA owns the entire intermediate row. The activation is written once to memory, and the down projection becomes an ordinary GEMM with no split-K reduction.

3. Getting data into the tensor cores

The structure is settled now, so we want to feed our very hungry 5-th gen tensor cores. We read both operands from shared memory through 64-bit descriptors and expect them as core matrices - 8 rows x 8 bf16 = 128 contiguous bytes. For a tile of height T, the required element index is:

idx(row, k) = (T * (k / 8) + row) * 8 + (k % 8)   // [k8][row][8]

Outermost we have the K-chunk, the row in the middle and 8 K-values innermost. A tile written that way is conflict-free by construction: each submatrix is 128 bytes, so every read touches all 32 banks exactly once. One could say that such a tile is pre-swizzled at load time.

Normally we would use swizzling here, which exists to stop the MMA's core-matrix reads from colliding on shared-memory banks when a tile sits in row-major order. The kernel uses both approaches, and which one it uses depends on who writes the tile:

For the TMA-loaded operands where we select a swizzled tensor-map layout, the swizzle is performed as part of the TMA transfer rather than by explicit kernel instructions and the MMA descriptor for that operand carries a matching swizzle-mode layout field. The epilogue writes the activation with ordinary stores, so that one keeps the idx(row, k) layout and a plain descriptor - and so does anything a hand-rolled cp.async producer writes, which on the MXFP8 path is the tokens and, on both paths, every scale tile.

For the bf16 weights: one folded rank-4 descriptor

w1 needs four dimensions, and Trick 1's fold is what makes them fit. The swizzle mode handles the core-matrix split, so no rank has to be spent on a k-major decomposition:

constexpr auto W_SWZ = BK == 64 ? CU_TENSOR_MAP_SWIZZLE_128B : CU_TENSOR_MAP_SWIZZLE_64B;
CUtensorMap map_w = init_tmap<4>("w", CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, w,
    /*shape  */ {K, 128ull, num_experts<<1, H>>7},
    /*strides*/ {K*sizeof(__nv_bfloat16), H*K*sizeof(__nv_bfloat16), 128ull*K*sizeof(__nv_bfloat16)},
    /*box    */ {BK, 128u, 2u, 1u},
    /*elem   */ {1u, 1u, 1u, 1u}, W_SWZ);

TMA fills innermost-first, so the box lands in shared memory as [parity][row][BK] - 256 rows of BK, gate half first, up half second, which is the interleave the accumulator wants. Both MMA operands are therefore read through swizzle-mode descriptors:

uint64_t a_desc = tcgen05::encode_smem_desc_swz<BK*2>(a_ptr);
uint64_t b_desc = tcgen05::encode_smem_desc_swz<BK*2>(b_ptr);

BK*2 is the row length in bytes: 128 for BK=64, 64 for BK=32. Swizzle mode and descriptor layout field track BK together.

For the tokens: TMA gathers them too

Rows of a token tile are a gather - row r comes from row_tok[r], an arbitrary row of x - and TMA's plain tile mode cannot express that. Doing it by hand costs a page of address arithmetic and a producer warp per eight rows. Blackwell's TMA has a gather mode that expresses it directly: one PTX instruction takes four independent row coordinates and one column coordinate:

cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4
    .mbarrier::complete_tx::bytes.L2::cache_hint
    [dst], [tmap, {col, r0, r1, r2, r3}], [mbar], policy;

The descriptor for x is therefore rank-2 - {K, M}, box {BK, 1}, one row tall - and a single warp fetches the whole 128-row token tile, 32 lanes with four rows each:

if (threadIdx.x < 32) {
    #pragma unroll
    for (int i=0; i < 4; ++i) g4_rows[i] = max(row_tok[threadIdx.x*4+i], 0);
    g4_dst = __cvta_generic_to_shared(smem_up.x) + threadIdx.x*4*BK*sizeof(__nv_bfloat16);
}
...
cp_async::gather4(g4_dst + smem_stage*XS*sizeof(__nv_bfloat16),
                  &map_x, *bar_copy[smem_stage], offs,
                  g4_rows[0], g4_rows[1], g4_rows[2], g4_rows[3]);

That puts PRODUCER_THREADS at 32: one producer warp plus warp_n consumer warps, 288 threads in the (block_n, warp_n) = (32, 8) configuration. The token tile is swizzled by hardware like the weights, so x and w share a descriptor encoder. And padding rows are free - a dead row is clamped to row 0 and dropped in the epilogue, so the gather never branches.

One caveat that section 6 will make sense of: gather4 is used by the fused pipeline only. The split one runs its gather kernel up front, after which the up kernel reads its token tile with an ordinary 128-row TMA tile load.

MXFP8 still hand-rolls its tokens

MXFP8 builds the same folded rank-4 descriptor as bf16, hardware-swizzled at 128B. What it does by hand is the tokens - 256 producer threads issuing cp.async 16-byte chunks into the [k16][row][16] submatrix layout, with the up-MMA reading them through a plain unswizzled descriptor while its w1 operand uses a swizzled one.

The reason to leave them there is the scales. Those are plain cp.async traffic no matter what the data does (section 7), so the wide producer warps have to exist anyway and there is no warp budget to reclaim by moving the tokens onto TMA. The split pipeline is the exception: its tokens come pre-sorted out of the gather kernel and go in through TMA with hardware swizzle, exactly like bf16.

4. The pipeline

With 288 threads (bf16, warp_n=8) or 512 (MXFP8) over a STAGES-deep ring, consumers just can't use __syncthreads() - producers are in a different loop and would never arrive, which means that we need a barrier, just for the consumer set:

auto con_sync = [=]() -> void { asm volatile ("bar.sync 1, %0;\n" :: "n"(CONSUMER_THREADS)); };

One barrier per stage

A stage isn't ready until everything in it has landed: the weight tile and the token tile, and on the MXFP8 path the two kinds of scale tile as well. Rather than one barrier per mechanism and a join, they all signal the same one, and the transaction count covers the whole stage payload:

bar_copy[smem_stage].expect_nb((CPC*WS + XS*BPC)*sizeof(__nv_bfloat16));

The consumer waits on exactly one bar_copy[stage].await(phase) per stage, however many copy engines contributed to it.

One ring, both GEMMs

Easy to oversee, but deliberate: smem_stage and phase are not reset between the up loop and the down loop.

if (smem_stage == STAGES) { phase^=1; smem_stage = 0; }
// ^ the only reset, on wrap ^

The down loop simply continues using the same ring. This is possible because a w2 tile (128x128) has exactly the same element count as a w1 tile (64x256) so both are WS = 16384. So the same slots, the same barriers and the same phase parity carry across the GEMM boundary, and the producer starts prefetching w2 while consumers are still draining the last w1 MMAs.

Two phases share their shared memory

smem_up and smem_down are two views of the same allocation, and the launcher takes the max so one allocation works for both phases:

constexpr size_t SMEM_SIZE = std::max(sizeof(smem_up_pod<...>), sizeof(smem_down_pod<...>));

Once the token tile has been consumed, its shared-memory storage is reused for the activation tile. This is safe because the token tile is dead by the time the activation is written. Consumers do not write activations until every up-MMA has retired through bar_mma, while the producer’s w2 loads are gated by bar_recycle and never touch the token region. The activation tile costs nothing on the shared-memory budget - it lives in the corpse of the token tile.

On the MXFP8 path the aliasing is load-bearing enough to check at compile time. The weight and scale regions of the two views have to land at byte-identical offsets:

template <typename Up, typename Down>
consteval bool mxfp8_pod_alias_impl() noexcept {
  return offsetof(Up, w) == offsetof(Down, w) && sizeof(Up::w) == sizeof(Down::w)
      && offsetof(Up, sfw) == offsetof(Down, sfw) && sizeof(Up::sfw) == sizeof(Down::sfw)
      && offsetof(Up, sfx) == offsetof(Down, sfx) && sizeof(Up::sfx) >= sizeof(Down::sfx);
}
static_assert(mxfp8_pod_alias<4, 8, 128, 128, 32>);

The down loop overlaps its own epilogue

The down GEMM pipelines keeps its epilogue one stage behind the MMA work. The epilogue for stage n-1 is issued after stage n's MMAs are committed and before they are awaited, so the tmem reads, the bf16 conversion and the reduction all run underneath live tensor-core work:

tcgen05::commit_mbarrier(*bar_mma[smem_stage]);
if (stage) down_epilogue(stage-1);
bar_mma[smem_stage].await(phase);

That only works because the two stages write different halves of the accumulator. The down GEMM's output tile is BM x DN = 128 x 128, half the 256 columns the up GEMM needed, so the accumulator ping-pongs:

tcgen05::mma_f16(tmem_acc+b*ACC+((stage&1)*DN), a_desc, b_desc, idesc, cg != 0 || k16 != 0);
...
int half = (j&1)*DN;   // the epilogue reads the other one

Stage n fills columns [0,128) while the epilogue for stage n-1 drains [128,256), so the MMA never overwrites a value that has not been read yet. The MXFP8 down loop does not do this - it awaits bar_mma and runs its epilogue in the same iteration, strictly serialized.

5. The two epilogues

SwiGLU - which never leaves the SM

Four warps map one thread to one row - which is the tcgen05 Layout-D lane mapping. Each iteration pulls 32 gate columns and their 32 up partners out of one accumulator with two tcgen05.ld...x32 instructions, and writes 32 activations back:

float gate[32], up[32];
tcgen05::ld_32x32b_x32(gate, tcgen05::tmem_addr(tmem_acc+slot*ACC, tmem_row, (c32<<5)));
tcgen05::ld_32x32b_x32(up, tcgen05::tmem_addr(tmem_acc+slot*ACC, tmem_row, BK2+(c32<<5)));
tcgen05::await_ld();
#pragma unroll
for (int i=0; i < 4; ++i) {
    auto *dst = act_smem+slot*DXS+((((c32<<2)+i)*BM + tmem_row)<<3);
    #pragma unroll
    for (int e=0; e < 8; ++e)
        dst[e] = __float2bfloat16(swiglu(gate[(i<<3)+e], up[(i<<3)+e]));
}

The store index (((c32<<2)+i)*BM + row)<<3 is idx(row, k) for the third time. The 8-column inner granularity of the submatrix and the 8-element inner group of the tmem load are the same 8 - so the epilogue writes the activation already in MMA-ready layout and GEMM-2 just consumes it as is.

That's the intermediate tensor's entire lifetime in the fused pipeline: fp32 in tensor memory -> registers -> bf16 in smem -> tensor core. It never reaches L2, let alone HBM. act_smem is where the two pipelines part ways - fused points it at the down GEMM's operand buffer and the tile is consumed in place, split points it at the same allocation and then stores it out. Same epilogue, same layout, different fate.

Three jobs - one instruction

cuda::ptx::cp_reduce_async_bulk(
    cuda::ptx::space_global, cuda::ptx::space_shared, cuda::ptx::op_add,
    out + tok_src[b]*N2 + j*DN,              // scatter target
    smem_down.out + PAD*tmem_row,            // one row per thread
    DN*sizeof(__nv_bfloat16));

cp.reduce.async.bulk ... .add.bf16 performs a Top-K reduction over experts, the split-K reduction of intermediate slices and the scatter from sorted order back to token order - all at once in the memory system instead of the SM. This instruction performs the entire reduction in hardware, eliminating a separate reduction kernel.

The staging buffer needs ordering on both sides, because ordinary stores (generic proxy) and TMA's reads of smem (async proxy) have no defined implicit ordering between them. Publishing is a fence.proxy.async.shared::cta after the stores and before the reduce. Reuse is the other direction - the previous stage's TMA read has to be off the buffer before this stage overwrites the row - and bf16 gets that from the bulk group instead:

asm volatile("cp.async.bulk.wait_group.read 0;\n");   // last stage's TMA is off the buffer
...                                                   // write this stage's row
cuda::ptx::fence_proxy_async(cuda::ptx::space_shared);
cuda::ptx::cp_reduce_async_bulk(...);

MXFP8 does fence on both sides, because its epilogue and its reduce are separated by a con_sync() across all consumer warps rather than being one thread's straight-line code.

The same three jobs without the staging buffer

That staging buffer costs BM*(BK*2+8) bf16, and at one or two stages the shared memory budget doesn't have it:

alignas(16) __nv_bfloat16 out[STAGES <= 2 ? 8 : BM*(BK*2+8)];

So the shallow configurations reduce straight out of registers instead - no smem round-trip, no bulk group, no fences:

__device__ __forceinline__ void red_global_add_bf16x8(__nv_bfloat16 *dst, const float *v, float scale) {
    uint32_t pk[4];
    #pragma unroll
    for (int i=0; i < 4; ++i) {
        __nv_bfloat162 q = __float22bfloat162_rn({scale*v[i<<1], scale*v[(i<<1)+1]});
        pk[i] = *reinterpret_cast<uint32_t *>(&q);
    }
    asm volatile(
        "createpolicy.fractional.L2::evict_last.b64 policy, 1.0;\n"
        "red.global.add.L2::cache_hint.noftz.v4.bf16x2 [%0], {%1, %2, %3, %4}, policy;\n"
        ...);
}

The instruction performs the same reduction directly from registers, processing 16 bytes per instruction instead of staging an entire row. The evict_last policy matters here: out is a global accumulator that C different CTAs will come back to, and the default policy is happy to evict it in between.

Which of the two you get is worth knowing: cp.reduce.async.bulk is the MXFP8 fused epilogue always, and the bf16 fused epilogue only at stages >= 3. The split pipeline's down kernel reduces from registers unconditionally, and the benchmark's bf16 default is stages=2 - so the register path is what runs in most configurations.

BPC and CPC

Two tiling knobs exist in the template. BPC lets one CTA process 2 row-tiles of 128 tokens against the same weight tiles, so the weights get reused. CPC gives one CTA 2 column groups of 256, which halves the split-K factor C on the down projection, so fewer partial rows land on the same output row. Both are pinned to 1 in the dispatcher - the fused kernel throws on anything else - because the split pipeline attacks the same problem without having to fit a second accumulator into tensor memory. The BPC idea does come back in the split pipeline's down kernel, where there is no SwiGLU epilogue competing for tmem; section 6 gets to that.

6. Paying the split-K tax once

Trick 2 bought the fusion, and section 5 made the reduction as cheap as one instruction can make it. The tax is still there though, and it scales with tokens.

Count it for the shape the benchmark defaults to, K = 2048 and H = 1024. Each CTA owns 256 gate/up columns, which is 128 intermediate columns, so C = H/128 = 8. Every routed (token, expert) pair therefore causes eight full K-wide bf16 partial rows to be added into out: 8 x 4 KB = 32 KB of reduction traffic per pair. The activation the whole kernel refuses to materialize is 1024 values - 2 KB in bf16, 1 KB in mxfp8.

For small token counts the fused kernel can win because the reduction rides in L2, there is one launch, and nothing round-trips. Past that the arithmetic flips and it is cheaper to write the activation down once. That is the split pipeline:

• A small gather kernel materializes x in sorted order first, so the up kernel can read its token tile with a plain TMA tile load instead of gather4. One extra pass over x.
• The up kernel runs the up projection and the SwiGLU epilogue as before, then instead of starting the down loop it TMA-stores its activation tile to a workspace and exits.
• A third kernel does the down projection, each CTA looping over the entire intermediate dimension. No split-K at all. The only reduction left in out is the top-k one, which no tiling can remove.

C goes from 8 to 1 - the down projection stops being a split-K GEMM at all.

The split path also wants the simple tiling, since its point is that one CTA owns a complete output tile:

static constexpr bool SP = SPLIT && BPC == 1 && CPC == 1 && STAGES <= 4;

The down-projection kernel

The down-projection kernel is a plain dense K-major GEMM with the same producer/consumer ring as the main one, three stages deep. With split-K removed, the kernel reduces directly from registers using red.global.add.v4.bf16x2. Only the top-k reduction remains.

The bf16 grid is half as tall because of that PAIR=2, which is the BPC idea from section 5 in the one place it fits. With no SwiGLU epilogue to host, the down kernel has all 512 tmem columns free, so it holds two 256-column accumulators and runs two row blocks of 128 tokens against one weight tile:

int nb = min(PAIR, total_padded/BM - rb0);
bool same = nb < 2 || expert_idxs[rb0] == expert_idxs[rb0+1];

Both row blocks have to belong to the same expert, which is not guaranteed - the second can straddle an expert boundary in the sorted order. When it does, same is false and the CTA walks the two blocks one after the other with their own weights: correct, just without the reuse. Two accumulators also cost occupancy, one CTA per SM.

Handing over with PDL

The up and down kernels are chained with programmatic dependent launch, so the down kernel's prologue - routing loads, tmem allocation, barrier init, and the first w2 and scale-tile prefetch, none of which depend on the activation - runs while the up grid is still draining. Only activation-dependent loads wait for the synchronization point:

// end of the up kernel
cp_async_extra::store3d(&map_act, act_smem, 0, block_base*BM, bx_col*(TC_N>>4));
cuda::ptx::cp_async_bulk_commit_group();
...
asm volatile("cp.async.bulk.wait_group 0;\n");
cudaTriggerProgrammaticLaunchCompletion();

// start of the down kernel's producer
cudaGridDependencySynchronize();

cudaLaunchAttribute dattr {};
dattr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
dattr.val.programmaticStreamSerializationAllowed = 1;
cudaLaunchKernelEx(&dcfg, dkernel, map_act_ld, map_w2d, out, ...);

The activation is stored and reloaded through two tensor maps over the same buffer, and the descriptors do the layout change for “free” in both directions. The workspace itself is plain row-major [sorted_num, H]. The store map is a rank-3 unswizzled descriptor whose box is {8, BM, 16} - exactly the [k8][row][8] core-matrix layout the epilogue left in shared memory - so the write de-interleaves on the way out. The load map is a rank-2 tile with 128B swizzle, so the read re-swizzles on the way in. Neither side spends an instruction on it, and neither pays for a transpose.

Zeroing without a bubble

out is still an accumulator and still has to start at zero, and a memset of M*K bf16 in front of the up kernel is a serial bubble for no reason. This pipeline is the one place where that bubble can be hidden: nothing touches out until the third kernel, so the zero-fill goes on a side stream and overlaps the gather and the entire up projection, with two events pinning the order:

cudaEventRecord(g_e_head, stream);
cudaStreamWaitEvent(g_zstream, g_e_head, 0);
cudaMemsetAsync(out, 0, M*K*sizeof(__nv_bfloat16), g_zstream);
cudaEventRecord(g_e_zero, g_zstream);
...
cudaStreamWaitEvent(stream, g_e_zero, 0);   // out zeroed before the down reds

So there are two shapes of the same kernel, and the choice between them is a shape decision. split=True is the default because the split-K reduction grows with tokens while the extra HBM round trip does not: once the reduction stops fitting in L2 the split pipeline is ahead and stays ahead. split=False is the interesting one below that, where a single launch and a cache-resident accumulator beat one extra pass over the activation.

7. The MXFP8 variant

There is a second one that runs the same structure in MXFP8: e4m3 data with one e8m0 scale per 32 contiguous elements along K. The pipeline, the split-K, the shared ring, the cp.reduce.async.bulk epilogue and the split variant all carry over - what changes is how the tensor cores are fed, plus two things the bf16 kernel has that this one does not: its down loop does not overlap its epilogue, and it always stages through shared memory rather than reducing from registers.

A different instruction descriptor

The MMA becomes tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale, and it does not take the instruction descriptor we built for bf16. Block-scaled kinds use a second layout - "format 2" - which is easy to get wrong because the fields do not merely move, they change meaning:

• there is no accumulator-type field,
• there is no saturate field,
• bits [4,6) and [29,31) hold the scale-factor selectors for B and A instead.

Since those selectors change per MMA, the descriptor is built once and then patched:

uint32_t idesc = tcgen05::encode_idesc_block_scaled_e4m3(BM, TC_N);
...
tcgen05::idesc_with_sf_id(idesc, k32)

Scales are piped through tensor memory

The scale factors are not MMA operands in the usual sense - they land in tensor memory, next to the accumulator. Producers cp.async them into smem as 512-byte tiles (one tile per 128 MN x 128 K), and the consumer moves each tile into tmem with a UTCCP copy:

tcgen05::cp_sf_32x128b(tmem_acc+TM_SFB+(j<<2),
    tcgen05::encode_sf_smem_desc(smem_up.sfw+smem_stage*NSFW*SFT+j*SFT));

TM_SFA and TM_SFB sit at tmem column 256 and up, past the 256 accumulator columns. This is also the reason the MXFP8 kernel uses BK = 128 where the bf16 kernel uses 64: one stage has to cover exactly one UTCCP group of four MMA-K blocks, which the kernel pins down with a static_assert.

This means that bpc = 1 is the only option here. 256 accumulator columns plus the scale-factor columns already fill the 512-column tensor memory, so there is no room for a second accumulator.

The scales are the one thing TMA does not carry

The weights use their tensor map, so their gate/up interleave is free. The scale tiles do not - they are plain cp.async. Which means that if we did nothing, every weight block would be scaled by its partner's exponents.

This works because the granularity lines up exactly. The interleave step is 128 rows, which is exactly one scale tile, and the tile index is the outermost axis of the blocked scale layout. So undoing the interleave is index arithmetic, not data movement:

int32_t pb = blockIdx.x*NSFW+j;                          // packed tile within the expert
int32_t sb = (1&pb)*(N>>8)+(pb>>1);                      // -> source tile
int64_t tile = (static_cast<int64_t>(expert_idx)*(N>>7)+sb)*(K>>7)+stage;

This is only possible because the two operations commute. MX quantization is performed along K, while the interleave permutes only MN. As a result, the scale values are reordered but unchanged.

The epilogue also quantizes

In the bf16 kernel the SwiGLU epilogue goes fp32 tmem -> registers -> bf16 smem. In MXFP8 the same epilogue additionally computes a per-block amax, derives the e8m0 exponent and stores e4m3, emitting a fresh scale tile for GEMM-2 as it goes:

mx_block_scale(amax, sf, inv);
...
smem_down.sfx[((row&31)<<4)+((row>>5)<<2)+g] = static_cast<uint8_t>(sf);
store_fp8x16(smem_down.x+((k16*BM+row)<<4), act+(h<<4), inv);

So the intermediate is quantized on the fly, inside the SM. It is never materialized in HBM and GEMM-2 consumes it with its scales already in the blocked layout the UTCCP copy expects.

The split variant, in MXFP8

When the up kernel hands off, the activation goes out through a tensor map, but its freshly computed scale tiles have to be written by hand - one 512-byte tile per (row block, column block), into a workspace that sits right behind the activation buffer in the same allocation:

reinterpret_cast<uint32_t *>(sf_ws + (block_base*gridDim.x + blockIdx.x)*SFT)[i] =
reinterpret_cast<const uint32_t *>(smem_down.sfx)[i]

And the down kernel rotates its scale-factor columns in tensor memory. UTCCP copies for later stages are issued while earlier stages' MMAs are still in flight, and the scales are read by the MMA rather than copied into it, so the destination columns have to move:

uint32_t sfa = 256u + (stage&3)*16u;   // rotating scale columns:
uint32_t sfb = sfa + 4u;               // in-flight stages never overwrite each other's SFs

Four sets of 16 columns, not two - the ring is three stages deep and the MMA queue adds more slack on top. They fit because the down kernel's accumulator is 256 columns and tensor memory has 512, and the same rotation runs in the split up kernel for the same reason.

8. Benchmarks on a B200

The primary baseline is the realistic "no fused kernel" option - one grouped GEMM per projection:

gate_up = grouped_mm(x, w1, offs)         # one launch across every expert
gate, up = gate_up.chunk(2, dim=-1)
act = F.silu(gate) * up
expert_out = grouped_mm(act, w2, offs)    # one launch across every expert
out = routing_weight * expert_out

torch.nn.functional.grouped_mm for bf16, scaled_grouped_mm with OCP MX (1x32 e8m0) scales on both operands for MXFP8. The MXFP8 baseline requantizes the intermediate activation exactly like the kernel does, otherwise the comparison is not fair - its groups also have to start on a 128-row swizzle-tile boundary, so each expert's slice is padded up to a multiple of 128 rows, which bf16 does not need. The per-expert loop from section 1 is timed as a third series, so both of the implementations above appear in the plots.

Routing, sorting, the gather and the final scatter happen once during setup and sit outside the timed region for both sides. Both sides are timed with CUDA graph replays, so the numbers are GPU time and not Python dispatch time.

Shape: E=32, top_k=4, K=2048, N=2048, H=1024, (block_n, warp_n) = (64, 4), stages=2 for bf16 and 4 for MXFP8, median of 50 runs.

Benchmark in bfloat16

Benchmark in MXFP8

For MXFP8 we have a second plot with the intermediate requantization switched off - the up projection stays MXFP8, but the SwiGLU output goes straight into a bf16 grouped_mm. This baseline intentionally measures a different question. It skips the intermediate requantization pass, isolating the cost of moving the activation through HBM. The fused kernel does perform the intermediate requantization, so this baseline both skips a full quantization pass over the activation. It is still worth plotting, because the gap between the two baselines is exactly the price of the requantization round trip through HBM, which is the work the fused kernel absorbs on chip between the two MMAs.

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

Benchmark Environment

• GPU: 8× NVIDIA B200 (P0, 1965 MHz SM, 3996 MHz HBM)
• Driver: 580.159.03
• PyTorch: 2.13.0+cu130
• CUDA (PyTorch build): 13.0
• cuDNN: 9.19.0
• Python: 3.14.6
• OS: Ubuntu 22.04.5 LTS
• CPU: Intel Xeon 6760P

9. Summary

The fusion rests on two tricks. Folding gate and up into the tensor map 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. Making the down projection a split-K GEMM over each CTA's intermediate slice lets a CTA that owns only a slice of j finish useful work anyway.

After that it is mostly about not moving data. In the fused pipeline the intermediate goes fp32 tmem -> registers -> smem in MMA-ready layout -> tensor core, quantized on the fly on the MXFP8 path, and never reaches L2. The reductions - top-k, split-K and the sorted-to-token scatter - all happen in the memory system, in one instruction, either cp.reduce.async.bulk .add.bf16 or red.global.add.v4.bf16x2 when there is no shared memory left to stage into.

Once the split-K reduction costs more than the tensor it was avoiding, split=True writes the activation to HBM exactly once and drops the split-K entirely - and because the reduction scales with tokens while that round trip does not, it is the default.

Citation

@article{primeintellect2026primeflashmoe,
author = {Mario Sieg},
title = {Prime Flash MoE - Faster MoE Kernels optimized for Blackwell},
journal = {Prime Intellect Blog},
year = {2026},
month = {August},
note = {https://www.primeintellect.ai/blog/prime-flash-moe}
}