执行摘要

• 一句话：新增 MoE 输出契约，为 top-k 归约与尾部融合铺路
• 推荐动作：值得精读，特别是关注 MoE 内核或性能优化的工程师。设计上有四点值得借鉴：契约分层（UnfinalizedMoEOutput 与 MoEOutput 各自职责清晰）；守卫条件集中到配置层并按需读取；0-token 边界处理；针对 FlashInfer 多版本返回形态的兼容收口。建议同时跟踪后续消费方 PR 与测试补充；若团队计划接入新的 MoE 内核，应先熟悉此契约。

功能与动机

PR body 未填写具体描述（仅为模板占位），动机需从标题、代码注释与 review 讨论中还原。标题「MoE tail fusion」与 UnfinalizedMoEOutput 的 docstring 表明：当前一体化 MoE 内核在 GEMM2 结束后必须立即终结（do_finalize=True），top-k 归约作为独立 kernel 运行，随后还要单独执行共享专家相加与 TP all-reduce。本 PR 允许内核停在 GEMM2 后，把 permuted 未加权输出、路由权重与 permute map 交回上层，使 top-k 归约能与后续算子融合，减少 kernel 启动次数与中间张量读写。zyongye 在 review 中进一步明确了边界：「We only want to defer this when we are using TP and if flashinfer_trtllm kernel is selected.」

实现拆解

1. 定义输出契约（新增 moe_output.py）：新文件定义两个 core dataclass——UnfinalizedMoEOutput 封装 GEMM2 后未终结的三元组（gemm2_permuted、expert_weights、expanded_idx_to_permuted_idx），并声明三条硬约束：gemm2_permuted 在 hidden_dim 维必须紧凑连续（行数为 autotuner 相关填充值，永不被引用）；expert_weights 必须精确匹配激活 dtype；permute map 为 [num_tokens, top_k] 的 int32 且 -1 表示非本 rank 专家。MoEOutput 表示归约仍开放的层输出，由 reduce_results=False 的层返回，消费方（如下一层 RMSNorm）可将 TP all-reduce 融入自身。该文件最初定义在 modular_kernel.py 中，经 review 后独立成文件，并将代号从 DeferredMoEOutput 改为 UnfinalizedMoEOutput。
2. 增加配置开关与守卫（config.py）：FusedMoEConfig 新增 defer_moe_finalize: bool = False 字段，以及只读属性 use_deferred_moe_finalize，返回 self.defer_moe_finalize and self.tp_size > 1。属性按需读取而不是在 __post_init__ 中固化，原因与 skip_final_all_reduce 一致：该标志由拥有融合消费方的层在构造之后设置。TP=1 时不存在可供 top-k 归约融合的 all-reduce，推迟形态没有收益，故以 tp_size > 1 作为守卫。不具备推迟能力的内核会忽略该字段，能力收敛在具体内核实现中。
3. 适配两个 TRTLLM 内核（trtllm_bf16_moe.py / trtllm_nvfp4_moe.py）：两个 apply 的返回类型从 torch.Tensor 放宽为 torch.Tensor | UnfinalizedMoEOutput，并以 use_deferred_moe_finalize and num_tokens > 0 计算 defer——空闲 rank 的 0-token dummy forward 必须保留 finalized 空形态，因为 runner 在 host 端按 token 数取平均。defer 时向 flashinfer 传递 do_finalize=False，内核返回三元组被包装为 UnfinalizedMoEOutput，其中扁平 permute map 被转换为 [num_tokens, top_k] 的 int32 视图（依赖 self.topk）；非 defer 路径保持原行为，bf16 路径额外兼容了 do_finalize=True 在部分 FlashInfer 版本返回裸 tensor、另一些返回单元素 list 的差异。
4. 测试与落地配套：本 PR 未包含直接对应的测试文件变更，UnfinalizedMoEOutput 的实际消费方（融合 top-k 归约的 runner / RMSNorm 逻辑）也未在本 PR 落地，属于系列工程的第一步。默认开关关闭，现有模型推理路径完全不受影响。

