Mixture-of-Experts Across the Stack: From a 20-Line Reference to Wide Expert Parallelism

Code permalinks are pinned so line numbers stay valid: huggingface/transformers at commit effde20, vllm-project/vllm and NVIDIA/Megatron-LM at main (these two move fast — I cite by file + class/function name, which is stable, rather than line number). Where a framework was mid-refactor when I read it, I say so.

Why this post

Every frontier open model in 2025–2026 — DeepSeek-V3, Qwen3-MoE, Llama-4, Kimi, GLM — is a Mixture-of-Experts model. The idea is almost embarrassingly simple: replace one big FFN with many small “expert” FFNs and a router that sends each token to a few of them. But the gap between that one-sentence idea and a model running at EP320 across 320 GPUs is enormous, and it’s where most of the systems work in modern LLM infra now lives.

This post builds one mental model of MoE and then traces the exact same computation through four codebases at increasing levels of optimization:

1. a plain-PyTorch reference (what the math is),
2. HuggingFace Qwen3-MoE (the readable, single-GPU implementation),
3. vLLM (fused grouped-GEMM kernels for inference),
4. Megatron-Core (the training path, with grouped GEMM + token dispatchers).

Then we get to the part that actually matters at scale — expert parallelism (EP) and its all-to-all — and finish with the open problems people are publishing on right now.

---

1. The reference: what an MoE layer actually computes

An MoE layer replaces a transformer block’s single FFN with E expert FFNs plus a lightweight router. Three pieces:

Router (gate). A bias-free linear gate: hidden → E produces per-token logits. Those become routing weights via:

• scoring — softmax over the experts dim (Switch / Mixtral / Qwen3) or sigmoid per expert (DeepSeek-V3, where each expert’s score is independent rather than a distribution),
• top-k selection — keep the k highest-scoring experts per token,
• renormalize — optionally divide the kept weights by their sum so they sum to 1 (the norm_topk_prob flag).

Per-expert FFN (SwiGLU). Each expert is a gated MLP with three projections — gate_proj, up_proj (hidden → intermediate), and down_proj (intermediate → hidden):

\[\text{expert}(x) = W_{\text{down}}\,\Big(\,\mathrm{SiLU}(W_{\text{gate}}\,x)\,\odot\,W_{\text{up}}\,x\,\Big)\]

where $W_{\text{gate}}, W_{\text{up}}$ are gate_proj/up_proj and $W_{\text{down}}$ is down_proj.

Token grouping. The naive-but-correct pattern: flatten tokens, then for each expert, gather the tokens routed to it, run that expert’s MLP, scale by the routing weight, and scatter-add back. This is the entire game — and every optimization later is just a faster way to do exactly this.

import torch, torch.nn as nn, torch.nn.functional as F

class Expert(nn.Module):                       # SwiGLU FFN
    def __init__(self, h, i):
        super().__init__()
        self.gate_proj = nn.Linear(h, i, bias=False)
        self.up_proj   = nn.Linear(h, i, bias=False)
        self.down_proj = nn.Linear(i, h, bias=False)
    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

class ReferenceMoE(nn.Module):
    def __init__(self, h, i, E, k, norm_topk_prob=True):
        super().__init__()
        self.k, self.E, self.norm = k, E, norm_topk_prob
        self.gate    = nn.Linear(h, E, bias=False)
        self.experts = nn.ModuleList(Expert(h, i) for _ in range(E))

    def forward(self, x):                              # x: (B, T, H)
        B, T, H = x.shape
        x = x.view(-1, H)                              # (N, H)

        # --- router ---
        w = F.softmax(self.gate(x), dim=-1, dtype=torch.float)   # (N, E)
        w, idx = torch.topk(w, self.k, dim=-1)                   # (N, k)
        if self.norm:
            w = w / w.sum(dim=-1, keepdim=True)
        w = w.to(x.dtype)

        out  = torch.zeros_like(x)
        mask = F.one_hot(idx, self.E).permute(2, 1, 0)           # (E, k, N)
        for e in range(self.E):
            slot, tok = torch.where(mask[e])           # which top-k slot, which token
            if tok.numel() == 0:
                continue
            y = self.experts[e](x[tok]) * w[tok, slot, None]
            out.index_add_(0, tok, y.to(out.dtype))    # scatter-add
        return out.view(B, T, H)

That for e in range(E) loop with torch.where + index_add_ is the shape you’ll see literally in HuggingFace next. A few routing variants worth naming, because they show up constantly:

