TensorRT-LLM 源码完全拆解（十三）：MoE 层前向执行——路由、Fused Kernel 与专家并行的三段拼装

旧路径的问题

重构之前，TensorRT-LLM 的 MoE 是”一个后端一个类”：CutlassFusedMoE 是一种实现，TritonFusedMoE 是另一种，每个类自己写 forward，自己管理权重加载、自己调通信。这种做法在只有两三个后端时还行，但仓库里的后端数量涨到了十几个：CUTLASS、TRTLLM-Gen、DeepGEMM、CuTe DSL、DenseGEMM、Marlin……而且每个后端还要和”EP 开不开""EPLB 开不开""用哪种通信策略”交叉组合。

继承体系在这种维度下会迅速失控。假设有 10 个后端、4 种通信策略、2 种 EPLB 状态，如果每一种组合都要一个类或者一层 if-else，代码就会变成组合爆炸。更糟的是，通信逻辑和对权重加载、量化的处理本来就没有关系，混在一个类里以后没法单独测试。

新的拆法

TensorRT-LLM 的答案是把”编排”和”计算”分开。官方开发者文档 MOE_DEVELOPER_GUIDE.md 里给的职责划分是这样的：

1
ConfigurableMoE
2
├── Backend           (纯计算: routing → quantize → FC1 → act → FC2)
3
├── Communication     (分布式, 可选: dispatch tokens → compute → combine)
4
├── EPLB              (可选: 专家在 GPU 间动态迁移)
5
├── DWDP              (可选: 绑定 composite-VA 权重; 不做 EP 通信)
6
└── MoEScheduler      (前向执行策略: 分块、EPLB 钩子顺序、通信编排;
7
                       由 backend.scheduler_kind 选择)

这张图是理解整个模块的关键，逐项说：

Backend 是纯计算。 它只负责”给我一批 token 和它们的路由结果，我把专家算完”。它不做任何跨卡通信，也不知道 EPLB 的存在。这条约束在文档里被反复强调，因为它带来一个非常重要的性质：后端的实现可以脱离分布式去单元测试。

Communication 是可插拔的。 dispatch（把 token 送到专家所在的卡）和 combine（把结果收回来）由独立的策略类实现，后端不感知。CommunicationFactory 在运行时根据硬件和配置选一个策略出来。

Scheduler 是前向执行策略。 它决定”要不要把这一批 token 切成几块分别算""EPLB 的统计钩子插在哪个位置""会不会用第二条 stream 重叠”。注意 Scheduler 是无状态的——它只持有对 ConfigurableMoE 的反向引用，读包装器的状态但不写。状态都在包装器上。

EPLB 是可选的专家迁移器。 它解决的是”路由器学出来的专家分布天然不均，冷的卡闲着、热的卡排队”这个问题，需要动态地把专家从热卡搬到冷卡上。

组合的边界：为什么必须是”任意 Backend × 任意 Communication”

文档里列了五条核心设计原则，其中第二条是：

Any Backend × Any Communication × EPLB On/Off —— 所有合法组合都应该能工作（受 can_implement 和 scheduler_kind 约束）。

这条原则听起来像口号，但它有非常具体的工程含义：当一个后端加进来时，它不需要为每一种通信策略各写一遍适配。通信策略的接口（dispatch / combine）是统一的，后端只负责算。反过来说，一个新的通信策略（比如某个硬件厂商的自定义 all-to-all）加进来时，也不需要动任何一个后端。

这个组合空间不是免费的。下面几节会看到，正是这个”任意组合”的要求，逼出了 can_implement 的纯函数约束、MoEResolutionReport 这样的可解释选择机制，以及”降级允许、静默不允许”这条规则。

forward_impl 有多薄

编排者自己的 forward_impl 被压缩到了极薄的一层，文档里给的原样是这样的：

1
def forward_impl(self, x, router_logits, ...):
2
    if self.enable_dwdp:
3
        self.dwdp_manager.wait_and_bind(self.backend, self.layer_idx)
4
    outputs = self.scheduler.forward(x, router_logits, ...)
5
    if self.enable_dwdp:
6
        self.dwdp_manager.record_compute_and_prefetch_next(self.layer_idx)
7
    self.repeat_idx = (self.repeat_idx + 1) % self.repeat_count
8
    return outputs

真正的活全在 self.scheduler.forward 里。包装器只做三件事：DWDP 的绑定与记账（一个把跨卡共享的专家权重页映射到本卡地址空间的机制，只在特定硬件路径上启用）、repeat_idx 的推进（EPLB 用到的”同一份权重被连续用了几次”的计数）、以及返回值的搬运。

