[PERF] Ascend NPU: MindIE-SD attention backend and VAE-parallel enablement

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()

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

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

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)

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",
 ]

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,

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}