From 2027ec3a2abb9416181a49890c750e6fa5a32f70 Mon Sep 17 00:00:00 2001 From: Shreyas-jk Date: Fri, 8 May 2026 15:43:36 -0700 Subject: [PATCH 1/2] Fix NaN attention scores on MPS when attention_mask is None (#11229) --- src/diffusers/models/attention_processor.py | 13 +++++++--- tests/models/test_attention_processor.py | 28 +++++++++++++++++++++ 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 1b923e7496639..8561387b6ad0b 100755 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -676,9 +676,16 @@ def get_attention_scores( key = key.float() if attention_mask is None: - baddbmm_input = torch.empty( - query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device - ) + if query.device.type == "mps": + # MPS' baddbmm does not short-circuit on beta=0, so an + # uninitialized input from torch.empty() can propagate NaN. + baddbmm_input = torch.zeros( + query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device + ) + else: + baddbmm_input = torch.empty( + query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device + ) beta = 0 else: baddbmm_input = attention_mask diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py index a2b02b56692c6..7bdf0ca35ad61 100644 --- a/tests/models/test_attention_processor.py +++ b/tests/models/test_attention_processor.py @@ -132,3 +132,31 @@ def test_conversion_when_using_device_map(self): assert np.allclose(pre_conversion, conversion, atol=1e-3) assert np.allclose(conversion, after_conversion, atol=1e-3) + + +@pytest.mark.skipif(torch_device != "mps", reason="test exercises an MPS-specific code path") +def test_no_nan_when_attention_mask_is_none_on_mps(): + # Regression test: torch.empty() on MPS can return non-finite values, + # and MPS' baddbmm does not short-circuit on beta=0, so an unmasked + # call to get_attention_scores used to propagate NaN into the output. + torch.manual_seed(0) + heads, dim_head, seq_len = 4, 32, 256 + attn = Attention( + query_dim=heads * dim_head, + heads=heads, + dim_head=dim_head, + bias=False, + ).to(torch_device, torch.float16) + + for _ in range(20): + # Pollute the MPS allocator pool with non-finite values so that a + # subsequent torch.empty() is likely to return NaN-filled memory. + polluter = torch.full((heads, seq_len, seq_len), float("nan"), device=torch_device, dtype=torch.float16) + del polluter + + query = torch.randn(1, seq_len, heads * dim_head, device=torch_device, dtype=torch.float16) + key = torch.randn(1, seq_len, heads * dim_head, device=torch_device, dtype=torch.float16) + scores = attn.get_attention_scores( + attn.head_to_batch_dim(query), attn.head_to_batch_dim(key), attention_mask=None + ) + assert not torch.isnan(scores).any().item(), "attention scores contain NaN on MPS" From dd10d2d5d4b52b3ffdaf3865ae0677b83c49bbe9 Mon Sep 17 00:00:00 2001 From: Shreyas-jk Date: Sat, 27 Jun 2026 08:28:29 -0700 Subject: [PATCH 2/2] Address review: gate MPS baddbmm workaround on torch<2.14 (pytorch#187522), port test to pytest Co-Authored-By: Claude Opus 4.8 (1M context) --- src/diffusers/models/attention_processor.py | 7 ++++--- tests/models/test_attention_processor.py | 6 +++++- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 8561387b6ad0b..14cfe3288379b 100755 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -676,9 +676,10 @@ def get_attention_scores( key = key.float() if attention_mask is None: - if query.device.type == "mps": - # MPS' baddbmm does not short-circuit on beta=0, so an - # uninitialized input from torch.empty() can propagate NaN. + if query.device.type == "mps" and is_torch_version("<", "2.14.0"): + # Before torch 2.14 (pytorch#187522), MPS' baddbmm did not + # short-circuit on beta=0, so an uninitialized input from + # torch.empty() could propagate NaN. Fixed upstream from 2.14. baddbmm_input = torch.zeros( query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device ) diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py index 7bdf0ca35ad61..df32afb3381ce 100644 --- a/tests/models/test_attention_processor.py +++ b/tests/models/test_attention_processor.py @@ -9,7 +9,7 @@ from diffusers import DiffusionPipeline from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor -from ..testing_utils import torch_device +from ..testing_utils import is_torch_version, torch_device class TestAttnAddedKVProcessor: @@ -135,6 +135,10 @@ def test_conversion_when_using_device_map(self): @pytest.mark.skipif(torch_device != "mps", reason="test exercises an MPS-specific code path") +@pytest.mark.skipif( + is_torch_version(">=", "2.14.0"), + reason="baddbmm beta=0 NaN fixed upstream in pytorch#187522 (torch>=2.14); MPS workaround no longer applied", +) def test_no_nan_when_attention_mask_is_none_on_mps(): # Regression test: torch.empty() on MPS can return non-finite values, # and MPS' baddbmm does not short-circuit on beta=0, so an unmasked