这么薄的代价是 forward_impl 里看不出任何技术细节；好处是任何新的执行策略都只需要写一个 Scheduler 子类，不需要碰权重加载、量化、EPLB 的生命周期代码。这就是”组合优于继承”在真实代码里的样子。

路由方法枚举

每个模型用的路由算法都不一样，代码里用 RoutingMethodType 枚举把它们编码下来：

1
class RoutingMethodType(IntEnum):
2
    # Default: Softmax -> TopK
3
    Default = 0,
4
    # Renormalize: TopK -> Softmax
5
    Renormalize = 1,
6
    # DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups -> Top8 experts
7
    DeepSeekV3 = 2,
8
    # Llama4: Top1 -> Sigmoid
9
    Llama4 = 3,
10
    # Qwen3: Softmax -> TopK -> Renormalize
11
    RenormalizeNaive = 4,
12
    # MiniMaxM2: Sigmoid -> RoutingBiasAdd -> TopK -> Renormalize(without bias)
13
    MiniMax2 = 5,
14
    # SigmoidRenorm: Sigmoid -> TopK -> Renormalize
15
    SigmoidRenorm = 6,
16
    # DeepSeek-V4
17
    DeepSeekV4 = 7,
18
    # Unspecified
19
    Unspecified = 8,

枚举上方的注释就是算法本身，这比写成散文更容易核对：Default 是”先 softmax 再取 top-k”，Renormalize 是”先取 top-k 再在选中的专家上归一化”，两者的差别只在归一化的分母包含哪些专家。

BaseMoeRoutingMethod 是所有路由方法的基类，它的核心接口是 apply：

1
def apply(self, _router_logits, _input_ids) -> tuple[torch.Tensor, torch.Tensor]:
2
    """
3
    返回 (token_selected_experts: int32, token_final_scales: float32)
4
    """

返回的两张张量形状都是 [num_tokens, top_k]：前者是专家 ID，后者是对应的权重系数。路由方法在文档里被明确要求是 nn.Module 的纯计算，不碰通信——路由结果只告诉”该去哪些专家”，至于那些专家在本卡还是别的卡上，是通信层的事。

DeepSeek-V3 的 noaux_tc：分组路由的算术

看一个最有代表性的实现。DeepSeek-V3 的路由叫 noaux_tc（No Auxiliary Loss, Top-K with Correction bias），它和”softmax + top-k”完全不是一回事。

打分不用 softmax 用 sigmoid：

1
@staticmethod
2
@torch.compile(options={"max-autotune": True})
3
def get_scores(logits, e_score_correction_bias):
4
    scores = F.sigmoid(logits)
5
    scores_with_bias = scores + e_score_correction_bias
6
    return scores, scores_with_bias

注意这里产生了两份分数：scores 是纯粹的 sigmoid 输出，scores_with_bias 是加上偏置之后的。偏置只用来决定”选谁”，不用来算权重——这一点是 noaux_tc 的核心设计（下面会解释为什么）。

接着是分组选择。256 个专家分成 n_group 组，每组内先取分数最高的 2 个求和作为组分数，然后从组分数里取 topk_group 个组，只在这些组里选最终的 top_k 个专家：

