Better MoE Routing Through Algo-Sys Co-Design, Field Notes for Pretraining

I keep this as a working diary of my own thoughts and experiments that I'd want to run for MoE model design. The constraint is pretraining-only where we desire higher throughput at large global batch and bf16/mxfp8, not single-token decode latency. It alternates between articles and code throughout. A lot of the notes here is me talking to AIs and noting them, so there is a fair bit of slop.

Why I think routing, and not just expert computation in MoEs deserves the kernel attention

1. The router is three operators, not one. A small GEMM into $E$ logits, a bandwidth-bound top- $K$ , and a permute/unpermute pair that streams the activation through HBM 4 times. These are all bandwidth-bound on B200, and B200's compute scaled faster than its HBM, so the bandwidth side is hurting more, not less.
2. Dispatch / combine all-to-all is fused with quantization, scale-factor shuffling, and grouped-GEMM layout in the modern stack. "Improving the routing layer" is no longer about the router, it is about the whole pipeline gating → topk → permute → A2A-dispatch → grouped-GEMM → A2A-combine → unpermute, and the wins come from collapsing the boundaries between those stages.

What follows is a tour through the recent literature with one constant question: what changes in megatron/core/transformer/moe/ do these papers and blog posts imply, and what is the smallest experiment that could run on top of the current code to validate one of those changes?

1. The Megatron MoE path these blogs are read against

For grounding, here is the actual call chain in the public NVIDIA/Megatron-LM tree, megatron/core/transformer/moe/moe_layer.py. MoELayer is defined at l. 213, and its forward (l. 598) wraps a nested custom_forward (l. 642) that runs the pipeline route → preprocess → dispatch → routed_experts_compute → combine → postprocess:

# MoELayer.custom_forward, moe_layer.py:642, each stage a method on MoELayer
probs, routing_map     = self.route(hidden_states, padding_mask)          # :438  → apply_module(self.router)(...)
hidden_states, probs   = self.preprocess(hidden_states, probs, routing_map)  # :448
dispatched_input, probs = self.dispatch(hidden_states, probs)             # :489  → token_dispatcher.token_dispatch(...)
output, mlp_bias       = self.routed_experts_compute(dispatched_input, probs)  # :529
output                 = self.combine(output)                             # :557  → token_dispatcher.token_combine(...)
output                 = self.postprocess(output, shared_expert_output)   # :566  → combine_postprocess(...)

(Line numbers are against main and will drift. The older token_permutation / token_unpermutation dispatcher entry points have been split into dispatch_preprocess / token_dispatch / dispatch_postprocess and combine_preprocess / token_combine / combine_postprocess.) Unrolled, that is:

A few non-obvious things stand out on a careful read:

• The router path is several small kernels per microbatch. TopKRouter.routing (router.py:741) goes to fp32, optionally applies apply_input_jitter (router.py:706, only when moe_input_jitter_eps is set), computes the score function (sigmoid/softmax), then topk via topk_routing_with_score_function, plus the aux-loss and expert-bias passes. That is at least 4–5 launches before the permute even appears. Each is small but they bracket every MoE layer. On mbs=1 traces there are 383k kernel events per step, and routing is a non-trivial slice of that.
• permute uses argsort + index_select. With moe_permute_fusion=True it falls through to TE's fused_permute. Without it, the path eats a transpose, an argsort, an index_select, and the gradient path.
• router_gating_linear keeps the gating GEMM in fp32 by default (moe_router_dtype = 'fp32'). That is correct for stability and wrong for throughput on B200, where fp32 tensor cores are 16× slower than bf16. The default deliberately leaves compute on the table for numerics.
• No mxfp8 anywhere on the routing/dispatch path. Experts are bf16 grouped-GEMM. Quantization + dequantization at the GEMM boundary is exactly the "almost 40% of matmul time" figure Cursor reports.

That is the code to be changed. Now the literature.

2. The literature, condensed

2a. Nous Research, MoE Scaling Field Notes

The single most useful thing in Nous's notes is what they observed while profiling, not the architectural take.

• DeepEP intranode kernels, cached_notify_combine specifically, scale poorly with EP: 38% of GPU time at EP=4 vs 19% at EP=2. Root cause was the launch shape: "1 + num_channels blocks with one warp per rank," so doubling EP doubled work-per-block without adding parallelism. Bumping num_sms from 24 → 128 yielded 56–84% kernel speedups and 2.3–2.6× on dispatch/combine bandwidth.
• ScatterAddBackward0 + FusedDispatch were ~53% of step time, of which 73–85% was CPU launch overhead, not GPU work. This is the same kernel-launch-tax phenomenon that bites at small mbs in PP.
• Router gate weight init matters where zero-initialized gates collapse all tokens to expert 0 in early steps. Easy to miss and pathological for early load balancing.
• torch._grouped_mm gives garbage gradients for experts with zero tokens. Workaround: pad each expert to a minimum of 8 tokens.
• bf16 rounding-mode differences between Triton and PyTorch produce ~1-ULP gradient differences.

The scaling number worth remembering: 15,057 tokens/sec/GPU on B200 → "57% of B200 theoretical bandwidth." That is the order-of-magnitude headroom in question, not 2×, but every percentage of bandwidth utilization counted.

2b. Cursor, Kernels for MoE training

This Cursor blog post is the closest thing to a recipe in the public literature. The numbers and where they came from:

