SGLang 源码完全拆解（二十一）：从对齐表到分组 GEMM——FusedMoE 如何把每个专家的几行 token 塞满 Tensor Core

1.1 三次 launch 到底省了什么

如果不做融合，朴素实现是：

1
for e in range(E):
2
    idx = (topk_ids == e).nonzero()            # 一遍扫描 + 主机同步
3
    if idx.numel() == 0:
4
        continue
5
    x_e = hidden_states[idx] * weights[idx]    # 高级索引：一次 gather
6
    h = act(x_e @ w13[e]) @ w2[e]              # 两次 GEMM，e 个

问题有三层：

1. launch 次数：E 次循环意味着最多 2E + 1 次内核启动，DeepSeek-V3 有 256 个路由专家，一步 decode 就要发几百次 launch，每次 launch 在 Hopper 上约 3–5 µs 的调度开销，几乎等于一次小 GEMM 的计算时长；
2. 形状太碎：decode 时每个专家往往只分到几行甚至零行 token，cuBLAS 对 M < 16 的矩阵效率极差，Tensor Core 的 m16n8k16 / m64n*k16 指令需要至少 16/64 行才有意义；
3. 主机同步：nonzero() 要读回设备上的结果才能在 Python 层决定形状，这在 CUDA Graph 捕获时是致命的——图里不允许主机依赖。

融合方案把这三点一起解决：先排序再补零，让所有专家的 token 连续排列在同一个「虚拟矩阵」里，用一个大 kernel 按块处理，块内 expert_ids[block] 指出这块该用哪个专家的权重；k 份结果写进 intermediate_cache3，最后一次 moe_sum_reduce 求和。

1.2 为什么不是「让每个专家各发一个 grid」

分组 GEMM 的另一种写法是「每个专家一个 grid」，即 grid 的第一维是专家号，每个专家自己决定要处理多少行。这在 EP（专家并行）场景下更自然，DeepGEMM 的 grouped_gemm_nt_f8f8bf16_contig 就是这样：传入 m_indices 标出每行属于哪个专家，内核内部按专家分块。我们后面会看到，SGLang 在 deep_gemm 后端上正是这么做的。

但 Triton 单机路径选择「单一 grid + 对齐表」，因为：

• 对齐表让 grid 的 shape 是静态可推导的（M_pad / BLOCK_SIZE_M × N / BLOCK_SIZE_N），可以挂在 CUDA Graph 里；
• 每个 program 只看自己那几行，天然适合 Triton 的 program-level 并行模型，不需要额外处理「专家之间行数不均」的动态负载；
• 零行专家不需要单独跳过——对齐表根本不给它块，expert_ids 中出现的专家号已经过滤掉了空桶。

代价在下一篇会展开：补齐的零行会被白算。这个代价的大小直接由 BLOCK_SIZE_M 决定，所以 BLOCK_SIZE_M 的选择是这一层最关键的调参旋钮。

2.1 三个 dispatch 相关的字段

FusedMoE 还存了三个与单机无关、但会影响 forward 分支的字段：

• self.dispatcher：由 create_moe_dispatcher(moe_runner_config, quant_method) 构造。单机时是 StandardDispatcher，EP 时是 DeepEPDispatcher 等。
• self.runner：由 quant_method.runner 直接取出来，也是 MoeRunner。
• self.moe_runner_config：一个 dataclass，把上面那堆参数打包成一份「执行合同」。

其中 moe_runner_config 是贯穿三层抽象的物件。它定义在 moe_runner/base.py：

1
@dataclass
2
class MoeRunnerConfig:
3
    num_experts: Optional[int] = None
4
    num_local_experts: Optional[int] = None
5
    hidden_size: Optional[int] = None
6
    intermediate_size_per_partition: Optional[int] = None
7
    layer_id: Optional[int] = None
8
    top_k: Optional[int] = None
9
    ...
10
    runner_backend: Optional[MoeRunnerBackendLike] = None
11
    activation: str = "silu"
12
    is_gated: bool = True