• Top-k softmax (Mixtral, Qwen3): softmax → top-k → optional renorm.
• Sigmoid scoring (DeepSeek-V3): per-expert sigmoid gate; weights renormalized separately.
• Group-limited / noaux_tc (DeepSeek-V3): experts are partitioned into groups; a bias-corrected score picks the top groups first, then top-k experts within them. This bounds how many devices a token can be dispatched to — a routing trick that is really a systems trick.
• Shared experts (DeepSeek, Qwen2-MoE): one always-on expert added to every token. Note: Qwen3-MoE dropped the shared expert that Qwen2-MoE had.

---

2. HuggingFace Qwen3-MoE: the readable implementation

HuggingFace is where you go to read an architecture. (Qwen3’s dense trunk — attention with q/k-norm + GQA, RoPE, RMSNorm, SwiGLU — gets a line-by-line treatment in a companion post; here we only add the MoE block.) The Qwen3-MoE block was recently refactored on main into dedicated router/expert classes (older tutorials describing one fused Qwen3MoeSparseMoeBlock.forward are now stale):

Class	File	Role
Qwen3MoeTopKRouter	modeling_qwen3_moe.py	bias-free gate → softmax(dim=-1) → topk → optional norm_topk_prob
Qwen3MoeExperts	modeling_qwen3_moe.py	the per-expert compute (eager loop or grouped GEMM)
Qwen3MoeSparseMoeBlock	modeling_qwen3_moe.py	reshape → router → experts → reshape

Two implementation details matter:

Weights are stored as 3D tensors, with gate+up fused. Instead of E separate nn.Linears, Qwen3MoeExperts holds gate_up_proj of shape (E, 2*intermediate, hidden) (gate and up concatenated, split later with .chunk(2, dim=-1)) and down_proj of shape (E, hidden, intermediate). Fusing gate+up into one tensor is the first step toward a single grouped GEMM.

The eager path is exactly the reference loop. Qwen3MoeExperts.forward builds one_hot(top_k_index).permute(2,1,0), finds which experts were hit, then loops torch.where(expert_mask[e]) → gather → MLP → scale → index_add_. If you understood §1, you’ve read it.

There is also an optimized path. A @use_experts_implementation decorator dispatches forward through an ExpertsInterface based on config._experts_implementation. The non-eager path (in integrations/moe.py) does the real-systems version: sort tokens by expert, compute per-expert offsets via histc+cumsum, run two grouped GEMMs (torch._grouped_mm, gated on torch>=2.9 and SM80/90), and accumulate with a deterministic view(N, k, H).sum(dim=1) instead of index_add_ (avoiding non-deterministic atomic-add). That is precisely the bridge to what vLLM and Megatron do in CUDA.

Qwen3-MoE’s defaults (configuration_qwen3_moe.py): num_experts=128, num_experts_per_tok=8, moe_intermediate_size=768, norm_topk_prob=False, decoder_sparse_step=1 (every eligible layer is MoE), and no shared_expert_* fields.

---

3. vLLM: one grouped-GEMM kernel for all experts

For inference, the per-expert Python loop is a non-starter — you want a single kernel. vLLM’s MoE subsystem (heavily refactored recently into experts/, router/, prepare_finalize/, runner/ subpackages — treat older flat-path docs as stale) is built around the FusedMoE layer, which owns all expert weights as stacked tensors w13 (gate+up) and w2 (down) and runs them through one grouped GEMM.

The trick that makes a single kernel possible is block-aligned expert sorting, in moe_align_block_size:

1. Sort the flattened (token, expert) assignments by expert.
2. Pad each expert’s token count up to a multiple of BLOCK_SIZE_M (padding slots get a sentinel token id).
3. Emit sorted_token_ids, expert_ids (one expert id per block), and num_tokens_post_padded.

Now fused_moe_kernel (a @triton.jit grouped GEMM) processes one BLOCK_SIZE_M-row block at a time; each block reads its expert_id to pick which expert’s weight slice to multiply. Padding blocks are skipped. Under expert parallelism, blocks for non-local experts get expert_id = -1 and are skipped — that one detail is how the same kernel serves both single-GPU and EP.

