Repository: huggingface/transformers
Pull request: Fix SDPA implementation in Qwen2-VL (issues with torch==2.6.0)

# 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

Changed files:
filename: src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py
patch: @@ -316,8 +316,10 @@ def forward(
         q = q.transpose(0, 1)
         k = k.transpose(0, 1)
         v = v.transpose(0, 1)
-        attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
-        attn_output = attn_output.transpose(0, 1)
+        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)
         attn_output = attn_output.reshape(seq_length, -1)
         attn_output = self.proj(attn_output)
         return attn_output

filename: src/transformers/models/qwen2_vl/modeling_qwen2_vl.py
patch: @@ -416,8 +416,10 @@ def forward(
         q = q.transpose(0, 1)
         k = k.transpose(0, 1)
         v = v.transpose(0, 1)
-        attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
-        attn_output = attn_output.transpose(0, 1)
+        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)
         attn_output = attn_output.reshape(seq_length, -1)
         attn_output = self.proj(attn_output)
         return attn_output