13
    apply_router_weight_on_input: bool = False
14
    inplace: bool = True
15
    no_combine: bool = False
16
    routed_scaling_factor: Optional[float] = None
17
    gate_up_interleaved: bool = True
18
    layer: Optional[torch.nn.Module] = None
19
    silu_mul_keep_fp32: bool = False

它有四个属性其实是从 layer 上读的：

1
@property
2
def moe_tp_size(self) -> int:
3
    return self._layer_placement("moe_tp_size")

为什么不在构造时直接拷一份？源码注释里写得很直白——DWDP 的 bind_full_expert_weights 会在权重加载完之后「塌缩」层的 EP 视图，也就是把「本地专家」改回「全部专家」，而 MoeRunnerConfig 可能在那之前就构造了。保持一个地址、每次读层，才不会出现两份不一致的账。

3.1 MoeRunner 的选择逻辑

MoeRunner.__init__ 是一棵很长的分支树。按优先级大致是：

1
自定义注册工厂
2
  > triton
3
  > ascend (NPU)
4
  > triton_kernels
5
  > deep_gemm
6
  > humming
7
  > aiter (ROCm)
8
  > marlin / flashinfer_* / cutlass / hpc_ops  ← 这些 runner_core 为 None，只走 fused 路径

每个分支背后是一个 MoeRunnerCore 子类，实现 run(runner_input, quant_info, running_state)。但更关键的是构造函数的末尾这段：

1
if not lora_enabled:
2
    a2a_backend_name = get_moe_a2a_backend().value
3
    runner_backend_name = runner_backend.value
4
    self.fused_func = FusedOpPool.get_fused_func(a2a_backend_name, runner_backend_name)
5

6
if self.runner_core is None and self.fused_func is None:
7
    raise NotImplementedError(...)

也就是说，存在两套并行的执行模式：

• fused 模式：(a2a_backend, runner_backend) 这一对在 FusedOpPool 里注册过，就直接调用它，跳过所有格式转换；
• core 模式：没有 fused 函数时，走 pre_permute → runner_core.run → post_permute 三段式。

FusedOpPool 的 key 为什么是两维而不是一维？因为融合函数通常把「分发格式」和「计算布局」焊死在一起。例如 Triton 后端的 fused_experts_none_to_triton，它要求 dispatch 输出是 StandardDispatchOutput 且数据从未被重排过——它直接在原始 hidden_states 上开始算。如果换成 DeepEP 的 dispatch，输出已经是按目标 rank 分好的 buffer，这时就必须走 deep_gemm 的 grouped GEMM，而不能复用 Triton 的实现。

3.2 PermuteMethodPool 与三段式

三段式中的 pre/post permute 用 (输出格式, runner backend) 作为 key：

1
@register_pre_permute("standard", "triton")
2
def pre_permute_standard_to_triton(dispatch_output, quant_info, runner_config, running_state):
3
    config, down_config, down_moe_use_tma, up_moe_use_tma, \
4
        sorted_token_ids, expert_ids, num_tokens_post_padded = _prepare_fused_moe_run(...)
5

6
    running_state["config"] = config
7
    running_state["down_config"] = down_config
8
    running_state["down_moe_use_tma"] = down_moe_use_tma
9
    running_state["up_moe_use_tma"] = up_moe_use_tma
10

11
    return TritonRunnerInput(hidden_states=..., topk_weights=..., topk_ids=...,
12
                             sorted_token_ids=..., expert_ids=..., num_tokens_post_padded=...)

注意 pre_permute 并不只做「排列」，它同时做了配置解析与对齐表构造。running_state 是一个在三个阶段之间传递可变字典——它承载了 config、TMA 开关，以及 CUDA Graph 场景下的 down_gemm_overlap_args。这是一种轻量的「side channel」：避免把这些字段塞进 RunnerInput 从而污染 dataclass 的语义。

TritonRunnerInput 本身很朴素，六个字段：

1
@dataclass
2
class TritonRunnerInput(RunnerInput):
3
    hidden_states: torch.Tensor
