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),
• faster than both alternatives at SDXL slice dimensions (fp16, 10×4096×64, M5 Pro, torch 2.13.0; median of 5 trials): 2.07 ms, vs 3.16 ms for baddbmm+empty (the buggy path, -34.7%) and 4.56 ms for baddbmm+zeros (the correct alternative, -54.7%). Zero-initialising costs 44% on top of the buggy path: it writes a 320 MiB scores-sized tensor that is then entirely overwritten (~1.6 ms of pure waste), and bmm is additionally the faster kernel for this shape even when the buffer is already allocated (2.03 ms vs 3.19 ms),
• 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 per the contribution policy: AI-assisted debugging and drafting; all experiments were run and verified by me 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).

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

[dineshyadav03]

Verified the diff: the MPS branch replaces baddbmm(torch.empty(...), beta=0, alpha=self.scale) with a buffer-free torch.bmm(query * self.scale, key.mT), applied identically in both call sites. Root cause and repro are convincing — MPS not honoring beta=0's "input ignored" contract is a real PyTorch-side gap.

One question: the PR frames the 35% speedup as bmm vs. the buggy empty+baddbmm path. Do you have numbers comparing bmm against a zeros-init baddbmm rather than against the broken baseline? That's the actually relevant comparison for justifying the larger code-path split.

Also: both regression tests are skipif(torch_device != "mps"), and HF CI doesn't appear to run MPS runners — so this fix has no automated regression coverage. Is there an MPS runner in CI?

Approve, with the zeros-vs-bmm perf question as non-blocking follow-up.

[RudraMantri123]

Thanks @dineshyadav03 — both questions are fair, and the first one is a good catch. Answering with fresh numbers rather than the ones in the description.

1. bmm vs a zeros-init baddbmm (the comparison that actually justifies the split). You're right that benchmarking against the broken path proves the wrong thing. Re-measured all three variants at SDXL slice dimensions (10×4096×64, fp16, M5 Pro, torch 2.13.0; 10 warmup + 50 reps, median of 5 trials):

variant	median	vs this PR
baddbmm + torch.empty (buggy baseline)	3.16 ms	—
baddbmm + torch.zeros (correct alternative)	4.56 ms	—
bmm(query * scale) (this PR)	2.07 ms	—

So bmm is 34.7% faster than the empty baseline and 54.7% faster than the zeros variant. Zeroing costs 44% on top of the buggy path, because it writes a full scores-sized tensor (256 MB at these dimensions) that is then completely overwritten.

That is the real argument for the code-path split, and it is stronger than the one I put in the description: the alternative fix is not merely slower than this one, it is slower than the bug. bmm also skips the allocation entirely, which matters on the memory-constrained devices that enable_attention_slicing() exists to serve. I'll update the description to lead with this comparison instead.

2. No MPS in CI — correct, and worth stating plainly. These tests provide no automated regression coverage in HF CI. There is no MPS runner, and adding a hosted macOS one would not fix it: I hit this directly while building a separate MPS test harness, where GitHub's macos-14 runners report torch.backends.mps.is_available() == True but then fail every allocation with MPS backend out of memory (MPS allocated: 0 bytes). So the flag is unreliable there and the checks cannot actually execute — real coverage needs a self-hosted Apple Silicon runner.

Given that, what the tests are for: they fail deterministically on main for anyone running the suite on Apple hardware, they document the invariant at the call site, and they let a reviewer without a Mac verify the claim by reading rather than trusting me. I would rather ship that than nothing, but I don't want to overstate it as CI protection — it isn't.

If maintainers would prefer these gated differently (e.g. marked as hardware-dependent, or moved to a slow/nightly suite), happy to adjust.

[RudraMantri123]

Correcting two things in my comment above — I checked my own numbers more carefully and one figure was wrong, and the explanation was incomplete in a way that undersells the result.

The tensor size was wrong. I wrote 256 MB. The scores tensor at these dimensions is 10 × 4096 × 4096 fp16 = 167,772,160 elements = 320 MiB (336 MB), not 256 MB.

And the cost is not only the zero-fill. I attributed the whole gap to writing a buffer that gets overwritten. Decomposing it properly, with the buffer hoisted out of the timed region so allocation and kernel are separated (same setup as before, median of 3 trials × 40 reps):

step	median
torch.zeros(10,4096,4096) — allocate + fill	1.62 ms
torch.empty(10,4096,4096) — allocate only	~0.00 ms
baddbmm on a pre-allocated zeros buffer	3.11 ms
baddbmm on a pre-allocated empty buffer	3.19 ms
bmm(query * scale)	2.03 ms

Two things fall out of this:

1. The zero-fill really does cost ~1.6 ms, and it is pure waste — baddbmm takes the same time whether the buffer it ignores was zeroed or not (3.11 vs 3.19 ms), which is the clearest possible demonstration that the buffer's contents are irrelevant to the computation and only its bytes are being paid for.
2. The rest of the win is the kernel, not the allocation. bmm at 2.03 ms beats baddbmm at 3.19 ms even when the buffer is already allocated and free. So of the ~2.5 ms saved versus the zeros variant, roughly 1.6 ms is skipping the fill and roughly 1.1 ms is bmm simply being the faster kernel for this shape.

The end-to-end numbers in my previous comment stand (2.07 / 3.16 / 4.56 ms); it is the attribution that was too simple. Net effect on the argument: unchanged in direction, stronger in substance — the buffer-free path wins on two independent counts rather than one.

Separately, since both PRs have been open a while: I re-verified both against current main (77d9eb4). Each merges cleanly, and their tests pass on the merged tree — including with both PRs applied together, in case they land in either order.