investigate compile flag in flex attention

[felipemello1]

context: https://x.com/cHHillee/status/1851418255749169419

We may be leaving some performance on table

[calvinpelletier]

We should also compare using _compile=True vs torch.compile. Horace recommends using _compile=True but someone else found that torch.compile was faster

Note: might be affected by this related issue which was fixed by this recent PR