feat(npu): Ascend multi-card Ulysses SP with CFG/FSDP support

diffsynth_engine/configs/pipeline.py

@@ -35,6 +35,7 @@ class AttnImpl(Enum):
     SAGE = "sage"  # Sage Attention
     SPARGE = "sparge"  # Sparge Attention
     VSA = "vsa"  # Video Sparse Attention
+    MINDIE = "mindie"  # Mindie Attention
 
 
 @dataclass

diffsynth_engine/models/basic/attention.py

@@ -15,6 +15,7 @@
     SPARGE_ATTN_AVAILABLE,
     VIDEO_SPARSE_ATTN_AVAILABLE,
     AITER_AVAILABLE,
+    MINDIE_AVAILABLE,
 )
 from diffsynth_engine.utils.platform import DTYPE_FP8
 
@@ -107,6 +108,25 @@ def sparge_attn(
         distributed_video_sparse_attn,
     )
 
+if MINDIE_AVAILABLE:
+    from mindiesd.layers.flash_attn.attention_forward import attention_forward
+
+    def mindie_attn(q, k, v, attn_mask=None, scale=None):
+        #return attention_forward(
+        #    query=q, key=k, value=v,
+        #   attn_mask=attn_mask, scale=scale,
+        #    fused=True, head_first=False,
+        #)
+
+        return attention_forward(
+            query=q, key=k, value=v,
+            attn_mask=attn_mask, scale=scale,
+            fused=True, head_first=False,
+            opt_mode="manual",
+            op_type="fused_attn_score",
+            layout="BSND",
+        )
+
 
 def eager_attn(q, k, v, attn_mask=None, scale=None):
     q = q.transpose(1, 2)
@@ -152,6 +172,7 @@ def attention(
         "sage",
         "sparge",
         "vsa",
+        "mindie",
     ]
     flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM
     if attn_impl is None or attn_impl == "auto":
@@ -192,6 +213,8 @@ def attention(
                 )
         if XFORMERS_AVAILABLE:
             return xformers_attn(q, k, v, attn_mask=attn_mask, scale=scale)
+        if MINDIE_AVAILABLE:
+            return mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale)
         if SDPA_AVAILABLE:
             return sdpa_attn(q, k, v, attn_mask=attn_mask, scale=scale)
         if FLASH_ATTN_2_AVAILABLE:
@@ -263,6 +286,8 @@ def attention(
                 cdfthreshd=kwargs.get("cdfthreshd", 0.98),
                 pvthreshd=kwargs.get("pvthreshd", 50),
             )
