Flash Attention fails with non aligned position_ids
Author: alessiodevoto
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.54.1
- Platform: Linux-6.1.123+-x86_64-with-glibc2.39
- Python version: 3.12.3
- Huggingface_hub version: 0.34.3
- Safetensors version: 0.5.3
- 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?: <fill in>
- GPU type: NVIDIA H100 80GB HBM3

### Who can help?

@winglian 

### Information

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

### Tasks

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

### Reproduction

### Overview
Hi! In the latest release (v4.54.1), there was a change in [how the max_length is computed when using flash attention](https://github.com/huggingface/transformers/blob/cb289ad243a5aa4c76719f4df1d4c07171e338da/src/transformers/modeling_flash_attention_utils.py#L243). This raises an error if we forward a sequence where no position_ids == 0. 

### Code to reproduce
This is minimal code to reproduce.  This **works fine** when using` attn_implementation="eager"`, but fails with Flash Attention.

```
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Use Llama-3.2-1B-Instruct for testing, but it applies to all models
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="flash_attention_2")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")

# create a dummy input
input_ids = tokenizer.encode("All good here how are you?", return_tensors="pt").to(model.device)
# the position ids start from 1 instead of 0
position_ids = torch.arange(1, input_ids.shape[1]+1).unsqueeze(0).to(model.device)

output = model(input_ids, position_ids=position_ids) # Fails
# RuntimeError: max(): Expected reduction dim to be specified for input.numel() == 0. Specify the reduction dim with the 'dim' argument.
```

I think the problem is that when we call `diff()` in [this line](https://github.com/huggingface/transformers/blob/cb289ad243a5aa4c76719f4df1d4c07171e338da/src/transformers/modeling_flash_attention_utils.py#L243), we get an empty tensor (as no position_ids equal 0) and then `max()` fails.

### Why it matters 
Passing position_ids without zero elements makes sense in all those cases where you have a KV Cache and want to generate starting from there. We maintain [NVIDIA/KVPress](https://github.com/NVIDIA/kvpress),  a library for KV Cache compression,  and rely on this for our pipeline.

### Expected behavior

No errors.

Comments:
Hm, I think we don't support using packed FA2 when generating on purpose. Otherwise we would need to handle arbitrary position ids as new tokens are generated

Hi @zucchini-nlp, thanks for your quick response!

This bug is primarily concerned with using `past_key_values` in the model's forward pass; e.g. 

```
context_length = past_key_values.get_seq_length()
position_ids = torch.arange(context_length, context_length + input_ids.shape[1], device=model.device).unsqueeze(0)

outputs = model(
    input_ids=input_ids.to(model.device),
    past_key_values=past_key_values,
    position_ids=position_ids,
)
```
will break with FA2 because `position_ids` don't contain 0 as an element. `position_ids` is still a sequence of consecutive ints.



As a reference, `model.generate` updated the code as follows
- compute `cu_seq_lens_q/cu_seq_lens_k` ([here](https://github.com/huggingface/transformers/blame/main/src/transformers/generation/utils.py#L692)) and set them as additional input kwargs.
- As these values are now given, the problematic function is [not called](https://github.com/huggingface/transformers/blob/cb289ad243a5aa4c76719f4df1d4c07171e338da/src/transformers/modeling_flash_attention_utils.py#L422)


**Note that both functions compute `cu_seq_lens_q/cu_seq_lens_k` from `position_ids` and the first one is working with `position_ids` not containing 0.**

_____________
In our opinion, it should still be possible to use aforementioned code snippet without being concerned with the internals of the attention implementation (computing `cu_seq_lens_q/cu_seq_lens_k`). 

It would e.g. be possible to 
- reuse the same code that is used in `.generate` to compute `cu_seq_lens_q/cu_seq_lens_k` in `modeling_flash_attention_utils.py`
- Remove [code](https://github.com/huggingface/transformers/blame/main/src/transformers/generation/utils.py#L692) in `.generate` and move (working) code over to `modeling_flash_attention_utils.py`.

Ahhh, I see, might have missed last updated in FA2. 

> t would e.g. be possible to
> reuse the same code that is used in .generate to compute cu_seq_lens_q/cu_seq_lens_k in modeling_flash_attention_utils.py
> Remove [code](https://github.com/huggingface/transformers/blame/main/src/transformers/generation/utils.py#L692) in .generate and move (working) code over to modeling_flash_attention_utils.py.

If we just replace code for inferring `cu_seq_len` with the one from generation block, it will break FA2 for users with use packed sequences. In that case `position_ids[:, -1].cumsum(0)` isn't representative of actual `cu_seq_len`. IIUC in your use-case there will be cache re-use and thus `query_length != key_length`. I think we can move the mentioned code under FA2 for cases `if query_length != key_length`

@ArthurZucker WDYT?

yes, happy to have a patch. There are many casses and we were a bit too convinced that pos ids are usually only passed by generate I think

Oke, i will submit a PR today. Indeed FA2 has a lot of edge cases 🙈 

Sorry for posting under a closed issue but I think this might be the best place to ask.

I'm trying to use a model with packed sequences and a shared kv cache among the sequences.
I managed to do that using 4D mask with eager attention. However, it seems like the flash attention code does not allow attention_mask to be used with packed sequences. It assumes that there are pad tokens in the sequence and tries to unpad the sequence, leading to a CUDA error. So I switched to using position_ids instead. But the outputs are not the same as using 4D mask.

Versions
- transformers                             4.55.4
- torch                                    2.6.0


Here's a script to reproduce
```python
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
import torch


def position_ids_to_seq_ids(position_ids):
    """
    Convert position_ids to sequence IDs by detecting where positions reset.

    Args:
        position_ids: Tensor of shape [1, seq_len]

    Returns:
        seq_ids: Tensor of same shape with sequence IDs (0, 1, 2, ...)
    """
    if position_ids.dim() == 2:
        position_ids = position_ids.squeeze(0)

    # Detect where position resets (new sequence starts)
    # A reset occurs when position[i] <= position[i-1]
    diffs = torch.diff(
        position_ids, prepend=torch.tensor([-1], device=position_ids.device)
    )
    resets = (diffs <= 0).long()

    # Cumsum gives us the sequence ID
    seq_ids = torch.cumsum(resets, dim=0)
    return seq_ids


if __name__ == "__main__":
    device = "cuda"
    model_name = "google/gemma-2-2b-it"
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        # attn_implementation="flash_attention_2",
        device_map=device,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    texts = [
        "Hello, my dog is cute. He is very small.",
        "Hello, my dog is cute. She loves playing fetch with her toys.",  # Different length
    ]

    # Tokenize without padding - returns list of tensors
    input_ids_list = [
        tokenizer(t, return_tensors="pt").input_ids.squeeze(0).to(model.device)
        for t in texts
    ]
    seq_lengths = [ids.shape[0] for ids in input_ids_list]

    cache_size = 5
    print(f"Sequence lengths: {seq_lengths}")

    # Ground truth: run each sequence separately
    gt_logits_list = []
    for ids in input_ids_list:
        gt_out = model(input_ids=ids.unsqueeze(0))
        gt_logits_list.append(gt_out.logits)

    # Create cache from shared prefix
    outputs = model(
        input_ids=input_ids_list[0][:cache_size].unsqueeze(0),
        use_cache=True,
        past_key_values=DynamicCache(),
    )
    past_kv = outputs.past_key_values

    # Pack remaining tokens (after cache) from all sequences
    new_input_ids = torch.cat([ids[cache_size:] for ids in input_ids_list]).unsqueeze(0)
    n_new_tokens = new_input_ids.shape[1]

    # Build position_ids for packed sequences
    position_ids = torch.cat(
        [
            torch.arange(seq_len - cache_size, device=model.device)
            for seq_len in seq_lengths
        ]
    ).unsqueeze(0)

    seq_ids = position_ids_to_seq_ids(position_ids)
    print(f"Position IDs: {position_ids}")
    print(f"Sequence IDs: {seq_ids}")

    kv_seq_ids = torch.cat(
        [torch.full((cache_size,), -1, device=model.device), seq_ids]
    )

    def mask_func(_, _h, q_idx, kv_idx):
        # this assumes monotonic position_ids
        return (kv_seq_ids[kv_idx] == -1) | (
            (seq_ids[q_idx] == kv_seq_ids[kv_idx]) & (q_idx + cache_size >= kv_idx)
        )

    attn_mask_4d = torch.nn.attention.flex_attention.create_mask(
        mask_func,
        B=1,
        H=1,
        Q_LEN=n_new_tokens,
        KV_LEN=n_new_tokens + cache_size,
        device=model.device,
    )
    attn_mask_4d = torch.where(
        attn_mask_4d,
        torch.tensor(0.0, dtype=model.dtype, device=model.device),
        torch.tensor(float("-inf"), dtype=model.dtype, device=model.device),
    )

    outputs = model(
        input_ids=new_input_ids,
        attention_mask=attn_mask_4d,  # outputs match with 4d attention mask w/ eager but does not support FA2
        # position_ids=position_ids + cache_size, # outputs dont match with position_ids
        past_key_values=past_kv,
        use_cache=True,
        cache_position=position_ids.squeeze(0) + cache_size,
    )

    # Compare with ground truth
    print("\n=== Comparison ===")
    offset = 0
    for seq_idx, seq_len in enumerate(seq_lengths):
        remaining_len = seq_len - cache_size
        packed_logits = outputs.logits[0, offset : offset + remaining_len]
        gt_logits = gt_logits_list[seq_idx][0, cache_size:]

        print(f"\nSequence {seq_idx}:")
        print(f"{packed_logits=}")
        print(f"{gt_logits=}")
        print(
            f"  Last token match: {torch.allclose(packed_logits[-1], gt_logits[-1], atol=1e-4)}"
        )
        print(f"  Max diff: {(packed_logits - gt_logits).abs().max().item()}")

        offset += remaining_len
```

@51616 FA2 with packed sequences does not expect any attention mask, and you;re doing right by passing packed position ids. I can see in your script that the model with FA2 gets a 4D attention mask, or is that supposed to be a comparison between FA2 model and flex attention model

In any case, can you open a separate PR ?