Prhub

#35290 [XPU] Lazily import tvm_ffi-dependent all_reduce kernel in minimax_m2

原始 PR 作者 koppulac 合并时间 2026-08-28 17:18 文件变更 1 提交数 7 评论 0 代码增减 +8 / -4

执行摘要

XPU 上延迟导入 tvm_ffi 相关 all_reduce 内核,修复 MiniMax M2 模型加载崩溃

在 XPU 环境中,导入 tvm_ffi 依赖的 all_reduce 内核会导致 MiniMax M2 模型模块加载失败,因为 XPU 上不支持 tvm_ffi。此 PR 旨在通过延迟导入机制,跳过 XPU 上的相关导入,从而让模型能够在 XPU 上正常运行,同时不影响其他后端的行为。

该 PR 简单且低风险,值得快速合并。不建议精读,但可关注 XPU 平台支持的类似模式,作为后续其他模型在 XPU 上兼容性修复的参考。

讨论亮点

无 review 评论,无讨论内容。

实现拆解

实现步骤如下:

  1. 在 sglang.srt.utils 导入中增加 is_xpu 函数。
  2. 模块加载时检测当前是否为 XPU 环境,并将 _is_xpu 作为模块级变量缓存。
  3. 将原本在模块顶层无条件导入的 fused_parallel_qknorm 和 get_fused_parallel_qknorm_max_occupancy 移入 if not _is_xpu: 条件块中,仅当非 XPU 时才导入。
  4. 保持 NPU 相关的 split_qkv_tp_rmsnorm_rope 导入逻辑不变,确保 NPU 后端行为不受影响。
文件 模块 状态 重要度
python/sglang/srt/models/minimax_m2.py 模型层 modified 6.64

关键符号

is_xpu

关键源码片段

python/sglang/srt/models/minimax_m2.py core-logic

该文件是 MiniMax M2 模型的入口,通过条件化导入修复了 XPU 环境下因 tvm_ffi 依赖导致的加载失败问题。

# 在 sglang.srt.utils 导入中新增 is_xpu
from sglang.srt.utils import (
    BumpAllocator,
    add_prefix,
    cpu_has_amx_support,
    get_bool_env_var,
    get_compiler_backend,
    is_cpu,
    is_cuda,
    is_non_idle_and_non_empty,
    is_npu,
    is_xpu, # 新增,用于检测 XPU 平台
    make_layers,
)# 在模块级别缓存平台检测结果
_is_cpu = is_cpu()
_is_amx_available = cpu_has_amx_support()
_is_cuda = is_cuda()
_is_npu = is_npu()
_is_xpu = is_xpu() # 新增 XPU 检测# 仅当非 XPU 时导入 tvm_ffi 依赖的 all_reduce 内核,避免 XPU 上导入失败
if not _is_xpu:
    from sglang.kernels.ops.communication.all_reduce import (
        fused_parallel_qknorm,
        get_fused_parallel_qknorm_max_occupancy,
    )# NPU 相关导入保持不变
if _is_npu:
    from sgl_kernel_npu.norm.split_qkv_tp_rmsnorm_rope import split_qkv_tp_rmsnorm_rope

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

风险较低。该 PR 仅涉及导入逻辑的条件化,不影响非 XPU 平台的行为。可能存在风险是:若 fused_parallel_qknorm 等函数在其他非 XPU 平台被间接依赖,延迟导入可能影响某些边缘场景,但模块内已明确调用这些函数,且不涉及执行路径的变更,风险可控。

影响范围限于 XPU 环境下的 MiniMax M2 模型加载,修复了启动崩溃问题。对 CPU、CUDA、NPU 等后端无影响,因为它们仍会导入相关内核。对团队而言,提升了 XPU 平台对 MiniMax M2 模型的支持能力。

导入路径条件化,低风险 无测试覆盖

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论