[Bug] Ascend NPU: RMSNorm crashes with elementwise_affine=False; _native_npu FA rejects [B, N, 1, Skv] masks (LTX-2)

[mengchengTang]

Describe the bug

When running LTX-2 with attn_backend=_native_npu on Ascend NPU, two incompatibilities show up:

1. RMSNorm – layers with elementwise_affine=False leave weight=None, but torch_npu.npu_rms_norm requires a gamma tensor → crash (gamma is None).

2. Fused attention mask – LTX cross-attn uses masks shaped [B, N, 1, Skv] (e.g. [1, 32, 1, 1024]). Ascend FA does not broadcast the singleton query-length dim the way SDPA does, so _maybe_modify_attn_mask_npu must expand it to [B, N, Sq, Skv]. Today only [B, 1, 1, Skv] is expanded.

Reproduction

Requires Ascend NPU + torch_npu. Minimal sketches of both failure modes:

import torch
from diffusers.models.normalization import RMSNorm
from diffusers.models.attention_dispatch import (
    AttentionBackendName,
    attention_backend,
    dispatch_attention_fn,
)

# --- Bug 1: RMSNorm with elementwise_affine=False ---
norm = RMSNorm(dim=64, eps=1e-6, elementwise_affine=False).to("npu")
x = torch.randn(2, 16, 64, device="npu", dtype=torch.float16)
# Crashes: npu_rms_norm called with weight=None
y = norm(x)

# --- Bug 2: FA mask [B, N, 1, Skv] under _native_npu ---
B, Sq, Skv, N, D = 1, 384, 1024, 32, 64
q = torch.randn(B, Sq, N, D, device="npu", dtype=torch.float16)
k = torch.randn(B, Skv, N, D, device="npu", dtype=torch.float16)
v = torch.randn(B, Skv, N, D, device="npu", dtype=torch.float16)
attn_mask = torch.ones(B, N, 1, Skv, device="npu", dtype=torch.float16)  # LTX-style

with attention_backend(AttentionBackendName._NATIVE_NPU):
    # Fails: Ascend FA expects Sq on dim=-2, not a singleton 1
    out = dispatch_attention_fn(q, k, v, attn_mask=attn_mask)

Logs

Error 1

Error 2

System Info

• OS: Linux aarch64
• Hardware: Ascend NPU
• Python: 3.11
• torch / torch_npu / CANN: 9.0.0
• diffusers: main

Who can help?

@yiyixuxu @sayakpaul

[mengchengTang]

I have a fix ready in #14288