Fix Ascend NPU RMSNorm and fused-attention mask shapes

Author: xfeng

src/diffusers/models/attention_dispatch.py

@@ -1893,13 +1893,21 @@ def _maybe_modify_attn_mask_npu(query: torch.Tensor, key: torch.Tensor, attn_mas
     if attn_mask is not None and torch.all(attn_mask != 0):
         attn_mask = None
 
-    # Reshape Attention Mask: [batch_size, seq_len_k] or [batch_size, 1, 1, seq_len_k] -> [batch_size, 1, sqe_len_q, seq_len_k]
+    # Reshape Attention Mask for Ascend fused attention, which does not broadcast a
+    # singleton query-length dim the way SDPA does.
+    # Supported expansions:
+    #   [B, Skv] / [B, 1, 1, Skv] / [B, N, 1, Skv] -> [B, 1|N, Sq, Skv]
     # https://www.hiascend.com/document/detail/zh/Pytorch/730/apiref/torchnpuCustomsapi/docs/context/torch_npu-npu_fusion_attention.md
     if attn_mask is not None:
         if attn_mask.ndim == 2 and attn_mask.shape[0] == query.shape[0] and attn_mask.shape[1] == key.shape[1]:
             batch_size, seq_len_q, seq_len_kv = attn_mask.shape[0], query.shape[1], key.shape[1]
             attn_mask = attn_mask.unsqueeze(1).expand(batch_size, seq_len_q, seq_len_kv).unsqueeze(1).contiguous()
-        elif attn_mask.ndim == 4 and attn_mask.shape[1:3] == (1, 1):
+        elif (
+            attn_mask.ndim == 4
+            and attn_mask.shape[2] == 1
+            and attn_mask.shape[-1] == key.shape[1]
+        ):
+            # Expand Sq for masks such as [B, 1, 1, Skv] or LTX cross-attn [B, N, 1, Skv].
             attn_mask = attn_mask.expand(-1, -1, query.shape[1], -1).contiguous()
 
         attn_mask = ~attn_mask.to(torch.bool)

src/diffusers/models/normalization.py

@@ -539,13 +539,15 @@ def __init__(self, dim, eps: float, elementwise_affine: bool = True, bias: bool
                 self.bias = nn.Parameter(torch.zeros(dim))
 
     def forward(self, hidden_states):
-        if is_torch_npu_available():
+        # `npu_rms_norm` requires a gamma tensor. Models that disable the affine
+        # transform (e.g. LTX-2 with `elementwise_affine=False`) keep
+        # `self.weight is None` and must fall back to the pure PyTorch path.
+        if is_torch_npu_available() and self.weight is not None:
             import torch_npu
 
-            if self.weight is not None:
-                # convert into half-precision if necessary
-                if self.weight.dtype in [torch.float16, torch.bfloat16]:
-                    hidden_states = hidden_states.to(self.weight.dtype)
+            # convert into half-precision if necessary
+            if self.weight.dtype in [torch.float16, torch.bfloat16]:
+                hidden_states = hidden_states.to(self.weight.dtype)
             hidden_states = torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=self.eps)[0]
             if self.bias is not None:
                 hidden_states = hidden_states + self.bias