Enabling Flash Attention 2

@varadhbhatnagar were you able to figure out the difference?