`RuntimeError` in `Siglip2Model` Attention with NaFlex when `actual_patches != max_num_patches`
Author: Enferlain
Pull request: no
Labels: name: bug

### System Info

*   Transformers version: 4.52.0.dev0 but same occurs on 4.51.3
*   Platform: Win 11
*   Python version: 3.11.5
*   PyTorch version: 2.6.0+cu126
*   Using GPU? Yes/No: Yes
*   GPU type: NVIDIA RTX 3090
*   CUDA/cuDNN version: cuda 12.6, cudnn 9.8
*   Using `accelerate`? No

### Issue description

When using `Siglip2Model` with the `Siglip2ImageProcessorFast` for NaFlex processing (`google/siglip2-so400m-patch16-naflex`), a `RuntimeError` occurs inside `torch.nn.functional.scaled_dot_product_attention` if the actual number of patches derived from the image and `spatial_shapes` is *different* from the `max_num_patches` value used during processing.

The error seems to stem from an inconsistency between the tensors passed to the model's attention layers:
- The `hidden_states` input to the encoder layers correctly reflects the *actual* number of patches (`L_actual`, derived from `spatial_shapes`).
- However, the `attention_mask` passed alongside it is derived from the mask returned by `Siglip2ImageProcessorFast`, which appears to be padded/truncated to `max_num_patches` (`L_max`).
- This leads to a shape mismatch inside `scaled_dot_product_attention`, e.g., `Target sizes: [B, H, L_actual, L_actual]. Tensor sizes: [B, 1, L_mask, L_mask]`, where `L_actual != L_mask`.

I encountered it after updating either torch/cuda (e.g., 2.6.0+cu126 tested, don't recall what version I was on, maybe cu124) on an identical dataset with the same processing approach that worked fine previously. Seems like it might have been silently ignored or handled differently in older versions. The error occurs regardless of whether `max_num_patches` is set higher (e.g., 2048) or lower (e.g., 1024) than the actual number of patches generated for a given image.

**Traceback (Original error with `max_num_patches=1024`):**
```
Selected Model: google/siglip2-so400m-patch16-naflex (hf type)
Selected Preprocessing Mode: naflex_resize
(Using HF Processor logic with target max_num_patches=1024)
Embeddings will be saved in: data\siglip2_so400m_patch16_naflex_Naflex_Proc1024
Found source subfolders: ['0', '1']
Processing: ./anatomy\0 -> data\siglip2_so400m_patch16_naflex_Naflex_Proc1024\0
Folder '0':   0%|                                                                        | 0/1500 [00:00<?, ?image/s]D:\CityClassifiers\generate_embeddings.py:270: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
model_call_kwargs = {"pixel_values": pixel_values.to(device=device, dtype=dtype), "attention_mask": attention_mask.to(device=device), "spatial_shapes": torch.tensor(spatial_shapes, dtype=torch.long).to(device=device)}
DEBUG: After Embeddings - hidden_states shape: torch.Size([1, 12288, 1152])
DEBUG: After Embeddings - spatial_shapes: tensor([[ 96, 128]], device='cuda:0')
DEBUG: After Embeddings - attention_mask shape: torch.Size([1, 1024])
DEBUG: Start of Encoder - hidden_states shape: torch.Size([1, 12288, 1152])
DEBUG: Start of Encoder - attention_mask shape: torch.Size([1, 1, 1024, 1024])
DEBUG: Layer 0 - After LayerNorm1 - hidden_states shape: torch.Size([1, 12288, 1152])
DEBUG: Layer 0 - Before SelfAttn - hidden_states shape: torch.Size([1, 12288, 1152])
DEBUG: Layer 0 - Before SelfAttn - attention_mask shape: torch.Size([1, 1, 1024, 1024])
DEBUG: Attention - Query shape: torch.Size([1, 16, 12288, 72])
DEBUG: Attention - Key shape: torch.Size([1, 16, 12288, 72])
DEBUG: Attention - Value shape: torch.Size([1, 16, 12288, 72])
DEBUG: Attention - Mask shape fed to interface: torch.Size([1, 1, 1024, 1024])
Error during get_embedding (v4.3.0) for 000-01-noob9-0.658.png (Mode: naflex_resize, Type: hf):
Traceback (most recent call last):
File "D:\CityClassifiers\generate_embeddings.py", line 273, in get_embedding
if vision_model_component: emb = vision_model_component(**model_call_kwargs).pooler_output
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\torch\nn\modules\module.py", line 1739, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\torch\nn\modules\module.py", line 1750, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\transformers\utils\generic.py", line 965, in wrapper
output = func(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\transformers\models\siglip2\modeling_siglip2.py", line 573, in forward
encoder_outputs: BaseModelOutput = self.encoder(
^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\torch\nn\modules\module.py", line 1739, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\torch\nn\modules\module.py", line 1750, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\transformers\utils\generic.py", line 965, in wrapper
output = func(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\transformers\models\siglip2\modeling_siglip2.py", line 489, in forward
layer_outputs = encoder_layer(
^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\torch\nn\modules\module.py", line 1739, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\torch\nn\modules\module.py", line 1750, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\transformers\models\siglip2\modeling_siglip2.py", line 400, in forward
hidden_states, attn_weights = self.self_attn(
^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\torch\nn\modules\module.py", line 1739, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\torch\nn\modules\module.py", line 1750, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\transformers\models\siglip2\modeling_siglip2.py", line 331, in forward
attn_output, attn_weights = attention_interface(
^^^^^^^^^^^^^^^^^^^^
File "D:\CityClassifiers\venv\Lib\site-packages\transformers\integrations\sdpa_attention.py", line 54, in sdpa_attention_forward
attn_output = torch.nn.functional.scaled_dot_product_attention(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: The expanded size of the tensor (12288) must match the existing size (1024) at non-singleton dimension 3.  Target sizes: [1, 16, 12288, 12288].  Tensor sizes: [1, 1, 1024, 1024]
```

**Example after some attempts at workarounds, still `max_num_patches=1024`):**
```
RuntimeError: The expanded size of the tensor (1024) must match the existing size (1008) at non-singleton dimension 3.  Target sizes: [1, 16, 1024, 1024].  Tensor sizes: [1, 1, 1008, 1008]
Traceback (most recent call last):
  ... (Traceback points through Siglip2VisionModel -> Siglip2VisionTransformer -> Siglip2Encoder -> Siglip2EncoderLayer -> Siglip2Attention -> attention_interface -> scaled_dot_product_attention) ...
  File ".../transformers/models/siglip2/modeling_siglip2.py", line 331, in forward
    attn_output, attn_weights = attention_interface( ... attention_mask=mask_shape_1008x1008 ... )
  File ".../transformers/integrations/sdpa_attention.py", line 54, in sdpa_attention_forward
    attn_output = torch.nn.functional.scaled_dot_product_attention( # Expects mask compatible with Q/K/V derived from L=1024 hidden_state?
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
```
*(The exact target/tensor sizes in the error can vary slightly depending on where the check fails, but the core mismatch between `L_actual` and `L_mask` persists).*

