[Attention] Abstract the MLA prefill backends

[MatthewBonanni]

Purpose

Abstracts the MLA prefill backends to simplify mla_attention.py and introduces a selection mechanism similar to that of the decode backends, via --attention-config.mla_prefill_backend. Old AttentionConfig arguments (use_cudnn_prefill, use_trtllm_ragged_deepseek_prefill, and disable_flashinfer_prefill) are retained (with deprecation warnings) for backwards compatibility.

Test Plan

tests/v1/attention/test_mla_prefill_selector.py

(introduced by this PR) should pass in CI (part of V1 attention (H100) and V1 attention (B200))

Test Result

TBD

---

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.

[gemini-code-assist]

Code Review

This pull request introduces a well-designed abstraction for MLA prefill backends, which significantly simplifies mla_attention.py and improves modularity. The new selection mechanism via --attention-config.mla_prefill_backend is a great addition, and the backward compatibility for old flags is handled correctly. The refactoring moves backend-specific logic into separate, well-organized files, making the code cleaner and more maintainable. I've found one issue related to a hardcoded device that could affect non-CUDA platforms, which I've commented on. Overall, this is an excellent refactoring effort.

[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

[mergify]

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

[mgoin]

Looks really clean to me, nice work! I'll defer to Lucas and would like someone from AMD to review as well cc @tjtanaa

[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

[mergify]

Hi @MatthewBonanni, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[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

[mergify]

Hi @MatthewBonanni, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @MatthewBonanni, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10