[Attention] Use FA4 for MLA prefill

[MatthewBonanni]

Purpose

Depends on vllm-project/flash-attention#123 and vllm-project/flash-attention#126

Enable using FA4 as an MLA prefill (forward_mha) backend. This PR makes FA4 the default MLA prefill backend on Blackwell due to its improved performance over TRT-LLM, especially for extend operations relevant to the chunked prefill performed in vLLM

Test Plan

python benchmarks/attention_benchmarks/benchmark.py --config benchmarks/attention_benchmarks/configs/mla_prefill.yaml

Test Result

Accuracy

DeepSeek-V2-Lite-Chat

FA4 (PR)

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6732|±  |0.0129|
|     |       |strict-match    |     5|exact_match|↑  |0.6657|±  |0.0130|

TRTLLM (Reference)

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

---

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.

[mergify]

Documentation preview: https://vllm--34732.org.readthedocs.build/en/34732/

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @MatthewBonanni.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

Code Review

This pull request introduces support for FlashAttention-4 (FA4) for MLA prefill, which is a significant enhancement. The changes are extensive, touching the core attention logic, build system, and benchmarking infrastructure. The implementation is well-structured, including fallbacks for hardware limitations and clear configuration options. The benchmark suite has been notably improved to allow for detailed comparison of different prefill backends. I've identified one potential issue in the CMake configuration that could lead to incorrect files being installed.

[LopezCastroRoberto]

LGTM, nice analysis!