执行摘要

• 一句话：回滚 #50879 模块级导入，修复 ROCm 误加载 tilelang 导致 CI 失败
• 推荐动作：值得快速合入以恢复 CI 稳定性，但不必精读实现细节；关注点应放在后续如何在不触发 tilelang 加载的前提下重新实现 #50879 的优化目标。对该回滚本身，建议记录到关联 issue #51151 中作为临时缓解措施，并推动一个带 ROCm 平台隔离与回归测试的正式修复。

功能与动机

PR body 明确说明：This PR was merged while AMD CI was failing ... due to it, even though it appears seemingly unrelated。关联 issue #51151 指出 #50879 causes wrongful hipMalloc to be loaded on ROCm due to tilelang load, causing tests and allocation failures，错误堆栈最终落到 RuntimeError: invalid argument for fmha_v3_varlen_fwd。因此需要立即回滚以恢复 AMD CI 和 ROCm 运行时的稳定性。

实现拆解

1. 变更入口：唯一改动文件 vllm/model_executor/layers/fused_moe/all2all_utils.py，该文件是 MoE 专家并行（EP）all2all 调度与 prepare_finalize 后端选择的公共工具模块。
2. 核心回滚动作：删除 #50879 新增的模块级条件导入块（if current_platform.is_cuda_alike(): 下的 has_deep_ep() / has_deep_ep_v2() / has_mori() / has_nixl_ep() 判断及对应 from .prepare_finalize.* import ...），恢复 maybe_roundup_layer_hidden_size 与 maybe_make_prepare_finalize 内部原有的函数级惰性导入写法。
3. 行为差异：恢复后，DeepEP、Mori、NixlEP 等后端只在实际命中对应 use_*_kernels 分支时才导入对应模块，避免在 vllm serve 启动阶段通过 has_nixl_ep() 等探测间接加载 tilelang，从而规避 ROCm 上 hipMalloc 被错误触发的回归。
4. 测试与配套：本 PR 未附带测试变更；属于对已合并回归的直接回滚。后续需要为 #50879 的目标（减少启动导入开销）设计带平台隔离的正确实现，并在 ROCm CI 中补充覆盖。

关键文件：

• vllm/model_executor/layers/fused_moe/all2all_utils.py（模块 MoE 通信；类别 source；类型 revert；符号 maybe_roundup_layer_hidden_size, maybe_make_prepare_finalize, get_ep_all2all_manager）: 唯一变更文件，回滚 #50879 的模块级条件导入，恢复 DeepEP/Mori/NixlEP 后端的函数内惰性导入，解决 ROCm 上 tilelang 误加载导致的 hipMalloc 与 fmha 崩溃。

关键符号：maybe_roundup_layer_hidden_size, maybe_make_prepare_finalize

关键源码片段

vllm/model_executor/layers/fused_moe/all2all_utils.py

唯一变更文件，回滚 #50879 的模块级条件导入，恢复 DeepEP/Mori/NixlEP 后端的函数内惰性导入，解决 ROCm 上 tilelang 误加载导致的 hipMalloc 与 fmha 崩溃。

# vllm/model_executor/layers/fused_moe/all2all_utils.py
# 回滚 #50879 后，prepare_finalize 各后端恢复为“函数内惰性导入”。
# 这样可避免在每次 `vllm serve` 配置 MoE 时无条件加载
# DeepEP / Mori / NixlEP / tilelang 等可选依赖，
# 防止 ROCm 环境因 tilelang 导入链误触发 hipMalloc 分配失败。
​
from typing import Any
​
import torch
​
from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEConfig,
    FusedMoEParallelConfig,
    FusedMoEQuantConfig,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import (
    FusedMoEPrepareAndFinalize,
)
​
​
def maybe_roundup_layer_hidden_size(
    hidden_size: int,
    act_dtype: torch.dtype,
    moe_parallel_config: FusedMoEParallelConfig,
) -> int:
    """根据 MoE 并行配置与 all2all 后端决定是否将 hidden_size 向上取整。"""
    if moe_parallel_config.use_deepep_ht_kernels:
        # 仅在真正使用 DeepEP HT kernel 时才导入，避免启动阶段加载无关依赖
        from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize
        hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size(
            hidden_size, act_dtype
        )
​
    if moe_parallel_config.use_deepep_ll_kernels:
        from .prepare_finalize.deepep_ll import DeepEPLLPrepareAndFinalize
        hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size(
            hidden_size
        )
​
    if moe_parallel_config.use_deepep_v2_kernels:
        from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize
        hidden_size = DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size(
            hidden_size, act_dtype
        )
​
    if moe_parallel_config.use_nixl_ep_kernels:
        from .prepare_finalize.nixl_ep import NixlEPPrepareAndFinalize
        hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size(
            hidden_size
        )
​
    return hidden_size

评论区精华

1. AMD CI 回归的根因定位（来自 PR body 与 issue #51151）：#50879 看似与 AMD 无关，但其中 has_nixl_ep() 等探测函数的调用链导致 tilelang 被加载，进而在 ROCm 上触发错误的 hipMalloc，最终表现为 fmha_v3_varlen_fwd 的 invalid argument。结论是直接回滚，恢复 AMD CI。
2. 自动化 review 与批准：claude[bot] 因 PR 来自 fork 而跳过自动审查；AndreasKaratzas 触发 Buildkite CI #82513 后与 mgoin 均批准该回滚，未提出额外修改意见。

• AMD CI 回归与 tilelang 误加载的根因 (correctness): 直接回滚 #50879，恢复原有函数内惰性导入；根因仍需后续修复。
• CI 触发与审批流程 (other): PR 获批并合入，CI 验证通过。

风险与影响

• 风险：
  1. 导入开销回归：回滚后恢复 #50879 想消除的启动期依赖探测路径，maybe_roundup_layer_hidden_size 和 maybe_make_prepare_finalize 在首次进入对应分支时会重复执行局部导入，但这与历史行为一致，实际开销只在匹配到 EP 内核分支时才产生。
  2. 根因未彻底解决：tilelang 被误加载的根本问题（ROCm 环境下 hipMalloc 与 aiter/tilelang 的交互）仍然存在，未来若有人重新尝试模块级预导入，可能再次踩坑。
  3. 缺少回归测试：该回滚没有配套测试，无法防止未来类似导入策略变更再次引入 ROCm 回归；建议后续补一个 CI 层防护（如在 ROCm buildkite 上校验关键 kernel 路径的启动行为）。
     - 影响：对用户的影响主要落在 ROCm 平台：回滚后 AMD CI 恢复稳定，ROCm 上 MoE 模型（尤其使用 EP/all2all 后端时）不再因 tilelang 误加载而崩溃。对系统而言，all2all_utils.py 的导入方式回到函数内惰性导入，降低了启动阶段的动态依赖风险。对团队而言，这是一个紧急恢复动作，需要尽快合入 main 并后续跟踪 #50879 的正确修复方向。
     - 风险标记：导入开销回归, 根因未修复, 缺少回归测试

关联脉络

• PR #50879 [Misc] Avoid importing nixl_ep on every vllm serve config: 被本 PR 回滚的对象；其模块级导入策略在 ROCm 上触发 tilelang 误加载并导致 AMD CI 失败。
• PR #51151 [Bug][Regression] #50879 causes wrongful hipMalloc to be loaded on ROCm due to tilelang load, causing tests and allocation failures: 关联 issue，记录了 #50879 导致的 ROCm 回归细节，是本 PR 回滚的直接依据。