# PR #39009 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[MoE] Move remaining PrepareAndFinalize to prepare finalize folder
- 合并时间：2026-04-24 08:00
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/39009

---

# 执行摘要

- 一句话：将 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 类从旧位置迁移至此，是本次重构的核心。

```python
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 移至子目录。