Add use_flash_attn for better FA + FA2 feature gating

[alvarobartt]

What does this PR do?

This PR adds use_flash_attn as a function to check whether a given model architecture supports Flash Attention based on the active features + compute capability; and also capturing the value for the env var USE_FLASH_ATTENTION in an unified way.

This PR was inspired by both #778 and #809, which create similar functions.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the documentation guidelines.
• Did you write any new necessary tests? If applicable, did you include or update the insta snapshots?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.