[P] Speed Up Stable Diffusion by ~50% Using Flash Attention

We got close to 50% speedup on A6000 by replacing most of cross attention operations in the U-Net with flash attention

Annotated Implementation: [https://nn.labml.ai/diffusion/stable\_diffusion/model/unet\_attention.html#section-45](https://nn.labml.ai/diffusion/stable_diffusion/model/unet_attention.html#section-45)

Github: [https://github.com/labmlai/annotated\_deep\_learning\_paper\_implementations/blob/master/labml\_nn/diffusion/stable\_diffusion/model/unet\_attention.py#L192](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/diffusion/stable_diffusion/model/unet_attention.py#L192)

We used this to speed up our stable diffusion playground: [promptart.labml.ai](https://promptart.labml.ai/)

attention diffusion machinelearning stable diffusion