Add flash attention 2

[adhikjoshi]

Flash attention 2 helps in faster speeds, here is birch labs implementation in diffusers

https://gist.github.com/Birch-san/4315701264b72bb72e8eac5a529ee93a

@patrickvonplaten

[sayakpaul]

Can we have some inference timing comparison of diffusers with

• SDPA
• xformers
• Flash attention 2?

Without a good handle on how much gains we can expect from using Flash Attention 2, it's hard for us to decide here.

Cc: @williamberman

[patrickvonplaten]

Very cool! I think we should support as soon as it's in xformers - think we're close to getting a new pypi version for xfromers: facebookresearch/xformers#795 (comment)

[williamberman]

Yeah once the xformers release is cut, you should have access to it. The api is the same so we shouldn't have to update the diffusers code. I'm going to close the issue since I don't think we need to make any changes to diffusers source :)