From 92bd5b5f3dcd9a400ea0a41def277ef789f0fdcb Mon Sep 17 00:00:00 2001 From: Rudra Mantri <189433012+RudraMantri123@users.noreply.github.com> Date: Wed, 12 Aug 2026 23:11:57 +0530 Subject: [PATCH 1/2] Fix NaN attention scores on MPS from uninitialized baddbmm buffer On MPS, torch.baddbmm does not honor the documented beta=0 semantics: NaN/Inf present in the input buffer propagate to the output. Both copies of get_attention_scores (Attention and AttentionModuleMixin) pass torch.empty() as that buffer, so recycled allocator pages containing NaN poison the attention scores, producing all-black images with SlicedAttnProcessor (e.g. SDXL + enable_model_cpu_offload + enable_attention_slicing). Use a buffer-free scaled bmm on MPS when there is no attention mask. This avoids relying on the beta=0 contract, skips the scores-sized buffer allocation entirely (lower peak memory on the memory-constrained devices the sliced path targets), and benchmarks ~35% faster than the baddbmm+empty path on Apple Silicon. Other backends are unchanged. Fixes #14438 --- src/diffusers/models/attention.py | 39 +++++++------ src/diffusers/models/attention_processor.py | 39 +++++++------ tests/models/test_attention_processor.py | 61 +++++++++++++++++++++ 3 files changed, 107 insertions(+), 32 deletions(-) diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index 5d9490503974..6d54e607c4f7 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -420,23 +420,30 @@ def get_attention_scores( query = query.float() 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 - ) - beta = 0 + if attention_mask is None and query.device.type == "mps": + # On MPS, baddbmm does not honor the documented beta=0 semantics: NaN/Inf in the + # uninitialized `input` buffer propagate to the output, producing NaN attention + # scores (https://github.com/huggingface/diffusers/issues/14438). A buffer-free + # scaled bmm avoids relying on that contract (and skips the buffer allocation). + attention_scores = torch.bmm(query * self.scale, key.transpose(-1, -2)) else: - baddbmm_input = attention_mask - beta = 1 - - attention_scores = torch.baddbmm( - baddbmm_input, - query, - key.transpose(-1, -2), - beta=beta, - alpha=self.scale, - ) - del baddbmm_input + if attention_mask is None: + 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 + beta = 1 + + attention_scores = torch.baddbmm( + baddbmm_input, + query, + key.transpose(-1, -2), + beta=beta, + alpha=self.scale, + ) + del baddbmm_input if self.upcast_softmax: attention_scores = attention_scores.float() diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 1b923e749663..22ab65e815b5 100755 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -675,23 +675,30 @@ def get_attention_scores( query = query.float() 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 - ) - beta = 0 + if attention_mask is None and query.device.type == "mps": + # On MPS, baddbmm does not honor the documented beta=0 semantics: NaN/Inf in the + # uninitialized `input` buffer propagate to the output, producing NaN attention + # scores (https://github.com/huggingface/diffusers/issues/14438). A buffer-free + # scaled bmm avoids relying on that contract (and skips the buffer allocation). + attention_scores = torch.bmm(query * self.scale, key.transpose(-1, -2)) else: - baddbmm_input = attention_mask - beta = 1 - - attention_scores = torch.baddbmm( - baddbmm_input, - query, - key.transpose(-1, -2), - beta=beta, - alpha=self.scale, - ) - del baddbmm_input + if attention_mask is None: + 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 + beta = 1 + + attention_scores = torch.baddbmm( + baddbmm_input, + query, + key.transpose(-1, -2), + beta=beta, + alpha=self.scale, + ) + del baddbmm_input if self.upcast_softmax: attention_scores = attention_scores.float() diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py index a2b02b56692c..5ca91311960e 100644 --- a/tests/models/test_attention_processor.py +++ b/tests/models/test_attention_processor.py @@ -132,3 +132,64 @@ 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) + + +class TestGetAttentionScoresMPS: + # Regression tests for https://github.com/huggingface/diffusers/issues/14438. + # On MPS, baddbmm propagates NaN/Inf from `input` even with beta=0, so + # get_attention_scores must not pass uninitialized memory as the buffer. + # Poison the allocator pool so a subsequent torch.empty of the same shape + # recycles NaN-bearing pages, then verify scores stay finite and correct. + + batch, tokens, dim_head = 8, 4096, 32 + + def _make_qk(self): + query = torch.randn(self.batch, self.tokens, self.dim_head, device="mps", dtype=torch.float16) + key = torch.randn(self.batch, self.tokens, self.dim_head, device="mps", dtype=torch.float16) + return query, key + + def _poison_pool(self): + # Fill and free a buffer of exactly the attention-scores shape so the + # allocator hands its NaN-bearing pages to the next torch.empty call. + junk = torch.full((self.batch, self.tokens, self.tokens), float("nan"), device="mps", dtype=torch.float16) + del junk + + @pytest.mark.skipif(torch_device != "mps", reason="regression test for an MPS-specific baddbmm issue") + def test_get_attention_scores_no_nan_from_recycled_buffer(self): + from types import SimpleNamespace + + from diffusers.models.attention import AttentionModuleMixin + + # Exercise both duplicated implementations of get_attention_scores in one + # process: allocator page-recycling on MPS is position-dependent, so a + # sequence of poisoned calls across both paths is what detects the leak + # deterministically. + attn = Attention(query_dim=64, heads=2, dim_head=32) + holder = SimpleNamespace(upcast_attention=False, upcast_softmax=False, scale=0.125) + query, key = self._make_qk() + + score_fns = [ + ("Attention", lambda: attn.get_attention_scores(query, key, attention_mask=None)), + ( + "AttentionModuleMixin", + lambda: AttentionModuleMixin.get_attention_scores(holder, query, key, attention_mask=None), + ), + ] + for round_idx in range(3): + for name, scores_fn in score_fns: + self._poison_pool() + scores = scores_fn() + assert not torch.isnan(scores).any(), ( + f"NaN leaked from uninitialized baddbmm buffer on MPS ({name}, round {round_idx})" + ) + + @pytest.mark.skipif(torch_device != "mps", reason="regression test for an MPS-specific baddbmm issue") + def test_get_attention_scores_matches_cpu_reference(self): + # The MPS path must stay numerically equivalent to the CPU baddbmm path, + # not merely NaN-free. + attn = Attention(query_dim=64, heads=2, dim_head=32) + query, key = self._make_qk() + self._poison_pool() + probs_mps = attn.get_attention_scores(query, key, attention_mask=None).cpu() + probs_cpu = attn.get_attention_scores(query.cpu(), key.cpu(), attention_mask=None) + assert torch.allclose(probs_mps, probs_cpu, atol=2e-3), "MPS attention probs diverge from CPU reference" From edbe78a939aaaae07a73f2178991d802d2f76cfe Mon Sep 17 00:00:00 2001 From: Rudra Mantri <189433012+RudraMantri123@users.noreply.github.com> Date: Thu, 27 Aug 2026 13:16:31 +0530 Subject: [PATCH 2/2] Add a device-agnostic equivalence test for get_attention_scores The existing regression tests here are gated on torch_device == "mps", so they are skipped on every runner in CI and give a reviewer without Apple hardware no signal at all. Add a companion test that runs on any backend and pins the property that matters off MPS: get_attention_scores must match an independently computed softmax(scale * q @ k^T) reference, across masked/unmasked inputs and all four combinations of upcast_attention / upcast_softmax, for both duplicated implementations. This is not merely decorative on CPU. Reintroducing the underlying bug in the non-MPS path -- flipping the unmasked branch's beta back to 1 so the uninitialized torch.empty buffer is added into the scores -- fails 4 of the 8 CPU cases. Corrupting the new MPS bmm branch fails 5 cases on MPS. So the two classes together cover both branches. --- tests/models/test_attention_processor.py | 56 ++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py index 5ca91311960e..ecf6322f959b 100644 --- a/tests/models/test_attention_processor.py +++ b/tests/models/test_attention_processor.py @@ -193,3 +193,59 @@ def test_get_attention_scores_matches_cpu_reference(self): probs_mps = attn.get_attention_scores(query, key, attention_mask=None).cpu() probs_cpu = attn.get_attention_scores(query.cpu(), key.cpu(), attention_mask=None) assert torch.allclose(probs_mps, probs_cpu, atol=2e-3), "MPS attention probs diverge from CPU reference" + + +class TestGetAttentionScoresEquivalence: + # Device-agnostic companion to TestGetAttentionScoresMPS. The MPS-specific + # regression tests above can only run on Apple hardware, so this class pins the + # property a reviewer on any other backend cares about: routing the unmasked + # case through a scaled `bmm` must not change what `get_attention_scores` + # returns, on any device, for any combination of the upcast flags. + + @staticmethod + def _reference(query, key, scale, attention_mask, upcast_attention, upcast_softmax): + dtype = query.dtype + if upcast_attention: + query, key = query.float(), key.float() + scores = scale * torch.bmm(query, key.transpose(-1, -2)) + if attention_mask is not None: + scores = scores + attention_mask + if upcast_softmax: + scores = scores.float() + return scores.softmax(dim=-1).to(dtype) + + @pytest.mark.parametrize("masked", [False, True]) + @pytest.mark.parametrize("upcast_attention", [False, True]) + @pytest.mark.parametrize("upcast_softmax", [False, True]) + def test_matches_reference_on_current_device(self, masked, upcast_attention, upcast_softmax): + from types import SimpleNamespace + + from diffusers.models.attention import AttentionModuleMixin + + batch, tokens, dim_head = 4, 16, 8 + torch.manual_seed(0) + query = torch.randn(batch, tokens, dim_head, device=torch_device) + key = torch.randn(batch, tokens, dim_head, device=torch_device) + attention_mask = torch.randn(batch, tokens, tokens, device=torch_device) if masked else None + + attn = Attention( + query_dim=dim_head * 2, + heads=2, + dim_head=dim_head, + upcast_attention=upcast_attention, + upcast_softmax=upcast_softmax, + ) + expected = self._reference(query, key, attn.scale, attention_mask, upcast_attention, upcast_softmax) + + # Both duplicated implementations must agree with the reference. + holder = SimpleNamespace(upcast_attention=upcast_attention, upcast_softmax=upcast_softmax, scale=attn.scale) + actual = { + "Attention": attn.get_attention_scores(query, key, attention_mask=attention_mask), + "AttentionModuleMixin": AttentionModuleMixin.get_attention_scores( + holder, query, key, attention_mask=attention_mask + ), + } + for name, probs in actual.items(): + assert probs.shape == (batch, tokens, tokens), name + assert torch.isfinite(probs).all(), name + assert torch.allclose(probs, expected, atol=1e-6), f"{name} diverges from the reference formulation"