Flex attention support with arbitrary 4d mask for LlamaModel
Author: alex-hh
Pull request: no
Labels: name: Feature request

name: Flash Attention

### Feature request

It would be nice to combine the benefits of flex attention and 4d masking.

Perhaps the llama model could be a first case, allowing arbitrary 4d masks to be handled via an efficient flex attention path.

### Motivation

Custom attention masking/biasing patterns lead to considerable improvements in flexibility, and are central to state-of-the-art models like AlphaFold and recent multimodal models.

4d attention masking in Transformers already provides the user with the flexibility to define custom biases, however performance is limited by the fact that 4d masking is incompatible with flash attention.

A 4d mask supporting flex attention attention path would retain full flexibility while maintaining performance. As far as I understand, nothing comparable exists in Transformers currently.

### Your contribution

Not very familiar with what this would take.

Comments:
Thanks for your request! cc @ArthurZucker 

Sounds super good! If a contributor wants to have a go happy to review, otherwise planned for sure! See #32877

@ArthurZucker I might be interested in having a go - I think it might involve creating a separate attention_bias argument (corresponding to current 4d attention mask and flex attention's score_mod); and allowing a binary/boolean attention_mask to determine mask_mod. Would you support this kind of change?

This way it seems like it should be possible to generate appropriate score_mod and mask_mod functions automatically I think, which might be nicer API wise than allowing the user to pass these explicitly, but I'm not sure and would be interested to know if you have thoughts on this before attempting anything? Don't know if it's a bit suboptimal to explicitly instantiate attention masks then convert back to flex attention functions for example.

In general we most probably won't allow passing the `score_mod` and `mask_mod`, we might allow defining it for `Gemma2Attention` for example: this way it's easy to overwrite. 

Automatic creation sounds a bit tricky + we'd rather have reference implementation for each model IMO! 🤗 

Ok makes sense I will maybe have a think about a reference implementation for llama.

By automatic creation I was specifically thinking of (actually an example in the flex attention post):

```python
def create_bias_mod(bias)
    def bias_mod(score, b, h, q_idx, kv_idx):
        return score + bias[b, h, q_idx, kv_idx]
    return bias_mod
```

I don't know enough to understand how useful this would be performance-wise

Will link a PR to this issue if I start working on this before then! 🤗 

Any updates on this?

I think this is fixed you can pass a BlockMask to the models 

Hey @alex-hh, sorry for the very slow follow-up here, we're catching up on the backlog.

Good news: this is now supported on current `main`. `LlamaModel` sets `_supports_flex_attn = True`, and the flex path handles arbitrary 4D masks two ways: you can pass a `torch.nn.attention.flex_attention.BlockMask` directly (used as the sparse `block_mask`), or pass a custom 4D float tensor as `attention_mask`, which is applied additively as a bias inside `score_mod`. Exactly the `create_bias_mod` pattern you sketched, see https://github.com/huggingface/transformers/blob/150eb7c9ed4091294c829fa0e9466b090cb0f87f/src/transformers/integrations/flex_attention.py#L282-L294

Any already-4D mask is passed straight through to the attention interface https://github.com/huggingface/transformers/blob/150eb7c9ed4091294c829fa0e9466b090cb0f87f/src/transformers/masking_utils.py#L817-L818

Thanks for the detailed proposal! (cc @minpeter, who asked for an update.)