### Information

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

### Tasks

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

### Reproduction

**To Reproduce**

Steps to reproduce the behavior:

1.  Use a recent PyTorch version (I tried 2.6.0+cu126) with CUDA enabled.
2.  Load the `google/siglip2-so400m-patch16-naflex` model and processor.
3.  Choose an input image that, when processed with a specific `max_num_patches` (e.g., 1024), results in `spatial_shapes` implying an `L_actual` different from `max_num_patches`. (Example image dimensions causing `L_actual=1008` when `max_num_patches=1024` might be needed, e.g., `336x768`).
4.  Call the processor: `inputs = processor(images=[img], return_tensors="pt", max_num_patches=1024)`
5.  Prepare inputs for the model, **using the `attention_mask` directly from the `inputs` dictionary**:
    ```python
    pixel_values = inputs["pixel_values"].to(device, dtype)
    attention_mask = inputs["pixel_attention_mask"].to(device) # Shape [1, 1024]
    spatial_shapes = inputs["spatial_shapes"].to(device) # Shape [[L_h, L_w]] implying L_actual=1008
    model_inputs = {"pixel_values": pixel_values, "attention_mask": attention_mask, "spatial_shapes": spatial_shapes.long()}
    ```
6.  Call the model's vision component: `outputs = model.vision_model(**model_inputs)`
7.  Observe the `RuntimeError` during the forward pass.

### Expected behavior

**Expected behavior**
The `Siglip2VisionTransformer` should ideally receive consistent inputs or internally handle the discrepancy between the sequence length implied by `spatial_shapes` (used for hidden states) and the length of the `attention_mask` provided by the processor. Either the processor should return an unpadded mask matching `spatial_shapes`, or the model should use `spatial_shapes` to correctly slice/interpret the padded mask.

**Temporary Workaround I tried:**
A workaround involves ignoring the `attention_mask` from the processor, calculating `L_actual` from `spatial_shapes`, creating a new `correct_attention_mask` of shape `[B, L_actual]`, un-padding `pixel_values` to length `L_actual`, and passing these consistent tensors to the model.

Comments:
maybe related https://github.com/pytorch/pytorch/issues/151558

Somehow my processor settings got set to do_resize = False, so that's why the size was the original img size