4
    topk_weights: torch.Tensor
5
    topk_ids: torch.Tensor
6
    sorted_token_ids: torch.Tensor
7
    expert_ids: torch.Tensor
8
    num_tokens_post_padded: torch.Tensor

TritonMoeQuantInfo 则是量化态的载体：w13_weight、w2_weight、各种 use_fp8_w8a8/use_int4_w4a16 开关、w13_scale、block_shape、fuse_swiglu_interleaved。它从 quant_method.get_moe_quant_info(layer, runner_backend) 里来——因此更换量化方案时，只需要换一个 TritonMoeQuantInfo 构造器。

3.3 一段「兜底断言」的工程含义

MoeRunner.__init__ 里还有三条启动期校验，很值得单独读，因为它们是跨模块耦合的显式契约：

1
if get_moe_runner_backend().is_hpc_ops() and not runner_backend.is_hpc_ops():
2
    raise ValueError(
3
        "--moe-runner-backend hpc_ops ... would consume global ids as local ones and misroute "
4
        "tokens under EP>1, so fail loudly at startup instead."
5
    )
6

7
if get_moe_a2a_backend().is_deepep_v2() and not runner_backend.is_deep_gemm():
8
    raise ValueError("--moe-a2a-backend deepep_v2 requires the deep_gemm MoE runner ...")

第一条说明了一个很容易踩的坑：hpc_ops 的 dispatcher 会让 StandardDispatcher 保留全局专家号（skip_local_expert_mapping），因为它自己会做 EP 映射。如果某一层因为量化方法退化到 Triton runner，Triton 会把全局号当成本地号用，在 EP>1 下静默错路由。源码选择「启动即报错」而不是「运行期出错」，这类 fail-fast 在大规模部署里比任何日志都可靠。

第二条是 deepep_v2 与 deep_gemm 之间的耦合：deepep_v2 的输出直接是 DeepGEMM 要求的 m_indices 布局（FP8 激活 + scale，或未量化专家的 BF16 激活），换成 Triton runner 就没法消费。

3.4 FusedOpPool 的实际成员

了解哪些组合有 fused 函数，就能知道「常规部署走的是哪条路」：

• ("none", "triton") → fused_experts_none_to_triton
• 还有 ("deepep", "deep_gemm")、("flashinfer", "flashinfer_trtllm")、("mooncake", ...)、("nixl", ...)、("pplx", ...) 等组合。
• marlin / cutlass / hpc_ops / flashinfer_* 的 runner_core 是 None，要求它们必须能找到 fused 函数。

所以「单机 Triton」这条路的名字有点误导：它并不是没有分发器，而是分发器什么都不做，然后直接调一个融合函数。

5.1 补齐行到底长什么样

moe_align_block_size 的文档字符串里有一个完整例子：

1
topk_ids = [[2,3,4],[1,2,4],[1,3,4],[1,2,3]], block_size = 4, num_experts = 4
2
flatten → [2,3,4, 1,2,4, 1,3,4, 1,2,3]
3
每专家 3 行，补到 4 → 各补 1 个占位 token（编号 = numel = 12）
4
排序后 token_ids = [3,6,9,12, 0,4,10,12, 1,7,11,12, 2,5,8,12]

其中 12 是「不存在的 token」，它在 sorted_token_ids 里的值超过了 num_valid_tokens。Triton 内核就是靠这个阈值判断「这块的这一行是补齐行，跳过」：

1
offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
2
token_mask = offs_token < num_valid_tokens

这个阈值 num_valid_tokens = topk_ids.numel()，是内核从 Python 侧作为标量拿到的，不需要每行单独判断。所以补齐行的开销只是「块内多算了几个无效元素」，而不是「额外的分支」。

5.2 什么时候 ALIGN 会变成大头

考虑 DeepSeek-V3 的 decode：T = 128，k = 8，E = 256（不含共享专家）。numel = 1024。

• 随机路由下，每个专家的平均 token 数为 4，BLOCK_SIZE_M = 16，每专家补到 16，M_pad ≈ 256 × 16 = 4096。
• 也就是说，真实要算的 1024 行，被扩成了 4096 行——4 倍放大。