关键文件：

• vllm/model_executor/layers/fused_moe/moe_output.py（模块 MoE 内核；类别 source；类型 data-contract；符号 UnfinalizedMoEOutput, MoEOutput）: 本 PR 的核心交付物：定义 MoE 层与消费者之间的输出契约（UnfinalizedMoEOutput / MoEOutput），是 tail fusion 系列的数据规范基础。
• vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py（模块 MoE 内核；类别 source；类型 data-contract；符号 TrtLlmBf16ExpertsBase.apply）: BF16 TRTLLM-Gen 内核的 apply 路径适配：支持 do_finalize=False 并返回 UnfinalizedMoEOutput，同时兼容 FlashInfer 多版本返回形态。
• vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py（模块 MoE 内核；类别 source；类型 data-contract；符号 TrtLlmNvFp4ExpertsBase.apply）: NVFP4 TRTLLM-Gen 内核的 apply 路径同步适配，使 FP4 量化 MoE 也能走未终结输出路径。
• vllm/model_executor/layers/fused_moe/config.py（模块 MoE 配置；类别 source；类型 configuration；符号 defer_moe_finalize, use_deferred_moe_finalize）: 新增 defer_moe_finalize 字段与 use_deferred_moe_finalize 守卫属性，把推迟终结能力收敛到配置层并要求 TP > 1。

关键符号：UnfinalizedMoEOutput, MoEOutput, TrtLlmBf16ExpertsBase.apply, TrtLlmNvFp4ExpertsBase.apply, FusedMoEConfig.use_deferred_moe_finalize

关键源码片段

vllm/model_executor/layers/fused_moe/moe_output.py

本 PR 的核心交付物：定义 MoE 层与消费者之间的输出契约（UnfinalizedMoEOutput / MoEOutput），是 tail fusion 系列的数据规范基础。

# SPDX-License-Identifier: Apache-2.0
"""MoE 层与消费方之间的输出契约，支撑 MoE tail fusion。
​
一体化 MoE 内核若能在 GEMM2 之后停下（TRTLLM-Gen 的 ``do_finalize=False``
路径），就把 permuted 且未加权的输出、路由权重和 permute map 原样交回上层，
让 top-k 归约能够与后续的共享专家相加以及 TP all-reduce 融合，而不是单独
运行一个 kernel。
"""
​
from dataclasses import dataclass
​
import torch
​
​
@dataclass
class UnfinalizedMoEOutput:
    """未终结（unfinalized）的 MoE 内核输出。
​
    消费方按行索引 ``gemm2_permuted``：它必须在 ``hidden_dim`` 维上紧凑连续
    排布，行数是 autotuner 依赖的填充值，永远不会被引用。
    """
​
    # [num_permuted_rows, hidden_dim]，permuted 且未加权的 GEMM2 输出。
    gemm2_permuted: torch.Tensor
    # [num_tokens, top_k] 路由权重，必须与激活 dtype 精确匹配，否则按更宽
    # dtype 读取会被当成垃圾数据；对把 routed_scaling_factor 折进权重的
    # 路由方法，权重已经携带该系数。
    expert_weights: torch.Tensor
    # [num_tokens, top_k] int32 permute map，-1 表示该 expert 不在本 rank。
    expanded_idx_to_permuted_idx: torch.Tensor
​
​
@dataclass
class MoEOutput:
    """MoE 层输出，归约步骤仍然保持开放。
​
    由以 ``reduce_results=False`` 运行 MoE 的层返回，消费方（通常是下一层的
    RMSNorm）可以把 TP all-reduce 融合进自身，省掉一次独立 kernel。当 routed
    输出仍是未终结形态时，top-k 归约同样保持开放，可以并入同一个 kernel。
    生产方只有在确定存在可消费的融合路径时才返回未终结形态，因此消费方看到
    ``UnfinalizedMoEOutput`` 就意味着融合路径适用，无需重新推导。
    """
