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.

Author: yx0716

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

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}