这个 4 倍放大就是 decode 阶段 MoE 的主要冗余来源。它换来的是：每个 program 都能吃满 m16n8k16 的指令，而且不需要动态 shape。prefill 时 numel 远大于 E × BLOCK_SIZE_M，对齐表几乎无冗余，收益就非常明显。

这里有一个容易忽略的细节：每个专家的补齐块数是 ceil(count / BLOCK_SIZE_M)，如果某专家 count = 0，它就没有块。expert_ids 的长度因此是 sum(ceil(count_e / BM))，而不是 E * ceil(...)。空专家完全被跳过，这与「朴素实现里要 continue」是对偶的，只是发生在 GPU 侧、不需要主机同步。

6.1 第一步：program 映射到哪个块

1
pid = tl.program_id(axis=0)
2
num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
3
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
4
num_pid_in_group = GROUP_SIZE_M * num_pid_n
5
group_id = pid // num_pid_in_group
6
first_pid_m = group_id * GROUP_SIZE_M
7
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
8
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
9
pid_n = (pid % num_pid_in_group) // group_size_m

这是 GROUP_SIZE_M 排序技巧（也见于 cuBLAS 的「swizzle」概念）：把 grid 切成一段一段的列，每段里面先沿 m 方向走，让相邻 program 覆盖同一个专家块的连续行，从而提高 L2 里 B[expert] 的复用。在 MoE 里它还有另一层含义——同一个 pid_m 段内的块可能对应同一个专家（对齐表按专家排），所以段内更有可能命中同一份 B。

6.2 第二步：这块属于谁、要处理谁

1
num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
2
if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
3
    return
4
offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
5
offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
6
token_mask = offs_token < num_valid_tokens
7

8
off_experts_i32 = tl.load(expert_ids_ptr + pid_m)
9
off_experts = off_experts_i32.to(tl.int64)

offs_token_id 是排序后的位置，offs_token 是原始的 token*k+slot 编码，off_experts 是这块负责的专家号。三者合在一起就能回答：「我要从原始 A 里取哪几行、用 B 的哪一页、写到 C 的哪些位置」。

这里的 token*k+slot 编码是为了解决一个细节：同一个 token 被复制 k 次，每次对应一个专家，但路由权重不同。用 token*k+slot 编码可以同时知道原始 token 和它在本行中的位置，从而从 topk_weights 里取出正确的权重：

1
moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)

topk_weights 是 [T, k] 展平后的一维数组，索引恰好是 token*k+slot。

6.3 第三步：构造 A、B 的指针

1
# A 的指针：跳过补齐行、跳过 slot，取原始 hidden 的那一行
2
a_ptrs = a_ptr + (
3
    offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
4
)
5

6
# B 的指针：按专家号取权重页
7
b_ptrs = (
8
    b_ptr
9
    + off_experts * stride_be
10
    + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
11
)

注意 offs_token[:, None] // top_k：这一步把 token*k+slot 还原成原始 token 号。所以同一个 token 的不同 slot，读的是同一行 A，但 B 不同、写入位置不同。

这一段还体现了 TMA 与非 TMA 的分叉：

1
if a_desc is not None:
2
    start_offs_m = pid_m * BLOCK_SIZE_M
3
else:
4
    ...
5
if b_desc is not None:
6
    start_offs_n = pid_n * BLOCK_SIZE_N
7
else:
8
    ...

TMA 模式不通过指针算术访问，而是把整块描述成 TensorDescriptor，由硬件 DMA 引擎按描述符搬运。描述符的坐标 [start_offs_m, k_start] 必须是块的起点，所以 sorted_token_ids 的取值不能被直接用在描述符上——c_sorted 这个开关就是干这个的。

6.4 第四步：K 方向累加与量化还原

1
for k_start in range(0, K, BLOCK_SIZE_K):
2
    if even_Ks:
3
        a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0)
4
        b = tl.load(b_ptrs)
