Native NPU attention incorrectly converts additive attention masks to boolean masks

[mengchengTang]

Describe the bug

_native_npu_attention currently passes every non-None attention mask to _maybe_modify_attn_mask_npu.

The helper assumes that the input is a boolean/binary keep mask, converts it to torch.bool, and then inverts it because PyTorch SDPA and npu_fusion_attention use opposite boolean-mask polarities.

This is correct for boolean keep masks:

• PyTorch SDPA: True means attend.
• npu_fusion_attention: True means discard.

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

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

Casting such a mask to boolean loses its additive values, and the subsequent inversion reverses the intended mask:

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, positions intended to be attended may be discarded, while masked positions may be attended.
I intend to submit a PR that:
Routes floating-point additive masks through the existing native SDPA implementation.
Keeps boolean masks on the npu_fusion_attention path.
Rejects unsupported mask dtypes instead of silently converting them to boolean.
Affected code:
https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_dispatch.py

Reproduction

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)

Logs

System Info

OS: Linux aarch64
Hardware: Ascend NPU
Python: 3.11
torch / torch_npu / CANN: 9.0.0
diffusers: main

Who can help?

@DN6 @yiyixuxu