feat: make flash-attn configurabele

[theissenhelen]

This adds a flag for flash attention.

See ecmwf/anemoi-models#47

[mchantry]

LGTM

[FussyDuck]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ theissenhelen
❌ Charlie Hebert-Pinard

---

Charlie Hebert-Pinard seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}