5
    else:
6
        a = tl.load(a_ptrs, mask=token_mask[:, None] & (offs_k[None, :] < K - k_start), other=0.0)
7
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k_start, other=0.0)
8

9
    if use_int8_w8a16:
10
        accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
11
    elif use_fp8_w8a8 or use_int8_w8a8:
12
        if group_k > 0 and group_n > 0:
13
            # 块量化：每个 BLOCK_SIZE_K 块各自乘一组 scale
14
            offs_ks = k_start // group_k
15
            a_scale = tl.load(a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0)
16
            b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)
17
            accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :]
18
        else:
19
            accumulator = tl.dot(a, b, acc=accumulator)
20
    else:
21
        accumulator += tl.dot(a, b)

这里最值得读的是 FP8 块量化的三处细节：

1. even_Ks：如果 K % BLOCK_SIZE_K == 0，A、B 都不需要在 K 方向做 mask，这样可以省掉两次逐元素比较指令。这就是为什么 padded_size 允许在权重末尾补几列——补到能被 BLOCK_SIZE_K 整除，换取内层循环里更轻的加法。
2. 块 scale 的读取位置：offs_bsn = pid_n * BLOCK_SIZE_N // group_n（当 BLOCK_SIZE_N <= group_n）或 offs_bsn = offs_bn // group_n（当 BLOCK_SIZE_N > group_n）。这是 DeepSeek-V3 的「128×128 块」，group_k = 128、group_n = 128。在内层 K 循环中，offs_ks = k_start // 128 标出当前 K 块属于哪一组。
3. 累加方式的分叉：当 BLOCK_SIZE_N <= group_n 时，一个 B 块只属于一列组，b_scale 是标量，于是写法是 a_scale[:, None] * b_scale；当 BLOCK_SIZE_N > group_n 时，一个 B 块横跨多列组，b_scale 是向量，于是 a_scale[:, None] * b_scale[None, :]。这两条路径本质是同一件事，区别只在广播维度。

还有一个 swap_ab 分支，只在 FP8 且配置里指定时启用：

1
if swap_ab:
2
    a, b = tl.trans(b, (1, 0)), tl.trans(a, (1, 0))
3
    a_scale, b_scale = b_scale, a_scale

它把 [BM, BK] × [BK, BN] 换成 [BN, BK] × [BK, BM]，主要是为了在 BLOCK_SIZE_M 很小时让 tl.dot 的形状更符合 Tensor Core 的偏好。对 decode（M 只有 16）特别有用：M 维太短，把 B 放到「A 位置」会让 W 矩阵的 N 维（intermediate size，通常 2048）占据 M 槽位，指令利用率更高。

6.5 第五步：epilogue——乘路由权重、写回

1
if bias_ptr is not None:
2
    accumulator += bias
3

4
if MUL_ROUTED_WEIGHT:
5
    moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
6
    accumulator *= moe_weight[:, None]
7

8
accumulator = accumulator.to(compute_type)

MUL_ROUTED_WEIGHT 在 GEMM1 里通常为 False、在 GEMM2 里为 True（除了 apply_router_weight_on_input 的情况）。为什么不统一乘在 GEMM2？因为 GEMM1 的输出要经过 SiLU(gate) * up，这个非线性操作与路由权重相乘的顺序不同——w·SiLU(g) * up ≠ SiLU(w·g) * (w·up)。所以必须把路由权重放在 GEMM2 的最后，而 apply_router_weight_on_input 是另一种模式：先把权重乘到输入上，这样 GEMM1 与 GEMM2 都不用再乘。

写回有三个变体，其中 FUSE_SWIGLU 那一段最长，值得单独拆：

