Add packed tensor format support for flex/sdpa/eager through the mask!
Author: Cyrilvallez
Pull request: yes
Labels: name: for patch

# What does this PR do?

As per the title.

```python
import torch
from transformers import AutoModelForCausalLM
from transformers.masking_utils import create_causal_mask

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", torch_dtype=torch.float16)


batch_size = 1
sequence_length = 10
cache_position = torch.arange(sequence_length)
position_ids = torch.tensor([[0,1,2,3,0,1,0,1,2,3]])  # This corresponds to 3 packed sequences

attention_mask = create_causal_mask(
    config=model.config,
    # we only need batch size, seq_length and dtype here - we don't care about the values of the embeddings
    input_embeds=torch.empty((batch_size, sequence_length), dtype=model.dtype),
    attention_mask=None,
    cache_position=cache_position,
    past_key_values=None,
    position_ids=position_ids,
)
attention_mask

>>> tensor([[[[ True, False, False, False, False, False, False, False, False, False],
          [ True,  True, False, False, False, False, False, False, False, False],
          [ True,  True,  True, False, False, False, False, False, False, False],
          [ True,  True,  True,  True, False, False, False, False, False, False],
          [False, False, False, False,  True, False, False, False, False, False],
          [False, False, False, False,  True,  True, False, False, False, False],
          [False, False, False, False, False, False,  True, False, False, False],
          [False, False, False, False, False, False,  True,  True, False, False],
          [False, False, False, False, False, False,  True,  True,  True, False],
          [False, False, False, False, False, False,  True,  True,  True,  True]]]])
```

<img width="542" alt="Screenshot 2025-07-03 at 11 56 22" src="https://github.com/user-attachments/assets/d6b2b62f-5937-4afb-a36c-e9f510d4d6cd" />



Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_39194). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: arcee, aria, bitnet, cohere, cohere2, csm, deepseek_v3, dia, diffllama, dots1, emu3, gemma, gemma2, gemma3, gemma3n, glm

Hey @Cyrilvallez the [docstring](https://github.com/huggingface/transformers/pull/39194/files#diff-fca2adc00a38d9fa729672b0d32752b4252a35b1c296397c208ff0f1c764e07bR1040) of `create_masks_for_generate` says the argument is optional:

>         position_ids (`torch.Tensor`, optional)
>            A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences.

but it's actually a required argument, breaking existing code that calls this function. Is it intentional that it's required or should it have a default of `None`?

Hi, I would like to inquire whether you could implement the attention_mask with the pattern [1,1,2,2,2,3,3,3], and support packed tensors with FlashAttention for scenarios requiring a sparse mask.

This approach would enable us to leverage a universal method (2D attention mask/position IDs) to handle variable-length attention via masking. Additionally, we could extend support to 4D masks for more complex cases, building upon SDPA (Scaled Dot-Product Attention) and eager attention.

Hey @BenjaminBossan! It's optional in the sense that it can be None, but indeed I did not provide a default of None to force the models to pass the argument to always allow packed format (same as the past_key_values).
We could however rethink the default values maybe (e.g. it could make sense to allow cache_position to be None as well when the kv are as well, and construct them on the fly for external usage of the mask functions). Let me know your thoughts!

@Snowdar for your usage of FA2, you should not pass any mask mask but forward the seqlens directly 🤗

> It's optional in the sense that it can be None, but indeed I did not provide a default of None to force the models to pass the argument to always allow packed format (same as the past_key_values).
> We could however rethink the default values maybe (e.g. it could make sense to allow cache_position to be None as well when the kv are as well, and construct them on the fly for external usage of the mask functions). Let me know your thoughts!

Thanks for explaining. I think in this case, it would have been better to provide a default, given that the signature was changed in a backwards incompatible way and then the change was released as a patch release, where the expectation as a user is that I can always upgrade without fear of breakage. I'm not sure if this function is considered "private", but even so, I think providing a default when there is a reasonable one would have been better. Now that the patch release is out, it's too late so I don't have any strong opinion either way.

Yep I think we need to patch again to have a default @Cyrilvallez