mini-sglang from Scratch Stage 4 · Faster and bigger

MoE and fused MoE

An MoE model like Qwen3-30B-A3B has hundreds of experts per layer and activates only a few per token. The hard part is that the tokens in a batch go to different experts, and computing expert by expert would launch hundreds of small matmuls. Fused MoE sorts all the "(token, expert) pairs" by expert and pads each expert's run to a multiple of the block size, so a whole MoE layer needs only two kernel launches. This chapter implements the MoE layer, a reference backend and a fused MoE written in Triton, and verifies it on a CPU with the Triton interpreter.

Self-test: if you can answer these, skip the chapter

1. How is MoE routing computed? What does norm_topk_prob mean?
2. After sorting the (token, expert) pairs by expert, why is each expert's run padded to a multiple of BLOCK_M?
3. Under tensor parallelism, how are the expert weights sharded? What communication is needed?
4. How do you run a Triton kernel on a machine with no GPU?

Answers (try it yourself first, then expand)

1. The router computes a score per expert, and after a softmax the top-k experts and their weights are taken; norm_topk_prob means renormalizing those k weights (dividing by their sum) so they add up to 1.
2. In the fused MoE kernel each output tile (BLOCK_M rows) can belong to only one expert and multiplies by that expert's weight, so padding each expert's run to a multiple of BLOCK_M keeps a tile from straddling two experts.
3. Every expert's weights split along the intermediate dimension: gate / up by column, down by row, so each rank holds part of every expert. The router is not split (every rank computes the full routing). One all-reduce at the end of the layer adds the partial sums.
4. Use Triton's interpreter mode (TRITON_INTERPRET=1): the kernel runs tile by tile on the CPU with numpy and the same semantics, which verifies correctness (performance means nothing there).

Files you will write: layers/moe.py, moe/base.py, moe/torch_backend.py, moe/fused.py, moe/__init__.py, and MoEMLP in models/utils.py.

Files so far: 79 (5 new, 1 modified, 0 brought in early)

minisgl/
├── __init__.py
├── core.py
├── env.py
├── __main__.py
├── shell.py
├── attention/  (6 files, unchanged)
├── distributed/  (3 files, unchanged)
├── engine/  (5 files, unchanged)
├── kernel/  (3 files, unchanged)
├── kvcache/  (5 files, unchanged)
├── layers/
│   ├── __init__.py
│   ├── base.py
│   ├── linear.py
│   ├── embedding.py
│   ├── norm.py
│   ├── rotary.py
│   ├── activation.py
│   ├── attention.py
│   ├── moe.py  ← new in this chapter
├── llm/  (2 files, unchanged)
├── message/  (5 files, unchanged)
├── models/
│   ├── config.py
│   ├── base.py
│   ├── utils.py  ← modified in this chapter
│   ├── decoder.py
│   ├── __init__.py
│   ├── register.py
│   ├── weight.py
├── moe/
│   ├── __init__.py  ← new in this chapter
│   ├── base.py  ← new in this chapter
│   ├── torch_backend.py  ← new in this chapter
│   ├── fused.py  ← new in this chapter
├── scheduler/  (9 files, unchanged)
├── server/  (4 files, unchanged)
├── tokenizer/  (4 files, unchanged)
├── utils/  (8 files, unchanged)

This step's main: examples/ch20_moe.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

The MoE layer

minisgl/models/utils.py · MoEMLP

class MoEMLP(BaseOP):
    def __init__(self, config: ModelConfig):
        from minisgl.layers.moe import MoELayer  # chapter 20's file: imported only when MoE is used

        self.experts = MoELayer(
            num_experts=config.num_experts,
            top_k=config.num_experts_per_tok,
            hidden_size=config.hidden_size,
            intermediate_size=config.moe_intermediate_size,
            renormalize=config.norm_topk_prob,
        )
        self.gate = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.experts.forward(hidden_states=x, router_logits=self.gate.forward(x))

gate is an unsharded linear layer (a full copy per rank) that scores every token against every expert. MoELayer holds all the experts' packed weights and leaves the computation to the MoE backend in the Context:

minisgl/layers/moe.py · MoELayer

