FlashAttention2 implementation for OpenELM model

[GorkaUrbizu]

Model description

OpenELM is already available on HF🤗 and transformers, but the model lacks support for flashAttention/FlashAttention2.

I'd love to have the flash attention available for OpenELM in the transformers enviroment.

thanks in advance.

Open source status

• The model implementation is available
• The model weights are available

Provide useful links for the implementation

No response

[Rocketknight1]

Hi @GorkaUrbizu, OpenELM uses custom code, and so we can't actually add support in Transformers! However, I believe it already supports SDPA, which can dispatch to FlashAttention on newer versions of PyTorch.