Adding Flash attention

[adi-kmt]

Will improve the training and inference speed by a large margin!!

[apapiu]

Hey the model uses https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html which should already use flash attention.