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 e9b5fdd0..35be7453 100644 --- a/diffsynth_engine/layers/attention/backends/abstract.py +++ b/diffsynth_engine/layers/attention/backends/abstract.py @@ -26,6 +26,7 @@ class AttentionType(str, enum.Enum): SAGE2 = "sage2" SAGE3 = "sage3" SPARGE = "sparge" + MINDIE = "mindie" def __str__(self) -> str: return self.value diff --git a/diffsynth_engine/layers/attention/backends/mindie_attn.py b/diffsynth_engine/layers/attention/backends/mindie_attn.py new file mode 100644 index 00000000..35682197 --- /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 73509cbe..ffa720bf 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 00000000..f8521b03 --- /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 00000000..5e37928f --- /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 00000000..cce56ebe --- /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 b6b9f4e5..84f7fa0f 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 575f5738..2ff4a955 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 a2f91d77..9f9b01e5 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 13bf802c..b06a6c3d 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 1bde4430..cad43dc4 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 ac3e00fb..53085bfb 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 895c2e31..d9d946c4 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 ffc2b45b..27ac9453 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 cad43dc4..3896039a 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 d9d946c4..f255b887 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 90b68dc32553244ab87c6ce9ca09f96d37b5acbf Mon Sep 17 00:00:00 2001 From: root Date: Fri, 14 Aug 2026 11:28:53 +0800 Subject: [PATCH 5/5] fix: npu all2all comm stream 60 -> 1 --- diffsynth_engine/layers/attention/layer.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/diffsynth_engine/layers/attention/layer.py b/diffsynth_engine/layers/attention/layer.py index 3896039a..f57c5f4b 100644 --- a/diffsynth_engine/layers/attention/layer.py +++ b/diffsynth_engine/layers/attention/layer.py @@ -174,6 +174,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, @@ -214,8 +218,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())