执行摘要

• 一句话：将 MoE 剩余 PrepareAndFinalize 文件移至独立子目录
• 推荐动作：建议合入，此重构提升了代码可维护性。后续可考虑解决 review 中提出的 in-place 修改问题，但非阻塞。

功能与动机

PR 目标是延续 fused_moe 根目录下文件向子目录迁移的清理工作，将剩余的三个 PrepareAndFinalize 相关实现统一放在 prepare_finalize/ 下，提高代码可维护性。

实现拆解

1. 创建 prepare_finalize/batched.py：将原来在 fused_batched_moe.py 中的 BatchedPrepareAndFinalize 类完整复制到新文件，类定义、方法 prepare、属性等均保持不变。
2. 从 fused_batched_moe.py 删除原类：移除该文件中 BatchedPrepareAndFinalize 类及其相关导入（如 normalize_scales_shape）。
3. 移动并重命名 mori_prepare_finalize.py → prepare_finalize/mori.py：将 mori 的 prepare/finalize 实现移动到子目录，文件名缩写为 mori.py。
4. 移动并重命名 nixl_ep_prepare_finalize.py → prepare_finalize/nixl_ep.py：类似操作，文件名缩写为 nixl_ep.py。
5. 更新导入路径：在 all2all_utils.py 中将原来对 .mori_prepare_finalize 和 .nixl_ep_prepare_finalize 的引用改为 .prepare_finalize.mori 和 .prepare_finalize.nixl_ep。并在 prepare_finalize/__init__.py 中添加 BatchedPrepareAndFinalize 的导出。
6. 测试文件调整：修改 test_batched_deepgemm.py、tests/kernels/moe/utils.py、modular_kernel_tools/mk_objects.py 中的导入路径，指向新的子包位置。

关键文件：

• vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py（模块 Batched PF；类别 source；类型 data-contract；符号 BatchedPrepareAndFinalize, init, activation_format, max_num_tokens_per_rank）: 新文件，将 BatchedPrepareAndFinalize 类从旧位置迁移至此，是本次重构的核心。
• vllm/model_executor/layers/fused_moe/fused_batched_moe.py（模块 Fused Batched；类别 source；类型 data-contract；符号 BatchedPrepareAndFinalize）: 删除了 BatchedPrepareAndFinalize 类及相关导入，是该类搬走的源位置。
• vllm/model_executor/layers/fused_moe/all2all_utils.py（模块 All2All Utils；类别 source；类型 rename-or-move）: 更新了 mori 和 nixl_ep 的导入路径，从旧位置改为 prepare_finalize 子包。
• vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py（模块 Init；类别 source；类型 rename-or-move）: 添加 BatchedPrepareAndFinalize 的导出，使其可通过子包访问。
• vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py（模块 Mori EP；类别 source；类型 rename-or-move）: 从 fused_moe/mori_prepare_finalize.py 重命名并移动至此。
• vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py（模块 Nixl EP；类别 source；类型 rename-or-move）: 从 fused_moe/nixl_ep_prepare_finalize.py 重命名并移动至此。
• tests/kernels/moe/test_batched_deepgemm.py（模块 Test；类别 test；类型 test-coverage）: 更新了 BatchedPrepareAndFinalize 的导入路径。
• tests/kernels/moe/utils.py（模块 Test Utils；类别 test；类型 test-coverage）: 更新了 BatchedPrepareAndFinalize 的导入路径。
• tests/kernels/moe/modular_kernel_tools/mk_objects.py（模块 MK Objects；类别 test；类型 test-coverage）: 更新了 BatchedPrepareAndFinalize 的导入路径。

关键符号：BatchedPrepareAndFinalize.init, BatchedPrepareAndFinalize.prepare, BatchedPrepareAndFinalize.activation_format, BatchedPrepareAndFinalize.max_num_tokens_per_rank, BatchedPrepareAndFinalize.topk_indices_dtype, BatchedPrepareAndFinalize.num_dispatchers, BatchedPrepareAndFinalize.output_is_reduced

关键源码片段

vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py

新文件，将 BatchedPrepareAndFinalize 类从旧位置迁移至此，是本次重构的核心。

class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
    """
    # 将 tokens 重新组织为专家批处理格式 E x max_num_tokens x K
    """
​
    def __init__(self, max_num_tokens: int, num_local_experts: int, num_dispatchers: int, rank: int):
        self.max_num_tokens = max_num_tokens
        self.num_local_experts = num_local_experts
        self.rank = rank
        self.num_dispatchers_ = num_dispatchers
​
    @property
    def activation_format(self) -> mk.FusedMoEActivationFormat:
        return mk.FusedMoEActivationFormat.BatchedExperts
​
    # ... 其他属性
​
    def prepare(self, a1, topk_weights, topk_ids, num_experts, expert_map, apply_router_weight_on_input, quant_config, defer_input_quant=False):
        # defer_input_quant 被暂时禁止
        # 原地乘权重（可能有风险，见 review）
        if apply_router_weight_on_input:
            a1.mul_(topk_weights.to(a1.dtype))
        # 创建批量化张量 b_a1
        tokens_per_expert = torch.zeros(num_experts, dtype=torch.int, device=a1.device)
        # 循环处理本地专家
        for expert_id in range(first_expert, last_expert):
            # ...

评论区精华

Gemini Code Assist 审阅：指出两处潜在问题：
1. prepare() 方法中原地修改输入 a1 可能影响残差连接，建议改为 a1 = a1 * topk_weights.to(a1.dtype) 避免副作用。
2. tokens_per_expert 张量应分配为 self.num_local_experts 大小，而非全局 num_experts，以减少冗余 kernel 块。
作者回应：同意第一点但表示是原代码遗留问题，暂不修改；第二点未回复。

• 原地修改输入张量潜在风险 (correctness): 作者同意但表示这是原代码遗留问题，暂不修改。
• tokens_per_expert 大小优化 (performance): 未收到作者回复，暂未修改。

风险与影响

• 风险：风险极低。所有变更为文件移动和导入调整，无核心逻辑变更。测试文件同步更新且 CI 通过。唯一风险是：如果存在未识别到的第三方代码或外部脚本直接引用旧路径，可能导致 ImportError，但概率极低。
• 影响：影响范围小。仅影响 MoE 模块内部开发者及该模块的后续维护。对用户无功能影响。代码组织更清晰，便于后续扩展。
• 风险标记：导入路径变更

关联脉络

• PR #40671 [MoE Refactor] Rename FusedMoE.make_expert_params_mapping to fused_moe_make_expert_params_mapping: 同属 MoE 模块化重构系列，均涉及文件移动或重命名。
• PR #40568 [MoE] Move xpu moe to fused_moe/experts/: 类似的文件迁移重构，将 XPU MoE 移至子目录。