[Attention] Implement universal BACKEND_MAP #25900

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Merged

simon-mo merged 14 commits into vllm-project:main from MatthewBonanni:registry_utils

Oct 8, 2025

Contributor

MatthewBonanni commented Sep 29, 2025 •

edited by github-actions bot

Loading

Purpose

This PR continues the attention backend selection refactor #24794, which is now being split up. Here, we allow backends to register themselves in a BACKEND_MAP which will be used later in the refactor. This also enables out-of-tree backends for other hardware to plug in neatly.

Test Plan

CI should be sufficient

Test Result

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.

MatthewBonanni requested review from ApostaC, LucasWilkinson, NickLucche, WoosukKwon, bigPYJ1151, gshtras, jikunshang, mgoin, sighingnow, tdoublep, tlrmchlsmth and yewentao256 as code owners

September 29, 2025 19:44

mergify bot added qwen rocm speculative-decoding v1 tpu kv-connector labels

gemini-code-assist bot reviewed

View reviewed changes

Contributor

gemini-code-assist bot left a comment

Code Review

This pull request refactors the attention backend selection mechanism by introducing a centralized registry. Backends now register themselves using a decorator, which is a significant improvement in modularity and removes hardcoded mappings. The changes are extensive and touch many files, mostly for updating imports and applying the new registration decorator.

The core logic is sound, but I've identified one high-severity issue related to testing on ROCm. The refactoring removed a platform-specific hack from a test utility without updating the tests that relied on it. This will lead to tests being skipped on ROCm, reducing test coverage for that platform. My review comment provides details on this issue.

tests/v1/attention/utils.py Outdated Show resolved Hide resolved

MatthewBonanni changed the title ~~Backends register themselves in BACKEND_MAP~~ [Attention] Backends register themselves in BACKEND_MAP

mergify bot commented Sep 30, 2025

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 bot added the needs-rebase label

MatthewBonanni marked this pull request as draft

September 30, 2025 14:37

MatthewBonanni mentioned this pull request

[Attention] Move Backend enum into registry #25893

Merged

5 tasks

MatthewBonanni added 2 commits

October 3, 2025 13:24


          attention backends register themselves in the mapping

3ba013e

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>


          rename

55ad364

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

MatthewBonanni force-pushed the registry_utils branch from 20a75cb to 55ad364 Compare

October 3, 2025 13:31

mergify bot removed the needs-rebase label


          fix rocm

950f955

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

MatthewBonanni marked this pull request as ready for review

October 3, 2025 13:39

Contributor Author

MatthewBonanni commented Oct 3, 2025

cc @ILikeIneine this has the registration mechanism you proposed in #24794!

MatthewBonanni added 4 commits

October 6, 2025 11:17


          Merge branch 'main' into registry_utils

2602f5c

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>


          fix pre-commit

8f3d4d9

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>


          Merge branch 'main' into registry_utils

e2dbb43

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>


          register aiter unified attn

21a0e10

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

MatthewBonanni mentioned this pull request

[Attention][UX] Add AttentionConfig and change attention backend to CLI argument #26315

Open

5 tasks

NickLucche approved these changes

View reviewed changes

Collaborator

NickLucche left a comment

LGTM , thanks @MatthewBonanni !

Left a few minor comments, also please check with AMD folks if the backend renaming is fine (just to be on the safe side).

vllm/attention/backends/registry.py Show resolved Hide resolved

vllm/attention/backends/registry.py Outdated Show resolved Hide resolved


          update comment

715f965

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Contributor Author

MatthewBonanni commented Oct 8, 2025

@NickLucche Thanks for your review! I've addressed your comments.

NickLucche added the ready label

mergify bot commented Oct 8, 2025

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 bot added the needs-rebase label


          Merge branch 'main' into registry_utils

ff18c0b

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

mergify bot removed the needs-rebase label

MatthewBonanni added 2 commits