vLLM also has a modular-kernel abstraction (modular_kernel.py) that factors the layer into a pipeline — [Router] → [Quantize/Dispatch] → [Permute · Experts · Unpermute] → [Combine] — so any communication scheme (the all-to-all backends in §5) composes with any expert backend (Triton, CUTLASS, DeepGEMM, Marlin, FlashInfer, …) without a combinatorial explosion of code. The orchestrator class is FusedMoEKernel; routing strategies live behind a factory in router/ (fused_topk, DeepSeek-style grouped_topk, bias routers, …).

---

4. Megatron-Core: the training path

Megatron-Core’s MoE (megatron/core/transformer/moe/) is the same computation again, but staged for training and large-scale parallelism. MoELayer’s forward is explicitly route → preprocess → dispatch → compute → combine → postprocess, with the stages broken out so each can be CUDA-graphed and so the EP communication can overlap with expert compute.

Router — TopKRouter: gating() (fp32 router linear, marked sequence-parallel) → routing() producing probs + a boolean routing_map of shape [tokens, experts], plus token dropping and the load-balancing losses.

Experts — experts.py has two real implementations (a correction to a lot of stale docs: there is no standalone GroupedMLP class on main anymore — only a GroupedMLPSubmodules spec dataclass):

• TEGroupedMLP — the grouped-GEMM path via Transformer Engine GroupedLinear, keyed by tokens_per_expert, with a fused weighted-SwiGLU activation. Selected when moe_grouped_gemm=True.
• SequentialMLP — a Python loop running each local expert as its own MLP (the fallback, and the Local-backend default).

Token dispatchers — token_dispatcher.py is where the parallelism lives (§5). Three flavors: MoEAllGatherTokenDispatcher (allgather over TP×EP), MoEAlltoAllTokenDispatcher (the standard EP>1 path), and the newer MoEFlexTokenDispatcher with pluggable _DeepepManager (DeepSeek DeepEP) and _HybridEPManager (NVIDIA HybridEP for GB200/MNNVL) comm backends.

Load balancing is a first-class config: moe_router_load_balancing_type ∈ {aux_loss, seq_aux_loss, global_aux_loss, sinkhorn, none} (or a list to combine). The shared kernel switch_load_balancing_loss_func computes the Switch-Transformer loss $\mathcal{L} = E\sum_i f_i \cdot P_i$ (fraction of tokens to expert $i$ × mean router prob for $i$), injected into the graph via MoEAuxLossAutoScaler. There’s also DeepSeek-V3-style aux-loss-free balancing via moe_router_enable_expert_bias, and capacity/dropless control via moe_expert_capacity_factor (None ⇒ dropless) and moe_pad_expert_input_to_capacity (static shapes for CUDA graphs).

---

5. Expert parallelism: the all-to-all dance

So far everything was one GPU. Expert parallelism (EP) shards the experts themselves across GPUs: with E experts and EP size W, each rank owns E/W experts. The catch is that a token’s chosen experts usually live on other GPUs. So every MoE layer needs two collectives:

            ┌─────────── dispatch (all-to-all) ──────────┐
 router →  send each token to the GPU(s) owning its experts
            → local expert GEMMs (gate/up → SiLU → down)
            → send results back ──── combine (all-to-all) ─┘ → weighted sum

This dispatch / combine all-to-all is the defining cost of MoE at scale — and it sits in two different regimes: prefill/training is bandwidth-bound (huge token batches that can saturate NVLink/RDMA), while decode is latency-bound (a handful of tokens per step, dominated by launch + round-trip latency). It’s also a hard data dependency: you cannot start the expert GEMMs until the dispatch lands.

EP in vLLM

Enabled with --enable-expert-parallel. The shard map is computed by determine_expert_map → a global→local table with -1 for non-local experts (the same -1 the grouped-GEMM kernel skips). The all-to-all itself is a pluggable backend, selected by ParallelConfig.all2all_backend and built by a maybe_make_prepare_finalize factory into a FusedMoEPrepareAndFinalize:

• deepep_high_throughput — DeepEP “normal” kernels, contiguous/permuted layout, max bandwidth, for prefill.
• deepep_low_latency — DeepEP low-latency kernels (RDMA, masked/batched layout, compiled per hidden size), for decode.
• allgather_reducescatter (default), plus mori_*, nixl_ep, flashinfer_nvlink_*.
• (Historical note: the pplx backend was removed; the literal now remaps to allgather_reducescatter. If you read an older blog mentioning pplx_prepare_finalize.py, it’s gone.)

vLLM also implements EPLB (Expert-Parallel Load Balancing) — EplbState tracks per-expert load and periodically rearranges experts, replicating hot ones as redundant experts across ranks. More on why below.

