[MPS] Upstream correctness issue in attention when value head dim differs from query

快速结论：该问题出现在 Apple MPS 后端上，当 value head dim 与 query head dim 不一致时，PyTorch 的 scaled_dot_product_attention 会返回未初始化的内存以及错误形状，从而触发 [MPS] Upstream correctness issue in attention when value head dim differs from query 相关症状。优先确认 PyTorch 版本，并升级到 torch>=2.12。

适用环境：Issue 中确认的环境为 Apple MPS 后端；未在正文中明确列出具体 macOS、Python、CUDA 或显卡型号，因此不补充推测版本。

最快修复方案：升级 PyTorch 到 2.12 或更高版本。上游修复 pytorch/pytorch#176843 已随 torch 2.12.0 发布，2.13 也已可用，因此维护者决定不再在 sdpa_attention_forward 中为旧版 torch 保留 MPS 版本相关的工作区。

注意事项：Issue 中曾讨论过对 value tensor 做 padding、调整输出形状或 strides、以及在 sdpa 集成中增加 MPS 专用修复或装饰器，但这些方案最终未被采用，因为升级 torch 即可解决。若无法升级到 torch>=2.12，则没有已验证的一步修复方案；相关变通方法均属于“可优先尝试”，且未在 Issue 中获得最终验证。

问题场景

在 Apple MPS 后端上运行 PyTorch 的 F.scaled_dot_product_attention，且 value 的 head dim 与 query 的 head dim 不同。Issue 中的复现示例使用 q 形状为 (1, 1, 8, 4)、v 形状为 (1, 1, 8, 2)，期望输出最后一维为 2，但 MPS 输出最后一维为 4，且两次运行结果不一致。这个问题可能在 Transformers 的 SDPA attention 集成中影响使用 MLA 的模型，例如 DeepSeek-V2/V3 等 query/key head dim 与 value head dim 不同的模型。

报错原文

[MPS] Upstream correctness issue in attention when value head dim differs from query

y_mps1-y_mps2 = tensor([[[[ 0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.0000],
          [-0.0957, -0.1705, -0.1087, -0.1648],
          [-0.1155, -0.1248, -0.0946, -0.1255],
          [-0.1089, -0.1662, -0.1261, -0.1559],
          [-0.0904, -0.1262, -0.0927, -0.1335],
          [-0.1080, -0.1644, -0.1226, -0.1611],
          [-0.0849, -0.1324, -0.0984, -0.1281]]]], device='mps:0')
y_cpu1-y_cpu2 = tensor([[[[0., 0.],
          [0., 0.],
          [0., 0.],
          [0., 0.],
          [0., 0.],
          [0., 0.],
          [0., 0.],
          [0., 0.]]]])
y_mps1.shape = torch.Size([1, 1, 8, 4])
y_cpu1.shape = torch.Size([1, 1, 8, 2])

原因分析

根本原因是 PyTorch MPS 后端在 value head dim 与 query head dim 不同时，scaled_dot_product_attention 返回了未初始化的内存，并且输出形状错误。Issue 指出该问题已在上游 PyTorch 修复，对应 PR 为 pytorch/pytorch#176843，但该修复直到 PyTorch 2.12 才发布。因此，在旧版 PyTorch 上，MPS 后端可能返回错误的输出宽度，导致数值不一致，或在使用 MLA 的模型中让后续 output projection 因宽度不匹配而报错。

环境排查

确认当前 PyTorch 版本是否低于 2.12；这是最关键的一项。
确认是否使用 Apple MPS 后端，例如 device="mps" 或模型被移到 MPS。
确认出现问题的 attention 调用中，value 的最后一维是否与 query 的最后一维不同。
确认是否使用 Transformers 的 SDPA attention 集成，尤其是 DeepSeek-V2/V3、MLA 等 query/key head dim 与 value head dim 不同的模型。
Issue 未提供 macOS、Python、CUDA 或具体显卡型号信息，这些项目无需作为必须排查项。

解决步骤

检查当前 torch 版本，确认是否低于 2.12。
将 PyTorch 升级到 2.12 或更高版本，例如 torch 2.12.0 或 2.13。
升级后，用 Issue 中的最小复现脚本在 MPS 上再次运行，确认 y_mps1-y_mps2 全为 0，且 y_mps1.shape 的最后一维等于 value 的 head dim，而不是 query 的 head dim。
如果因环境限制无法升级 torch，可优先尝试 Issue 中讨论过的变通方法：对 value tensor 做 padding 后再调用 attention，例如将 v 从最后一维 Ev padding 到 E，计算后再截回 y[..., :Ev]。但 Issue 最终未采用该方案，因此只能作为临时尝试，且需自行验证正确性。
不建议在 Transformers 中为旧版 torch 添加 MPS 专用 workaround；维护者已明确倾向直接要求升级 torch。

验证方法

重新运行 Issue 中的 MPS 复现脚本，确认两次 MPS 调用的差值为全 0，且 MPS 输出形状与 CPU 输出形状一致，即最后一维为 2，而不是 4。对于 DeepSeek-V2/V3 等模型，还应确认 attention 输出进入 output projection 时宽度正确，不再因形状不匹配而报错。