Flash Attention 2 for audio/musicgen

[yoinked-h]

Feature request

Add flash attention 2 to musicgen models (and or audio gen)

Motivation

Musicgen is a really large model, and its very hard to run on consumer GPUs, adding flash attention could make the model much lighter and possibly faster.

Your contribution

I can test on both WSL and windows 11 on a 3060

[amyeroberts]

cc @sanchit-gandhi @ylacombe

[sanchit-gandhi]

Would indeed be cool having FA2 support for MusicGen! Since MusicGen copies the attention mechanism from Bart, you can copy the changes from this PR: #27203

Would you like to have a go at this integration @yoinked-h? It should be quite a fun PR where you get quite big speed-ups for not too much additional effort!

[yoinked-h]

I've tried to implement the Bart flash attention as best as i could, but i get some error (cu_seqlens_k must have shape (batch_size + 1)) from the module itself, will try tomorrow!

[staghado]

I would like to help if possible

[ylacombe]

Hey @yoinked-h and @staghado, as @sanchit-gandhi pointed it should be fun to add!
Feel free to open a PR so that you can have our help on this

[amosyou]

we can close this due to #29939

[yoinked-h]

amazing!