feat(npu): add MindIE-SD compile fusion and MINDIE attention backend

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)
 

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)
 

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

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,
+        )

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(

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",
+]

diffsynth_engine/platform/npu/compilation.py

@@ -0,0 +1,19 @@
+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, 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 CompilationConfig, MindieSDBackend
+
+    CompilationConfig.fusion_patterns.enable_fast_gelu = False
+    return torch.compile(model, backend=MindieSDBackend())

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}")

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():