Repository: transformers
Issue: Transformers is incompatible with flash attention version 2.7.3
Labels: bug

### System Info


Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.48.0
- Platform: Linux-5.13.0-35-generic-x86_64-with-glibc2.31
- Python version: 3.10.11
- Huggingface_hub version: 0.27.1
- Safetensors version: 0.5.2
- Accelerate version: 0.29.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.3.1+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: no
- Using GPU in script?: yes
- GPU type: NVIDIA A100-SXM4-80GB


Also
```
python -c 'import flash_attn; print(flash_attn.__version__)'
2.7.3
```

### Who can help?

_No response_

### 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

The path `/mnt/models/google/gemma2-2b` points to [Google Gemma2 2b model](https://huggingface.co/google/gemma-2-2b) from Huggingface.

```python
import torch

from transformers import AutoTokenizer
from transformers.tokenization_utils import PaddingStrategy
from transformers.models.gemma2 import Gemma2ForCausalLM


def main():
    model_path = '/mnt/models/google/gemma2-2b'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = Gemma2ForCausalLM.from_pretrained(
        model_path,
        attn_implementation='flash_attention_2',
        torch_dtype=torch.bfloat16,
    ).to('cuda')

    texts = ['have a nice day', 'you too']
    input_ids = tokenizer(
        texts,
        return_tensors='pt',
        padding_side='left',
        truncation=True,
        max_length=256,
        padding=PaddingStrategy.MAX_LENGTH,
    ).input_ids.to(model.device)

    model.generate(input_ids, max_new_tokens=2)


if __name__ == '__main__':
    main()
```

Running this script gives 
```python
python reproduce.py
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
Loading checkpoint shards: 100%|███████████████████████████████| 3/3 [00:01<00:00,  2.68it/s]
The 'batch_size' attribute of HybridCache is deprecated and will be removed in v4.49. Use the more precisely named 'self.max_batch_size' attribute instead.
Traceback (most recent call last):
  File "/my_code/turbo-alignment/reproduce.py", line 31, in <module>
    main()
  File "/my_code/turbo-alignment/reproduce.py", line 27, in main
    model.generate(input_ids, max_new_tokens=2)
  File "/usr/local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/transformers/generation/utils.py", line 2255, in generate
    result = self._sample(
  File "/usr/local/lib/python3.10/site-packages/transformers/generation/utils.py", line 3254, in _sample
    outputs = self(**model_inputs, return_dict=True)
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 842, in forward
    outputs = self.model(
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 629, in forward
    layer_outputs = decoder_layer(
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 299, in forward
    hidden_states, self_attn_weights = self.self_attn(
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 236, in forward
    attn_output, attn_weights = attention_interface(
  File "/usr/local/lib/python3.10/site-packages/transformers/integrations/flash_attention.py", line 50, in flash_attention_forward
    attn_output = _flash_attention_forward(
  File "/usr/local/lib/python3.10/site-packages/transformers/modeling_flash_attention_utils.py", line 305, in _flash_attention_forward
    query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = _upad_input(
  File "/usr/local/lib/python3.10/site-packages/transformers/modeling_flash_attention_utils.py", line 124, in _upad_input
    query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
ValueError: too many values to unpack (expected 4)
```

### Expected behavior

Script runs smoothly