Better MoE model inference with warp decode

TL;DR

Cursor describes “warp decode,” a Blackwell-targeted MoE inference kernel that flips the standard parallelism axis: instead of organizing decode-time computation around experts (gather-tokens-per-expert, run a batched GEMM, scatter+combine), each warp is assigned one output scalar and streams the weight rows it needs across all top-k routed experts into a single FP32 running accumulator. The reorganization compresses the MoE compute layer to two kernels (moe_gate_up_3d_batched, moe_down_3d_batched), eliminates five “bookkeeping” stages (padding, gather, scatter, combine, intermediate-buffer round-trips), and on B200 yields 1.84× throughput vs the expert-centric path with outputs 1.4× closer to FP32 ground truth. Reaches 58% of B200’s measured peak contiguous-read bandwidth (3.95 / 6.8 TB/s). Speedups apply to MoE decode (single-token, low-batch); expert-centric is still better for prefill / large-batch where shared work per expert amortizes the layout overhead.

Key claims

The traditional expert-centric MoE inference path has eight pipeline stages, five of which are bookkeeping (padding to 128-byte boundaries, gathering activations into expert-major layout, scattering per-expert outputs, combining via reduction) — at decode-time with batch ≈ 1 these are non-amortizable overhead with no real computation [§The conventional MoE path].
Warp decode reorganizes parallelism so each warp owns one output value (neuron) for its entire lifetime; the warp streams weight rows for all top-k routed experts and aggregates into a single FP32 register accumulator [§Reorganized parallelism].
The full MoE compute layer compresses to two fused kernels: moe_gate_up_3d_batched (each CTA = 8 warps; each warp owns one intermediate neuron per (token, expert) pair; MXFP8 weights upcast to FP32 on the fly; gate+up dot products share one read of the activation vector) and moe_down_3d_batched (each warp owns one output dim per token; routing weight folded into the FP32 accumulator) [§How the two kernels work].
The final cross-lane reduction inside a warp uses __shfl_xor_sync → PTX shfl.sync.bfly, a hardware butterfly that exchanges registers between lanes without touching shared memory, L1, or barriers [§How the two kernels work].
Eliminating per-expert intermediate buffers saves 32 KB of intermediate traffic per token in BF16 (8 experts × 2048 hidden × 2 bytes) — that L2 capacity is reclaimed for the weight rows that actually determine performance [§Buffer elimination].
Warp independence is total: input activations are read-only, accumulator is in private registers, output writes are to unique addresses; with thousands of independent warps across the B200’s 148 SMs the scheduler hides memory latency behind useful work from other warps [§Warp Independence].
End-to-end throughput on a Qwen-3-style MoE on B200: 1.84× over the expert-centric baseline, flat across all context-length buckets — confirming it’s a pure generation-time improvement independent of prompt length [§End-to-end decode throughput at scale].
Activations stay in BF16 throughout and accumulators in FP32, so the BF16→MXFP8→BF16 quantization round-trip used by the traditional path is skipped; the result is 1.4× closer to a full FP32 reference [§Improved accuracy].
Sustained memory bandwidth at B=32 is 3.95 TB/s — 58% of B200’s measured peak 6.8 TB/s for contiguous reads. The remaining gap is attributed to non-contiguous expert-row access patterns (e.g. routing a token to experts 5, 8, 14, 19) [§Hardware Efficiency].
Correctness vs reference: minimum cosine similarity > 0.999996, maximum absolute difference 0.001953 across all batch sizes [§Hardware Efficiency].
Warp decode is explicitly not a general replacement: prefill and large-batch inference still benefit from expert-centric packing because many tokens share each expert and the organize-cost amortizes. The technique wins specifically when there isn’t enough shared work per expert to justify the bookkeeping [§Warp decode and Composer training].
The kernel is part of Cursor’s Composer training/inference flywheel — faster inference accelerates research iteration, not just user-facing latency [§Warp decode and Composer training].

Method

