🚨🚨[core] Completely rewrite the masking logic for all attentions

[Cyrilvallez]

What does this PR do?

As per the title. The goal is to properly separate masking logic from modeling code itself, to continue our objective of simplifying the library.

• Code is much simpler to understand
• Much more general: always work for all lengths, and all attention implementations, e.g.:
  • flex attention now works with sliding/hybrid models (not the case before)
  • FA2 now works with static caches (including models with default hybrid structures) (was only the case for hybrid models before)
• All models can use all Cache classes (e.g. models with Hybrid structure can default back to use DynamicCache)
• Extremely scalable in the future: any pattern of layers can be taken into account WITHOUT ANY CHANGE to modeling or masking. A new masking pattern (e.g. the recently introduced chunked attention for Llama4) can be added with minimal efforts (just add a new mask_mod to describe it, and voila!)
• A single truth: mask creation was copied over and over again, but sometimes with slight changes to account for sliding windows or similar. This would eventually lead to mistakes or inefficiencies as things would be "forced to fit", and a lot of maintenance burden
• compile compatible: the new mask creation is technically compile compatible - it should however stay outside what is compiled in the forward to avoid recompilations as it's being done in generate
• Allow external mask creation: In case someone passes their custom attention implementation, they may need their own mask creation function, which is now supported
• TGI/vLLM backend should be even more efficient now, as we don't waste compute on creating a useless mask (would previously create a 4d mask as for sdpa, which would not be used)

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

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

Looks very vey nice!
One thing I want to consider is to rather call the sliding, causal and chuncking directly in the modeling.
For example:

• llama only need causal_mask, under the hood the causal mask should do an and with sddpa or cflash or flex
• gemma need sliding_causal: same
• llama4 needs chuncked causal
  I want the modeling to call an explicit function, rather than the mega general one!

This would keep our philosophy, as we don't want too general stuff hapenning when not needed (ex: llama should never care about sliding in codepathes)

Also misssing doc about how to add a new func!

[ydshieh]

Wow!!!!!!!! 🚀

This PR seems worth a manually full CI. Ping me when it's time you think this PR is ready to trigger CI.

[ArthurZucker]

Damn nice

[ArthurZucker]

REview for the core logic, IMO can be simplified! BUt the modeling part is absolutely perfect!
For the visualization, I'll see how we could just overwrite the repr without affecting other operations!

[vasqu]

Just a quick question on this refactor: If I understand the code correctly, then the focus is currently on causal masks only, correct?

Would be nice to add a non-causal alternative which should only use a padding mask and expand respectively to the q_len and kv_len. That's more food for thought :D I dont want to make this PR even harder than it is.

[Cyrilvallez]

For now it's mostly on causal masks because they are the one we need, but the idea is that it can be extended super easily from a set of mask primitives!

[ArthurZucker]

Mega nice!
TODO before merging:

1. move the causal_mask_mapping to a class attribute!
2. show example of how to register a new function, but minimal (without sdpa correction for example)
3. Make sure full graph training is not broken maybe? or at least fa2 training

That should be i

[ArthurZucker]

Let's go!

[ArthurZucker]

let's rather show how to do something like the paligemma or document masking here, something relevant!

[Cyrilvallez]

Those are a bit different, it's modifying the mask pattern vs adding a new mask format for the attention itself (both are complementary)

[BenjaminBossan]

Hi @Cyrilvallez I noticed that after this PR, calling prepare_inputs_for_generation can return an attention_mask that is a dict instead of a tensor. Is this expected? If yes, I need to update PEFT.

Reproducer:

from transformers import AutoModelForCausalLM
from transformers.cache_utils import HybridCache
import torch

model_id = 'hf-internal-testing/tiny-random-Gemma3ForCausalLM'
model = AutoModelForCausalLM.from_pretrained(model_id)

inputs = torch.arange(6).view(2, 3)
attention_mask = torch.ones_like(inputs)
# cache is required, w/o cache a tensor is returned as expected
cache = HybridCache(model.config, max_batch_size=2, max_cache_len=3)

model_kwargs = model.prepare_inputs_for_generation(
    inputs, attention_mask=attention_mask, past_key_values=cache, cache_position=torch.arange(3)
)
mask = model_kwargs['attention_mask']
assert isinstance(mask, torch.Tensor), f"expected attention mask to be tensor, got {mask}"

Before the PR (f8630c778c9220defecf1e3026d3438108b0baba), this passes. After the PR (163138a911c1fb4451ec4b32edaee20918a59def), it fails with:

AssertionError: expected attention mask to be tensor, got {'sliding_attention': tensor([[[[ True, False, False],
          [ True,  True, False],
          [ True,  True,  True]]],


        [[[ True, False, False],
          [ True,  True, False],
          [ True,  True,  True]]]])}

[guangy10]

@Cyrilvallez What is layer_types? I'm concerning whether changes here are backwards compatible. For existing models on Hub like google/gemma-3-1b, it doesn't seem to come with the layer_types so it will fallback to the static cache which doesn't look correct.

[kimishpatel]

One comment i have is that the way mask calculation is incorporated in most models is that the calculation of mask happens at model level. e.g. here https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma3/modeling_gemma3.py#L565-L566, however, different cache implementations may imply different attention masks. Different layers may have different cache impl, for example some layers can have sliding window of different size, others may use attention sink to keep say first few or some tokens. I feel the best way for the custom mask is at the attention layer so that the said layer can pass in all the information, including kv cache, to the custom mask function (e.g. layer_index).

[Cyrilvallez]

Hey, sorry all I was on vacations!

@BenjaminBossan indeed, this is expected. Models using different types of attention in different layers (i.e. gemma3) will now have a dict returned by prepare_inputd_for_generation (one dict entry per attention type). Sorry this broke your tests! 😬

@guangy10 if you look at the configs, e.g. here, you'll see that the attribute was added in a BC manner for all models that were refactored! Let me know if you notice any issue though!

@kimishpatel In transformers, Cache are not at the layer level, so as of now only some configurations are acceptable (though I've had in mind to change that for some time, to make it more modular). And computing the mask at the AttentionLayer-level is not only redundant (most layers will create the same mask, wasting precious time), but it breaks compile completely, as we cannot pre-compute the masks anymore. For now, there are no known models with sliding windows of different sizes for different layers, so we decided to make it as simple as possible. This was taken into account when doing this refactor though, no worries, we definitely thought about it to scale easily in the future should this scenario happen

[kimishpatel]

And computing the mask at the AttentionLayer-level is not only redundant (most layers will create the same mask, wasting precious time),

To be fair, I doubt attention mask calculation has that much impact on performance for the most models. I have implemented ring buffer based kv cache, that needs a very different way of calculating mask and that mask calculation, while redundant, happens at attention layer. I have not observed any significant amount of time spent in there.

Although I think for block mask in flex attention, you might be right. That one is non-trivial.

Cache are not at the layer level,

how so? cache's update functions accept layer_idx, so they do have to know what layer the update belongs to.

I do understand though that transformers is not exactly providing building blocks for model authoring so from that perspective composability and modularity has limited value i suppose

[guangy10]

@guangy10 if you look at the configs, e.g. here, you'll see that the attribute was added in a BC manner for all models that were refactored! Let me know if you notice any issue though!

@Cyrilvallez Let's follow up in #38646