diff --git a/.gitignore b/.gitignore index 82691f8..13a6f17 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,11 @@ dist/ CLAUDE.md .claude/ .kiro/ + +# Performance benchmarks & results (internal use only) +benchmarks/ +results/ + +# PR artifacts +PR_BODY.md +.pr272.diff diff --git a/diffsynth_engine/args.py b/diffsynth_engine/args.py index d445255..495e76b 100644 --- a/diffsynth_engine/args.py +++ b/diffsynth_engine/args.py @@ -117,6 +117,11 @@ def parse_cli_args() -> Dict[str, Any]: action="store_true", help="Compile repeated transformer blocks with torch.compile", ) + optimization_group.add_argument( + "--compile-ffn", + action="store_true", + help="Compile only FFN (MLP) blocks with torch.compile (finer-grained than --use-torch-compile)", + ) # Parallelism configuration group parallel_group = parser.add_argument_group("Parallelism Configuration") @@ -184,6 +189,7 @@ def parse_cli_args() -> Dict[str, Any]: # Optimization configuration args_dict["use_torch_compile"] = args.use_torch_compile + args_dict["compile_ffn"] = args.compile_ffn # Parallelism configuration args_dict["parallelism"] = args.parallelism diff --git a/diffsynth_engine/configs/base.py b/diffsynth_engine/configs/base.py index 226dff2..1143e16 100644 --- a/diffsynth_engine/configs/base.py +++ b/diffsynth_engine/configs/base.py @@ -6,6 +6,7 @@ from diffsynth_engine.layers.attention import AttentionType from diffsynth_engine.registry import get_attn_backend from diffsynth_engine.utils import logging +from diffsynth_engine.platforms import get_device_type, resolve_platform logger = logging.get_logger(__name__) @@ -27,7 +28,7 @@ class PipelineConfig: model_dtype: torch.dtype = torch.bfloat16 text_encoder_dtype: torch.dtype = torch.bfloat16 vae_dtype: torch.dtype = torch.float32 - device: str | torch.device = "cuda" + device: str | torch.device = "auto" pipeline_class_name: str | None = None @@ -42,6 +43,7 @@ class PipelineConfig: # optimization use_torch_compile: bool = False + compile_ffn: bool = False # parallelism parallelism: int = 1 @@ -62,6 +64,7 @@ def __post_init__(self): self.attn_type = str(self.attn_type) init_parallel_config(self) validate_attn_config(self) + init_device_config(self) def init_parallel_config(config: PipelineConfig): @@ -109,3 +112,11 @@ def validate_attn_config(config: PipelineConfig): if config.sp_ring_degree is not None and config.sp_ring_degree > 1: if not attn_backend.supports_ring_attention(): raise ValueError(f"Attention backend {config.attn_type!r} does not support ring attention.") + + +def init_device_config(config: PipelineConfig): + if config.device is None or (isinstance(config.device, str) and config.device.lower() in ("auto", "")): + config.device = get_device_type() + return + # Validate that the explicit device type is registered (fail fast at construction). + resolve_platform(config.device) diff --git a/diffsynth_engine/distributed/parallel_state.py b/diffsynth_engine/distributed/parallel_state.py index a2f91d7..9f9b01e 100644 --- a/diffsynth_engine/distributed/parallel_state.py +++ b/diffsynth_engine/distributed/parallel_state.py @@ -13,7 +13,6 @@ import torch import torch.distributed -from torch.cuda import device_count, set_device from diffsynth_engine.distributed.group_coordinator import ( GroupCoordinator, @@ -22,7 +21,11 @@ ) from diffsynth_engine.utils import logging from diffsynth_engine.utils.constants import IDLE_TIMEOUT_SEC -from diffsynth_engine.utils.platform import get_torch_distributed_backend +from diffsynth_engine.utils.platform import ( + device_count, + get_torch_distributed_backend, + set_device, +) logger = logging.get_logger(__name__) @@ -425,10 +428,22 @@ def init_distributed_environment( distributed_init_method, backend, ) + # local_rank is not available in torch ProcessGroup, + # see https://github.com/pytorch/pytorch/issues/122816 + if local_rank == -1: + # local rank not set, this usually happens in single-node + # setting, where we can use rank as local rank + if distributed_init_method == "env://": + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + else: + local_rank = rank if rank >= 0 else 0 + if not torch.distributed.is_initialized(): assert distributed_init_method is not None, ( "distributed_init_method must be provided when initializing distributed environment" ) + # Bind device before init_process_group (required by HCCL on Ascend). + set_device(local_rank % max(device_count(), 1)) # this backend is used for WORLD torch.distributed.init_process_group( backend=backend, @@ -436,17 +451,7 @@ def init_distributed_environment( world_size=world_size, rank=rank, ) - set_device(torch.distributed.get_rank() % device_count()) - # set the local rank - # local_rank is not available in torch ProcessGroup, - # see https://github.com/pytorch/pytorch/issues/122816 - if local_rank == -1: - # local rank not set, this usually happens in single-node - # setting, where we can use rank as local rank - if distributed_init_method == "env://": - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - else: - local_rank = rank + set_device(torch.distributed.get_rank() % max(device_count(), 1)) global _WORLD if _WORLD is None: ranks = list(range(torch.distributed.get_world_size())) diff --git a/diffsynth_engine/engine.py b/diffsynth_engine/engine.py index 13bf802..67f79d9 100644 --- a/diffsynth_engine/engine.py +++ b/diffsynth_engine/engine.py @@ -1,7 +1,6 @@ from typing import Any import torch.multiprocessing as mp -from torch.cuda import set_device from diffsynth_engine.configs import PipelineConfig from diffsynth_engine.registry import ( @@ -9,6 +8,7 @@ get_pipeline_class_name, ) from diffsynth_engine.utils import logging +from diffsynth_engine.utils.platform import set_device from diffsynth_engine.utils.torch_profiler import TorchProfiler from diffsynth_engine.worker import run_worker_loop diff --git a/diffsynth_engine/layers/__init__.py b/diffsynth_engine/layers/__init__.py index e69de29..3205671 100644 --- a/diffsynth_engine/layers/__init__.py +++ b/diffsynth_engine/layers/__init__.py @@ -0,0 +1,5 @@ +from diffsynth_engine.layers.transformer_helper import RMSNorm + +__all__ = [ + "RMSNorm", +] diff --git a/diffsynth_engine/layers/attention/__init__.py b/diffsynth_engine/layers/attention/__init__.py index ffc2b45..0fe03d1 100644 --- a/diffsynth_engine/layers/attention/__init__.py +++ b/diffsynth_engine/layers/attention/__init__.py @@ -1,9 +1,13 @@ from .backends.abstract import AttentionMetadata, AttentionType +from .factory import create_parallel_attention from .layer import LocalAttention, USPAttention +from .ascend_long_context import AscendLongContextAttention __all__ = [ "AttentionType", "AttentionMetadata", "LocalAttention", "USPAttention", + "AscendLongContextAttention", + "create_parallel_attention", ] diff --git a/diffsynth_engine/layers/attention/ascend_long_context.py b/diffsynth_engine/layers/attention/ascend_long_context.py new file mode 100644 index 0000000..138db7e --- /dev/null +++ b/diffsynth_engine/layers/attention/ascend_long_context.py @@ -0,0 +1,358 @@ +# Adapted from https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +from torch import distributed as dist +import torch +import torch.nn as nn + +from diffsynth_engine.distributed.comm import SeqAllToAll4D +from diffsynth_engine.distributed.parallel_state import ( + get_sp_group, +) +from diffsynth_engine.forward_context import ForwardContext, get_forward_context +from diffsynth_engine.registry import get_attn_backend + + +class AscendLongContextAttention(nn.Module): + # Single dedicated communication stream shared by all instances (one per transformer + # block, e.g. 60 in Qwen-Image), so only one `stream2` is allocated per device. + _shared_comm_stream = None + + def __init__( + self, + num_heads: int = 24, + head_size: int = 128, + softmax_scale: float | None = None, + causal: bool = False, + num_kv_heads: int | None = None, + attn_type: str = "mindie", + scatter_idx: int = 2, + gather_idx: int = 1, + fa_head_loop: int | None = None, + **extra_impl_args, + ) -> None: + super().__init__() + from diffsynth_engine.platforms import AscendPlatform + + if num_kv_heads is None: + num_kv_heads = num_heads + + self.scatter_idx = scatter_idx + self.gather_idx = gather_idx + self.num_heads = num_heads + self.head_size = head_size + self.num_kv_heads = num_kv_heads + + self.ulysses_pg = get_sp_group().ulysses_group + self.sp_ulysses_degree = get_sp_group().ulysses_world_size + self.sp_ring_degree = get_sp_group().ring_world_size + + + self.fa_alltoall_overlap = AscendPlatform.fa_alltoall_overlap + self.fa_alltoall_cut = AscendPlatform.fa_alltoall_cut + if fa_head_loop is not None: + self.fa_head_loop = fa_head_loop + elif self.fa_alltoall_cut > 1: + self.fa_head_loop = self.fa_alltoall_cut + elif self.fa_alltoall_overlap > 1: + self.fa_head_loop = self.fa_alltoall_overlap + else: + self.fa_head_loop = self.num_heads // self.sp_ulysses_degree + + if self.fa_alltoall_overlap > 1 and self.fa_alltoall_cut <= 1: + if AscendLongContextAttention._shared_comm_stream is None: + AscendLongContextAttention._shared_comm_stream = torch.npu.Stream() + self.stream2 = AscendLongContextAttention._shared_comm_stream + self.event = [] + for i in range(self.fa_head_loop): + self.event.append(torch.npu.Event()) + + self.attn_type = str(attn_type) if attn_type is not None else None + attn_backend = get_attn_backend(attn_type) + if not attn_backend.supports_head_size(head_size): + raise ValueError(f"Attention backend {attn_type!r} does not support head size {head_size}.") + + impl_cls = attn_backend.get_impl_cls() + self.attn_impl = impl_cls( + num_heads=num_heads, + head_size=head_size, + softmax_scale=softmax_scale, + causal=causal, + num_kv_heads=num_kv_heads, + **extra_impl_args, + ) + + # TODO: currunt MindIE only support Ulysses + if self.sp_ring_degree > 1: + raise RuntimeError( + "NPU MindIE attention currently supports Ulysses only " + f"(sp_ring_degree must be 1, got {self.sp_ring_degree})" + ) + + + def _run_attention(self, q, k, v, **attn_kwargs): + return self.attn_impl.forward(q, k, v, **attn_kwargs) + + @staticmethod + def all_to_all_4D_pre(input: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None): + assert ( + input.dim() == 4 + ), f"input must be 4D tensor, got {input.dim()} and shape {input.shape}" + + seq_world_size = dist.get_world_size(group) + + if scatter_idx == 2 and gather_idx == 1: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs) + bs, shard_seqlen, hc, hs = input.shape + seqlen = shard_seqlen * seq_world_size + shard_hc = hc // seq_world_size + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs) + input_t = ( + input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs) + .transpose(0, 2) + .contiguous() + ) + + return input_t + + elif scatter_idx == 1 and gather_idx == 2: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs) + bs, seqlen, shard_hc, hs = input.shape + hc = shard_hc * seq_world_size + shard_seqlen = seqlen // seq_world_size + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs) + input_t = ( + input.reshape(bs, seq_world_size, shard_seqlen, shard_hc, hs) + .transpose(0, 3) + .transpose(0, 1) + .contiguous() + .reshape(seq_world_size, shard_hc, shard_seqlen, bs, hs) + ) + + return input_t + else: + raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2") + + @staticmethod + def all_to_all_4D_after(input: torch.tensor, output: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, + group=None): + seq_world_size = dist.get_world_size(group) + + if scatter_idx == 2 and gather_idx == 1: + bs, shard_seqlen, hc, hs = input.shape + + seqlen = shard_seqlen * seq_world_size + shard_hc = hc // seq_world_size + + output = output.reshape(seqlen, bs, shard_hc, hs) + + # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs) + output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs) + return output + elif scatter_idx == 1 and gather_idx == 2: + bs, seqlen, shard_hc, hs = input.shape + hc = shard_hc * seq_world_size + shard_seqlen = seqlen // seq_world_size + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(hc, shard_seqlen, bs, hs) + + # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs) + output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs) + return output + else: + raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2") + + + @staticmethod + def split_qkv_by_head(query, key, value, sp_ulysses_degree, loop_time): + """Split Q/K/V along head dim into chunks for insertcomm / blockattn.""" + _, _, head_count, _ = query.shape + if head_count % sp_ulysses_degree != 0: + raise ValueError( + f"head_count must be divisible by ulysses world size, " + f"got head_count={head_count}, sp_ulysses_degree={sp_ulysses_degree}" + ) + heads_per_rank = head_count // sp_ulysses_degree + if heads_per_rank % loop_time != 0: + raise ValueError( + f"heads_per_rank must be divisible by loop_time={loop_time}, " + f"got heads_per_rank={heads_per_rank}" + ) + global_chunk_heads = heads_per_rank // loop_time * sp_ulysses_degree + return ( + query.split(global_chunk_heads, dim=2), + key.split(global_chunk_heads, dim=2), + value.split(global_chunk_heads, dim=2), + ) + + + @torch.compiler.disable + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + """forward + + Arguments: + query (torch.Tensor): query input to the layer + key (torch.Tensor): key input to the layer + value (torch.Tensor): value input to the layer + + Returns: + * output (torch.Tensor): context output + """ + # Check input shapes + assert query.dim() == 4 and key.dim() == 4 and value.dim() == 4, "Expected 4D tensors" + + forward_context: ForwardContext = get_forward_context() + attn_metadata = forward_context.attn_metadata + + attn_kwargs = {"attn_metadata": attn_metadata} + attn_kwargs.update(kwargs) + + output = None + + if self.fa_alltoall_cut <= 1 and self.fa_alltoall_overlap <= 1: + # baseline: both 0 / both 1 / a single 1 all mean "not enabled" (1 chunk = no split) + # 3 X (bs, seq_len/N, head_cnt, head_size) -> 3 X (bs, seq_len, head_cnt/N, head_size) + # scatter 2, gather 1 + query_layer = SeqAllToAll4D.apply( + self.ulysses_pg, query, self.scatter_idx, self.gather_idx + ) + key_layer = SeqAllToAll4D.apply( + self.ulysses_pg, key, self.scatter_idx, self.gather_idx + ) + value_layer = SeqAllToAll4D.apply( + self.ulysses_pg, value, self.scatter_idx, self.gather_idx + ) + + out = self._run_attention(query_layer, key_layer, value_layer, **attn_kwargs) + # (bs, seq_len, head_cnt/N, head_size) -> (bs, seq_len/N, head_cnt, head_size) + # scatter 1, gather 2 + output = SeqAllToAll4D.apply( + self.ulysses_pg, out, self.gather_idx, self.scatter_idx + ) + elif self.fa_alltoall_cut > 1: # fa_alltoall_cut + # Split heads into chunks (loop_time = fa_alltoall_cut), full Ulysses round-trip per chunk. + q_chunks, k_chunks, v_chunks = self.split_qkv_by_head( + query, key, value, self.sp_ulysses_degree, self.fa_head_loop + ) + output_chunks = [] + for q_chunk, k_chunk, v_chunk in zip(q_chunks, k_chunks, v_chunks): + query_layer = SeqAllToAll4D.apply( + self.ulysses_pg, q_chunk, self.scatter_idx, self.gather_idx + ) + key_layer = SeqAllToAll4D.apply( + self.ulysses_pg, k_chunk, self.scatter_idx, self.gather_idx + ) + value_layer = SeqAllToAll4D.apply( + self.ulysses_pg, v_chunk, self.scatter_idx, self.gather_idx + ) + out = self._run_attention(query_layer, key_layer, value_layer, **attn_kwargs) + out = SeqAllToAll4D.apply( + self.ulysses_pg, out, self.gather_idx, self.scatter_idx + ) + output_chunks.append(out) + output = torch.cat(output_chunks, dim=2) + elif self.fa_alltoall_overlap > 1 : # fa_alltoall_overlap + # B, S/sp, N/tp, D + # Refresh the current stream here: __init__ runs at model-build time and may capture a + # different stream than the one forward actually executes on (e.g. under a stream context, + # pipeline/CFG stream switching, or CUDA-graph capture). Pinning the build-time stream would + # break event/stream synchronization in the overlap pipeline. + self.current_stream = torch.npu.current_stream() + query_layer_list, key_layer_list, value_layer_list = self.split_qkv_by_head( + query, key, value, self.sp_ulysses_degree, self.fa_head_loop + ) + for_loop = len(query_layer_list) + + # scatter 2, gather 1 + output_fa = [] + q_event = torch.npu.Event() + k_event = torch.npu.Event() + v_event = torch.npu.Event() + q_lists, k_lists, v_lists, kv_lists = [], [], [], [] + + for i in range(0, for_loop): + input_q = self.all_to_all_4D_pre(query_layer_list[i], self.scatter_idx, self.gather_idx, + self.ulysses_pg) + q_event.record() + with torch.npu.stream(self.stream2): + self.stream2.wait_event(q_event) + query_layer = torch.empty_like(input_q) + dist.all_to_all_single(query_layer, input_q, group=self.ulysses_pg) + + input_k = self.all_to_all_4D_pre(key_layer_list[i], self.scatter_idx, self.gather_idx, + self.ulysses_pg) + input_v = self.all_to_all_4D_pre(value_layer_list[i], self.scatter_idx, self.gather_idx, + self.ulysses_pg) + v_event.record() + + with torch.npu.stream(self.stream2): + self.stream2.wait_event(v_event) + key_layer = torch.empty_like(input_k) + dist.all_to_all_single(key_layer, input_k, group=self.ulysses_pg) + + value_layer = torch.empty_like(input_v) + dist.all_to_all_single(value_layer, input_v, group=self.ulysses_pg) + k_event.record() + + q_lists.append(query_layer) + k_lists.append(key_layer) + v_lists.append(value_layer) + + k_lists[i] = self.all_to_all_4D_after(key_layer_list[i], k_lists[i], self.scatter_idx, + self.gather_idx, self.ulysses_pg) + v_lists[i] = self.all_to_all_4D_after(value_layer_list[i], v_lists[i], self.scatter_idx, + self.gather_idx, self.ulysses_pg) + q_event.record() + with torch.npu.stream(self.stream2): + self.stream2.wait_event(q_event) + self.event[i].record() + q_lists[i] = self.all_to_all_4D_after(query_layer_list[i], q_lists[i], self.scatter_idx, self.gather_idx, self.ulysses_pg) + + + for i in range(0, for_loop): + # fa + self.current_stream.wait_event(self.event[i]) + + out = self._run_attention(q_lists[i], k_lists[i], v_lists[i], **attn_kwargs) + kv_lists.append(out) + input_t = self.all_to_all_4D_pre(out, self.gather_idx, self.scatter_idx, self.ulysses_pg) + q_event.record() + + with torch.npu.stream(self.stream2): + self.stream2.wait_event(q_event) + output = torch.empty_like(input_t) + dist.all_to_all_single(output, input_t, group=self.ulysses_pg) + self.event[i].record() + output_fa.append(output) + + for i in range(for_loop): + self.current_stream.wait_event(self.event[i]) + output_fa[i] = self.all_to_all_4D_after(kv_lists[i], output_fa[i], self.gather_idx, self.scatter_idx, self.ulysses_pg) + output = torch.cat(output_fa, dim=2) + else: + raise RuntimeError( + f"Invalid configuration: fa_alltoall_cut={self.fa_alltoall_cut}, fa_alltoall_overlap={self.fa_alltoall_overlap}" + ) + return output + +_ASCEND_LONGCTX_ATTN: Optional[AscendLongContextAttention] = None + + +def _get_ascend_long_context_attn() -> AscendLongContextAttention: + global _ASCEND_LONGCTX_ATTN + if _ASCEND_LONGCTX_ATTN is None: + _ASCEND_LONGCTX_ATTN = AscendLongContextAttention() + return _ASCEND_LONGCTX_ATTN diff --git a/diffsynth_engine/layers/attention/backends/abstract.py b/diffsynth_engine/layers/attention/backends/abstract.py index e9b5fdd..35be745 100644 --- a/diffsynth_engine/layers/attention/backends/abstract.py +++ b/diffsynth_engine/layers/attention/backends/abstract.py @@ -26,6 +26,7 @@ class AttentionType(str, enum.Enum): SAGE2 = "sage2" SAGE3 = "sage3" SPARGE = "sparge" + MINDIE = "mindie" def __str__(self) -> str: return self.value diff --git a/diffsynth_engine/layers/attention/backends/mindie_attn.py b/diffsynth_engine/layers/attention/backends/mindie_attn.py new file mode 100644 index 0000000..81bb243 --- /dev/null +++ b/diffsynth_engine/layers/attention/backends/mindie_attn.py @@ -0,0 +1,117 @@ +import torch + +from diffsynth_engine.layers.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder, + AttentionType, +) +from diffsynth_engine.utils import logging + +logger = logging.get_logger(__name__) + + +class MindieAttentionMetadataBuilder(AttentionMetadataBuilder): + def __init__(self) -> None: + pass + + def build(self, **kwargs) -> AttentionMetadata: + return AttentionMetadata() + + +class MindieAttentionBackend(AttentionBackend): + @staticmethod + def check_availability() -> None: + from diffsynth_engine.platforms import AscendPlatform + + if not AscendPlatform.supports("device"): + error_msg = "MindIE attention requires an available Ascend NPU device." + logger.error(error_msg) + raise RuntimeError(error_msg) + if not AscendPlatform.supports("mindie_attention"): + error_msg = ( + "MindIE attention backend is not available. " + "Install MindIE-SD 3.x matching the current torch_npu and CANN versions, " + "and ensure mindiesd.layers.flash_attn.attention_forward works on NPU." + ) + logger.error(error_msg) + raise RuntimeError(error_msg) + + @staticmethod + def get_type() -> str: + return str(AttentionType.MINDIE) + + @staticmethod + def get_impl_cls() -> type["AttentionImpl"]: + return MindieAttentionImpl + + @staticmethod + def get_metadata_cls() -> type["AttentionMetadata"]: + return AttentionMetadata + + @staticmethod + def get_builder_cls() -> type["AttentionMetadataBuilder"]: + return MindieAttentionMetadataBuilder + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [] + + @classmethod + def supports_ring_attention(cls) -> bool: + return False + + +class MindieAttentionImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float | None = None, + causal: bool = False, + num_kv_heads: int | None = None, + **extra_impl_args, + ) -> None: + if num_kv_heads is None: + num_kv_heads = num_heads + self.num_kv_groups = num_heads // num_kv_heads + self.causal = causal + self.softmax_scale = softmax_scale + self.num_heads = num_heads + self.head_size = head_size + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + attn_metadata: AttentionMetadata | None = None, + **kwargs, + ) -> torch.Tensor: + from mindiesd.layers.flash_attn.attention_forward import attention_forward + + # MindIE attention_forward 没有 causal 参数，只接受 attn_mask（布尔张量， + # True 表示保留/参与计算，False 表示被屏蔽）。因此当请求 causal 且调用方未显式 + # 传入 attn_mask 时，手动构造下三角 causal mask 并通过 attn_mask 传入。 + # layout 为 "BSND"，故 q_seqlen = query.shape[1]、kv_seqlen = key.shape[1]。 + if self.causal and attn_mask is None: + q_seqlen = query.shape[1] + kv_seqlen = key.shape[1] + attn_mask = torch.tril( + torch.ones(q_seqlen, kv_seqlen, dtype=torch.bool, device=query.device) + ) + + return attention_forward( + query=query, + key=key, + value=value, + attn_mask=attn_mask, + scale=self.softmax_scale, + fused=True, + head_first=False, + opt_mode="manual", + op_type="fused_attn_score", + layout="BSND", + ) diff --git a/diffsynth_engine/layers/attention/factory.py b/diffsynth_engine/layers/attention/factory.py new file mode 100644 index 0000000..c376a66 --- /dev/null +++ b/diffsynth_engine/layers/attention/factory.py @@ -0,0 +1,76 @@ +"""Attention module factory for platform-aware parallel attention creation.""" + +import torch.nn as nn + + +def create_parallel_attention( + num_heads: int, + head_size: int, + attn_type: str | None = None, + softmax_scale: float | None = None, + causal: bool = False, + num_kv_heads: int | None = None, + scatter_idx: int = 2, + gather_idx: int = 1, + **extra_impl_args, +) -> nn.Module: + """ + 根据平台能力和并行配置创建合适的序列并行 attention 模块。 + + - NPU MindIE (attn_type == "mindie") + SP initialized + Ulysses-only (ring degree == 1): + AscendLongContextAttention + - 其他: USPAttention + + AscendLongContextAttention 目前仅支持 MindIE backend 且仅支持 Ulysses 序列并行 + (sp_ring_degree == 1)。因此只有在调用方显式请求 "mindie" 且未启用 ring 并行时才路由到 + NPU 长上下文实现，其余情况一律 fallback 到 USPAttention。 + + Args: + num_heads: attention head 数量 + head_size: 每个 head 的维度 + attn_type: attention backend 类型 (如 "mindie", "sdpa", "fa2" 等) + softmax_scale: softmax 缩放系数 + causal: 是否使用因果 attention + num_kv_heads: KV head 数量 (GQA) + scatter_idx: Ulysses scatter 维度索引 + gather_idx: Ulysses gather 维度索引 + **extra_impl_args: 传递给底层 attention 实现的额外参数 + + Returns: + nn.Module: 配置好的 attention 模块 + """ + # Lazy imports to avoid circular dependencies + from diffsynth_engine.distributed.parallel_state import ( + get_ring_parallel_world_size, + is_sp_group_initialized, + ) + from diffsynth_engine.utils.platform import is_mindie_sd_available + + common_kwargs = dict( + num_heads=num_heads, + head_size=head_size, + softmax_scale=softmax_scale, + causal=causal, + num_kv_heads=num_kv_heads, + attn_type=attn_type, + scatter_idx=scatter_idx, + gather_idx=gather_idx, + **extra_impl_args, + ) + + # AscendLongContextAttention 只服务 MindIE backend，且只支持 Ulysses（ring degree == 1）。 + # 因此必须校验 attn_type，并确认 ring 配置未启用，否则 fallback 到 USPAttention。 + # 注意短路顺序：get_ring_parallel_world_size() 依赖 SP 已初始化，必须放在 is_sp_group_initialized() 之后。 + if ( + is_mindie_sd_available() + and is_sp_group_initialized() + and str(attn_type) == "mindie" + and get_ring_parallel_world_size() == 1 + ): + from diffsynth_engine.layers.attention.ascend_long_context import AscendLongContextAttention + + return AscendLongContextAttention(**common_kwargs) + else: + from diffsynth_engine.layers.attention.layer import USPAttention + + return USPAttention(**common_kwargs) diff --git a/diffsynth_engine/layers/attention/layer.py b/diffsynth_engine/layers/attention/layer.py index 1bde443..8b86c51 100644 --- a/diffsynth_engine/layers/attention/layer.py +++ b/diffsynth_engine/layers/attention/layer.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 +from torch import distributed as dist import torch import torch.nn as nn @@ -101,6 +102,7 @@ def __init__( self.num_kv_heads = num_kv_heads self.scatter_idx = scatter_idx self.gather_idx = gather_idx + self.attn_type = str(attn_type) if attn_type is not None else None attn_backend = get_attn_backend(attn_type) if not attn_backend.supports_head_size(head_size): @@ -147,6 +149,12 @@ def forward( ulysses_parallel_world_size = get_ulysses_parallel_world_size() if is_sp_group_initialized() else 1 ring_parallel_world_size = get_ring_parallel_world_size() if is_sp_group_initialized() else 1 + if ring_parallel_world_size > 1 and self.attn_type == "mindie": + raise RuntimeError( + "NPU MindIE attention currently supports Ulysses only " + f"(sp_ring_degree must be 1, got {ring_parallel_world_size})" + ) + if ulysses_parallel_world_size > 1: q = SeqAllToAll4D.apply(get_sp_group().ulysses_group, q, self.scatter_idx, self.gather_idx) k = SeqAllToAll4D.apply(get_sp_group().ulysses_group, k, self.scatter_idx, self.gather_idx) @@ -160,4 +168,4 @@ def forward( if ulysses_parallel_world_size > 1: output = SeqAllToAll4D.apply(get_sp_group().ulysses_group, output, self.gather_idx, self.scatter_idx) - return output + return output \ No newline at end of file diff --git a/diffsynth_engine/layers/lora/linear.py b/diffsynth_engine/layers/lora/linear.py index 17d633b..e3cfbaf 100644 --- a/diffsynth_engine/layers/lora/linear.py +++ b/diffsynth_engine/layers/lora/linear.py @@ -2,6 +2,7 @@ import torch.nn as nn from diffsynth_engine.utils import logging +from diffsynth_engine.utils.platform import pin_memory logger = logging.get_logger(__name__) @@ -180,9 +181,7 @@ def _save_original_weight(self): if self._original_weight is not None: return weight = self.base_layer.weight.data - self._original_weight = weight.detach().cpu().clone() - if not torch.backends.mps.is_available(): - self._original_weight = self._original_weight.pin_memory() + self._original_weight = pin_memory(weight.detach().cpu().clone()) def merge_loras(self, chunked: bool = False, high_precision: bool = True) -> list[str]: """Merge active LoRA models into the wrapped base layer weight. diff --git a/diffsynth_engine/layers/transformer_helper.py b/diffsynth_engine/layers/transformer_helper.py new file mode 100644 index 0000000..aaaeb5d --- /dev/null +++ b/diffsynth_engine/layers/transformer_helper.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reusable transformer helper layers (fused RMSNorm for NPU, manual fallback elsewhere).""" + +import torch +import torch.nn as nn + +from diffsynth_engine.platforms.ops import fused_rms_norm + + +class RMSNorm(nn.Module): + """RMSNorm over the last dim. + + API-compatible with `diffusers.models.normalization.RMSNorm` + (dim, eps, elementwise_affine), so existing checkpoints (weight key) load + unchanged. When `elementwise_affine` is enabled, delegates to + `fused_rms_norm` which automatically dispatches to the platform-optimal + implementation; otherwise falls back to reference fp32 math. + """ + + def __init__(self, dim, eps=1e-6, elementwise_affine=True): + super().__init__() + self.dim = dim + self.eps = eps + self.elementwise_affine = elementwise_affine + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim)) + else: + self.register_parameter("weight", None) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + if self.elementwise_affine: + return fused_rms_norm(x, self.weight, self.eps) + + output = self._norm(x.float()).type_as(x) + return output + + def extra_repr(self) -> str: + return f"dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}" \ No newline at end of file diff --git a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py index 895c2e3..6755462 100644 --- a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py +++ b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py @@ -24,7 +24,7 @@ from diffusers.configuration_utils import register_to_config from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.modeling_outputs import Transformer2DModelOutput -from diffusers.models.normalization import AdaLayerNormContinuous, RMSNorm +from diffusers.models.normalization import AdaLayerNormContinuous from diffsynth_engine.distributed.parallel_state import ( get_tensor_model_parallel_world_size, @@ -32,9 +32,11 @@ ) from diffsynth_engine.distributed.utils import sequence_parallel_shard, sequence_parallel_unshard from diffsynth_engine.forward_context import get_forward_context -from diffsynth_engine.layers.attention import USPAttention +from diffsynth_engine.layers import RMSNorm +from diffsynth_engine.layers.attention.factory import create_parallel_attention from diffsynth_engine.layers.tensor_parallel import ColumnParallelLinear, RowParallelLinear, TPFeedForward from diffsynth_engine.models.base import DiffusionModel +from diffsynth_engine.platforms.ops import fused_layernorm_scale_shift, fused_rotary_embedding from diffsynth_engine.utils import logging logger = logging.get_logger(__name__) @@ -60,32 +62,7 @@ def apply_rotary_emb_qwen( Returns: Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. """ - if use_real: - cos, sin = freqs_cis # [S, D] - cos = cos[None, None] - sin = sin[None, None] - cos, sin = cos.to(x.device), sin.to(x.device) - - if use_real_unbind_dim == -1: - # Used for flux, cogvideox, hunyuan-dit - x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] - x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) - elif use_real_unbind_dim == -2: - # Used for Stable Audio, OmniGen, CogView4 and Cosmos - x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, S, H, D//2] - x_rotated = torch.cat([-x_imag, x_real], dim=-1) - else: - raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.") - - out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) - - return out - else: - x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) - freqs_cis = freqs_cis.unsqueeze(1) - x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) - - return x_out.type_as(x) + return fused_rotary_embedding(x, freqs_cis, use_real=use_real, use_real_unbind_dim=use_real_unbind_dim) def compute_text_seq_len_from_mask( @@ -462,7 +439,7 @@ def __init__( # USPAttention for joint attention computation forward_context = get_forward_context() - self.usp_attn = USPAttention( + self.usp_attn = create_parallel_attention( num_heads=self.heads, head_size=attention_head_dim, attn_type=forward_context.attn_type, @@ -584,32 +561,25 @@ def __init__( self.zero_cond_t = zero_cond_t - def _modulate(self, x, mod_params, index=None): - """Apply modulation to input tensor""" - # x: b l d, shift: b d, scale: b d, gate: b d + def _split_mod_params(self, mod_params, index=None): + """Split modulation params into shift/scale/gate, optionally indexed for CFG.""" shift, scale, gate = mod_params.chunk(3, dim=-1) if index is not None: - # Assuming mod_params batch dim is 2*actual_batch (chunked into 2 parts) - # So shift, scale, gate have shape [2*actual_batch, d] actual_batch = shift.size(0) // 2 - shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:] # each: [actual_batch, d] + shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:] scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:] gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:] - # index: [b, l] where b is actual batch size - # Expand to [b, l, 1] to match feature dimension - index_expanded = index.unsqueeze(-1) # [b, l, 1] + index_expanded = index.unsqueeze(-1) - # Expand chunks to [b, 1, d] then broadcast to [b, l, d] - shift_0_exp = shift_0.unsqueeze(1) # [b, 1, d] - shift_1_exp = shift_1.unsqueeze(1) # [b, 1, d] + shift_0_exp = shift_0.unsqueeze(1) + shift_1_exp = shift_1.unsqueeze(1) scale_0_exp = scale_0.unsqueeze(1) scale_1_exp = scale_1.unsqueeze(1) gate_0_exp = gate_0.unsqueeze(1) gate_1_exp = gate_1.unsqueeze(1) - # Use torch.where to select based on index shift_result = torch.where(index_expanded == 0, shift_0_exp, shift_1_exp) scale_result = torch.where(index_expanded == 0, scale_0_exp, scale_1_exp) gate_result = torch.where(index_expanded == 0, gate_0_exp, gate_1_exp) @@ -618,7 +588,19 @@ def _modulate(self, x, mod_params, index=None): scale_result = scale.unsqueeze(1) gate_result = gate.unsqueeze(1) - return x * (1 + scale_result) + shift_result, gate_result + return shift_result, scale_result, gate_result + + def _modulate(self, x, mod_params, index=None): + """Apply modulation to input tensor""" + shift_result, scale_result, gate_result = self._split_mod_params(mod_params, index) + # x*(1+scale)+shift == x + x*scale + shift — prefer addcmul over Adds+Mul+Add + return torch.addcmul(x, x, scale_result) + shift_result, gate_result + + def _norm_modulate(self, norm: nn.LayerNorm, x: torch.Tensor, mod_params, index=None): + """LayerNorm + Ada modulate. Fuses via platform ops when enabled.""" + shift_result, scale_result, gate_result = self._split_mod_params(mod_params, index) + out = fused_layernorm_scale_shift(norm, x, scale_result, shift_result) + return out, gate_result def forward( self, @@ -641,13 +623,8 @@ def forward( img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] - # Process image stream - norm1 + modulation - img_normed = self.img_norm1(hidden_states) - img_modulated, img_gate1 = self._modulate(img_normed, img_mod1, modulate_index) - - # Process text stream - norm1 + modulation - txt_normed = self.txt_norm1(encoder_hidden_states) - txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1) + img_modulated, img_gate1 = self._norm_modulate(self.img_norm1, hidden_states, img_mod1, modulate_index) + txt_modulated, txt_gate1 = self._norm_modulate(self.txt_norm1, encoder_hidden_states, txt_mod1) # Use QwenDoubleStreamAttention for joint attention computation # This directly implements the DoubleStreamLayerMegatron logic: @@ -665,21 +642,17 @@ def forward( image_rotary_emb=image_rotary_emb, ) - # Apply attention gates and add residual (like in Megatron) - hidden_states = hidden_states + img_gate1 * img_attn_output - encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output + # addcmul: residual + gate * out — prefer single op over Mul+Add + hidden_states = torch.addcmul(hidden_states, img_gate1, img_attn_output) + encoder_hidden_states = torch.addcmul(encoder_hidden_states, txt_gate1, txt_attn_output) - # Process image stream - norm2 + MLP - img_normed2 = self.img_norm2(hidden_states) - img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2, modulate_index) + img_modulated2, img_gate2 = self._norm_modulate(self.img_norm2, hidden_states, img_mod2, modulate_index) img_mlp_output = self.img_mlp(img_modulated2) - hidden_states = hidden_states + img_gate2 * img_mlp_output + hidden_states = torch.addcmul(hidden_states, img_gate2, img_mlp_output) - # Process text stream - norm2 + MLP - txt_normed2 = self.txt_norm2(encoder_hidden_states) - txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2) + txt_modulated2, txt_gate2 = self._norm_modulate(self.txt_norm2, encoder_hidden_states, txt_mod2) txt_mlp_output = self.txt_mlp(txt_modulated2) - encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output + encoder_hidden_states = torch.addcmul(encoder_hidden_states, txt_gate2, txt_mlp_output) # Clip to prevent overflow for fp16 if encoder_hidden_states.dtype == torch.float16: diff --git a/diffsynth_engine/pipelines/base.py b/diffsynth_engine/pipelines/base.py index 73509cb..cb149d0 100644 --- a/diffsynth_engine/pipelines/base.py +++ b/diffsynth_engine/pipelines/base.py @@ -19,6 +19,7 @@ from diffsynth_engine.forward_context import set_forward_context from diffsynth_engine.utils import logging from diffsynth_engine.utils.load_utils import fix_state_dict_key, load_model_weights, prepare_model_weights +from diffsynth_engine.utils.platform import get_compile_kwargs logger = logging.get_logger(__name__) @@ -41,10 +42,11 @@ def compile_transformer_blocks(model: nn.Module) -> nn.Module: if not repeated_blocks: raise ValueError(f"`_repeated_blocks` is not defined for {type(model).__name__}") + compile_kwargs = get_compile_kwargs() has_compiled_region = False for submodule in model.modules(): if submodule.__class__.__name__ in repeated_blocks: - submodule.compile() + submodule.compile(**compile_kwargs) has_compiled_region = True if not has_compiled_region: @@ -53,6 +55,42 @@ def compile_transformer_blocks(model: nn.Module) -> nn.Module: ) return model + @staticmethod + def compile_ffn_blocks(model: nn.Module) -> nn.Module: + """Compile only FFN (MLP) submodules within transformer blocks. + + This is a finer-grained alternative to compile_transformer_blocks that + targets only the feed-forward networks (img_mlp / txt_mlp) while leaving + attention and modulation untouched. + """ + import torch + + compile_kwargs = get_compile_kwargs() + compiled_count = 0 + for name, submodule in model.named_modules(): + if name.endswith((".img_mlp", ".txt_mlp")): + compiled_module = torch.compile(submodule, **compile_kwargs) + # Replace the submodule in parent + parts = name.rsplit(".", 1) + if len(parts) == 2: + parent_name, attr_name = parts + parent = dict(model.named_modules())[parent_name] + else: + parent = model + attr_name = parts[0] + setattr(parent, attr_name, compiled_module) + compiled_count += 1 + logger.info(f"Compiled FFN block: {name}") + + if compiled_count == 0: + logger.warning( + f"No FFN blocks (img_mlp/txt_mlp) found in {type(model).__name__}; " + "compile_ffn had no effect." + ) + else: + logger.info(f"Compiled {compiled_count} FFN blocks in {type(model).__name__}") + return model + @classmethod def init_transformer( cls, @@ -135,6 +173,8 @@ def init_transformer( del state_dict if pipeline_config.use_torch_compile: model = cls.compile_transformer_blocks(model) + elif pipeline_config.compile_ffn: + model = cls.compile_ffn_blocks(model) return model @staticmethod diff --git a/diffsynth_engine/platforms/__init__.py b/diffsynth_engine/platforms/__init__.py new file mode 100644 index 0000000..affd1d1 --- /dev/null +++ b/diffsynth_engine/platforms/__init__.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import Type + +import torch + +from .ascend import ( + AscendPlatform, + probe_ascend_capabilities, + probe_ascend_feature, + reset_ascend_capability_cache, +) +from .base import PlatformBackend, PlatformCapabilities + + +class CPUPlatform(PlatformBackend): + name = "cpu" + device_type = "cpu" + + @classmethod + def distributed_backend(cls) -> str: + raise NotImplementedError("Unsupported device type") + + @classmethod + def pin_memory(cls, tensor: torch.Tensor) -> torch.Tensor: + return tensor + + +class CUDAPlatform(PlatformBackend): + name = "cuda" + device_type = "cuda" + + @classmethod + def is_available(cls) -> bool: + return torch.cuda.is_available() + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + torch.cuda.set_device(index) + + @classmethod + def device_count(cls) -> int: + return torch.cuda.device_count() + + @classmethod + def get_device(cls, local_rank: int) -> torch.device: + return torch.device(cls.device_type, local_rank) + + @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" + + @classmethod + def fp8_dtype(cls) -> torch.dtype: + if not cls.is_available(): + return torch.float8_e4m3fn + properties = torch.cuda.get_device_properties(0) + if "gfx94" in properties.gcnArchName: + return torch.float8_e4m3fnuz + return torch.float8_e4m3fn + + +class MPSPlatform(PlatformBackend): + name = "mps" + 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() + + @classmethod + def pin_memory(cls, tensor: torch.Tensor) -> torch.Tensor: + return tensor + + +_PLATFORM_REGISTRY: dict[str, Type[PlatformBackend]] = { + "cpu": CPUPlatform, + "cuda": ROCmPlatform if torch.version.hip else CUDAPlatform, + "mps": MPSPlatform, + "npu": AscendPlatform, +} + + +def register_platform(device_type: str, platform_cls: Type[PlatformBackend], *, overwrite: bool = False) -> None: + if device_type in _PLATFORM_REGISTRY and not overwrite: + raise ValueError(f"Platform for device type {device_type!r} is already registered") + _PLATFORM_REGISTRY[device_type] = platform_cls + + +@lru_cache(maxsize=None) +def auto_detect_device() -> str: + """auto detect device type in order of cuda(gpu/rocm), npu, mps, cpu""" + for device_type in ("cuda", "npu", "mps", "cpu"): + try: + if resolve_platform(device_type).is_available(): + return device_type + except Exception: + continue + return "cpu" + + +def get_device_type(device: str | torch.device | None = None) -> str: + if device is None or (isinstance(device, str) and device.lower() in ("auto", "")): + return auto_detect_device() + if isinstance(device, torch.device): + return device.type + return str(device).split(":", 1)[0].lower() + + +def resolve_platform(device: str | torch.device) -> Type[PlatformBackend]: + device_type = get_device_type(device) + try: + return _PLATFORM_REGISTRY[device_type] + except KeyError as exc: + available = ", ".join(sorted(_PLATFORM_REGISTRY)) + raise ValueError(f"Unsupported device type {device_type!r}. Registered device types: {available}") from exc + + +# The platform backend for the auto-detected process accelerator, resolved once +current_platform: Type[PlatformBackend] = resolve_platform(get_device_type()) + + +__all__ = [ + "AscendPlatform", + "CPUPlatform", + "CUDAPlatform", + "MPSPlatform", + "PlatformBackend", + "PlatformCapabilities", + "ROCmPlatform", + "auto_detect_device", + "current_platform", + "get_device_type", + "probe_ascend_capabilities", + "probe_ascend_feature", + "register_platform", + "reset_ascend_capability_cache", + "resolve_platform", +] diff --git a/diffsynth_engine/platforms/ascend.py b/diffsynth_engine/platforms/ascend.py new file mode 100644 index 0000000..00fccc5 --- /dev/null +++ b/diffsynth_engine/platforms/ascend.py @@ -0,0 +1,246 @@ +from __future__ import annotations + +import importlib +import os +from functools import lru_cache +from typing import Any + +import torch + +from .base import PlatformBackend, PlatformCapabilities +from diffsynth_engine.utils import logging + +logger = logging.get_logger(__name__) + + +def _import_torch_npu(): + try: + return importlib.import_module("torch_npu") + except (ImportError, OSError) as exc: + raise RuntimeError( + "Ascend device requested, but torch_npu is not installed. " + "Install the torch_npu wheel matching the PyTorch and CANN versions." + ) from exc + + +def _import_mindie_sd(): + try: + return importlib.import_module("mindiesd") + except (ImportError, OSError) as exc: + raise RuntimeError( + "This Ascend feature requires MindIE-SD. Install a MindIE-SD 3.x wheel " + "matching the current torch_npu and CANN versions." + ) from exc + + +def _has_callable(obj: Any, name: str) -> bool: + return callable(getattr(obj, name, None)) + + +def _probe_npu_runtime(npu: Any) -> bool: + try: + probe = torch.zeros(1, device="npu:0") + probe.add_(1) + npu.synchronize() + return True + except Exception: + return False + + +@lru_cache(maxsize=1) +def _probe_ascend_device() -> bool: + try: + torch_npu = _import_torch_npu() + npu = getattr(torch_npu, "npu", getattr(torch, "npu", None)) + device_available = bool(npu is not None and _has_callable(npu, "is_available") and npu.is_available()) + except Exception: + return False + + if device_available: + # is_available() can stay true while ACL initialization is failing. + device_available = _probe_npu_runtime(npu) + return device_available + + +@lru_cache(maxsize=1) +def _probe_mindie_installation() -> bool: + if not _probe_ascend_device(): + return False + + try: + _import_mindie_sd() + except Exception: + return False + return True + + +def _feature_api_available(feature: str) -> bool: + if feature == "mindie_attention": + module = importlib.import_module("mindiesd.layers.flash_attn.attention_forward") + return _has_callable(module, "attention_forward") + if feature == "mindie_compile": + module = importlib.import_module("mindiesd.compilation") + return callable(getattr(module, "MindieSDBackend", None)) + + raise ValueError(f"Unknown Ascend capability: {feature}") + + +def _tensor_probe_succeeded(output: torch.Tensor) -> bool: + torch_npu = _import_torch_npu() + torch_npu.npu.synchronize() + return bool(torch.isfinite(output).all().cpu().item()) + + +def _probe_mindie_attention_operation() -> bool: + module = importlib.import_module("mindiesd.layers.flash_attn.attention_forward") + query = torch.randn(1, 128, 8, 128, device="npu:0", dtype=torch.bfloat16) + with torch.no_grad(): + output = module.attention_forward( + query=query, + key=query, + value=query, + attn_mask=None, + scale=None, + fused=True, + head_first=False, + ) + return output.shape == query.shape and _tensor_probe_succeeded(output) + + +def _probe_mindie_compile_operation() -> bool: + compilation_module = importlib.import_module("mindiesd.compilation") + + def probe_fn(value): + return torch.nn.functional.gelu(value + 1) + + compiled_fn = torch.compile(probe_fn, backend=compilation_module.MindieSDBackend(), fullgraph=False) + value = torch.randn(8, 32, device="npu:0", dtype=torch.bfloat16) + with torch.no_grad(): + output = compiled_fn(value) + return output.shape == value.shape and _tensor_probe_succeeded(output) + + + + +_OPERATION_PROBES = { + "mindie_attention": _probe_mindie_attention_operation, + "mindie_compile": _probe_mindie_compile_operation, +} + + +@lru_cache(maxsize=None) +def probe_ascend_feature(feature: str) -> bool: + if feature == "device": + return _probe_ascend_device() + if feature == "mindie": + return _probe_mindie_installation() + if feature not in _OPERATION_PROBES: + raise ValueError(f"Unknown Ascend capability: {feature}") + if not _probe_mindie_installation(): + return False + + try: + if not _feature_api_available(feature): + return False + return bool(_OPERATION_PROBES[feature]()) + except Exception: + return False + + +def probe_ascend_capabilities() -> PlatformCapabilities: + device = probe_ascend_feature("device") + if not device: + return PlatformCapabilities() + mindie = probe_ascend_feature("mindie") + if not mindie: + return PlatformCapabilities(device=True) + + return PlatformCapabilities( + device=True, + mindie=True, + mindie_attention=probe_ascend_feature("mindie_attention"), + mindie_compile=probe_ascend_feature("mindie_compile"), + ) + + +def reset_ascend_capability_cache() -> None: + _probe_ascend_device.cache_clear() + _probe_mindie_installation.cache_clear() + probe_ascend_feature.cache_clear() + + +class AscendPlatform(PlatformBackend): + name = "ascend" + device_type = "npu" + + # Ascend-specific tuning knobs (seeded from environment for now). + op_fusion = os.environ.get("USE_MINDIESD_FUSE", "False").lower() == "true" + fa_alltoall_overlap = int(os.environ.get("FA_ALLTOALL_OVERLAP", 1)) + fa_alltoall_cut = int(os.environ.get("FA_ALLTOALL_CUT", 1)) + + @classmethod + def is_available(cls) -> bool: + return probe_ascend_feature("device") + + @classmethod + def normalize_device(cls, device: str | torch.device) -> torch.device: + _import_torch_npu() + return torch.device(device) + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + torch_npu = _import_torch_npu() + torch_npu.npu.set_device(index) + + @classmethod + def device_count(cls) -> int: + _import_torch_npu() + return torch.npu.device_count() + + @classmethod + def get_device(cls, local_rank: int) -> torch.device: + return torch.device(cls.device_type, local_rank) + + @classmethod + def synchronize(cls) -> None: + torch_npu = _import_torch_npu() + 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) -> Any | None: + if not cls.supports("mindie_compile"): + logger.warning( + "MindIE-SD compile backend is unavailable; falling back to default torch.compile backend" + ) + return None + from mindiesd.compilation import CompilationConfig, MindieSDBackend + + CompilationConfig.fusion_patterns.enable_fast_gelu = False + return MindieSDBackend() + + @classmethod + def capabilities(cls) -> PlatformCapabilities: + return probe_ascend_capabilities() + + @classmethod + def supports(cls, capability: str) -> bool: + return probe_ascend_feature(capability) diff --git a/diffsynth_engine/platforms/base.py b/diffsynth_engine/platforms/base.py new file mode 100644 index 0000000..a72f559 --- /dev/null +++ b/diffsynth_engine/platforms/base.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from abc import ABC +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass(frozen=True) +class PlatformCapabilities: + device: bool = False + mindie: bool = False + mindie_attention: bool = False + mindie_compile: bool = False + + +class PlatformBackend(ABC): + name = "unknown" + device_type = "cpu" + + op_fusion = False + + @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 get_device(cls, local_rank: int) -> torch.device: + return torch.device(cls.device_type) + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + return None + + @classmethod + def device_count(cls) -> int: + return 1 + + @classmethod + def fp8_dtype(cls) -> torch.dtype: + return torch.float8_e4m3fn + + @classmethod + def synchronize(cls) -> None: + return None + + @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/ops.py b/diffsynth_engine/platforms/ops.py new file mode 100644 index 0000000..6c6dada --- /dev/null +++ b/diffsynth_engine/platforms/ops.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +"""统一平台融合算子接口。 + +GPU 路径直通原逻辑，NPU 路径调用 mindiesd / torch_npu 融合算子。 +调用方无需判断平台类型，直接调用本模块函数即可自动分发。 +""" + +from typing import Tuple, Union + +import torch +import torch.nn as nn + +from diffsynth_engine.platforms import current_platform +from diffsynth_engine.utils.platform import is_mindie_sd_available + + +def _is_npu_fused(x: torch.Tensor) -> bool: + """判断是否走 NPU 融合路径。""" + return current_platform.op_fusion and is_mindie_sd_available() and x.device.type == "npu" + + +# --------------------------------------------------------------------------- +# fused_rotary_embedding +# --------------------------------------------------------------------------- + + +def fused_rotary_embedding( + x: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + use_real: bool = True, + use_real_unbind_dim: int = -1, +) -> torch.Tensor: + """平台无关 RoPE 旋转位置编码。 + + - NPU (op_fusion + mindiesd): mindiesd.rotary_position_embedding 融合算子 + - GPU / 其他: 原始数学实现 (view_as_complex/polar 或 real-valued cos/sin) + + Args: + x: 输入张量，形状 [B, S, H, D]。 + freqs_cis: 预计算频率张量。use_real=True 时为 (cos, sin) 元组； + use_real=False 时为复数张量 [S, D/2]。 + use_real: 是否使用实数形式的 cos/sin（适用于 flux/cogvideox 等）。 + use_real_unbind_dim: use_real=True 时，拆分维度 (-1 或 -2)。 + + Returns: + 应用 RoPE 后的张量，形状与输入一致。 + """ + if use_real: + # Real-valued cos/sin 模式 (flux, cogvideox, hunyuan-dit, stable audio, etc.) + cos, sin = freqs_cis # [S, D] + cos = cos[None, None] + sin = sin[None, None] + cos, sin = cos.to(x.device), sin.to(x.device) + + if use_real_unbind_dim == -1: + # Used for flux, cogvideox, hunyuan-dit + x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) + elif use_real_unbind_dim == -2: + # Used for Stable Audio, OmniGen, CogView4 and Cosmos + x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, S, H, D//2] + x_rotated = torch.cat([-x_imag, x_real], dim=-1) + else: + raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.") + + out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + return out + + else: + # Complex-valued 模式 (qwen_image 等) + if _is_npu_fused(x): + from mindiesd import rotary_position_embedding + + # Cache expanded cos/sin on the tensor object itself. + cached = getattr(freqs_cis, "_rope_expanded", None) + if cached is None: + cos = freqs_cis.real # (s, d/2) + sin = freqs_cis.imag + cos = cos.reshape(1, -1, 1, cos.shape[-1]) # (1, S, 1, D/2) + sin = sin.reshape(1, -1, 1, sin.shape[-1]) + cos = cos.unsqueeze(-1).expand(-1, -1, -1, -1, 2).flatten(start_dim=-2) # (1, S, 1, D) + sin = sin.unsqueeze(-1).expand(-1, -1, -1, -1, 2).flatten(start_dim=-2) + cos, sin = cos.to(x.device), sin.to(x.device) + cached = (cos, sin) + freqs_cis._rope_expanded = cached + cos, sin = cached + return rotary_position_embedding( + x, + cos, + sin, + rotated_mode="rotated_interleaved", + head_first=False, + fused=True, + ) + + # GPU fallback: complex 乘法 + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(1) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + return x_out.type_as(x) + + +# --------------------------------------------------------------------------- +# fused_layernorm_scale_shift +# --------------------------------------------------------------------------- + + +def fused_layernorm_scale_shift( + norm_layer: nn.LayerNorm, + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + """融合 LayerNorm + Scale + Shift。 + + - NPU (op_fusion + mindiesd): mindiesd.layernorm_scale_shift 融合算子 + - GPU / 其他: norm_layer(x) * (1 + scale) + shift + + Args: + norm_layer: nn.LayerNorm 层实例。 + x: 输入张量。 + scale: 缩放因子 (对应 Ada modulate 中的 scale)。 + shift: 偏移量 (对应 Ada modulate 中的 shift)。 + + Returns: + 融合归一化+调制后的张量。 + """ + if _is_npu_fused(x): + from mindiesd import layernorm_scale_shift + + return layernorm_scale_shift(norm_layer, x, scale, shift, fused=True) + + # GPU fallback: 手动计算 + return norm_layer(x) * (1 + scale) + shift + + +# --------------------------------------------------------------------------- +# fused_rms_norm +# --------------------------------------------------------------------------- + + +def fused_rms_norm( + x: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, +) -> torch.Tensor: + """融合 RMSNorm。 + + - NPU (op_fusion + mindiesd + elementwise_affine): torch_npu.npu_rms_norm + - GPU / 其他: fp32 手动计算 + + Args: + x: 输入张量。 + weight: RMSNorm 权重参数。 + eps: 数值稳定性 epsilon。 + + Returns: + 归一化后的张量。 + """ + if _is_npu_fused(x): + import torch_npu + + return torch_npu.npu_rms_norm(x, weight, epsilon=eps)[0] + + # GPU fallback: fp32 精度手动计算 + output = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps) + return (output * weight).type_as(x) diff --git a/diffsynth_engine/registry.py b/diffsynth_engine/registry.py index b6b9f4e..84f7fa0 100644 --- a/diffsynth_engine/registry.py +++ b/diffsynth_engine/registry.py @@ -33,6 +33,7 @@ "fa3": "diffsynth_engine.layers.attention.backends.flash_attn_3:FlashAttention3Backend", "fa3_fp8": "diffsynth_engine.layers.attention.backends.flash_attn_3:FlashAttention3FP8Backend", "fa4": "diffsynth_engine.layers.attention.backends.flash_attn_4:FlashAttention4Backend", + "mindie": "diffsynth_engine.layers.attention.backends.mindie_attn:MindieAttentionBackend", "sage2": "diffsynth_engine.layers.attention.backends.sage_attn_2:SageAttention2Backend", "sage3": "diffsynth_engine.layers.attention.backends.sage_attn_3:SageAttention3Backend", "sdpa": "diffsynth_engine.layers.attention.backends.sdpa:SDPABackend", diff --git a/diffsynth_engine/utils/platform.py b/diffsynth_engine/utils/platform.py index 575f573..f59d30d 100644 --- a/diffsynth_engine/utils/platform.py +++ b/diffsynth_engine/utils/platform.py @@ -1,46 +1,49 @@ import torch -def _is_cuda() -> bool: - return torch.version.cuda is not None +from diffsynth_engine.platforms import ( + AscendPlatform, + current_platform, +) -def _is_rocm() -> bool: - return torch.version.hip is not None +def is_npu_available() -> bool: + return AscendPlatform.is_available() -def _is_mps() -> bool: - return torch.backends.mps.is_available() +def is_mindie_sd_available() -> bool: + return AscendPlatform.supports("mindie") -def get_device(local_rank: int) -> torch.device: - if _is_cuda() or _is_rocm(): - return torch.device("cuda", local_rank) - if _is_mps(): - return torch.device("mps") - else: - return torch.device("cpu") +def get_device_type() -> str: + return current_platform.device_type -def get_device_type() -> str: - if _is_cuda() or _is_rocm(): - return "cuda" - if _is_mps(): - return "mps" - else: - return "cpu" +def get_device(local_rank: int) -> torch.device: + return current_platform.get_device(local_rank) def get_torch_distributed_backend() -> str: - if _is_cuda() or _is_rocm(): - return "nccl" - if _is_mps(): - return "gloo" - else: - raise NotImplementedError("Unsupported device type") + return current_platform.distributed_backend() + + +def device_count() -> int: + return current_platform.device_count() + + +def set_device(index: int | str | torch.device) -> None: + current_platform.set_device(index) + + +def pin_memory(tensor: torch.Tensor) -> torch.Tensor: + return current_platform.pin_memory(tensor) + + +def get_compile_kwargs() -> dict: + return current_platform.compile_kwargs() -DTYPE_FP8 = torch.float8_e4m3fnuz if _is_rocm() else torch.float8_e4m3fn +DTYPE_FP8 = current_platform.fp8_dtype() DTYPE_MAP: dict[str, torch.dtype] = { # Integer dtypes diff --git a/diffsynth_engine/worker.py b/diffsynth_engine/worker.py index ac3e00f..4dab635 100644 --- a/diffsynth_engine/worker.py +++ b/diffsynth_engine/worker.py @@ -38,6 +38,7 @@ def __init__( os.environ["LOCAL_RANK"] = str(local_rank) os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) + init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank) cfg_degree = 2 if pipeline_config.use_cfg_parallel else 1 diff --git a/docs/npu_adaptation.md b/docs/npu_adaptation.md new file mode 100644 index 0000000..bffe158 --- /dev/null +++ b/docs/npu_adaptation.md @@ -0,0 +1,179 @@ +# Qwen Image NPU 适配技术文档 + +## 1. 概述 + +本文档描述 Qwen Image 模型在华为 Atlas 950 ProR (Ascend 910B) NPU 上的完整适配方案。 + +### 支持场景 + +- **text-to-image**: 文本生成图像 +- **image-edit**: 图像编辑 +- **image-edit-plus**: 增强图像编辑 +- **layered-generation**: 分层生成 + +## 2. 环境要求 + +### 硬件 + +| 项目 | 规格 | +|------|------| +| 加速卡 | Atlas 950 ProR (Ascend 910B) | +| 卡数 | 8 卡 | + +### 软件栈 + +| 组件 | 版本 | +|------|------| +| CANN | 9.1.0 | +| PyTorch | 2.10.0 (with torch_npu) | +| MindIE SDK | mindiesd | +| Python | 3.11+ | + +### 环境变量 + +```bash +# 启用 MindIE 融合算子 +export USE_MINDIESD_FUSE=true + +# 指向项目根目录 +export PYTHONPATH=/path/to/DiffSynth-Engine:$PYTHONPATH +``` + +## 3. 快速启动 + +### 单卡推理 + +```python +from diffsynth_engine import DiffSynthEngine, QwenImagePipelineConfig + +config = QwenImagePipelineConfig( + model_path="Qwen/Qwen-Image", + device="npu", + attn_type="mindie", +) +engine = DiffSynthEngine(config) +image = engine("A cat sitting on a windowsill", num_inference_steps=28) +image.save("output.png") +``` + +### 多卡并行 (Ulysses SP) + +```bash +torchrun --nproc_per_node=4 examples/qwen_image/run_text_to_image.py \ + --model-path Qwen/Qwen-Image \ + --device npu \ + --attn-type mindie \ + --parallelism 4 \ + --sp-ulysses-degree 4 +``` + +## 4. 架构设计 + +### 4.1 统一平台 ops 接口 + +**文件**: `diffsynth_engine/platforms/ops.py` + +提供统一的算子接口，根据运行平台自动选择最优实现: + +| 接口 | NPU 实现 | GPU 路径 | +|------|----------|----------| +| `fused_rotary_embedding()` | `mindiesd.rotary_position_embedding` | 零开销直通原始实现 | +| `fused_layernorm_scale_shift()` | `mindiesd.layernorm_scale_shift` | 零开销直通原始实现 | +| `fused_rms_norm()` | `torch_npu.npu_rms_norm` | 零开销直通原始实现 | + +设计原则: +- NPU 路径利用 MindIE SDK 融合算子获取加速 +- GPU 路径保持零开销直通，不引入额外调度延迟 +- 通过环境变量 `USE_MINDIESD_FUSE` 控制是否启用融合 + +### 4.2 Attention 工厂 + +**文件**: `diffsynth_engine/layers/attention/factory.py` + +`create_parallel_attention()` 根据平台和并行配置自动选择 Attention 实现: + +``` +┌─────────────────────────────────────────────┐ +│ create_parallel_attention() │ +├─────────────────────────────────────────────┤ +│ NPU + SP → AscendLongContextAttention │ +│ (Ulysses SP + AllToAll overlap) │ +│ GPU/其他 → USPAttention │ +└─────────────────────────────────────────────┘ +``` + +### 4.3 AscendLongContextAttention + +**文件**: `diffsynth_engine/layers/attention/ascend_long_context.py` + +从原 `layer.py` 提取为独立模块，专门针对昇腾 NPU 优化: + +- **Ulysses Sequence Parallelism**: 将长序列切分到多卡并行处理 +- **AllToAll 通信计算重叠** (overlap mode): 隐藏通信延迟 +- **切分优化** (cut mode): 减少显存占用 + +## 5. 配置参数 + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `device` | `str` | `"auto"` | 设备类型: `"npu"`, `"cuda"`, `"auto"` | +| `attn_type` | `str` | `"sdpa"` | Attention 后端: `"sdpa"`, `"flash"`, `"mindie"` | +| `op_fusion` | `bool` | `True` | 是否启用算子融合 | +| `compile_ffn` | `bool` | `False` | 是否编译 FFN (实验性，当前 NPU 无收益) | +| `parallelism` | `int` | `1` | 并行卡数 | +| `sp_ulysses_degree` | `int` | `1` | Ulysses SP 并行度 | + +## 6. 性能数据 (单卡, Atlas 950 ProR) + +| 场景 | Steps | 耗时 (s) | 每步 (ms) | +|------|-------|----------|-----------| +| text-to-image 1024×1024 | 28 | 15.85 | 565.9 | +| image-edit | 50 | 78.29 | 1565.7 | +| image-edit-plus | 50 | 70.38 | 1407.5 | +| layered-generation | 50×3 | 33.56 | 671.2 | + +## 7. 已知限制 + +1. **torch.compile 对 FFN 无优化效果** + - 在当前 CANN 9.1.0 上，torch.compile 对 FFN 模块无加速效果且存在精度退化 + - 建议保持 `compile_ffn=False` + +2. **AscendLongContextAttention 与 dynamo 不兼容** + - 内部使用 stream/event 进行通信计算重叠 + - 已通过 `@torch.compiler.disable` 装饰器规避 + +3. **多卡模式需要 HCCL 初始化** + - 框架内部自动处理，无需手动配置 + - 需确保所有 NPU 设备可见 + +4. **多卡性能数据待补充** + - 当前仅验证单卡场景 + - 多卡 Ulysses SP 性能数据待后续补充 + +## 8. 测试 + +### NPU 单卡测试 + +```bash +python -m pytest tests/test_pipelines/test_qwen_image_npu.py -v +``` + +### NPU 多卡测试 (4 卡) + +```bash +torchrun --nproc_per_node=4 -m pytest tests/test_pipelines/test_qwen_image_npu_parallel.py -v +``` + +### GPU 回归测试 + +```bash +python -m pytest tests/test_pipelines/test_qwen_image.py -v +``` + +## 9. 故障排查 + +| 错误信息 | 原因 | 解决方案 | +|----------|------|----------| +| `MindIE SDK not found` | mindiesd 未安装或不在搜索路径 | 确认 mindiesd 已安装且 `PYTHONPATH` 正确 | +| `HCCL init failed` | HCCL 通信环境异常 | 检查 NCCL/HCCL 环境，确认所有 NPU 可见 | +| 精度异常 (SSIM < 0.95) | 融合算子未启用或版本不匹配 | 检查 `USE_MINDIESD_FUSE` 环境变量是否为 `true` | diff --git a/tests/test_npu/__init__.py b/tests/test_npu/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_npu/test_attention_factory.py b/tests/test_npu/test_attention_factory.py new file mode 100644 index 0000000..0d1ccbf --- /dev/null +++ b/tests/test_npu/test_attention_factory.py @@ -0,0 +1,90 @@ +"""Unit tests for the attention factory function.""" + +from unittest.mock import patch + +import pytest + +from diffsynth_engine.layers.attention.factory import create_parallel_attention +from diffsynth_engine.layers.attention.layer import USPAttention + + +class TestCreateParallelAttention: + """Tests for create_parallel_attention factory function.""" + + def test_import_from_factory_module(self): + """Verify direct import from factory module works.""" + from diffsynth_engine.layers.attention.factory import create_parallel_attention as fn + + assert callable(fn) + + def test_import_from_package(self): + """Verify import from package __init__ works.""" + from diffsynth_engine.layers.attention import create_parallel_attention as fn + + assert callable(fn) + + @patch("diffsynth_engine.utils.platform.is_mindie_sd_available", return_value=False) + @patch("diffsynth_engine.distributed.parallel_state.is_sp_group_initialized", return_value=False) + def test_returns_usp_attention_when_no_mindie(self, mock_sp, mock_mindie): + """On GPU/CPU (no MindIE), factory should return USPAttention.""" + attn = create_parallel_attention( + num_heads=24, + head_size=128, + attn_type=None, + ) + assert isinstance(attn, USPAttention) + + @patch("diffsynth_engine.utils.platform.is_mindie_sd_available", return_value=True) + @patch("diffsynth_engine.distributed.parallel_state.is_sp_group_initialized", return_value=False) + def test_returns_usp_attention_when_sp_not_initialized(self, mock_sp, mock_mindie): + """Even if MindIE is available, without SP group we fall back to USPAttention.""" + attn = create_parallel_attention( + num_heads=24, + head_size=128, + attn_type=None, + ) + assert isinstance(attn, USPAttention) + + @patch("diffsynth_engine.utils.platform.is_mindie_sd_available", return_value=False) + @patch("diffsynth_engine.distributed.parallel_state.is_sp_group_initialized", return_value=True) + def test_returns_usp_attention_when_mindie_unavailable(self, mock_sp, mock_mindie): + """If SP is initialized but MindIE is not available, return USPAttention.""" + attn = create_parallel_attention( + num_heads=24, + head_size=128, + attn_type=None, + ) + assert isinstance(attn, USPAttention) + + @patch("diffsynth_engine.utils.platform.is_mindie_sd_available", return_value=False) + @patch("diffsynth_engine.distributed.parallel_state.is_sp_group_initialized", return_value=False) + def test_parameters_passed_correctly(self, mock_sp, mock_mindie): + """Verify that parameters are correctly forwarded to USPAttention.""" + attn = create_parallel_attention( + num_heads=32, + head_size=64, + attn_type=None, + num_kv_heads=8, + scatter_idx=2, + gather_idx=1, + ) + assert isinstance(attn, USPAttention) + assert attn.num_heads == 32 + assert attn.head_size == 64 + assert attn.num_kv_heads == 8 + assert attn.scatter_idx == 2 + assert attn.gather_idx == 1 + + @patch("diffsynth_engine.utils.platform.is_mindie_sd_available", return_value=False) + @patch("diffsynth_engine.distributed.parallel_state.is_sp_group_initialized", return_value=False) + def test_default_parameters(self, mock_sp, mock_mindie): + """Test factory with minimal required parameters.""" + attn = create_parallel_attention( + num_heads=16, + head_size=128, + ) + assert isinstance(attn, USPAttention) + assert attn.num_heads == 16 + assert attn.head_size == 128 + # num_kv_heads defaults to num_heads when None + assert attn.num_kv_heads == 16 diff --git a/tests/test_npu/test_platform_ops.py b/tests/test_npu/test_platform_ops.py new file mode 100644 index 0000000..38dae2a --- /dev/null +++ b/tests/test_npu/test_platform_ops.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for diffsynth_engine.platforms.ops on CPU (GPU-fallback paths).""" + +import math + +import torch +import torch.nn as nn +import pytest + +from diffsynth_engine.platforms.ops import ( + fused_layernorm_scale_shift, + fused_rms_norm, + fused_rotary_embedding, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def assert_tensor_equal(actual: torch.Tensor, expected: torch.Tensor, atol=1e-6, rtol=1e-6): + """断言两个 tensor 在给定容差内相等。""" + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# fused_rms_norm tests +# --------------------------------------------------------------------------- + + +class TestFusedRmsNorm: + """验证 fused_rms_norm GPU fallback 与手动计算一致。""" + + @pytest.mark.parametrize("shape", [(2, 8, 64), (1, 128), (4, 16, 32)]) + def test_output_matches_manual(self, shape): + torch.manual_seed(42) + x = torch.randn(*shape) + weight = torch.randn(shape[-1]) + eps = 1e-6 + + # 手动计算参考结果 + x_fp32 = x.float() + rms = torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + eps) + expected = (x_fp32 * rms * weight).to(x.dtype) + + result = fused_rms_norm(x, weight, eps) + assert_tensor_equal(result, expected) + + def test_preserves_dtype_bfloat16(self): + torch.manual_seed(0) + x = torch.randn(2, 16, dtype=torch.bfloat16) + weight = torch.randn(16, dtype=torch.bfloat16) + + result = fused_rms_norm(x, weight, eps=1e-6) + assert result.dtype == torch.bfloat16 + + def test_preserves_dtype_float16(self): + torch.manual_seed(0) + x = torch.randn(2, 16, dtype=torch.float16) + weight = torch.randn(16, dtype=torch.float16) + + result = fused_rms_norm(x, weight, eps=1e-6) + assert result.dtype == torch.float16 + + +# --------------------------------------------------------------------------- +# fused_layernorm_scale_shift tests +# --------------------------------------------------------------------------- + + +class TestFusedLayernormScaleShift: + """验证 fused_layernorm_scale_shift GPU fallback 与手动计算一致。""" + + @pytest.mark.parametrize("shape", [(2, 8, 64), (1, 4, 128)]) + def test_output_matches_manual(self, shape): + torch.manual_seed(42) + dim = shape[-1] + norm = nn.LayerNorm(dim) + x = torch.randn(*shape) + scale = torch.randn(*shape) + shift = torch.randn(*shape) + + # 手动参考: norm(x) * (1 + scale) + shift + expected = norm(x) * (1 + scale) + shift + + result = fused_layernorm_scale_shift(norm, x, scale, shift) + assert_tensor_equal(result, expected) + + def test_zero_scale_shift(self): + """scale=0, shift=0 应等同于 norm(x)。""" + torch.manual_seed(7) + dim = 32 + norm = nn.LayerNorm(dim) + x = torch.randn(2, 4, dim) + scale = torch.zeros(2, 4, dim) + shift = torch.zeros(2, 4, dim) + + expected = norm(x) + result = fused_layernorm_scale_shift(norm, x, scale, shift) + assert_tensor_equal(result, expected) + + +# --------------------------------------------------------------------------- +# fused_rotary_embedding tests +# --------------------------------------------------------------------------- + + +class TestFusedRotaryEmbedding: + """验证 fused_rotary_embedding GPU fallback 与原始 apply_rotary_emb_qwen 一致。""" + + def _make_complex_freqs(self, seq_len: int, dim: int) -> torch.Tensor: + """生成 complex 格式的 freqs_cis [S, D/2]。""" + half_dim = dim // 2 + freqs = torch.randn(seq_len, half_dim) + # 转为 complex: e^(i*theta) 形式 + angles = torch.randn(seq_len, half_dim) + freqs_cis = torch.polar(torch.ones_like(angles), angles) + return freqs_cis + + def _make_real_freqs(self, num_positions: int, dim: int): + """生成 real 格式的 (cos, sin) freqs [num_positions, D]。 + + 注意: use_real 路径中 cos[None, None] 形成 [1,1,N,D]， + 与 x [B,S,H,D] 做广播时 N 须等于 H。 + """ + cos = torch.randn(num_positions, dim) + sin = torch.randn(num_positions, dim) + return (cos, sin) + + def test_complex_mode_matches_reference(self): + """use_real=False: 与 view_as_complex 参考实现一致。""" + torch.manual_seed(42) + B, S, H, D = 2, 8, 4, 64 + x = torch.randn(B, S, H, D) + freqs_cis = self._make_complex_freqs(S, D) + + # 参考实现 + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + fc = freqs_cis.unsqueeze(1) + expected = torch.view_as_real(x_rotated * fc).flatten(3).type_as(x) + + result = fused_rotary_embedding(x, freqs_cis, use_real=False) + assert_tensor_equal(result, expected) + + def test_real_mode_unbind_neg1(self): + """use_real=True, use_real_unbind_dim=-1 与参考实现一致。""" + torch.manual_seed(42) + B, S, H, D = 2, 8, 4, 64 + x = torch.randn(B, S, H, D) + # cos/sin 第一维须等于 H 以满足 cos[None,None] 与 x 的广播 + cos, sin = self._make_real_freqs(H, D) + freqs_cis = (cos, sin) + + # 参考实现 + c = cos[None, None].to(x.device) + s = sin[None, None].to(x.device) + x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) + expected = (x.float() * c + x_rotated.float() * s).to(x.dtype) + + result = fused_rotary_embedding(x, freqs_cis, use_real=True, use_real_unbind_dim=-1) + assert_tensor_equal(result, expected) + + def test_real_mode_unbind_neg2(self): + """use_real=True, use_real_unbind_dim=-2 与参考实现一致。""" + torch.manual_seed(42) + B, S, H, D = 2, 8, 4, 64 + x = torch.randn(B, S, H, D) + # cos/sin 第一维须等于 H 以满足 cos[None,None] 与 x 的广播 + cos, sin = self._make_real_freqs(H, D) + freqs_cis = (cos, sin) + + # 参考实现 + c = cos[None, None].to(x.device) + s = sin[None, None].to(x.device) + x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) + x_rotated = torch.cat([-x_imag, x_real], dim=-1) + expected = (x.float() * c + x_rotated.float() * s).to(x.dtype) + + result = fused_rotary_embedding(x, freqs_cis, use_real=True, use_real_unbind_dim=-2) + assert_tensor_equal(result, expected) + + def test_invalid_unbind_dim_raises(self): + """无效的 use_real_unbind_dim 应抛出 ValueError。""" + x = torch.randn(1, 4, 2, 8) + freqs_cis = (torch.randn(4, 8), torch.randn(4, 8)) + + with pytest.raises(ValueError, match="use_real_unbind_dim"): + fused_rotary_embedding(x, freqs_cis, use_real=True, use_real_unbind_dim=0) + + def test_output_shape_preserved(self): + """输出形状与输入一致。""" + B, S, H, D = 1, 16, 8, 128 + x = torch.randn(B, S, H, D) + freqs_cis = self._make_complex_freqs(S, D) + + result = fused_rotary_embedding(x, freqs_cis, use_real=False) + assert result.shape == x.shape diff --git a/tests/test_pipelines/test_qwen_image_npu.py b/tests/test_pipelines/test_qwen_image_npu.py new file mode 100644 index 0000000..ff18e49 --- /dev/null +++ b/tests/test_pipelines/test_qwen_image_npu.py @@ -0,0 +1,210 @@ +"""NPU 单卡 Qwen Image 全场景集成测试""" +import os +import unittest + +import torch + +from diffsynth_engine import DiffSynthEngine +from diffsynth_engine.configs import QwenImagePipelineConfig +from diffsynth_engine.utils.download import fetch_model +from tests.common.test_case import ImageTestCase + +# NPU 可用性检查 +try: + import torch_npu # noqa: F401 + + NPU_AVAILABLE = torch.npu.is_available() +except ImportError: + NPU_AVAILABLE = False + + +@unittest.skipUnless(NPU_AVAILABLE, "NPU not available") +class TestQwenImageNPU(ImageTestCase): + """NPU 单卡 text-to-image 测试""" + + @classmethod + def setUpClass(cls): + os.environ["USE_MINDIESD_FUSE"] = "true" + cls.model_path = fetch_model("Qwen/Qwen-Image") + config = QwenImagePipelineConfig( + model_path=cls.model_path, + device="npu", + attn_type="mindie", + model_dtype=torch.bfloat16, + ) + cls.engine = DiffSynthEngine.from_pretrained(config) + + @classmethod + def tearDownClass(cls): + cls.engine.shutdown() + del cls.engine + torch.npu.empty_cache() + + def test_txt2img(self): + prompt = "A painting of a cat in a zen garden" + negative_prompt = "ugly, blurry, low quality" + output = self.engine.generate( + prompt=prompt, + negative_prompt=negative_prompt, + true_cfg_scale=4.0, + width=1328, + height=1328, + num_inference_steps=28, + generator=torch.Generator(device="cpu").manual_seed(42), + ) + image = output.images[0] + self.assertImageEqualAndSaveFailed(image, "qwen_image/qwen_image.png", threshold=0.95) + + +@unittest.skipUnless(NPU_AVAILABLE, "NPU not available") +class TestQwenImageEditNPU(ImageTestCase): + """NPU 单卡 image-edit 测试""" + + @classmethod + def setUpClass(cls): + os.environ["USE_MINDIESD_FUSE"] = "true" + cls.model_path = fetch_model("Qwen/Qwen-Image-Edit") + config = QwenImagePipelineConfig( + model_path=cls.model_path, + device="npu", + attn_type="mindie", + model_dtype=torch.bfloat16, + ) + cls.engine = DiffSynthEngine.from_pretrained(config) + + @classmethod + def tearDownClass(cls): + cls.engine.shutdown() + del cls.engine + torch.npu.empty_cache() + + def test_single_image_edit(self): + """Test single image editing on NPU""" + input_image = self.get_input_image("qwen_image_edit_input.png") + prompt = "Replace '通义千问' with '呜哩AI'" + negative_prompt = " " + + output = self.engine.generate( + image=input_image, + prompt=prompt, + negative_prompt=negative_prompt, + true_cfg_scale=4.0, + num_inference_steps=50, + generator=torch.Generator(device="cpu").manual_seed(42), + ) + image = output.images[0] + self.assertImageEqualAndSaveFailed(image, "qwen_image/qwen_image_edit.png", threshold=0.95) + + +@unittest.skipUnless(NPU_AVAILABLE, "NPU not available") +class TestQwenImageEditPlusNPU(ImageTestCase): + """NPU 单卡 edit-plus 测试""" + + @classmethod + def setUpClass(cls): + os.environ["USE_MINDIESD_FUSE"] = "true" + cls.model_path = fetch_model("Qwen/Qwen-Image-Edit-2511") + config = QwenImagePipelineConfig( + model_path=cls.model_path, + device="npu", + attn_type="mindie", + model_dtype=torch.bfloat16, + ) + cls.engine = DiffSynthEngine.from_pretrained(config) + + @classmethod + def tearDownClass(cls): + cls.engine.shutdown() + del cls.engine + torch.npu.empty_cache() + + def test_single_image_edit(self): + """Test single image editing with Edit Plus pipeline on NPU""" + input_image = self.get_input_image("qwen_image_edit_input.png") + prompt = "Replace '通义千问' with '呜哩AI'" + negative_prompt = " " + + output = self.engine.generate( + image=input_image, + prompt=prompt, + negative_prompt=negative_prompt, + true_cfg_scale=4.0, + num_inference_steps=50, + generator=torch.Generator(device="cpu").manual_seed(42), + ) + image = output.images[0] + self.assertImageEqualAndSaveFailed(image, "qwen_image/qwen_image_edit_plus_single_2511.png", threshold=0.95) + + def test_multi_image_edit(self): + """Test multiple images editing with Edit Plus pipeline on NPU""" + input_images = [ + self.get_input_image("qwen_image_edit_input_1.png").convert("RGB"), + self.get_input_image("qwen_image_edit_input_2.png").convert("RGB"), + ] + prompt = "根据这图1中女性和图2中的男性，生成一组结婚照，并遵循以下描述：新郎穿着红色的中式马褂，新娘穿着精致的秀禾服，头戴金色凤冠。他们并肩站立在古老的朱红色宫墙前，背景是雕花的木窗。光线明亮柔和，构图对称，氛围喜庆而庄重。" + negative_prompt = " " + + output = self.engine.generate( + image=input_images, + prompt=prompt, + negative_prompt=negative_prompt, + true_cfg_scale=4.0, + num_inference_steps=40, + generator=torch.Generator(device="cpu").manual_seed(42), + ) + image = output.images[0] + self.assertImageEqualAndSaveFailed(image, "qwen_image/qwen_image_edit_plus_multi_2511.png", threshold=0.95) + + +@unittest.skipUnless(NPU_AVAILABLE, "NPU not available") +class TestQwenImageLayeredNPU(ImageTestCase): + """NPU 单卡 layered 测试""" + + @classmethod + def setUpClass(cls): + os.environ["USE_MINDIESD_FUSE"] = "true" + cls.model_path = fetch_model("Qwen/Qwen-Image-Layered") + config = QwenImagePipelineConfig( + model_path=cls.model_path, + device="npu", + attn_type="mindie", + model_dtype=torch.bfloat16, + ) + cls.engine = DiffSynthEngine.from_pretrained(config) + + @classmethod + def tearDownClass(cls): + cls.engine.shutdown() + del cls.engine + torch.npu.empty_cache() + + def test_image_layered(self): + """Test layered image generation on NPU""" + input_image = self.get_input_image("qwen_image_layered_input.png").convert("RGBA") + prompt = "" + + output = self.engine.generate( + image=input_image, + prompt=prompt, + num_inference_steps=50, + true_cfg_scale=4.0, + layers=3, + resolution=640, + cfg_normalize=False, + use_en_prompt=True, + generator=torch.Generator(device="cpu").manual_seed(42), + ) + + images = output.images[0] + self.assertEqual(len(images), 3) + + for i, layer_image in enumerate(images): + self.assertImageEqualAndSaveFailed( + layer_image, + f"qwen_image/qwen_image_layered_{i}.png", + threshold=0.95, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pipelines/test_qwen_image_npu_parallel.py b/tests/test_pipelines/test_qwen_image_npu_parallel.py new file mode 100644 index 0000000..d93f6ba --- /dev/null +++ b/tests/test_pipelines/test_qwen_image_npu_parallel.py @@ -0,0 +1,105 @@ +"""NPU 多卡并行 Qwen Image 测试（Ulysses SP）""" +import os +import unittest + +import torch + +try: + import torch_npu + + NPU_AVAILABLE = torch.npu.is_available() + NPU_COUNT = torch.npu.device_count() if NPU_AVAILABLE else 0 +except ImportError: + NPU_AVAILABLE = False + NPU_COUNT = 0 + +from diffsynth_engine import DiffSynthEngine +from diffsynth_engine.configs import QwenImagePipelineConfig +from diffsynth_engine.utils.download import fetch_model +from tests.common.test_case import ImageTestCase + + +@unittest.skipUnless(NPU_AVAILABLE and NPU_COUNT >= 4, "Need at least 4 NPUs") +class TestQwenImageNPU4Card(ImageTestCase): + """4 卡 Ulysses SP 测试""" + + @classmethod + def setUpClass(cls): + os.environ["USE_MINDIESD_FUSE"] = "true" + model_path = fetch_model("Qwen/Qwen-Image") + config = QwenImagePipelineConfig( + model_path=model_path, + model_dtype=torch.bfloat16, + device="npu", + attn_type="mindie", + parallelism=4, + sp_ulysses_degree=4, + sp_ring_degree=1, + ) + cls.engine = DiffSynthEngine.from_pretrained(config) + + @classmethod + def tearDownClass(cls): + cls.engine.shutdown() + del cls.engine + torch.npu.empty_cache() + + def test_txt2img_ulysses_4card(self): + prompt = "A painting of a cat in a zen garden" + negative_prompt = "ugly, blurry, low quality" + output = self.engine.generate( + prompt=prompt, + negative_prompt=negative_prompt, + true_cfg_scale=4.0, + width=1328, + height=1328, + num_inference_steps=28, + generator=torch.Generator(device="cpu").manual_seed(42), + ) + image = output.images[0] + self.assertImageEqualAndSaveFailed(image, "qwen_image/qwen_image.png", threshold=0.97) + + +@unittest.skipUnless(NPU_AVAILABLE and NPU_COUNT >= 8, "Need at least 8 NPUs") +class TestQwenImageNPU8Card(ImageTestCase): + """8 卡 Ulysses SP 测试""" + + @classmethod + def setUpClass(cls): + os.environ["USE_MINDIESD_FUSE"] = "true" + model_path = fetch_model("Qwen/Qwen-Image") + config = QwenImagePipelineConfig( + model_path=model_path, + model_dtype=torch.bfloat16, + device="npu", + attn_type="mindie", + parallelism=8, + sp_ulysses_degree=8, + sp_ring_degree=1, + ) + cls.engine = DiffSynthEngine.from_pretrained(config) + + @classmethod + def tearDownClass(cls): + cls.engine.shutdown() + del cls.engine + torch.npu.empty_cache() + + def test_txt2img_ulysses_8card(self): + prompt = "A painting of a cat in a zen garden" + negative_prompt = "ugly, blurry, low quality" + output = self.engine.generate( + prompt=prompt, + negative_prompt=negative_prompt, + true_cfg_scale=4.0, + width=1328, + height=1328, + num_inference_steps=28, + generator=torch.Generator(device="cpu").manual_seed(42), + ) + image = output.images[0] + self.assertImageEqualAndSaveFailed(image, "qwen_image/qwen_image.png", threshold=0.97) + + +if __name__ == "__main__": + unittest.main()