+        if attn_impl == "mindie":
+            return mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale)
         if attn_impl == "vsa":
             return video_sparse_attn(
                 q,
@@ -324,6 +349,44 @@ def forward(
         return self.to_out(out)
 
 
+def _npu_ulysses_mindie_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+):
+    """Ulysses SP on NPU: SeqAllToAll4D + MindIE local attn. ring_degree>1 not supported."""
+    from yunchang.comm.all_to_all import SeqAllToAll4D
+
+    from diffsynth_engine.utils.process_group import get_sp_ring_world_size, get_sp_ulysses_group
+
+    if q.device.type != "npu":
+        raise RuntimeError("mindie long-context attention is only supported on NPU")
+    if not MINDIE_AVAILABLE:
+        raise RuntimeError(
+            "NPU Ulysses sequence parallel requires MindIE attention, but MindIE-SD is not available"
+        )
+    if get_sp_ring_world_size() > 1:
+        raise RuntimeError(
+            "NPU long-context attention currently supports Ulysses only "
+            f"(sp_ring_degree must be 1, got {get_sp_ring_world_size()})"
+        )
+    assert attn_mask is None, "long context attention does not support attention mask"
+
+    # scatter heads (dim=2), gather sequence (dim=1) — same as video_sparse / v1 USP
+    scatter_idx, gather_idx = 2, 1
+    group = get_sp_ulysses_group()
+    q = SeqAllToAll4D.apply(group, q, scatter_idx, gather_idx)
+    k = SeqAllToAll4D.apply(group, k, scatter_idx, gather_idx)
+    v = SeqAllToAll4D.apply(group, v, scatter_idx, gather_idx)
+
+    # Must not call attention() here — it is patched to long_context_attention under SP.
+    out = mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale)
+    out = SeqAllToAll4D.apply(group, out, gather_idx, scatter_idx)
+    return out
+
+
 def long_context_attention(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -354,10 +417,14 @@ def long_context_attention(
         "sage",
         "sparge",
         "vsa",
+        "mindie",
     ]
     assert attn_mask is None, "long context attention does not support attention mask"
     flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM
     if attn_impl is None or attn_impl == "auto":
+        # NPU has no FA/yunchang TORCH_EFFICIENT kernel; pick MindIE Ulysses when available.
+        if q.device.type == "npu" and MINDIE_AVAILABLE:
+            return _npu_ulysses_mindie_attention(q, k, v, attn_mask=attn_mask, scale=scale)
         if FLASH_ATTN_3_AVAILABLE:
             if flash_attn3_compatible:
                 return LongContextAttention(attn_type=AttnType.FA3)(q, k, v, softmax_scale=scale)
@@ -378,6 +445,8 @@ def long_context_attention(
             return LongContextAttention(attn_type=AttnType.FA)(q, k, v, softmax_scale=scale)
         raise ValueError("No available long context attention implementation")
     else:
+        if attn_impl == "mindie":
+            return _npu_ulysses_mindie_attention(q, k, v, attn_mask=attn_mask, scale=scale)
         if attn_impl == "fa3" or attn_impl == "fa3_fp8":
             if not flash_attn3_compatible:
                 raise RuntimeError(

diffsynth_engine/models/basic/transformer_helper.py

@@ -3,6 +3,9 @@
 import torch.nn.functional as F
 import math
 
+from diffsynth_engine.utils.flag import MINDIE_AVAILABLE
+import os
+USE_MINDIESD_FUSE = os.environ.get("USE_MINDIESD_FUSE", "0") == "1"
 
 def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
     return x * (1 + scale) + shift
@@ -80,6 +83,10 @@ def norm(self, x):
         return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
 
     def forward(self, x):
+        if USE_MINDIESD_FUSE and MINDIE_AVAILABLE and x.device.type == "npu" and self.elementwise_affine:
+            import torch_npu
+            return torch_npu.npu_rms_norm(x, self.weight, epsilon=self.eps)[0]
+
         norm_result = self.norm(x.float()).to(x.dtype)
         if self.elementwise_affine:
             return norm_result * self.weight

diffsynth_engine/models/qwen_image/qwen_image_dit.py

@@ -1,3 +1,8 @@
+import os
+from diffsynth_engine.utils.flag import MINDIE_AVAILABLE
+
+USE_MINDIESD_FUSE = os.environ.get("USE_MINDIESD_FUSE", "0") == "1"
+
 import torch
 import torch.nn as nn
 from typing import Any, Dict, List, Tuple, Union, Optional
@@ -156,6 +161,30 @@ def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
 
 
 def apply_rotary_emb_qwen(x: torch.Tensor, freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]]):
+    if USE_MINDIESD_FUSE and MINDIE_AVAILABLE and x.device.type == "npu":
+        from mindiesd import rotary_position_embedding
+        # Cache expanded cos/sin on the tensor object itself.
+        # Python object identity avoids data_ptr collision across different-length slices.
+        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
+        output = rotary_position_embedding(
+            x, cos, sin,
+            rotated_mode="rotated_interleaved",
+            head_first=False,
+            fused=True,
+        )
+        return output
+
     x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))  # (b, s, h, d) -> (b, s, h, d/2, 2)
     x_out = torch.view_as_real(x_rotated * freqs_cis.unsqueeze(1)).flatten(3)  # (b, s, h, d/2, 2) -> (b, s, h, d)
     return x_out.type_as(x)
