Fix SDPA implementation in Qwen2-VL (issues with torch==2.6.0)
Author: ManuelFay
Pull request: yes

# What does this PR do?

**Fix SDPA call in Qwen2VL for compatibility with PyTorch 2.6.0 on MPS**

### Overview

This PR fixes an `IndexError` that occurs when running the `scaled_dot_product_attention` call in the Qwen2VL model implementation on **MPS devices** with **PyTorch 2.6.0**.

### Problem

Starting in **PyTorch 2.6.0**, `torch.nn.functional.scaled_dot_product_attention` appears to have stricter input shape validation, especially on **MPS backends**. When passing 3D tensors, the following error occurs:

```
IndexError: Dimension out of range (expected to be in range of [-3, 2], but got 3)
```

The error is specific to the use of SDPA in **Qwen2VL**, and seems to stem from the function expecting inputs of shape `(batch, num_heads, seq_len, head_dim)` but not properly handling 3D inputs without inferred layout or metadata.

From the documentation, this function expects 4D tensors. However, this was **not an issue in PyTorch 2.5.1**, where the same code ran without error. 

I ran a few ablations and realized 3D tensors didn't run either on 2.5.1 except when at least one of q,k or v came out of an attention module, in which case, I guess some extra info is added to the tensor, and the function is able to cast the operation correctly. 


### Fix

I'm opening this more as a discussion, because I am not completely sure of what's going on behind the scenes. 
I did succesfully patch the SDPA call as follows, and it runs faster in 2.6.0 than in 2.5.1 without this patch, for the same results.

```python
attn_output = F.scaled_dot_product_attention(
    q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0),
    attention_mask, dropout_p=0.0
)
attn_output = attn_output.squeeze(0).transpose(0, 1)
```

This explicitly adds an extra dimension and restores the original shape after the operation. It ensures the attention call runs consistently across PyTorch versions and on MPS.

### Notes

- Fix is limited in scope to Qwen2VL's attention implementation (SDPA).
- This change maintains compatibility with PyTorch 2.5.1 while resolving the issue in 2.6.0.
- Reproducible on MPS (Apple Silicon) devices.


Fixes # (issue)
[#36413](https://github.com/huggingface/transformers/issues/36413)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?
@qubvel @ArthurZucker @zucchini-nlp 


Specs: mps, any recent transformers version, torch==2.6.0

Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page).

The repository consistency CI fail I don't quite get...
Modular imports the SDPA attention in Qwen2_5 from Qwen2 which is the file I just changed so I guess the update should propagate automatically. Not sure what more I need to do.

Edit: fixed by modifying the code that would have been modified by the CI. Note, the comment at the top says not to do that explicitely, so maybe it might be better to modify it and explain to do it but just as a failsafe ?

Great ! This PR does update Qwen2-VL (and 2.5 VL), maybe you were thinking of other models ?