class MoELayer(BaseOP):
    """Every expert's weights packed into two 3-dimensional tensors; how to compute them is left to the
    MoE backend in the Context.

    gate_up_proj: [E, 2 * I_local, H] (each expert's gate and up concatenated)
    down_proj:    [E, H, I_local]
    Under tensor parallelism each expert's intermediate dimension I is split (as in a dense MLP),
    with an all-reduce at the end.
    """

    def __init__(self, num_experts: int, top_k: int, hidden_size: int, intermediate_size: int,
                 renormalize: bool = True):
        self.num_experts = num_experts
        self.top_k = top_k
        self.renormalize = renormalize
        self._comm = DistributedCommunicator()
        self.tp_size = get_tp_info().size
        local = div_even(intermediate_size, self.tp_size)
        self.gate_up_proj = torch.empty(num_experts, 2 * local, hidden_size)
        self.down_proj = torch.empty(num_experts, hidden_size, local)

    def forward(self, hidden_states: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor:
        out = get_global_ctx().moe_backend.forward(
            hidden_states=hidden_states,
            w1=self.gate_up_proj,
            w2=self.down_proj,
            gating_output=router_logits,
            topk=self.top_k,
            renormalize=self.renormalize,
        )
        return self._comm.all_reduce(out) if self.tp_size > 1 else out

gate_up_proj has shape [experts, 2 × intermediate, hidden], with each expert's gate and up concatenated (just like a dense MLP). Under tensor parallelism every expert's intermediate dimension is split, exactly as in a dense MLP: column parallel first, row parallel second, one all-reduce at the end. The weight loader reads the separately stored experts.0.gate_proj, experts.1.gate_proj and so on from the checkpoint and packs them into a three-dimensional tensor (chapter 3).

Routing

minisgl/moe/base.py · select_experts

def select_experts(gating_output: torch.Tensor, topk: int,
                   renormalize: bool) -> Tuple[torch.Tensor, torch.Tensor]:
    """Routing: take the top-k experts after a softmax, optionally renormalizing those k weights to sum to 1."""
    probs = torch.softmax(gating_output.float(), dim=-1)
    topk_weights, topk_ids = probs.topk(topk, dim=-1)
    if renormalize:
        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
    return topk_weights, topk_ids

Take the top k experts after a softmax; with norm_topk_prob=True (Qwen3-MoE) those k weights are renormalized to sum to 1.

The reference backend

minisgl/moe/torch_backend.py · TorchMoeBackend

class TorchMoeBackend(BaseMoeBackend):
    """The reference implementation: grouped by expert, with each expert handling only the tokens routed to it."""

    def forward(self, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,
                gating_output: torch.Tensor, topk: int, renormalize: bool) -> torch.Tensor:
        topk_weights, topk_ids = select_experts(gating_output, topk, renormalize)
        out = torch.zeros_like(hidden_states)
        flat_ids = topk_ids.flatten()                       # [T * k]
        token_idx = torch.arange(hidden_states.shape[0], device=hidden_states.device)
        token_idx = token_idx.repeat_interleave(topk)       # which token element i belongs to
        flat_weights = topk_weights.flatten().to(hidden_states.dtype)
        for e in flat_ids.unique().tolist():
            sel = (flat_ids == e).nonzero().squeeze(1)
            tokens = token_idx[sel]
            h = silu_and_mul(hidden_states[tokens] @ w1[e].t())
            y = (h @ w2[e].t()) * flat_weights[sel].unsqueeze(1)
            out.index_add_(0, tokens, y)
        return out

Grouped by expert: for each selected expert, pick out the tokens routed to it, run one SwiGLU MLP, multiply by the routing weight, and accumulate back onto the tokens with index_add_. Straightforward, but as many rounds of small matmuls as there are experts.

Fused MoE

sort by expert, pad, then two fused GEMMs

minisgl/moe/fused.py · moe_align_block_size

def moe_align_block_size(topk_ids: torch.Tensor, block_m: int,
                         num_experts: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """Sort the (token, expert) pairs by expert and pad each expert's run to a multiple of block_m.

    Returns sorted_ids (indices into topk_ids.flatten(), with num_pairs as the padding value) and
    each tile's expert id.
    """
    flat = topk_ids.flatten().long()
    num_pairs = flat.numel()
    counts = torch.bincount(flat, minlength=num_experts)
    padded = (counts + block_m - 1) // block_m * block_m
    seg_start = torch.cumsum(padded, 0) - padded  # where each expert's run starts in the output
    order = torch.argsort(flat, stable=True)  # the pair indices after sorting by expert
    sorted_expert = flat[order]
    rank = torch.arange(num_pairs, device=flat.device) - (torch.cumsum(counts, 0) - counts)[sorted_expert]
    sorted_ids = torch.full((int(padded.sum()),), num_pairs, dtype=torch.int32, device=flat.device)
    sorted_ids[seg_start[sorted_expert] + rank] = order.to(torch.int32)
    expert_ids = torch.repeat_interleave(torch.arange(num_experts, device=flat.device),
                                         padded // block_m).to(torch.int32)
    return sorted_ids, expert_ids

The first step sorts all the (token, expert) pairs by expert and pads each expert's run to a multiple of BLOCK_M, filling the padding with a placeholder (equal to the pair count). Every BLOCK_M consecutive pairs then belong to one expert, so a thread block can handle the matmul of those BLOCK_M tokens against one expert's weights.

minisgl/moe/fused.py · fused_moe_kernel

@triton.jit
def fused_moe_kernel(
    a_ptr, b_ptr, c_ptr, topk_weights_ptr, sorted_token_ids_ptr, expert_ids_ptr,
    N, K, num_valid_tokens,
    stride_am, stride_ak, stride_be, stride_bn, stride_bk, stride_cm, stride_cn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    MUL_ROUTED_WEIGHT: tl.constexpr, TOP_K: tl.constexpr,
):
    """C[t, :] = A[t // TOP_K, :] @ B[expert].T (optionally times the routing weight); t indexes the sorted (token, expert) pairs."""
    pid = tl.program_id(0)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    pid_m, pid_n = pid // num_pid_n, pid % num_pid_n
    offs_token = tl.load(sorted_token_ids_ptr + pid_m * BLOCK_M + tl.arange(0, BLOCK_M))
    token_mask = offs_token < num_valid_tokens  # the padding placeholders take no part in the computation
    expert = tl.load(expert_ids_ptr + pid_m)  # which expert this tile belongs to
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    a_ptrs = a_ptr + (offs_token[:, None] // TOP_K) * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = b_ptr + expert * stride_be + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        k_remain = K - k * BLOCK_K
        a = tl.load(a_ptrs, mask=token_mask[:, None] & (offs_k[None, :] < k_remain), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remain) & (offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk
    if MUL_ROUTED_WEIGHT:
        w = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)
        acc = acc * w[:, None]
    c_ptrs = c_ptr + offs_token[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=token_mask[:, None] & (offs_n[None, :] < N))

Each program instance owns one [BLOCK_M, BLOCK_N] tile of the output: it reads that tile's expert id, gathers BLOCK_M input rows by the sorted indices (in the first GEMM a token corresponds to TOP_K pairs, so the row is pair index // TOP_K), and does a tiled matmul against that expert's weights, with the placeholder rows masked out. The second GEMM multiplies by the routing weight along the way.

minisgl/moe/fused.py · FusedMoeBackend.forward

def forward(self, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,
            gating_output: torch.Tensor, topk: int, renormalize: bool) -> torch.Tensor:
    num_tokens, hidden = hidden_states.shape
    topk_weights, topk_ids = select_experts(gating_output, topk, renormalize)
    num_pairs = topk_ids.numel()
    sorted_ids, expert_ids = moe_align_block_size(topk_ids, self.block_m, w1.shape[0])
    weights = topk_weights.flatten().to(hidden_states.dtype).contiguous()
    cache1 = hidden_states.new_empty(num_pairs, w1.shape[1])  # [T*k, 2I]
    _invoke(hidden_states, w1, cache1, weights, sorted_ids, expert_ids, False, topk,
            num_pairs, self.block_m)
    cache2 = silu_and_mul(cache1)  # [T*k, I]
    cache3 = hidden_states.new_empty(num_pairs, hidden)  # [T*k, H], already times the routing weight
    _invoke(cache2, w2, cache3, weights, sorted_ids, expert_ids, True, 1,
            num_pairs, self.block_m)
    return cache3.view(num_tokens, topk, hidden).sum(dim=1)

A whole MoE layer: sort (in PyTorch), the first fused GEMM (w1), SiLU × up, the second fused GEMM (w2, times the routing weight), then sum the k results per token.

Running Triton on a CPU

Triton has an interpreter mode (TRITON_INTERPRET=1): the kernel is interpreted program instance by program instance on the CPU with NumPy, with the same semantics as on a GPU, only very slow. It must be set before triton is imported.

examples/ch20_moe.py

"""Chapter 20: MoE routing, sorting and padding by expert, and the Triton fused MoE against the reference.

The Triton kernel runs on a CPU in interpreter mode (TRITON_INTERPRET=1 must be set before triton
is imported).
"""

import os

os.environ["TRITON_INTERPRET"] = "1"

import torch  # noqa: E402
from minisgl.moe.base import select_experts  # noqa: E402
from minisgl.moe.fused import FusedMoeBackend, moe_align_block_size  # noqa: E402
from minisgl.moe.torch_backend import TorchMoeBackend  # noqa: E402

torch.manual_seed(0)
T, H, I, E, K = 5, 64, 48, 8, 2
x, gating = torch.randn(T, H), torch.randn(T, E)
weights, ids = select_experts(gating, K, renormalize=True)
print("the experts each token picked:", ids.tolist())
print("routing weights (after renormalization):", [[round(w, 3) for w in row] for row in weights.tolist()])
sorted_ids, expert_ids = moe_align_block_size(ids, block_m=4, num_experts=E)
print("sorted by expert, padded to a multiple of 4 (10 is the placeholder):", sorted_ids.tolist())
print("which expert each tile belongs to:", expert_ids.tolist())

x, gating = torch.randn(37, H), torch.randn(37, E)
w1, w2 = torch.randn(E, 2 * I, H) * 0.1, torch.randn(E, H, I) * 0.1
a = FusedMoeBackend().forward(x, w1, w2, gating, K, renormalize=True)
b = TorchMoeBackend().forward(x, w1, w2, gating, K, renormalize=True)
print(f"37 tokens: fused (Triton interpreter) against the reference, max error {(a - b).abs().max().item():.1e}")

What it prints

the experts each token picked: [[4, 3], [6, 1], [1, 7], [6, 4], [5, 4]]
routing weights (after renormalization): [[0.692, 0.308], [0.545, 0.455], [0.93, 0.07], [0.711, 0.289], [0.809, 0.191]]
sorted by expert, padded to a multiple of 4 (10 is the placeholder): [3, 4, 10, 10, 1, 10, 10, 10, 0, 7, 9, 10, 8, 10, 10, 10, 2, 6, 10, 10, 5, 10, 10, 10]
which expert each tile belongs to: [1, 3, 4, 5, 6, 7]
37 tokens: fused (Triton interpreter) against the reference, max error 2.4e-07

5 tokens choosing 2 experts each make 10 pairs. After sorting, expert 1's two pairs (the 3rd and 4th) form the first tile with two placeholders added; expert 3 has only one pair and fills a tile as well, for 6 tiles in all. A bigger block wastes more on padding; a smaller one makes the matmul less efficient. The fused MoE on 37 tokens differs from the reference by about float32 rounding.

End to end

A small two-layer, 8-expert Qwen3-MoE model built from random weights is compared against the Hugging Face implementation (the reference backend is the default on a CPU):

tests/test_ch20_moe.py · test_tiny_qwen3_moe_matches_hf

def test_tiny_qwen3_moe_matches_hf(tiny_models):
    path = tiny_models["qwen3moe"]
    llm = build_llm(path)
    prompts = ["hello world", "The quick brown fox"]
    out = llm.generate(prompts, greedy(10))
    assert [o["token_ids"] for o in out] == hf_greedy(path, prompts, 10)
    llm.shutdown()

The official implementation

• layers/moe.py · MoELayer (lines 9-59), models/qwen3_moe.py
• fused MoE: moe/fused.py and kernel/triton/fused_moe.py. The routing (topk_softmax) and the sorting and padding (moe_align_block_size) use CUDA kernels from sgl_kernel, and the block size is one of two choices by shape: small tiles (16 × 32 × 64) when the token count does not exceed the expert count, large ones (64 × 64 × 32) otherwise, which is exactly what exercise 1 is about

Differences from upstream

Our Triton kernel is a simplified version: the block sizes are fixed (BLOCK_M=16, BLOCK_N=32, BLOCK_K=32) rather than chosen by shape, there is no GROUP_SIZE_M grouping (which improves L2 hit rate) as in the official kernel, and the sorting and padding use PyTorch rather than a CUDA kernel. The torch reference backend is ours and is the default on a CPU.

Tests

tests/test_ch20_moe.py · test_fused_triton_moe_in_interpreter

@pytest.mark.skipif(importlib.util.find_spec("triton") is None, reason="没有安装 Triton（macOS 上没有 Triton 的安装包）")
def test_fused_triton_moe_in_interpreter():
    code = """
import torch
from minisgl.moe.fused import FusedMoeBackend, moe_align_block_size
from minisgl.moe.torch_backend import TorchMoeBackend
torch.manual_seed(0)
x, g = torch.randn(37, 64), torch.randn(37, 8)
w1, w2 = torch.randn(8, 96, 64) * 0.1, torch.randn(8, 64, 48) * 0.1
a = FusedMoeBackend().forward(x, w1, w2, g, 2, True)
b = TorchMoeBackend().forward(x, w1, w2, g, 2, True)
print((a - b).abs().max().item())
"""
    env = dict(os.environ, TRITON_INTERPRET="1", PYTHONPATH=str(ROOT / "python"))
    r = subprocess.run([sys.executable, "-c", code], env=env, capture_output=True, text=True, timeout=600)
    assert r.returncode == 0, r.stderr
    assert float(r.stdout.strip()) < 1e-5

tests/test_ch20_moe.py also verifies the reference backend against the most literal "per token, per expert" loop, and Llama 3's long-context RoPE from chapter 2 (again a small random-weight model compared with Hugging Face).

How to explain it

On MoE: routing is softmax, then top-k, then an optional renormalization (norm_topk_prob makes the k selected weights sum to 1). Fused MoE sorts the (token, expert) pairs by expert and pads each expert's run to a multiple of BLOCK_M, so every output tile belongs to one expert and a single kernel computes every expert's GEMM (two launches per layer, gate_up and down) instead of one launch per expert. Under tensor parallelism each expert is split along the intermediate dimension, like a dense MLP, with one all-reduce; expert parallelism splits by expert instead and needs an all-to-all. Without a GPU, Triton's interpreter mode runs the kernel on a CPU with the same semantics to verify correctness.

Exercises

1. With BLOCK_M=64, 128 experts, and only 32 tokens per decode step (each choosing 8 experts), how many pairs are there after padding? What fraction is wasted? What does that mean for decode?
2. Expert parallelism (EP) puts different experts on different ranks, so every token has to travel to its expert's rank. How does its communication pattern differ from this chapter's tensor-parallel split? (See the expert-parallelism chapter of the Inference Systems handbook.)
3. Fuse SiLU × up into the first GEMM's output: the kernel has to compute both the gate and the up tile. How would you arrange offs_n so that one program instance gets the matching gate and up columns?

Answers

1. 256 pairs spread over at most 128 experts, so most experts have only 1 to 3 pairs and each pads to 64: at worst about 128 × 64 = 8192 rows for 256 useful ones, around 97% wasted. The fused MoE kernel is very inefficient during decode, which is why decode often uses smaller tiles or kernels designed for small batches (a GEMV per expert, for instance).
2. Tensor parallelism gives every rank part of every expert and communicates with one all-reduce (whose volume is proportional to tokens × hidden); expert parallelism gives every rank all of some experts and needs an all-to-all to send tokens to their experts' ranks and gather the results back, with a volume that depends on the routing distribution and can be badly unbalanced.
3. Have the program instance own a stretch [n, n + BLOCK_N) of the intermediate dimension, load the gate columns [n, n+BLOCK_N) and the up columns [I + n, I + n + BLOCK_N) together, do two tl.dots, and write out [BLOCK_M, BLOCK_N] after silu(gate) * up.

Summary

• The MoE layer: an unsharded router plus expert weights packed into a three-dimensional tensor; under tensor parallelism each expert splits along the intermediate dimension with one all-reduce.
• Routing: softmax, top-k, optional renormalization.
• Fused MoE: sort by expert and pad to the block size so every output tile belongs to one expert, giving two kernel launches per layer.
• The Triton interpreter runs the kernel on a CPU with the same semantics, which is enough to verify correctness.

Exercises for this chapter (1)

Write the code and run the tests in your browser, or judge locally with python practice/judge.py. The exercises are in Chinese for now.

• 203. moe_align_block_size 与分组 GEMM (hard)

Found a mistake or a better way to say it? Edit this page on GitHub · Report an issue