Prhub

#21633 [Diffusion][NPU] Add support for MOVA

原始 PR 作者 LLThomas 合并时间 2026-04-03 10:33 文件变更 3 提交数 3 评论 19 代码增减 +54 / -14

npu diffusion feature jit-kernel

执行摘要

为 MOVA 扩散模型添加 NPU 硬件支持，优化 RoPE 和内存格式兼容性。

根据PR body，动机是'This PR adds NPU support to the MOVA model.'，并引用Issue #18967 '[Roadmap] [NPU] Sglang Diffusion on Ascend'，作为NPU支持路线图的一部分，旨在扩展硬件覆盖并解决现有NPU平台上的执行失败问题。

建议涉及NPU开发或Diffusion模型的工程师精读此PR，关注平台检测设计（如current_platform.is_npu）、RoPE优化策略（torch_npu.npu_rotary_mul集成）以及如何处理硬件不兼容性（如channels_last_3d和complex128临时修复）。这些设计决策值得学习，以应对类似跨平台支持场景。

讨论亮点

review中核心讨论点：1) gemini-code-assist[bot]建议提取共享函数以减少visual_freqs和audio_freqs准备的冗余代码，但未直接采纳；2) ping1jing2和Makcum888e建议减少代码中的is_npu检查，推动改动到npu_fallback.py中，作者LLThomas响应并修改相关代码；3) Makcum888e要求添加注释说明临时修复（如'remove it when torch.cat on NPU is supported torch.complex128'），作者已添加；4) ping1jing2关注CI触发问题，指出Diffusion PR意外触发LLM CI，合并后计划优化CI工作流。讨论结论是问题已解决，代码改进并合并。

实现拆解

实现分三部分：1) 在python/sglang/jit_kernel/diffusion/triton/npu_fallback.py中添加NPU特定的Rotary Embedding优化，引入torch_npu.npu_rotary_mul以提高性能，并设置常量限制；2) 在python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py中修改patchify函数，根据current_platform.is_npu选择contiguous方式，避免NPU不支持的torch.channels_last_3d；3) 在python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py中调整visual_freqs和audio_freqs准备逻辑，处理NPU上torch.cat对torch.complex128的不兼容，临时使用torch.complex64。

文件 模块 状态 重要度
python/sglang/jit_kernel/diffusion/triton/npu_fallback.py JIT-kernel/Diffusion modified 7.0
python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py Diffusion/Models modified 6.0
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py Diffusion/Pipelines modified 7.0

分析完成后，这里会展示 LLM 生成的相对完整源码片段和详细注释。

关键符号

apply_rotary_embedding_native patchify inference_single_step

评论区精华

代码冗余和共享函数提取 设计

gemini-code-assist[bot] 建议提取 visual_freqs 和 audio_freqs 准备的共享函数以减少冗余，但未直接采纳。

结论：代码已修改，但未实现共享函数，保持原有逻辑。 · 已解决

减少 is_npu 检查并集中优化 设计

ping1jing2 和 Makcum888e 建议将改动集中到 npu_fallback.py，避免在多个文件中分散 is_npu 检查，以提升代码可维护性。

结论：作者响应并修改了相关代码，推动优化到统一位置。 · 已解决

添加注释和代码一致性 documentation

Makcum888e 要求添加注释说明临时修复（如 torch.cat 不兼容），ping1jing2 确保注释与代码逻辑一致。

结论：作者添加了所需注释，提高了代码清晰度。 · 已解决

CI 工作流触发问题 infra

ping1jing2 指出 Diffusion PR 因修改 jit_kernel 文件意外触发 LLM CI 测试，需优化 CI 配置以避免不合理触发。

结论：PR 已合并，但问题未完全解决，计划在后续 PR 中优化 CI 工作流。 · 待处理

风险与影响

技术风险包括：1) NPU特定优化（如torch_npu.npu_rotary_mul）可能引入GPU兼容性问题，依赖条件判断准确性；2) 临时使用complex64处理complex128不兼容可能影响数值精度，需未来移除；3) 添加的平台检测代码增加维护负担，注释中TODO表示临时修复；4) 性能优化依赖外部NPU库版本，可能不稳定；5) CI工作流因修改jit_kernel文件意外触发LLM测试，需后续优化。

影响范围：1) 用户：MOVA模型现支持NPU硬件，扩展了部署选项和硬件生态；2) 系统：Diffusion模块增加NPU支持代码，提升硬件覆盖率，但引入平台特定逻辑，可能增加复杂度；3) 团队：推进NPU路线图，为后续NPU优化奠定基础，促进跨团队协作。影响程度中等，主要限于MOVA模型和NPU平台。

平台特定代码 临时修复依赖 性能优化不确定性 CI 工作流影响

关联 Issue

#18967 [Roadmap] [NPU] Sglang Diffusion on Ascend

