Wrong function signature introduced in PR #40029 breaking flash attention
Author: maxjeblick
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.56.0.dev0
- Platform: Linux-6.14.0-27-generic-x86_64-with-glibc2.39
- Python version: 3.12.3
- Huggingface_hub version: 0.34.3
- Safetensors version: 0.6.1
- Accelerate version: 1.9.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.7.1+cu126 (CUDA)
- 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?: <fill in>
- Using GPU in script?: Yes, with flash-attention
- GPU type: NVIDIA RTX A6000

### Who can help?

@zucchini-nlp 

### Information

- [ ] 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

`prepare_fa_kwargs_from_position_ids` returns `(cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k)` (a tuple of length 2, but `_prepare_from_posids` expects this to be `cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k`.

Bug may have been introduced in #40029

<img width="2179" height="1850" alt="Image" src="https://github.com/user-attachments/assets/6e464691-0de5-4ff2-a2a9-4595d0fa9513" />

### Expected behavior

No error thrown

Comments:
Will be fixed in #40002 thx for the quick find tho!

FYI:

After fixing `(cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) -> cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k` locally, the following code fails with `RuntimeError: CUDA error: an illegal memory access was encountered`

```
from transformers import AutoTokenizer, AutoModelForCausalLM


def generate_with_fa2():
    max_new_tokens = 20
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", torch_dtype="auto",
                                                 attn_implementation="flash_attention_2").eval().cuda()
    # model.set_attn_implementation("flash_attention_2")
    inputs_dict = tokenizer("Hello, how are you?", return_tensors="pt")
    generation_out = model.generate(inputs_dict["input_ids"].cuda(), do_sample=False, max_new_tokens=max_new_tokens)

    generated_tokens = generation_out.sequences[0]
    print(generated_tokens)
    original_length = inputs_dict["input_ids"].shape[1]

    for i in range(max_new_tokens):
        token_id = generated_tokens[original_length + i].cpu().item()
        token_text = tokenizer.decode(token_id)
        print(f"Token {i + 1}: '{token_text}'")

    print("\nGenerate method complete sequence:")
    print(f"'{tokenizer.decode(generated_tokens, skip_special_tokens=True)}'")


generate_with_fa2()

```


Slightly changed the code:
```python
from transformers import AutoTokenizer, AutoModelForCausalLM


def generate_with_fa2():
    max_new_tokens = 20
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", torch_dtype="auto",
                                                 attn_implementation="flash_attention_2").eval().cuda()
    # model.set_attn_implementation("flash_attention_2")
    inputs_dict = tokenizer("Hello, how are you?", return_tensors="pt")
    generation_out = model.generate(inputs_dict["input_ids"].cuda(), do_sample=False, max_new_tokens=max_new_tokens)

    generated_tokens = generation_out[0]
    print(generated_tokens)
    original_length = inputs_dict["input_ids"].shape[1]

    for i in range(max_new_tokens):
        token_id = generated_tokens[original_length + i].cpu().item()
        token_text = tokenizer.decode(token_id)
        print(f"Token {i + 1}: '{token_text}'")

    print("\nGenerate method complete sequence:")
    print(f"'{tokenizer.decode(generated_tokens, skip_special_tokens=True)}'")


generate_with_fa2()
```

Seems like the prepare function is not working as expected, I will revert it in the PR I linked - very much appreciated ❤️ 
(ie try to remove https://github.com/huggingface/transformers/blob/f6b6e177198264bf5d24ac4889cdc15ed9429fc9/src/transformers/generation/utils.py#L681-L690)

Can you try with #40002 again? Should work now

Works perfect, thanks again for your fix!

Closing this while pointing to #40002 but feel free to reopen if you discover this issue again (or other ones)!