huggingface / huggingface/diffusers
[Bug] Ascend NPU: RMSNorm crashes with elementwise_affine=False; _native_npu FA rejects [B, N, 1, Skv] masks (LTX-2)
- Vorherrschende Sprache
- Python
- Sterne
- 34.5k
- Forks
- 7.3k
- Ø Merge
- 3 T. 3 Std.
- Gemergte PRs (30 T.)
- 91
Beschreibung
### 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.
These bugs were discovered while running LTX-2.3 text-to-audio-video FlowGRPO training on Ascend NPU using [verl-omni](https://github.com/verl-project/verl-omni/blob/main/examples/flowgrpo_trainer/ltx2/run_ltx2_3_t2av_lora_npu.sh).
### Reproduction
Requires Ascend NPU + `torch_npu`. Minimal sketches of both failure modes:
```python
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.zeros(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: 2.10.0
- torch_npu: 2.10.0
- CANN: 9.0.0
- diffusers: main
### Who can help?
@yiyixuxu @sayakpaul
Beitragsleitfaden
Rechercherichtung
Beginne mit diffusers.models.normalization.RMSNorm und dem _maybe_modify_attn_mask_npu-Pfad in diffusers.models.attention_dispatch und führe anschließend die beiden bereitgestellten Ascend NPU-Reproduktionsskizzen aus. Als abgeschlossen gilt die Aufgabe, wenn RMSNorm mit elementwise_affine=False funktioniert und _native_npu LTX-artige Masken im Format [B, N, 1, Skv] akzeptiert, indem die Query-Längen-Dimension korrekt behandelt wird.
Vom Indexierungsmodell aus dem Issue-Text verfasst.
Bewertung
- Tech-Stack
- python, pytorch
- Bereich
- machine-learning
- Issue-Typ
- Bug
- Schwierigkeit
- 4/5
- Geschätzter Aufwand
- 3-5 Tage
- Aktivitätsstatus
- Ruhig
- Klarheit
- Größtenteils klar
- Anfängerfreundlichkeit
- 52/100