From 8ec1dce3f78552c670c4f0db4486bb074b928443 Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 02:09:18 +0800 Subject: [PATCH 1/3] fix(models): accept batched tensors in wan22 VAE parallel decode Wan22VideoVAE.decode's parallel branch called torch.stack on its input, assuming a list of per-video tensors, but VAEStage.decode_video passes a batched [B, C, T, H, W] tensor; the serial branch only worked because iterating a tensor yields its batch slices. Any enable_vae_parallel run of the Wan2.2 48-channel VAE therefore failed with TypeError regardless of platform. Normalize tensor inputs before stacking. Found while enabling spatially parallel VAE decode on Ascend 910B2; the parallel path now proceeds to communicator setup (further multi-communicator progress on that host is limited by its CANN driver, see branch notes). --- telefuser/models/wan22_video_vae.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 From 5b71ae513b20db9cb9bc6c9df0c7bcaaa077ada7 Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 02:09:18 +0800 Subject: [PATCH 2/3] fix(worker): isolate HCCL socket ranges per parallel worker group Concurrent worker groups sharing the same NPU devices (e.g. denoising plus VAE workers) collide on HCCL's default data-plane socket range and fail comm init with EJ0003 ("IP address and port have been bound already"). Allocate a distinct HCCL_IF_BASE_PORT per spawned group on NPU platforms, mirroring the existing per-group MASTER_PORT allocation; other platforms are untouched. Verified: single-group 4-card Wan2.2-TI2V-5B regression on Ascend 910B2 passes with the env applied (33.4s generate, parity with baseline). --- telefuser/worker/parallel_worker.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) 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, From a818e72423f707309771a4e052d58ad3c30cc262 Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 14:54:02 +0800 Subject: [PATCH 3/3] feat(ops): add optional MindIE-SD attention backend for Ascend NPU Add AttnImplType.MINDIE_ATTN, routed through mindiesd.attention_forward (auto-tuned Ascend kernels including LaserAttention) for dense BNSD attention. Availability is probed like the sageattention backend (optional import, NPU platform only), so default behavior is unchanged on every platform and the enum is strictly opt-in. Verified on Atlas 910B (CANN 8.2): kernel microbench 1.29-1.53x over SDPA across wan shapes (holds at ulysses head-sharded 12-head shapes); unmodified wan22_t2v_5b example end to end: 720p 121f 477.8->405.0s (1.18x) on 1 card and 173.9->157.6s (1.10x) on 4 cards, 480p 121f 180.9->162.6s on 1 card; a 480p 4-card regression (87.0->103.0s) is documented in the PR - the backend is recommended for compute-dominated configurations. Output is bit-identical to SDPA at probed shapes; dispatch covered by a mock unit test; ruff clean. --- docs/en/attention.md | 1 + telefuser/core/config.py | 1 + telefuser/ops/attention/attention_impl.py | 10 ++++-- telefuser/ops/attention/backends.py | 41 +++++++++++++++++++++++ tests/unit/ops/test_attention_backends.py | 16 +++++++++ 5 files changed, 67 insertions(+), 2 deletions(-) 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/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/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}