From e808792c554905a8263d3e1de0e9b1cd6982c8e4 Mon Sep 17 00:00:00 2001 From: "fanxingpeng.fxp" Date: Mon, 20 Jul 2026 03:15:57 +0000 Subject: [PATCH] feat: add extensible Ascend runtime support --- README.md | 3 +- diffsynth_engine/__init__.py | 2 + diffsynth_engine/configs/__init__.py | 2 + diffsynth_engine/configs/pipeline.py | 17 + diffsynth_engine/models/basic/attention.py | 375 +++++++++++------- .../models/basic/attention_backend.py | 95 +++++ .../models/qwen_image/qwen_image_dit.py | 21 +- .../qwen_image/qwen_image_dit_fbcache.py | 188 ++++++--- diffsynth_engine/pipelines/base.py | 10 +- diffsynth_engine/pipelines/qwen_image.py | 33 +- diffsynth_engine/platforms/__init__.py | 129 ++++++ diffsynth_engine/platforms/ascend.py | 293 ++++++++++++++ diffsynth_engine/platforms/ascend_qwen.py | 142 +++++++ diffsynth_engine/platforms/base.py | 71 ++++ diffsynth_engine/platforms/qwen.py | 49 +++ diffsynth_engine/platforms/runtime.py | 92 +++++ diffsynth_engine/utils/offload.py | 10 +- diffsynth_engine/utils/platform.py | 58 ++- docs/ascend.md | 148 +++++++ examples/qwen_image_ascend.py | 23 ++ tests/test_models/basic/__init__.py | 0 .../basic/test_attention_backend.py | 126 ++++++ .../test_qwen_attention_processor.py | 63 +++ .../qwen_image/test_qwen_fbcache.py | 38 ++ tests/test_platforms/__init__.py | 0 .../test_platforms/test_ascend_integration.py | 95 +++++ tests/test_platforms/test_platforms.py | 158 ++++++++ tests/test_platforms/test_runtime_adapters.py | 258 ++++++++++++ 28 files changed, 2271 insertions(+), 228 deletions(-) create mode 100644 diffsynth_engine/models/basic/attention_backend.py create mode 100644 diffsynth_engine/platforms/__init__.py create mode 100644 diffsynth_engine/platforms/ascend.py create mode 100644 diffsynth_engine/platforms/ascend_qwen.py create mode 100644 diffsynth_engine/platforms/base.py create mode 100644 diffsynth_engine/platforms/qwen.py create mode 100644 diffsynth_engine/platforms/runtime.py create mode 100644 docs/ascend.md create mode 100644 examples/qwen_image_ascend.py create mode 100644 tests/test_models/basic/__init__.py create mode 100644 tests/test_models/basic/test_attention_backend.py create mode 100644 tests/test_models/qwen_image/test_qwen_attention_processor.py create mode 100644 tests/test_models/qwen_image/test_qwen_fbcache.py create mode 100644 tests/test_platforms/__init__.py create mode 100644 tests/test_platforms/test_ascend_integration.py create mode 100644 tests/test_platforms/test_platforms.py create mode 100644 tests/test_platforms/test_runtime_adapters.py diff --git a/README.md b/README.md index c797b6ba..66139cf9 100644 --- a/README.md +++ b/README.md @@ -75,7 +75,8 @@ image = pipe(prompt="a girl, qipao") image.save("image.png") ``` -For more details, please refer to our tutorials ([English](./docs/tutorial.md), [中文](./docs/tutorial_zh.md)). +For more details, please refer to our tutorials ([English](./docs/tutorial.md), [中文](./docs/tutorial_zh.md)) +and the [Ascend A5 / Ascend 950 guide](./docs/ascend.md). ## Showcase diff --git a/diffsynth_engine/__init__.py b/diffsynth_engine/__init__.py index f11688d6..690d4b01 100644 --- a/diffsynth_engine/__init__.py +++ b/diffsynth_engine/__init__.py @@ -18,6 +18,7 @@ AttnImpl, SpargeAttentionParams, VideoSparseAttentionParams, + QuantizationConfig, LoraConfig, ControlNetParams, ControlType, @@ -73,6 +74,7 @@ "AttnImpl", "SpargeAttentionParams", "VideoSparseAttentionParams", + "QuantizationConfig", "LoraConfig", "ControlNetParams", "ControlType", diff --git a/diffsynth_engine/configs/__init__.py b/diffsynth_engine/configs/__init__.py index 4134173c..f809bc68 100644 --- a/diffsynth_engine/configs/__init__.py +++ b/diffsynth_engine/configs/__init__.py @@ -24,6 +24,7 @@ AttnImpl, SpargeAttentionParams, VideoSparseAttentionParams, + QuantizationConfig, LoraConfig, ) from .controlnet import ( @@ -59,6 +60,7 @@ "AttnImpl", "SpargeAttentionParams", "VideoSparseAttentionParams", + "QuantizationConfig", "LoraConfig", "ControlType", "ControlNetParams", diff --git a/diffsynth_engine/configs/pipeline.py b/diffsynth_engine/configs/pipeline.py index 37da60a1..32308d34 100644 --- a/diffsynth_engine/configs/pipeline.py +++ b/diffsynth_engine/configs/pipeline.py @@ -35,6 +35,7 @@ class AttnImpl(Enum): SAGE = "sage" # Sage Attention SPARGE = "sparge" # Sparge Attention VSA = "vsa" # Video Sparse Attention + MINDIE = "mindie" # MindIE-SD attention for Ascend NPU @dataclass @@ -56,12 +57,28 @@ class AttentionConfig: attn_params: Optional[SpargeAttentionParams | VideoSparseAttentionParams] = None +@dataclass +class QuantizationConfig: + backend: Literal["auto", "native", "nunchaku", "mindie"] = "auto" + linear: Literal["none", "fp8", "int4"] = "none" + attention: Literal["none", "fp8"] = "none" + + def __post_init__(self): + if self.backend not in {"auto", "native", "nunchaku", "mindie"}: + raise ValueError(f"Unsupported quantization backend: {self.backend}") + if self.linear not in {"none", "fp8", "int4"}: + raise ValueError(f"Unsupported linear quantization mode: {self.linear}") + if self.attention not in {"none", "fp8"}: + raise ValueError(f"Unsupported attention quantization mode: {self.attention}") + + @dataclass class OptimizationConfig: use_fp8_linear: bool = False use_fbcache: bool = False fbcache_relative_l1_threshold: float = 0.05 use_torch_compile: bool = False + quantization: Optional[QuantizationConfig] = None @dataclass diff --git a/diffsynth_engine/models/basic/attention.py b/diffsynth_engine/models/basic/attention.py index 3d3d49cf..d7627c6a 100644 --- a/diffsynth_engine/models/basic/attention.py +++ b/diffsynth_engine/models/basic/attention.py @@ -17,6 +17,12 @@ AITER_AVAILABLE, ) from diffsynth_engine.utils.platform import DTYPE_FP8 +from diffsynth_engine.models.basic.attention_backend import ( + AttentionBackend, + AttentionRequest, + register_attention_backend, + resolve_attention_backend, +) FA3_MAX_HEADDIM = 256 @@ -122,6 +128,228 @@ def eager_attn(q, k, v, attn_mask=None, scale=None): return out.transpose(1, 2) +def _supports_head_dim_256(request: AttentionRequest, *, allow_mask: bool) -> tuple[bool, str | None]: + if request.q.shape[-1] > FA3_MAX_HEADDIM: + return False, f"head_dim={request.q.shape[-1]} exceeds the supported maximum {FA3_MAX_HEADDIM}" + if not allow_mask and request.attn_mask is not None: + return False, "attention masks are not supported" + return True, None + + +def _run_fa4(request: AttentionRequest) -> torch.Tensor: + output = flash_attn4(request.q, request.k, request.v, softmax_scale=request.scale) + return output[0] if isinstance(output, tuple) else output + + +def _run_fa3(request: AttentionRequest, *, fp8: bool = False) -> torch.Tensor: + if not fp8: + return flash_attn3(request.q, request.k, request.v, softmax_scale=request.scale) + origin_dtype = request.q.dtype + output = flash_attn3( + request.q.to(dtype=DTYPE_FP8), + request.k.to(dtype=DTYPE_FP8), + request.v.to(dtype=DTYPE_FP8), + softmax_scale=request.scale, + ) + return output.to(dtype=origin_dtype) + + +def _run_aiter(request: AttentionRequest, *, fp8: bool = False) -> torch.Tensor: + if not fp8: + return aiter_flash_attn(request.q, request.k, request.v, softmax_scale=request.scale) + origin_dtype = request.q.dtype + output = aiter_flash_attn_fp8( + request.q.to(dtype=DTYPE_FP8), + request.k.to(dtype=DTYPE_FP8), + request.v.to(dtype=DTYPE_FP8), + softmax_scale=request.scale, + ) + return output.to(dtype=origin_dtype) + + +def _run_mindie(request: AttentionRequest) -> torch.Tensor: + from mindiesd.layers.flash_attn.attention_forward import attention_forward + + return attention_forward( + query=request.q, + key=request.k, + value=request.v, + attn_mask=request.attn_mask, + scale=request.scale, + fused=True, + head_first=False, + ) + + +def _mindie_available() -> bool: + from diffsynth_engine.platforms import AscendPlatform + + return AscendPlatform.supports("mindie_attention") + + +def _mindie_supports(request: AttentionRequest) -> tuple[bool, str | None]: + if request.q.ndim != 4 or request.k.ndim != 4 or request.v.ndim != 4: + return False, "MindIE attention requires 4D Q/K/V tensors" + return True, None + + +def _register_builtin_attention_backends() -> None: + cuda = frozenset({"cuda"}) + all_devices = None + register_attention_backend( + AttentionBackend( + "fa4", + _run_fa4, + cuda, + 700, + lambda: FLASH_ATTN_4_AVAILABLE, + lambda r: _supports_head_dim_256(r, allow_mask=False), + ) + ) + register_attention_backend( + AttentionBackend( + "fa3", + _run_fa3, + cuda, + 600, + lambda: FLASH_ATTN_3_AVAILABLE, + lambda r: _supports_head_dim_256(r, allow_mask=False), + ) + ) + register_attention_backend( + AttentionBackend( + "fa3_fp8", + lambda r: _run_fa3(r, fp8=True), + cuda, + 0, + lambda: FLASH_ATTN_3_AVAILABLE, + lambda r: _supports_head_dim_256(r, allow_mask=False), + False, + ) + ) + register_attention_backend( + AttentionBackend( + "aiter", + _run_aiter, + cuda, + 500, + lambda: AITER_AVAILABLE, + lambda r: _supports_head_dim_256(r, allow_mask=False), + auto_supports=lambda r: _supports_head_dim_256(r, allow_mask=True), + ) + ) + register_attention_backend( + AttentionBackend( + "aiter_fp8", + lambda r: _run_aiter(r, fp8=True), + cuda, + 0, + lambda: AITER_AVAILABLE, + lambda r: _supports_head_dim_256(r, allow_mask=False), + False, + ) + ) + register_attention_backend( + AttentionBackend( + "xformers", + lambda r: xformers_attn(r.q, r.k, r.v, r.attn_mask, r.scale), + cuda, + 400, + lambda: XFORMERS_AVAILABLE, + ) + ) + register_attention_backend( + AttentionBackend( + "sdpa", + lambda r: sdpa_attn(r.q, r.k, r.v, r.attn_mask, scale=r.scale), + all_devices, + 300, + lambda: SDPA_AVAILABLE, + ) + ) + register_attention_backend( + AttentionBackend( + "fa2", + lambda r: flash_attn2(r.q, r.k, r.v, softmax_scale=r.scale), + cuda, + 200, + lambda: FLASH_ATTN_2_AVAILABLE, + ) + ) + register_attention_backend( + AttentionBackend("eager", lambda r: eager_attn(r.q, r.k, r.v, r.attn_mask, r.scale), all_devices, 100) + ) + register_attention_backend( + AttentionBackend( + "sage", + lambda r: sage_attn(r.q, r.k, r.v, r.attn_mask, r.scale), + cuda, + 0, + lambda: SAGE_ATTN_AVAILABLE, + auto_select=False, + ) + ) + register_attention_backend( + AttentionBackend( + "sparge", + lambda r: sparge_attn( + r.q, + r.k, + r.v, + attn_mask=r.attn_mask, + scale=r.scale, + smooth_k=r.kwargs.get("smooth_k", True), + simthreshd1=r.kwargs.get("simthreshd1", 0.6), + cdfthreshd=r.kwargs.get("cdfthreshd", 0.98), + pvthreshd=r.kwargs.get("pvthreshd", 50), + ), + cuda, + 0, + lambda: SPARGE_ATTN_AVAILABLE, + auto_select=False, + ) + ) + register_attention_backend( + AttentionBackend( + "vsa", + lambda r: video_sparse_attn( + r.q, + r.k, + r.v, + r.g, + sparsity=r.kwargs.get("sparsity"), + num_tiles=r.kwargs.get("num_tiles"), + total_seq_length=r.kwargs.get("total_seq_length"), + tile_partition_indices=r.kwargs.get("tile_partition_indices"), + reverse_tile_partition_indices=r.kwargs.get("reverse_tile_partition_indices"), + variable_block_sizes=r.kwargs.get("variable_block_sizes"), + non_pad_index=r.kwargs.get("non_pad_index"), + ), + cuda, + 0, + lambda: VIDEO_SPARSE_ATTN_AVAILABLE, + auto_select=False, + ) + ) + register_attention_backend( + AttentionBackend( + "mindie", + _run_mindie, + frozenset({"npu"}), + 800, + _mindie_available, + _mindie_supports, + unavailable_reason=( + "MindIE attention is unavailable; install a compatible MindIE-SD 3.x package " + "or use attn_impl='auto' to allow SDPA/eager fallback" + ), + ) + ) + + +_register_builtin_attention_backends() + + def attention( q: torch.Tensor, k: torch.Tensor, @@ -137,147 +365,9 @@ def attention( k: [B, Lk, Nk, C1] v: [B, Lk, Nk, C2] """ - assert attn_impl in [ - None, - "auto", - "eager", - "fa2", - "fa3", - "fa3_fp8", - "fa4", - "aiter", - "aiter_fp8", - "xformers", - "sdpa", - "sage", - "sparge", - "vsa", - ] - flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM - if attn_impl is None or attn_impl == "auto": - if FLASH_ATTN_4_AVAILABLE: - # FA4 also has the same max-head-256 limitation as FA3 - if flash_attn3_compatible and attn_mask is None: - attn_out = flash_attn4(q, k, v, softmax_scale=scale) - if isinstance(attn_out, tuple): - attn_out = attn_out[0] - return attn_out - else: - if not flash_attn3_compatible: - logger.warning( - f"head_dim={q.shape[-1]}, but flash_attn_4 only supports head dimension at most {FA3_MAX_HEADDIM}, will use fallback attention implementation" - ) - else: - logger.debug( - "flash_attn_4 does not support attention mask, will use fallback attention implementation" - ) - if FLASH_ATTN_3_AVAILABLE: - if flash_attn3_compatible and attn_mask is None: - return flash_attn3(q, k, v, softmax_scale=scale) - else: - if not flash_attn3_compatible: - logger.warning( - f"head_dim={q.shape[-1]}, but flash_attn_3 only supports head dimension at most {FA3_MAX_HEADDIM}, will use fallback attention implementation" - ) - else: - logger.debug( - "flash_attn_3 does not support attention mask, will use fallback attention implementation" - ) - if AITER_AVAILABLE: - if flash_attn3_compatible: - return aiter_flash_attn(q, k, v, softmax_scale=scale) - else: - logger.warning( - f"head_dim={q.shape[-1]}, but aiter_flash_attn only supports head dimension at most {FA3_MAX_HEADDIM}, will use fallback attention implementation" - ) - if XFORMERS_AVAILABLE: - return xformers_attn(q, k, v, attn_mask=attn_mask, scale=scale) - if SDPA_AVAILABLE: - return sdpa_attn(q, k, v, attn_mask=attn_mask, scale=scale) - if FLASH_ATTN_2_AVAILABLE: - return flash_attn2(q, k, v, softmax_scale=scale) - return eager_attn(q, k, v, attn_mask=attn_mask, scale=scale) - else: - if attn_impl == "eager": - return eager_attn(q, k, v, attn_mask=attn_mask, scale=scale) - if attn_impl == "fa3" or attn_impl == "fa3_fp8": - if not flash_attn3_compatible: - raise RuntimeError( - f"head_dim={q.shape[-1]}, but flash_attn_3 only supports head dimension at most {FA3_MAX_HEADDIM}" - ) - if attn_mask is not None: - raise RuntimeError("flash_attn_3 does not support attention mask") - if attn_impl == "fa3": - return flash_attn3(q, k, v, softmax_scale=scale) - else: - origin_dtype = q.dtype - q = q.to(dtype=DTYPE_FP8) - k = k.to(dtype=DTYPE_FP8) - v = v.to(dtype=DTYPE_FP8) - out = flash_attn3(q, k, v, softmax_scale=scale) - return out.to(dtype=origin_dtype) - if attn_impl == "aiter" or attn_impl == "aiter_fp8": - if not flash_attn3_compatible: - raise RuntimeError( - f"head_dim={q.shape[-1]}, but aiter_flash_attn only supports head dimension at most {FA3_MAX_HEADDIM}" - ) - if attn_mask is not None: - raise RuntimeError("aiter_flash_attn does not support attention mask") - if attn_impl == "aiter": - return aiter_flash_attn(q, k, v, softmax_scale=scale) - else: - origin_dtype = q.dtype - q = q.to(dtype=DTYPE_FP8) - k = k.to(dtype=DTYPE_FP8) - v = v.to(dtype=DTYPE_FP8) - out = aiter_flash_attn_fp8(q, k, v, softmax_scale=scale) - return out.to(dtype=origin_dtype) - if attn_impl == "fa4": - if not flash_attn3_compatible: - raise RuntimeError( - f"head_dim={q.shape[-1]}, but flash_attn_4 only supports head dimension at most {FA3_MAX_HEADDIM}" - ) - if attn_mask is not None: - raise RuntimeError("flash_attn_4 does not support attention mask") - attn_out = flash_attn4(q, k, v, softmax_scale=scale) - if isinstance(attn_out, tuple): - attn_out = attn_out[0] - return attn_out - if attn_impl == "fa2": - return flash_attn2(q, k, v, softmax_scale=scale) - if attn_impl == "xformers": - return xformers_attn(q, k, v, attn_mask=attn_mask, scale=scale) - if attn_impl == "sdpa": - return sdpa_attn(q, k, v, attn_mask=attn_mask, scale=scale) - if attn_impl == "sage": - return sage_attn(q, k, v, attn_mask=attn_mask, scale=scale) - if attn_impl == "sparge": - return sparge_attn( - q, - k, - v, - attn_mask=attn_mask, - scale=scale, - smooth_k=kwargs.get("smooth_k", True), - simthreshd1=kwargs.get("simthreshd1", 0.6), - cdfthreshd=kwargs.get("cdfthreshd", 0.98), - pvthreshd=kwargs.get("pvthreshd", 50), - ) - if attn_impl == "vsa": - return video_sparse_attn( - q, - k, - v, - g, - sparsity=kwargs.get("sparsity"), - num_tiles=kwargs.get("num_tiles"), - total_seq_length=kwargs.get("total_seq_length"), - tile_partition_indices=kwargs.get("tile_partition_indices"), - reverse_tile_partition_indices=kwargs.get("reverse_tile_partition_indices"), - variable_block_sizes=kwargs.get("variable_block_sizes"), - non_pad_index=kwargs.get("non_pad_index"), - ) - raise ValueError(f"Invalid attention implementation: {attn_impl}") + request = AttentionRequest(q=q, k=k, v=v, g=g, attn_mask=attn_mask, scale=scale, kwargs=kwargs) + backend = resolve_attention_backend(attn_impl or "auto", request) + return backend.forward(request) class Attention(nn.Module): @@ -354,7 +444,10 @@ def long_context_attention( "sage", "sparge", "vsa", + "mindie", ] + if attn_impl == "mindie": + raise RuntimeError("MindIE long-context attention is not supported in the single-NPU release") assert attn_mask is None, "long context attention does not support attention mask" flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM if attn_impl is None or attn_impl == "auto": diff --git a/diffsynth_engine/models/basic/attention_backend.py b/diffsynth_engine/models/basic/attention_backend.py new file mode 100644 index 00000000..0023002a --- /dev/null +++ b/diffsynth_engine/models/basic/attention_backend.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Optional + +import torch + + +@dataclass(frozen=True) +class AttentionRequest: + q: torch.Tensor + k: torch.Tensor + v: torch.Tensor + g: Optional[torch.Tensor] = None + attn_mask: Optional[torch.Tensor] = None + scale: Optional[float] = None + kwargs: dict[str, Any] = field(default_factory=dict) + + +AvailabilityCheck = Callable[[], bool] +SupportCheck = Callable[[AttentionRequest], tuple[bool, str | None]] +AttentionForward = Callable[[AttentionRequest], torch.Tensor] + + +@dataclass(frozen=True) +class AttentionBackend: + name: str + forward: AttentionForward + devices: frozenset[str] | None = None + priority: int = 0 + is_available: AvailabilityCheck = lambda: True + supports: SupportCheck = lambda request: (True, None) + auto_select: bool = True + unavailable_reason: str | None = None + auto_supports: SupportCheck | None = None + + +_ATTENTION_BACKENDS: dict[str, AttentionBackend] = {} + + +def register_attention_backend(backend: AttentionBackend, *, overwrite: bool = False) -> None: + if backend.name in _ATTENTION_BACKENDS and not overwrite: + raise ValueError(f"Attention backend {backend.name!r} is already registered") + _ATTENTION_BACKENDS[backend.name] = backend + + +def get_attention_backend(name: str) -> AttentionBackend: + try: + return _ATTENTION_BACKENDS[name] + except KeyError as exc: + available = ", ".join(sorted(_ATTENTION_BACKENDS)) + raise ValueError(f"Invalid attention implementation {name!r}. Available backends: {available}") from exc + + +def _check_backend( + backend: AttentionBackend, + request: AttentionRequest, + *, + auto: bool = False, +) -> tuple[bool, str | None]: + device_type = request.q.device.type + if backend.devices is not None and device_type not in backend.devices: + return False, f"backend {backend.name!r} does not support device type {device_type!r}" + if not backend.is_available(): + reason = backend.unavailable_reason or f"backend {backend.name!r} is not available in the current environment" + return False, reason + support_check = backend.auto_supports if auto and backend.auto_supports is not None else backend.supports + return support_check(request) + + +def resolve_attention_backend(name: str, request: AttentionRequest) -> AttentionBackend: + if name not in {"auto", None}: + backend = get_attention_backend(name) + supported, reason = _check_backend(backend, request) + if not supported: + raise RuntimeError(reason or f"Attention backend {name!r} does not support this request") + return backend + + candidates = sorted( + (backend for backend in _ATTENTION_BACKENDS.values() if backend.auto_select), + key=lambda backend: backend.priority, + reverse=True, + ) + failures = [] + for backend in candidates: + supported, reason = _check_backend(backend, request, auto=True) + if supported: + return backend + if reason: + failures.append(reason) + raise RuntimeError("No attention backend supports this request: " + "; ".join(failures)) + + +def list_attention_backends() -> tuple[str, ...]: + return tuple(sorted(_ATTENTION_BACKENDS)) diff --git a/diffsynth_engine/models/qwen_image/qwen_image_dit.py b/diffsynth_engine/models/qwen_image/qwen_image_dit.py index 9b5933c4..eeffbf2d 100644 --- a/diffsynth_engine/models/qwen_image/qwen_image_dit.py +++ b/diffsynth_engine/models/qwen_image/qwen_image_dit.py @@ -189,6 +189,10 @@ def __init__( self.to_out = nn.Linear(dim_a, dim_a, device=device, dtype=dtype) self.to_add_out = nn.Linear(dim_b, dim_b, device=device, dtype=dtype) + self.attention_processor = None + + def set_attention_processor(self, processor) -> None: + self.attention_processor = processor def forward( self, @@ -224,7 +228,22 @@ def forward( joint_v = torch.cat([txt_v, img_v], dim=1) attn_kwargs = attn_kwargs if attn_kwargs is not None else {} - joint_attn_out = attention_ops.attention(joint_q, joint_k, joint_v, attn_mask=attn_mask, **attn_kwargs) + if self.attention_processor is not None: + joint_attn_out = self.attention_processor( + joint_q, + joint_k, + joint_v, + attn_mask=attn_mask, + **attn_kwargs, + ) + else: + joint_attn_out = attention_ops.attention( + joint_q, + joint_k, + joint_v, + attn_mask=attn_mask, + **attn_kwargs, + ) joint_attn_out = rearrange(joint_attn_out, "b s h d -> b s (h d)").to(joint_q.dtype) diff --git a/diffsynth_engine/models/qwen_image/qwen_image_dit_fbcache.py b/diffsynth_engine/models/qwen_image/qwen_image_dit_fbcache.py index 641168b5..d3291392 100644 --- a/diffsynth_engine/models/qwen_image/qwen_image_dit_fbcache.py +++ b/diffsynth_engine/models/qwen_image/qwen_image_dit_fbcache.py @@ -1,10 +1,16 @@ import torch -from typing import Any, Dict, Optional +from math import prod +from typing import Any, Dict, List, Optional from diffsynth_engine.models.qwen_image import QwenImageDiT from diffsynth_engine.utils.gguf import gguf_inference from diffsynth_engine.utils.fp8_linear import fp8_inference -from diffsynth_engine.utils.parallel import cfg_parallel, cfg_parallel_unshard +from diffsynth_engine.utils.parallel import ( + cfg_parallel, + cfg_parallel_unshard, + sequence_parallel, + sequence_parallel_unshard, +) class QwenImageDiTFBCache(QwenImageDiT): @@ -14,11 +20,22 @@ def __init__( device: str = "cuda:0", dtype: torch.dtype = torch.bfloat16, relative_l1_threshold: float = 0.05, + zero_cond_t: bool = False, ): - super().__init__(num_layers=num_layers, device=device, dtype=dtype) + super().__init__(num_layers=num_layers, device=device, dtype=dtype, zero_cond_t=zero_cond_t) self.relative_l1_threshold = relative_l1_threshold self.step_count = 0 self.num_inference_steps = 0 + self._cache_stream = 0 + self._cache_states = [self._new_cache_state()] + + @staticmethod + def _new_cache_state(): + return { + "step_count": 0, + "prev_first_hidden_states_residual": None, + "previous_residual": None, + } def is_relative_l1_below_threshold(self, prev_residual, residual, threshold): if threshold <= 0.0: @@ -32,16 +49,29 @@ def is_relative_l1_below_threshold(self, prev_residual, residual, threshold): diff = mean_diff / mean_prev_residual return diff.item() < threshold - def refresh_cache_status(self, num_inference_steps): + def refresh_cache_status(self, num_inference_steps, num_cache_streams=1): self.step_count = 0 self.num_inference_steps = num_inference_steps + self._cache_stream = 0 + self._cache_states = [self._new_cache_state() for _ in range(num_cache_streams)] + + def set_cache_stream(self, stream_index): + if stream_index < 0 or stream_index >= len(self._cache_states): + raise ValueError(f"Invalid FB-cache stream index: {stream_index}") + self._cache_stream = stream_index + self.step_count = self._cache_states[stream_index]["step_count"] def forward( self, image: torch.Tensor, + edit: torch.Tensor = None, text: torch.Tensor = None, timestep: torch.LongTensor = None, - txt_seq_lens: torch.LongTensor = None, + text_seq_lens: torch.LongTensor = None, + context_latents: Optional[torch.Tensor] = None, + entity_text: Optional[List[torch.Tensor]] = None, + entity_seq_lens: Optional[List[torch.LongTensor]] = None, + entity_masks: Optional[List[torch.Tensor]] = None, attn_kwargs: Optional[Dict[str, Any]] = None, ): h, w = image.shape[-2:] @@ -53,65 +83,115 @@ def forward( cfg_parallel( ( image, - text, + *(edit if edit is not None else ()), timestep, - txt_seq_lens, + text, + text_seq_lens, + *(entity_text if entity_text is not None else ()), + *(entity_seq_lens if entity_seq_lens is not None else ()), + *(entity_masks if entity_masks is not None else ()), + context_latents, ), use_cfg=use_cfg, ), ): + if self.zero_cond_t: + timestep = torch.cat([timestep, timestep * 0], dim=0) + modulate_index = None conditioning = self.time_text_embed(timestep, image.dtype) - video_fhw = (1, h // 2, w // 2) # frame, height, width - max_length = txt_seq_lens.max().item() - image_rotary_emb = self.pos_embed(video_fhw, max_length, image.device) - + video_fhw = [(1, h // 2, w // 2)] + text_seq_len = text_seq_lens.max().item() image = self.patchify(image) - image = self.img_in(image) - text = self.txt_in(self.txt_norm(text[:, :max_length])) - - # first block - original_hidden_states = image - text, image = self.transformer_blocks[0]( - image=image, - text=text, - temb=conditioning, - image_rotary_emb=image_rotary_emb, - attn_kwargs=attn_kwargs, - ) - first_hidden_states_residual = image - original_hidden_states - - if self.step_count == 0 or self.step_count == (self.num_inference_steps - 1): - should_calc = True - else: - skip = self.is_relative_l1_below_threshold( - first_hidden_states_residual, - self.prev_first_hidden_states_residual, - threshold=self.relative_l1_threshold, + image_seq_len = image.shape[1] + if context_latents is not None: + context_latents = context_latents.to(dtype=image.dtype) + context_latents = self.patchify(context_latents) + image = torch.cat([image, context_latents], dim=1) + video_fhw += [(1, h // 2, w // 2)] + if edit is not None: + for img in edit: + img = img.to(dtype=image.dtype) + edit_h, edit_w = img.shape[-2:] + img = self.patchify(img) + image = torch.cat([image, img], dim=1) + video_fhw += [(1, edit_h // 2, edit_w // 2)] + if self.zero_cond_t: + modulate_index = torch.tensor( + [[0] * prod(sample[0]) + [1] * sum([prod(s) for s in sample[1:]]) for sample in [video_fhw]], + device=timestep.device, + dtype=torch.int, ) - should_calc = not skip - self.step_count += 1 - - if not should_calc: - image += self.previous_residual - else: - self.prev_first_hidden_states_residual = first_hidden_states_residual - first_hidden_states = image.clone() - - for block in self.transformer_blocks[1:]: - text, image = block( - image=image, - text=text, - temb=conditioning, - image_rotary_emb=image_rotary_emb, - attn_kwargs=attn_kwargs, - ) + modulate_index = modulate_index.unsqueeze(-1) + rotary_emb = self.pos_embed(video_fhw, text_seq_len, image.device) - previous_residual = image - first_hidden_states - self.previous_residual = previous_residual + image = self.img_in(image) + text = self.txt_in(self.txt_norm(text[:, :text_seq_len])) - image = self.norm_out(image, conditioning) - image = self.proj_out(image) + attn_mask = None + if entity_text is not None: + text, rotary_emb, attn_mask = self.process_entity_masks( + text, + text_seq_lens, + rotary_emb, + video_fhw, + entity_text, + entity_seq_lens, + entity_masks, + image.device, + image.dtype, + ) + img_freqs, txt_freqs = rotary_emb + with sequence_parallel((image, text, img_freqs, txt_freqs, modulate_index), seq_dims=(1, 1, 0, 0, 1)): + cache_state = self._cache_states[self._cache_stream] + rotary_emb = (img_freqs, txt_freqs) + original_hidden_states = image + text, image = self.transformer_blocks[0]( + image=image, + text=text, + temb=conditioning, + rotary_emb=rotary_emb, + attn_mask=attn_mask, + attn_kwargs=attn_kwargs, + modulate_index=modulate_index, + ) + first_hidden_states_residual = image - original_hidden_states + + if cache_state["step_count"] == 0 or cache_state["step_count"] == (self.num_inference_steps - 1): + should_calc = True + else: + skip = self.is_relative_l1_below_threshold( + first_hidden_states_residual, + cache_state["prev_first_hidden_states_residual"], + threshold=self.relative_l1_threshold, + ) + should_calc = not skip + cache_state["step_count"] += 1 + self.step_count = cache_state["step_count"] + + if not should_calc: + image += cache_state["previous_residual"] + else: + cache_state["prev_first_hidden_states_residual"] = first_hidden_states_residual + first_hidden_states = image.clone() + for block in self.transformer_blocks[1:]: + text, image = block( + image=image, + text=text, + temb=conditioning, + rotary_emb=rotary_emb, + attn_mask=attn_mask, + attn_kwargs=attn_kwargs, + modulate_index=modulate_index, + ) + cache_state["previous_residual"] = image - first_hidden_states + + if self.zero_cond_t: + conditioning = conditioning.chunk(2, dim=0)[0] + image = self.norm_out(image, conditioning) + image = self.proj_out(image) + (image,) = sequence_parallel_unshard((image,), seq_dims=(1,), seq_lens=(image_seq_len,)) + image = image[:, :image_seq_len] image = self.unpatchify(image, h, w) (image,) = cfg_parallel_unshard((image,), use_cfg=use_cfg) @@ -125,12 +205,14 @@ def from_state_dict( dtype: torch.dtype, num_layers: int = 60, relative_l1_threshold: float = 0.05, + use_zero_cond_t: bool = False, ): model = cls( device="meta", dtype=dtype, num_layers=num_layers, relative_l1_threshold=relative_l1_threshold, + zero_cond_t=use_zero_cond_t, ) model = model.requires_grad_(False) model.load_state_dict(state_dict, assign=True) diff --git a/diffsynth_engine/pipelines/base.py b/diffsynth_engine/pipelines/base.py index c63abbf2..14f3c61e 100644 --- a/diffsynth_engine/pipelines/base.py +++ b/diffsynth_engine/pipelines/base.py @@ -342,7 +342,7 @@ def _enable_model_cpu_offload(self): for model_name in self.model_names: model = getattr(self, model_name) if model is not None: - self._offload_param_dict[model_name] = offload_model_to_dict(model) + self._offload_param_dict[model_name] = offload_model_to_dict(model, self.device) self.offload_mode = "cpu_offload" def _enable_sequential_cpu_offload(self): @@ -376,8 +376,7 @@ def load_models_to_device(self, load_model_names: List[str] | None = None): if not self.offload_mode: return if self.offload_mode == "sequential_cpu_offload": - # fresh the cuda cache - empty_cache() + empty_cache(self.device) return # offload unnecessary models to cpu @@ -395,8 +394,7 @@ def load_models_to_device(self, load_model_names: List[str] | None = None): ) if model is not None and (p := next(model.parameters(), None)) is not None and p.device.type != self.device: model.to(self.device) - # fresh the cuda cache - empty_cache() + empty_cache(self.device) def model_lifecycle_finish(self, model_names: List[str] | None = None): if not self.offload_to_disk or self.offload_mode is None: @@ -409,7 +407,7 @@ def model_lifecycle_finish(self, model_names: List[str] | None = None): setattr(self, model_name, None) print(f"model {model_name} has been deleted from memory") logger.info(f"model {model_name} has been deleted from memory") - empty_cache() + empty_cache(self.device) def compile(self): raise NotImplementedError(f"{self.__class__.__name__} does not support compile") diff --git a/diffsynth_engine/pipelines/qwen_image.py b/diffsynth_engine/pipelines/qwen_image.py index 342cb70b..85ddd6f6 100644 --- a/diffsynth_engine/pipelines/qwen_image.py +++ b/diffsynth_engine/pipelines/qwen_image.py @@ -36,9 +36,9 @@ ) from diffsynth_engine.utils.parallel import ParallelWrapper from diffsynth_engine.utils import logging -from diffsynth_engine.utils.fp8_linear import enable_fp8_linear from diffsynth_engine.utils.download import fetch_model from diffsynth_engine.utils.flag import NUNCHAKU_AVAILABLE +from diffsynth_engine.platforms.runtime import get_runtime_adapter logger = logging.get_logger(__name__) @@ -145,6 +145,7 @@ def __init__( dtype=config.model_dtype, ) self.config = config + self.runtime_adapter = get_runtime_adapter(config.device, "qwen_image") # qwen image self.prompt_template_encode = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" self.prompt_template_encode_start_idx = 34 @@ -216,12 +217,17 @@ def from_pretrained(cls, model_path_or_config: str | QwenImagePipelineConfig) -> else: config = model_path_or_config + runtime_adapter = get_runtime_adapter(config.device, "qwen_image") + if runtime_adapter.platform.name == "ascend": + runtime_adapter.validate_config(config) + logger.info(f"loading state dict from {config.model_path} ...") model_state_dict = cls.load_model_checkpoint( config.model_path, device="cpu", dtype=config.model_dtype, convert_dtype=False ) config = cls._setup_nunchaku_config(model_state_dict, config) + get_runtime_adapter(config.device, "qwen_image").validate_config(config) # for svd quant model fp4/int4 linear layers, do not convert dtype here if not config.use_nunchaku: @@ -263,6 +269,7 @@ def from_pretrained(cls, model_path_or_config: str | QwenImagePipelineConfig) -> @classmethod def from_state_dict(cls, state_dicts: QwenImageStateDicts, config: QwenImagePipelineConfig) -> "QwenImagePipeline": config = cls._setup_nunchaku_config(state_dicts.model, config) + get_runtime_adapter(config.device, "qwen_image").validate_config(config) if config.parallelism > 1: pipe = ParallelWrapper( @@ -312,6 +319,7 @@ def _from_state_dict(cls, state_dicts: QwenImageStateDicts, config: QwenImagePip device=("cpu" if config.use_fsdp else init_device), dtype=config.model_dtype, relative_l1_threshold=config.fbcache_relative_l1_threshold, + use_zero_cond_t=config.use_zero_cond_t, ) elif config.use_nunchaku: if not NUNCHAKU_AVAILABLE: @@ -338,9 +346,6 @@ def _from_state_dict(cls, state_dicts: QwenImageStateDicts, config: QwenImagePip dtype=config.model_dtype, use_zero_cond_t=config.use_zero_cond_t, ) - if config.use_fp8_linear and not config.use_nunchaku: - enable_fp8_linear(dit) - pipe = cls( config=config, tokenizer=tokenizer, @@ -350,6 +355,7 @@ def _from_state_dict(cls, state_dicts: QwenImageStateDicts, config: QwenImagePip vae=vae, ) pipe.eval() + pipe.dit = pipe.runtime_adapter.prepare_component("dit", pipe.dit, config) if config.offload_mode is not None: pipe.enable_cpu_offload(config.offload_mode, config.offload_to_disk) @@ -376,9 +382,10 @@ def update_weights(self, state_dicts: QwenImageStateDicts) -> None: self.update_component(self.vae, state_dicts.vae, self.config.device, self.config.vae_dtype) def compile(self): - self.dit.compile_repeated_blocks() + self.dit = self.runtime_adapter.compile_component("dit", self.dit) def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False): + self.runtime_adapter.validate_dynamic_weights(self.config) assert self.config.tp_degree is None or self.config.tp_degree == 1, ( "load LoRA is not allowed when tensor parallel is enabled; " "set tp_degree=None or tp_degree=1 during pipeline initialization" @@ -515,6 +522,8 @@ def predict_noise_with_cfg( batch_cfg: bool = False, ): if cfg_scale <= 1.0 or negative_prompt_emb is None: + if hasattr(self.dit, "set_cache_stream"): + self.dit.set_cache_stream(0) return self.predict_noise( latents, image_latents, @@ -529,6 +538,8 @@ def predict_noise_with_cfg( if not batch_cfg: # cfg by predict noise one by one h, w = latents.shape[-2:] + if hasattr(self.dit, "set_cache_stream"): + self.dit.set_cache_stream(0) positive_noise_pred = self.predict_noise( latents, image_latents, @@ -540,6 +551,8 @@ def predict_noise_with_cfg( entity_prompt_emb_masks=entity_prompt_emb_masks, entity_masks=entity_masks, ) + if hasattr(self.dit, "set_cache_stream"): + self.dit.set_cache_stream(1) negative_noise_pred = self.predict_noise( latents, image_latents, @@ -558,6 +571,8 @@ def predict_noise_with_cfg( return noise_pred else: # cfg by predict noise in one batch + if hasattr(self.dit, "set_cache_stream"): + self.dit.set_cache_stream(0) bs, _, h, w = latents.shape prompt_emb = pad_and_concat(prompt_emb, negative_prompt_emb) prompt_emb_mask = pad_and_concat(prompt_emb_mask, negative_prompt_emb_mask) @@ -699,6 +714,8 @@ def __call__( if not isinstance(controlnet_params, list): controlnet_params = [controlnet_params] + if controlnet_params: + self.runtime_adapter.validate_dynamic_weights(self.config) context_latents = None for param in controlnet_params: @@ -757,8 +774,14 @@ def __call__( self.model_lifecycle_finish(["encoder"]) self.load_models_to_device(["dit"]) + if hasattr(self.dit, "refresh_cache_status"): + cache_streams = ( + 2 if cfg_scale > 1.0 and negative_prompt_emb is not None and not self.config.batch_cfg else 1 + ) + self.dit.refresh_cache_status(len(timesteps), cache_streams) hide_progress = dist.is_initialized() and dist.get_rank() != 0 for i, timestep in enumerate(tqdm(timesteps, disable=hide_progress)): + self.runtime_adapter.before_denoise_step(i) timestep = timestep.unsqueeze(0).to(dtype=self.dtype) noise_pred = self.predict_noise_with_cfg( latents=latents, diff --git a/diffsynth_engine/platforms/__init__.py b/diffsynth_engine/platforms/__init__.py new file mode 100644 index 00000000..30d0081b --- /dev/null +++ b/diffsynth_engine/platforms/__init__.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import platform as host_platform +from typing import Type + +import torch + +from .ascend import ( + AscendPlatform, + probe_ascend_capabilities, + probe_ascend_feature, + reset_ascend_capability_cache, +) +from .base import PlatformBackend, PlatformCapabilities + + +class CPUPlatform(PlatformBackend): + name = "cpu" + device_type = "cpu" + + +class CUDAPlatform(PlatformBackend): + name = "cuda" + device_type = "cuda" + + @classmethod + def is_available(cls) -> bool: + return torch.cuda.is_available() + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + torch.cuda.set_device(index) + + @classmethod + def synchronize(cls) -> None: + torch.cuda.synchronize() + + @classmethod + def empty_cache(cls) -> None: + torch.cuda.empty_cache() + + @classmethod + def distributed_backend(cls) -> str: + return "nccl" + + +class ROCmPlatform(CUDAPlatform): + name = "rocm" + + +class MPSPlatform(PlatformBackend): + name = "mps" + device_type = "mps" + + @classmethod + def is_available(cls) -> bool: + return torch.backends.mps.is_available() + + @classmethod + def synchronize(cls) -> None: + torch.mps.synchronize() + + @classmethod + def empty_cache(cls) -> None: + torch.mps.empty_cache() + + +_PLATFORM_REGISTRY: dict[str, Type[PlatformBackend]] = { + "cpu": CPUPlatform, + "cuda": ROCmPlatform if torch.version.hip else CUDAPlatform, + "mps": MPSPlatform, + "npu": AscendPlatform, +} + + +def register_platform(device_type: str, platform_cls: Type[PlatformBackend], *, overwrite: bool = False) -> None: + if device_type in _PLATFORM_REGISTRY and not overwrite: + raise ValueError(f"Platform for device type {device_type!r} is already registered") + _PLATFORM_REGISTRY[device_type] = platform_cls + + +def get_device_type(device: str | torch.device) -> str: + if isinstance(device, torch.device): + return device.type + return str(device).split(":", 1)[0].lower() + + +def resolve_platform(device: str | torch.device) -> Type[PlatformBackend]: + device_type = get_device_type(device) + try: + return _PLATFORM_REGISTRY[device_type] + except KeyError as exc: + available = ", ".join(sorted(_PLATFORM_REGISTRY)) + raise ValueError(f"Unsupported device type {device_type!r}. Registered device types: {available}") from exc + + +def get_preferred_fp8_dtype(device: str | torch.device = "cuda") -> torch.dtype: + platform_cls = resolve_platform(device) + if platform_cls is ROCmPlatform and platform_cls.is_available(): + properties = torch.cuda.get_device_properties(0) + if "gfx94" in properties.gcnArchName: + return torch.float8_e4m3fnuz + return torch.float8_e4m3fn + + +def pin_memory(tensor: torch.Tensor, device: str | torch.device | None = None) -> torch.Tensor: + if host_platform.system() != "Linux": + return tensor + platform_cls = resolve_platform(device or "cuda") + return platform_cls.pin_memory(tensor) + + +__all__ = [ + "AscendPlatform", + "CPUPlatform", + "CUDAPlatform", + "MPSPlatform", + "PlatformBackend", + "PlatformCapabilities", + "ROCmPlatform", + "get_device_type", + "get_preferred_fp8_dtype", + "pin_memory", + "probe_ascend_capabilities", + "probe_ascend_feature", + "register_platform", + "reset_ascend_capability_cache", + "resolve_platform", +] diff --git a/diffsynth_engine/platforms/ascend.py b/diffsynth_engine/platforms/ascend.py new file mode 100644 index 00000000..7d56447d --- /dev/null +++ b/diffsynth_engine/platforms/ascend.py @@ -0,0 +1,293 @@ +from __future__ import annotations + +import importlib +from functools import lru_cache +from typing import Any + +import torch +import torch.nn as nn + +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)) + + quantization_module = importlib.import_module("mindiesd.quantization") + quant_mode_module = importlib.import_module("mindiesd.quantization.mode") + quant_algorithm = getattr(quant_mode_module, "QuantAlgorithm", None) + if not all( + ( + _has_callable(quantization_module, "quantize"), + getattr(quantization_module, "OnlineQuantConfig", None) is not None, + quant_algorithm is not None, + ) + ): + return False + if feature == "mindie_mxfp8_linear": + return hasattr(quant_algorithm, "W8A8_MXFP8") + if feature == "mindie_w4a4_linear": + return hasattr(quant_algorithm, "W4A4_MXFP4_DYNAMIC") + if feature == "mindie_fp8_attention": + quant_layer_module = importlib.import_module("mindiesd.quantization.layer") + return hasattr(quant_algorithm, "FP8_DYNAMIC") and getattr( + quant_layer_module, "FP8RotateQuantFA", None + ) is not 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) + + +def _probe_mindie_linear_operation(algorithm_name: str) -> bool: + quantization_module = importlib.import_module("mindiesd.quantization") + quant_mode_module = importlib.import_module("mindiesd.quantization.mode") + algorithm = getattr(quant_mode_module.QuantAlgorithm, algorithm_name) + model = nn.Sequential(nn.Linear(256, 256, device="npu:0", dtype=torch.bfloat16)).eval() + model = quantization_module.quantize( + model, + online_config=quantization_module.OnlineQuantConfig(quant_type=algorithm), + dtype=torch.bfloat16, + ) + value = torch.randn(1, 256, device="npu:0", dtype=torch.bfloat16) + with torch.no_grad(): + output = model(value) + return output.shape == value.shape and _tensor_probe_succeeded(output) + + +def _probe_mindie_fp8_attention_operation() -> bool: + quantization_module = importlib.import_module("mindiesd.quantization") + quant_mode_module = importlib.import_module("mindiesd.quantization.mode") + + class ProbeAttention(nn.Module): + def __init__(self): + super().__init__() + self.head_dim = 128 + self.register_buffer("_device_anchor", torch.empty(0, device="npu:0")) + + model = nn.ModuleDict({"attn": ProbeAttention()}).eval() + model = quantization_module.quantize( + model, + online_config=quantization_module.OnlineQuantConfig( + quant_type=quant_mode_module.QuantAlgorithm.W8A8_DYNAMIC, + fa_layers=("ProbeAttention",), + fa_quant_type=quant_mode_module.QuantAlgorithm.FP8_DYNAMIC, + ), + dtype=torch.bfloat16, + ) + query = torch.randn(1, 128, 8, 128, device="npu:0", dtype=torch.bfloat16) + with torch.no_grad(): + output = model["attn"].fa_quant(query, query, query, layout="BSND") + return output.shape == query.shape and _tensor_probe_succeeded(output) + + +_OPERATION_PROBES = { + "mindie_attention": _probe_mindie_attention_operation, + "mindie_compile": _probe_mindie_compile_operation, + "mindie_mxfp8_linear": lambda: _probe_mindie_linear_operation("W8A8_MXFP8"), + "mindie_w4a4_linear": lambda: _probe_mindie_linear_operation("W4A4_MXFP4_DYNAMIC"), + "mindie_fp8_attention": _probe_mindie_fp8_attention_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"), + mindie_mxfp8_linear=probe_ascend_feature("mindie_mxfp8_linear"), + mindie_w4a4_linear=probe_ascend_feature("mindie_w4a4_linear"), + mindie_fp8_attention=probe_ascend_feature("mindie_fp8_attention"), + ) + + +def reset_ascend_capability_cache() -> None: + _probe_ascend_device.cache_clear() + _probe_mindie_installation.cache_clear() + probe_ascend_feature.cache_clear() + + +class AscendPlatform(PlatformBackend): + name = "ascend" + device_type = "npu" + + @classmethod + def is_available(cls) -> bool: + return probe_ascend_feature("device") + + @classmethod + def normalize_device(cls, device: str | torch.device) -> torch.device: + _import_torch_npu() + return torch.device(device) + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + torch_npu = _import_torch_npu() + torch_npu.npu.set_device(index) + + @classmethod + def synchronize(cls) -> None: + torch_npu = _import_torch_npu() + torch_npu.npu.synchronize() + + @classmethod + def empty_cache(cls) -> None: + torch_npu = _import_torch_npu() + torch_npu.npu.empty_cache() + + @classmethod + def pin_memory(cls, tensor: torch.Tensor) -> torch.Tensor: + _import_torch_npu() + try: + return tensor.pin_memory(device="npu") + except (RuntimeError, TypeError): + # Pageable CPU memory is slower but remains correct on runtimes that + # do not expose an NPU-specific pinned allocator. + return tensor + + @classmethod + def distributed_backend(cls) -> str: + return "hccl" + + @classmethod + def compile_backend(cls): + if not cls.supports("mindie_compile"): + raise RuntimeError( + "MindIE-SD compilation was requested, but MindieSDBackend is unavailable " + "in the installed MindIE-SD package." + ) + compilation_module = importlib.import_module("mindiesd.compilation") + return compilation_module.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/ascend_qwen.py b/diffsynth_engine/platforms/ascend_qwen.py new file mode 100644 index 00000000..5931ae96 --- /dev/null +++ b/diffsynth_engine/platforms/ascend_qwen.py @@ -0,0 +1,142 @@ +from __future__ import annotations + +import importlib +from typing import Any + +import torch.nn as nn + +from diffsynth_engine.configs import QuantizationConfig + +from .runtime import ModelRuntimeAdapter + + +class MindIEFP8AttentionProcessor: + def __init__(self, fa_quant): + self.fa_quant = fa_quant + + def __call__(self, query, key, value, *, attn_mask=None, **kwargs): + if attn_mask is not None: + raise RuntimeError("MindIE FP8 attention does not support Qwen entity attention masks") + return self.fa_quant(query, key, value, layout="BSND") + + +class AscendQwenImageRuntimeAdapter(ModelRuntimeAdapter): + def __init__(self, platform, model_family: str): + super().__init__(platform, model_family) + self._quantization_enabled = False + + @staticmethod + def _quantization_config(config: Any) -> QuantizationConfig: + quantization = getattr(config, "quantization", None) + if quantization is None and getattr(config, "use_fp8_linear", False): + quantization = QuantizationConfig(linear="fp8") + return quantization or QuantizationConfig() + + @classmethod + def _uses_quantization(cls, config: Any) -> bool: + quantization = cls._quantization_config(config) + return quantization.linear != "none" or quantization.attention != "none" + + def _supports(self, capability: str) -> bool: + supports = getattr(self.platform, "supports", None) + if callable(supports): + return supports(capability) + return bool(getattr(self.platform.capabilities(), capability)) + + def validate_config(self, config: Any) -> None: + if not self._supports("device"): + raise RuntimeError( + "Ascend device requested, but no available NPU was detected. Check torch_npu and CANN installation." + ) + if getattr(config, "parallelism", 1) > 1: + raise RuntimeError("Multi-NPU execution is not supported in the first Ascend release") + + attn_impl = getattr(getattr(config, "dit_attn_impl", "auto"), "value", "auto") + if attn_impl == "mindie" and not self._supports("mindie_attention"): + raise RuntimeError( + "MindIE attention was explicitly requested, but the installed MindIE-SD package " + "does not provide a compatible attention_forward implementation" + ) + + quantization = self._quantization_config(config) + uses_quantization = self._uses_quantization(config) + if quantization.backend == "nunchaku" or getattr(config, "use_nunchaku", False): + raise RuntimeError( + "Nunchaku SVDQ/AWQ checkpoints contain CUDA-specific packed weights and cannot run on Ascend. " + "Use an unquantized Qwen-Image checkpoint with QuantizationConfig(backend='mindie', ...)." + ) + if quantization.backend not in {"auto", "native", "mindie"}: + raise ValueError(f"Unsupported Ascend quantization backend: {quantization.backend}") + + use_compile = getattr(config, "use_torch_compile", False) + offload_mode = getattr(config, "offload_mode", None) + if use_compile and offload_mode not in {None, "disable"}: + raise ValueError("MindIE-SD compilation cannot be combined with CPU offload") + if uses_quantization and offload_mode not in {None, "disable"}: + raise ValueError("MindIE-SD native quantization cannot be combined with CPU offload") + if uses_quantization and use_compile: + raise ValueError("MindIE-SD native quantization and compilation are not supported together") + + if use_compile and not self._supports("mindie_compile"): + raise RuntimeError("MindIE-SD compilation support is unavailable in the current Ascend environment") + if quantization.linear == "fp8" and not self._supports("mindie_mxfp8_linear"): + raise RuntimeError("MindIE-SD MXFP8 linear support is unavailable on this Ascend runtime") + if quantization.linear == "int4" and not self._supports("mindie_w4a4_linear"): + raise RuntimeError("MindIE-SD W4A4 linear support is unavailable on this Ascend runtime") + if quantization.attention == "fp8" and not self._supports("mindie_fp8_attention"): + raise RuntimeError("MindIE-SD FP8 attention support is unavailable on this Ascend runtime") + + def prepare_component(self, component_name: str, module: nn.Module, config: Any) -> nn.Module: + if component_name != "dit" or not self._uses_quantization(config): + return module + + quantization = self._quantization_config(config) + quantization_module = importlib.import_module("mindiesd.quantization") + quant_mode_module = importlib.import_module("mindiesd.quantization.mode") + QuantAlgorithm = quant_mode_module.QuantAlgorithm + + if quantization.linear == "fp8": + quant_type = QuantAlgorithm.W8A8_MXFP8 + fallback_layers = None + elif quantization.linear == "int4": + quant_type = QuantAlgorithm.W4A4_MXFP4_DYNAMIC + fallback_layers = None + else: + quant_type = QuantAlgorithm.W8A8_DYNAMIC + fallback_layers = { + name: QuantAlgorithm.W16A16 + for name, layer in module.named_modules() + if isinstance(layer, nn.Linear) + } + + fa_quant_type = QuantAlgorithm.FP8_DYNAMIC if quantization.attention == "fp8" else None + online_config = quantization_module.OnlineQuantConfig( + quant_type=quant_type, + fallback_layers=fallback_layers, + fa_layers=("transformer_blocks.*.attn", "QwenDoubleStreamAttention") if fa_quant_type else None, + fa_quant_type=fa_quant_type, + ) + module = quantization_module.quantize(module, online_config=online_config, dtype=config.model_dtype) + if fa_quant_type is not None: + for layer in module.modules(): + set_processor = getattr(layer, "set_attention_processor", None) + fa_quant = getattr(layer, "fa_quant", None) + if callable(set_processor) and fa_quant is not None: + set_processor(MindIEFP8AttentionProcessor(fa_quant)) + self._quantization_enabled = True + return module.eval() + + def before_denoise_step(self, step_index: int) -> None: + if not self._quantization_enabled: + return + quantization_module = importlib.import_module("mindiesd.quantization") + quantization_module.TimestepManager.set_timestep_idx(step_index) + + def validate_dynamic_weights(self, config: Any) -> None: + if getattr(config, "use_torch_compile", False): + raise ValueError("Dynamic LoRA/ControlNet loading is not supported with MindIE-SD compilation") + if self._uses_quantization(config): + raise ValueError("Dynamic LoRA/ControlNet loading is not supported with MindIE-SD native quantization") + + +__all__ = ["AscendQwenImageRuntimeAdapter", "MindIEFP8AttentionProcessor"] diff --git a/diffsynth_engine/platforms/base.py b/diffsynth_engine/platforms/base.py new file mode 100644 index 00000000..bf4a7f5e --- /dev/null +++ b/diffsynth_engine/platforms/base.py @@ -0,0 +1,71 @@ +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 + mindie_mxfp8_linear: bool = False + mindie_w4a4_linear: bool = False + mindie_fp8_attention: bool = False + + +class PlatformBackend(ABC): + name = "unknown" + device_type = "cpu" + + @classmethod + def is_available(cls) -> bool: + return True + + @classmethod + def normalize_device(cls, device: str | torch.device) -> torch.device: + return torch.device(device) + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + return None + + @classmethod + def synchronize(cls) -> None: + return None + + @classmethod + def empty_cache(cls) -> None: + return None + + @classmethod + def pin_memory(cls, tensor: torch.Tensor) -> torch.Tensor: + return tensor.pin_memory() + + @classmethod + def distributed_backend(cls) -> str: + return "gloo" + + @classmethod + def compile_backend(cls) -> Any | None: + return None + + @classmethod + def compile_kwargs(cls) -> dict[str, Any]: + backend = cls.compile_backend() + return {} if backend is None else {"backend": backend} + + @classmethod + def supports(cls, capability: str) -> bool: + capabilities = cls.capabilities() + if not hasattr(capabilities, capability): + raise ValueError(f"Unknown platform capability: {capability}") + return bool(getattr(capabilities, capability)) + + @classmethod + def capabilities(cls) -> PlatformCapabilities: + return PlatformCapabilities(device=cls.is_available()) diff --git a/diffsynth_engine/platforms/qwen.py b/diffsynth_engine/platforms/qwen.py new file mode 100644 index 00000000..b7972ec0 --- /dev/null +++ b/diffsynth_engine/platforms/qwen.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from typing import Any + +import torch.nn as nn + +from diffsynth_engine.configs import QuantizationConfig +from diffsynth_engine.utils.fp8_linear import enable_fp8_linear + +from .runtime import ModelRuntimeAdapter + + +class CudaQwenImageRuntimeAdapter(ModelRuntimeAdapter): + @staticmethod + def _quantization_config(config: Any) -> QuantizationConfig | None: + return getattr(config, "quantization", None) + + def validate_config(self, config: Any) -> None: + quantization = self._quantization_config(config) + if quantization is None: + return + if quantization.backend == "mindie": + raise ValueError("MindIE quantization requires an Ascend NPU device") + if quantization.attention != "none": + raise ValueError( + "QuantizationConfig attention quantization is currently Ascend-only; " + "use AttnImpl.FA3_FP8 or AttnImpl.AITER_FP8 on CUDA/ROCm" + ) + if quantization.backend == "nunchaku": + if not getattr(config, "use_nunchaku", False): + raise ValueError("Nunchaku quantization requires a Nunchaku-packed Qwen checkpoint") + return + if quantization.linear == "int4": + raise ValueError("INT4 Qwen execution on CUDA/ROCm requires a Nunchaku-packed checkpoint") + + def prepare_component(self, component_name: str, module: nn.Module, config: Any) -> nn.Module: + if component_name != "dit" or getattr(config, "use_nunchaku", False): + return module + + quantization = self._quantization_config(config) + use_fp8_linear = getattr(config, "use_fp8_linear", False) if quantization is None else False + if quantization is not None: + use_fp8_linear = quantization.linear == "fp8" and quantization.backend in {"auto", "native"} + if use_fp8_linear: + enable_fp8_linear(module) + return module + + +__all__ = ["CudaQwenImageRuntimeAdapter"] diff --git a/diffsynth_engine/platforms/runtime.py b/diffsynth_engine/platforms/runtime.py new file mode 100644 index 00000000..53fa518e --- /dev/null +++ b/diffsynth_engine/platforms/runtime.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import importlib +from typing import Any, Type + +import torch.nn as nn + +from .base import PlatformBackend +from . import resolve_platform + + +class ModelRuntimeAdapter: + def __init__(self, platform: Type[PlatformBackend], model_family: str): + self.platform = platform + self.model_family = model_family + + def validate_config(self, config: Any) -> None: + return None + + def prepare_component(self, component_name: str, module: nn.Module, config: Any) -> nn.Module: + return module + + def before_denoise_step(self, step_index: int) -> None: + return None + + def compile_component(self, component_name: str, module: nn.Module) -> nn.Module: + if hasattr(module, "compile_repeated_blocks"): + module.compile_repeated_blocks(**self.platform.compile_kwargs()) + else: + module.compile(**self.platform.compile_kwargs()) + return module + + def validate_dynamic_weights(self, config: Any) -> None: + return None + + +class DefaultRuntimeAdapter(ModelRuntimeAdapter): + pass + + +_RUNTIME_ADAPTERS: dict[tuple[str, str], Type[ModelRuntimeAdapter]] = {} +_LAZY_RUNTIME_ADAPTERS = { + ("cuda", "qwen_image"): ( + "diffsynth_engine.platforms.qwen", + "CudaQwenImageRuntimeAdapter", + ), + ("rocm", "qwen_image"): ( + "diffsynth_engine.platforms.qwen", + "CudaQwenImageRuntimeAdapter", + ), + ("ascend", "qwen_image"): ( + "diffsynth_engine.platforms.ascend_qwen", + "AscendQwenImageRuntimeAdapter", + ), +} + + +def register_runtime_adapter( + platform_name: str, + model_family: str, + adapter_cls: Type[ModelRuntimeAdapter], + *, + overwrite: bool = False, +) -> None: + key = (platform_name, model_family) + if key in _RUNTIME_ADAPTERS and not overwrite: + raise ValueError(f"Runtime adapter for {key!r} is already registered") + _RUNTIME_ADAPTERS[key] = adapter_cls + + +def _load_lazy_adapter(key: tuple[str, str]) -> None: + if key not in _LAZY_RUNTIME_ADAPTERS or key in _RUNTIME_ADAPTERS: + return + module_name, class_name = _LAZY_RUNTIME_ADAPTERS[key] + module = importlib.import_module(module_name) + _RUNTIME_ADAPTERS[key] = getattr(module, class_name) + + +def get_runtime_adapter(device: str, model_family: str) -> ModelRuntimeAdapter: + platform = resolve_platform(device) + key = (platform.name, model_family) + _load_lazy_adapter(key) + adapter_cls = _RUNTIME_ADAPTERS.get(key, DefaultRuntimeAdapter) + return adapter_cls(platform, model_family) + + +__all__ = [ + "DefaultRuntimeAdapter", + "ModelRuntimeAdapter", + "get_runtime_adapter", + "register_runtime_adapter", +] diff --git a/diffsynth_engine/utils/offload.py b/diffsynth_engine/utils/offload.py index a5e7f7a2..63516f21 100644 --- a/diffsynth_engine/utils/offload.py +++ b/diffsynth_engine/utils/offload.py @@ -29,11 +29,11 @@ def _forward_pre_hook(module: nn.Module, input_): buffer.data = buffer.data.to(device=device) return tuple(x.to(device=device) if isinstance(x, torch.Tensor) else x for x in input_) for name, param in module.named_parameters(recurse=recurse): - param.data = pin_memory(param.data) + param.data = pin_memory(param.data, device) offload_param_dict[name] = param.data param.data = param.data.to(device=device) for name, buffer in module.named_buffers(recurse=recurse): - buffer.data = pin_memory(buffer.data) + buffer.data = pin_memory(buffer.data, device) offload_param_dict[name] = buffer.data buffer.data = buffer.data.to(device=device) setattr(module, "_offload_param_dict", offload_param_dict) @@ -55,14 +55,14 @@ def _forward_hook(module: nn.Module, input_, output_): setattr(module, "_cpu_offload_enabled", True) -def offload_model_to_dict(module: nn.Module) -> Dict[str, torch.Tensor]: +def offload_model_to_dict(module: nn.Module, device: str = "cuda") -> Dict[str, torch.Tensor]: module = module.to("cpu") offload_param_dict = {} for name, param in module.named_parameters(recurse=True): - param.data = pin_memory(param.data) + param.data = pin_memory(param.data, device) offload_param_dict[name] = param.data for name, buffer in module.named_buffers(recurse=True): - buffer.data = pin_memory(buffer.data) + buffer.data = pin_memory(buffer.data, device) offload_param_dict[name] = buffer.data return offload_param_dict diff --git a/diffsynth_engine/utils/platform.py b/diffsynth_engine/utils/platform.py index 49a69680..82d14abb 100644 --- a/diffsynth_engine/utils/platform.py +++ b/diffsynth_engine/utils/platform.py @@ -1,25 +1,51 @@ -# cross-platform definitions and utilities -import torch import gc -import platform +from typing import Optional + +import torch + +from diffsynth_engine.platforms import ( + AscendPlatform, + get_device_type as _get_device_type, + get_preferred_fp8_dtype, + pin_memory as _pin_memory, + resolve_platform, +) -# data type -# AMD only supports float8_e4m3fnuz -# https://onnx.ai/onnx/technical/float8.html -if torch.version.hip and "gfx94" in torch.cuda.get_device_properties(0).gcnArchName: - DTYPE_FP8 = torch.float8_e4m3fnuz -else: - DTYPE_FP8 = torch.float8_e4m3fn +DTYPE_FP8 = get_preferred_fp8_dtype("cuda") -def empty_cache(): +def empty_cache(device: Optional[str | torch.device] = None): gc.collect() + if device is not None: + platform_cls = resolve_platform(device) + if platform_cls.is_available(): + platform_cls.empty_cache() + return + + # No-argument calls preserve the historical CUDA/MPS behavior. NPU callers + # pass their explicit device so a coexisting NPU never affects CUDA cleanup. if torch.cuda.is_available(): - torch.cuda.empty_cache() - if torch.mps.is_available(): - torch.mps.empty_cache() + resolve_platform("cuda").empty_cache() + if torch.backends.mps.is_available(): + resolve_platform("mps").empty_cache() + + +def pin_memory(tensor: torch.Tensor, device: Optional[str | torch.device] = None): + return _pin_memory(tensor, device) + + +def is_npu_available() -> bool: + return AscendPlatform.supports("device") + + +def is_mindie_sd_available() -> bool: + return AscendPlatform.supports("mindie") + + +def get_device_type(device: str | torch.device) -> str: + return _get_device_type(device) -def pin_memory(tensor: torch.Tensor): - return tensor.pin_memory() if platform.system() == "Linux" else tensor +def get_torch_distributed_backend(device: str | torch.device) -> str: + return resolve_platform(device).distributed_backend() diff --git a/docs/ascend.md b/docs/ascend.md new file mode 100644 index 00000000..516d41e2 --- /dev/null +++ b/docs/ascend.md @@ -0,0 +1,148 @@ +# Ascend A5 / Ascend 950 推理指南 + +DiffSynth-Engine 首批支持在单张 Ascend A5 / Ascend 950 上运行 Qwen-Image、 +Qwen-Image-Edit-2509 和 Qwen-Image-Edit-2511。NPU 不会自动替换默认 CUDA 路径， +必须显式设置 `device="npu:0"`。 + +## 环境 + +请使用同一发布组合中的 CANN、PyTorch、torch-npu 和 MindIE-SD 3.x。torch-npu 与 +MindIE-SD 不属于 DiffSynth-Engine 的核心 PyPI 依赖，需要根据华为发布矩阵单独安装。 +硬件 CI 镜像应记录实际 wheel 版本；项目不在 `pyproject.toml` 中固定这些平台依赖。 + +| 组件 | 要求 | +| --- | --- | +| 硬件 | Ascend A5 / Ascend 950 | +| CANN | 与 torch-npu wheel 配套的版本 | +| torch-npu | 与当前 PyTorch 和 CANN 配套的版本 | +| MindIE-SD | 3.x，与当前 torch-npu/CANN 配套 | + +普通 BF16 NPU 推理只要求 torch-npu。MindIE-SD 不存在时，`AttnImpl.AUTO` 会从 +MindIE attention 回退到 PyTorch SDPA，再回退到 eager attention。显式选择 MindIE +attention、MindIE compile 或原生量化时不会静默降级，而是在模型分配前报告缺失能力。 + +安装后可检查能力： + +```python +from diffsynth_engine.platforms import probe_ascend_capabilities + +print(probe_ascend_capabilities()) +``` + +该调用会为基础设备、MindIE attention、compile、MXFP8、W4A4 和 FP8 attention +分别执行一次最小算子并缓存结果，不只检查 Python 模块能否导入。正常推理则按需探测所使用的 +功能；例如 BF16 `AUTO` 配置不会在初始化阶段触发量化或 compile 探测。 + +## BF16 推理 + +完整示例见 `examples/qwen_image_ascend.py`，核心配置如下： + +```python +from diffsynth_engine import AttnImpl, QwenImagePipeline, QwenImagePipelineConfig + +config = QwenImagePipelineConfig.basic_config( + model_path=model_path, + encoder_path=encoder_path, + vae_path=vae_path, + device="npu:0", + parallelism=1, +) +config.dit_attn_impl = AttnImpl.AUTO +pipe = QwenImagePipeline.from_pretrained(config) +``` + +如需强制使用 MindIE-SD attention，可设置： + +```python +config.dit_attn_impl = AttnImpl.MINDIE +``` + +该模式缺少 MindIE-SD 或 `attention_forward` API 不兼容时会立即失败。首批版本不支持 +NPU 长上下文 attention，也不支持 `parallelism > 1`；多 NPU 配置会 fail-fast。 + +## MindIE compile + +compile 保持显式开启，只编译 Qwen DiT 中重复的 Transformer blocks，不编译 encoder +或 VAE： + +```python +config.use_torch_compile = True +``` + +NPU 使用 `MindieSDBackend()`，CUDA/ROCm 的空参数 `torch.compile` 行为不变。MindIE +compile 需要静态权重，不能与 CPU/model/sequential offload、动态 LoRA 或 ControlNet、 +原生量化组合。FB-cache 的控制流位于已编译 block 外，可以与 compile 组合验证。 + +## 原生量化 + +原生量化从未量化的 Qwen 权重在线构建，不接受 Nunchaku/SVDQ/AWQ packed checkpoint。 + +```python +from diffsynth_engine import QuantizationConfig + +# W8A8_MXFP8 linear +config.quantization = QuantizationConfig(backend="mindie", linear="fp8") + +# W4A4_MXFP4_DYNAMIC linear +config.quantization = QuantizationConfig(backend="mindie", linear="int4") + +# FP8_DYNAMIC attention，可与上述任一 linear 配置合并 +config.quantization = QuantizationConfig( + backend="mindie", + linear="fp8", + attention="fp8", +) +``` + +Qwen denoise 循环会在每一步更新 MindIE `TimestepManager`。如果运行时没有对应的 +MXFP8、W4A4 或 FP8 attention 算子，初始化直接报错，不会退回 BF16。原生量化不能与 +offload、compile、动态 LoRA 或 ControlNet 组合。 + +旧配置 `use_fp8_linear=True` 仍然保留：CUDA/ROCm 继续走现有 `_scaled_mm` 路径，NPU +映射为 MindIE `W8A8_MXFP8`。推荐新代码使用 `QuantizationConfig`。 + +## 支持矩阵 + +| 场景 | 状态 | 约束 | +| --- | --- | --- | +| Qwen-Image / Edit-2509 / Edit-2511 BF16 | 支持 | 单 NPU，AUTO 或 MindIE attention | +| LoRA / ControlNet | 支持 | BF16；可与 offload 或 FB-cache 组合 | +| FB-cache | 支持 | 可分别与 BF16、compile 或原生量化组合验证 | +| CPU/model/sequential offload | 支持 | 不与 compile 或原生量化组合 | +| MindIE compile | 支持 | 静态权重；不含动态 LoRA/ControlNet | +| MXFP8 / W4A4 linear、FP8 attention | 支持 | A5 算子能力探测必须通过 | +| Nunchaku/SVDQ/AWQ checkpoint | NPU 不支持 | 使用未量化权重和 MindIE 原生量化 | +| 多 NPU | 暂不支持 | `parallelism > 1` 明确报错 | + +## 硬件测试 + +PR [#259](https://github.com/modelscope/DiffSynth-Engine/pull/259) 的 v1 实现已经在单张 +Ascend 950 上完成 Qwen-Image-Edit-2511 实测，可作为主线融合的硬件基线。测试使用 +`device="npu:0"`、MindIE attention、`parallelism=1`、预热后 10 个推理步： + +| 输入 | H20 端到端 | Ascend 950 端到端 | H20 / Ascend | +| --- | ---: | ---: | ---: | +| 1 image, 1024x1024 | 29.10s | 16.97s | 1.71x | +| 2 images, 1024x1024 | 49.12s | 30.65s | 1.60x | +| 4 images, 1280x720 | 98.35s | 54.68s | 1.80x | + +三组设备 kernel 总时间比分别为 1.97x、1.99x、1.97x；主要热点是 GEMM/Linear 和 +FlashAttention。该结果证明 v1 硬件路径可用，但主线重构后的最终发布仍需在同一镜像中 +重新执行下列回归，不能用 v1 结果替代主线 golden、compile 和量化验收。 + +先运行最小 NPU 与 MindIE attention 探测： + +```bash +RUN_ASCEND_TESTS=1 python -m unittest tests.test_platforms.test_ascend_integration +``` + +量化发布验收还需启用完整能力断言，并执行 Qwen golden tensor、文生图、Edit、 +ControlNet、LoRA、FB-cache、offload、compile 和三种量化图像用例： + +```bash +RUN_ASCEND_TESTS=1 RUN_ASCEND_QUANT_TESTS=1 \ + python -m unittest tests.test_platforms.test_ascend_integration +``` + +图像阈值沿用现有测试：基础 Qwen 0.99、Edit/ControlNet 0.95、量化 0.90。compile +输出还需与同机 eager 对齐，并确认 MindIE 融合 graph 实际生成。 diff --git a/examples/qwen_image_ascend.py b/examples/qwen_image_ascend.py new file mode 100644 index 00000000..35dea30d --- /dev/null +++ b/examples/qwen_image_ascend.py @@ -0,0 +1,23 @@ +from diffsynth_engine import AttnImpl, QwenImagePipeline, QwenImagePipelineConfig, fetch_model + + +if __name__ == "__main__": + config = QwenImagePipelineConfig.basic_config( + model_path=fetch_model("MusePublic/Qwen-image", revision="v1", path="transformer/*.safetensors"), + encoder_path=fetch_model("MusePublic/Qwen-image", revision="v1", path="text_encoder/*.safetensors"), + vae_path=fetch_model("MusePublic/Qwen-image", revision="v1", path="vae/*.safetensors"), + device="npu:0", + parallelism=1, + ) + config.dit_attn_impl = AttnImpl.AUTO + + pipe = QwenImagePipeline.from_pretrained(config) + image = pipe( + prompt="A red panda reading a book beside a sunlit window", + negative_prompt=" ", + width=1328, + height=1328, + num_inference_steps=30, + seed=42, + ) + image.save("qwen_image_ascend.png") diff --git a/tests/test_models/basic/__init__.py b/tests/test_models/basic/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_models/basic/test_attention_backend.py b/tests/test_models/basic/test_attention_backend.py new file mode 100644 index 00000000..3fd0d64c --- /dev/null +++ b/tests/test_models/basic/test_attention_backend.py @@ -0,0 +1,126 @@ +import sys +import types +import unittest +from unittest import mock + +import torch + +from diffsynth_engine.models.basic import attention as attention_ops +from diffsynth_engine.models.basic import attention_backend +from diffsynth_engine.models.basic.attention_backend import ( + AttentionBackend, + AttentionRequest, + register_attention_backend, + resolve_attention_backend, +) + + +class FakeTensor: + def __init__(self, device_type="cpu", shape=(1, 8, 2, 4)): + self.device = types.SimpleNamespace(type=device_type) + self.shape = shape + self.ndim = len(shape) + + +def make_request(device_type="cpu"): + tensor = FakeTensor(device_type) + return AttentionRequest(tensor, tensor, tensor) + + +class TestAttentionBackendResolution(unittest.TestCase): + def test_npu_auto_order_is_mindie_sdpa_eager(self): + backends = ( + AttentionBackend("mindie", mock.Mock(), frozenset({"npu"}), 800), + AttentionBackend("sdpa", mock.Mock(), None, 300), + AttentionBackend("eager", mock.Mock(), None, 100), + ) + with mock.patch.dict(attention_backend._ATTENTION_BACKENDS, {}, clear=True): + for backend in backends: + register_attention_backend(backend) + self.assertEqual(resolve_attention_backend("auto", make_request("npu")).name, "mindie") + + register_attention_backend( + AttentionBackend("mindie", mock.Mock(), frozenset({"npu"}), 800, lambda: False), + overwrite=True, + ) + self.assertEqual(resolve_attention_backend("auto", make_request("npu")).name, "sdpa") + + register_attention_backend( + AttentionBackend("sdpa", mock.Mock(), None, 300, lambda: False), + overwrite=True, + ) + self.assertEqual(resolve_attention_backend("auto", make_request("npu")).name, "eager") + + def test_explicit_unavailable_backend_does_not_fall_back(self): + backend = AttentionBackend("optional", mock.Mock(), None, 1, lambda: False) + with mock.patch.dict(attention_backend._ATTENTION_BACKENDS, {"optional": backend}, clear=True): + with self.assertRaisesRegex(RuntimeError, "not available"): + resolve_attention_backend("optional", make_request()) + + def test_auto_can_preserve_backend_specific_legacy_capabilities(self): + explicit_support = mock.Mock(return_value=(False, "explicit request rejected")) + auto_support = mock.Mock(return_value=(True, None)) + backend = AttentionBackend( + "legacy", + mock.Mock(), + supports=explicit_support, + auto_supports=auto_support, + ) + with mock.patch.dict(attention_backend._ATTENTION_BACKENDS, {"legacy": backend}, clear=True): + self.assertIs(resolve_attention_backend("auto", make_request()), backend) + with self.assertRaisesRegex(RuntimeError, "explicit request rejected"): + resolve_attention_backend("legacy", make_request()) + + def test_cuda_auto_priority_matches_mainline(self): + expected = ["fa4", "fa3", "aiter", "xformers", "sdpa", "fa2", "eager"] + priorities = { + name: attention_backend._ATTENTION_BACKENDS[name].priority + for name in expected + } + self.assertEqual( + [name for name, _ in sorted(priorities.items(), key=lambda item: item[1], reverse=True)], + expected, + ) + + +class TestMindieAttention(unittest.TestCase): + def test_forwards_public_attention_arguments(self): + attention_forward = mock.Mock(return_value=torch.ones(1)) + module_name = "mindiesd.layers.flash_attn.attention_forward" + fake_module = types.ModuleType(module_name) + fake_module.attention_forward = attention_forward + q = torch.randn(1, 4, 2, 8) + k = torch.randn(1, 4, 2, 8) + v = torch.randn(1, 4, 2, 8) + mask = torch.randn(1, 1, 4, 4) + request = AttentionRequest(q, k, v, attn_mask=mask, scale=0.125) + + with mock.patch.dict(sys.modules, {module_name: fake_module}): + output = attention_ops._run_mindie(request) + + self.assertEqual(output.shape, (1,)) + attention_forward.assert_called_once_with( + query=q, + key=k, + value=v, + attn_mask=mask, + scale=0.125, + fused=True, + head_first=False, + ) + + def test_mindie_rejects_non_4d_inputs(self): + request = make_request("npu") + request.q.ndim = 3 + supported, reason = attention_ops._mindie_supports(request) + self.assertFalse(supported) + self.assertIn("4D", reason) + + def test_explicit_mindie_has_actionable_dependency_error(self): + with mock.patch("diffsynth_engine.platforms.AscendPlatform.supports", return_value=False): + with self.assertRaisesRegex(RuntimeError, "install a compatible MindIE-SD 3.x"): + resolve_attention_backend("mindie", make_request("npu")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_models/qwen_image/test_qwen_attention_processor.py b/tests/test_models/qwen_image/test_qwen_attention_processor.py new file mode 100644 index 00000000..144cce5d --- /dev/null +++ b/tests/test_models/qwen_image/test_qwen_attention_processor.py @@ -0,0 +1,63 @@ +import unittest +from unittest import mock + +import torch +import torch.nn as nn + +from diffsynth_engine.models.qwen_image.qwen_image_dit import QwenDoubleStreamAttention +from diffsynth_engine.platforms.ascend_qwen import MindIEFP8AttentionProcessor + + +class FakeFAQuant(nn.Module): + def __init__(self): + super().__init__() + self.call = mock.Mock() + + def forward(self, query, key, value, **kwargs): + self.call(query, key, value, **kwargs) + return query + + +class TestQwenAttentionProcessor(unittest.TestCase): + def setUp(self): + self.attention = QwenDoubleStreamAttention( + dim_a=8, + dim_b=8, + num_heads=2, + head_dim=4, + device="cpu", + dtype=torch.float32, + ) + self.image = torch.randn(1, 3, 8) + self.text = torch.randn(1, 2, 8) + + def test_fa_quant_processor_uses_bsnd_layout(self): + fa_quant = FakeFAQuant() + self.attention.set_attention_processor(MindIEFP8AttentionProcessor(fa_quant)) + image, text = self.attention(self.image, self.text) + + self.assertEqual(image.shape, self.image.shape) + self.assertEqual(text.shape, self.text.shape) + self.assertEqual(fa_quant.call.call_args.kwargs, {"layout": "BSND"}) + query = fa_quant.call.call_args.args[0] + self.assertEqual(query.shape, (1, 5, 2, 4)) + + def test_default_path_still_uses_common_attention(self): + with mock.patch( + "diffsynth_engine.models.qwen_image.qwen_image_dit.attention_ops.attention", + side_effect=lambda q, k, v, **kwargs: q, + ) as common_attention: + self.attention(self.image, self.text, attn_kwargs={"attn_impl": "eager"}) + + self.assertEqual(common_attention.call_count, 1) + self.assertEqual(common_attention.call_args.kwargs["attn_impl"], "eager") + + def test_fa_quant_processor_rejects_entity_mask(self): + self.attention.set_attention_processor(MindIEFP8AttentionProcessor(FakeFAQuant())) + mask = torch.zeros(1, 1, 5, 5) + with self.assertRaisesRegex(RuntimeError, "entity attention masks"): + self.attention(self.image, self.text, attn_mask=mask) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_models/qwen_image/test_qwen_fbcache.py b/tests/test_models/qwen_image/test_qwen_fbcache.py new file mode 100644 index 00000000..e7b82c80 --- /dev/null +++ b/tests/test_models/qwen_image/test_qwen_fbcache.py @@ -0,0 +1,38 @@ +import inspect +import unittest + +from diffsynth_engine.models.qwen_image.qwen_image_dit_fbcache import QwenImageDiTFBCache + + +class TestQwenFBCacheContract(unittest.TestCase): + def test_forward_matches_current_qwen_conditioning_contract(self): + parameters = inspect.signature(QwenImageDiTFBCache.forward).parameters + for name in ( + "edit", + "text_seq_lens", + "context_latents", + "entity_text", + "entity_seq_lens", + "entity_masks", + "attn_kwargs", + ): + self.assertIn(name, parameters) + + def test_cache_state_resets_for_each_pipeline_call(self): + model = QwenImageDiTFBCache.__new__(QwenImageDiTFBCache) + model.step_count = 9 + model.num_inference_steps = 9 + model.refresh_cache_status(30, num_cache_streams=2) + self.assertEqual(model.step_count, 0) + self.assertEqual(model.num_inference_steps, 30) + + model._cache_states[0]["step_count"] = 4 + model._cache_states[1]["step_count"] = 7 + model.set_cache_stream(0) + self.assertEqual(model.step_count, 4) + model.set_cache_stream(1) + self.assertEqual(model.step_count, 7) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_platforms/__init__.py b/tests/test_platforms/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_platforms/test_ascend_integration.py b/tests/test_platforms/test_ascend_integration.py new file mode 100644 index 00000000..5b4d528f --- /dev/null +++ b/tests/test_platforms/test_ascend_integration.py @@ -0,0 +1,95 @@ +import os +import unittest + +import torch + + +RUN_ASCEND_TESTS = os.getenv("RUN_ASCEND_TESTS", "0") == "1" +RUN_ASCEND_QUANT_TESTS = os.getenv("RUN_ASCEND_QUANT_TESTS", "0") == "1" + + +@unittest.skipUnless(RUN_ASCEND_TESTS, "RUN_ASCEND_TESTS is not set") +class TestAscendOperatorIntegration(unittest.TestCase): + @classmethod + def setUpClass(cls): + import torch_npu # noqa: F401 + + from diffsynth_engine.platforms import probe_ascend_feature + + cls.probe_feature = staticmethod(probe_ascend_feature) + if not cls.probe_feature("device"): + raise unittest.SkipTest("No available Ascend NPU") + + def test_basic_npu_execution(self): + from diffsynth_engine.platforms import resolve_platform + + platform = resolve_platform("npu:0") + platform.set_device("npu:0") + lhs = torch.randn(32, 32, device="npu:0", dtype=torch.bfloat16) + rhs = torch.randn(32, 32, device="npu:0", dtype=torch.bfloat16) + output = lhs @ rhs + platform.synchronize() + self.assertEqual(output.device.type, "npu") + self.assertEqual(output.shape, (32, 32)) + + def test_mindie_attention_operator(self): + if not self.probe_feature("mindie_attention"): + self.skipTest("MindIE attention is unavailable") + from diffsynth_engine.models.basic.attention import attention + + q = torch.randn(1, 128, 8, 128, device="npu:0", dtype=torch.bfloat16) + output = attention(q, q, q, attn_impl="mindie") + self.assertEqual(output.shape, q.shape) + self.assertEqual(output.device.type, "npu") + + @unittest.skipUnless(RUN_ASCEND_QUANT_TESTS, "RUN_ASCEND_QUANT_TESTS is not set") + def test_native_linear_quantization_operators(self): + self.assertTrue(self.probe_feature("mindie_mxfp8_linear")) + self.assertTrue(self.probe_feature("mindie_w4a4_linear")) + import torch.nn as nn + from mindiesd.quantization import OnlineQuantConfig, quantize + from mindiesd.quantization.mode import QuantAlgorithm + + for algorithm in (QuantAlgorithm.W8A8_MXFP8, QuantAlgorithm.W4A4_MXFP4_DYNAMIC): + with self.subTest(algorithm=algorithm): + model = nn.Sequential(nn.Linear(256, 256, device="npu:0", dtype=torch.bfloat16)) + model = quantize( + model, + online_config=OnlineQuantConfig(quant_type=algorithm), + dtype=torch.bfloat16, + ) + output = model(torch.randn(1, 256, device="npu:0", dtype=torch.bfloat16)) + self.assertEqual(output.shape, (1, 256)) + self.assertTrue(torch.isfinite(output).all().cpu()) + + @unittest.skipUnless(RUN_ASCEND_QUANT_TESTS, "RUN_ASCEND_QUANT_TESTS is not set") + def test_native_fp8_attention_operator(self): + self.assertTrue(self.probe_feature("mindie_fp8_attention")) + import torch.nn as nn + from mindiesd.quantization import OnlineQuantConfig, quantize + from mindiesd.quantization.mode import QuantAlgorithm + + class ProbeAttention(nn.Module): + def __init__(self): + super().__init__() + self.head_dim = 128 + self.register_buffer("_device_anchor", torch.empty(0, device="npu:0")) + + model = nn.ModuleDict({"attn": ProbeAttention()}).to("npu:0") + model = quantize( + model, + online_config=OnlineQuantConfig( + quant_type=QuantAlgorithm.W8A8_DYNAMIC, + fa_layers=("ProbeAttention",), + fa_quant_type=QuantAlgorithm.FP8_DYNAMIC, + ), + dtype=torch.bfloat16, + ) + q = torch.randn(1, 128, 8, 128, device="npu:0", dtype=torch.bfloat16) + output = model["attn"].fa_quant(q, q, q, layout="BSND") + self.assertEqual(output.shape, q.shape) + self.assertTrue(torch.isfinite(output).all().cpu()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_platforms/test_platforms.py b/tests/test_platforms/test_platforms.py new file mode 100644 index 00000000..98c603f5 --- /dev/null +++ b/tests/test_platforms/test_platforms.py @@ -0,0 +1,158 @@ +import types +import unittest +from unittest import mock + +import torch + +from diffsynth_engine.platforms import ( + get_device_type, + resolve_platform, +) +from diffsynth_engine.platforms import ascend +from diffsynth_engine.utils import platform as compatibility_platform + + +class TestPlatformResolution(unittest.TestCase): + def test_resolution_only_uses_explicit_device(self): + with mock.patch.object(ascend.AscendPlatform, "is_available", return_value=True): + self.assertIn(resolve_platform("cuda").name, {"cuda", "rocm"}) + self.assertEqual(resolve_platform("npu:0").name, "ascend") + self.assertEqual(get_device_type(torch.device("cpu")), "cpu") + + def test_auto_device_is_not_silently_resolved(self): + with self.assertRaisesRegex(ValueError, "Unsupported device type 'auto'"): + resolve_platform("auto") + + def test_resolution_does_not_import_torch_npu(self): + with mock.patch.object(ascend, "_import_torch_npu") as import_torch_npu: + self.assertEqual(resolve_platform("npu:0").name, "ascend") + import_torch_npu.assert_not_called() + + def test_legacy_cache_cleanup_does_not_implicitly_touch_npu(self): + with ( + mock.patch.object(torch.cuda, "is_available", return_value=False), + mock.patch.object(torch.backends.mps, "is_available", return_value=False), + mock.patch.object(ascend.AscendPlatform, "empty_cache") as npu_empty_cache, + ): + compatibility_platform.empty_cache() + npu_empty_cache.assert_not_called() + + +class TestAscendCapabilities(unittest.TestCase): + def tearDown(self): + ascend.reset_ascend_capability_cache() + + @staticmethod + def _fake_torch_npu(): + return types.SimpleNamespace( + npu=types.SimpleNamespace(is_available=lambda: True, get_device_name=lambda index: "Ascend"), + npu_dynamic_mx_quant=lambda *args, **kwargs: None, + npu_quant_matmul=lambda *args, **kwargs: None, + npu_dynamic_block_quant=lambda *args, **kwargs: None, + npu_fused_infer_attention_score_v2=lambda *args, **kwargs: None, + float4_e2m1fn_x2=object(), + float8_e4m3fn=object(), + ) + + def test_capability_probe_checks_apis_and_operators(self): + class QuantAlgorithm: + W8A8_MXFP8 = object() + W4A4_MXFP4_DYNAMIC = object() + FP8_DYNAMIC = object() + + modules = { + "mindiesd.layers.flash_attn.attention_forward": types.SimpleNamespace( + attention_forward=lambda *args, **kwargs: None + ), + "mindiesd.compilation": types.SimpleNamespace(MindieSDBackend=lambda: object()), + "mindiesd.quantization": types.SimpleNamespace( + quantize=lambda *args, **kwargs: None, + OnlineQuantConfig=object, + ), + "mindiesd.quantization.mode": types.SimpleNamespace(QuantAlgorithm=QuantAlgorithm), + "mindiesd.quantization.layer": types.SimpleNamespace(FP8RotateQuantFA=object), + } + operation_probes = {name: mock.Mock(return_value=True) for name in ascend._OPERATION_PROBES} + + ascend.reset_ascend_capability_cache() + with ( + mock.patch.object(ascend, "_import_torch_npu", return_value=self._fake_torch_npu()), + mock.patch.object(ascend, "_probe_npu_runtime", return_value=True), + mock.patch.object(ascend, "_import_mindie_sd", return_value=object()), + mock.patch.object(ascend.importlib, "import_module", side_effect=modules.__getitem__), + mock.patch.dict(ascend._OPERATION_PROBES, operation_probes), + ): + capabilities = ascend.probe_ascend_capabilities() + cached_capabilities = ascend.probe_ascend_capabilities() + + self.assertTrue(capabilities.device) + self.assertTrue(capabilities.mindie_attention) + self.assertTrue(capabilities.mindie_compile) + self.assertTrue(capabilities.mindie_mxfp8_linear) + self.assertTrue(capabilities.mindie_w4a4_linear) + self.assertTrue(capabilities.mindie_fp8_attention) + self.assertEqual(capabilities, cached_capabilities) + for operation_probe in operation_probes.values(): + operation_probe.assert_called_once_with() + + def test_failed_operator_probe_disables_only_that_feature(self): + ascend.reset_ascend_capability_cache() + with ( + mock.patch.object(ascend, "_probe_ascend_device", return_value=True), + mock.patch.object(ascend, "_probe_mindie_installation", return_value=True), + mock.patch.object(ascend, "_feature_api_available", return_value=True), + mock.patch.dict( + ascend._OPERATION_PROBES, + {"mindie_attention": mock.Mock(side_effect=RuntimeError("operator unavailable"))}, + ), + ): + self.assertFalse(ascend.probe_ascend_feature("mindie_attention")) + + def test_missing_torch_npu_reports_no_device(self): + ascend.reset_ascend_capability_cache() + with mock.patch.object(ascend, "_import_torch_npu", side_effect=RuntimeError("missing")): + capabilities = ascend.probe_ascend_capabilities() + self.assertFalse(capabilities.device) + + def test_mindie_missing_keeps_basic_npu_available(self): + ascend.reset_ascend_capability_cache() + with ( + mock.patch.object(ascend, "_import_torch_npu", return_value=self._fake_torch_npu()), + mock.patch.object(ascend, "_probe_npu_runtime", return_value=True), + mock.patch.object(ascend, "_import_mindie_sd", side_effect=RuntimeError("missing")), + ): + capabilities = ascend.probe_ascend_capabilities() + + self.assertTrue(capabilities.device) + self.assertFalse(capabilities.mindie) + self.assertFalse(capabilities.mindie_attention) + + def test_incompatible_mindie_api_disables_feature_without_breaking_npu(self): + ascend.reset_ascend_capability_cache() + with ( + mock.patch.object(ascend, "_import_torch_npu", return_value=self._fake_torch_npu()), + mock.patch.object(ascend, "_probe_npu_runtime", return_value=True), + mock.patch.object(ascend, "_import_mindie_sd", return_value=object()), + mock.patch.object(ascend.importlib, "import_module", side_effect=ValueError("incompatible API")), + ): + capabilities = ascend.probe_ascend_capabilities() + + self.assertTrue(capabilities.device) + self.assertTrue(capabilities.mindie) + self.assertFalse(capabilities.mindie_attention) + self.assertFalse(capabilities.mindie_compile) + self.assertFalse(capabilities.mindie_mxfp8_linear) + + def test_runtime_probe_failure_overrides_is_available(self): + ascend.reset_ascend_capability_cache() + with ( + mock.patch.object(ascend, "_import_torch_npu", return_value=self._fake_torch_npu()), + mock.patch.object(ascend, "_probe_npu_runtime", return_value=False), + ): + capabilities = ascend.probe_ascend_capabilities() + + self.assertFalse(capabilities.device) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_platforms/test_runtime_adapters.py b/tests/test_platforms/test_runtime_adapters.py new file mode 100644 index 00000000..2146e00c --- /dev/null +++ b/tests/test_platforms/test_runtime_adapters.py @@ -0,0 +1,258 @@ +import types +import unittest +from unittest import mock + +import torch +import torch.nn as nn + +from diffsynth_engine.configs import QuantizationConfig +from diffsynth_engine.platforms import PlatformCapabilities +from diffsynth_engine.platforms import ascend_qwen +from diffsynth_engine.platforms.ascend_qwen import AscendQwenImageRuntimeAdapter +from diffsynth_engine.platforms.qwen import CudaQwenImageRuntimeAdapter +from diffsynth_engine.platforms.runtime import DefaultRuntimeAdapter + + +class FakePlatform: + name = "ascend" + current_capabilities = PlatformCapabilities( + device=True, + mindie=True, + mindie_attention=True, + mindie_compile=True, + mindie_mxfp8_linear=True, + mindie_w4a4_linear=True, + mindie_fp8_attention=True, + ) + + @classmethod + def capabilities(cls): + return cls.current_capabilities + + @classmethod + def supports(cls, capability): + return bool(getattr(cls.current_capabilities, capability)) + + @classmethod + def compile_kwargs(cls): + return {"backend": "mindie"} + + +def make_config(**overrides): + values = { + "device": "npu:0", + "parallelism": 1, + "dit_attn_impl": types.SimpleNamespace(value="auto"), + "quantization": None, + "use_fp8_linear": False, + "use_nunchaku": False, + "use_torch_compile": False, + "offload_mode": None, + "model_dtype": torch.bfloat16, + } + values.update(overrides) + return types.SimpleNamespace(**values) + + +class TestAscendQwenRuntimeAdapter(unittest.TestCase): + def setUp(self): + FakePlatform.current_capabilities = PlatformCapabilities( + device=True, + mindie=True, + mindie_attention=True, + mindie_compile=True, + mindie_mxfp8_linear=True, + mindie_w4a4_linear=True, + mindie_fp8_attention=True, + ) + self.adapter = AscendQwenImageRuntimeAdapter(FakePlatform, "qwen_image") + + def test_rejects_first_release_incompatibilities(self): + invalid = ( + ({"parallelism": 2}, "Multi-NPU"), + ({"use_nunchaku": True}, "cannot run on Ascend"), + ({"use_torch_compile": True, "offload_mode": "cpu_offload"}, "compilation cannot be combined"), + ( + {"quantization": QuantizationConfig(linear="fp8"), "offload_mode": "cpu_offload"}, + "quantization cannot be combined", + ), + ( + {"quantization": QuantizationConfig(linear="fp8"), "use_torch_compile": True}, + "quantization and compilation", + ), + ) + for overrides, message in invalid: + with self.subTest(overrides=overrides), self.assertRaisesRegex((ValueError, RuntimeError), message): + self.adapter.validate_config(make_config(**overrides)) + + def test_explicit_mindie_fails_before_model_loading(self): + FakePlatform.current_capabilities = PlatformCapabilities(device=True) + config = make_config(dit_attn_impl=types.SimpleNamespace(value="mindie")) + with self.assertRaisesRegex(RuntimeError, "explicitly requested"): + self.adapter.validate_config(config) + + def test_basic_bf16_allows_mindie_to_be_optional(self): + FakePlatform.current_capabilities = PlatformCapabilities(device=True) + self.adapter.validate_config(make_config()) + + def test_quantization_capability_is_not_silently_downgraded(self): + FakePlatform.current_capabilities = PlatformCapabilities(device=True, mindie=True) + with self.assertRaisesRegex(RuntimeError, "MXFP8"): + self.adapter.validate_config(make_config(quantization=QuantizationConfig(linear="fp8"))) + + def test_quantization_modes_map_to_mindie_algorithms(self): + class QuantAlgorithm: + W8A8_MXFP8 = "mxfp8" + W4A4_MXFP4_DYNAMIC = "mxfp4" + W8A8_DYNAMIC = "dynamic" + W16A16 = "bf16" + FP8_DYNAMIC = "fp8_attention" + + records = [] + + class OnlineQuantConfig: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + records.append(self) + + quantization_module = types.SimpleNamespace( + OnlineQuantConfig=OnlineQuantConfig, + quantize=lambda module, **kwargs: module, + ) + mode_module = types.SimpleNamespace(QuantAlgorithm=QuantAlgorithm) + + for linear, expected in (("fp8", "mxfp8"), ("int4", "mxfp4")): + with self.subTest(linear=linear), mock.patch.object( + ascend_qwen.importlib, + "import_module", + side_effect=lambda name: mode_module if name.endswith(".mode") else quantization_module, + ): + module = nn.Sequential(nn.Linear(4, 4)) + self.adapter.prepare_component( + "dit", + module, + make_config(quantization=QuantizationConfig(linear=linear)), + ) + self.assertEqual(records[-1].quant_type, expected) + + def test_attention_only_quantization_keeps_linear_layers_unquantized(self): + class QuantAlgorithm: + W8A8_DYNAMIC = "dynamic" + W16A16 = "bf16" + FP8_DYNAMIC = "fp8_attention" + + record = {} + + class OnlineQuantConfig: + def __init__(self, **kwargs): + record.update(kwargs) + + quantization_module = types.SimpleNamespace( + OnlineQuantConfig=OnlineQuantConfig, + quantize=lambda module, **kwargs: module, + ) + mode_module = types.SimpleNamespace(QuantAlgorithm=QuantAlgorithm) + module = nn.Sequential(nn.Linear(4, 4), nn.Sequential(nn.Linear(4, 4))) + with mock.patch.object( + ascend_qwen.importlib, + "import_module", + side_effect=lambda name: mode_module if name.endswith(".mode") else quantization_module, + ): + self.adapter.prepare_component( + "dit", + module, + make_config(quantization=QuantizationConfig(attention="fp8")), + ) + + self.assertEqual(record["fallback_layers"], {"0": "bf16", "1.0": "bf16"}) + self.assertEqual(record["fa_quant_type"], "fp8_attention") + + def test_fp8_attention_is_installed_through_processor_slot(self): + class QuantAlgorithm: + W8A8_DYNAMIC = "dynamic" + W16A16 = "bf16" + FP8_DYNAMIC = "fp8_attention" + + class AttentionLayer(nn.Module): + def __init__(self): + super().__init__() + self.fa_quant = mock.Mock() + self.processor = None + + def set_attention_processor(self, processor): + self.processor = processor + + class OnlineQuantConfig: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + quantization_module = types.SimpleNamespace( + OnlineQuantConfig=OnlineQuantConfig, + quantize=lambda module, **kwargs: module, + ) + mode_module = types.SimpleNamespace(QuantAlgorithm=QuantAlgorithm) + module = AttentionLayer() + with mock.patch.object( + ascend_qwen.importlib, + "import_module", + side_effect=lambda name: mode_module if name.endswith(".mode") else quantization_module, + ): + self.adapter.prepare_component( + "dit", + module, + make_config(quantization=QuantizationConfig(attention="fp8")), + ) + + self.assertIsInstance(module.processor, ascend_qwen.MindIEFP8AttentionProcessor) + + def test_timestep_and_dynamic_weight_hooks(self): + timestep_manager = mock.Mock() + quantization_module = types.SimpleNamespace(TimestepManager=timestep_manager) + self.adapter._quantization_enabled = True + with mock.patch.object(ascend_qwen.importlib, "import_module", return_value=quantization_module): + self.adapter.before_denoise_step(7) + timestep_manager.set_timestep_idx.assert_called_once_with(7) + + with self.assertRaisesRegex(ValueError, "Dynamic LoRA/ControlNet"): + self.adapter.validate_dynamic_weights(make_config(use_torch_compile=True)) + with self.assertRaisesRegex(ValueError, "Dynamic LoRA/ControlNet"): + self.adapter.validate_dynamic_weights( + make_config(quantization=QuantizationConfig(linear="fp8")) + ) + + def test_compile_routes_backend_to_repeated_blocks(self): + module = mock.Mock() + module.compile_repeated_blocks = mock.Mock() + self.assertIs(self.adapter.compile_component("dit", module), module) + module.compile_repeated_blocks.assert_called_once_with(backend="mindie") + + +class TestOtherRuntimeAdapters(unittest.TestCase): + def test_default_compile_preserves_empty_torch_compile_kwargs(self): + platform = types.SimpleNamespace(compile_kwargs=lambda: {}) + module = mock.Mock() + module.compile_repeated_blocks = mock.Mock() + adapter = DefaultRuntimeAdapter(platform, "qwen_image") + self.assertIs(adapter.compile_component("dit", module), module) + module.compile_repeated_blocks.assert_called_once_with() + + def test_cuda_quantization_config_uses_existing_fp8_path(self): + adapter = CudaQwenImageRuntimeAdapter(types.SimpleNamespace(), "qwen_image") + module = nn.Linear(4, 4) + config = make_config( + device="cuda", + quantization=QuantizationConfig(backend="native", linear="fp8"), + ) + with mock.patch("diffsynth_engine.platforms.qwen.enable_fp8_linear") as enable_fp8: + self.assertIs(adapter.prepare_component("dit", module, config), module) + enable_fp8.assert_called_once_with(module) + + def test_cuda_rejects_mindie_quantization(self): + adapter = CudaQwenImageRuntimeAdapter(types.SimpleNamespace(), "qwen_image") + config = make_config(device="cuda", quantization=QuantizationConfig(backend="mindie")) + with self.assertRaisesRegex(ValueError, "requires an Ascend"): + adapter.validate_config(config) + + +if __name__ == "__main__": + unittest.main()