手写 mini-sglang 阶段四 · 更快、更大

MoE 与 fused MoE¶

Qwen3-30B-A3B 这样的 MoE 模型，每层有上百个专家，每个 token 只激活其中几个。难点在于：一个 batch 里的 token 各自去往不同的专家，按专家逐个计算会启动上百个小矩阵乘。fused MoE 把所有"(token, 专家) 对"按专家排好序，每个专家的那一段补齐到块大小的整数倍，于是整层 MoE 只需两次 kernel 启动。这一章实现 MoE 层、参考后端和一个 Triton 写的 fused MoE，并在 CPU 上用 Triton 解释器验证。

自测：能答上来就可以跳过本章

MoE 的路由是怎么算的？norm_topk_prob 是什么意思？
为什么把 (token, 专家) 对按专家排序之后，每个专家的段要补齐到 BLOCK_M 的整数倍？
张量并行时，MoE 的专家权重怎么切？需要什么通信？
在没有 GPU 的机器上怎么运行 Triton kernel？

自测参考答案（先自己答，再展开对照）

路由层算出每个专家的分数，softmax 之后选 top-k 个专家和权重；norm_topk_prob 表示把选出的 k 个权重重新归一化（除以它们的和），让它们加起来是 1。
fused MoE kernel 里每个输出块（BLOCK_M 行）只能属于一个专家、用这个专家的权重做矩阵乘；把每个专家的段补齐到 BLOCK_M 的整数倍，块才不会跨两个专家。
每个专家的权重都按中间维切：gate / up 按列切、down 按行切，每个 rank 持有所有专家的一部分；路由层不切（每个 rank 都算完整的路由）。每层最后一次 all-reduce 把部分和加起来。
用 Triton 的解释器模式（TRITON_INTERPRET=1）：kernel 在 CPU 上用 numpy 以相同的语义逐块运行，用来验证正确性（性能没有意义）。

本章要写的文件：layers/moe.py、moe/base.py、moe/torch_backend.py、moe/fused.py、moe/__init__.py，以及 models/utils.py 中的 MoEMLP。

到这一章为止的文件：79 个（新建 5，修改 1，提前引入 0）

minisgl/
├── __init__.py
├── core.py
├── env.py
├── __main__.py
├── shell.py
├── attention/  （6 个文件，不变）
├── distributed/  （3 个文件，不变）
├── engine/  （5 个文件，不变）
├── kernel/  （3 个文件，不变）
├── kvcache/  （5 个文件，不变）
├── layers/
│   ├── __init__.py
│   ├── base.py
│   ├── linear.py
│   ├── embedding.py
│   ├── norm.py
│   ├── rotary.py
│   ├── activation.py
│   ├── attention.py
│   ├── moe.py  ← 本章新建
├── llm/  （2 个文件，不变）
├── message/  （5 个文件，不变）
├── models/
│   ├── config.py
│   ├── base.py
│   ├── utils.py  ← 本章修改
│   ├── decoder.py
│   ├── __init__.py
│   ├── register.py
│   ├── weight.py
├── moe/
│   ├── __init__.py  ← 本章新建
│   ├── base.py  ← 本章新建
│   ├── torch_backend.py  ← 本章新建
│   ├── fused.py  ← 本章新建
├── scheduler/  （9 个文件，不变）
├── server/  （4 个文件，不变）
├── tokenizer/  （4 个文件，不变）
├── utils/  （8 个文件，不变）

这一步的 main：examples/ch20_moe.py——它只用到上面这些文件；python tools/steps.py check 会逐章搭出这棵树、跑这个 main。

MoE 层¶

minisgl/models/utils.py · MoEMLP

class MoEMLP(BaseOP):
    def __init__(self, config: ModelConfig):
        from minisgl.layers.moe import MoELayer  # 第 20 章的文件：用到 MoE 时才导入

        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 是一个不切分的线性层（每个 rank 一份完整的），算出每个 token 对每个专家的打分。MoELayer 持有所有专家打包后的权重，计算交给 Context 里的 MoE 后端：

minisgl/layers/moe.py · MoELayer

gate_up_proj 形状 [专家数, 2 × 中间维, hidden]，每个专家的 gate n>

