diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 9414c151fd67..cfb3af501071 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -920,6 +920,11 @@ def _cudnn_attention_forward_op( if enable_gqa: raise ValueError("`enable_gqa` is not yet supported for cuDNN attention.") + # The aten op takes an additive bias, so a boolean mask has to be converted the same way + # `F.scaled_dot_product_attention` does before dispatching to it. + if attn_mask is not None and attn_mask.dtype == torch.bool: + attn_mask = torch.zeros_like(attn_mask, dtype=query.dtype).masked_fill_(attn_mask.logical_not(), float("-inf")) + tensors_to_save = () # Contiguous is a must here! Calling cuDNN backend with aten ops produces incorrect results diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py new file mode 100644 index 000000000000..67f9e17b52bb --- /dev/null +++ b/tests/models/test_attention_dispatch.py @@ -0,0 +1,48 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers.models.attention_dispatch import _cudnn_attention_forward_op + +from ..testing_utils import assert_tensors_close, is_attention, require_torch_gpu, torch_device + + +@is_attention +@require_torch_gpu +class TestCudnnAttentionForwardOp: + @pytest.mark.parametrize("mask_type", ["partial", "fully_masked_row"]) + def test_boolean_attn_mask_matches_sdpa(self, mask_type): + batch_size, num_heads, seq_len, head_dim = 1, 2, 16, 64 + torch.manual_seed(0) + + # the forward op takes `(batch_size, seq_len, num_heads, head_dim)` + query, key, value = ( + torch.randn(batch_size, seq_len, num_heads, head_dim, device=torch_device, dtype=torch.bfloat16) + for _ in range(3) + ) + attn_mask = torch.ones(batch_size, num_heads, seq_len, seq_len, device=torch_device, dtype=torch.bool) + if mask_type == "partial": + attn_mask[..., 3, 5:] = False + else: + attn_mask[..., 7, :] = False + + out = _cudnn_attention_forward_op(None, query, key, value, attn_mask=attn_mask, _save_ctx=False) + expected = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), attn_mask=attn_mask + ).transpose(1, 2) + + assert_tensors_close(out, expected, atol=1e-2, rtol=1e-2, msg=f"cuDNN forward op with {mask_type} mask")