flex_attention support for Qwen2.5/Gemma is broken
Author: flukeskywalker
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.50.3
- Platform: Linux-6.8.0-52-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.30.1
- Safetensors version: 0.5.3
- Accelerate version: not installed
- Accelerate config: not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0+cu124 (True)
- 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?: No
- Using GPU in script?: Yes
- GPU type: NVIDIA GeForce RTX 4090

### Who can help?

@ArthurZucker 

### Information

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

### Tasks

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

### Reproduction

```python
import torch
from torch.nn.attention.flex_attention import create_block_mask
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "Qwen/Qwen2.5-0.5B" # AttributeError: 'BlockMask' object has no attribute 'dim'
# model_name = "meta-llama/Llama-3.2-1B" # works

model = AutoModelForCausalLM.from_pretrained(model_name, attn_implementation="flex_attention").to("cuda")
tokenizer = AutoTokenizer.from_pretrained(model_name)
print(model._supports_flex_attn)

tokens = tokenizer("Hello model, this is human. ")["input_ids"]

def causal_mask(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

block_mask = create_block_mask(causal_mask, None, None, len(tokens), len(tokens), device="cuda")
model(input_ids=torch.tensor(tokens, dtype=torch.long, device="cuda").unsqueeze(0), attention_mask=block_mask)
```

### Expected behavior

This snippet should run without error, like it does for `meta-llama/Llama-3.2-1B`, since Qwen2.5 model is based on Llama arch and both support flex_attention (`_supports_flex_attention=True`).

The error occurs because `Qwen2Model._update_causal_mask()` doesn't handle the case when flex_attention is enabled and the block mask is passed in as the `attention_mask`. This is handled in `LlamaModel._update_causal_mask()`:

```python
        if self.config._attn_implementation == "flex_attention":
            if isinstance(attention_mask, torch.Tensor):
                attention_mask = make_flex_block_causal_mask(attention_mask)
            if isinstance(attention_mask, BlockMask):
                return attention_mask
```

IIUC, adding the same handling to Qwen2Model should fix the issue, and indeed this works on my local fork. But Qwen2Model code is auto-generated, so it must be fixed elsewhere.

Comments:
I haven't used Gemma yet but noticed that its implementation also omits handling the flex_attention case in the `_update_causal_mask()` method [here](https://github.com/huggingface/transformers/blob/d1b92369ca193da49f9f7ecd01b08ece45c2c9aa/src/transformers/models/gemma3/modeling_gemma3.py#L754). 

Please correct me if I'm wrong, but this appears to be a wider bug in the model code generation logic that misses handling flex_attention when the model has `_supports_flex_attention=True`.

Yes, this is required indeed ! Do you want to open a PR for a fix? 

Yes I'd like to! Will do once I'm sure what the fix is.

I'm new to internals of `transformers` and the details of the model codegen aren't fully clear to me, but digging more into the codebase I think the Qwen2 issue might be caused due to inheriting from `MistralModel` instead of `LlamaModel`. 

Now `MistralModel` itself inherits from `LlamaModel` so this would be okay but it also overrides `_update_causal_mask()` without handling the block mask. I think this means that the issue also affects Mistral models. But I'm not sure if there is a good reason why `Qwen2Model` inherits from `MistralModel` or why `MistralModel` overrides `_update_causal_mask()`.

Potential fix ideas:
1. Make `Qwen2Model` inherit directly from `LlamaModel` to fix Qwen2
2. Remove overriden methods `_update_causal_mask()` and `_prepare_4d_causal_attention_mask_with_cache_position()` from `MistralModel` since these are already mostly defined by `LlamaModel`. 

Technically 2 alone would solve the Qwen2 issue too, but 1 would still simplify things.

---

Similarly, the Gemma3 text model inheritance goes llke: `Gemma3TextModel` <- `Gemma2Model` <- `GemmaModel` <- `LlamaModel`. This would in principle work with all attention interfaces, except that `Gemma2Model` overrides `_update_causal_mask()` causing break in support for flex_attention. 

3. The fix here could be to simply remove `Gemma2Model._update_causal_mask()`. Edit: Like the Mistral case, we'd need to handle the case of Hybrid/Static cache 

---

It comes down to keeping all masking related logic contained in `LlamaModel` which would hopefully fix flex_attention support for Qwen2, Gemma2, Gemma3 and Mistral models. What do you think about fixes 1, 2 and 3?

Edit: the above fixes would require some additions to `LlamaModel._update_causal_mask()` to handle/check the different cache types supported by different models. It is possible that you do not want to bloat `LlamaModel` like this, so  in that case the fixes would consist of checking the `attn_implementation` and creating/passing along the block mask in `MistralModel` to fix Mistral/Qwen2, and `Gemma2Model` to fix Gemma2 and 3. Do you prefer this simpler fix?

PS: why does `Qwen2Model` inherit from `MistralModel`?

just update the overwritter `_update_causal_mask` where needed, the simpler the better. FLex should be on most / all models if possible! 

@Cyrilvallez is reworking our attention creation api