🐛 Describe the bug

The meta kernel for aten._scaled_dot_product_attention_math_for_mps branches on query_seq_len <= 8, which creates a data-dependent guard that prevents torch.export(strict=True) from supporting dynamic sequence lengths.

Introduced in #176843 (eccc7be).

Minimal reproduction

import torch
from torch.export import Dim, export

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("k", torch.zeros(1, 4, 64, 128))
        self.register_buffer("v", torch.zeros(1, 4, 64, 128))

    def forward(self, q, mask):
        out, _ = torch.ops.aten._scaled_dot_product_attention_math_for_mps(
            q, self.k, self.v, mask, 0.0, False, None
        )
        return out

model = M().to(torch.bfloat16).eval()
seq = Dim("seq", min=1, max=64)

q = torch.randn(1, 4, 4, 128, dtype=torch.bfloat16)
mask = torch.zeros(4, 64, dtype=torch.bfloat16)

# Fails: guard on seq_len <= 8 specializes the trace
ep = export(model, (q, mask), dynamic_shapes={"q": {2: seq}, "mask": {0: seq}}, strict=True)

Error:

torch._dynamo.exc.UserError: Constraints violated (seq)!
  - Not all values of seq = L['q'].size()[2] in the specified range seq <= 64
    satisfy the generated guard 2 <= L['q'].size()[2] and L['q'].size()[2] <= 8
Suggested fixes:
  seq = Dim('seq', max=8)

Expected behavior

Export succeeds with seq as a fully dynamic dimension covering [1, 64].

Root cause

In torch/_meta_registrations.py, meta__scaled_dot_product_attention_math_for_mps has:

query_seq_len = q_.size(2)
supports_sdpa_vector = (
    (query_seq_len <= 8)  # <-- data-dependent guard on dynamic dim
    and ...
)

The three code paths (sdpa_general_mps, sdpa_vector_2pass_mps, sdpa_vector_fast_mps) return different output shapes. During tracing, the branch on query_seq_len specializes, creating guards that reject all other seq_len values.

Proposed fix

Always return sdpa_general_mps shapes from the meta kernel — (B, H, q_size, value_head_size) for output, (B, H, q_size, max_seq_length) for attention. The runtime kernel still dispatches to the optimal implementation; the meta kernel only needs correct shapes for tracing.

This is safe for the first return value (the actual output) because the vector paths require query_head_dim == value_head_dim, so value_head_size is always correct. The second return value is never consumed — all C++ callsites use std::get<0>(...).

TestSDPAMetaDispatchMode compares both return values against the real kernel and will need updating for the 2pass case: the 2pass runtime path returns a 5D intermediate (B,H,q,32,D) vs the meta kernel's 4D (B,H,q,max_seq_length). Since the second return value is universally discarded and already inconsistent across paths (general→4D attn weights, 2pass→5D intermediate, full_attention→duplicate of output), the test should only compare the first return value.

Versions

Regression bisected to eccc7be10726ebabb8f064e8e9990984a83226ec (#176843).

Tested on torch 2.12.0.dev20260312.

cc @chauhang @penguinwu @avikchaudhuri @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @malfet @hvaara