1
group_scores = torch.sum(torch.topk(
2
    scores_with_bias.view(scores_shape[:-1] + [n_group, scores_shape[-1] // n_group]),
3
    k=2, dim=-1, largest=True, sorted=True)[0], dim=-1)
4
_, group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1,
5
                          largest=True, sorted=True)
6
group_mask = torch.zeros_like(group_scores)
7
group_mask.scatter_(-1, group_idx, 1)

组外专家的分数被置成负无穷，之后的 top-k 就不可能选中它们：

1
scores_with_bias = torch.where(score_mask.bool(), scores_with_bias,
2
                               torch.tensor(float('-inf'), ...))
3
_, topk_idx = torch.topk(scores_with_bias, k=self.top_k, dim=-1, largest=True)

最后一步是权重计算，也是最能体现设计意图的地方：

1
new_mask = torch.zeros_like(scores)
2
new_mask.scatter_(-1, topk_idx, 1)
3
scores = scores * new_mask           # 用无偏置分数，靠 mask 筛出选中的专家
4
score_sum = torch.sum(scores, dim=-1, keepdim=True) + 1e-20
5
scores = scores / score_sum * self.routed_scaling_factor

为什么选择用带偏置的分数、而权重用不带偏置的分数？ 因为 e_score_correction_bias 是一个负载均衡手段：训练时它被调整成让冷门专家更容易被选中（不需要额外的辅助损失项，这就是 “noaux” 的含义）。如果这个偏置进入了权重计算，那么一个被偏置”抬进”top-k 的专家会同时拿到虚高的权重，输出的语义就被污染了。所以偏置只做”入场券”，权重仍由模型原本学出来的 sigmoid 分数决定。这是整个 noaux_tc 里最容易写错的地方。

融合算子与回退阈值

上面这段 PyTorch 逻辑只在”形状不受支持”时才会跑。正常路径是调一个融合算子：

1
topk_values, topk_indices = torch.ops.trtllm.noaux_tc_op(
2
    logits, e_score_correction_bias, n_group, self.topk_group,
3
    self.top_k, self.routed_scaling_factor)

融合与否则由一个明确的阈值决定：

1
if self.n_group > 1:
2
    experts_per_group = num_experts // n_group
3
    if (self.top_k > 8 or num_experts > 256 or experts_per_group > 32
4
            or experts_per_group * self.topk_group > 256):
5
        self.is_fused = False

也就是说，融合 kernel 只在”top_k ≤ 8、专家数 ≤ 256、每组 ≤ 32 个专家、且参与最终选择的专家总数 ≤ 256”时才可用。超过这些边界就退回 PyTorch 实现并打一条 warning。为什么要设这些阈值？因为融合 kernel 的分组选择通常是把专家的候选集合放进 shared memory 或寄存器里做一次归约，容量是硬约束；超过容量就必须走通用实现。这类”融合版有形状上限、超了静默降级”的模式在 TensorRT-LLM 里到处都是，理解它比记住具体数字更重要。

路由与后端的解耦：requires_separated_routing

对绝大多数后端，路由是”贴着”专家计算的——因为路由结果本来就要按专家重新排列 token，这个重排操作顺手就做了。但对某些后端不行，比如某个 kernel 要求输入已经按专家排好序、或者要求路由分两阶段做。代码里用一个属性表达这件事：

1
class DeepSeekV4MoeRoutingMethod(BaseMoeRoutingMethod):
2
    @property
3
    def requires_separated_routing(self) -> bool:
4
        ...

它和 TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING 这个环境变量配合，让”路由是否必须和计算分开”成为一个可以由后端声明、也可以由用户强制的事实。这是”任意后端都能和任意路由方法配对”这一目标的必要补丁：解耦不能只靠接口统一，还得给那些有特殊约束的成员留一个诚实的声明位置。

选择依据是”通信在哪”

后端通过 scheduler_kind 类属性声明自己属于哪一类，默认是 EXTERNAL_COMM：

1
class MoESchedulerKind(Enum):
2
    """
3
    - EXTERNAL_COMM: comm 在 MoE kernel 边界之外; 调度器从 host 发起
4
      Communication.dispatch / Communication.combine, 支持 per-chunk EPLB 钩子
5
      与可选的多 stream 分块重叠 (Cutlass, DeepGemm, CuteDSL, DenseGEMM, TRTLLMGen).
6
    - FUSED_COMM: comm 通过 NVLink SymmBuffer 融进后端的融合 kernel
7
      (DeepGEMM fp8_fp4_mega_moe). 无 host 通信; lockstep 分块启动;
8
      EPLB 统计更新用 ignore_allreduce=False.
9
    """
10

11
    EXTERNAL_COMM = "external_comm"
12
    FUSED_COMM = "fused_comm"

两个调度器的完整执行序列，文档里是这么写的：

1
ExternalCommMoEScheduler._forward_chunk_impl 每个 chunk 执行:
2
[EPLB start_wait_gpu] → routing → [EPLB done_wait_gpu + update_statistic + route]
3
  → [comm.prepare_dispatch (NVLink2-sided)] → quantize/dispatch (自适应顺序)
4
  → backend.run_moe → [EPLB start_set_cpu] → [unfinalized_combine_fn]
5
  → comm.combine → [EPLB done_set_cpu]
6

7
FusedCommMoEScheduler._forward_chunk 每个 chunk 执行:
8
[EPLB start_wait_gpu] → routing → [EPLB done_wait_gpu + update_statistic + route]
9
  → backend.quantize_input → backend.run_moe (融合 dispatch+GEMM+act+GEMM+combine)
10
  → [EPLB start_set_cpu + done_set_cpu]

差别很清楚：外部通信版里，dispatch 和 combine 是 host 上发起的两次独立调用，中间夹着 run_moe；融合版里这三步都在 kernel 内部完成，host 侧只看到一次 run_moe 调用。

图源：TensorRT-LLM 官方技术博客的 MoE 数据流图（docs/source/blogs/media/tech_blog18_dispatch_moe_combine.png）。图里把”一次 MoE 前向”拆成三段：dispatch 把 token 按目标专家送到各 rank，MoE 段在本卡持有的专家上完成 FC1→激活→FC2，combine 再把每个专家产出的 partial 结果送回 token 的原属 rank 并按路由权重相加。本文反复提到的”通信与计算的分界”，指的就是这两个虚线边界——EXTERNAL_COMM 后端把边界做成 host 发起的真实通信调用，FUSED_COMM 后端则把整张图塞进一个 kernel。

“自适应顺序”里的一个细节

上面序列里那句 “quantize/dispatch (自适应顺序)” 值得展开：

1
Adaptive quantize/dispatch order (gated by comm.supports_post_quant_dispatch()):
2
  Post-quant flow: quantize_input() → comm.dispatch()   (发送量化后的数据)
3
  Pre-quant flow:  comm.dispatch() → quantize_input()   (发送原始数据, 本地量化)

为什么会有两种顺序？ 直觉上应该先量化再发——量化后数据量小 4 到 8 倍，通信量跟着降。但这不是无条件成立的：dispatch 要求发送方把数据写成接收方期望的内存布局，如果量化本身就要做布局变换（比如 NVFP4 要配 scale factor 的 swizzle 布局），那么”先量化再 dispatch”意味着量化 kernel 得按远端布局写；反过来”先 dispatch 原始 BF16、到达后再本地量化”则让量化 kernel 只需面对本卡布局。

选择哪一种，取决于通信策略的 supports_post_quant_dispatch() 和量化方法的实现代价。这是典型的”哪个 kernel 更贵就让它少做事”权衡，而不是一条普遍规则。

两条路径的约定是刻意相反的

这是全篇最反直觉的一段。文档里明确写着：

两条路径有刻意相反的不变量（use_dp_padding 被遵守 vs 被忽略、ADP padding 保留 vs 剥离、空 chunk 被替换 vs 零 token kernel 启动、多 stream 重叠允许 vs 禁止）。

逐条理解：

DP padding（注意力数据并行下的 token 数对齐）。 开了 attention DP 以后，不同 rank 上的 token 数不一样；做集合通信前必须把它们补齐到相同长度。外部通信路径尊重调用方传进来的 use_dp_padding 语义；融合通信路径则忽略它——因为融合 kernel 里的通信是 lockstep 的，每个 rank 都必须在同一个时刻进入 kernel 里的 NVLink barrier，padding 与否由 kernel 内部的协议决定。

空 chunk 的处理方式相反，是这类差异里最有趣的一条。 外部通信路径遇到某个 rank 在某一块上一个 token 都没有时，会用”替代”手法：给这个 rank 塞一个零 token 的占位，让它在集合通信里仍然参与，从而不会把别的 rank 卡住。融合通信路径则相反——它会照样启动 kernel，哪怕一个 token 都没有。为什么？因为融合 kernel 里有一个 in-kernel 的 NVLink barrier，如果这个 rank 不启动 kernel，别的 rank 会在 barrier 上永远等下去。文档里这句话把原因说得很直白：

Zero-token chunks still launch the kernel so peer EP ranks can cross the in-kernel NVLink barrier.

多 stream 重叠允许 vs 禁止。 外部通信路径可以在 enable_alltoall 关闭且存在 aux stream 时开启分块重叠——第 $k$ 块的 GEMM 和第 $k+1$ 块的 dispatch 跑在不同的 stream 上。融合路径不行：通信已经进了 kernel，分块必须 lockstep 地顺序启动。

这组”相反的不变量”提示了一个更普遍的道理：当两个实现共享同一个抽象接口，但底层同步模型完全不同时，不要为了接口好看而把差异藏起来。TensorRT-LLM 的做法是把差异诚实写进各自的类文档字符串（moe_scheduler.py 的 class docstring 和 MOE_SCHEDULER_DESIGN.md），让后来者一眼能看出”这里的行为和那边不一样，而且是故意的”。

外部通信调度器的真实前向：先把 padding 修好

看 ExternalCommMoEScheduler.forward 的真实代码，它的步骤编号在 docstring 里就写清楚了：

1
def forward(self, x, router_logits, *, do_finalize, output_dtype,
2
            all_rank_num_tokens, use_dp_padding, input_ids=None,
3
            lora_params=None, routing_aux=None, input_ready_event=None):
4
    moe = self.moe
5

6
    # ========== Step 1: 处理 padding ==========
7
    if all_rank_num_tokens is None:
8
        all_rank_num_tokens = [x.shape[0]]
9
    all_rank_max_num_tokens = max(all_rank_num_tokens)
10
    if use_dp_padding:
11
        all_rank_num_tokens_padded = [all_rank_max_num_tokens] * len(all_rank_num_tokens)
12
    else:
13
        all_rank_num_tokens_padded = all_rank_num_tokens

接着有一个零 token 死锁修复，注释写得很详细：

1
# ========== 0-token rank deadlock fix ==========
2
# 当某些 rank 在单 chunk 前向里 0 token 时, 这些 rank 会在
3
# moe.comm.dispatch() 之前就卡在 CUDA kernel 里 (例如 0 行的 NVFP4
4
# quantize_input), 导致非零 rank 在 NCCL AllGather 上死锁。
5
# 修复: 在所有 rank 上统一打开 DP padding, 让每个 rank 都用 sizes=None
6
# (统一 allgather) 并把 x/router_logits 补齐到 max_tokens。
7
if (moe.comm is not None and moe.use_dp
8
        and all_rank_max_num_tokens > 0 and not use_dp_padding
9
        and any(t == 0 for t in all_rank_num_tokens_padded)):
10
    use_dp_padding = True
11
    all_rank_num_tokens_padded = [all_rank_max_num_tokens] * len(all_rank_num_tokens)
12
    local_n = x.shape[0]
13
    if local_n < all_rank_max_num_tokens:
14
        input_ready_event = _wait_input_ready(input_ready_event)
15
        pad = all_rank_max_num_tokens - local_n
16
        x = torch.cat([x, x.new_zeros((pad, x.shape[1]))], dim=0)
17
        router_logits = torch.cat([router_logits,
18
                                   router_logits.new_zeros((pad, router_logits.shape[1]))], dim=0)

这段代码的价值在于它暴露了一个真实世界里非常常见的死锁模式：一个 rank 提前死在 kernel 里，另一个 rank 在集合通信上等它。注意死锁不是发生在通信里，而是发生在通信之前——0 行的量化 kernel 本身就把 rank 卡住了。修复办法是让所有 rank 的输入形状一致，让”0 行”这种特殊形状彻底不出现。

还有一段关于 stream 同步的设计：

1
"""``input_ready_event``: ``x`` 是在另一条 stream 上写好的。
2
   把 join 推迟到第一个真正读 ``x`` 的 kernel 上, 这样路由可以与之重叠。
3
"""

也就是说，调度器不会一进门就 wait_event，而是把等待推迟到路由算完之后——路由只需要 router_logits，不需要 x。这是流水线重叠里很典型的一个手法：把同步点尽量往后压，压到最后一个还不读数据的算子之后。

图源：TensorRT-LLM 官方技术博客的 all-to-all 数据流图（docs/source/blogs/media/tech_blog14_alltoall_dataflow.png）。这张图解释了为什么 MoE 的通信记账不能只看”本 rank 有多少 token”：开了 attention DP 以后，dispatch 会把所有 DP rank 的 token 汇总过来，所以一次前向实际要处理的 token 行数是 $N_{\mathrm{dp}} \times \max_i n_i$，而不是本地的 $n_r$。本文后面 calculate_num_chunks 里那个 num_rows 就是按这个口径算的。

分块：calculate_num_chunks 与并行策略的分工

单块还是多块，由 moe.calculate_num_chunks(all_rank_num_tokens_padded) 决定；多块走 _forward_multiple_chunks。这个函数的真实实现把”要处理多少行 token”算得很讲究：

1
def calculate_num_chunks(self, all_rank_num_tokens: List[int]) -> int:
2
    if self.use_dp and self.comm is not None:
3
        num_rows = self._dp_padded_num_rows(all_rank_num_tokens)
4
    elif self.enable_dwdp:
5
        num_rows = max(all_rank_num_tokens)
6
    else:
7
        assert len(all_rank_num_tokens) == 1
8
        num_rows = all_rank_num_tokens[0]
9
    return (num_rows + self.moe_max_num_tokens - 1) // self.moe_max_num_tokens

三种情形对应三种”这批到底有多少行”的口径：DP 通信开着时是 len(rank) * max(tokens)（所有 DP rank 的 token 会被汇总过来）；DWDP 模式不搬 token 只预取权重，所以只按本卡自己的行数算；其余情况就是本地行数。最后向上取整除以 moe_max_num_tokens——分块的动机不是”把大 batch 拆小”，而是把每个块的 token 行数压在一个配置上限内，因为 fused kernel 的 workspace、路由的分组归约容量、以及量化 kernel 的 scale 布局都是按这个上限设计的。

分块的第二个收益才是重叠：多块时不同块的 dispatch/GEMM/combine 可以并行推进。

通信策略的兜底也在这一步发生。 调度器算完块数后调 determine_communication_method，它只做两件事：问当前策略”这个负载你能不能干”，不能干就换成 AllGather：

1
def determine_communication_method(self, all_rank_num_tokens, num_chunks) -> None:
2
    if self.comm is None or isinstance(self.comm, AllGatherReduceScatter):
3
        return
4

5
    feasible_workload = self.comm.is_workload_feasible(all_rank_num_tokens, num_chunks)
6

7
    if not feasible_workload:
8
        logger.info(
9
            f"Communication strategy {self.comm.__class__.__name__} "
10
            f"cannot be used (num_chunks={num_chunks}, max_num_tokens={...}). "
11
            f"Falling back to AllGatherReduceScatter.")
12
        self.comm.destroy()
13
        self.comm = AllGatherReduceScatter(mapping=self.mapping)

注意两点：它不负责创建策略，只负责校验和降级（创建在别处惰性发生），以及旧的策略要先 destroy() 再换新的——通信策略可能持有 DeepEP buffer、symm memory 之类的原生资源，直接丢弃会泄漏。ConfigurableMoE 对此还有一段专门的析构说明：DeepEP 的 Buffer.__del__ 内部有一个集合通信式的 barrier，如果各 rank 的 GC 时机不一致，就会出现一部分 rank 进了 barrier、另一部分继续跑，结果是无限挂起。所以策略必须显式、同步地销毁，ConfigurableMoE 也因此实现了上下文管理器协议（with 语句）来保证这一点。

有一处值得留意的约束：

1
if num_chunks > 1 and moe.backend.capabilities.supports_moe_lora and moe.backend._moe_lora_active(lora_params):
2
    raise_moe_lora_multichunk_unsupported(num_chunks)

开了 MoE LoRA 就不允许分块。原因是 LoRA adapter 的选择是 per-request 的，分块以后 adapter 与块的对应关系会变得需要额外记账；框架的选择是直接拒绝并给出明确错误，而不是实现一套复杂的记账。这类”明确拒绝优于默默降级”的取舍在 MoE 模块里出现频率很高——它和前面 can_implement 必须纯、“降级要写进报告”是同一种工程哲学：宁可让调用方看到一条清晰的失败，也不要让它在性能上悄悄差一半。

EP 的 rank 坐标系

MoE 层需要一套和 TP 不同的并行坐标。代码里从 mapping 上取的是专门的一组字段：

1
self.cluster_rank = model_config.mapping.moe_cluster_rank
2
self.cluster_size = model_config.mapping.moe_cluster_size
3
self.smart_router = True if self.cluster_size > 1 else False
4

5
self.tp_rank = model_config.mapping.moe_tp_rank
6
self.tp_size = model_config.mapping.moe_tp_size
7

8
self.ep_size = model_config.mapping.moe_ep_size
9
self.ep_rank = model_config.mapping.moe_ep_rank

上一篇讲注意力层的张量并行时提到过：注意力侧用的是 tp_size，MoE 侧用的是 moe_tp_size / moe_ep_size，两者在同一组进程上可以被解释成不同的划分。moe_cluster_* 是更外一层的概念（多机集群层级），smart_router 表示当存在多个集群时路由器需要感知拓扑。

_compute_ep_partition：整除之外的分布

每张卡分几个专家？最朴素的答案是 num_experts / ep_size。但 256 个专家配 6 张卡呢？代码里的回答是：

1
def _compute_ep_partition(num_experts: int, ep_size: int, ep_rank: int) -> tuple:
2
    """计算本 rank 的专家数与槽位边界。
3

4
    用 ceil/floor 分布: rank 0..remainder-1 持有 (base+1) 个专家,
5
    其余 rank 持有 base 个。即使 num_experts % ep_size != 0 也能覆盖全部专家。
6

7
    Returns: (expert_size, slot_start, slot_end)
8
    """
9
    if num_experts < ep_size:
10
        raise ValueError(
11
            f"num_experts ({num_experts}) must be >= ep_size ({ep_size}); "
12
            f"configurations producing ranks with zero local experts are not supported.")
13
    base = num_experts // ep_size
14
    remainder = num_experts % ep_size
15
    expert_size = base + (1 if ep_rank < remainder else 0)
16
    slot_start = ep_rank * base + min(ep_rank, remainder)
17
    return expert_size, slot_start, slot_start + expert_size

三件事值得说：

向上不向下。 num_experts < ep_size 直接报错，而不是让某些卡分到 0 个专家。理由是”零专家 rank 没有任何后端和通信策略能端到端支持”——一个持有 0 个专家的卡在 dispatch 时收不到任何 token，在 combine 时也没有东西可发，各种集合通信的语义都会变成特例。注释里还提到下游的 all-to-all 算子本来有个 numExperts >= epSize 的兜底检查，但 AllGatherReduceScatter 回退路径绕过了它，所以必须在这里提前拦住。

用 ceil/floor 而不是整除。 slot_start = ep_rank * base + min(ep_rank, remainder) 这一行把前 remainder 个 rank 各自多拿一个专家的偏移算得干干净净。这保证了专家槽位是连续且无重叠的一段，槽位边界可以直接用来算指针偏移。

“非整除 EP”是一项后端要自己认领的能力。 前面 _init_load_balancer 里有这段：

1
if (self.num_experts % self.ep_size != 0
2
        and not type(self)._supports_non_divisible_ep):
3
    raise ValueError(
4
        f"{type(self).__name__} does not support non-divisible EP: "
5
        f"num_experts ({self.num_experts}) must be divisible by ep_size ({self.ep_size}). "
6
        f"Enable EPLB with num_slots divisible by ep_size, or override "
7
        f"`_supports_non_divisible_ep = True` on the subclass after verifying "
8
        f"the kernel/comm path handles ceil/floor partitioning.")

注意错误信息本身就是一个决策树：要么开 EPLB（它有统一的 slot 划分），要么在后端上显式声明”我验证过我的 kernel 和通信路径能处理 ceil/floor 划分”。它不给你一条”自动帮你对齐”的路——因为对齐意味着 pad 或 reshape，那是对正确性的改动，必须由后端的作者确认。

EPLB 路径下为什么反而要求整除

对照一下 EPLB 打开时的检查：

1
if moe_load_balancer_config.num_slots % self.ep_size != 0:
2
    raise ValueError(
3
        f"{type(self).__name__}: with EPLB enabled, num_slots "
4
        f"({moe_load_balancer_config.num_slots}) must be divisible by ep_size ({self.ep_size}) "
5
        f"so each rank holds the same number of slots.")

也就是说，非整除 EP 在 EPLB 下变成了硬性禁止。原因是 EPLB 引入了”槽位”（slot）这个比”专家”更抽象的概念：一张卡持有 num_slots / ep_size 个槽位，每个槽位里放哪个专家是可以动态换的。所有后端 kernel 看到的都是”本卡的槽位数”这个统一形状，专家数是否整除 ep_size 对 kernel 来说已经不可见了。所以要求从”专家数整除”降级成了”槽位数整除”——约束更弱、更好满足，代价是需要 num_slots >= num_experts 的额外显存（多出来的槽位用来做迁移的缓冲）。

这是一个很漂亮的设计：EPLB 不只是个性能开关，它同时改写了并行划分的算术基础。

允许重复的除法与 KV 显存估算的联动

回到上一篇里出现过的一条对照。注意力层里 KV head 用的是向上取整：

1
self.num_key_value_heads = (self.num_key_value_heads + tp_size - 1) // tp_size

MoE 侧的专家切分在没有 EPLB 且后端声明支持非整除时用 ceil/floor，在其余情况下用整除。两者形式不同，但背后的动机一致：框架宁可在少量 rank 上重复持有，也不肯让某张卡空着。

这里的”重复持有”在 MoE 里有一个更实质的含义。当 ep_size > num_experts（比如 4 个专家配 8 张卡）时，向上取整会让多张卡持有同一个专家。这时 token 的 dispatch 需要决定”把 token 发给哪一个副本”——多集群拓扑就是为了处理”同一专家在多个位置上都有副本时该去哪边”的问题。

smart_router 正是这个场景的开关。它的定义只有一处、语义也很直白：

1
self.smart_router = True if self.cluster_size > 1 else False

当 moe_cluster_size > 1（例如一个 GB200 NVL72 被划成多个机柜级集群）时，框架要求路由器”感知拓扑”——它能按集群就近度排序候选专家，尽量把 token 留在本集群内，避免跨集群的远端流量。代价是所有后端都必须实现这条路径：参考实现 VanillaMoE 在 can_implement 里明确以 TOPOLOGY_UNSUPPORTED 拒绝了 smart_router 的部署，理由写在错误信息里——“VanillaMoE has no smart-router path”。这是”后端声明能力、选择器按能力筛选”这套机制的一个真实用例：一条只对多集群拓扑成立的硬约束，通过纯函数的拒绝理由暴露出来，而不是等到运行时算错。

还有一条相关的约束藏在负载均衡器的创建处：maybe_create_moe_load_balancer 只在”模型架构受支持、EP > 1、且没有使用 smart router”时才真正创建 EPLB。也就是说，多集群拓扑下的 MoE 走的是另一条均衡路线，EPLB 的开与关并不是完全独立的开关，而是和拓扑选择绑在一起的。

激活值的”语义常数 → kernel 寄存器”

激活函数这块有一个特别值得学的抽象。模型构建时给的是一个语义描述：

1
self.activation = activation
2
# ActivationType, 而不是裸 int: 作为 IntEnum 它仍然能被 op schema 和
3
# tuning key 当 int 读, 而拒绝信息里可以用 .name.
4
self.activation_type = ActivationType(activation.kind)

而”具体 kernel 需要哪几个寄存器/哪个常量”这件事被推迟到一个专门的转换点。文档里对 activation.py 的描述是：

激活词汇表——模型构建的 MoEActivation 载体、后端声明的 MoEActivationSupport，以及 install_activation_params / materialize_activation_params，这是语义常数变成 kernel 寄存器的唯一地方。

为什么要把这一步单拎出来？ 因为不同后端的融合 kernel 对激活的支持方式完全不同：有的在 kernel 里硬编码了 SwiGLU 的 epilogue，有的支持一个可配置的 clamp 参数（比如 SiTU 的边界），有的只支持几种固定组合。如果让每个后端自己从 activation.kind 翻译，就会出现”同一个模型在不同后端上激活参数含义不同”的隐蔽 bug。把它收成唯一的转换点，并让后端通过 MoEActivationSupport 声明”我支持哪些、我接受哪几个参数”，语义就只在一个地方落地。

通信超时：为什么要有专门的守卫

MoE 的几类通信 kernel 有一个共性——它们会在 kernel 内部等待对端，等待超时后直接 trap（陷入），而 trap 是一个粘性 CUDA 错误：一旦触发，所有在等的 rank 一起死。文档里的原话是”the trap is a sticky CUDA error, so every waiting rank dies”。

问题在于 warmup 阶段：JIT 编译、自动调优、模块加载这些每 rank 只做一次的工作可能让一个健康的 rank 延迟好几分钟，而这期间没有任何集合通信。如果超时是固定的，正常启动就会被判成故障。

moe_comm_timeout_guard.py 就是解决这个的：

MoECommTimeoutGuard 解析两个预算：TRTLLM_MOE_COMM_WARMUP_TIMEOUT_SEC（默认 1800 秒）与 TRTLLM_MOE_COMM_TIMEOUT_SEC（默认取各库原生超时），每次切换时把当前值推给所有注册的 MoECommTimeoutProxy。

各后端的原生超时差别很大，文档列了一张表：

还有一个非常容易出错的细节：

超时是名义值：NVLinkOneSided、DeepGEMM 和 DeepEP 按假定的 2 GHz SM 时钟把秒换算成周期，而 NcclEP 用的是设备实际时钟频率。

为什么这个细节重要？ 因为”把 30 秒换算成 cycle 数”这件事依赖时钟频率假设。如果 kernel 里用的是固定假设，而实际硬件降频或超频，真实的墙钟超时就会偏离配置值。这是一个典型的需要读源码才能发现的坑：你在环境变量里写了 60，实际生效的可能是 40 或 90。

另外，CUDA Graph 捕获期间必须冻结超时：

模型引擎的 is_warmup setter 调用 set_moe_comm_warmup。捕获运行的 CUDA-graph pass 在 moe_comm_serving_timeouts() 里进行，因为被捕获的图会保留捕获时刻的超时值。捕获过程中切换会抛异常。

原因很直接：CUDA Graph 把 kernel 启动参数录进图里，超时值就是参数之一。如果捕获时用的是 warmup 的 1800 秒，那么重放的时候每一次都会带 1800 秒——线上真出事时你要等半小时才看到报错。所以在捕获时切到服务态的值，并在捕获期间拒绝切换。