完整报告

Markdown PDF

执行摘要

  • 一句话：为MOVA扩散模型添加NPU硬件支持，优化RoPE和内存格式兼容性。
  • 推荐动作：建议涉及NPU开发或Diffusion模型的工程师精读此PR，关注平台检测设计（如current_platform.is_npu）、RoPE优化策略（torch_npu.npu_rotary_mul集成）以及如何处理硬件不兼容性（如channels_last_3d和complex128临时修复）。这些设计决策值得学习，以应对类似跨平台支持场景。

功能与动机

根据PR body，动机是'This PR adds NPU support to the MOVA model.'，并引用Issue #18967 '[Roadmap] [NPU] Sglang Diffusion on Ascend'，作为NPU支持路线图的一部分，旨在扩展硬件覆盖并解决现有NPU平台上的执行失败问题。

实现拆解

实现分三部分：1) 在python/sglang/jit_kernel/diffusion/triton/npu_fallback.py中添加NPU特定的Rotary Embedding优化，引入torch_npu.npu_rotary_mul以提高性能，并设置常量限制；2) 在python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py中修改patchify函数，根据current_platform.is_npu选择contiguous方式，避免NPU不支持的torch.channels_last_3d；3) 在python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py中调整visual_freqs和audio_freqs准备逻辑，处理NPU上torch.cat对torch.complex128的不兼容，临时使用torch.complex64。

关键文件：

  • python/sglang/jit_kernel/diffusion/triton/npu_fallback.py（模块 JIT-kernel/Diffusion）: 核心优化点，添加NPU-specific Rotary Embedding fallback使用torch_npu.npu_rotary_mul，提高性能并设置常量限制。
  • python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py（模块 Diffusion/Models）: 处理NPU上channels_last_3d不兼容，修改patchify函数以避免性能下降，影响模型前向逻辑。
  • python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py（模块 Diffusion/Pipelines）: 调整频率准备逻辑以处理NPU上complex128不兼容，关键推理步骤，涉及visual和audio频率张量处理。

关键符号：apply_rotary_embedding_native, patchify, inference_single_step

评论区精华

review中核心讨论点：1) gemini-code-assist[bot]建议提取共享函数以减少visual_freqs和audio_freqs准备的冗余代码，但未直接采纳；2) ping1jing2和Makcum888e建议减少代码中的is_npu检查，推动改动到npu_fallback.py中，作者LLThomas响应并修改相关代码；3) Makcum888e要求添加注释说明临时修复（如'remove it when torch.cat on NPU is supported torch.complex128'），作者已添加；4) ping1jing2关注CI触发问题，指出Diffusion PR意外触发LLM CI，合并后计划优化CI工作流。讨论结论是问题已解决，代码改进并合并。

  • 代码冗余和共享函数提取 (design): 代码已修改，但未实现共享函数，保持原有逻辑。
  • 减少is_npu检查并集中优化 (design): 作者响应并修改了相关代码，推动优化到统一位置。
  • 添加注释和代码一致性 (documentation): 作者添加了所需注释，提高了代码清晰度。
  • CI工作流触发问题 (infra): PR已合并，但问题未完全解决，计划在后续PR中优化CI工作流。

风险与影响

  • 风险：技术风险包括：1) NPU特定优化（如torch_npu.npu_rotary_mul）可能引入GPU兼容性问题，依赖条件判断准确性；2) 临时使用complex64处理complex128不兼容可能影响数值精度，需未来移除；3) 添加的平台检测代码增加维护负担，注释中TODO表示临时修复；4) 性能优化依赖外部NPU库版本，可能不稳定；5) CI工作流因修改jit_kernel文件意外触发LLM测试，需后续优化。
  • 影响：影响范围：1) 用户：MOVA模型现支持NPU硬件，扩展了部署选项和硬件生态；2) 系统：Diffusion模块增加NPU支持代码，提升硬件覆盖率，但引入平台特定逻辑，可能增加复杂度；3) 团队：推进NPU路线图，为后续NPU优化奠定基础，促进跨团队协作。影响程度中等，主要限于MOVA模型和NPU平台。
  • 风险标记：平台特定代码, 临时修复依赖, 性能优化不确定性, CI工作流影响

关联脉络

  • PR #21408 [NPU] Support GLM-4.7-Flash on NPU: 同为NPU支持功能，涉及硬件后端和注意力适配，共享NPU优化策略。
  • PR #21955 [diffusion] chore: fix stage profiler for multi-stage denoising: 同属Diffusion模块的bugfix，相关性能分析和阶段计时，可能影响整体性能评估。
  • PR #21922 Revert "Rollback flashmla to older version [1/2]": 涉及JIT-kernel和CUDA内核支持，与本PR的jit-kernel优化有技术关联。

参与讨论