• Original observation: MoE layer ≈ 53% of forward, 27% of backward. Quantization itself was 40% of matmul time on Hopper.
• They built MXFP8 grouped GEMMs (Fprop / Dgrad / Wgrad) with FP8E4M3 elements + FP8E8M0 scales, 32-element block scaling.
• Critically: a fused MXFP8 quantization kernel that produces layout directly compatible with tcgen05.mma, eliminating the intermediate scale-shuffle that everyone else does.
• Fused MXFP8 SwiGLU with dequant in prologue and re-quant in the epilogue.
• "Expert-wise supergrouping" for L2 cache: tile the grouped-GEMM so each expert's submatrix lives in cache as a unit.
• 2-CTA collaborative MMA for 15–20% extra, warp specialization between data-movement and compute warps, and a persistent-grid pattern (1 threadblock per SM).
• Result: 3.5× MoE-layer speedup, 1.5× end-to-end training, 2× over the original Hopper setup. ~2,650 TFLOP/s grouped MXFP8, only 4% drop from non-grouped.

Two things stand out here. First, the quantization layout is the bottleneck, not the matmul. Second, B200 ≠ Hopper: Blackwell's TMEM and tcgen05.mma require a different quantization layout than Hopper. Anything ported from H100 needs to be re-checked on B200.