The standard MoE decode kernel works expert-by-expert: for each of the 128 (or however many) experts, gather the tokens it was routed to, run a batched GEMM on that slice, then scatter the outputs back into token order, then combine across the 8 active experts per token. At batch=1 with one new token per request, every “gather” is a copy of data that already exists, every per-expert GEMM does almost no real work, and 32 KB per token of intermediate buffer is written-and-immediately-read for the combine step. Cursor frames this as five bookkeeping stages around three real compute stages.

Warp decode inverts the loop: pick an output scalar (one neuron in gate_up or one output dimension in down_proj), assign it to a single warp, and have that warp pull in the weight rows it needs across all top-k experts that the relevant token was routed to. The activation vector lives in registers; the accumulator is FP32; the routing-gate weight is folded directly into the accumulator inside the loop, so the per-expert intermediate outputs never materialize. At the end the 32 lane partials are reduced with a butterfly shuffle (__shfl_xor_sync → shfl.sync.bfly), which is a single PTX instruction that exchanges registers between lanes via a lane mask — no shared memory, no L1, no barrier. The whole MoE layer is two such kernels, fused so the activation read is shared between gate and up. Every warp is independent of every other, so the GPU scheduler treats the entire output dim × batch dim space as a flat namespace of work items it can issue in any order.

Results

Throughput: 1.84× end-to-end decode on a Qwen-3-style MoE on B200, flat across all context-length buckets [§End-to-end decode throughput at scale].
Accuracy: outputs 1.4× closer to FP32 ground truth than the classical path, by avoiding the BF16↔MXFP8 round-trip on activations [§Improved accuracy].
Hardware utilization: 3.95 TB/s sustained memory bandwidth at B=32, vs the B200’s 6.8 TB/s measured peak for contiguous reads (58%); gap attributed to non-contiguous per-token expert routing patterns [§Hardware Efficiency].
Correctness: min cosine sim > 0.999996, max abs diff 0.001953 across all batch sizes [§Hardware Efficiency].
Stage count: 8 → 2 kernels (5 bookkeeping stages removed); per-token intermediate buffer 32 KB → 0 [§Buffer elimination].

Why it’s interesting

This is a textbook instance of the IO-Aware Kernel Design template applied to MoE decode: identify an N × K intermediate (the per-expert output buffer; 8 experts × 2048 hidden) that exists only to be reduced along K (the expert axis), then rewrite the dataflow so the producer and consumer fuse into a single streaming kernel with on-chip state. The algebraic decomposition that allows it here is the linearity of the top-k expert combination — the routing weight commutes through into the accumulator. Same trick as FlashAttention (softmax-over-blocks), Flash-KMeans (online argmin + segmented sum), and FlashSampling (argmax-over-partitions) on IO-Aware Kernel Design; the new ingredient is that the parallelism axis itself is flipped — outputs not experts — which is the natural restatement of “fuse producer + consumer” for MoE.

Cursor positions this explicitly as the inference-side companion to its Improving Composer through real-time RL / Composer 2 Technical Report stack: real-time RL ships a new Composer checkpoint roughly every five hours, and the bottleneck on that cycle is decode throughput on the training cluster’s evaluator and on production serving. The 1.84× and the higher numerical accuracy are both compute-amortized over the iteration cadence, not just a serving-latency win. Together with FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling (FlashAttention-4 retuned for B200) this is the second filed Blackwell-specific kernel rewrite, both Cursor- or Dao-AILab-adjacent, both targeting the same diagnosis that B200’s bandwidth and MUFU units shift the bottleneck off arithmetic intensity onto SMEM traffic and intermediate-buffer staging.

It also sharpens an open question on LLM Inference Efficiency that has so far been about speculative decoding and KV-cache compression: those are target-model optimizations that take the kernel as given. Warp decode is the orthogonal lever — speed up the target model’s own decode forward pass by fixing the kernel. A natural follow-up is whether warp decode composes with SD draft models that themselves use MoE (e.g. EAGLE-3 drafts of Kimi K2.5 in TorchSpec: Speculative Decoding Training at Scale), and whether the same output-axis parallelism flip applies to other batch-1 expert-routed structures (DeepSeek-V3 / V4 fine-grained MoE, MoE attention experts in Path-Constrained Mixture-of-Experts).