feat(npu): add MindIE-SD compile fusion and MINDIE attention backend by Chitandaaaaa · Pull Request #259 · modelscope/DiffSynth-Engine

Chitandaaaaa · 2026-07-13T04:18:00Z

Summary

This PR adds NPU support on top of the v1 attention registry architecture.

Platform detection: Add is_npu_available() and is_mindie_sd_available() in utils/platform.py, and route device selection / distributed backend to NPU (npu) and HCCL when available.
MINDIE attention backend: Register a new mindie backend in registry.py and implement mindie_attn.py, which delegates attention computation to MindIE-SD's fused flash attention kernel.
Auto attention selection: When --attn-type is not specified, automatically choose mindie on NPU and sdpa otherwise.
MindIE-SD compile fusion: Add platform/npu/compilation.py and apply torch.compile(backend=MindieSDBackend()) in DiffusionModel.from_pretrained when MindIE-SD is available, enabling RMSNorm / RoPE / AdaLayerNorm / FastGELU / MulAdd fusion on NPU.

Made with Cursor

gemini-code-assist

Code Review

This pull request introduces NPU support and integrates the MindIE-SD attention backend, enabling optimized attention execution and compilation fusion patterns on NPU devices. It also adds a multi-block profiling mechanism for the Qwen Image transformer. The feedback highlights critical improvements for robustness and flexibility: replacing NPU availability checks with specific MindIE-SD library checks to prevent runtime crashes when the library is missing, addressing thread-safety issues with global profiling variables, and providing an environment-variable-based opt-out mechanism for forced model compilation.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-13T04:19:46Z

+from diffsynth_engine.utils.platform import is_npu_available
+
+
+class MindieAttentionBackend(AttentionBackend):
+    @staticmethod
+    def check_availability() -> None:
+        if not is_npu_available():
+            raise RuntimeError("NPU is not available, cannot use MINDIE attention backend")


The check_availability method currently checks is_npu_available(), but the backend relies on the mindiesd library which is imported dynamically during the forward pass. If torch_npu is available but mindiesd is not installed, check_availability will pass, but the forward pass will crash with an ImportError. It is safer to check is_mindie_sd_available() instead.

Suggested change

from diffsynth_engine.utils.platform import is_npu_available

class MindieAttentionBackend(AttentionBackend):

@staticmethod

def check_availability() -> None:

if not is_npu_available():

raise RuntimeError("NPU is not available, cannot use MINDIE attention backend")

from diffsynth_engine.utils.platform import is_mindie_sd_available

class MindieAttentionBackend(AttentionBackend):

@staticmethod

def check_availability() -> None:

if not is_mindie_sd_available():

raise RuntimeError("MindIE-SD is not available, cannot use MINDIE attention backend")

确实应该判断mindie才对

gemini-code-assist · 2026-07-13T04:19:46Z


    if attn_type is None:
-        attn_type = "sdpa"
+        attn_type = "mindie" if is_npu_available() else "sdpa"


When attn_type is None (auto-detect), defaulting to "mindie" solely based on is_npu_available() will cause a runtime crash if the mindiesd library is not installed on the NPU system. It is safer to default to "mindie" only if is_mindie_sd_available() is True, falling back to "sdpa" otherwise.

Suggested change

attn_type = "mindie" if is_npu_available() else "sdpa"

attn_type = "mindie" if is_mindie_sd_available() else "sdpa"

gemini-code-assist · 2026-07-13T04:19:46Z

 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.constants import MODEL_INDEX_NAME
 from diffsynth_engine.utils.import_utils import LazyImport
+from diffsynth_engine.utils.platform import is_npu_available


Import is_mindie_sd_available to safely check if the MindIE-SD library is installed before defaulting to the mindie attention backend.

Suggested change

from diffsynth_engine.utils.platform import is_npu_available

from diffsynth_engine.utils.platform import is_mindie_sd_available, is_npu_available

gemini-code-assist · 2026-07-13T04:19:46Z

+_PROFILE_BLOCK_COUNT = 4
+_blocks_profiled = 0
+_profile_limit_reached = False
+_multi_block_profiler = None


Using global variables (_blocks_profiled, _profile_limit_reached, _multi_block_profiler) to track profiling state during model execution is not thread-safe. In multi-threaded environments (e.g., concurrent inference requests in a web server), this can lead to race conditions, corrupted profiling metrics, or unexpected crashes. Consider using thread-local storage (threading.local()) or passing the profiling state through the context/call stack to ensure thread safety.

gemini-code-assist · 2026-07-13T04:19:46Z

+import torch
+
+from diffsynth_engine.utils.platform import is_mindie_sd_available
+
+
+def apply_mindie_sd_compile(model: torch.nn.Module) -> torch.nn.Module:
+    """Apply MindIE-SD torch.compile backend fusion patterns.
+
+    Enables RMSNorm, RoPE, AdaLayerNorm, FastGELU, and MulAdd op fusion
+    for NPU execution. On non-NPU devices or when mindiesd is unavailable,
+    returns the model unchanged.
+    """
+    if not is_mindie_sd_available():
+        return model
+
+    from mindiesd.compilation import MindieSDBackend
+
+    return torch.compile(model, backend=MindieSDBackend())


