add flash attn 2

[williamberman]

Marginal improvements over xformers, note that xformers itself isn't a tremendous improvement over vanilla attention.

We have relatively small sequence lengths of 256 for self attention and 256x77 for cross attention. Our head dimensions of 16 heads of 64 inner dim as well are on the smaller side.

Only linking f8 runs, measurements are same for f16 as we have the same attention block configs/sequence lengths https://wandb.ai/williamberman/muse/runs/9heg02rd?workspace=user-williamberman https://wandb.ai/williamberman/muse/runs/2tc3ueir?workspace=user-williamberman https://wandb.ai/williamberman/muse/runs/f4e8r72x/workspace?workspace=user-williamberman