huggingface / huggingface/diffusers

Native NPU attention incorrectly converts additive attention masks to boolean masks

Abierto
#14,397 2 comentarios 0 reacciones 0 asignados Ver en GitHub
bug models
Lenguaje dominante
Python
Estrellas
34.5k
Forks
7.3k
Merge medio
3 d 3 h
PR fusionados (30 d)
91

Descripción

### Describe the bug

I found this issue while testing `dg845/LTX-2.3-Diffusers` with the [verl-omni LTX-2.3 NPU training script](https://github.com/verl-project/verl-omni/blob/main/examples/flowgrpo_trainer/ltx2/run_ltx2_3_t2av_lora_npu.sh).

This issue has only been observed and validated with LTX-2.3. During testing, the train-inference consistency error was around `1e-3`, while the usual empirical value in verl-omni is around `1e-5`. Further investigation traced the discrepancy to the handling of floating-point attention masks in `_native_npu_attention`.

Currently, every non-`None` attention mask is passed to `_maybe_modify_attn_mask_npu`. This helper assumes that the input is a boolean keep mask, converts it to `torch.bool`, and then inverts it because PyTorch SDPA and `npu_fusion_attention` use opposite boolean-mask polarities:

- PyTorch SDPA: `True` means attend.
- `npu_fusion_attention`: `True` means discard.

However, floating-point masks in PyTorch SDPA are additive biases:

- `0.0` means attend without changing the attention score.
- `-10000.0` or `-inf` means discard.

Casting an additive mask to boolean and then inverting it reverses its semantics:

```text
Additive mask: [0.0, 0.0, -10000.0, -10000.0]
Current NPU result: [True, True, False, False]
Expected block mask: [False, False, True, True]
```

As a result, valid positions may be discarded while masked positions may be attended, causing a silent train-inference consistency error instead of a runtime failure.

Affected code:

https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_dispatch.py

### Proposed fix

- Route floating-point additive masks through the existing PyTorch SDPA path.
- Keep boolean masks on the `npu_fusion_attention` path.
- Reject unsupported mask dtypes instead of silently converting them to boolean.

### Reproduction

The following isolated example demonstrates the incorrect conversion in the current helper. After the proposed routing change, floating-point masks will not be passed to this helper.

```python
import torch

from diffusers.models.attention_dispatch import _maybe_modify_attn_mask_npu

query = torch.zeros(1, 2, 1, 8)
key = torch.zeros(1, 4, 1, 8)

# PyTorch SDPA additive mask:
# 0.0 = attend, -10000.0 = discard
additive_mask = torch.tensor([[[[0.0, 0.0, -10000.0, -10000.0]]]])

actual = _maybe_modify_attn_mask_npu(query, key, additive_mask)
expected = torch.tensor(
[[[[False, False, True, True], [False, False, True, True]]]]
)

print("Actual:")
print(actual)

print("Expected:")
print(expected)

torch.testing.assert_close(actual, expected)
```

The assertion fails because the current result is:

```text
[[[[ True, True, False, False],
[ True, True, False, False]]]]
```

### System Info

- OS: Linux aarch64
- Hardware: Ascend NPU
- Python: 3.11
- PyTorch: 2.10.0
- torch-npu: 2.10.0
- CANN: 9.0.0
- diffusers: `main`

### Who can help?

@DN6 @yiyixuxu

Guía de contribución

Abrir la guía de contribución

Línea de trabajo

Empieza en src/diffusers/models/attention_dispatch.py leyendo _native_npu_attention y _maybe_modify_attn_mask_npu, y luego ejecuta la reproducción aislada de la máscara aditiva del issue. Se considera terminado cuando las máscaras de punto flotante utilizan la ruta de PyTorch SDPA, las máscaras booleanas conservan la ruta de NPU y los dtypes de máscara no compatibles se rechazan en lugar de convertirse silenciosamente.

Escrito por el modelo de indexación a partir del texto del issue.

Evaluación

Stack tecnológico
python, pytorch
Área
machine-learning
Tipo de issue
Error
Dificultad
3/5
Tiempo estimado
1-2 días
Estado de actividad
Tranquilo
Claridad
Bien especificado
Aptitud para principiantes
68/100

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.