[Attention] Implement universal BACKEND_MAP #25900
Conversation
MatthewBonanni
requested review from
ApostaC,
LucasWilkinson,
NickLucche,
WoosukKwon,
bigPYJ1151,
gshtras,
jikunshang,
mgoin,
sighingnow,
tdoublep,
tlrmchlsmth and
yewentao256
as code owners
mergify
bot
added
qwen
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this 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.
MatthewBonanni
changed the title
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Matthew Bonanni <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]>
MatthewBonanni
force-pushed
the
registry_utils
branch
from
20a75cb to
55ad364
Compare
Signed-off-by: Matthew Bonanni <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]>
NickLucche
left a comment
NickLucche
left a comment
There was a problem hiding this 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).
Signed-off-by: Matthew Bonanni <[email protected]>
|
@NickLucche Thanks for your review! I've addressed your comments. |
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Matthew Bonanni <[email protected]>
… to selection. Leave decorator to enable out-of-tree backend plugins Signed-off-by: Matthew Bonanni <[email protected]>
MatthewBonanni
changed the title
Signed-off-by: Matthew Bonanni <[email protected]>
| } | ||
|
|
||
|
|
||
| def register_attn_backend(backend: _Backend, class_path: Optional[str] = None): |
There was a problem hiding this comment.
How does this work for OOT backends, wouldn't they have to extend the _Backend enum?
There was a problem hiding this comment.
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.
Signed-off-by: Matthew Bonanni <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]> Signed-off-by: xuebwang-amd <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]> Signed-off-by: Dhruvil Bhatt <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]> Signed-off-by: xuebwang-amd <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]> Signed-off-by: 0xrushi <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]> Signed-off-by: 0xrushi <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]>
Signed-off-by: Matthew Bonanni <[email protected]>
4 participants
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_MAPwhich 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
supported_models.mdandexamplesfor a new model.