From bf49ef1be4272d1c2ad4cab077df2123a16104fb Mon Sep 17 00:00:00 2001 From: hammer Date: Thu, 13 Aug 2026 12:07:23 +0800 Subject: [PATCH 1/5] feat(npu): add Ascend platform abstraction and MindIE attention Add platforms registry with Ascend capability probes, extend utils/platform for npu/hccl/compile kwargs, and register a MindIE attention backend for Qwen-Image on the v1 layout. Co-authored-by: Cursor --- .../layers/attention/backends/abstract.py | 1 + .../layers/attention/backends/mindie_attn.py | 106 ++++++++ diffsynth_engine/pipelines/base.py | 4 +- diffsynth_engine/platforms/__init__.py | 153 ++++++++++++ diffsynth_engine/platforms/ascend.py | 233 ++++++++++++++++++ diffsynth_engine/platforms/base.py | 72 ++++++ diffsynth_engine/registry.py | 1 + diffsynth_engine/utils/platform.py | 104 +++++++- 8 files changed, 664 insertions(+), 10 deletions(-) create mode 100644 diffsynth_engine/layers/attention/backends/mindie_attn.py 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/layers/attention/backends/abstract.py b/diffsynth_engine/layers/attention/backends/abstract.py index e9b5fdd..35be745 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 0000000..3568219 --- /dev/null +++ b/diffsynth_engine/layers/attention/backends/mindie_attn.py @@ -0,0 +1,106 @@ +import torch + +from diffsynth_engine.layers.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder, + AttentionType, +) +from diffsynth_engine.utils import logging + +logger = logging.get_logger(__name__) + + +class MindieAttentionMetadataBuilder(AttentionMetadataBuilder): + def __init__(self) -> None: + pass + + def build(self, **kwargs) -> AttentionMetadata: + return AttentionMetadata() + + +class MindieAttentionBackend(AttentionBackend): + @staticmethod + def check_availability() -> None: + from diffsynth_engine.platforms import AscendPlatform + + if not AscendPlatform.supports("device"): + error_msg = "MindIE attention requires an available Ascend NPU device." + logger.error(error_msg) + raise RuntimeError(error_msg) + if not AscendPlatform.supports("mindie_attention"): + error_msg = ( + "MindIE attention backend is not available. " + "Install MindIE-SD 3.x matching the current torch_npu and CANN versions, " + "and ensure mindiesd.layers.flash_attn.attention_forward works on NPU." + ) + logger.error(error_msg) + raise RuntimeError(error_msg) + + @staticmethod + def get_type() -> str: + return str(AttentionType.MINDIE) + + @staticmethod + def get_impl_cls() -> type["AttentionImpl"]: + return MindieAttentionImpl + + @staticmethod + def get_metadata_cls() -> type["AttentionMetadata"]: + return AttentionMetadata + + @staticmethod + def get_builder_cls() -> type["AttentionMetadataBuilder"]: + return MindieAttentionMetadataBuilder + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [] + + @classmethod + def supports_ring_attention(cls) -> bool: + return False + + +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: + if num_kv_heads is None: + num_kv_heads = num_heads + self.num_kv_groups = num_heads // num_kv_heads + self.causal = causal + self.softmax_scale = softmax_scale + self.num_heads = num_heads + self.head_size = head_size + + 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.softmax_scale, + fused=True, + head_first=False, + opt_mode="manual", + op_type="fused_attn_score", + layout="BSND", + ) diff --git a/diffsynth_engine/pipelines/base.py b/diffsynth_engine/pipelines/base.py index 73509cb..ffa720b 100644 --- a/diffsynth_engine/pipelines/base.py +++ b/diffsynth_engine/pipelines/base.py @@ -19,6 +19,7 @@ from diffsynth_engine.forward_context import set_forward_context from diffsynth_engine.utils import logging from diffsynth_engine.utils.load_utils import fix_state_dict_key, load_model_weights, prepare_model_weights +from diffsynth_engine.utils.platform import get_compile_kwargs logger = logging.get_logger(__name__) @@ -41,10 +42,11 @@ def compile_transformer_blocks(model: nn.Module) -> nn.Module: if not repeated_blocks: raise ValueError(f"`_repeated_blocks` is not defined for {type(model).__name__}") + compile_kwargs = get_compile_kwargs() has_compiled_region = False for submodule in model.modules(): if submodule.__class__.__name__ in repeated_blocks: - submodule.compile() + submodule.compile(**compile_kwargs) has_compiled_region = True if not has_compiled_region: diff --git a/diffsynth_engine/platforms/__init__.py b/diffsynth_engine/platforms/__init__.py new file mode 100644 index 0000000..f8521b0 --- /dev/null +++ b/diffsynth_engine/platforms/__init__.py @@ -0,0 +1,153 @@ +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 device_count(cls) -> int: + return torch.cuda.device_count() + + @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 parse_device_type(device: str | torch.device | None = None) -> str: + """Parse a device spec, or auto-detect when ``device`` is None/auto.""" + 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 = parse_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(parse_device_type(device)) + return platform_cls.pin_memory(tensor) + + +__all__ = [ + "AscendPlatform", + "CPUPlatform", + "CUDAPlatform", + "MPSPlatform", + "PlatformBackend", + "PlatformCapabilities", + "ROCmPlatform", + "auto_detect_device", + "get_preferred_fp8_dtype", + "parse_device_type", + "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 0000000..5e37928 --- /dev/null +++ b/diffsynth_engine/platforms/ascend.py @@ -0,0 +1,233 @@ +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 device_count(cls) -> int: + _import_torch_npu() + return torch.npu.device_count() + + @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 0000000..cce56eb --- /dev/null +++ b/diffsynth_engine/platforms/base.py @@ -0,0 +1,72 @@ +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 device_count(cls) -> int: + return 1 + + @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/registry.py b/diffsynth_engine/registry.py index b6b9f4e..84f7fa0 100644 --- a/diffsynth_engine/registry.py +++ b/diffsynth_engine/registry.py @@ -33,6 +33,7 @@ "fa3": "diffsynth_engine.layers.attention.backends.flash_attn_3:FlashAttention3Backend", "fa3_fp8": "diffsynth_engine.layers.attention.backends.flash_attn_3:FlashAttention3FP8Backend", "fa4": "diffsynth_engine.layers.attention.backends.flash_attn_4:FlashAttention4Backend", + "mindie": "diffsynth_engine.layers.attention.backends.mindie_attn:MindieAttentionBackend", "sage2": "diffsynth_engine.layers.attention.backends.sage_attn_2:SageAttention2Backend", "sage3": "diffsynth_engine.layers.attention.backends.sage_attn_3:SageAttention3Backend", "sdpa": "diffsynth_engine.layers.attention.backends.sdpa:SDPABackend", diff --git a/diffsynth_engine/utils/platform.py b/diffsynth_engine/utils/platform.py index 575f573..2ff4a95 100644 --- a/diffsynth_engine/utils/platform.py +++ b/diffsynth_engine/utils/platform.py @@ -1,7 +1,12 @@ import torch +from diffsynth_engine.utils import logging + +logger = logging.get_logger(__name__) + def _is_cuda() -> bool: + # Historical: torch build has CUDA, not necessarily a visible GPU. return torch.version.cuda is not None @@ -13,31 +18,112 @@ def _is_mps() -> bool: return torch.backends.mps.is_available() +def _active_platform(): + """Resolve PlatformBackend for the process-preferred accelerator.""" + from diffsynth_engine.platforms import resolve_platform + + return resolve_platform(get_device_type()) + + +def is_npu_available() -> bool: + from diffsynth_engine.platforms import AscendPlatform + + return AscendPlatform.is_available() + + +def is_mindie_sd_available() -> bool: + from diffsynth_engine.platforms import AscendPlatform + + return AscendPlatform.supports("mindie") + + def get_device(local_rank: int) -> torch.device: if _is_cuda() or _is_rocm(): return torch.device("cuda", local_rank) + if is_npu_available(): + return torch.device("npu", local_rank) if _is_mps(): return torch.device("mps") - else: - return torch.device("cpu") + return torch.device("cpu") def get_device_type() -> str: + """Preferred accelerator for this process (no-arg, v1 public API). + + Priority matches historical utils behavior: cuda/rocm build > npu > mps > cpu. + Differs from ``platforms.auto_detect_device`` which uses ``is_available()``. + """ if _is_cuda() or _is_rocm(): return "cuda" + if is_npu_available(): + return "npu" if _is_mps(): return "mps" - else: - return "cpu" + return "cpu" def get_torch_distributed_backend() -> str: - if _is_cuda() or _is_rocm(): - return "nccl" - if _is_mps(): - return "gloo" - else: + device_type = get_device_type() + if device_type == "cpu": raise NotImplementedError("Unsupported device type") + return _active_platform().distributed_backend() + + +def device_count() -> int: + return _active_platform().device_count() + + +def set_device(index: int | str | torch.device) -> None: + """Bind the current process to a local device (cuda or npu).""" + _active_platform().set_device(index) + + +def align_config_device(config_device: str | torch.device, target_type: str | None = None) -> str: + """Rewrite historical CUDA placeholder to NPU when Ascend is the active accelerator. + + Leaves other placeholders alone (e.g. default ``cuda`` on a CPU laptop). + Explicit ``npu`` on a non-NPU machine raises. + """ + if target_type is None: + target_type = get_device_type() + device_str = str(config_device) + current_type = device_str.split(":", 1)[0].lower() + if current_type == target_type: + return device_str + if target_type == "npu" and current_type == "cuda": + return "npu" + if current_type == "npu" and target_type != "npu": + raise RuntimeError( + f"config.device={config_device!r} does not match available device_type={target_type!r}" + ) + return device_str + + +def bind_rank_device(config_device: str | torch.device, local_rank: int) -> str: + """Worker-only: bind config.device to this rank's local device (e.g. npu:0).""" + device_type = str(config_device).split(":", 1)[0].lower() + if device_type in ("cpu", "mps"): + return str(config_device) + return f"{device_type}:{local_rank}" + + +def get_compile_kwargs() -> dict: + """Return kwargs for ``nn.Module.compile`` / ``torch.compile``. + + On Ascend with MindIE compile available, injects MindieSDBackend. + Otherwise returns ``{}`` so the default inductor path is used. + """ + if not is_npu_available(): + return {} + + from diffsynth_engine.platforms import AscendPlatform + + if not AscendPlatform.supports("mindie_compile"): + logger.warning( + "MindIE-SD compile backend is unavailable; falling back to default torch.compile backend" + ) + return {} + return AscendPlatform.compile_kwargs() DTYPE_FP8 = torch.float8_e4m3fnuz if _is_rocm() else torch.float8_e4m3fn From 8c1e61cf325eb46be5db62f63d5190441aceefd7 Mon Sep 17 00:00:00 2001 From: hammer Date: Thu, 13 Aug 2026 12:07:23 +0800 Subject: [PATCH 2/5] feat(npu): enable Ulysses SP with CFG parallel on Ascend Bind Ascend devices before HCCL init_process_group, align/bind config.device across engine workers, and reject MindIE with ring SP. Co-authored-by: Cursor --- .../distributed/parallel_state.py | 31 +++++++++++-------- diffsynth_engine/engine.py | 3 +- diffsynth_engine/layers/attention/layer.py | 7 +++++ diffsynth_engine/worker.py | 5 +++ 4 files changed, 32 insertions(+), 14 deletions(-) diff --git a/diffsynth_engine/distributed/parallel_state.py b/diffsynth_engine/distributed/parallel_state.py index a2f91d7..9f9b01e 100644 --- a/diffsynth_engine/distributed/parallel_state.py +++ b/diffsynth_engine/distributed/parallel_state.py @@ -13,7 +13,6 @@ import torch import torch.distributed -from torch.cuda import device_count, set_device from diffsynth_engine.distributed.group_coordinator import ( GroupCoordinator, @@ -22,7 +21,11 @@ ) from diffsynth_engine.utils import logging from diffsynth_engine.utils.constants import IDLE_TIMEOUT_SEC -from diffsynth_engine.utils.platform import get_torch_distributed_backend +from diffsynth_engine.utils.platform import ( + device_count, + get_torch_distributed_backend, + set_device, +) logger = logging.get_logger(__name__) @@ -425,10 +428,22 @@ def init_distributed_environment( distributed_init_method, backend, ) + # local_rank is not available in torch ProcessGroup, + # see https://github.com/pytorch/pytorch/issues/122816 + if local_rank == -1: + # local rank not set, this usually happens in single-node + # setting, where we can use rank as local rank + if distributed_init_method == "env://": + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + else: + local_rank = rank if rank >= 0 else 0 + if not torch.distributed.is_initialized(): assert distributed_init_method is not None, ( "distributed_init_method must be provided when initializing distributed environment" ) + # Bind device before init_process_group (required by HCCL on Ascend). + set_device(local_rank % max(device_count(), 1)) # this backend is used for WORLD torch.distributed.init_process_group( backend=backend, @@ -436,17 +451,7 @@ def init_distributed_environment( world_size=world_size, rank=rank, ) - set_device(torch.distributed.get_rank() % device_count()) - # set the local rank - # local_rank is not available in torch ProcessGroup, - # see https://github.com/pytorch/pytorch/issues/122816 - if local_rank == -1: - # local rank not set, this usually happens in single-node - # setting, where we can use rank as local rank - if distributed_init_method == "env://": - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - else: - local_rank = rank + set_device(torch.distributed.get_rank() % max(device_count(), 1)) global _WORLD if _WORLD is None: ranks = list(range(torch.distributed.get_world_size())) diff --git a/diffsynth_engine/engine.py b/diffsynth_engine/engine.py index 13bf802..b06a6c3 100644 --- a/diffsynth_engine/engine.py +++ b/diffsynth_engine/engine.py @@ -1,7 +1,6 @@ from typing import Any import torch.multiprocessing as mp -from torch.cuda import set_device from diffsynth_engine.configs import PipelineConfig from diffsynth_engine.registry import ( @@ -9,6 +8,7 @@ get_pipeline_class_name, ) from diffsynth_engine.utils import logging +from diffsynth_engine.utils.platform import align_config_device, set_device from diffsynth_engine.utils.torch_profiler import TorchProfiler from diffsynth_engine.worker import run_worker_loop @@ -19,6 +19,7 @@ class DiffSynthEngine: @classmethod def from_pretrained(cls, model_path_or_config: str | PipelineConfig, **kwargs): pipeline_config = _resolve_pipeline_config(model_path_or_config) + pipeline_config.device = align_config_device(pipeline_config.device) num_workers = pipeline_config.parallelism master_addr = kwargs.get("master_addr", "localhost") master_port = kwargs.get("master_port", 29500) diff --git a/diffsynth_engine/layers/attention/layer.py b/diffsynth_engine/layers/attention/layer.py index 1bde443..cad43dc 100644 --- a/diffsynth_engine/layers/attention/layer.py +++ b/diffsynth_engine/layers/attention/layer.py @@ -101,6 +101,7 @@ def __init__( self.num_kv_heads = num_kv_heads self.scatter_idx = scatter_idx self.gather_idx = gather_idx + self.attn_type = str(attn_type) if attn_type is not None else None attn_backend = get_attn_backend(attn_type) if not attn_backend.supports_head_size(head_size): @@ -147,6 +148,12 @@ def forward( ulysses_parallel_world_size = get_ulysses_parallel_world_size() if is_sp_group_initialized() else 1 ring_parallel_world_size = get_ring_parallel_world_size() if is_sp_group_initialized() else 1 + if ring_parallel_world_size > 1 and self.attn_type == "mindie": + raise RuntimeError( + "NPU MindIE attention currently supports Ulysses only " + f"(sp_ring_degree must be 1, got {ring_parallel_world_size})" + ) + if ulysses_parallel_world_size > 1: q = SeqAllToAll4D.apply(get_sp_group().ulysses_group, q, self.scatter_idx, self.gather_idx) k = SeqAllToAll4D.apply(get_sp_group().ulysses_group, k, self.scatter_idx, self.gather_idx) diff --git a/diffsynth_engine/worker.py b/diffsynth_engine/worker.py index ac3e00f..53085bf 100644 --- a/diffsynth_engine/worker.py +++ b/diffsynth_engine/worker.py @@ -11,6 +11,7 @@ ) from diffsynth_engine.registry import get_pipeline_class from diffsynth_engine.utils import logging +from diffsynth_engine.utils.platform import bind_rank_device from diffsynth_engine.utils.torch_profiler import TorchProfiler logger = logging.get_logger(__name__) @@ -38,6 +39,10 @@ def __init__( os.environ["LOCAL_RANK"] = str(local_rank) os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) + + # Bind config.device to this rank's local device before HCCL init / model load. + self.pipeline_config.device = bind_rank_device(self.pipeline_config.device, local_rank) + init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank) cfg_degree = 2 if pipeline_config.use_cfg_parallel else 1 From 0cf05b6c8e6641266820051f0bae550786fbfaf2 Mon Sep 17 00:00:00 2001 From: hammer Date: Thu, 13 Aug 2026 12:07:23 +0800 Subject: [PATCH 3/5] feat(npu): fuse Qwen DiT addcmul, MindIE RoPE, and LN modulate Use addcmul for gated residual/modulate micro-opts, and gate MindIE RoPE / layernorm_scale_shift behind USE_MINDIESD_FUSE. Co-authored-by: Cursor --- .../qwen_image/transformer_qwenimage.py | 94 ++++++++++++------- 1 file changed, 62 insertions(+), 32 deletions(-) diff --git a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py index 895c2e3..d9d946c 100644 --- a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py +++ b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py @@ -15,6 +15,7 @@ # limitations under the License. import functools +import os from math import prod from typing import Any, Dict, List, Optional, Tuple, Union @@ -36,9 +37,12 @@ from diffsynth_engine.layers.tensor_parallel import ColumnParallelLinear, RowParallelLinear, TPFeedForward from diffsynth_engine.models.base import DiffusionModel from diffsynth_engine.utils import logging +from diffsynth_engine.utils.platform import is_mindie_sd_available logger = logging.get_logger(__name__) +USE_MINDIESD_FUSE = os.environ.get("USE_MINDIESD_FUSE", "0") == "1" + def apply_rotary_emb_qwen( x: torch.Tensor, @@ -81,6 +85,32 @@ def apply_rotary_emb_qwen( return out else: + if USE_MINDIESD_FUSE and is_mindie_sd_available() and x.device.type == "npu": + from mindiesd import rotary_position_embedding + + # Cache expanded cos/sin on the tensor object itself. + # Python object identity avoids data_ptr collision across different-length slices. + cached = getattr(freqs_cis, "_rope_expanded", None) + if cached is None: + cos = freqs_cis.real # (s, d/2) + sin = freqs_cis.imag + cos = cos.reshape(1, -1, 1, cos.shape[-1]) # (1, S, 1, D/2) + sin = sin.reshape(1, -1, 1, sin.shape[-1]) + cos = cos.unsqueeze(-1).expand(-1, -1, -1, -1, 2).flatten(start_dim=-2) # (1, S, 1, D) + sin = sin.unsqueeze(-1).expand(-1, -1, -1, -1, 2).flatten(start_dim=-2) + cos, sin = cos.to(x.device), sin.to(x.device) + cached = (cos, sin) + freqs_cis._rope_expanded = cached + cos, sin = cached + return rotary_position_embedding( + x, + cos, + sin, + rotated_mode="rotated_interleaved", + head_first=False, + fused=True, + ) + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) freqs_cis = freqs_cis.unsqueeze(1) x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) @@ -584,32 +614,25 @@ def __init__( self.zero_cond_t = zero_cond_t - def _modulate(self, x, mod_params, index=None): - """Apply modulation to input tensor""" - # x: b l d, shift: b d, scale: b d, gate: b d + def _split_mod_params(self, mod_params, index=None): + """Split modulation params into shift/scale/gate, optionally indexed for CFG.""" shift, scale, gate = mod_params.chunk(3, dim=-1) if index is not None: - # Assuming mod_params batch dim is 2*actual_batch (chunked into 2 parts) - # So shift, scale, gate have shape [2*actual_batch, d] actual_batch = shift.size(0) // 2 - shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:] # each: [actual_batch, d] + shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:] scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:] gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:] - # index: [b, l] where b is actual batch size - # Expand to [b, l, 1] to match feature dimension - index_expanded = index.unsqueeze(-1) # [b, l, 1] + index_expanded = index.unsqueeze(-1) - # Expand chunks to [b, 1, d] then broadcast to [b, l, d] - shift_0_exp = shift_0.unsqueeze(1) # [b, 1, d] - shift_1_exp = shift_1.unsqueeze(1) # [b, 1, d] + shift_0_exp = shift_0.unsqueeze(1) + shift_1_exp = shift_1.unsqueeze(1) scale_0_exp = scale_0.unsqueeze(1) scale_1_exp = scale_1.unsqueeze(1) gate_0_exp = gate_0.unsqueeze(1) gate_1_exp = gate_1.unsqueeze(1) - # Use torch.where to select based on index shift_result = torch.where(index_expanded == 0, shift_0_exp, shift_1_exp) scale_result = torch.where(index_expanded == 0, scale_0_exp, scale_1_exp) gate_result = torch.where(index_expanded == 0, gate_0_exp, gate_1_exp) @@ -618,7 +641,23 @@ def _modulate(self, x, mod_params, index=None): scale_result = scale.unsqueeze(1) gate_result = gate.unsqueeze(1) - return x * (1 + scale_result) + shift_result, gate_result + return shift_result, scale_result, gate_result + + def _modulate(self, x, mod_params, index=None): + """Apply modulation to input tensor""" + shift_result, scale_result, gate_result = self._split_mod_params(mod_params, index) + # x*(1+scale)+shift == x + x*scale + shift — prefer addcmul over Adds+Mul+Add + return torch.addcmul(x, x, scale_result) + shift_result, gate_result + + def _norm_modulate(self, norm: nn.LayerNorm, x: torch.Tensor, mod_params, index=None): + """LayerNorm + Ada modulate. Fuses via mindiesd.layernorm_scale_shift when enabled.""" + if USE_MINDIESD_FUSE and is_mindie_sd_available() and x.device.type == "npu": + from mindiesd import layernorm_scale_shift + + shift_result, scale_result, gate_result = self._split_mod_params(mod_params, index) + out = layernorm_scale_shift(norm, x, scale_result, shift_result, fused=True) + return out, gate_result + return self._modulate(norm(x), mod_params, index) def forward( self, @@ -641,13 +680,8 @@ def forward( img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] - # Process image stream - norm1 + modulation - img_normed = self.img_norm1(hidden_states) - img_modulated, img_gate1 = self._modulate(img_normed, img_mod1, modulate_index) - - # Process text stream - norm1 + modulation - txt_normed = self.txt_norm1(encoder_hidden_states) - txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1) + img_modulated, img_gate1 = self._norm_modulate(self.img_norm1, hidden_states, img_mod1, modulate_index) + txt_modulated, txt_gate1 = self._norm_modulate(self.txt_norm1, encoder_hidden_states, txt_mod1) # Use QwenDoubleStreamAttention for joint attention computation # This directly implements the DoubleStreamLayerMegatron logic: @@ -665,21 +699,17 @@ def forward( image_rotary_emb=image_rotary_emb, ) - # Apply attention gates and add residual (like in Megatron) - hidden_states = hidden_states + img_gate1 * img_attn_output - encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output + # addcmul: residual + gate * out — prefer single op over Mul+Add + hidden_states = torch.addcmul(hidden_states, img_gate1, img_attn_output) + encoder_hidden_states = torch.addcmul(encoder_hidden_states, txt_gate1, txt_attn_output) - # Process image stream - norm2 + MLP - img_normed2 = self.img_norm2(hidden_states) - img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2, modulate_index) + img_modulated2, img_gate2 = self._norm_modulate(self.img_norm2, hidden_states, img_mod2, modulate_index) img_mlp_output = self.img_mlp(img_modulated2) - hidden_states = hidden_states + img_gate2 * img_mlp_output + hidden_states = torch.addcmul(hidden_states, img_gate2, img_mlp_output) - # Process text stream - norm2 + MLP - txt_normed2 = self.txt_norm2(encoder_hidden_states) - txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2) + txt_modulated2, txt_gate2 = self._norm_modulate(self.txt_norm2, encoder_hidden_states, txt_mod2) txt_mlp_output = self.txt_mlp(txt_modulated2) - encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output + encoder_hidden_states = torch.addcmul(encoder_hidden_states, txt_gate2, txt_mlp_output) # Clip to prevent overflow for fp16 if encoder_hidden_states.dtype == torch.float16: From 67b93a161c0f164aaae003febb68c4afa4b8dcf5 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 13 Aug 2026 20:28:57 +0800 Subject: [PATCH 4/5] Add AscendLongContextAttention: NPU/MindIE long-context attention under Ulysses SP, hiding all-to-all latency behind FA compute (compute-communication overlap). - Ulysses flow via all_to_all_4D_pre/single/after (SeqAllToAll4D). - FA_ALLTOALL_CUT/OVERLAP count semantics (0/1 off, >1 on, CUT wins): baseline single round trip; cut = head-chunked round trips; overlap = all-to-all on side stream2 overlapped with FA on main stream via npu Event/Stream sync. - current_stream refreshed at forward time. - Ulysses-only; falls back to USPAttention when not available. --- diffsynth_engine/layers/attention/__init__.py | 3 +- diffsynth_engine/layers/attention/layer.py | 338 ++++++++++++++++++ .../qwen_image/transformer_qwenimage.py | 24 +- 3 files changed, 359 insertions(+), 6 deletions(-) diff --git a/diffsynth_engine/layers/attention/__init__.py b/diffsynth_engine/layers/attention/__init__.py index ffc2b45..27ac945 100644 --- a/diffsynth_engine/layers/attention/__init__.py +++ b/diffsynth_engine/layers/attention/__init__.py @@ -1,9 +1,10 @@ from .backends.abstract import AttentionMetadata, AttentionType -from .layer import LocalAttention, USPAttention +from .layer import LocalAttention, USPAttention, AscendLongContextAttention __all__ = [ "AttentionType", "AttentionMetadata", "LocalAttention", "USPAttention", + "AscendLongContextAttention", ] diff --git a/diffsynth_engine/layers/attention/layer.py b/diffsynth_engine/layers/attention/layer.py index cad43dc..3896039 100644 --- a/diffsynth_engine/layers/attention/layer.py +++ b/diffsynth_engine/layers/attention/layer.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 +import os +from torch import distributed as dist import torch import torch.nn as nn @@ -168,3 +170,339 @@ def forward( if ulysses_parallel_world_size > 1: output = SeqAllToAll4D.apply(get_sp_group().ulysses_group, output, self.gather_idx, self.scatter_idx) return output + + +from typing import Optional +class AscendLongContextAttention(nn.Module): + def __init__( + self, + num_heads: int = 24, + head_size: int = 128, + softmax_scale: float | None = None, + causal: bool = False, + num_kv_heads: int | None = None, + attn_type: str = "mindie", + scatter_idx: int = 2, + gather_idx: int = 1, + fa_head_loop: int | None = None, + **extra_impl_args, + ) -> None: + super().__init__() + if num_kv_heads is None: + num_kv_heads = num_heads + + self.scatter_idx = scatter_idx + self.gather_idx = gather_idx + self.num_heads = num_heads + self.head_size = head_size + self.num_kv_heads = num_kv_heads + + self.ulysses_pg = get_sp_group().ulysses_group + self.sp_ulysses_degree = get_sp_group().ulysses_world_size + self.sp_ring_degree = get_sp_group().ring_world_size + + + self.fa_alltoall_overlap = int(os.getenv('FA_ALLTOALL_OVERLAP', 1)) + self.fa_alltoall_cut = int(os.getenv('FA_ALLTOALL_CUT', 1)) + if fa_head_loop is not None: + self.fa_head_loop = fa_head_loop + elif self.fa_alltoall_cut > 1: + self.fa_head_loop = self.fa_alltoall_cut + elif self.fa_alltoall_overlap > 1: + self.fa_head_loop = self.fa_alltoall_overlap + else: + self.fa_head_loop = self.num_heads // self.sp_ulysses_degree + + if self.fa_alltoall_overlap > 1 and self.fa_alltoall_cut <= 1: + self.current_stream = torch.npu.current_stream() + self.stream2 = torch.npu.Stream() + self.event = [] + for i in range(self.fa_head_loop): + self.event.append(torch.npu.Event()) + + self.attn_type = str(attn_type) if attn_type is not None else None + attn_backend = get_attn_backend(attn_type) + if not attn_backend.supports_head_size(head_size): + raise ValueError(f"Attention backend {attn_type!r} does not support head size {head_size}.") + + impl_cls = attn_backend.get_impl_cls() + self.attn_impl = impl_cls( + num_heads=num_heads, + head_size=head_size, + softmax_scale=softmax_scale, + causal=causal, + num_kv_heads=num_kv_heads, + **extra_impl_args, + ) + + # TODO: currunt MindIE only support Ulysses + if self.sp_ring_degree > 1: + raise RuntimeError( + "NPU MindIE attention currently supports Ulysses only " + f"(sp_ring_degree must be 1, got {self.sp_ring_degree})" + ) + + + def _run_attention(self, q, k, v, **attn_kwargs): + return self.attn_impl.forward(q, k, v, **attn_kwargs) + + @staticmethod + def all_to_all_4D_pre(input: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None): + assert ( + input.dim() == 4 + ), f"input must be 4D tensor, got {input.dim()} and shape {input.shape}" + + seq_world_size = dist.get_world_size(group) + + if scatter_idx == 2 and gather_idx == 1: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs) + bs, shard_seqlen, hc, hs = input.shape + seqlen = shard_seqlen * seq_world_size + shard_hc = hc // seq_world_size + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs) + input_t = ( + input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs) + .transpose(0, 2) + .contiguous() + ) + + return input_t + + elif scatter_idx == 1 and gather_idx == 2: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs) + bs, seqlen, shard_hc, hs = input.shape + hc = shard_hc * seq_world_size + shard_seqlen = seqlen // seq_world_size + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs) + input_t = ( + input.reshape(bs, seq_world_size, shard_seqlen, shard_hc, hs) + .transpose(0, 3) + .transpose(0, 1) + .contiguous() + .reshape(seq_world_size, shard_hc, shard_seqlen, bs, hs) + ) + + return input_t + else: + raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2") + + @staticmethod + def all_to_all_4D_after(input: torch.tensor, output: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, + group=None): + seq_world_size = dist.get_world_size(group) + + if scatter_idx == 2 and gather_idx == 1: + bs, shard_seqlen, hc, hs = input.shape + + seqlen = shard_seqlen * seq_world_size + shard_hc = hc // seq_world_size + + output = output.reshape(seqlen, bs, shard_hc, hs) + + # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs) + output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs) + return output + elif scatter_idx == 1 and gather_idx == 2: + bs, seqlen, shard_hc, hs = input.shape + hc = shard_hc * seq_world_size + shard_seqlen = seqlen // seq_world_size + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(hc, shard_seqlen, bs, hs) + + # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs) + output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs) + return output + else: + raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2") + + + @staticmethod + def split_qkv_by_head(query, key, value, sp_ulysses_degree, loop_time): + """Split Q/K/V along head dim into chunks for insertcomm / blockattn.""" + _, _, head_count, _ = query.shape + if head_count % sp_ulysses_degree != 0: + raise ValueError( + f"head_count must be divisible by ulysses world size, " + f"got head_count={head_count}, sp_ulysses_degree={sp_ulysses_degree}" + ) + heads_per_rank = head_count // sp_ulysses_degree + if heads_per_rank % loop_time != 0: + raise ValueError( + f"heads_per_rank must be divisible by loop_time={loop_time}, " + f"got heads_per_rank={heads_per_rank}" + ) + global_chunk_heads = heads_per_rank // loop_time * sp_ulysses_degree + return ( + query.split(global_chunk_heads, dim=2), + key.split(global_chunk_heads, dim=2), + value.split(global_chunk_heads, dim=2), + ) + + + @torch.compiler.disable + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + """forward + + Arguments: + query (torch.Tensor): query input to the layer + key (torch.Tensor): key input to the layer + value (torch.Tensor): value input to the layer + + Returns: + * output (torch.Tensor): context output + """ + # Check input shapes + assert query.dim() == 4 and key.dim() == 4 and value.dim() == 4, "Expected 4D tensors" + + forward_context: ForwardContext = get_forward_context() + attn_metadata = forward_context.attn_metadata + + attn_kwargs = {"attn_metadata": attn_metadata} + attn_kwargs.update(kwargs) + + output = None + + if self.fa_alltoall_cut <= 1 and self.fa_alltoall_overlap <= 1: + # baseline: both 0 / both 1 / a single 1 all mean "not enabled" (1 chunk = no split) + # 3 X (bs, seq_len/N, head_cnt, head_size) -> 3 X (bs, seq_len, head_cnt/N, head_size) + # scatter 2, gather 1 + query_layer = SeqAllToAll4D.apply( + self.ulysses_pg, query, self.scatter_idx, self.gather_idx + ) + key_layer = SeqAllToAll4D.apply( + self.ulysses_pg, key, self.scatter_idx, self.gather_idx + ) + value_layer = SeqAllToAll4D.apply( + self.ulysses_pg, value, self.scatter_idx, self.gather_idx + ) + + out = self._run_attention(query_layer, key_layer, value_layer, **attn_kwargs) + # (bs, seq_len, head_cnt/N, head_size) -> (bs, seq_len/N, head_cnt, head_size) + # scatter 1, gather 2 + output = SeqAllToAll4D.apply( + self.ulysses_pg, out, self.gather_idx, self.scatter_idx + ) + elif self.fa_alltoall_cut > 1: # 0415 fa_alltoall_cut + # Split heads into chunks (loop_time = fa_alltoall_cut), full Ulysses round-trip per chunk. + q_chunks, k_chunks, v_chunks = self.split_qkv_by_head( + query, key, value, self.sp_ulysses_degree, self.fa_head_loop + ) + output_chunks = [] + for q_chunk, k_chunk, v_chunk in zip(q_chunks, k_chunks, v_chunks): + query_layer = SeqAllToAll4D.apply( + self.ulysses_pg, q_chunk, self.scatter_idx, self.gather_idx + ) + key_layer = SeqAllToAll4D.apply( + self.ulysses_pg, k_chunk, self.scatter_idx, self.gather_idx + ) + value_layer = SeqAllToAll4D.apply( + self.ulysses_pg, v_chunk, self.scatter_idx, self.gather_idx + ) + out = self._run_attention(query_layer, key_layer, value_layer, **attn_kwargs) + out = SeqAllToAll4D.apply( + self.ulysses_pg, out, self.gather_idx, self.scatter_idx + ) + output_chunks.append(out) + output = torch.cat(output_chunks, dim=2) + elif self.fa_alltoall_overlap > 1 : # 0415 fa_alltoall_overlap + # B, S/sp, N/tp, D + # Refresh the current stream here: __init__ runs at model-build time and may capture a + # different stream than the one forward actually executes on (e.g. under a stream context, + # pipeline/CFG stream switching, or CUDA-graph capture). Pinning the build-time stream would + # break event/stream synchronization in the overlap pipeline. + self.current_stream = torch.npu.current_stream() + query_layer_list, key_layer_list, value_layer_list = self.split_qkv_by_head( + query, key, value, self.sp_ulysses_degree, self.fa_head_loop + ) + for_loop = len(query_layer_list) + + # scatter 2, gather 1 + output_fa = [] + q_event = torch.npu.Event() + k_event = torch.npu.Event() + v_event = torch.npu.Event() + q_lists, k_lists, v_lists, kv_lists = [], [], [], [] + + for i in range(0, for_loop): + input_q = self.all_to_all_4D_pre(query_layer_list[i], self.scatter_idx, self.gather_idx, + self.ulysses_pg) + q_event.record() + with torch.npu.stream(self.stream2): + self.stream2.wait_event(q_event) + query_layer = torch.empty_like(input_q) + dist.all_to_all_single(query_layer, input_q, group=self.ulysses_pg) + + input_k = self.all_to_all_4D_pre(key_layer_list[i], self.scatter_idx, self.gather_idx, + self.ulysses_pg) + input_v = self.all_to_all_4D_pre(value_layer_list[i], self.scatter_idx, self.gather_idx, + self.ulysses_pg) + v_event.record() + + with torch.npu.stream(self.stream2): + self.stream2.wait_event(v_event) + key_layer = torch.empty_like(input_k) + dist.all_to_all_single(key_layer, input_k, group=self.ulysses_pg) + + value_layer = torch.empty_like(input_v) + dist.all_to_all_single(value_layer, input_v, group=self.ulysses_pg) + k_event.record() + + q_lists.append(query_layer) + k_lists.append(key_layer) + v_lists.append(value_layer) + + k_lists[i] = self.all_to_all_4D_after(key_layer_list[i], k_lists[i], self.scatter_idx, + self.gather_idx, self.ulysses_pg) + v_lists[i] = self.all_to_all_4D_after(value_layer_list[i], v_lists[i], self.scatter_idx, + self.gather_idx, self.ulysses_pg) + q_event.record() + with torch.npu.stream(self.stream2): + self.stream2.wait_event(q_event) + self.event[i].record() + q_lists[i] = self.all_to_all_4D_after(query_layer_list[i], q_lists[i], self.scatter_idx, self.gather_idx, self.ulysses_pg) + + + for i in range(0, for_loop): + # fa + self.current_stream.wait_event(self.event[i]) + + out = self._run_attention(q_lists[i], k_lists[i], v_lists[i], **attn_kwargs) + kv_lists.append(out) + input_t = self.all_to_all_4D_pre(out, self.gather_idx, self.scatter_idx, self.ulysses_pg) + q_event.record() + + with torch.npu.stream(self.stream2): + self.stream2.wait_event(q_event) + output = torch.empty_like(input_t) + dist.all_to_all_single(output, input_t, group=self.ulysses_pg) + self.event[i].record() + output_fa.append(output) + + for i in range(for_loop): + self.current_stream.wait_event(self.event[i]) + output_fa[i] = self.all_to_all_4D_after(kv_lists[i], output_fa[i], self.gather_idx, self.scatter_idx, self.ulysses_pg) + output = torch.cat(output_fa, dim=2) + else: + raise RuntimeError( + f"Invalid configuration: fa_alltoall_cut={self.fa_alltoall_cut}, fa_alltoall_overlap={self.fa_alltoall_overlap}" + ) + return output + +_ASCEND_LONGCTX_ATTN: Optional[AscendLongContextAttention] = None + + +def _get_ascend_long_context_attn() -> AscendLongContextAttention: + global _ASCEND_LONGCTX_ATTN + if _ASCEND_LONGCTX_ATTN is None: + _ASCEND_LONGCTX_ATTN = AscendLongContextAttention() + return _ASCEND_LONGCTX_ATTN \ No newline at end of file diff --git a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py index d9d946c..f255b88 100644 --- a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py +++ b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py @@ -29,6 +29,7 @@ from diffsynth_engine.distributed.parallel_state import ( get_tensor_model_parallel_world_size, + is_sp_group_initialized, is_tp_group_initialized, ) from diffsynth_engine.distributed.utils import sequence_parallel_shard, sequence_parallel_unshard @@ -492,11 +493,24 @@ def __init__( # USPAttention for joint attention computation forward_context = get_forward_context() - self.usp_attn = USPAttention( - num_heads=self.heads, - head_size=attention_head_dim, - attn_type=forward_context.attn_type, - ) + + # AscendLongContextAttention calls get_sp_group() unconditionally in __init__, so it can + # only be built when the sequence-parallel group is initialized; otherwise fall back to + # USPAttention, which safely degrades to world_size=1 when SP is not set up. + if is_mindie_sd_available() and is_sp_group_initialized(): + from diffsynth_engine.layers.attention import AscendLongContextAttention + + self.usp_attn = AscendLongContextAttention( + num_heads=self.heads, + head_size=attention_head_dim, + attn_type=forward_context.attn_type, + ) + else: + self.usp_attn = USPAttention( + num_heads=self.heads, + head_size=attention_head_dim, + attn_type=forward_context.attn_type, + ) def forward( self, From 0840a681a506625be601b120ce9c94f8f3f0f10d Mon Sep 17 00:00:00 2001 From: gaoyuanyuanqiqi <2535180690@qq.com> Date: Tue, 18 Aug 2026 16:01:39 +0800 Subject: [PATCH 5/5] feat(npu): long-context attention, fused ops, unified platform abstraction - Add Ascend long-context attention under Ulysses SP with a single shared all-to-all comm stream - Add fused RMSNorm and use it in the Qwen-Image transformer - Unify platform abstraction via current_platform - Centralize Ascend tuning knobs (op_fusion, fa_alltoall_overlap/cut) as platform attributes --- diffsynth_engine/configs/base.py | 12 +- diffsynth_engine/engine.py | 3 +- diffsynth_engine/layers/__init__.py | 5 + diffsynth_engine/layers/attention/layer.py | 20 +-- diffsynth_engine/layers/lora/linear.py | 5 +- diffsynth_engine/layers/transformer_helper.py | 50 ++++++++ .../qwen_image/transformer_qwenimage.py | 12 +- diffsynth_engine/platforms/__init__.py | 55 +++++---- diffsynth_engine/platforms/ascend.py | 21 +++- diffsynth_engine/platforms/base.py | 10 ++ diffsynth_engine/utils/platform.py | 115 +++--------------- diffsynth_engine/worker.py | 4 - 12 files changed, 161 insertions(+), 151 deletions(-) create mode 100644 diffsynth_engine/layers/transformer_helper.py diff --git a/diffsynth_engine/configs/base.py b/diffsynth_engine/configs/base.py index 226dff2..adcb4f8 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.platforms import get_device_type, resolve_platform logger = logging.get_logger(__name__) @@ -27,7 +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 @@ -62,6 +63,7 @@ def __post_init__(self): self.attn_type = str(self.attn_type) init_parallel_config(self) validate_attn_config(self) + init_device_config(self) def init_parallel_config(config: PipelineConfig): @@ -109,3 +111,11 @@ def validate_attn_config(config: PipelineConfig): if config.sp_ring_degree is not None and config.sp_ring_degree > 1: if not attn_backend.supports_ring_attention(): raise ValueError(f"Attention backend {config.attn_type!r} does not support ring attention.") + + +def init_device_config(config: PipelineConfig): + if config.device is None or (isinstance(config.device, str) and config.device.lower() in ("auto", "")): + config.device = get_device_type() + return + # Validate that the explicit device type is registered (fail fast at construction). + resolve_platform(config.device) diff --git a/diffsynth_engine/engine.py b/diffsynth_engine/engine.py index b06a6c3..67f79d9 100644 --- a/diffsynth_engine/engine.py +++ b/diffsynth_engine/engine.py @@ -8,7 +8,7 @@ get_pipeline_class_name, ) from diffsynth_engine.utils import logging -from diffsynth_engine.utils.platform import align_config_device, set_device +from diffsynth_engine.utils.platform import set_device from diffsynth_engine.utils.torch_profiler import TorchProfiler from diffsynth_engine.worker import run_worker_loop @@ -19,7 +19,6 @@ class DiffSynthEngine: @classmethod def from_pretrained(cls, model_path_or_config: str | PipelineConfig, **kwargs): pipeline_config = _resolve_pipeline_config(model_path_or_config) - pipeline_config.device = align_config_device(pipeline_config.device) num_workers = pipeline_config.parallelism master_addr = kwargs.get("master_addr", "localhost") master_port = kwargs.get("master_port", 29500) diff --git a/diffsynth_engine/layers/__init__.py b/diffsynth_engine/layers/__init__.py index e69de29..3205671 100644 --- a/diffsynth_engine/layers/__init__.py +++ b/diffsynth_engine/layers/__init__.py @@ -0,0 +1,5 @@ +from diffsynth_engine.layers.transformer_helper import RMSNorm + +__all__ = [ + "RMSNorm", +] diff --git a/diffsynth_engine/layers/attention/layer.py b/diffsynth_engine/layers/attention/layer.py index 3896039..c9d6a68 100644 --- a/diffsynth_engine/layers/attention/layer.py +++ b/diffsynth_engine/layers/attention/layer.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 -import os from torch import distributed as dist import torch import torch.nn as nn @@ -174,6 +173,10 @@ def forward( from typing import Optional class AscendLongContextAttention(nn.Module): + # Single dedicated communication stream shared by all instances (one per transformer + # block, e.g. 60 in Qwen-Image), so only one `stream2` is allocated per device. + _shared_comm_stream = None + def __init__( self, num_heads: int = 24, @@ -188,6 +191,8 @@ def __init__( **extra_impl_args, ) -> None: super().__init__() + from diffsynth_engine.platforms import AscendPlatform + if num_kv_heads is None: num_kv_heads = num_heads @@ -202,8 +207,8 @@ def __init__( self.sp_ring_degree = get_sp_group().ring_world_size - self.fa_alltoall_overlap = int(os.getenv('FA_ALLTOALL_OVERLAP', 1)) - self.fa_alltoall_cut = int(os.getenv('FA_ALLTOALL_CUT', 1)) + self.fa_alltoall_overlap = AscendPlatform.fa_alltoall_overlap + self.fa_alltoall_cut = AscendPlatform.fa_alltoall_cut if fa_head_loop is not None: self.fa_head_loop = fa_head_loop elif self.fa_alltoall_cut > 1: @@ -214,8 +219,9 @@ def __init__( self.fa_head_loop = self.num_heads // self.sp_ulysses_degree if self.fa_alltoall_overlap > 1 and self.fa_alltoall_cut <= 1: - self.current_stream = torch.npu.current_stream() - self.stream2 = torch.npu.Stream() + if AscendLongContextAttention._shared_comm_stream is None: + AscendLongContextAttention._shared_comm_stream = torch.npu.Stream() + self.stream2 = AscendLongContextAttention._shared_comm_stream self.event = [] for i in range(self.fa_head_loop): self.event.append(torch.npu.Event()) @@ -392,7 +398,7 @@ def forward( output = SeqAllToAll4D.apply( self.ulysses_pg, out, self.gather_idx, self.scatter_idx ) - elif self.fa_alltoall_cut > 1: # 0415 fa_alltoall_cut + elif self.fa_alltoall_cut > 1: # fa_alltoall_cut # Split heads into chunks (loop_time = fa_alltoall_cut), full Ulysses round-trip per chunk. q_chunks, k_chunks, v_chunks = self.split_qkv_by_head( query, key, value, self.sp_ulysses_degree, self.fa_head_loop @@ -414,7 +420,7 @@ def forward( ) output_chunks.append(out) output = torch.cat(output_chunks, dim=2) - elif self.fa_alltoall_overlap > 1 : # 0415 fa_alltoall_overlap + elif self.fa_alltoall_overlap > 1 : # fa_alltoall_overlap # B, S/sp, N/tp, D # Refresh the current stream here: __init__ runs at model-build time and may capture a # different stream than the one forward actually executes on (e.g. under a stream context, diff --git a/diffsynth_engine/layers/lora/linear.py b/diffsynth_engine/layers/lora/linear.py index 17d633b..e3cfbaf 100644 --- a/diffsynth_engine/layers/lora/linear.py +++ b/diffsynth_engine/layers/lora/linear.py @@ -2,6 +2,7 @@ import torch.nn as nn from diffsynth_engine.utils import logging +from diffsynth_engine.utils.platform import pin_memory logger = logging.get_logger(__name__) @@ -180,9 +181,7 @@ def _save_original_weight(self): if self._original_weight is not None: return weight = self.base_layer.weight.data - self._original_weight = weight.detach().cpu().clone() - if not torch.backends.mps.is_available(): - self._original_weight = self._original_weight.pin_memory() + self._original_weight = pin_memory(weight.detach().cpu().clone()) def merge_loras(self, chunked: bool = False, high_precision: bool = True) -> list[str]: """Merge active LoRA models into the wrapped base layer weight. diff --git a/diffsynth_engine/layers/transformer_helper.py b/diffsynth_engine/layers/transformer_helper.py new file mode 100644 index 0000000..d899e11 --- /dev/null +++ b/diffsynth_engine/layers/transformer_helper.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reusable transformer helper layers (fused RMSNorm for NPU, manual fallback elsewhere).""" + +import torch +import torch.nn as nn + +from diffsynth_engine.utils.platform import current_platform, is_mindie_sd_available + + +class RMSNorm(nn.Module): + """RMSNorm over the last dim. + + API-compatible with `diffusers.models.normalization.RMSNorm` + (dim, eps, elementwise_affine), so existing checkpoints (weight key) load + unchanged. On Ascend with `current_platform.op_fusion` enabled the norm is + fused into a single `torch_npu.npu_rms_norm` op; otherwise it falls back to the + reference fp32 math so numerics match diffusers exactly. + """ + + def __init__(self, dim, eps=1e-6, elementwise_affine=True): + super().__init__() + self.dim = dim + self.eps = eps + self.elementwise_affine = elementwise_affine + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim)) + else: + self.register_parameter("weight", None) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + if ( + current_platform.op_fusion + and is_mindie_sd_available() + and self.elementwise_affine + and x.device.type == "npu" + ): + import torch_npu + + return torch_npu.npu_rms_norm(x, self.weight, self.eps)[0] + + output = self._norm(x.float()).type_as(x) + if self.weight is not None: + output = output * self.weight + return output + + def extra_repr(self) -> str: + return f"dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}" \ No newline at end of file diff --git a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py index f255b88..f089941 100644 --- a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py +++ b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py @@ -15,7 +15,6 @@ # limitations under the License. import functools -import os from math import prod from typing import Any, Dict, List, Optional, Tuple, Union @@ -25,7 +24,7 @@ from diffusers.configuration_utils import register_to_config from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.modeling_outputs import Transformer2DModelOutput -from diffusers.models.normalization import AdaLayerNormContinuous, RMSNorm +from diffusers.models.normalization import AdaLayerNormContinuous from diffsynth_engine.distributed.parallel_state import ( get_tensor_model_parallel_world_size, @@ -34,16 +33,15 @@ ) from diffsynth_engine.distributed.utils import sequence_parallel_shard, sequence_parallel_unshard from diffsynth_engine.forward_context import get_forward_context +from diffsynth_engine.layers import RMSNorm from diffsynth_engine.layers.attention import USPAttention from diffsynth_engine.layers.tensor_parallel import ColumnParallelLinear, RowParallelLinear, TPFeedForward from diffsynth_engine.models.base import DiffusionModel from diffsynth_engine.utils import logging -from diffsynth_engine.utils.platform import is_mindie_sd_available +from diffsynth_engine.utils.platform import current_platform, is_mindie_sd_available logger = logging.get_logger(__name__) -USE_MINDIESD_FUSE = os.environ.get("USE_MINDIESD_FUSE", "0") == "1" - def apply_rotary_emb_qwen( x: torch.Tensor, @@ -86,7 +84,7 @@ def apply_rotary_emb_qwen( return out else: - if USE_MINDIESD_FUSE and is_mindie_sd_available() and x.device.type == "npu": + if current_platform.op_fusion and is_mindie_sd_available() and x.device.type == "npu": from mindiesd import rotary_position_embedding # Cache expanded cos/sin on the tensor object itself. @@ -665,7 +663,7 @@ def _modulate(self, x, mod_params, index=None): def _norm_modulate(self, norm: nn.LayerNorm, x: torch.Tensor, mod_params, index=None): """LayerNorm + Ada modulate. Fuses via mindiesd.layernorm_scale_shift when enabled.""" - if USE_MINDIESD_FUSE and is_mindie_sd_available() and x.device.type == "npu": + if current_platform.op_fusion and is_mindie_sd_available() and x.device.type == "npu": from mindiesd import layernorm_scale_shift shift_result, scale_result, gate_result = self._split_mod_params(mod_params, index) diff --git a/diffsynth_engine/platforms/__init__.py b/diffsynth_engine/platforms/__init__.py index f8521b0..affd1d1 100644 --- a/diffsynth_engine/platforms/__init__.py +++ b/diffsynth_engine/platforms/__init__.py @@ -1,6 +1,5 @@ from __future__ import annotations -import platform as host_platform from functools import lru_cache from typing import Type @@ -19,6 +18,14 @@ class CPUPlatform(PlatformBackend): name = "cpu" device_type = "cpu" + @classmethod + def distributed_backend(cls) -> str: + raise NotImplementedError("Unsupported device type") + + @classmethod + def pin_memory(cls, tensor: torch.Tensor) -> torch.Tensor: + return tensor + class CUDAPlatform(PlatformBackend): name = "cuda" @@ -36,6 +43,10 @@ def set_device(cls, index: int | str | torch.device) -> None: def device_count(cls) -> int: return torch.cuda.device_count() + @classmethod + def get_device(cls, local_rank: int) -> torch.device: + return torch.device(cls.device_type, local_rank) + @classmethod def synchronize(cls) -> None: torch.cuda.synchronize() @@ -52,6 +63,15 @@ def distributed_backend(cls) -> str: class ROCmPlatform(CUDAPlatform): name = "rocm" + @classmethod + def fp8_dtype(cls) -> torch.dtype: + if not cls.is_available(): + return torch.float8_e4m3fn + properties = torch.cuda.get_device_properties(0) + if "gfx94" in properties.gcnArchName: + return torch.float8_e4m3fnuz + return torch.float8_e4m3fn + class MPSPlatform(PlatformBackend): name = "mps" @@ -69,6 +89,10 @@ def synchronize(cls) -> None: def empty_cache(cls) -> None: torch.mps.empty_cache() + @classmethod + def pin_memory(cls, tensor: torch.Tensor) -> torch.Tensor: + return tensor + _PLATFORM_REGISTRY: dict[str, Type[PlatformBackend]] = { "cpu": CPUPlatform, @@ -96,8 +120,7 @@ def auto_detect_device() -> str: return "cpu" -def parse_device_type(device: str | torch.device | None = None) -> str: - """Parse a device spec, or auto-detect when ``device`` is None/auto.""" +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): @@ -106,7 +129,7 @@ def parse_device_type(device: str | torch.device | None = None) -> str: def resolve_platform(device: str | torch.device) -> Type[PlatformBackend]: - device_type = parse_device_type(device) + device_type = get_device_type(device) try: return _PLATFORM_REGISTRY[device_type] except KeyError as exc: @@ -114,23 +137,8 @@ def resolve_platform(device: str | torch.device) -> Type[PlatformBackend]: 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(parse_device_type(device)) - return platform_cls.pin_memory(tensor) +# The platform backend for the auto-detected process accelerator, resolved once +current_platform: Type[PlatformBackend] = resolve_platform(get_device_type()) __all__ = [ @@ -142,9 +150,8 @@ def pin_memory( "PlatformCapabilities", "ROCmPlatform", "auto_detect_device", - "get_preferred_fp8_dtype", - "parse_device_type", - "pin_memory", + "current_platform", + "get_device_type", "probe_ascend_capabilities", "probe_ascend_feature", "register_platform", diff --git a/diffsynth_engine/platforms/ascend.py b/diffsynth_engine/platforms/ascend.py index 5e37928..00fccc5 100644 --- a/diffsynth_engine/platforms/ascend.py +++ b/diffsynth_engine/platforms/ascend.py @@ -1,12 +1,16 @@ from __future__ import annotations import importlib +import os from functools import lru_cache from typing import Any import torch from .base import PlatformBackend, PlatformCapabilities +from diffsynth_engine.utils import logging + +logger = logging.get_logger(__name__) def _import_torch_npu(): @@ -169,6 +173,11 @@ class AscendPlatform(PlatformBackend): name = "ascend" device_type = "npu" + # Ascend-specific tuning knobs (seeded from environment for now). + op_fusion = os.environ.get("USE_MINDIESD_FUSE", "False").lower() == "true" + fa_alltoall_overlap = int(os.environ.get("FA_ALLTOALL_OVERLAP", 1)) + fa_alltoall_cut = int(os.environ.get("FA_ALLTOALL_CUT", 1)) + @classmethod def is_available(cls) -> bool: return probe_ascend_feature("device") @@ -188,6 +197,10 @@ def device_count(cls) -> int: _import_torch_npu() return torch.npu.device_count() + @classmethod + def get_device(cls, local_rank: int) -> torch.device: + return torch.device(cls.device_type, local_rank) + @classmethod def synchronize(cls) -> None: torch_npu = _import_torch_npu() @@ -213,12 +226,12 @@ def distributed_backend(cls) -> str: return "hccl" @classmethod - def compile_backend(cls): + def compile_backend(cls) -> Any | None: if not cls.supports("mindie_compile"): - raise RuntimeError( - "MindIE-SD compilation was requested, but MindieSDBackend is unavailable " - "in the installed MindIE-SD package." + logger.warning( + "MindIE-SD compile backend is unavailable; falling back to default torch.compile backend" ) + return None from mindiesd.compilation import CompilationConfig, MindieSDBackend CompilationConfig.fusion_patterns.enable_fast_gelu = False diff --git a/diffsynth_engine/platforms/base.py b/diffsynth_engine/platforms/base.py index cce56eb..a72f559 100644 --- a/diffsynth_engine/platforms/base.py +++ b/diffsynth_engine/platforms/base.py @@ -19,6 +19,8 @@ class PlatformBackend(ABC): name = "unknown" device_type = "cpu" + op_fusion = False + @classmethod def is_available(cls) -> bool: return True @@ -27,6 +29,10 @@ def is_available(cls) -> bool: def normalize_device(cls, device: str | torch.device) -> torch.device: return torch.device(device) + @classmethod + def get_device(cls, local_rank: int) -> torch.device: + return torch.device(cls.device_type) + @classmethod def set_device(cls, index: int | str | torch.device) -> None: return None @@ -35,6 +41,10 @@ def set_device(cls, index: int | str | torch.device) -> None: def device_count(cls) -> int: return 1 + @classmethod + def fp8_dtype(cls) -> torch.dtype: + return torch.float8_e4m3fn + @classmethod def synchronize(cls) -> None: return None diff --git a/diffsynth_engine/utils/platform.py b/diffsynth_engine/utils/platform.py index 2ff4a95..f59d30d 100644 --- a/diffsynth_engine/utils/platform.py +++ b/diffsynth_engine/utils/platform.py @@ -1,132 +1,49 @@ import torch -from diffsynth_engine.utils import logging -logger = logging.get_logger(__name__) - - -def _is_cuda() -> bool: - # Historical: torch build has CUDA, not necessarily a visible GPU. - return torch.version.cuda is not None - - -def _is_rocm() -> bool: - return torch.version.hip is not None - - -def _is_mps() -> bool: - return torch.backends.mps.is_available() - - -def _active_platform(): - """Resolve PlatformBackend for the process-preferred accelerator.""" - from diffsynth_engine.platforms import resolve_platform - - return resolve_platform(get_device_type()) +from diffsynth_engine.platforms import ( + AscendPlatform, + current_platform, +) def is_npu_available() -> bool: - from diffsynth_engine.platforms import AscendPlatform - return AscendPlatform.is_available() def is_mindie_sd_available() -> bool: - from diffsynth_engine.platforms import AscendPlatform - return AscendPlatform.supports("mindie") -def get_device(local_rank: int) -> torch.device: - if _is_cuda() or _is_rocm(): - return torch.device("cuda", local_rank) - if is_npu_available(): - return torch.device("npu", local_rank) - if _is_mps(): - return torch.device("mps") - return torch.device("cpu") - - def get_device_type() -> str: - """Preferred accelerator for this process (no-arg, v1 public API). + return current_platform.device_type - Priority matches historical utils behavior: cuda/rocm build > npu > mps > cpu. - Differs from ``platforms.auto_detect_device`` which uses ``is_available()``. - """ - if _is_cuda() or _is_rocm(): - return "cuda" - if is_npu_available(): - return "npu" - if _is_mps(): - return "mps" - return "cpu" + +def get_device(local_rank: int) -> torch.device: + return current_platform.get_device(local_rank) def get_torch_distributed_backend() -> str: - device_type = get_device_type() - if device_type == "cpu": - raise NotImplementedError("Unsupported device type") - return _active_platform().distributed_backend() + return current_platform.distributed_backend() def device_count() -> int: - return _active_platform().device_count() + return current_platform.device_count() def set_device(index: int | str | torch.device) -> None: - """Bind the current process to a local device (cuda or npu).""" - _active_platform().set_device(index) - - -def align_config_device(config_device: str | torch.device, target_type: str | None = None) -> str: - """Rewrite historical CUDA placeholder to NPU when Ascend is the active accelerator. - - Leaves other placeholders alone (e.g. default ``cuda`` on a CPU laptop). - Explicit ``npu`` on a non-NPU machine raises. - """ - if target_type is None: - target_type = get_device_type() - device_str = str(config_device) - current_type = device_str.split(":", 1)[0].lower() - if current_type == target_type: - return device_str - if target_type == "npu" and current_type == "cuda": - return "npu" - if current_type == "npu" and target_type != "npu": - raise RuntimeError( - f"config.device={config_device!r} does not match available device_type={target_type!r}" - ) - return device_str - - -def bind_rank_device(config_device: str | torch.device, local_rank: int) -> str: - """Worker-only: bind config.device to this rank's local device (e.g. npu:0).""" - device_type = str(config_device).split(":", 1)[0].lower() - if device_type in ("cpu", "mps"): - return str(config_device) - return f"{device_type}:{local_rank}" - + current_platform.set_device(index) -def get_compile_kwargs() -> dict: - """Return kwargs for ``nn.Module.compile`` / ``torch.compile``. - On Ascend with MindIE compile available, injects MindieSDBackend. - Otherwise returns ``{}`` so the default inductor path is used. - """ - if not is_npu_available(): - return {} +def pin_memory(tensor: torch.Tensor) -> torch.Tensor: + return current_platform.pin_memory(tensor) - from diffsynth_engine.platforms import AscendPlatform - if not AscendPlatform.supports("mindie_compile"): - logger.warning( - "MindIE-SD compile backend is unavailable; falling back to default torch.compile backend" - ) - return {} - return AscendPlatform.compile_kwargs() +def get_compile_kwargs() -> dict: + return current_platform.compile_kwargs() -DTYPE_FP8 = torch.float8_e4m3fnuz if _is_rocm() else torch.float8_e4m3fn +DTYPE_FP8 = current_platform.fp8_dtype() DTYPE_MAP: dict[str, torch.dtype] = { # Integer dtypes diff --git a/diffsynth_engine/worker.py b/diffsynth_engine/worker.py index 53085bf..4dab635 100644 --- a/diffsynth_engine/worker.py +++ b/diffsynth_engine/worker.py @@ -11,7 +11,6 @@ ) from diffsynth_engine.registry import get_pipeline_class from diffsynth_engine.utils import logging -from diffsynth_engine.utils.platform import bind_rank_device from diffsynth_engine.utils.torch_profiler import TorchProfiler logger = logging.get_logger(__name__) @@ -40,9 +39,6 @@ def __init__( os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) - # Bind config.device to this rank's local device before HCCL init / model load. - self.pipeline_config.device = bind_rank_device(self.pipeline_config.device, local_rank) - init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank) cfg_degree = 2 if pipeline_config.use_cfg_parallel else 1