2c. PyTorch ao, MoE mxfp8 next steps (issue #3379)

The PyTorch ao tracking issue is the clearest statement of what an end-to-end mxfp8 MoE pipeline should look like.

• Stay in mxfp8 across the boundary: all-to-all dispatch, token shuffle, grouped GEMM. The naive path quantizes-on-entry / dequantizes-on-exit at every kernel boundary. The optimal path keeps the FP8 elements + E8M0 scales together end to end.
• mxfp8 all-to-all as a drop-in replacement for the bf16 autograd all_to_all is the explicit ask. That cuts the dispatch tensor by 2× immediately.
• Triton kernel for token shuffling that simultaneously permutes the scaling factors so they remain aligned with their tokens. This is the bit nobody implements correctly the first time.
• Current 3D expert weight quantization kernel runs at 65–70% peak HBM bandwidth, with a target of 85%+. That is bandwidth-bound code that should be fixable.
• Small-K dimensions (DeepSeek-v3 / Kimi shapes) are still bad. This is a real concern for any model with $d_{\text{expert}}$ on the small side.

2d. SGLang, Large-scale Expert Parallelism (May 2025)

The LMSYS post is decode-leaning, but a few things matter for pretraining:

• DeepEP has two dispatch modes (Normal for prefill, Low-Latency for decode). For pretraining the right mode is Normal, the throughput-optimized path, not the latency-optimized one.
• Two-Batch Overlap (TBO): split the microbatch in two and overlap the comm of one half with the compute of the other, 27–35% throughput. This composes with the pipeline schedule. It does not replace it.
• DeepGEMM with contiguous and masked layouts plus a Triton permutation kernel that transforms the dispatch output into contiguous-GEMM format is the single Triton kernel most worth looking at hardest. It directly replaces the permute + sort_chunks_by_idxs pair in moe_utils.py.

2e. NVIDIA, Hybrid Expert Parallel

The Hybrid-EP blog is the cleanest description of what topology-aware EP looks like.

• Two operators (dispatch / combine) implemented as block-level pipelines with warp groups for RDMA, G2S, S2G, and reduce.
• Hierarchical communication, where dispatch routes hierarchically and combine does intra-node reduction → RDMA → inter-node consolidation. This is the inter-node fraction $f_{\text{inter}}$ from the cost model, made explicit in the kernel.
• DGX Hopper saturates NVLink with only 8 SMs. A 32-GPU cluster needs ~4 SMs to fill NIC bandwidth. Grace Blackwell needs 16 SMs to fill NVLink. Numbers worth remembering when budgeting SMs to dispatch.
• Native FP8/BF16 in dispatch, BF16 in combine. They are quantizing to FP8 for the network leg and dequantizing on receive, a different choice than Cursor's "mxfp8 all the way."
• Reported model speedups are DeepSeek-V3 MXFP8 14% over prior, Qwen 3 235B MXFP8 10% / BF16 5.5%, and Megatron-FSDP MXFP8 8%. Already integrated into Megatron Core.

2f. Poolside, Laguna M.1 / XS.2 (May 2026)

The Laguna M.1/XS.2 technical report has the cleanest description of MoE comm-compute overlap in a recent foundation-model report, and crucially it's H200-specific (8× H200 over NVLink, EP=8 within a node). Their Compute-communication Overlap for MoE recipe is worth recording in detail because it's directly portable to the Megatron path:

• Two A2A collectives per layer (dispatch + combine) get fused directly into the CUTLASS-based grouped GEMM, not run as separate kernels. Inspired by ParallelKittens. The point is to keep the Tensor Cores never-idle by treating dispatch as a producer to the tile scheduler rather than a barrier in front of it.
• Dispatch, 8 of 132 SMs are dedicated to copying tokens from peer GPUs to local HBM in increasing expert order, using vectorized 128-bit loads over NVLink. As soon as all 8 source GPUs have copied the tokens for expert N, an HBM flag is set. The remaining 124 SMs run a modified grouped-GEMM whose tile scheduler waits on the per-expert flag before issuing tiles. So expert 0 starts MMA the moment its 8 copies finish, while experts 1, 2, ... are still being copied, the overlap is at the expert-flag granularity, not at the kernel-boundary granularity.
• Combine, they modify the grouped-GEMM epilogue to skip the usual "store output tile from registers to HBM," and instead rearrange the tile in shared memory into a 128-bit-aligned layout and send each token directly over NVLink to its owner. Epilogues have no inter-tile data deps, so all SMs run combines in parallel.
• NCCL also gets a dedicated SM budget alongside the dispatch SMs: 4 SMs reserved for FSDP AllGather/ReduceScatter and 1 SM for aux-loss aggregation. So per H200, 8 SMs to MoE dispatch + 5 SMs to NCCL = 13 of 132 SMs (~10%) carved out for comm, leaving 119 for MMA.
• Network-class separation matters. Their dispatch/combine uses the scale-up network (NVLink) while DDP/FSDP collectives use the scale-out network (IB), so the two streams do not contend for bandwidth, it's not just a scheduler trick, it's a topology-aware split.

A few things from this writeup that change how to think about the experiments here:

1. The fusion target is the grouped GEMM tile scheduler, not the permute kernel. Their dispatch is a separate set of SMs running concurrently with the grouped GEMM, gated by a flag-based barrier per expert. That is a more invasive change than a fused-permute Triton kernel, but it eliminates the dispatch from the critical path entirely rather than just compressing it.
2. 128-bit vectorized loads over NVLink is the right granularity on H200, it matches the LDGSTS / NVLink transaction size and gives you the bandwidth without needing TMA-style descriptors.
3. The scheme is dtype-agnostic as written (bf16 grouped GEMM in their report), but the same flag-gated tile-scheduler design composes with mxfp8 grouped GEMM as long as the dispatch SMs also write the E8M0 scale factor alongside each FP8 tile. This is not a "mxfp8 all the way" report, it is a "bf16 dispatch + grouped GEMM overlap perfectly" report, and the mxfp8 layer is orthogonal.
4. EP=8 within a node (no inter-node EP) is what makes this clean: every dispatch is over NVLink, never over IB, so there is no $f_{\text{inter}}$ correction. For multi-node EP groups (common in real workloads), Laguna's intra-node scheme must be composed with Hybrid-EP's hierarchical inter-node dispatch.

2g. DeepSeek V3 + V4, DualPipe → MegaMoE wave-pipeline

The DeepSeek line is the longest-running public study of EP comm-overlap and is worth treating as two distinct contributions. V3 (Dec 2024) introduced DualPipe and the hierarchical IB+NVLink dispatch kernel. V4 (Jan 2026) replaced both with a fused per-wave EP mega-kernel that the team open-sourced as MegaMoE inside DeepGEMM PR #304.

DeepSeek V3, DualPipe and the IB↔NVLink kernel (H800)

• 671B total / 37B activated, 256 routed experts, 2048 H800 GPUs. PP=16, EP=64 across 8 nodes, ZeRO-1 DP.
• DualPipe pipeline schedule: overlaps forward and backward chunks of the same microbatch pair so that all-to-all and PP comm can be fully hidden during execution. Reduces the bubble vs 1F1B and ZB1P (their Table 2) at the cost of $1/PP$ extra peak activation memory. It only requires that PP and microbatch counts both be divisible by 2.
• Hierarchical dispatch kernel: to balance IB (50 GB/s) and NVLink (160 GB/s, ~3.2× IB), each token is limited to at most 4 destination nodes. A token first goes IB to the destination node, then is instantaneously forwarded over NVLink to the GPU that owns its target expert, "without being blocked by subsequently arriving tokens." This caps IB traffic and lets each token reach an average of 3.2 experts per node "without incurring additional overhead from NVLink."
• Warp specialization on 20 SMs / 10 channels. Dispatch warps are split between (1) IB sending, (2) IB-to-NVLink forwarding, (3) NVLink receiving. Combine warps split between (1) NVLink sending, (2) NVLink-to-IB forwarding & accumulation, (3) IB receiving & accumulation. The number of warps per task is dynamically adjusted based on observed workload. Of 132 SMs on H800, 20 are donated to dispatch+combine, about 15%.
• This explicitly enabled them to train without TP, which is a non-trivial systems claim: TP-free training removes the cross-rank attention-comm cost entirely.

DeepSeek V4, Fine-Grained EP Wave Pipeline (DeepGEMM MegaMoE)

• V4-Pro: 1.6T total / 49B activated, 384 routed experts + 1 shared (top-6, hidden 7168, 61 layers). V4-Flash: 284B total / 13B activated, 256 routed experts + 1 shared (top-6, hidden 4096). Pre-trained on 33T / 32T tokens. The first 3 MoE layers in both use Hash routing instead of learned top-K (a stability/cold-start choice that is cheap and worth noting).
• Routed expert weights are FP4 during post-training (their FP4 quantization-aware training). At training time the dispatch is FP8 + BF16 combine, the relevant choice for the pretraining-MFU question.
• The big idea is a wave pipeline. V4 splits the MoE layer into 4 stages (Dispatch, Linear-1 = up-proj, Linear-2 = down-proj, Combine) and partitions experts into waves of a small fraction of the full expert set. Once a wave's tokens are dispatched, its compute starts immediately without waiting for other waves. In steady state, wave $k$ is computing Linear-1/2 while wave $k+1$ is dispatching and wave $k-1$ is combining. The theoretical speedup is 1.92× vs naive serial and 1.42× for Comet-style "Dispatch overlapped only with Linear-1, Combine with Linear-2" (their Figure 5). Measured: 1.50–1.73× over non-fused baselines for general inference, up to 1.96× for latency-sensitive RL rollouts.
• The interconnect-bandwidth scaling result. They derive that for V4-Pro with 6 $hd$ FLOPs per token-expert pair and 3 $h$ bytes of comm (FP8 dispatch + BF16 combine), full overlap holds when $C/B\le 2d=6144$ FLOPs/Byte, i.e. each GB/s of interconnect bandwidth suffices to hide ~6.1 TFLOP/s of compute, after which more bandwidth has diminishing returns. This is the cleanest "stop optimizing comm, your kernels are now compute-bound" rule in any public report to date.
• Pull-based dispatch. Each GPU "actively reads activations from remote GPUs" rather than the source GPU pushing, this avoids the high notification latency that fine-grained push entails. They explicitly say push would be better if hardware had lower-latency cross-GPU signaling, which is a clear ask to vendors.
• Determinism via token-order pre-processing. They impose a single rank's per-expert send/accumulation order so that EP results are bit-reproducible despite floating-point non-associativity, same point as UniEP, but spelled out as an operational requirement, not a megakernel side effect.
• TileLang as the kernel DSL. Hundreds of fine-grained ATen ops were replaced with TileLang fused kernels. They state TileLang gave them iteration speed close to writing in PyTorch with performance close to hand-written CUTLASS, a real production data point, not a hype claim.

The compact way to read DeepSeek V3→V4 is that V3 made intra-step overlap work (DualPipe + hierarchical kernel), and V4 went one level deeper to make intra-layer overlap work (wave pipeline + fused mega-kernel). Both still apply to bf16/FP8, neither hinges on FP4 (which is post-training only) or B200-specific features.

2h. Microsoft AI, MAI-Thinking-1 / MAI-Base-1 (June 2026)

The MAI-Thinking-1 report is the first detailed public Microsoft frontier-model report. The base model MAI-Base-1 is 35B-active / 1T-total MoE pre-trained on 8K GB200 GPUs in Azure (Microsoft confirms an earlier blog mentioned ~15K H100s for an earlier MAI-1-preview, while this report covers a different, newer model on Blackwell).

• Routing: top-8 of 512 experts in a compressed latent space, they call this LatentMoE. The routing decision is made on the uncompressed representation, but the dispatched payload is the compressed one (compression factor 2×, expansion 3× inside the expert). This is a deliberate algo-sys choice: it cuts the dispatch tensor by 2× without changing the routing distribution. Routing decisions are based on the original representation, and each compressed representation is routed to 8 out of the 512 experts with softmax gating.
• 78 layers, hidden 6656, FFN 13312, 80 Q-heads / 8 KV-heads. 5:1 SWA-to-global attention ratio.
• EP=64 within the NVLink domain, ZeRO-2 across racks, TP=1, no PP (their default during pretraining, with CP only used for long-context mid-training).
• Numerical recipe (their term, §2.6.3): weights+activations BF16. FP8 E4M3 for forward GEMM, FP8 E5M2 for data-grad, BF16 weight-grad with FP32 accumulation. All FP8 ops use delayed scaling with one shared history per tensor. Sensitive activations (attention scores, MoE router logits, output logits, MoE combine, and the residual stream) stay BF16.
• Dispatch+compute+collect pipeline at the EP-group level. "We partition local experts into groups of configurable size and pipeline the dispatch → compute → collect phases across groups, so that all but the first dispatch and last collect overlap with expert computation." Exposed comm is then overlapped with global load-balance computation and (optionally) shared experts. Same structural idea as DeepSeek V4 waves and Laguna's dispatch SMs, but expressed at a coarser granularity.
• Static-memory dropless mode. To prevent imbalance-induced OOMs in dropless MoE, they run multiple capped dispatch→compute→collect rounds per group, each processing up to a fixed token capacity, and use per-expert-per-round fine-grained recompute in backward to avoid storing imbalanced activations. This is a memory-determinism trick not previously published explicitly.
• Dispatch kernels are CuTe DSL. "Custom CuTe DSL symmetric-memory kernels for device-initiated, variably-sized, high-throughput all-to-all communication over NVLink." Symmetric-memory = NVSHMEM-style addressing. This is the GB200-native equivalent of DeepSeek V3's hierarchical IB+NVLink kernel, but Microsoft's setup is intra-NVLink-domain only, so they only need the NVLink leg.
• DeepEP integration constraint as architecture pressure: they revised MAI-Base-1's expert input size from the L78 ladder spec specifically because DeepEP requires the all-to-all hidden dimension to be divisible by 512. This is a great example of system constraints feeding back into model architecture decisions.
• Determinism is a first-class requirement. They explicitly disable NVLink SHARP for collective determinism (cost: reduced perf) and use a stable sort in the top-K MoE routing kernel to avoid non-deterministic tie-breaking. Top-K stability is exactly the bit-equivalence concern raised for Experiment B.
• Cluster Launch Control (CLC). Their non-grouped quantization kernels use CLC, "a Blackwell-specific feature for dynamic kernel load balancing." First public mention of CLC use to date, worth flagging as a B200/GB200 lever.

2i. UniEP / FlashMoE / megakernel papers

UniEP and FlashMoE are the "single-megakernel for MoE" papers. The pitch is: fuse routing + expert GEMM + all-to-all into one persistent CUDA kernel using NVSHMEM, eliminate CPU-GPU sync entirely, get 2–4× over multi-kernel baselines.

The numbers are believable, but megakernels are set aside here as a last-resort approach for this setting:

• They are written in CUDA (or CUTLASS), not Triton. Iteration speed is much lower.
• The benefit is largest when CPU launch overhead dominates. As noted above, this is the regime at small mbs in PP. At larger mbs the share of the win narrows.
• Most of the megakernel benefit comes from fusing two specific boundaries, quantization↔dispatch and dispatch↔grouped-GEMM. Both can be captured by less ambitious Triton fusions.

The useful idea from these papers, even without writing a megakernel, is the mental model: deterministic token ordering. UniEP guarantees numerical consistency of overlap schedules by making the token order a first-class invariant. That can be replicated with one extra argsort and a stable scheduler. It matters for grad-bit-equivalence in load-balancing-loss training runs.

2j. Comparison table, recent MoE pretraining systems

The reports above span very different points in the design space (HW platform, EP scope, dtype, overlap granularity). It is useful to put them in one place, with only facts that are explicitly stated in the source PDFs. Empty cells are marked N/D (not disclosed in the available public report). This table is updated as new MoE pretraining reports drop.

A few patterns this table makes obvious:

• Every report converges on "experts as a pipeline." Laguna gates the grouped-GEMM tile scheduler on per-expert flags. DeepSeek V4 does it via wave partitioning. MAI does it at the EP-group level. The granularity differs but the structural idea, don't wait for the whole all-to-all, is universal.
• No public report has FP4 dispatch in pretraining. FP4 is post-training only (V4) or not used at all (everyone else). The "mxfp8 all the way" Cursor recipe is the most aggressive training-time quant in this set.
• Intra-NVLink EP is the hot path. Laguna, MAI, and the implied Cursor setting are all EP-within-node. DeepSeek V3 and the Hybrid-EP system are the only two that explicitly tackle inter-node EP, and both build hierarchical kernels for it.
• The comm-kernel DSL story is bifurcating. Laguna writes CUTLASS, DeepSeek V4 writes TileLang, MAI writes CuTe DSL with symmetric memory. None ship Triton for the dispatch path itself, though SGLang / PyTorch ao have Triton dispatch-format permute kernels. This is useful framing for Experiment A: a Triton fused permute is the right level of ambition for an internal experiment, but the public production stacks have all moved to CUTLASS / TileLang / CuTe DSL for the dispatch-fused kernels themselves.

The table is a living artifact. As Qwen3.6, GLM-5, and the next DeepSeek release publish reports, their rows get added here.

2k. The broader 2025–2026 open-weight MoE landscape

The systems story above only makes sense against the architectures people are actually shipping. The table below is an open-weight MoE configuration sweep (adapted from internal scaling-laws notes), kept to the columns that matter for the comm/compute analysis. Every numeric cell here is taken from the official technical report or the model's published config.json. Cells that could not be verified are N/D.

As a definition reminder, sparsity = $1-(\text{active params}/\text{total params})$ , the fraction of parameters not touched per token. MLP ratio = expert intermediate dim / hidden dim (the per-expert FFN expansion, where smaller = more fine-grained). "Width/Depth" = hidden size / number of layers.

Cells filled in this revision come from primary sources. The entire DeepSeek V4-Pro / V4-Flash rows are from the V4 technical report §4.2.2 (1.6T/49B and 284B/13B, 384/256 routed experts top-6, hidden 7168/4096, expert intermediate 3072/2048, first 3 MoE layers Hash-routed). The Width/Depth, MLP-ratio, and routing fields for Qwen3, Qwen3-Next, Kimi K2, GLM 4.5, GPT-OSS, MiniMax M2 are from their published config.json files. MiniMax M2 and Mistral Large experts/top-k normalized to the #experts (top-k) convention. Laguna XS.2 from its tech report (8 of 256 + 1 shared).

What this landscape says, distilled to the parts that drive the kernel work:

• Fine-grained is the consensus. MLP ratios have collapsed from ~1.0 (GPT-OSS) toward 0.25–0.3 (Qwen3-Next, Kimi K2, DeepSeek). Smaller experts mean more of them, which means more dispatch destinations per token and smaller per-expert GEMMs, exactly the regime where the router/permute/dispatch overhead (buckets 1–4) dominates and the grouped-GEMM efficiency (bucket 5) suffers from small-M tiles. The kernel work in this post is a direct response to the fine-graining trend.
• Sparsity is climbing with scale. The largest models (Kimi K2, DeepSeek V4-Pro, Qwen3-Next) sit at 96–97% sparsity. Higher sparsity = a larger fraction of step time is moving tokens rather than computing on them, which raises the payoff of every byte saved in dispatch (mxfp8 dispatch, compressed-latent dispatch à la MAI).
• Shared experts are split ~50/50. DeepSeek, Kimi, GLM, Qwen3-Next, and Mistral keep 1 shared expert. Qwen3, GPT-OSS, and MiniMax drop it. The shared expert is the natural thing to overlap dispatch comm against (Megatron's shared_expert_overlap, Laguna/MAI both exploit this), so its presence or absence changes which overlap lever is available.
• sigmoid + expert-bias (aux-loss-free) routing is winning. DeepSeek V3/V4, MiniMax M2, and Mistral Large all use it. It's exactly the path in the Megatron topk_routing_with_score_function + expert-bias code targeted in Experiment B. The fused router kernel should assume sigmoid scoring with an additive bias as the common case.

For contrast, the 2023–2024 generation

It is worth keeping the previous generation in view, because the direction of every trend above is what makes the kernel work newly urgent. The early open MoEs were coarse-grained (few, large experts) with top-2/top-4 switch/GShard routing, a regime where the dispatch tensor was small and the per-expert GEMM was large enough to be efficient on its own. The routing layer simply was not the bottleneck it is today.

The contrast is the whole point. Going from Mixtral's 8 experts / top-2 / MLP-ratio ≈ 3.5 to Qwen3-Next's 512 experts / top-10 / MLP-ratio 0.25 moves roughly two orders of magnitude more routing decisions and dispatch destinations per token, while shrinking each expert GEMM by ~14×. The router, permute, and dispatch kernels that were rounding error in 2023 are the MFU-limiting path in 2026. (Sources for this table are the routing-industry-trends table in the Scaling Conditional Computation with Sparse LLM Pretraining notes and the DeepSeekMoE entry per Dai et al. 2024.)

3. A taxonomy, where the wins come from

Stripping all of the above to its skeleton, MoE-routing-layer wins come from one of seven places. Each is marked with what currently runs in the Megatron path.

Buckets 1–3 are router/permute. Buckets 4–7 are the rest of the MoE block.

Since the routing layer is the focus here, buckets 1–3 carry the weight in the experiment design and 4–7 are treated as context.

4. Hardware: B200 vs H200 and what it means for the kernels

A few hardware facts that change which buckets are leverage points on which platform.

B200 (Blackwell)

• 192 GB HBM3e per GPU at ~8 TB/s per GPU.
• Tensor cores expose tcgen05.mma and TMEM (tensor memory), distinct from SMEM. Kernels written for Hopper need re-layout, especially for FP8 scale factors.
• Native MXFP8 path (FP8E4M3 elements + FP8E8M0 group scales over 32-element blocks) at full tensor-core rate, this is the thing Cursor's kernels exploit.
• 2-CTA collaborative MMA, persistent threadblock pattern, warp specialization between DMA and compute warps.
• NVLink 5 delivers 1.8 TB/s per GPU, and 8-GPU NVL72 nodes have very high intra-node bandwidth. Inter-node is still 400 Gb/s class IB / Ethernet, so the intra/inter ratio is more skewed on B200 than on H100 and hierarchical dispatch matters more.
• SM count is higher than H100, so the fraction of SMs you donate to comm kernels is lower for the same absolute count. Hybrid-EP's "16 SMs for NVLink saturation on Grace Blackwell" is a useful budget.

H200 (Hopper refresh)

• 141 GB HBM3e at 4.8 TB/s. Bigger and faster than H100, same compute path.
• No tcgen05.mma and no TMEM. FP8 scale layout is the Hopper-style one (per-tile scales held in SMEM).
• The interesting thing about H200 vs H100 is purely the +76% HBM bandwidth. Everything in the cost model that lives on the bandwidth branch, norms, RoPE, residuals, GLU, top-K, permute, gets ~1.7× faster for free. The compute branch is unchanged.
• This means on H200, the relative cost share of the bandwidth-bound router-side ops drops, and the compute-bound expert GEMMs become a slightly larger share of the MoE block than on H100. Different bottleneck → different first kernel to attack.

Sequencing experiments by platform:

• B200 → mxfp8-everywhere (Cursor-style), with the layout for tcgen05.mma baked into the dispatch / permute path. The router itself is a smaller fraction of MoE-block time, but each saving is repeated many times per layer.
• H200 → fuse the router → permute → dispatch prologue more aggressively, because that's where the bandwidth-bound stack lives. mxfp8 helps here too but the router-side latency reduction is the easier win.
• H100 → either, but the "DeepEP num_sms is wrong" story from Nous is something to check first because it's free.

For both, across an EP group that spans nodes, the inter-node fraction in the all-to-all term dominates. Hybrid-EP-style hierarchical dispatch gets you the biggest single jump, and that part is independent of dtype.

5. Mapping the kernel changes onto the existing code

For each bucket, here is what code to change.

Bucket 1, Router GEMM

• File: moe_utils.py:RouterGatingLinearFunction (l. 1294, wrapper router_gating_linear at l. 1394)
• Current: TE GEMM in fp32 if moe_router_dtype='fp32'.
• Change: a Triton kernel that does the gating GEMM in bf16 with fp32 accumulation, fused with the sigmoid (or softmax) and noise + tau * raw_noise in score space. The output is clean scores and noisy logits in one pass. The cost is one HBM read of input and one of weight, writing logits/scores once. Saves at minimum 2 HBM round-trips per layer per microbatch.
• The reason this is a defensible numerics change is that the gating GEMM only feeds top-K. Top-K is robust to ULP-level noise in the logits (training already injects noise of order ${10}^{-2}$ in score space). fp32 accumulation is what stability actually needs. fp32 input dtype is overkill, and the existing code keeps it for "memory" reasons that don't apply once the kernel layout is under control.

Bucket 2, Top-k + bias + capacity

• File: moe_utils.py:topk_routing_with_score_function (l. 710, group path via group_limited_topk at l. 617)
• Current: this is Python + PyTorch ops orchestrated across at least 5 kernels (sigmoid, topk, gather, bias_add, index_*).
• Change: one Triton kernel from logits → (routing_map_bool, probs_topk, top_indices). Same output contract, fewer launches. With the fused router-GEMM above it, clean+noisy scores come from one kernel and routing_map from a second.

Bucket 3, Permute + scale shuffle

• File: moe_utils.py:permute (l. 337) and unpermute (l. 470)
• Current: argsort + index_select, falls through to TE's fused-permute when enabled.
• Change: Triton kernel that takes (tokens, routing_map, num_out_tokens) and emits (permuted_tokens, sorted_indices) in the layout grouped-GEMM expects (per-expert contiguous chunks, padded to alignment). For mxfp8 mode, also emits the scale factors permuted to match. This is exactly the "Triton permutation kernel that transforms dispatch output into contiguous-GEMM format" that SGLang built and PyTorch ao is asking for.
• This is the single highest-leverage Triton kernel because it sits between routing and grouped-GEMM and it touches both forward and backward.

Bucket 4, Dispatch all-to-all

• File: fused_a2a.py:FusedDispatch
• Current: DeepEP Buffer.dispatch, bf16, with num_sms defaulting to whatever DeepEP picks.
• Changes:
  • Sweep num_sms per platform, Nous's 24→128 number is platform-specific.
  • Try mxfp8 over the wire (Hybrid-EP supports this, and the open question is whether DeepEP can take FP8 + scales as a payload).
  • Hierarchical dispatch when EP spans nodes (Hybrid-EP).

Bucket 5, Grouped GEMM

• File: experts.py:TEGroupedMLP (l. 173) and SequentialMLP (l. 1165)
• Current: grouped_gemm package or TE grouped-GEMM in bf16.
• Change: switch to TE's mxfp8 grouped GEMM where supported on B200, or call out to a Triton mxfp8 grouped GEMM with the right layout. This is the Cursor recipe.

Bucket 6, GLU activation

• Already fused via weighted_bias_swiglu_impl (experts.py:739). The mxfp8 win here is dequant-in-prologue / requant-in-epilogue, relevant only when the surrounding GEMMs are mxfp8.

Bucket 7, Comm/compute overlap

• Code: MoEAlltoAllTokenDispatcher already has _maybe_dtoh_and_synchronize plumbing for stream sync. shared_expert_overlap is implemented. TBO is not.
• Change: split each microbatch into halves, run the second-half permute on a side stream while the first-half dispatch is in flight. This is a scheduling change in MoELayer.forward, not a kernel change.

6. If I were to run a couple experiments at first to internalize the above

The target is 1–2 experiments, in Triton, on top of the existing Megatron code, focused on reducing the MoE layer's impact on pretraining MFU. Here is the cut.

Experiment A, Fused mxfp8-aware permute+dispatch-prologue Triton kernel

This is the highest-leverage routing experiment because it sits exactly at the boundary between the bandwidth-bound router stack and the compute-bound expert grouped-GEMM.

What the kernel does, in one pass:

1. Read tokens : [N, H] (bf16) from HBM.
2. Read routing_map : [N, E] (bool) and probs : [N, K] (fp32).
3. For each token, look up its $k$ destination experts.
4. Quantize tokens[i] to mxfp8: 32-element groups, FP8E4M3 elements + FP8E8M0 scale.
5. Scatter the FP8 element + the per-group E8M0 scale into the per-expert contiguous output buffer, padded to grouped-GEMM tile alignment, in the layout tcgen05.mma expects on B200 (or the SMEM-tile layout on H200).
6. Emit tokens_per_expert and the sorted_indices needed by the backward unpermute.

Why a single kernel. The naive path does 4 HBM round-trips for these steps: bf16 permute, bf16→mxfp8 quant, scale-factor reshape, layout-shuffle into grouped-GEMM format. Fusing them collapses to one read of tokens + one read of routing_map + one write to the dispatch buffer. On the operator cost model that is a $4\times$ reduction in HBM traffic for this stage.

Where it plugs in. Replace permute(...) at token_dispatcher.py:522 with a Triton path that returns the dispatch buffer and indices directly. Skip the standalone TE quantization that today runs immediately after permute. The unpermute path needs a paired kernel that does mxfp8 dequant + scatter-add-with-probs in one pass.

Predicted gain. Cursor's report of "40% of matmul time" being quant overhead is the upper bound on what fusing these stages can save in the matmul-adjacent code. A more conservative estimate from the operator cost model: at $m\cdot d\cdot b_w$ bytes per HBM pass and 4 collapsed passes, on B200 (~8 TB/s) the saving is on the order of 0.4–0.8 ms per MoE layer per microbatch. With 47 MoE layers and 64 microbatches the per-step floor is ~1.2–2.5s, call it 2–4% MFU on the GBS=8192 / PP=16 / mbs=4 trace from the earlier post.

Falsification. If the new kernel does not show ≥1.5× HBM-bandwidth utilization vs the unfused path (measured via Nsight HBM throughput counter), the fusion is wrong. If end-to-end MFU does not move ≥1%, either bf16 grouped-GEMM is fast enough that the saved bytes are not on the critical path, or the fusion has not actually displaced the unfused TE quantization that runs after permute.

Risk. mxfp8 numerics on the routing path. The router itself stays in bf16. Only the tokens are mxfp8-quantized on the way to the experts. The permuted scales must match exactly across forward and backward, UniEP's "deterministic token ordering" idea is an explicit design constraint here. One extra argsort enforces stable order so backward and recompute see identical layouts.

Experiment B, Fused router: gating + noise + sigmoid + topk in one Triton kernel

What the kernel does:

1. logits = input @ weight.T in bf16 with fp32 accumulation. Output dim is small ( $E$ ), so this is a strip-mined GEMM, fine for Triton.
2. Apply tanh-cap (moe_router_logit_cap) and z-loss term in registers.
3. Apply the input jitter (apply_input_jitter, router.py:706) inline when moe_input_jitter_eps is set, from a counter-based RNG rather than a host-side torch.distributions.Uniform round-trip. (If a training run wants Gumbel/Gaussian logit noise instead of the stock uniform jitter, the same in-register RNG slot covers it.)
4. Compute clean_scores = sigmoid(logits) and noisy_scores = sigmoid(logits + tau * noise) (or the score-space variant) without rematerializing logits.
5. Add expert_bias to noisy_scores if present.
6. Top-K reduction across $E$ , emitting top_indices : [N, K] and gathered clean_scores for those indices.
7. Build routing_map : [N, E] bool by scatter_(top_indices, True) in the same kernel.

Where it plugs in. Replace RouterGatingLinearFunction.forward (moe_utils.py:1294) and most of topk_routing_with_score_function (moe_utils.py:710), at least the dropless / no-group-topk branch, which is the common case. The Sinkhorn / group-topk / capacity-padding branches stay as-is, routed through the legacy path. The forward becomes a one-call replacement. The backward keeps the existing RouterGatingLinearFunction.backward, since the gradient path need not be fused on round one.

Predicted gain. Each of these on its own is microseconds, but at small mbs the kernel-launch tax dominates idle time. From the GBS=512 / mbs=4 trace, GPU idle was 11% of step. Collapsing 5 kernels into 1 across 47 MoE layers and 256 microbatches per step removes ~50k kernel launches. At $\sim 5\mu s$ per launch that's 250 ms per step, order 0.5–1% MFU at GBS=8192 / mbs=4, more at small mbs.

Falsification. Nsight should show $E_{\text{kl,routing}}$ per step drop by ≥4×. If it doesn't, the dispatcher is launching them itself and the fusion is in the wrong place.

Risk. Numerics on the noise. Upstream apply_input_jitter draws from torch.distributions.Uniform on the host, so reproducing it exactly in a Triton kernel means matching its RNG stream. A counter-based RNG (Philox-4×32-10) seeded to track the same stream gives bit-identical fp32 floats. That is the path to take. A one-shot test sanity-checks that the new kernel's jitter matches the old path token-for-token when moe_input_jitter_eps > 0.

Why this and not "rewrite the whole MoE block in one Triton megakernel"? Diminishing returns and engineering cost. The two fusions above capture the two boundaries that dominate routing-side time.
The full megakernel (UniEP / FlashMoE) buys you the dispatch / grouped-GEMM boundary on top, which is where the expert time lives, not the routing time in focus here. That can be revisited once Experiment A is in.

References

• Nous Research, MoE Scaling Field Notes.
• Cursor, Kernels, MXFP8 grouped GEMM and Blackwell MoE.
• PyTorch ao, issue #3379, MoE mxfp8 next steps.
• LMSYS / SGLang, Large-Scale Expert Parallelism.
• Microsoft Azure, Achieving Optimal Performance for DeepEP on Azure.
• NVIDIA, Optimizing Communication for Mixture-of-Experts Training with Hybrid Expert-Parallel.
• Poolside, Laguna M.1 / XS.2 Technical Report, Section Compute-communication Overlap for MoE on flag-gated grouped-GEMM tile scheduling and epilogue-fused combine over NVLink.
• DeepSeek-AI, DeepSeek-V3 Technical Report, Sections 3.2.1 DualPipe and Computation-Communication Overlap and 3.2.2 Efficient Implementation of Cross-Node All-to-All Communication (20 SMs / 10 channels, warp-specialized IB↔NVLink kernel).
• DeepSeek-AI, DeepSeek-V4 Technical Report, Section 3.1 Fine-Grained Communication-Computation Overlap in Expert Parallelism (wave-pipeline MegaMoE open-sourced in DeepGEMM PR #304).
• DeepSeek-AI, DeepSeek-V3.2-Exp, DSA sparse attention, with no new MoE-comm material vs V3.
• Microsoft AI, MAI-Thinking-1, LatentMoE on 8K GB200, BF16/FP8 numerical recipe, CuTe DSL symmetric-memory NVLink dispatch, CLC use on Blackwell.
• FlashMoE: Fast Distributed MoE in a Single Kernel, arXiv:2506.04667.
• UniEP: Unified Expert Parallelism via MoE Megakernels, arXiv:2604.19241.
• DeepSeek, DeepSeek V4, Fine-Grained Communication-Computation Overlap in Expert Parallelism.
• Awesome ML-SYS Tutorial, RLHF system design notes 4 (English).
• Megatron-LM MoE source, megatron/core/transformer/moe/{moe_layer,token_dispatcher,router,moe_utils,fused_a2a,experts}.py from the public NVIDIA/Megatron-LM main tree (line numbers cited above are against main and will drift as the file evolves).