Add packed tensor format support for flex/sdpa/eager through the mask!

[Cyrilvallez]

What does this PR do?

As per the title.

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

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

Very nice, just missing a test 😉

[winglian]

There really shouldn't be a restriction on this. It should work too with 2D packed tensors.

[Cyrilvallez]

Humm alright, I can lift it easily - thought it was always packed with all sequences along a batch of 1

[github-actions]

[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

[winglian]

perfect! fixes the regression and flex + packing works again for us now

[BenjaminBossan]

Hey @Cyrilvallez the docstring 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?

[Snowdar]

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.

[Cyrilvallez]

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!

[Cyrilvallez]

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

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

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.

[ArthurZucker]

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