huggingface / huggingface/diffusers

Native NPU attention incorrectly converts additive attention masks to boolean masks

オープン
#14,397 コメント 2 件 リアクション 0 件 担当者 0 名 GitHub で見る
bug models
主要言語
Python
スター
34.5k
フォーク
7.3k
平均マージ
3日 3時間
マージ済み PR(30日)
91

説明

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

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

src/diffusers/models/attention_dispatch.py で _native_npu_attention と _maybe_modify_attn_mask_npu を読み始め、その後、issue にある加算マスクの分離再現を実行します。浮動小数点マスクが PyTorch SDPA パスを使用し、ブールマスクが NPU パスを維持し、サポートされていないマスクの dtype が暗黙的に変換されず拒否されれば完了です。

索引モデルが issue の本文から書いたものです。

評価

技術スタック
python, pytorch
領域
machine-learning
issue の種類
バグ
難易度
3/5
見積もり時間
1〜2日
活発さ
静か
明瞭さ
明確に書かれている
初心者へのやさしさ
68/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。