@@ -279,7 +308,7 @@ def __init__(
         self.txt_mlp = QwenFeedForward(dim=dim, dim_out=dim, device=device, dtype=dtype)
         self.zero_cond_t = zero_cond_t
 
-    def _modulate(self, x, mod_params, index=None):
+    def _split_mod_params(self, mod_params, index=None):
         shift, scale, gate = mod_params.chunk(3, dim=-1)
         if index is not None:
             actual_batch = shift.size(0) // 2
@@ -299,7 +328,23 @@ def _modulate(self, x, mod_params, index=None):
             shift_result = shift.unsqueeze(1)
             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):
+        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 mindiesd.layernorm_scale_shift when enabled."""
+        if USE_MINDIESD_FUSE and MINDIE_AVAILABLE and x.device.type == "npu":
+            from mindiesd import layernorm_scale_shift
+
+            shift_result, scale_result, gate_result = self._split_mod_params(mod_params, index)
+            # out = LN(x) * (1 + scale) + shift; gate stays separate for residual.
+            out = layernorm_scale_shift(norm, x, scale_result, shift_result, fused=True)
+            return out, gate_result
+        return self._modulate(norm(x), mod_params, index)
 
     def forward(
         self,
@@ -316,11 +361,8 @@ def forward(
             temb = torch.chunk(temb, 2, dim=0)[0]
         txt_mod_attn, txt_mod_mlp = self.txt_mod(temb).chunk(2, dim=-1)  # [B, 3*dim] each
 
-        img_normed = self.img_norm1(image)
-        img_modulated, img_gate = self._modulate(img_normed, img_mod_attn, modulate_index)
-
-        txt_normed = self.txt_norm1(text)
-        txt_modulated, txt_gate = self._modulate(txt_normed, txt_mod_attn)
+        img_modulated, img_gate = self._norm_modulate(self.img_norm1, image, img_mod_attn, modulate_index)
+        txt_modulated, txt_gate = self._norm_modulate(self.txt_norm1, text, txt_mod_attn)
 
         img_attn_out, txt_attn_out = self.attn(
             image=img_modulated,
@@ -329,20 +371,18 @@ def forward(
             attn_mask=attn_mask,
             attn_kwargs=attn_kwargs,
         )
-        image = image + img_gate * img_attn_out
-        text = text + txt_gate * txt_attn_out
-
-        img_normed_2 = self.img_norm2(image)
-        img_modulated_2, img_gate_2 = self._modulate(img_normed_2, img_mod_mlp, modulate_index)
+        # addcmul: residual + gate * out — prefer single op over Mul+Add on NPU
+        image = torch.addcmul(image, img_gate, img_attn_out)
+        text = torch.addcmul(text, txt_gate, txt_attn_out)
 
-        txt_normed_2 = self.txt_norm2(text)
-        txt_modulated_2, txt_gate_2 = self._modulate(txt_normed_2, txt_mod_mlp)
+        img_modulated_2, img_gate_2 = self._norm_modulate(self.img_norm2, image, img_mod_mlp, modulate_index)
+        txt_modulated_2, txt_gate_2 = self._norm_modulate(self.txt_norm2, text, txt_mod_mlp)
 
         img_mlp_out = self.img_mlp(img_modulated_2)
         txt_mlp_out = self.txt_mlp(txt_modulated_2)
 
-        image = image + img_gate_2 * img_mlp_out
-        text = text + txt_gate_2 * txt_mlp_out
+        image = torch.addcmul(image, img_gate_2, img_mlp_out)
+        text = torch.addcmul(text, txt_gate_2, txt_mlp_out)
 
         return text, image
 

diffsynth_engine/pipelines/qwen_image.py

@@ -376,7 +376,9 @@ def update_weights(self, state_dicts: QwenImageStateDicts) -> None:
         self.update_component(self.vae, state_dicts.vae, self.config.device, self.config.vae_dtype)
 
     def compile(self):
-        self.dit.compile_repeated_blocks()
+        from diffsynth_engine.platforms import resolve_platform
+        platform_cls = resolve_platform(self.config.device)
+        self.dit.compile_repeated_blocks(**platform_cls.compile_kwargs())
 
     def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
         assert self.config.tp_degree is None or self.config.tp_degree == 1, (