[MPS] Silent correctness issue in bidirectional attention
Author: hvaara
Pull request: no
Labels: name: WIP

name: bug

### System Info

A bug in PyTorch for the MPS backend (pytorch/pytorch#174861) results in a silent correctness issue in bidirectional attention under certain conditions:

- dtype != `torch.float` (eg. `float16` or `bfloat16`)
- non-masked or boolean mask
- non-causal
- query sequence length <= 8
- query sequence length <= key sequence length
- query head dim = value head dim
- query head dim in {64, 96, 128}
- key sequence length >= 1024 OR (number of key heads < number of query heads AND key sequence length >= 4096)

This is the case for several models in transformers. For example GLM-Image, but I've seen it in other vision and TTS models too, but I can't remember which ones top of mind.

The issue has been in PyTorch since `v2.8.0` released 2025-08-06. I've submitted a fix to PyTorch (pytorch/pytorch#174945), which is on track to make it in the `v2.11.0` release of PyTorch.

I recommend that we add a workaround for versions between `v2.7.1` and `v2.11.0` exclusive. For example by upcasting the qkv tensors to fp32 when calling sdpa under the conditions outlined above. The source for these conditions can be found in the PyTorch source. The issue is in `sdpa_vector_2pass_mps`, so conditions that takes the other code path through `sdpa_general_mps` would lead to correct output.

https://github.com/pytorch/pytorch/blob/07180141f03510ab0b53dd0a544b71a5d5bdba93/aten/src/ATen/native/mps/operations/Attention.mm#L488

I'm happy to provide a PR once we're aligned on mitigation steps.

### Who can help?

@vasqu @ArthurZucker @CyrilVallez @ivarflakstad 

### Information

- [x] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

#### MRE

```python
import torch
from diffusers.pipelines.glm_image import GlmImagePipeline

pipe = GlmImagePipeline.from_pretrained("zai-org/GLM-Image", torch_dtype=torch.bfloat16, device_map="mps")
prompt = "A beautifully designed modern food magazine style dessert recipe illustration, themed around a raspberry mousse cake. The overall layout is clean and bright, divided into four main areas: the top left features a bold black title 'Raspberry Mousse Cake Recipe Guide', with a soft-lit close-up photo of the finished cake on the right, showcasing a light pink cake adorned with fresh raspberries and mint leaves; the bottom left contains an ingredient list section, titled 'Ingredients' in a simple font, listing 'Flour 150g', 'Eggs 3', 'Sugar 120g', 'Raspberry puree 200g', 'Gelatin sheets 10g', 'Whipping cream 300ml', and 'Fresh raspberries', each accompanied by minimalist line icons (like a flour bag, eggs, sugar jar, etc.); the bottom right displays four equally sized step boxes, each containing high-definition macro photos and corresponding instructions, arranged from top to bottom as follows: Step 1 shows a whisk whipping white foam (with the instruction 'Whip egg whites to stiff peaks'), Step 2 shows a red-and-white mixture being folded with a spatula (with the instruction 'Gently fold in the puree and batter'), Step 3 shows pink liquid being poured into a round mold (with the instruction 'Pour into mold and chill for 4 hours'), Step 4 shows the finished cake decorated with raspberries and mint leaves (with the instruction 'Decorate with raspberries and mint'); a light brown information bar runs along the bottom edge, with icons on the left representing 'Preparation time: 30 minutes', 'Cooking time: 20 minutes', and 'Servings: 8'. The overall color scheme is dominated by creamy white and light pink, with a subtle paper texture in the background, featuring compact and orderly text and image layout with clear information hierarchy."
image = pipe(
    prompt=prompt,
    height=32 * 32,
    width=36 * 32,
    num_inference_steps=50,
    guidance_scale=1.5,
    generator=torch.Generator(device="mps").manual_seed(42),
).images[0]

image.save("output_t2i.png")
```

With PyTorch `v2.10.0` results in the following Python stacktrace

```
File ~/dev/huggingface/transformers/src/transformers/generation/utils.py:2884, in GenerationMixin._sample(self, input_ids, logits_processor, stopping_criteria, generation_config, synced_gpus, streamer, **model_kwargs)
   2882     probs = nn.functional.softmax(next_token_scores, dim=-1)
   2883     # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution
-> 2884     next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
   2885 else:
   2886     next_tokens = torch.argmax(next_token_scores, dim=-1)

RuntimeError: probability tensor contains either `inf`, `nan` or element < 0
```

### Expected behavior

I expected there to be no error and that the output image is generated successfully.

Comments:
Iiuc, then all the listed conditions must apply. Would this mean that if we were to never skip the mask creation, that we wouldn't encounter this issue?

Casting from half to full (and back) sounds quite expensive to me so I'd like to avoid it if possible. And MPS backend does not benefit from anything special between fp32 and half I assume or using no mask?

> Iiuc, then all the listed conditions must apply. 

Yes

> Would this mean that if we were to never skip the mask creation, that we wouldn't encounter this issue?

Yes, as long as it's not a bool mask,

https://github.com/pytorch/pytorch/blob/07180141f03510ab0b53dd0a544b71a5d5bdba93/aten/src/ATen/native/mps/operations/Attention.mm#L461-L462

would become false, and then you end up going down `sdpa_general_mps`, which is not affected by the correctness issue (although much slower).

> Casting from half to full (and back) sounds quite expensive to me so I'd like to avoid it if possible.

Attention is quite expensive in general on MPS. `sdpa_vector_2pass_mps` is about 40% faster than `sdpa_general_mps`, based on benchmarking stats from the PR that introduced the Metal kernels. I'm unsure what the perf impact would be of running this in fp32 vs using `sdpa_general_mps`. It's possible, but I'm not sure, it would actually be faster to upcast to fp32 and go down `sdpa_vector_2pass_mps`. Perhaps a good candidate for some quick benchmarking.

> And MPS backend does not benefit from anything special between fp32 and half I assume or using no mask?

Non-masked or bool mask is a prereq to use the Metal kernels (which are a lot faster as I mentioned above). Both Metal and MPS kernels support fp32/fp16/bf16.

Any memory impacts we could observe on the upcast. In general, I'm up for workarounds as long as we show (benchmark) which one is the better one, speed and memory-wise:
- Non-boolean mask forcing 
- Fp32 upcasting

At least these seem to be the current viable options. Fp32 upcasting is potentially expensive on the memory side, should have clarified a bit more.

I think forcing a mask is much simpler on our side. Upcasting to fp32 would require manual tweaks in every modeling code, whereas forcing mask creation can be simply done in the masking API (though it adds yet another code paths due to broen 3rd party unfortunately 🥲)

I think it could live in the sdpa integration for upcasting but we definitely have less coverage for the attention interface than the masks 👀 

Forcing a mask seems like the better option to me too, especially considering tight memory constraints on consumer hardware. Adding a warning might also be beneficial so those affected have a nag to raise awareness of the issue and resolution steps (update PyTorch).

Do you want to open a PR @hvaara then? You essentially need to fixup 2 things for this
- https://github.com/huggingface/transformers/blob/41888df59eed5379467262000a4f68493d3fbced/src/transformers/masking_utils.py#L994 --> force eager interface
- https://github.com/huggingface/transformers/blob/41888df59eed5379467262000a4f68493d3fbced/src/transformers/masking_utils.py#L997 --> force `allow_xxx_skip=False`

I just picked the bidirectional mask but would need to happen on every top level mask fn. Hard for me to do a PR because I don't use a Mac 😓 

Sure, I can take a look. I'll loop you in on the review if you don't mind 😃 

Sounds good! 

Looks like at least GLM-Image doesn't use the masking API. Would it be better to add the workaround directly in the attention interface (in `src/transformers/integrations/sdpa_attention.py`)? What's the preferred approach?

> Looks like at least GLM-Image doesn't use the masking API

Argh, yea the vision portions don't use a mask (so pretty much any of the VLMs of qwen, zai, and baidu). Could we not simply create a mask with `torch.zeros` in that case? To cover as much as we can, we need to 
1. Force mask creations where we can, i.e. in the mask API
2. In the sdpa integration, create a zeros mask if mps x affected torch versions x no mask

It's a bit of a rubber band tbh, but we don't have much choice. Case 1 should cover almost anything and case 2 for the special VLM cases. But @Cyrilvallez wdyt? Maybe you have a different idea

If we need to do 2. anyway, I'd say let's only do 2. then no? That way we keep it confined to the sdpa integration. Even a single function for both linkes issues with mps

Since it was not clear in #44359 and #44591: I am already working on this issue.

> If we need to do 2. anyway, I'd say let's only do 2. then no? That way we keep it confined to the sdpa integration. Even a single function for both linkes issues with mps

Yep, it was before we had the second issue hit us 😢 

> Since it was not clear in https://github.com/huggingface/transformers/pull/44359 and https://github.com/huggingface/transformers/pull/44591: I am already working on this issue.

Yep don't worry, we will keep an eye out for your PR. A lot of coding agents just randomly pick up issues/prs 😅 

Yeah no worries @hvaara, we will always prioritize the initial issue raiser if they are willing to contribute, then us, then all the random agent prs 😁

Yes we'll do provide a fix soon

On Thu, Mar 12, 2026, 10:42 AM Cyril Vallez ***@***.***>
wrote:

> *Cyrilvallez* left a comment (huggingface/transformers#44247)
> <https://github.com/huggingface/transformers/issues/44247#issuecomment-4045704951>
>
> Yeah no worries @hvaara <https://github.com/hvaara>, we will always
> prioritize the initial issue raiser if they are willing to contribute, then
> us, then all the random agent prs 😁
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/44247#issuecomment-4045704951>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/B7HTEELZKMHFMT6NCBSO7PD4QKIBJAVCNFSM6AAAAACV5KSGA6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHM2DANBVG4YDIOJVGE>
> .
> You are receiving this because you are subscribed to this thread.Message
> ID: ***@***.***>
>


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

I'll get around to finish it this week.

I have prepared a design + implementation for the MPS SDPA workaround following the discussion above (zeros mask in `src/transformers/integrations/sdpa_attention.py`, version-gated for PyTorch >= 2.8, < 2.11 for the bidirectional correctness bug, plus a separate fix for value head dim mismatch < 2.12). No changes to `masking_utils.py` per the maintainer decision.

@hvaara — are you still actively working on this? I don't want to duplicate effort. Happy to align with your approach or hand off mine as a reference. My implementation follows the approach discussed with @Cyrilvallez and @vasqu above: single function in sdpa_attention.py, device-gated on MPS only, no CPU/CUDA regression.

The implementation is ready for review. Per AGENTS.md policy, a human submitter will review every changed line before any PR is opened.