[Feature] Will there be any integration of using Flex-attention (and Paged attention)?
Author: jianan-gu
Pull request: no
Labels: name: Feature request

### Feature request

Using (https://pytorch.org/blog/flexattention/) Flex-attention (and [Paged attention](https://github.com/pytorch/pytorch/pull/121845/files)) to speedup transformers models and provide flexibility






### Motivation

FlexAttention was proposed as a performant attention implementation leveraging torch.compile with easy APIs for adding support for complex attention variants such as Causal, [Relative Positional Embeddings](https://paperswithcode.com/method/relative-position-encodings), [Alibi](https://paperswithcode.com/method/alibi), [Sliding Window Attention](https://mistral.ai/news/announcing-mistral-7b/), [PrefixLM](https://twitter.com/andersonbcdefg/status/1800907703688339569), https://github.com/pytorch/torchtune/pull/875, [Tanh Soft-Capping](https://twitter.com/LysandreJik/status/1807779471891538199), [PagedAttention](https://arxiv.org/abs/2309.06180), etc.





### Your contribution

Not sure.

Comments:
There is a start for flex attention in #34282 (gemma esp. due to the tanh softcapping) 

I guess it's planned for the future but as everything new, some stuff needs to be fleshed out and it takes time 

cc @ArthurZucker 

> There is a start for flex attention in #34282 (gemma esp. due to the tanh softcapping)
> 
> I guess it's planned for the future but as everything new, some stuff needs to be fleshed out and it takes time
> 
> cc @ArthurZucker

Thanks for pointing out this start! It looks great and it is exciting to have the integration. Will the integration be limited to specific models only or common for most of the models (like llama or others)?

Besides, is there any planned Transformers release version to begin to include this feature? And will there be also Paged-attention based integration, and related KV cache management for the model generation? 

Thanks!


Mind that I'm not a maintainer :) but here are my thoughts on the integration

Flex attention is only in the prototype phase (prototype -> beta -> stable) so I don't expect it to be officially supported until it's at least beta but just my wild guess. Why? Because there will be a lot of bugs (e.g. memory leaks) and fixes which would need to be accomodated here or referred to elsewhere.

When integration starts, I expect it to start with the most common ones e.g. llama, mistral, gemma etc. Afterwards, there will probably a community call for contributions as the number of models is simply too large to retroactively add it to all of them. It has been done for torch's sdpa too for example. 

As I mentioned before, i expect the full integration to start a bit later but maybe i'm wrong and since i'm not a maintainer my information is prolly not too trustworthy on this.

Paged attention is only available in main and also only really recently ( see https://github.com/pytorch/pytorch/pull/137164 ) so that will probably take a bit longer i guess. It will also need a new cache format that's working with current generation methods etc - i.e. more work. 

cc @LysandreJik if he knows more / can add more to this 😄 

Great question! Indeed, once ready we'll implement FlexAttention. @ArthurZucker will lead this effort.

@jianan-gu fyi #34809