PR #179592: Fix MPS SDPA correctness when score matrix exceeds 2^32 elements

• Repository: pytorch/pytorch
• Author: jbartolozzi
• State: open | merged: False
• Link: https://github.com/pytorch/pytorch/pull/179592

Description (problem / solution / changelog)

Summary

Fixes #179352. Apple's MPSGraph corrupts SDPA outputs when the score matrix B*H*Nq*Nkv exceeds 2^32 total elements: at the canary shape (B=1, H=8, Nq=16384, Nkv=65536, D=64), output diverges from CPU with cosine similarity ~0.44 and max relative error ~1.2. The threshold is independent of dtype, of is_causal, and of whether an explicit attn_mask is supplied — the bug is in the matmul/softmax sequence underneath all four MPSGraph codepaths.

The bug is reachable via sdpa_general_mps in aten/src/ATen/native/mps/operations/Attention.mm, which builds an MPSGraph of the form matrixMultiplicationWithPrimaryTensor (Q·Kᵀ) → optional causal/additive mask → softMaxWithTensor → matrixMultiplicationWithPrimaryTensor (·V). The threshold shape (element-counted, not byte-counted, exactly at 2^32) is consistent with a 32-bit flat element index somewhere inside Apple's matmul or softmax kernel, but the bug has not been confirmed against Apple source. PyTorch's own custom Metal kernels in aten/src/ATen/native/mps/kernels/Attention.metal are not involved for shapes that hit this bug — _scaled_dot_product_attention_math_mps only dispatches to those when query_seq_len <= 8, and the affected shapes have much larger Q.

Note that Attention.mm also defines a third path, sdpa_full_attention_mps, that wires up the custom attention Metal kernel and would not go through MPSGraph at all — but it is dead code on main, never called from the dispatcher (the test_fast_full_attention_* tests are correspondingly all skipped with "Full attention fast kernel not implemented yet"). Enabling it was therefore not a viable fix for this PR.

Since SDPA's softmax is row-wise over the Q axis, splitting Q into chunks and concatenating outputs is mathematically exact. sdpa_general_mps now factors its existing graph build/run into sdpa_general_mps_impl and wraps it: when the score matrix is small enough (≤ 2^32 elements) it calls the impl once, preserving the original codepath byte-for-byte; otherwise it loops over Q slices of size max(1, 2^32 / (B*H*Nkv)) and stitches the outputs back together. For is_causal=True, the wrapper synthesizes a per-chunk additive causal mask from the chunk's Q offset and passes it to the impl with is_causal=false (the in-graph causal path assumes the chunk starts at row 0). User-supplied attn_mask is narrowed along the Q axis. A TORCH_WARN_ONCE guards the unrealistic edge case where B*H*Nkv alone exceeds 2^32.

Suggested review order: read the new sdpa_general_mps_impl first (it's the unchanged old body), then the new sdpa_general_mps wrapper, then the test.

Test plan

• `python test/test_mps.py -k attention -v` — 96 tests, all pass (skipped=24, expected failures=10, identical to baseline)
• `python test/test_mps.py -k test_sdpa_large_score_matrix_179352 -v` — 4/4 pass on a 64 GB Mac (3 variants + meta)
• Empirical sweep across dtypes (fp32/fp16/bf16), `is_causal=True/False`, and bool/additive masks — all formerly failing shapes now match CPU to ~1e-5 max relative error
• Lint clean (`spin lint`)
• Skipped on CI by design (`total_memory >= 64 GB` + `IS_CI` gates)

Authored with the assistance of Claude (Anthropic).

Changed files

• aten/src/ATen/native/mps/operations/Attention.mm (modified, +86/-12)
• test/test_mps.py (modified, +38/-0)