Forcing torch.compile via apply_mindie_sd_compile whenever is_mindie_sd_available() is True leaves no option for the user to run in eager mode. This can be problematic for debugging, avoiding compilation overhead, or handling models that are incompatible with compilation. Consider adding an opt-out mechanism, such as checking an environment variable (e.g., DIFFSYNTH_DISABLE_COMPILE).

Suggested change

import torch

from diffsynth_engine.utils.platform import is_mindie_sd_available

def apply_mindie_sd_compile(model: torch.nn.Module) -> torch.nn.Module:

"""Apply MindIE-SD torch.compile backend fusion patterns.

Enables RMSNorm, RoPE, AdaLayerNorm, FastGELU, and MulAdd op fusion

for NPU execution. On non-NPU devices or when mindiesd is unavailable,

returns the model unchanged.

"""

if not is_mindie_sd_available():

return model

from mindiesd.compilation import MindieSDBackend

return torch.compile(model, backend=MindieSDBackend())

import os

import torch

from diffsynth_engine.utils.platform import is_mindie_sd_available

def apply_mindie_sd_compile(model: torch.nn.Module) -> torch.nn.Module:

"""Apply MindIE-SD torch.compile backend fusion patterns.

Enables RMSNorm, RoPE, AdaLayerNorm, FastGELU, and MulAdd op fusion

for NPU execution. On non-NPU devices or when mindiesd is unavailable,

returns the model unchanged.

"""

if not is_mindie_sd_available() or os.environ.get("DIFFSYNTH_DISABLE_COMPILE", "0") == "1":

return model

from mindiesd.compilation import MindieSDBackend

return torch.compile(model, backend=MindieSDBackend())

Enable torch.compile with MindieSDBackend on model load for NPU, register the mindie attention backend in the v1 registry, and auto-detect SDPA vs MINDIE based on platform availability. Co-authored-by: Cursor <cursoragent@cursor.com>

Pipeline now loads transformers via from_config + load_state_dict, so the compile hook on DiffusionModel.from_pretrained never runs. Re-apply it on the shared init path so fusion kernels show up in profiling again. Co-authored-by: Cursor <cursoragent@cursor.com>

akaitsuki-ii · 2026-07-17T03:09:48Z

+    try:
+        import torch_npu
+
+        return torch_npu.npu.is_available()


建议改成类似FLASH_ATTN_3_AVAILABLE的环境变量而不是每次try import

akaitsuki-ii · 2026-07-17T03:14:36Z

+        attn_metadata: AttentionMetadata | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        from mindiesd.layers.flash_attn.attention_forward import attention_forward


import统一放开头清晰一点

akaitsuki-ii · 2026-07-17T03:15:27Z

        model.load_state_dict(state_dict, strict=True, assign=True)
        model.to(device=device)
-        return model
+


实际初始化不走这个逻辑就删掉吧

akaitsuki-ii · 2026-07-17T03:18:57Z

+
+    from mindiesd.compilation import MindieSDBackend
+
+    return torch.compile(model, backend=MindieSDBackend())


这里整体torch.compile不会有graph break问题吗？一般实践中我们只compile repeated blocks

akaitsuki-ii · 2026-07-17T03:20:51Z

@@ -1 +1 @@
 import torch


platform已经是独立目录的话，考虑把这些平台相关的内容完整迁移过去？

akaitsuki-ii · 2026-07-17T03:30:28Z


    # attention
-    attn_type: AttentionType | str = AttentionType.SDPA
+    attn_type: AttentionType | str | None = None  # None = auto-detect


直接默认是"auto"是不是更合适点🤔

Co-authored-by: Cursor <cursoragent@cursor.com>

Pipeline loads transformers via init_transformer, so apply compile only there. Also clear unused platform package re-export. Co-authored-by: Cursor <cursoragent@cursor.com>

gemini-code-assist Bot reviewed Jul 13, 2026

View reviewed changes

Chitandaaaaa force-pushed the feat/npu-compile-rebased branch 2 times, most recently from 6583fa7 to d2ad4d8 Compare July 15, 2026 03:32

akaitsuki-ii reviewed Jul 17, 2026

View reviewed changes

fxp19950422 mentioned this pull request Jul 20, 2026

feat: add extensible Ascend A5 / Ascend 950 runtime support #260

Open

hammer and others added 2 commits July 21, 2026 16:24

fix(npu): disable MindIE-SD fast GELU fusion in compile backend

0b6dff2

Co-authored-by: Cursor <cursoragent@cursor.com>

refactor(npu): drop redundant MindIE compile from DiffusionModel

49ade03

Pipeline loads transformers via init_transformer, so apply compile only there. Also clear unused platform package re-export. Co-authored-by: Cursor <cursoragent@cursor.com>

	attn_type = "mindie" if is_npu_available() else "sdpa"
	attn_type = "mindie" if is_mindie_sd_available() else "sdpa"

	from diffsynth_engine.utils.platform import is_npu_available
	from diffsynth_engine.utils.platform import is_mindie_sd_available, is_npu_available


		from mindiesd.compilation import MindieSDBackend

		return torch.compile(model, backend=MindieSDBackend())

Uh oh!

Conversation

Chitandaaaaa commented Jul 13, 2026

Summary

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

akaitsuki-ii Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

akaitsuki-ii Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

akaitsuki-ii Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

akaitsuki-ii Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

akaitsuki-ii Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

akaitsuki-ii Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

akaitsuki-ii Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants