huggingface / huggingface/diffusers
[Bug] Ascend NPU: RMSNorm crashes with elementwise_affine=False; _native_npu FA rejects [B, N, 1, Skv] masks (LTX-2)
- Lingua principale
- Python
- Stelle
- 34.5k
- Fork
- 7.3k
- Merge medio
- 3g 3h
- PR unite (30g)
- 91
Descrizione
### 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
Guida per i contributori
Apri la guida per i contributori
Direzione di ricerca
Inizia con diffusers.models.normalization.RMSNorm e il percorso _maybe_modify_attn_mask_npu in diffusers.models.attention_dispatch, quindi esegui i due sketch di riproduzione Ascend NPU forniti. Il lavoro è completato quando RMSNorm funziona con elementwise_affine=False e _native_npu accetta maschere in stile LTX con forma [B, N, 1, Skv], gestendo correttamente la dimensione della lunghezza della query.
Scritto dal modello di indicizzazione a partire dal testo della issue.
Valutazione
- Stack tecnologico
- python, pytorch
- Ambito
- machine-learning
- Tipo di issue
- Bug
- Difficoltà
- 4/5
- Tempo stimato
- 3-5 giorni
- Stato di attività
- Tranquilla
- Chiarezza
- Abbastanza chiara
- Idoneità per principianti
- 52/100