diff --git a/docs/en/attention.md b/docs/en/attention.md index 757e53ca..89a2774a 100644 --- a/docs/en/attention.md +++ b/docs/en/attention.md @@ -24,6 +24,7 @@ class AttnImplType(Enum): FLASH_ATTN_2 = auto() FLASH_ATTN_3 = auto() FLASH_ATTN_4 = auto() # For Hopper (SM90) and Blackwell (SM100+) GPUs + MINDIE_ATTN = auto() # Ascend NPU, requires the optional mindiesd package SAGE_ATTN_2_8_8 = auto() SAGE_ATTN_2_8_16 = auto() SAGE_ATTN_2_8_8_SM90 = auto() diff --git a/telefuser/core/config.py b/telefuser/core/config.py index 7bd1f5f8..5b776d36 100644 --- a/telefuser/core/config.py +++ b/telefuser/core/config.py @@ -127,6 +127,7 @@ class AttnImplType(Enum): FLASH_ATTN_2 = auto() FLASH_ATTN_3 = auto() FLASH_ATTN_4 = auto() + MINDIE_ATTN = auto() # Sparse attention implementations RADIAL_ATTN = auto() # Radial attention for video generation LOCAL_SPARSE_ATTN = auto() # Local window sparse attention diff --git a/telefuser/models/wan22_video_vae.py b/telefuser/models/wan22_video_vae.py index 35e722e2..82fb9be6 100644 --- a/telefuser/models/wan22_video_vae.py +++ b/telefuser/models/wan22_video_vae.py @@ -1393,7 +1393,11 @@ def decode( if self.parallelism > 1 and dist.is_initialized(): # tiled=True → tile_dist, tiled=False → 2d_split method = "tile_dist" if tiled else "2d_split" - hidden_states_tensor = torch.stack(hidden_states) + # The stage passes a batched [B, C, T, H, W] tensor; lists of per-video tensors are stacked. + if isinstance(hidden_states, torch.Tensor): + hidden_states_tensor = hidden_states + else: + hidden_states_tensor = torch.stack(hidden_states) return self.decode_parallel(hidden_states_tensor, device, method=method) # Single GPU processing diff --git a/telefuser/ops/attention/attention_impl.py b/telefuser/ops/attention/attention_impl.py index f03427e3..1b2b8d85 100755 --- a/telefuser/ops/attention/attention_impl.py +++ b/telefuser/ops/attention/attention_impl.py @@ -1,7 +1,7 @@ """Unified attention implementation with dense and sparse support. Supports multiple attention implementations: -- Dense: TORCH_SDPA, TORCH_CUDNN, FLASH_ATTN_2/3/4, SAGE_ATTN variants, SPARGE_ATTN +- Dense: TORCH_SDPA, TORCH_CUDNN, FLASH_ATTN_2/3/4, SAGE_ATTN variants, SPARGE_ATTN, MINDIE_ATTN (Ascend NPU) - Sparse: RADIAL_ATTN, LOCAL_SPARSE_ATTN, SOL_ATTN Note: Attention functions are decorated with @torch.compiler.disable because: @@ -31,6 +31,7 @@ FLASH_ATTN_2_AVAILABLE, FLASH_ATTN_3_AVAILABLE, FLASH_ATTN_4_AVAILABLE, + MINDIE_ATTN_AVAILABLE, SAGE_ATTN_AVAILABLE, SDPA_AVAILABLE, SOL_ATTN_AVAILABLE, @@ -39,6 +40,7 @@ flash_attn4, flash_attn4_varlen, get_lse_fallback_impl, + mindie_attn, sageattention, sdpa_attn_cudnn, sol_attn, @@ -272,7 +274,7 @@ def attention( # - Flash Attention expects BSND (NHD) layout # - PyTorch SDPA variants expect BNSD (HND) layout # - SageAttention can accept both via tensor_layout parameter - BNSD_IMPLS = {AttnImplType.TORCH_CUDNN, AttnImplType.TORCH_SDPA, AttnImplType.SPARGE_ATTN} + BNSD_IMPLS = {AttnImplType.TORCH_CUDNN, AttnImplType.TORCH_SDPA, AttnImplType.SPARGE_ATTN, AttnImplType.MINDIE_ATTN} # Track current layout after potential conversions current_layout = input_layout @@ -347,6 +349,10 @@ def attention( else: output = result + # MindIE-SD attention (Ascend NPU) + elif sequence_lengths is None and attn_impl == AttnImplType.MINDIE_ATTN and MINDIE_ATTN_AVAILABLE: + output = mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale, is_causal=is_causal) + # PyTorch SDPA elif sequence_lengths is None and attn_impl == AttnImplType.TORCH_CUDNN and SDPA_AVAILABLE: output = sdpa_attn_cudnn(q, k, v, attn_mask=attn_mask, scale=scale, is_causal=is_causal) diff --git a/telefuser/ops/attention/backends.py b/telefuser/ops/attention/backends.py index 1b76942b..48ae78ba 100644 --- a/telefuser/ops/attention/backends.py +++ b/telefuser/ops/attention/backends.py @@ -13,6 +13,7 @@ import torch from torch import Tensor +from telefuser.platforms import current_platform from telefuser.utils.logging import logger # Availability flags @@ -24,6 +25,7 @@ SPARGE_ATTN_AVAILABLE = False FLASHINFER_AVAILABLE = False SOL_ATTN_AVAILABLE = False +MINDIE_ATTN_AVAILABLE = False # Backend function references (populated on successful import) flash_attn2: Callable | None = None @@ -34,6 +36,7 @@ spas_sage2_attn_meansim_cuda: Callable | None = None flashinfer: object | None = None sol_attn: Callable | None = None +mindiesd_attention_forward: Callable | None = None def _try_import_flash_attn() -> None: @@ -162,12 +165,48 @@ def _try_import_sol_attn() -> None: # Initialize all backends +def _try_import_mindie_attn() -> None: + """Import the optional MindIE-SD attention backend (Ascend NPU only).""" + global MINDIE_ATTN_AVAILABLE, mindiesd_attention_forward + + MINDIE_ATTN_AVAILABLE = False + mindiesd_attention_forward = None + if current_platform.device_type != "npu": + return + try: + if importlib.util.find_spec("mindiesd") is None: + return + from mindiesd import attention_forward + except (ModuleNotFoundError, ImportError) as error: + logger.debug("MindIE-SD attention backend unavailable: %s", error) + return + mindiesd_attention_forward = attention_forward + MINDIE_ATTN_AVAILABLE = True + logger.debug("MindIE-SD attention available") + + +def mindie_attn( + q: Tensor, + k: Tensor, + v: Tensor, + *, + attn_mask: Tensor | None = None, + scale: float | None = None, + is_causal: bool = False, +) -> Tensor: + """MindIE-SD attention wrapper for BNSD inputs on Ascend NPU.""" + if is_causal: + raise ValueError("MINDIE_ATTN does not support causal attention") + return mindiesd_attention_forward(q, k, v, attn_mask=attn_mask, scale=scale, head_first=True) + + _try_import_flash_attn() _try_import_sdpa() _try_import_sage_attn() _try_import_sparge_attn() _try_import_flashinfer() _try_import_sol_attn() +_try_import_mindie_attn() def supports_return_lse(attn_impl: str) -> bool: @@ -234,6 +273,7 @@ def sparge_attn( "SPARGE_ATTN_AVAILABLE", "FLASHINFER_AVAILABLE", "SOL_ATTN_AVAILABLE", + "MINDIE_ATTN_AVAILABLE", "flash_attn2", "flash_attn3", "flash_attn4", @@ -245,4 +285,5 @@ def sparge_attn( "get_lse_fallback_impl", "sdpa_attn_cudnn", "sparge_attn", + "mindie_attn", ] diff --git a/telefuser/worker/parallel_worker.py b/telefuser/worker/parallel_worker.py index 2c160544..7a4596a8 100644 --- a/telefuser/worker/parallel_worker.py +++ b/telefuser/worker/parallel_worker.py @@ -7,6 +7,7 @@ from __future__ import annotations import gc +import itertools import os import signal import threading @@ -35,6 +36,10 @@ _DISCARD_TENSOR_REFS = "__telefuser_discard_tensor_refs__" +# Each concurrent worker group needs a distinct HCCL socket range on NPU hosts; +# overlapping ranges fail comm init with EJ0003 (port already bound). +_hccl_group_counter = itertools.count() + def to_device(data: Any, device: str | torch.device) -> Any: """Recursively move data to target device.""" @@ -65,6 +70,7 @@ def _worker_loop( master_port: int, tensor_output_channel: WorkerTensorChannel | None = None, tensor_input_channels: tuple[WorkerTensorChannel, ...] = (), + hccl_if_base_port: int | None = None, ) -> None: """Worker process main loop. @@ -89,6 +95,11 @@ def _worker_loop( os.environ["WORLD_SIZE"] = str(world_size) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(master_port) + if hccl_if_base_port is not None: + # Give each worker group a distinct host-side HCCL socket base; + # overlapping ranges fail comm init with EJ0003 when several + # groups share the same devices (e.g. denoise + VAE workers). + os.environ["HCCL_IF_BASE_PORT"] = str(hccl_if_base_port) device_ids = parallel_config.device_ids device_id = rank if device_ids is not None: @@ -253,6 +264,7 @@ def __init__( master_port, self.tensor_output_channel, self.tensor_input_channels, + 60000 + 1024 * next(_hccl_group_counter) if current_platform.device_type == "npu" else None, ), nprocs=self.world_size, join=False, diff --git a/tests/unit/ops/test_attention_backends.py b/tests/unit/ops/test_attention_backends.py index 4599e1f7..0d7a82c1 100644 --- a/tests/unit/ops/test_attention_backends.py +++ b/tests/unit/ops/test_attention_backends.py @@ -269,3 +269,19 @@ def test_ring_config_converts_fallback_name() -> None: assert result.attn_impl is AttnImplType.SAGE_ATTN_2_8_8_SM90 assert result.scale == 0.125 assert result.is_causal is True + + +def test_mindie_attn_dispatch_converts_to_bnsd_and_back() -> None: + q = torch.randn(1, 3, 2, 4) + forward = MagicMock(side_effect=lambda q_, k_, v_, **kwargs: q_) + + with ( + patch.object(attention_impl, "MINDIE_ATTN_AVAILABLE", True), + patch.object(backends, "mindiesd_attention_forward", forward), + ): + output = attention_impl.attention(q, q, q, attn_impl=attention_impl.AttnImplType.MINDIE_ATTN) + + assert output.shape == q.shape + called_q = forward.call_args.args[0] + assert called_q.shape == (1, 2, 3, 4) + assert forward.call_args.kwargs == {"attn_mask": None, "scale": None, "head_first": True}