class MoELayer(BaseOP): """所有专家的权重打包成两个 3 维张量；具体怎么算交给 Context 里的 MoE 后端。 class="sd">    gate_up_proj: [E, 2 * I_local, H]（每个专家的 gate 和 up 拼在一起） class="sd">    down_proj:    [E, H, I_local] class="sd">    张量并行时每个专家的中间维 I 被切开（和普通 MLP 一样），最后 all-reduce。 class="sd">    """ 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 和 up 拼在一起（与稠密 MLP 相同）。张量并行时，每个专家的中间维都被切开，和稠密 MLP 一样：先列并行、再行并行，最后一次 all-reduce。权重加载器把 checkpoint 里分开存的 experts.0.gate_proj、experts.1.gate_proj……读齐后打包成三维张量（第 3 章）。
路由¶
minisgl/moe/base.py · select_expertsdef select_experts(gating_output: torch.Tensor, topk: int,
                   renormalize: bool) -> Tuple[torch.Tensor, torch.Tensor]:
    """路由：softmax 后取前 top-k 个专家，可选地把这 k 个权重重新归一化到和为 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

softmax 后取前 k 个专家；norm_topk_prob=True（Qwen3-MoE）时把这 k 个权重重新归一化到和为 1。
参考后端¶
minisgl/moe/torch_backend.py · TorchMoeBackendclass TorchMoeBackend(BaseMoeBackend):
    """参考实现：按专家分组，每个专家只处理路由到它的那些 token。"""

    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)       # 第 i 个元素属于哪个 token
        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

按专家分组：对每个被选中的专家，挑出路由到它的 token，做一次 SwiGLU MLP，乘上路由权重，用 index_add_ 累加回各自的 token。直白，但专家有多少个，就有多少轮小矩阵乘。
fused MoE¶
按专家排序、补齐，再做两次 fused GEMM
minisgl/moe/fused.py · moe_align_block_sizedef moe_align_block_size(topk_ids: torch.Tensor, block_m: int,
                         num_experts: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """把 (token, 专家) 对按专家排序，每个专家的段补齐到 block_m 的倍数。

    返回 sorted_ids（元素是 topk_ids.flatten() 中的下标，补齐位填 num_pairs）和每个块的专家号。
    """
    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  # 每个专家的段在输出中的起点
    order = torch.argsort(flat, stable=True)  # 按专家排序后的 pair 下标
    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

第一步把所有 (token, 专家) 对按专家排序，每个专家的段补齐到 BLOCK_M 的倍数，补齐的位置填一个占位值（等于对数）。于是每 BLOCK_M 个连续的对属于同一个专家，一个线程块可以处理这 BLOCK_M 个 token 与一个专家权重的矩阵乘。
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[专家].T（可选再乘路由权重），t 是排序后的 (token, 专家) 对编号。"""
    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  # 补齐用的占位符不参与计算
    expert = tl.load(expert_ids_ptr + pid_m)  # 这个块属于哪个专家
    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))

kernel 的每个程序实例负责输出的一个 [BLOCK_M, BLOCK_N] 块：读这个块对应的专家号，按排序后的下标取出 BLOCK_M 行输入（第一次 GEMM 时一个 token 对应 TOP_K 个对，所以行号是 对编号 // TOP_K），与该专家的权重做分块矩阵乘；占位符的行被掩掉。第二次 GEMM 时顺便乘上路由权重。
minisgl/moe/fused.py · FusedMoeBackend.forwarddef 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]，已乘路由权重
    _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)