​
    # 未归约的路由输出，可能是已终结或未终结形态。
    routed: torch.Tensor | UnfinalizedMoEOutput
    # 未归约的共享专家输出，由执行终结的一方负责并入。
    shared_output: torch.Tensor | None = None
    # 施加到路由输出上的缩放系数，在共享专家相加之前应用；对已经折进路由
    # 权重的方法，该值保持 1.0。
    routed_scaling_factor: float = 1.0

vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py

BF16 TRTLLM-Gen 内核的 apply 路径适配：支持 do_finalize=False 并返回 UnfinalizedMoEOutput，同时兼容 FlashInfer 多版本返回形态。

    num_tokens = hidden_states.shape[0]
    # runner 在 host 端按 token 数取平均，因此空闲 rank 的 0-token dummy
    # forward 必须保留 finalized（空）形态，不能走未终结路径。
    defer = self.moe_config.use_deferred_moe_finalize and num_tokens > 0
​
    routing_replay_out = self._maybe_make_routing_replay_buffer(
        num_tokens=num_tokens,
        device=hidden_states.device,
    )
    out = flashinfer.fused_moe.trtllm_bf16_moe(
        routing_logits=router_logits,
        routing_bias=e_score_correction_bias,
        hidden_states=hidden_states,
        gemm1_weights=w1,
        gemm2_weights=w2,
        num_experts=global_num_experts,
        top_k=self.topk,
        n_group=num_expert_group,
        topk_group=topk_group,
        intermediate_size=self.intermediate_size_per_partition,
        local_expert_offset=self.ep_rank * self.local_num_experts,
        local_num_experts=self.local_num_experts,
        routed_scaling_factor=routed_scaling_factor,
        routing_method_type=self.routing_method_type,
        activation_type=activation_to_flashinfer_int(activation),
        tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
        routing_replay_out=routing_replay_out,
        do_finalize=not defer, # 只有需要融合 top-k 归约时才停在 GEMM2 之后
    )
    self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens)
    if defer:
        # flashinfer 返回扁平 permute map；契约要求 [num_tokens, top_k] 形状，
        # 让消费方可以直接从 shape 读出 top_k。
        return UnfinalizedMoEOutput(
            gemm2_permuted=out[0],
            expert_weights=out[1],
            expanded_idx_to_permuted_idx=out[2]
            .to(torch.int32)
            .view(num_tokens, self.topk),
        )
    # do_finalize=True 在部分 FlashInfer 版本返回裸 tensor，另一些返回单元素
    # list，这里统一取出张量返回。
    return out[0] if isinstance(out, (list, tuple)) else out

vllm/model_executor/layers/fused_moe/config.py

新增 defer_moe_finalize 字段与 use_deferred_moe_finalize 守卫属性，把推迟终结能力收敛到配置层并要求 TP > 1。

class FusedMoEConfig:
    # 当为 True 时，允许能在 GEMM2 后停下的专家内核返回 UnfinalizedMoEOutput
    # 而不是终结态，把 top-k 归约留给消费方融合。该字段由具备融合消费方的
    # 层在对象构造之后设置，因此读取统一走 use_deferred_moe_finalize 属性
    # 完成守卫判断；不具备该能力的内核会忽略此设置。默认 False。
    defer_moe_finalize: bool = False
​
    # ... 其余字段省略 ...
​
    @property
    def use_deferred_moe_finalize(self) -> bool:
        """当前部署下专家是否允许返回未终结输出。
​
        按需读取而不是在 __post_init__ 中固化，是因为 defer_moe_finalize
        与 skip_final_all_reduce 一样在对象构造之后才被设置。没有 TP 时不
        存在可供 top-k 归约融合的 all-reduce，推迟形态没有收益，因此这里
        用 tp_size > 1 作为守卫。
        """
        return self.defer_moe_finalize and self.tp_size > 1

评论区精华

