From acffa4f871aa51a7a52571bd0c1e6badc62c4bda Mon Sep 17 00:00:00 2001 From: root Date: Tue, 28 Jul 2026 11:36:15 +0800 Subject: [PATCH] feat(npu): add Ascend 950 platform abstraction and Qwen-Image support 1. Add NPU 950 platform abstraction layer with device registry and auto_detect_device 2. Support mindiesd_attention and mindiesd_compile with capability detection and compile_backend 3. Adapt Qwen-Image model to use platform compile kwargs and NPU attention backend --- diffsynth_engine/configs/pipeline.py | 1 + diffsynth_engine/models/basic/attention.py | 19 ++ diffsynth_engine/pipelines/qwen_image.py | 4 +- diffsynth_engine/platforms/__init__.py | 148 +++++++++++++ diffsynth_engine/platforms/ascend.py | 228 +++++++++++++++++++++ diffsynth_engine/platforms/base.py | 68 ++++++ diffsynth_engine/utils/flag.py | 5 + diffsynth_engine/utils/platform.py | 49 +++-- 8 files changed, 505 insertions(+), 17 deletions(-) create mode 100644 diffsynth_engine/platforms/__init__.py create mode 100644 diffsynth_engine/platforms/ascend.py create mode 100644 diffsynth_engine/platforms/base.py diff --git a/diffsynth_engine/configs/pipeline.py b/diffsynth_engine/configs/pipeline.py index 37da60a1..5ea5f7f4 100644 --- a/diffsynth_engine/configs/pipeline.py +++ b/diffsynth_engine/configs/pipeline.py @@ -35,6 +35,7 @@ class AttnImpl(Enum): SAGE = "sage" # Sage Attention SPARGE = "sparge" # Sparge Attention VSA = "vsa" # Video Sparse Attention + MINDIE = "mindie" # Mindie Attention @dataclass diff --git a/diffsynth_engine/models/basic/attention.py b/diffsynth_engine/models/basic/attention.py index 3d3d49cf..d472b5fd 100644 --- a/diffsynth_engine/models/basic/attention.py +++ b/diffsynth_engine/models/basic/attention.py @@ -15,6 +15,7 @@ SPARGE_ATTN_AVAILABLE, VIDEO_SPARSE_ATTN_AVAILABLE, AITER_AVAILABLE, + MINDIE_AVAILABLE, ) from diffsynth_engine.utils.platform import DTYPE_FP8 @@ -107,6 +108,16 @@ def sparge_attn( distributed_video_sparse_attn, ) +if MINDIE_AVAILABLE: + from mindiesd.layers.flash_attn.attention_forward import attention_forward + + def mindie_attn(q, k, v, attn_mask=None, scale=None): + return attention_forward( + query=q, key=k, value=v, + attn_mask=attn_mask, scale=scale, + fused=True, head_first=False, + ) + def eager_attn(q, k, v, attn_mask=None, scale=None): q = q.transpose(1, 2) @@ -152,6 +163,7 @@ def attention( "sage", "sparge", "vsa", + "mindie", ] flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM if attn_impl is None or attn_impl == "auto": @@ -192,6 +204,8 @@ def attention( ) if XFORMERS_AVAILABLE: return xformers_attn(q, k, v, attn_mask=attn_mask, scale=scale) + if MINDIE_AVAILABLE: + return mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale) if SDPA_AVAILABLE: return sdpa_attn(q, k, v, attn_mask=attn_mask, scale=scale) if FLASH_ATTN_2_AVAILABLE: @@ -263,6 +277,8 @@ def attention( cdfthreshd=kwargs.get("cdfthreshd", 0.98), pvthreshd=kwargs.get("pvthreshd", 50), ) + if attn_impl == "mindie": + return mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale) if attn_impl == "vsa": return video_sparse_attn( q, @@ -354,7 +370,10 @@ def long_context_attention( "sage", "sparge", "vsa", + "mindie", ] + if attn_impl == "mindie": + raise RuntimeError("mindie long-context attention is not supported yet") assert attn_mask is None, "long context attention does not support attention mask" flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM if attn_impl is None or attn_impl == "auto": diff --git a/diffsynth_engine/pipelines/qwen_image.py b/diffsynth_engine/pipelines/qwen_image.py index 342cb70b..79741547 100644 --- a/diffsynth_engine/pipelines/qwen_image.py +++ b/diffsynth_engine/pipelines/qwen_image.py @@ -376,7 +376,9 @@ def update_weights(self, state_dicts: QwenImageStateDicts) -> None: self.update_component(self.vae, state_dicts.vae, self.config.device, self.config.vae_dtype) def compile(self): - self.dit.compile_repeated_blocks() + from diffsynth_engine.platforms import resolve_platform + platform_cls = resolve_platform(self.config.device) + self.dit.compile_repeated_blocks(**platform_cls.compile_kwargs()) def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False): assert self.config.tp_degree is None or self.config.tp_degree == 1, ( diff --git a/diffsynth_engine/platforms/__init__.py b/diffsynth_engine/platforms/__init__.py new file mode 100644 index 00000000..c90a953b --- /dev/null +++ b/diffsynth_engine/platforms/__init__.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import platform as host_platform +from functools import lru_cache +from typing import Type + +import torch + +from .ascend import ( + AscendPlatform, + probe_ascend_capabilities, + probe_ascend_feature, + reset_ascend_capability_cache, +) +from .base import PlatformBackend, PlatformCapabilities + + +class CPUPlatform(PlatformBackend): + name = "cpu" + device_type = "cpu" + + +class CUDAPlatform(PlatformBackend): + name = "cuda" + device_type = "cuda" + + @classmethod + def is_available(cls) -> bool: + return torch.cuda.is_available() + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + torch.cuda.set_device(index) + + @classmethod + def synchronize(cls) -> None: + torch.cuda.synchronize() + + @classmethod + def empty_cache(cls) -> None: + torch.cuda.empty_cache() + + @classmethod + def distributed_backend(cls) -> str: + return "nccl" + + +class ROCmPlatform(CUDAPlatform): + name = "rocm" + + +class MPSPlatform(PlatformBackend): + name = "mps" + device_type = "mps" + + @classmethod + def is_available(cls) -> bool: + return torch.backends.mps.is_available() + + @classmethod + def synchronize(cls) -> None: + torch.mps.synchronize() + + @classmethod + def empty_cache(cls) -> None: + torch.mps.empty_cache() + + +_PLATFORM_REGISTRY: dict[str, Type[PlatformBackend]] = { + "cpu": CPUPlatform, + "cuda": ROCmPlatform if torch.version.hip else CUDAPlatform, + "mps": MPSPlatform, + "npu": AscendPlatform, +} + + +def register_platform(device_type: str, platform_cls: Type[PlatformBackend], *, overwrite: bool = False) -> None: + if device_type in _PLATFORM_REGISTRY and not overwrite: + raise ValueError(f"Platform for device type {device_type!r} is already registered") + _PLATFORM_REGISTRY[device_type] = platform_cls + + +@lru_cache(maxsize=None) +def auto_detect_device() -> str: + """auto detect device type in order of cuda(gpu/rocm), npu, mps, cpu""" + for device_type in ("cuda", "npu", "mps", "cpu"): + try: + if resolve_platform(device_type).is_available(): + return device_type + except Exception: + continue + return "cpu" + + +def get_device_type(device: str | torch.device | None = None) -> str: + if device is None or (isinstance(device, str) and device.lower() in ("auto", "")): + return auto_detect_device() + if isinstance(device, torch.device): + return device.type + return str(device).split(":", 1)[0].lower() + + +def resolve_platform(device: str | torch.device) -> Type[PlatformBackend]: + device_type = get_device_type(device) + try: + return _PLATFORM_REGISTRY[device_type] + except KeyError as exc: + available = ", ".join(sorted(_PLATFORM_REGISTRY)) + raise ValueError(f"Unsupported device type {device_type!r}. Registered device types: {available}") from exc + + +def get_preferred_fp8_dtype(device: str | torch.device = "cuda") -> torch.dtype: + platform_cls = resolve_platform(device) + if platform_cls is ROCmPlatform and platform_cls.is_available(): + properties = torch.cuda.get_device_properties(0) + if "gfx94" in properties.gcnArchName: + return torch.float8_e4m3fnuz + return torch.float8_e4m3fn + + +def pin_memory( + tensor: torch.Tensor, + device: str | torch.device | None = None, +) -> torch.Tensor: + if host_platform.system() != "Linux": + return tensor + platform_cls = resolve_platform(get_device_type(device)) + return platform_cls.pin_memory(tensor) + + +__all__ = [ + "AscendPlatform", + "CPUPlatform", + "CUDAPlatform", + "MPSPlatform", + "PlatformBackend", + "PlatformCapabilities", + "ROCmPlatform", + "auto_detect_device", + "get_device_type", + "get_preferred_fp8_dtype", + "pin_memory", + "probe_ascend_capabilities", + "probe_ascend_feature", + "register_platform", + "reset_ascend_capability_cache", + "resolve_platform", +] diff --git a/diffsynth_engine/platforms/ascend.py b/diffsynth_engine/platforms/ascend.py new file mode 100644 index 00000000..bf3964a6 --- /dev/null +++ b/diffsynth_engine/platforms/ascend.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +import importlib +from functools import lru_cache +from typing import Any + +import torch + +from .base import PlatformBackend, PlatformCapabilities + + +def _import_torch_npu(): + try: + return importlib.import_module("torch_npu") + except (ImportError, OSError) as exc: + raise RuntimeError( + "Ascend device requested, but torch_npu is not installed. " + "Install the torch_npu wheel matching the PyTorch and CANN versions." + ) from exc + + +def _import_mindie_sd(): + try: + return importlib.import_module("mindiesd") + except (ImportError, OSError) as exc: + raise RuntimeError( + "This Ascend feature requires MindIE-SD. Install a MindIE-SD 3.x wheel " + "matching the current torch_npu and CANN versions." + ) from exc + + +def _has_callable(obj: Any, name: str) -> bool: + return callable(getattr(obj, name, None)) + + +def _probe_npu_runtime(npu: Any) -> bool: + try: + probe = torch.zeros(1, device="npu:0") + probe.add_(1) + npu.synchronize() + return True + except Exception: + return False + + +@lru_cache(maxsize=1) +def _probe_ascend_device() -> bool: + try: + torch_npu = _import_torch_npu() + npu = getattr(torch_npu, "npu", getattr(torch, "npu", None)) + device_available = bool(npu is not None and _has_callable(npu, "is_available") and npu.is_available()) + except Exception: + return False + + if device_available: + # is_available() can stay true while ACL initialization is failing. + device_available = _probe_npu_runtime(npu) + return device_available + + +@lru_cache(maxsize=1) +def _probe_mindie_installation() -> bool: + if not _probe_ascend_device(): + return False + + try: + _import_mindie_sd() + except Exception: + return False + return True + + +def _feature_api_available(feature: str) -> bool: + if feature == "mindie_attention": + module = importlib.import_module("mindiesd.layers.flash_attn.attention_forward") + return _has_callable(module, "attention_forward") + if feature == "mindie_compile": + module = importlib.import_module("mindiesd.compilation") + return callable(getattr(module, "MindieSDBackend", None)) + + raise ValueError(f"Unknown Ascend capability: {feature}") + + +def _tensor_probe_succeeded(output: torch.Tensor) -> bool: + torch_npu = _import_torch_npu() + torch_npu.npu.synchronize() + return bool(torch.isfinite(output).all().cpu().item()) + + +def _probe_mindie_attention_operation() -> bool: + module = importlib.import_module("mindiesd.layers.flash_attn.attention_forward") + query = torch.randn(1, 128, 8, 128, device="npu:0", dtype=torch.bfloat16) + with torch.no_grad(): + output = module.attention_forward( + query=query, + key=query, + value=query, + attn_mask=None, + scale=None, + fused=True, + head_first=False, + ) + return output.shape == query.shape and _tensor_probe_succeeded(output) + + +def _probe_mindie_compile_operation() -> bool: + compilation_module = importlib.import_module("mindiesd.compilation") + + def probe_fn(value): + return torch.nn.functional.gelu(value + 1) + + compiled_fn = torch.compile(probe_fn, backend=compilation_module.MindieSDBackend(), fullgraph=False) + value = torch.randn(8, 32, device="npu:0", dtype=torch.bfloat16) + with torch.no_grad(): + output = compiled_fn(value) + return output.shape == value.shape and _tensor_probe_succeeded(output) + + + + +_OPERATION_PROBES = { + "mindie_attention": _probe_mindie_attention_operation, + "mindie_compile": _probe_mindie_compile_operation, +} + + +@lru_cache(maxsize=None) +def probe_ascend_feature(feature: str) -> bool: + if feature == "device": + return _probe_ascend_device() + if feature == "mindie": + return _probe_mindie_installation() + if feature not in _OPERATION_PROBES: + raise ValueError(f"Unknown Ascend capability: {feature}") + if not _probe_mindie_installation(): + return False + + try: + if not _feature_api_available(feature): + return False + return bool(_OPERATION_PROBES[feature]()) + except Exception: + return False + + +def probe_ascend_capabilities() -> PlatformCapabilities: + device = probe_ascend_feature("device") + if not device: + return PlatformCapabilities() + mindie = probe_ascend_feature("mindie") + if not mindie: + return PlatformCapabilities(device=True) + + return PlatformCapabilities( + device=True, + mindie=True, + mindie_attention=probe_ascend_feature("mindie_attention"), + mindie_compile=probe_ascend_feature("mindie_compile"), + ) + + +def reset_ascend_capability_cache() -> None: + _probe_ascend_device.cache_clear() + _probe_mindie_installation.cache_clear() + probe_ascend_feature.cache_clear() + + +class AscendPlatform(PlatformBackend): + name = "ascend" + device_type = "npu" + + @classmethod + def is_available(cls) -> bool: + return probe_ascend_feature("device") + + @classmethod + def normalize_device(cls, device: str | torch.device) -> torch.device: + _import_torch_npu() + return torch.device(device) + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + torch_npu = _import_torch_npu() + torch_npu.npu.set_device(index) + + @classmethod + def synchronize(cls) -> None: + torch_npu = _import_torch_npu() + torch_npu.npu.synchronize() + + @classmethod + def empty_cache(cls) -> None: + torch_npu = _import_torch_npu() + torch_npu.npu.empty_cache() + + @classmethod + def pin_memory(cls, tensor: torch.Tensor) -> torch.Tensor: + _import_torch_npu() + try: + return tensor.pin_memory(device="npu") + except (RuntimeError, TypeError): + # Pageable CPU memory is slower but remains correct on runtimes that + # do not expose an NPU-specific pinned allocator. + return tensor + + @classmethod + def distributed_backend(cls) -> str: + return "hccl" + + @classmethod + def compile_backend(cls): + if not cls.supports("mindie_compile"): + raise RuntimeError( + "MindIE-SD compilation was requested, but MindieSDBackend is unavailable " + "in the installed MindIE-SD package." + ) + from mindiesd.compilation import CompilationConfig, MindieSDBackend + + CompilationConfig.fusion_patterns.enable_fast_gelu = False + return MindieSDBackend() + + @classmethod + def capabilities(cls) -> PlatformCapabilities: + return probe_ascend_capabilities() + + @classmethod + def supports(cls, capability: str) -> bool: + return probe_ascend_feature(capability) diff --git a/diffsynth_engine/platforms/base.py b/diffsynth_engine/platforms/base.py new file mode 100644 index 00000000..955c55d9 --- /dev/null +++ b/diffsynth_engine/platforms/base.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from abc import ABC +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass(frozen=True) +class PlatformCapabilities: + device: bool = False + mindie: bool = False + mindie_attention: bool = False + mindie_compile: bool = False + + +class PlatformBackend(ABC): + name = "unknown" + device_type = "cpu" + + @classmethod + def is_available(cls) -> bool: + return True + + @classmethod + def normalize_device(cls, device: str | torch.device) -> torch.device: + return torch.device(device) + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + return None + + @classmethod + def synchronize(cls) -> None: + return None + + @classmethod + def empty_cache(cls) -> None: + return None + + @classmethod + def pin_memory(cls, tensor: torch.Tensor) -> torch.Tensor: + return tensor.pin_memory() + + @classmethod + def distributed_backend(cls) -> str: + return "gloo" + + @classmethod + def compile_backend(cls) -> Any | None: + return None + + @classmethod + def compile_kwargs(cls) -> dict[str, Any]: + backend = cls.compile_backend() + return {} if backend is None else {"backend": backend} + + @classmethod + def supports(cls, capability: str) -> bool: + capabilities = cls.capabilities() + if not hasattr(capabilities, capability): + raise ValueError(f"Unknown platform capability: {capability}") + return bool(getattr(capabilities, capability)) + + @classmethod + def capabilities(cls) -> PlatformCapabilities: + return PlatformCapabilities(device=cls.is_available()) diff --git a/diffsynth_engine/utils/flag.py b/diffsynth_engine/utils/flag.py index 51aabc01..d0ee0e24 100644 --- a/diffsynth_engine/utils/flag.py +++ b/diffsynth_engine/utils/flag.py @@ -2,6 +2,7 @@ import torch from diffsynth_engine.utils import logging +from diffsynth_engine.utils.platform import is_mindie_sd_available logger = logging.get_logger(__name__) @@ -41,6 +42,10 @@ def check_module_available(module_path: str, module_name: str = None) -> bool: SPARGE_ATTN_AVAILABLE = check_module_available("spas_sage_attn", "Sparge attention") VIDEO_SPARSE_ATTN_AVAILABLE = check_module_available("vsa", "Video sparse attention") +# NPU +MINDIE_AVAILABLE = is_mindie_sd_available() + + NUNCHAKU_AVAILABLE = check_module_available("nunchaku", "Nunchaku") NUNCHAKU_IMPORT_ERROR = None if not NUNCHAKU_AVAILABLE: diff --git a/diffsynth_engine/utils/platform.py b/diffsynth_engine/utils/platform.py index 49a69680..060d7e7d 100644 --- a/diffsynth_engine/utils/platform.py +++ b/diffsynth_engine/utils/platform.py @@ -1,25 +1,42 @@ # cross-platform definitions and utilities -import torch import gc -import platform +import torch + +from diffsynth_engine.platforms import ( + AscendPlatform, + get_device_type as _get_device_type, + get_preferred_fp8_dtype, + pin_memory as _pin_memory, + resolve_platform, +) -# data type -# AMD only supports float8_e4m3fnuz -# https://onnx.ai/onnx/technical/float8.html -if torch.version.hip and "gfx94" in torch.cuda.get_device_properties(0).gcnArchName: - DTYPE_FP8 = torch.float8_e4m3fnuz -else: - DTYPE_FP8 = torch.float8_e4m3fn +DTYPE_FP8 = get_preferred_fp8_dtype("cuda") -def empty_cache(): +def empty_cache(device: str | torch.device | None = None): gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - if torch.mps.is_available(): - torch.mps.empty_cache() + resolve_platform(get_device_type(device)).empty_cache() + + +def pin_memory( + tensor: torch.Tensor, + device: str | torch.device | None = None, +): + return _pin_memory(tensor, device) + + +def is_npu_available() -> bool: + return AscendPlatform.supports("device") + + +def is_mindie_sd_available() -> bool: + return AscendPlatform.supports("mindie") + + +def get_device_type(device: str | torch.device | None = None) -> str: + return _get_device_type(device) -def pin_memory(tensor: torch.Tensor): - return tensor.pin_memory() if platform.system() == "Linux" else tensor +def get_torch_distributed_backend(device: str | torch.device) -> str: + return resolve_platform(device).distributed_backend()