Native NPU attention incorrectly converts additive attention masks to boolean masks

[mengchengTang]

Describe the bug

I found this issue while testing dg845/LTX-2.3-Diffusers with the verl-omni LTX-2.3 NPU training script.

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:

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.

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:

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

[zhtmike]

Which model do you encounter this issue with?

[mengchengTang]

Which model do you encounter this issue with?

I encountered this issue with dg845/LTX-2.3-Diffusers while running the LTX-2.3 text-to-audio-video FlowGRPO workload on Ascend NPU with verl-omni. I have updated the issue with the model and the background of the triggering scenario, and added the end-to-end and minimal reproduction test results to the PR.