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