huggingface / huggingface/diffusers

Native NPU attention incorrectly converts additive attention masks to boolean masks

Open
#14,397 2 comments 0 reactions 0 assignees View on GitHub
bug models
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

### 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

Contributor guide

Open the contributing guide

Research direction

Start in src/diffusers/models/attention_dispatch.py by reading _native_npu_attention and _maybe_modify_attn_mask_npu, then run the isolated additive-mask reproduction from the issue. Done means floating-point masks use the PyTorch SDPA path, boolean masks retain the NPU path, and unsupported mask dtypes are rejected rather than silently converted.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.