Source-linked AI summary
MonoMoE: An Efficient Fused Mega-kernel for Quantized MoE Decoding
Yu Gong, Kailash Budhathoki, Taeho Kim, Haipeng Li, Ashish Khetan
TL;DR
Sparse MoE decode remains costly because routed experts must fetch weights while auxiliary work can add exposed latency. MonoMoE uses a weight-major persistent kernel to fuse the routed path and overlap auxiliary work; on H200, it outperforms the supplied baselines on the routed-MoE operator.
Problem
Routed MoE layers account for up to 40% of end-to-end decoding latency, while expert-weight fetches make decode memory-bound despite modest per-token arithmetic.
Source: §1 · p. 1Method
MonoMoE partitions expert weights into output-row stripes, places the decode-step token tile on the tensor core's fine-grained N dimension, and fuses the quantized routed path in a persistent kernel.
Source: §1 · p. 2Results
On H200, MonoMoE is up to 1.54× faster than vLLM Triton Grouped GEMM on the same routed-MoE operator and 2.20–3.84× faster than FlashMoE-FP8 adaptation across models and batch sizes.
Source: §1 · p. 2Takeaways & Limitations
MonoMoE executes every evaluated serving configuration without FlashMoE's fixed expert-plus-dispatch SM-capacity limit.
Source: §1 · p. 2, §4.4 · p. 12In MonoMoE's target regime, terminal tensor-parallel all-reduce offers no useful overlap window, so separate TP collectives are preferable.
Source: §5.2 · p. 13
Abstract
from arXiv · showhide
Mixture-of-Experts (MoE) layers increase model capacity without proportionally increasing arithmetic, but their sparse expert computation is difficult to execute efficiently during autoregressive decode. Existing grouped and batched GEMMs are token-major: they construct expert-local token tiles and obtain parallelism from the token dimension. When few tokens reach each expert, this organization incurs tile padding and preprocessing, launches short-lived grids that underutilize memory bandwidth, and exposes quantization, activation, and reduction as separate stages. We present \textbf{MonoMoE}, a weight-major persistent megakernel for block-wise quantized MoE decode. MonoMoE places the complete decode-step token tile on the fine-grained tensor-core $N$ dimension and partitions CTAs over expert-weight tiles, eliminating expert-local token materialization and reducing padded arithmetic. A persistent grid fuses routing, top-$k$ selection, quantization, both expert projections, activation, and reduction in one launch; warp specialization and readiness flags overlap auxiliary work with the dominant expert-weight stream. MonoMoE is integrated with vLLM and supports multiple model shapes through generated kernel specializations and offline schedule tuning. On NVIDIA H200 GPUs, MonoMoE accelerates the complete routed-MoE operator by up to $\mathbf{1.54\times}$ over vLLM Triton Grouped GEMM, is $\mathbf{2.20}$--$\mathbf{3.84\times}$ faster than FlashMoE-FP8 adaptation across various models and batch sizes, and reduces end-to-end time per output token by up to $\mathbf{18.7\%}$ across the evaluated FP8 models, while preserving task accuracy. The MonoMoE implementation and supporting artifacts are open source and available in the \href{https://github.com/flashinfer-ai/flashinfer/tree/main/csrc/fused_moe/monomoe}{FlashInfer repository}.
1 INTRODUCTION
Sparse MoE decode remains costly because expert-weight movement and auxiliary stages burden the critical path, while token-major scheduling underuses tiles and memory bandwidth at low token counts. MonoMoE instead schedules over expert weights and fuses the routed path, improving operator and end-to-end serving performance.
1 INTRODUCTION
- Routed MoE layers account for up to 40% of end-to-end decoding latency, with expert-weight fetches making decode memory-bound despite modest per-token computation.Routing, activation quantization, and scale handling also remain exposed on the critical path.Source: §1 · p. 1
- Token-major grouped and batched GEMMs underfill expert-local token tiles at low decode counts, while separate kernels interrupt expert-weight streaming and expose auxiliary work.Grouped GEMM constructs expert-local work lists for short, low-concurrency grids; fusion alone, as in FlashMoE, retains token-major scheduling assumptions.Source: §1 · p. 1
- MonoMoE partitions expert weights into output-row stripes and places the complete decode-step token tile on tensor-core N, deriving parallelism from weights rather than expert-local token blocks.This mapping removes off-chip token sorting and activation dispatch and reduces padded tensor-core work.Source: §1 · p. 2
- A persistent CUDA kernel fuses routing, top-k selection, quantization, both projections, activation, and reduction, while warp specialization and readiness flags overlap auxiliary work with weight movement.Generated kernel specializations and offline tuning adapt the dataflow to different model shapes, tensor-parallel degrees, and routing policies.Source: §1 · p. 2
- MonoMoE is up to 1.54× faster than vLLM Triton Grouped GEMM and 2.20–3.84× faster than FlashMoE-FP8 adaptation; mean time per output token falls by up to 18.7% on evaluated models.For DeepSeek-V3.1’s eight-way tensor-parallel path, TPOT ranges from a 1.2% regression to a 3.7% reduction, while aggregate output throughput improves by up to 31.8%; GSM8K and HumanEval scores remain within reported uncertainty when MonoMoE’s mean is below baseline.Source: §1 · p. 2
2 BOTTLENECK ANALYSIS & RELATED WORK
The analysis estimates recoverable MoE decode latency, traces the deficit to token-major scheduling and staged execution, and uses these findings to motivate MonoMoE’s weight-major persistent design.
2.1 Entitlement: Quantifying the MonoMoe Kernel Opportunity
- At B=1–8, the monokernel can recover 32–36% of inter-token latency; at B=32, only 8% remains.For B=1, expert GEMMs dominate while sustaining at most 22% of peak DRAM bandwidth; auxiliary work is described as largely hideable under continuous weight streaming.Source: §2.1 · pp. 3–4, Table 1 · p. 4
2.2 Sources of the Deficit
- Token-major grouped and batched GEMMs derive parallelism from expert-local token tiles, which decode sparsity leaves poorly amortized.They work well when experts receive enough tokens to fill tiles, but decode distributes few tokens across many experts.Source: §2.2 · p. 4
2.2.1 Token Preprocessing Overhead
- Token-major preprocessing rounds expert-local rows to M-tile granularity, leaving padding; grouped and batched layouts can also materialize expert-ordered features.For B=8, k=8, and BLOCK M=16, grouped GEMM has 93.75% padding in the dispersed-routing case; batched GEMM has a 3.125% useful-row fraction when E=256 and k=8 in the stated concentrated-routing case.Source: §2.2.1 · p. 4
2.2.2 Low Effective Memory Throughput
- Small or short-lived token-major grids limit outstanding HBM requests, while separate kernels repeatedly restart weight streaming and incur pipeline transients.The Figure 4 caption attributes grouped-GEMM throughput limits to insufficient grid-level concurrency and memory-pipeline fill/drain transients.Source: §2.2.2 · pp. 4–5, Figure 4 · p. 6
2.2.3 Untapped Latency-Hiding Opportunity
- Decode latency is lower-bounded by routing and expert-weight fetch; quantization, top-k selection, activation, and reduction can largely overlap with weight loading.The shared-expert gate is already hidden on a concurrent stream in the measured baseline.Source: §2.2.3 · pp. 5–6
2.3 Existing Works
- Prior persistent designs retain token-major scheduling, whereas MonoMoE targets low-token decode with the original token tile and expert-weight-tile parallelism.The cited discussion says FlashMoE and Mega MoE consolidate stages but retain expert-local activation materialization and underfilled token tiles in low-token decode.Source: §2.3 · p. 6
3 MONOMOE DESIGN
MonoMoE turns routed MoE decode into one persistent kernel, partitioning work over expert-weight tiles while keeping the full decode-step token tile on the tensor-core N dimension. Its design combines padding-reducing token handling, persistent execution, asynchronous overlap, and shape-specific schedules.
3 MONOMOE DESIGN
- MonoMoE executes top-k, quantization, both expert projections, activation, and reduction in one persistent megakernel, with readiness flags handing data between stages without separate launches or grid-wide phase barriers.Its grid is partitioned over expert-weight tiles rather than expert-local token blocks.Source: §3 · p. 6
3.1 Weight-Major, Padding-Free Token Handling
- SwapAB places the token tile on the tensor core’s small-N side and partitions CTAs over weight tiles, exposing weight-matrix parallelism even when experts receive few tokens.The finer supported N shapes reduce padded tensor-core work; MonoMoE also keeps the original input tile on chip and avoids off-chip token sorting or dispatch.Source: §3.1 · p. 7
3.2 Persistent Single-Launch Execution
- A fixed resident CTA grid streams work across multiple experts and both projections, turning many short weight streams into fewer long-lived streams.Projection group sizes are selected independently, while readiness flags connect producer and consumer work; shared-memory buffers with disjoint lifetimes alias storage to keep the persistent footprint small.Source: Figure 5 · p. 7, §3.2 · p. 8
3.3 Latency Hiding via Asynchronous Overlap
- Dynamic warp specialization overlaps routing, input preparation, weight transfers, and deferred epilogues with expert GEMMs whenever dependencies permit.Multi-stage lookahead fetches future weight tiles, including across expert boundaries; only work without a successor, such as the final weighted reduction, remains exposed at pipeline drain.Source: §3.3 · pp. 8–9
3.4 Generic Kernel Interface and Shape-Specific Tuning
- Generated specializations keep token count, routing fanout, and policy as runtime parameters while specializing model shape and schedule at compile time.Candidate schedules are filtered for hardware and memory constraints, numerically validated, then tuned offline by CUDA-Graph latency over target token counts and representative routing distributions.Source: §3.4 · pp. 9–10
4 EVALUATION
On H200, the evaluation compares block-wise FP8 MonoMoE with vLLM Grouped GEMM and FlashMoE-FP8 across isolated routed-MoE and vLLM serving. MonoMoE improves kernel and TPOT performance, while downstream-task accuracy remains within reported uncertainty.
4.1 Experimental Setup
- Evaluation used H200 GPUs, four block-wise FP8 MoE configurations, and vLLM’s CUDA-Graph decode path, with production vLLM Grouped GEMM and an FP8 FlashMoE adaptation as baselines.The end-to-end Sonnet workload swept request concurrency over 1, 2, 4, and 8.Source: §4.1 · p. 10
4.3 Kernel Performance Analysis
- At B=1 on Qwen3.5-35B, removing padding and preprocessing saves 5.467 µs; effective throughput rises from 10.62% to 21.20% of peak.MonoMoE also overlaps routing, quantization, activation, requantization, and partial accumulation with the expert-weight stream.Source: §4.3 · p. 11
4.4 End-to-End Serving
- Against vLLM Grouped GEMM, MonoMoE reduces TPOT by 18.7% on Qwen3.5-35B at B=1, with reductions narrowing as concurrency grows.The reduction on DeepSeek-V3.1 falls from 12.8% to 3.5% across B=1, 2, 4, and 8.Source: §4.4 · p. 11
4.2 Kernel-Level Performance
- MonoMoE beats vLLM Grouped GEMM across evaluated kernel points and is 2.20–3.84× faster than FlashMoE-FP8 overall.Figure 8 normalizes each model and token-count group to MonoMoE latency.Source: §4.2 · p. 11, Figure 8 · p. 12
4.5 End-to-End Accuracy
- Across models and both tasks, MonoMoE introduces no measurable accuracy degradation through the complete serving stack.When MonoMoE’s mean score is lower, the difference is smaller than the combined standard error.Source: Table 3 · p. 12, §4.5 · p. 12
5 DISCUSSION
MonoMoE targets decode steps with at most eight concurrently processed tokens per GPU, where weight-tile parallelism is useful and token-major tiles are underfilled. Communication remains a separate concern: for small-batch single-node decode, the discussion favors separate collectives over fusion.
5.1 Target Operating Regime
- MonoMoE targets B≤8; as B grows, its advantage narrows and grouped or batched GEMM becomes increasingly appropriate for larger token sets.A longer decode context alone does not increase B; prefill and chunked prefill process many sequence positions together.Source: §5.1 · pp. 12–13, Figure 9 · p. 13
5.2 Communication in the MonoMoE
- For single-node, small-batch decode, separate TP collectives are preferable because cross-rank synchronization remains exposed and the persistent grid competes with expert-weight traffic.The all-reduce is terminal, while producing partial results earlier would add output traffic and a reduction pass.Source: §5.2 · p. 13
- Fused EP dispatch and combine performs at parity for small-batch single-node decode; larger batches, higher EP degrees, or scale-out links could create useful overlap windows.In the evaluated setting, the shared expert is already overlapped by the runtime.Source: §5.2 · p. 13