From d2ad4d8141d4ba6b7552e7e530f51c624a0ffae0 Mon Sep 17 00:00:00 2001 From: hammer Date: Mon, 13 Jul 2026 12:14:28 +0800 Subject: [PATCH 1/6] feat(npu): add MindIE-SD compile fusion and MINDIE attention backend 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 --- diffsynth_engine/args.py | 8 +-- diffsynth_engine/configs/base.py | 5 +- .../layers/attention/backends/abstract.py | 1 + .../layers/attention/backends/mindie_attn.py | 62 +++++++++++++++++++ diffsynth_engine/models/base.py | 5 +- diffsynth_engine/platform/__init__.py | 3 + diffsynth_engine/platform/npu/__init__.py | 8 +++ diffsynth_engine/platform/npu/compilation.py | 18 ++++++ diffsynth_engine/registry.py | 4 +- diffsynth_engine/utils/platform.py | 24 +++++++ 10 files changed, 130 insertions(+), 8 deletions(-) create mode 100644 diffsynth_engine/layers/attention/backends/mindie_attn.py create mode 100644 diffsynth_engine/platform/__init__.py create mode 100644 diffsynth_engine/platform/npu/__init__.py create mode 100644 diffsynth_engine/platform/npu/compilation.py diff --git a/diffsynth_engine/args.py b/diffsynth_engine/args.py index d73a030a..dda22915 100644 --- a/diffsynth_engine/args.py +++ b/diffsynth_engine/args.py @@ -18,7 +18,7 @@ def _parse_tuple(value: str) -> Tuple[int, int] | int: def _parse_attention_params( - attn_type: str, + attn_type: str | None, sparge_topk: float | None = None, ) -> AttentionParams | None: """Parse attention parameters based on attention type""" @@ -100,8 +100,8 @@ def parse_cli_args() -> Dict[str, Any]: attn_group.add_argument( "--attn-type", type=str, - default="sdpa", - help="Attention type (default: sdpa)", + default=None, + help="Attention type (default: auto, SDPA on GPU, mindie on NPU)", ) attn_group.add_argument( "--sparge-topk", @@ -171,7 +171,7 @@ def parse_cli_args() -> Dict[str, Any]: args_dict["vae_tile_stride"] = _parse_tuple(args.vae_tile_stride) # Attention configuration - attn_type = args.attn_type.lower() + attn_type = args.attn_type.lower() if args.attn_type is not None else None args_dict["attn_type"] = attn_type args_dict["attn_params"] = _parse_attention_params(attn_type, args.sparge_topk) diff --git a/diffsynth_engine/configs/base.py b/diffsynth_engine/configs/base.py index b44c472d..a3e032c5 100644 --- a/diffsynth_engine/configs/base.py +++ b/diffsynth_engine/configs/base.py @@ -37,7 +37,7 @@ class PipelineConfig: vae_tile_stride: int | Tuple[int, int] = (192, 192) # attention - attn_type: AttentionType | str = AttentionType.SDPA + attn_type: AttentionType | str | None = None # None = auto-detect attn_params: Optional[AttentionParams] = None # parallelism @@ -56,7 +56,8 @@ def from_dict(cls, args_dict: Dict[str, Any]) -> "PipelineConfig": return cls(**filtered_dict) def __post_init__(self): - self.attn_type = str(self.attn_type) + if self.attn_type is not None: + self.attn_type = str(self.attn_type) init_parallel_config(self) validate_attn_config(self) diff --git a/diffsynth_engine/layers/attention/backends/abstract.py b/diffsynth_engine/layers/attention/backends/abstract.py index e9b5fdd0..35be7453 100644 --- a/diffsynth_engine/layers/attention/backends/abstract.py +++ b/diffsynth_engine/layers/attention/backends/abstract.py @@ -26,6 +26,7 @@ class AttentionType(str, enum.Enum): SAGE2 = "sage2" SAGE3 = "sage3" SPARGE = "sparge" + MINDIE = "mindie" def __str__(self) -> str: return self.value diff --git a/diffsynth_engine/layers/attention/backends/mindie_attn.py b/diffsynth_engine/layers/attention/backends/mindie_attn.py new file mode 100644 index 00000000..e0c847c1 --- /dev/null +++ b/diffsynth_engine/layers/attention/backends/mindie_attn.py @@ -0,0 +1,62 @@ +import torch + +from diffsynth_engine.layers.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionType, +) +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") + + @staticmethod + def get_type() -> str: + return str(AttentionType.MINDIE) + + @staticmethod + def get_impl_cls() -> type["AttentionImpl"]: + return MindieAttentionImpl + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [] + + +class MindieAttentionImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float | None = None, + causal: bool = False, + num_kv_heads: int | None = None, + **extra_impl_args, + ) -> None: + self.scale = softmax_scale or (head_size**-0.5) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + attn_metadata: AttentionMetadata | None = None, + **kwargs, + ) -> torch.Tensor: + from mindiesd.layers.flash_attn.attention_forward import attention_forward + + return attention_forward( + query=query, + key=key, + value=value, + attn_mask=attn_mask, + scale=self.scale, + fused=True, + head_first=False, + ) diff --git a/diffsynth_engine/models/base.py b/diffsynth_engine/models/base.py index 20768ee1..d9b59d0c 100644 --- a/diffsynth_engine/models/base.py +++ b/diffsynth_engine/models/base.py @@ -64,7 +64,10 @@ def from_pretrained( model.load_state_dict(state_dict, strict=True, assign=True) model.to(device=device) - return model + + from diffsynth_engine.platform.npu import apply_mindie_sd_compile + + return apply_mindie_sd_compile(model) class AutoregressiveModel(nn.Module): diff --git a/diffsynth_engine/platform/__init__.py b/diffsynth_engine/platform/__init__.py new file mode 100644 index 00000000..4c9f275f --- /dev/null +++ b/diffsynth_engine/platform/__init__.py @@ -0,0 +1,3 @@ +from . import npu + +__all__ = ["npu"] diff --git a/diffsynth_engine/platform/npu/__init__.py b/diffsynth_engine/platform/npu/__init__.py new file mode 100644 index 00000000..c33ceef0 --- /dev/null +++ b/diffsynth_engine/platform/npu/__init__.py @@ -0,0 +1,8 @@ +from diffsynth_engine.utils.platform import is_mindie_sd_available + +from .compilation import apply_mindie_sd_compile + +__all__ = [ + "apply_mindie_sd_compile", + "is_mindie_sd_available", +] diff --git a/diffsynth_engine/platform/npu/compilation.py b/diffsynth_engine/platform/npu/compilation.py new file mode 100644 index 00000000..32f444d9 --- /dev/null +++ b/diffsynth_engine/platform/npu/compilation.py @@ -0,0 +1,18 @@ +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()) diff --git a/diffsynth_engine/registry.py b/diffsynth_engine/registry.py index b6b9f4e5..8b735a74 100644 --- a/diffsynth_engine/registry.py +++ b/diffsynth_engine/registry.py @@ -7,6 +7,7 @@ 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 if TYPE_CHECKING: from diffsynth_engine.layers.attention.backends.abstract import AttentionBackend @@ -37,6 +38,7 @@ "sage3": "diffsynth_engine.layers.attention.backends.sage_attn_3:SageAttention3Backend", "sdpa": "diffsynth_engine.layers.attention.backends.sdpa:SDPABackend", "sparge": "diffsynth_engine.layers.attention.backends.sparge_attn:SpargeAttentionBackend", + "mindie": "diffsynth_engine.layers.attention.backends.mindie_attn:MindieAttentionBackend", } PIPELINE_REGISTRY: dict[str, LazyImport] = {} @@ -118,7 +120,7 @@ def get_attn_backend(attn_type: str | None = None) -> type["AttentionBackend"]: _attention_backend_registry_initialized = True if attn_type is None: - attn_type = "sdpa" + attn_type = "mindie" if is_npu_available() else "sdpa" if attn_type not in ATTENTION_BACKEND_REGISTRY: available_backends = sorted(ATTENTION_BACKEND_REGISTRY) raise ValueError(f"Attention backend {attn_type!r} not found. Available backends: {available_backends}") diff --git a/diffsynth_engine/utils/platform.py b/diffsynth_engine/utils/platform.py index 575f5738..4e11c394 100644 --- a/diffsynth_engine/utils/platform.py +++ b/diffsynth_engine/utils/platform.py @@ -13,7 +13,27 @@ def _is_mps() -> bool: return torch.backends.mps.is_available() +def is_npu_available() -> bool: + try: + import torch_npu + + return torch_npu.npu.is_available() + except ImportError: + return False + + +def is_mindie_sd_available() -> bool: + try: + import mindiesd # noqa: F401 + + return is_npu_available() + except ImportError: + return False + + def get_device(local_rank: int) -> torch.device: + if is_npu_available(): + return torch.device("npu", local_rank) if _is_cuda() or _is_rocm(): return torch.device("cuda", local_rank) if _is_mps(): @@ -23,6 +43,8 @@ def get_device(local_rank: int) -> torch.device: def get_device_type() -> str: + if is_npu_available(): + return "npu" if _is_cuda() or _is_rocm(): return "cuda" if _is_mps(): @@ -32,6 +54,8 @@ def get_device_type() -> str: def get_torch_distributed_backend() -> str: + if is_npu_available(): + return "hccl" if _is_cuda() or _is_rocm(): return "nccl" if _is_mps(): From e94848c786ae73d945de5f527c957893aef678ce Mon Sep 17 00:00:00 2001 From: hammer Date: Wed, 15 Jul 2026 16:16:10 +0800 Subject: [PATCH 2/6] fix(npu): apply MindIE-SD compile in Pipeline.init_transformer 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 --- diffsynth_engine/pipelines/base.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/diffsynth_engine/pipelines/base.py b/diffsynth_engine/pipelines/base.py index 601e8957..5848ddd8 100644 --- a/diffsynth_engine/pipelines/base.py +++ b/diffsynth_engine/pipelines/base.py @@ -96,7 +96,12 @@ def init_transformer( model.to(device=pipeline_config.device) del state_dict - return model + + # Pipeline loads via from_config + load_state_dict, so compile must be + # applied here (DiffusionModel.from_pretrained is no longer on this path). + from diffsynth_engine.platform.npu import apply_mindie_sd_compile + + return apply_mindie_sd_compile(model) @staticmethod def init_text_encoder( From 0b6dff2a667e70b8fc9d627b655596373f18836b Mon Sep 17 00:00:00 2001 From: hammer Date: Tue, 21 Jul 2026 15:41:59 +0800 Subject: [PATCH 3/6] fix(npu): disable MindIE-SD fast GELU fusion in compile backend Co-authored-by: Cursor --- diffsynth_engine/platform/npu/compilation.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/diffsynth_engine/platform/npu/compilation.py b/diffsynth_engine/platform/npu/compilation.py index 32f444d9..4cb4a661 100644 --- a/diffsynth_engine/platform/npu/compilation.py +++ b/diffsynth_engine/platform/npu/compilation.py @@ -6,13 +6,14 @@ 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. + Enables RMSNorm, RoPE, AdaLayerNorm, and MulAdd op fusion for NPU execution. + FastGELU fusion is disabled (npu_fast_gelu was slower than native GELU in practice). + 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 + from mindiesd.compilation import CompilationConfig, MindieSDBackend + CompilationConfig.fusion_patterns.enable_fast_gelu = False return torch.compile(model, backend=MindieSDBackend()) From 49ade032817466fe1f3a8f1a06e897bc00206269 Mon Sep 17 00:00:00 2001 From: hammer Date: Wed, 22 Jul 2026 11:43:51 +0800 Subject: [PATCH 4/6] refactor(npu): drop redundant MindIE compile from DiffusionModel Pipeline loads transformers via init_transformer, so apply compile only there. Also clear unused platform package re-export. Co-authored-by: Cursor --- diffsynth_engine/models/base.py | 5 +---- diffsynth_engine/platform/__init__.py | 3 --- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/diffsynth_engine/models/base.py b/diffsynth_engine/models/base.py index d9b59d0c..20768ee1 100644 --- a/diffsynth_engine/models/base.py +++ b/diffsynth_engine/models/base.py @@ -64,10 +64,7 @@ def from_pretrained( model.load_state_dict(state_dict, strict=True, assign=True) model.to(device=device) - - from diffsynth_engine.platform.npu import apply_mindie_sd_compile - - return apply_mindie_sd_compile(model) + return model class AutoregressiveModel(nn.Module): diff --git a/diffsynth_engine/platform/__init__.py b/diffsynth_engine/platform/__init__.py index 4c9f275f..e69de29b 100644 --- a/diffsynth_engine/platform/__init__.py +++ b/diffsynth_engine/platform/__init__.py @@ -1,3 +0,0 @@ -from . import npu - -__all__ = ["npu"] From 30aa278a29250f33219e31e28ffd6750bbf25576 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 24 Jul 2026 17:47:30 +0800 Subject: [PATCH 5/6] fix(npu): auto-detect device and attention backend Auto-detect device (npu/cuda/mps/cpu) and attention backend (mindie when NPU+mindiesd available, sdpa otherwise). Refactor mindie_attn.py to use module-level import and combined NPU+mindiesd availability flag. --- diffsynth_engine/configs/base.py | 14 +++++++++----- .../layers/attention/backends/mindie_attn.py | 17 +++++++++++++---- diffsynth_engine/registry.py | 4 ++-- diffsynth_engine/utils/platform.py | 5 +++-- 4 files changed, 27 insertions(+), 13 deletions(-) diff --git a/diffsynth_engine/configs/base.py b/diffsynth_engine/configs/base.py index a3e032c5..4d2adfe1 100644 --- a/diffsynth_engine/configs/base.py +++ b/diffsynth_engine/configs/base.py @@ -6,6 +6,7 @@ from diffsynth_engine.layers.attention import AttentionType from diffsynth_engine.registry import get_attn_backend from diffsynth_engine.utils import logging +from diffsynth_engine.utils.platform import get_device_type, is_mindie_sd_available logger = logging.get_logger(__name__) @@ -27,8 +28,7 @@ class PipelineConfig: model_dtype: torch.dtype = torch.bfloat16 text_encoder_dtype: torch.dtype = torch.bfloat16 vae_dtype: torch.dtype = torch.float32 - device: str | torch.device = "cuda" - + device: str | torch.device = "auto" pipeline_class_name: str | None = None # vae @@ -37,7 +37,7 @@ class PipelineConfig: vae_tile_stride: int | Tuple[int, int] = (192, 192) # attention - attn_type: AttentionType | str | None = None # None = auto-detect + attn_type: AttentionType | str = "auto" attn_params: Optional[AttentionParams] = None # parallelism @@ -56,8 +56,12 @@ def from_dict(cls, args_dict: Dict[str, Any]) -> "PipelineConfig": return cls(**filtered_dict) def __post_init__(self): - if self.attn_type is not None: - self.attn_type = str(self.attn_type) + if self.device == "auto": + self.device = get_device_type() + if self.attn_type == "auto": + self.attn_type = "mindie" if is_mindie_sd_available() else "sdpa" + + self.attn_type = str(self.attn_type) init_parallel_config(self) validate_attn_config(self) diff --git a/diffsynth_engine/layers/attention/backends/mindie_attn.py b/diffsynth_engine/layers/attention/backends/mindie_attn.py index e0c847c1..45f85475 100644 --- a/diffsynth_engine/layers/attention/backends/mindie_attn.py +++ b/diffsynth_engine/layers/attention/backends/mindie_attn.py @@ -6,14 +6,25 @@ AttentionMetadata, AttentionType, ) +from diffsynth_engine.utils import logging from diffsynth_engine.utils.platform import is_npu_available +logger = logging.get_logger(__name__) + +try: + from mindiesd.layers.flash_attn.attention_forward import attention_forward + + MINDIESD_ATTN_AVAILABLE = is_npu_available() +except ImportError: + MINDIESD_ATTN_AVAILABLE = False class MindieAttentionBackend(AttentionBackend): @staticmethod def check_availability() -> None: - if not is_npu_available(): - raise RuntimeError("NPU is not available, cannot use MINDIE attention backend") + if not MINDIESD_ATTN_AVAILABLE: + error_msg = "MindiesdAttention backend is not available. Please visit https://gitcode.com/Ascend/MindIE-SD/blob/dev/docs/zh/features/core_layers.md and follow the installation instructions." + logger.error(error_msg) + raise RuntimeError(error_msg) @staticmethod def get_type() -> str: @@ -49,8 +60,6 @@ def forward( attn_metadata: AttentionMetadata | None = None, **kwargs, ) -> torch.Tensor: - from mindiesd.layers.flash_attn.attention_forward import attention_forward - return attention_forward( query=query, key=key, diff --git a/diffsynth_engine/registry.py b/diffsynth_engine/registry.py index 8b735a74..ebde9235 100644 --- a/diffsynth_engine/registry.py +++ b/diffsynth_engine/registry.py @@ -7,7 +7,7 @@ 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 +from diffsynth_engine.utils.platform import is_mindie_sd_available if TYPE_CHECKING: from diffsynth_engine.layers.attention.backends.abstract import AttentionBackend @@ -120,7 +120,7 @@ def get_attn_backend(attn_type: str | None = None) -> type["AttentionBackend"]: _attention_backend_registry_initialized = True if attn_type is None: - attn_type = "mindie" if is_npu_available() else "sdpa" + attn_type = "mindie" if is_mindie_sd_available() else "sdpa" if attn_type not in ATTENTION_BACKEND_REGISTRY: available_backends = sorted(ATTENTION_BACKEND_REGISTRY) raise ValueError(f"Attention backend {attn_type!r} not found. Available backends: {available_backends}") diff --git a/diffsynth_engine/utils/platform.py b/diffsynth_engine/utils/platform.py index 4e11c394..2208a461 100644 --- a/diffsynth_engine/utils/platform.py +++ b/diffsynth_engine/utils/platform.py @@ -1,4 +1,5 @@ import torch +from functools import cache def _is_cuda() -> bool: @@ -12,7 +13,7 @@ def _is_rocm() -> bool: def _is_mps() -> bool: return torch.backends.mps.is_available() - +@cache def is_npu_available() -> bool: try: import torch_npu @@ -21,7 +22,7 @@ def is_npu_available() -> bool: except ImportError: return False - +@cache def is_mindie_sd_available() -> bool: try: import mindiesd # noqa: F401 From 7553f937b94123cce12a33a771421130581223de Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 01:09:14 +0800 Subject: [PATCH 6/6] fix(npu): compile repeated blocks individually to avoid graph break Switch from full-model to per-block compile via _repeated_blocks. Add DIFFSYNTH_DISABLE_COMPILE=1 env var to disable compile. --- diffsynth_engine/platform/npu/compilation.py | 30 +++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/diffsynth_engine/platform/npu/compilation.py b/diffsynth_engine/platform/npu/compilation.py index 4cb4a661..7ce21874 100644 --- a/diffsynth_engine/platform/npu/compilation.py +++ b/diffsynth_engine/platform/npu/compilation.py @@ -1,3 +1,4 @@ +import os import torch from diffsynth_engine.utils.platform import is_mindie_sd_available @@ -9,11 +10,38 @@ def apply_mindie_sd_compile(model: torch.nn.Module) -> torch.nn.Module: Enables RMSNorm, RoPE, AdaLayerNorm, and MulAdd op fusion for NPU execution. FastGELU fusion is disabled (npu_fast_gelu was slower than native GELU in practice). On non-NPU devices or when mindiesd is unavailable, returns the model unchanged. + + Compilation strategy: + If the model defines ``_repeated_blocks`` (e.g. ``["QwenImageTransformerBlock"]``), + compiles only the repeated submodules individually to reduce compilation time + and graph break risk. Otherwise falls back to full-model ``torch.compile``. + Set ``DIFFSYNTH_DISABLE_COMPILE=1`` to disable. """ - if not is_mindie_sd_available(): + if not is_mindie_sd_available() or os.environ.get("DIFFSYNTH_DISABLE_COMPILE", "0") == "1": return model from mindiesd.compilation import CompilationConfig, MindieSDBackend CompilationConfig.fusion_patterns.enable_fast_gelu = False + + repeated_blocks = getattr(model, "_repeated_blocks", None) + if repeated_blocks: + paths = [ + name + for name, submodule in model.named_modules() + if submodule.__class__.__name__ in repeated_blocks + ] + # Sort by depth in descending order (submodules before parent modules) + paths.sort(key=lambda p: len(p.split(".")), reverse=True) + + backend = MindieSDBackend() + for name in paths: + parent_name, _, child_name = name.rpartition(".") + parent = model.get_submodule(parent_name) if parent_name else model + submodule = getattr(parent, child_name) + setattr(parent, child_name, torch.compile(submodule, backend=backend)) + + return model + + # Fall back to full-model compilation for models without `_repeated_blocks` defined. return torch.compile(model, backend=MindieSDBackend())