整层 MoE：排序（PyTorch）→ 第一次 fused GEMM（w1）→ SiLU × up → 第二次 fused GEMM（w2，乘路由权重）→ 按 token 把 k 个结果相加。
在 CPU 上运行 Triton¶
Triton 提供解释器模式（TRITON_INTERPRET=1）：kernel 用 NumPy 在 CPU 上逐个程序实例地解释执行，语义与 GPU 上相同，只是很慢。它必须在导入 triton 之前设置。
examples/ch20_moe.py"""第 20 章：MoE 的路由、按专家排序补齐，Triton fused MoE 与参考实现的比较。

Triton kernel 在 CPU 上用解释器模式运行（TRITON_INTERPRET=1 必须在导入 triton 之前设置）。
"""

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("每个 token 选中的专家:", ids.tolist())
print("路由权重（重新归一化后）:", [[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("按专家排序、补齐到 4 的倍数（10 是占位符）:", sorted_ids.tolist())
print("每个块属于哪个专家:", 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 个 token：fused（Triton 解释器）与参考实现的最大误差 {(a - b).abs().max().item():.1e}")

运行结果每个 token 选中的专家: [[4, 3], [6, 1], [1, 7], [6, 4], [5, 4]]
路由权重（重新归一化后）: [[0.692, 0.308], [0.545, 0.455], [0.93, 0.07], [0.711, 0.289], [0.809, 0.191]]
按专家排序、补齐到 4 的倍数（10 是占位符）: [3, 4, 10, 10, 1, 10, 10, 10, 0, 7, 9, 10, 8, 10, 10, 10, 2, 6, 10, 10, 5, 10, 10, 10]
每个块属于哪个专家: [1, 3, 4, 5, 6, 7]
37 个 token：fused（Triton 解释器）与参考实现的最大误差 2.4e-07

5 个 token、每个选 2 个专家，共 10 个对。排序后专家 1 的两个对（第 3、4 个）组成第一个块，补两个占位符；专家 3 只有一个对，也占满一个块……共 6 个块。块越大，补齐浪费越多；块越小，矩阵乘效率越低。37 个 token 的 fused MoE 与参考实现的误差在 float32 舍入的量级。
端到端¶
用随机权重构造一个两层、8 个专家的小 Qwen3-MoE 模型，与 Hugging Face 的实现对比（CPU 上默认用参考后端）：
tests/test_ch20_moe.py · test_tiny_qwen3_moe_matches_hfdef 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()


官方实现

layers/moe.py · MoELayer（第 9～59 行）、models/qwen3_moe.py
fused MoE：moe/fused.py 与 kernel/triton/fused_moe.py。路由（topk_softmax）和排序补齐（moe_align_block_size）用 sgl_kernel 里的 CUDA kernel；块大小按形状二选一：token 数不超过专家数时用小块（16 × 32 × 64），否则用大块（64 × 64 × 32）——正是练习 1 讨论的问题



与官方的差异
我们的 Triton kernel 是简化版：块大小固定（BLOCK_M=16、BLOCK_N=32、BLOCK_K=32），没有按形状选择，没有官方 kernel 里的 GROUP_SIZE_M 分组（提高 L2 命中率），排序补齐用 PyTorch 而不是 CUDA kernel。参考后端 torch 是我们加的，CPU 上默认使用。

测试¶
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 还用最直白的"逐 token、逐专家"循环验证了参考后端，以及第 2 章提到的 Llama 3 长上下文 RoPE（同样用随机权重的小模型与 Hugging Face 对比）。

怎么讲清楚
讲 MoE：路由是 softmax → top-k → 可选的重新归一化（norm_topk_prob 让选中的 k 个权重之和为 1）。fused MoE 把 (token, 专家) 对按专家排序，每个专家的段补齐到 BLOCK_M 的整数倍，这样每个输出块只属于一个专家、一个 kernel 就能算完所有专家的 GEMM（整层两次启动：gate_up 和 down），而不是每个专家各启动一次。张量并行时每个专家按中间维切（和稠密 MLP 一样），一次 all-reduce；专家并行则按专家切，需要 all-to-all。没有 GPU 时用 Triton 的解释器模式在 CPU 上以相同的语义运行 kernel，验证正确性。

练习¶

BLOCK_M=64、128 个专家、每步 decode 只有 32 个 token（每个选 8 个专家）时，补齐之后有多少个对？浪费的比例是多少？这对 decode 意味着什么？
专家并行（EP）把不同的专家放在不同的 rank 上，每个 token 要被送到它的专家所在的 rank。它与本章的张量并行切法相比，通信模式有什么不同？（参考推理系统手册的专家并行一章。）
把 SiLU × up 融合进第一次 GEMM 的输出：kernel 需要同时算出 gate 和 up 两个块。怎样安排 offs_n 才能让一个程序实例拿到配对的 gate 和 up 列？


参考答案

256 个对最多分布在 128 个专家上，大多数专家只有 1～3 个对，每个都要补到 64：最坏约 128 × 64 = 8192 行，有效只有 256 行，浪费约 97%。decode 时 MoE 的 fused kernel 效率很低，所以 decode 常用更小的块、或专门为小批量设计的 kernel（例如按专家做 GEMV）。
张量并行每个 rank 都有全部专家的一部分，通信是一次 all-reduce（数据量与 token 数 × hidden 成正比）；专家并行每个 rank 有一部分专家的全部，需要 all-to-all 把 token 分发到专家所在的 rank、算完再收回来，数据量取决于路由分布，还可能出现负载不均。
让程序实例负责中间维上的一段 [n, n + BLOCK_N)，同时加载 gate 列 [n, n+BLOCK_N) 和 up 列 [I + n, I + n + BLOCK_N)，做两次 tl.dot，再 silu(gate) * up 后写出 [BLOCK_M, BLOCK_N]。


小结¶

 MoE 层：不切分的路由层 + 打包成三维张量的专家权重；张量并行时每个专家按中间维切，一次 all-reduce。
 路由：softmax → top-k → 可选重新归一化。
 fused MoE：按专家排序并补齐到块大小，每个输出块只属于一个专家，整层两次 kernel 启动。
 Triton 解释器让 kernel 在 CPU 上以相同的语义运行，用来验证正确性。


本章练习题（1 道）
在浏览器里直接写代码、跑测试；也可以在本地用 python practice/judge.py 判题。

203. moe_align_block_size 与分组 GEMM（困难）









  
  






                
  
  
    
    发现错误或有更好的写法？在 GitHub 上编辑此页 · 提交问题