October 8, 2025 11:44


          revert to manual registration because backends are not imported prior…

7e31b74

… to selection. Leave decorator to enable out-of-tree backend plugins

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>


          Merge branch 'registry_utils' of github.com:MatthewBonanni/vllm into …

825fcee

…registry_utils

mergify bot removed the tpu label

MatthewBonanni changed the title ~~[Attention] Backends register themselves in BACKEND_MAP~~ [Attention] Implement universal BACKEND_MAP


          fix test expectation

40c7ea3

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

ProExpertProg reviewed

View reviewed changes

vllm/attention/backends/registry.py

    
              }

              def register_attn_backend(backend: _Backend, class_path: Optional[str] = None):

Collaborator

ProExpertProg Oct 8, 2025

How does this work for OOT backends, wouldn't they have to extend the _Backend enum?

Contributor Author

MatthewBonanni Oct 8, 2025 •

edited

Loading

The plan is for them to override an existing backend (see #24794 (comment)). We could extend this in the future (and find an alternative to the enum) to make it cleaner though.

simon-mo merged commit 76879cc into vllm-project:main

50 of 53 checks passed

mrasquinha-g pushed a commit to mrasquinha-g/vllm that referenced this pull request


          [Attention] Implement universal BACKEND_MAP (vllm-project#25900)

a3269e3

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

zhiyuan1i pushed a commit to zhiyuan1i/vllm that referenced this pull request


          [Attention] Implement universal BACKEND_MAP (vllm-project#25900)

798afcf

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

xuebwang-amd pushed a commit to xuebwang-amd/vllm that referenced this pull request


          [Attention] Implement universal BACKEND_MAP (vllm-project#25900)

426ccb4

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: xuebwang-amd <xuebwang@amd.com>

Dhruvilbhatt pushed a commit to Dhruvilbhatt/vllm that referenced this pull request


          [Attention] Implement universal BACKEND_MAP (vllm-project#25900)

0c82da1

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Dhruvil Bhatt <bhattdbh@amazon.com>

lywa1998 pushed a commit to lywa1998/vllm that referenced this pull request


          [Attention] Implement universal BACKEND_MAP (vllm-project#25900)

e8bd559

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

alhridoy pushed a commit to alhridoy/vllm that referenced this pull request


          [Attention] Implement universal BACKEND_MAP (vllm-project#25900)

40fa046

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

xuebwang-amd pushed a commit to xuebwang-amd/vllm that referenced this pull request


          [Attention] Implement universal BACKEND_MAP (vllm-project#25900)

9de664c

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: xuebwang-amd <xuebwang@amd.com>

0xrushi pushed a commit to 0xrushi/vllm that referenced this pull request


          [Attention] Implement universal BACKEND_MAP (vllm-project#25900)

6d0f954

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: 0xrushi <6279035+0xrushi@users.noreply.github.com>

0xrushi pushed a commit to 0xrushi/vllm that referenced this pull request


          [Attention] Implement universal BACKEND_MAP (vllm-project#25900)

cc00d24

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: 0xrushi <6279035+0xrushi@users.noreply.github.com>

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

ProExpertProg ProExpertProg left review comments

NickLucche NickLucche approved these changes

gshtras Awaiting requested review from gshtras

sighingnow Awaiting requested review from sighingnow

jikunshang Awaiting requested review from jikunshang

bigPYJ1151 Awaiting requested review from bigPYJ1151

mgoin Awaiting requested review from mgoin mgoin is a code owner

tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth is a code owner

WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner

yewentao256 Awaiting requested review from yewentao256 yewentao256 is a code owner

tdoublep Awaiting requested review from tdoublep

LucasWilkinson Awaiting requested review from LucasWilkinson LucasWilkinson is a code owner

ApostaC Awaiting requested review from ApostaC ApostaC is a code owner

+1 more reviewer

gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Labels

kv-connector qwen ready rocm speculative-decoding v1