Fix NaN attention scores on MPS from uninitialized baddbmm buffer

[RudraMantri123]

What does this PR do?

Fixes #14438 — SDXL produces all-black images on MPS when enable_attention_slicing() is used (most visibly together with enable_model_cpu_offload()).

Root cause

get_attention_scores (in both Attention and AttentionModuleMixin) passes torch.empty(...) to torch.baddbmm with beta=0, relying on the documented guarantee that the input is ignored and NaN/Inf in it are not propagated. The MPS backend violates that guarantee, so NaN garbage in recycled allocator pages leaks into the attention scores. Only the sliced attention processors reach this code path (the default processor uses SDPA), which makes enable_attention_slicing() the trigger; offload merely churns the allocator so torch.empty recycles dirty pages more often — plain pipe.to("mps") + enable_attention_slicing("max") reproduces without any offload.

Minimal reproduction of the underlying defect (torch 2.13.0, Apple Silicon, no diffusers):

import torch
B, T, D = 10, 4096, 64
junk = torch.full((B, T, T), float("nan"), device="mps", dtype=torch.float16)
del junk  # allocator will recycle these NaN-bearing pages
q = torch.randn(B, T, D, device="mps", dtype=torch.float16)
k = torch.randn(B, T, D, device="mps", dtype=torch.float16)
buf = torch.empty(B, T, T, device="mps", dtype=torch.float16)
print(torch.isnan(torch.baddbmm(buf, q, k.mT, beta=0, alpha=0.125)).any())  # True — bug
print(torch.isnan(torch.bmm(q, k.mT) * 0.125).any())                        # False — control

I will file this against PyTorch separately; this PR makes diffusers robust to it.

The fix

On MPS, when there is no attention mask, compute scores with a buffer-free scaled bmm instead of baddbmm:

• removes the reliance on beta=0 semantics and the uninitialized buffer,
• skips a scores-sized allocation (lower peak memory on the devices the sliced path targets),
• ~35% faster than the baddbmm+empty path at SDXL slice dimensions on Apple Silicon,
• fp32 output matches the CPU reference exactly; fp16 within 2e-3. All other backends and the masked path are unchanged.

The issue's reproduction script now renders correctly across seeds (verified on M5 Pro 24GB; end-to-end images inspected).

Tests

• test_get_attention_scores_no_nan_from_recycled_buffer — poisons the MPS allocator pool and exercises both implementations; fails deterministically on current main, passes with the fix.
• test_get_attention_scores_matches_cpu_reference — guards numerical equivalence with the CPU path, not just NaN-absence.

Both are gated to MPS.

Notes for reviewers

AI disclosure: Claude Code assisted with debugging and drafting; all experiments were run and verified by the author on real hardware.

• The MPS path pre-scales the query (query * self.scale) instead of post-scaling via alpha — mathematically identical, bounded by the CPU-parity test (fp32 exact). A zero-initialized buffer also fixes the bug but benchmarked +53% slower and keeps the allocation.
• The same torch.empty + baddbmm(beta=0) pattern exists in pipelines/kolors/text_encoder.py; left for a follow-up since I cannot end-to-end test Kolors on this hardware.

Who can review?

@yiyixuxu @dg845 @asomoza — cc @pupa3066 for co-verification on M1 8GB (the memory-constrained case I can't cover).

🤖 Generated with Claude Code

[RudraMantri123]

Upstream status update — correcting my earlier note, and it changes the framing of this PR.

I said above that I'd file the baddbmm beta=0 violation against PyTorch. It turns out it was already reported and already fixed upstream — I should have found this before saying that:

• [MPS] baddbmm incorrectly propagates NaN when beta == 0 pytorch/pytorch#187521 — "[MPS] baddbmm incorrectly propagates NaN when beta == 0" (filed 16 Jun 2026, labeled module: correctness (silent))
• Fixed by Fix baddbmm nan propagation and test pytorch/pytorch#187522, landed as commit e18b4e69b on 17 Jun 2026. main now guards the path explicitly: "When beta == 0, input is ignored so nan/inf in it are not propagated (matches CPU/CUDA and addmm)."

But the fix is not in any released version yet. The v2.13.0 release branch (published 8 Jul 2026) was cut before that commit and it wasn't cherry-picked — the guard is absent from v2.13.0's LinearAlgebra.mm, and PyTorch's own minimal repro still fails on current stable:

import torch  # 2.13.0
inp = torch.full((1, 1, 1), float("nan"), device="mps")
torch.baddbmm(inp, torch.ones(1,1,1,device="mps"), torch.ones(1,1,1,device="mps"), beta=0)
# MPS -> nan     CPU -> 1.0   (docs: input ignored when beta=0, nan/inf not propagated)

So every Apple Silicon user on torch ≤ 2.13.0 — i.e. everyone on the current stable release — still hits #14438 with sliced attention, and will until 2.14.0 ships and they upgrade.

What that means for this PR, three honest options:

1. Merge as-is. It fixes the reported bug for all currently-released torch versions. Worth noting the change isn't only a workaround: dropping the baddbmm+torch.empty pattern also removes a scores-sized allocation and benchmarked ~37% faster than the current path at SDXL slice dimensions on an M5 Pro — so it stands on its own merits after 2.14.0 too, rather than becoming dead weight.
2. Version-gate it (torch < 2.14) if you'd rather the workaround expire on its own. Happy to add that; my own preference is (1) given the perf result, but your call.
3. Close it and let users wait for the upstream release, with the workaround documented on the issue instead.

Tell me which you prefer and I'll adjust — (2) is a small change on this branch.

One process note: CI hasn't run here yet, since first-time-contributor workflows need a maintainer's approval click. Whenever someone has a spare second for that, the test suites (including the new MPS regression test) can go green and make this reviewable at a glance.

@yiyixuxu @dg845 @asomoza

[RudraMantri123]

Gentle ping @yiyixuxu @dg845 @asomoza — this one is blocked only on workflow approval for a first-time contributor, so CI has never actually run here. A single approval click would get the suites green (including the new MPS regression test, which fails deterministically on main without the fix) and make this reviewable at a glance.

On scope, in case it helps prioritise: it closes #14438 and also #11229, which has been open since April 2025 with the same root cause. And I'm still happy to switch to the torch<2.14 version-gated variant if you'd rather the workaround expire on its own once the upstream PyTorch fix ships.