[Hardware][AMD] Enable FlexAttention backend on ROCm

[mawong-amd]

Purpose

Enable FlexAttention on ROCm if it is specified e.g. via VLLM_ATTENTION_BACKEND=FLEX_ATTENTION. This makes progress towards batch invariant inference on AMD.

Test Plan

Comparing correctness of default attention backend vs FlexAttention
lm_eval --model local-completions --model_args model=meta-llama/Llama-3.1-8B,base_url=http://0.0.0.0:8000/v1/completions,num_concurrent=128,max_retries=5 --tasks gsm8k

Test Result

Default:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.5011|±  |0.0138|
|     |       |strict-match    |     5|exact_match|↑  |0.5011|±  |0.0138|

FlexAttention:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.4882|±  |0.0138|
|     |       |strict-match    |     5|exact_match|↑  |0.4875|±  |0.0138|

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.