1
if FUSE_SWIGLU:
2
    acc_pairs = tl.reshape(accumulator, (BLOCK_SIZE_M, BLOCK_SIZE_N // 2, 2))
3
    gate_b, up_b = tl.split(acc_pairs)
4
    gate_f = gate_b.to(tl.float32)
5
    exp_neg = tl.inline_asm_elementwise(
6
        "{ mul.ftz.f32 $0, $1, 0fBFB8AA3B; ex2.approx.ftz.f32 $0, $0; }",
7
        "=f,f", [gate_f], dtype=tl.float32, is_pure=True, pack=1,
8
    )
9
    silu_f = tl.inline_asm_elementwise(
10
        "div.approx.ftz.f32 $0, $1, $2;", "=f,f,f", [gate_f, 1.0 + exp_neg],
11
        dtype=tl.float32, is_pure=True, pack=1,
12
    )
13
    out_act = (silu_f * up_b.to(tl.float32)).to(compute_type)

这里有一个重要的工程细节，源码注释写得很清楚：

The asm below is the bit-parity contract with the silu_and_mul this replaces, not an optimization: that kernel uses the fast-math intrinsics __fdividef(x, 1 + __expf(-x)), while Triton’s operators lower to accurate expf and IEEE div.rn. The 1-2 ULP gap is enough to flip the stored bf16 (8 mantissa bits) on many inputs.

也就是说，这段内联 PTX 不是「更快」，而是「算得和原来那个独立 silu_and_mul 完全一样」。因为 bf16 只有 8 个尾数位，1–2 ULP 的差异会让很多值在舍入时翻转。对于「batch-invariant」（结果不随 batch 组成变化）这类严格要求，这种位级对齐是必须的。

0fBFB8AA3B 是 -log2(e) 的 float 位模式，前面把负号折叠进常量，这样 ex2.approx(x * -log2 e) 就等于 exp(-x) 的近似，少一条指令，且符号翻转精确。

6.6 第四阶段的 combine 去哪了

回到 _fused_moe_kernel_sequence。GEMM2 写完 intermediate_cache3 [T, k, H] 之后，调用 moe_combine_topk：

1
elif _is_cuda or _is_musa:
2
    if topk == 1 and routed_scaling_factor == 1.0 and not _use_intermediate:
3
        pass  # 直接写进 out_hidden_states
4
    elif topk == 2 and routed_scaling_factor == 1.0:
5
        torch.add(intermediate_cache3[:, 0], intermediate_cache3[:, 1], out=out_hidden_states).squeeze(dim=1)
6
    else:
7
        if _use_moe_sum_reduce_torch_compile(num_tokens):
8
            ... moe_sum_reduce_torch_compile / moe_sum_reduce_triton ...
9
        else:
10
            moe_sum_reduce(...)

这里有三个可以省的分支：

• topk == 1：不需要求和，GEMM2 直接写入输出。
• topk == 2：用 torch.add 而不是专用内核，因为两路求和用 add 就是最优路径。
• topk > 2：分两档，小 token 数用 torch.compile 的版本（编译器会展开为向量化加法），大 token 数用预编译的 moe_sum_reduce（减少编译开销与寄存器压力）。_use_moe_sum_reduce_torch_compile 的阈值就是这两档的分界线。

topk == 2 这条路径对 Mixtral 很关键：num_experts_per_tok = 2 是经典配置。

6.7 defer_sum 与「最后一次求和」外移

_fused_moe_kernel_sequence 有一个 defer_sum 参数：

1
if defer_sum:
2
    # Keep routing multiplication inside GEMM2, including its BF16 output
3
    # rounding. The consumer owns the FP32 sum and routed scaling.
4
    return intermediate_cache3

它允许 GEMM2 只产出加权后的 [T, k, H]，把求和挪到调用方。谁会用？在 DeepseekV2MoE 的分层边界（stage boundary）场景——当流水/PP 在层边界需要做 reduce_scatter 或与共享专家输出合并时，可以在求和之前先把共享专家那一路加进来，再一次性做 FP32 求和与 routed_scaling_factor。这比「先求和、再加共享专家」少一次读改写。

7.1 total_tokens 的两套口径

1
padded_tokens = (
2
    min(num_tokens * topk, E + 1) * (config["BLOCK_SIZE_M"] - 1)
3
    if down_moe_use_tma
4
    else 0
5
)
6
total_tokens = num_tokens * topk + padded_tokens

非 TMA 模式：中间缓冲区大小恰好是 T × k，因为对齐表里的「虚拟补零」写进 C 时靠 mask 处理，不会多分配。 TMA 模式：每行还可能因为 BLOCK_SIZE_M 补齐而多出现一次（同一行可能被拆到相邻两块的不同位置），所以额外多分配 min(T·k, E+1) × (BM-1)。这个上界来自「每块至多浪费 BM-1 行 × 最多 E+1 个块」。

7.2 use_fused_moe_sum_all_reduce

在 TP 模式下，MoE 的输出需要 all_reduce。普通做法是「求 k 路和 → 一次 all_reduce」，但在某些配置下还有一个更激进的选择：

1
use_fused_moe_sum_all_reduce = (
2
    get_exec().moe.enable_fused_moe_sum_all_reduce
3
    and (not no_combine) and (not defer_sum) and (topk > 2)
4
    and (not use_int8_w8a16) and (not use_int4_w4a16)
5
    and (not (hooks and hooks.after_down is not None))
6
)

启用时，GEMM2 不用 intermediate_cache3，而是直接用 tl.atomic_add 把每个专家的输出累加到输出缓冲区，再借 NCCL 的对称内存做一次「求和 + all_reduce」融合。topk > 2 是必要条件：两路用 add 更快；小于等于 2 时融合的收益被融合开销吃掉。

tl.atomic_add 为什么安全？因为每个 token 的 k 个 slot 由同一 pid_m 内的不同行处理，目标地址不重叠（不同 token 的目标不同），但 pid_n 方向上会有多个 program 争抢同一列。atomic_add 保证了同位置累加不丢；out_slice.zero_() 提前清零提供初始值。

8.1 三个内核怎么塞进一张图

CUDA Graph 捕获要求所有内核的 grid 与参数在捕获时可以静态决定。SGLang 的 FusedMoE.forward 做了这样的判断：

1
if is_in_tc_piecewise_cuda_graph():
2
    if TopKOutputChecker.format_is_bypassed(topk_output) and (...):
3
        return self.forward_impl(...)
4
    if TopKOutputChecker.format_is_standard(topk_output):
5
        return moe_forward_piecewise_cuda_graph_impl(...)
6
    elif TopKOutputChecker.format_is_bypassed(topk_output):
7
        return fused_moe_bypassed_piecewise_cuda_graph_impl(...)
8
    else:
9
        return self.forward_impl(...)

（tc 指 Tensor Core 分段图，在 SGLang 第九篇、vLLM 第十二篇里有更完整的铺垫。）

这里的分支逻辑在说：

• standard TopK：topk_weights/topk_ids 已物化，图的形状由「当前 batch 的 padding 桶」决定；此时三个内核都可以作为「固定 shape」的 op 被捕获。
• bypassed TopK：路由与专家执行需要同一个 op 捕获——因为中间结果 router_logits 与 topk_ids 的 shape 是动态的，不能跨 op 边界留在图里。fused_moe_bypassed_piecewise_cuda_graph_impl 就是把「路由 + 对齐 + GEMM + 求和」整段注册为一个 custom op。
• 带 padding 或 expert_location 元数据的：只能走非图路径。

8.2 TMA descriptor 的缓存

TMA 模式需要为 A 与 B 各构造 TensorDescriptor。A 的 descriptor 每次都要新建（shape 是当前 batch 的），B 的 descriptor 可以被缓存，因为 B 是静态权重：

1
if b_use_tma:
2
    b_desc = _get_b_tma_desc_cached(B, config["BLOCK_SIZE_N"], config["BLOCK_SIZE_K"])

_get_b_tma_desc_cached 用 (data_ptr, BN, BK) 做 key。这是一个细小的工程点：TMA descriptor 构造约几十微秒，若每层每步都新建，decode 的 CPU 开销会明显增长。权重是常驻的，所以缓存是安全的——除非 EP 重排（EPLB）会改变专家顺序，那时需要清缓存，源码里对应 clear_b_tma_desc_cache()。

8.3 PDL：核间流水

最后有一行很容易被忽略的 launch 参数：

1
pdl_kwargs = (
2
    {"USE_GDC": True, "launch_pdl": True, "GDC_EARLY": A.shape[0] <= 512}
3
    if is_arch_support_pdl() else {}
4
)

PDL（Programmatic Dependent Launch）是 Hopper 引入的机制：下个 kernel 可以在上个 kernel 还没完全退出时就开始启动，通过 gdc_wait/gdc_launch_dependents 同步。内核开头与结尾的：

1
if USE_GDC:
2
    tl.extra.cuda.gdc_wait()
3
    if GDC_EARLY:
4
        tl.extra.cuda.gdc_launch_dependents()
5
...
6
if USE_GDC and not GDC_EARLY:
7
    tl.extra.cuda.gdc_launch_dependents()

就是手写的握手。在 GEMM1 → 激活 → GEMM2 这条链上，PDL 能把三次 launch 的尾部气泡从约 2×3 µs 压到 1 µs 量级。

GDC_EARLY 由 A.shape[0] <= 512 决定：小 batch 时，前一个 kernel 的尾巴很长（每块浪费的补齐行比例高），提前放行后续 kernel 反而可能读到还没写的数据——这个开关的选择依据是「GEMM1 的 epilogue 是否已经能把数据安全地交给下一个阶段」。

10.1 inplace 与 no_combine 不能同时开

fused_experts 里：

1
if moe_runner_config.inplace:
2
    assert not moe_runner_config.no_combine, "no combine + inplace makes no sense"
3
    inplace_fused_experts(...)
4
    return hidden_states

inplace=True 时 GEMM2 直接写到 hidden_states，省掉再分配输出。但如果 no_combine=True（要求返回 [T, k, H] 中间结果），输出跟输入形状不同，原地写就不成立。

10.2 out_hidden_states.unsqueeze(0)

当 topk == 1 且 routed_scaling_factor 是 1.0、又没有 LoRA 钩子时，GEMM2 的目标指针写的是：

1
out_hidden_states if use_fused_moe_sum_all_reduce
2
else (intermediate_cache3 if _use_intermediate else out_hidden_states.unsqueeze(0))

unsqueeze(0) 是把 [T, H] 变成 [1, T, H]，从而把 topk 维看作长度 1。这样一个写入路径可以同时覆盖 [T, k, H] 与 [T, H]，且在内核里不需要多一个分支。

10.3 共享专家缩放因子的两个来源

DeepseekV2MoE 的构造函数里有一个计算：

1
fused_shared_experts_scaling_factor = None
2
if self.moe_ep_size > 1 and self.num_fused_shared_experts > 0 and not _uses_per_rank_shared_slots:
3
    fused_shared_experts_scaling_factor = 1.0 / float(self.moe_ep_size)

在 TP 模式下，每个 rank 都持有全部共享专家、都会算出同一份输出，最终 all_reduce 会把这份输出加 TP 次。于是在 EP 模式下需要乘 1/ep_size 来抵消这种重复相加。而 DeepEP/MegaMOE 用的是每 rank 物理槽位，共享专家只在「家 rank」算一次，就无需缩放。这个细节和第二十篇里的「共享专家列如何拼进 topk」是一致的：拓扑决定了账本。

10.4 MoE 的「最小形状」在推理框架里被两种方式处理

intermediate_size_per_partition 在 flashinfer_trtllm 分支下会被 round-up：

1
if (self.use_flashinfer_trtllm_moe
2
    and not isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod)
3
    and self.intermediate_size_per_partition % 128 != 0):
4
    self.intermediate_size_per_partition = round_up(self.intermediate_size_per_partition, 128)

TRT-LLM 的后端对 I_pp 有 128 对齐要求，缺的列在权重加载时补零。这与 padded_size（K 补零）是同一个思路：用权重末端的空白换取内核里无分支。Triton 路径则不做这个补齐，而是在 even_Ks 分支里用 mask 处理。