zyongye 的首轮 review 认可独立 dataclass 的形态，但提出三点收紧建议：defer 开关应放到 MoE 配置而非内核类；契约 dataclass 应独立成文件；DeferredMoEOutput 命名不佳。作者逐条回应并落地：开关移入 FusedMoEConfig（commit 748519c）、dataclass 移入新文件 moe_output.py、命名改为 UnfinalizedMoEOutput。此外 zyongye 明确要求加守卫——只有 TP 场景、且选中 flashinfer_trtllm 内核时才允许推迟终结，最终通过 use_deferred_moe_finalize 的 tp_size > 1 守卫以及 defer 逻辑仅存在于两个 TRTLLM 内核来实现。终审追加 APPROVED（LGTM）。

• defer 开关归属：内核类属性 vs MoE 配置 (design): 作者通过 commit 748519c 落实：开关下沉到 FusedMoEConfig，新增 defer_moe_finalize 字段与 use_deferred_moe_finalize 守卫属性，内核统一从 self.moe_config 读取。
• 输出契约 dataclass 是否独立成文件 (design): 作者将其移入新文件 moe_output.py，并顺带组织 MoEOutput 契约，modular_kernel.py 不再承载该数据结构。
• DeferredMoEOutput 命名 (style): 作者改为 UnfinalizedMoEOutput，强调「未终结」状态而非「推迟」动作，语义更贴合 GEMM2 之后的中间产物。
• defer 路径的守卫条件（TP 与内核选择） (design): 通过 use_deferred_moe_finalize 的 tp_size > 1 守卫实现 TP 限制；内核维度上，defer 逻辑只存在于 trtllm_bf16 与 trtllm_nvfp4 两个具体内核中，其余内核天然不受影响。

风险与影响

• 风险：
  1. 测试缺口：本次改动没有任何直接对应的测试文件变更，defer 分支、permute map 形状转换（扁平 → [num_tokens, top_k]）与 0-token 守卫均无自动化测试保护，回归风险主要靠后续消费方 PR 与人工验证兜底。
  2. 返回类型契约扩展：apply 返回值从 torch.Tensor 变为联合类型，所有调用方（runner、modular wrapper）必须能识别 UnfinalizedMoEOutput；当前默认关闭使风险受控，但未来开启后未适配的调用方会直接报错。
  3. FlashInfer 返回形态差异：bf16 路径显式兼容了 do_finalize=True 时裸 tensor / 单元素 list 两种形态；nvfp4 路径始终按 result[0] 下标解包，若未来 FlashInfer 版本改变返回结构可能出错。
  4. dtype 与形状敏感假设：expert_weights 必须精确匹配激活 dtype（更宽的 buffer 会被读出垃圾值）；permute map 的 view 依赖 self.topk，与内核实际 topk 不一致会静默产出错误结果。
     - 影响：对用户：默认开关关闭，无任何用户可感知行为变化。对系统：为 MoE tail fusion 打地基——top-k 归约可并入共享专家相加与 TP all-reduce，预期减少 kernel 启动次数与中间张量读写，但真实收益仍需后续消费方 PR 兑现。对团队：moe_output.py 成为 MoE 层输出契约的规范文件，后续新内核（如其他 backend 的 do_finalize=False 支持）与融合消费方（RMSNorm、runner）都必须遵循该契约；FusedMoEConfig 新增的开关字段为模型层提供显式 opt-in 入口。
     - 风险标记：缺少测试覆盖, 核心路径契约变更, 依赖 FlashInfer 返回形态, dtype 与形状敏感假设

关联脉络

• PR #51819 [Bugfix][MoE] Support GELU tanh in FlashInfer B12x MoE: 同属 fused_moe/experts 的 FlashInfer MoE 内核能力演进，与本 PR 共享 MoE 内核契约演进脉络。
• PR #49444 [Misc] Enable test_silu_mul_fp8_quant_deep_gemm on XPU: 同目录 MoE 内核修复与测试使能，反映 fused_moe 内核生态的持续迭代。
• PR #51473 [ROCm][DSV4] Preserve native MXFP4 TP8 shard allocation: 同为 fused_moe 下 MoE 内核/量化路径优化，与本 PR 减少不必要 kernel 开销的方向一致。