EP in Megatron

MoEAlltoAllTokenDispatcher’s documented workflow is preprocess → permute → A2A(EP) → AG(TP) → [experts] → RS(TP) → A2A(EP) → unpermute. The subtle part is preprocess: because all-to-all split sizes are dynamic (they depend on routing), it has to compute input_splits/output_splits from the per-expert token counts and carefully schedule the device→host syncs to read them back. Megatron also separates EP (expert_model_parallel_size) from expert tensor parallelism / ETP (expert_tensor_parallel_size), with experts using a distinct expt_tp process group from attention’s TP — so the expert GEMMs can be sharded independently of attention.

The common thread across both: dispatch and combine are explicit collectives, and hiding them behind compute is the whole optimization target.

---

6. The frontier: what the industry is fighting now

This is where MoE research has actually moved — not the architecture, but the systems.

(1) The all-to-all bottleneck → overlap and specialized kernels. Because dispatch/combine dominates, the work is in (a) faster collectives and (b) hiding them. DeepEP (DeepSeek’s EP communication library) ships separate normal (throughput) and low-latency (decode) dispatch/combine kernels — originally NVSHMEM-based, with the transport layer reworked in later releases. On top of that, Two-Batch Overlap (TBO) splits a batch into micro-batches so one micro-batch’s expert compute hides the other’s all-to-all; SGLang reported 27–35% prefill throughput from TBO. The endgame is fusing the collective into the GEMM entirely.

(2) Load balancing → auxiliary-loss-free. Routers collapse toward a few hot experts, wasting capacity and creating EP stragglers. The classic auxiliary load-balancing loss injects interference gradients that fight model quality. DeepSeek’s Loss-Free Balancing (arXiv:2408.15664) instead adds a per-expert bias to routing scores before top-k, nudged dynamically by recent utilization — no aux-loss term, and it’s what DeepSeek-V3 uses for stable training. At inference, EPLB duplicates hot experts as redundant replicas (hierarchical for prefill, global for decode); SGLang measured 1.49× prefill / 2.54× decode from it.

(3) Wide EP + prefill/decode disaggregation. Decode is weight-bandwidth-bound, so spreading experts across more GPUs (“wide EP”) shrinks per-GPU weight pressure — at the cost of more all-to-all latency and more imbalance. The enabling move is running prefill and decode on separate pools with different EP sizes. DeepSeek’s own deployment (hardware paper, arXiv:2505.09343) disaggregates with large-EP decode (their reported figures put decode around EP320); SGLang reproduced the economics at prefill EP32 / decode EP72 on 96 H100s for roughly $0.20 / 1M output tokens; and NVIDIA’s wide-EP on GB200 NVL72 reports up to 1.8× output tokens/s/GPU at EP32 vs EP8 under a fixed latency SLA.

(4) Fine-grained / shared experts + FP8. DeepSeekMoE (arXiv:2401.06066) introduced fine-grained expert segmentation (more, smaller experts → finer routing) plus shared-expert isolation (a few always-on experts absorb common knowledge). Once experts are tiny, quantization and kernel design dominate: DeepSeek-V3 trains in FP8 with block-wise scaling, and DeepGEMM implements the matching FP8 grouped GEMMs in contiguous (prefill) and masked (decode) layouts — reportedly up to ~1550 TFLOPS on H800. The kernel side of this — CuTe layouts, grouped GEMM, the Hopper/Blackwell instruction story — is the subject of the CUTLASS deep dive.

(5) Stability and memory. ST-MoE (arXiv:2202.08906) added the router z-loss to keep router logits from blowing up — still the foundational training-stability tool. And on the memory side, MoE pairs naturally with KV-cache compression (DeepSeek-V3’s MLA): sparse FFN params spread across GPUs, compressed KV per GPU.

---

The one-paragraph summary

MoE is “many FFNs + a router, run a few per token.” A reference implementation is a 20-line loop; HuggingFace’s eager path is that same loop; vLLM and Megatron replace the loop with a single block-aligned grouped GEMM; and the moment experts span GPUs, the entire game becomes the dispatch/combine all-to-all and how completely you can hide it behind compute. Everything the field is publishing right now — DeepEP, TBO, EPLB, wide-EP disaggregation, FP8 grouped GEMMs — is some angle of attack on that one collective.

---

Revisions read: transformers effde20, vLLM main, Megatron-LM main (June 2026). Framework MoE code moves fast; cite by class/function name and re-pin before relying on a line number.