[Model][Mamba] Add selector for mamba attention backend and make it pluggable for other device #26487

shen-shanshan · 2025-10-09T12:40:43Z

Purpose

Motivation:

Like attention and MLA attention modules, we want to use some device-specific kernels for mamba layers and customize the proccessing of mamba attn backend, i.e., this is significant for running some mamba-like models (e.g., Qwen3-Next) on Ascend platform.

Main changes:

Add get_mamba_attn_backend() to attention selector, and force all mamba layers to get their attention backend by calling this method.
Add get_mamba_attn_backend_cls() to platform, thus other device besides GPU can custom their mamba attention backend.

Backend select priority:

Select device-specific mamba backend acording to mamba_type. If no customization here, then comes to the default logic.
Get default backend according to mamba_type from the mamba_type_to_backend_map.

Mamba layer:

Pass mamba_type to get_mamba_attn_backend() method to get its backend when initialization.

Update 2025/10/28:

Add _MambaBackend and MAMBA_BACKEND_MAP in registry.py.

Update 2025/11/10:

Add another argument linear_attn_type, together with mamba_type, to specify some special linear attention backend, e.g., GDNAttention.

linear_attn_type is optional:

None: use normal linear attention or other mamba backend.
not None: use mamba_type + linear_attn_type to get related backend.

Other changes:

Refactor KimiDeltaAttention to use get_mamba_attn_backend().
Update the doc about adding a new mamba type layer.

Update 2025/11/17:

Refactor follow #24794.

Test Plan

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.

gemini-code-assist

Code Review

This pull request introduces a pluggable selector for Mamba attention backends, refactoring the model layers to use this new centralized mechanism instead of hardcoded imports. This is a good architectural improvement for modularity. My review focuses on the robustness of the new selector. I've identified a potential issue in the error handling within vllm/attention/selector.py where the check for a valid backend class could be more robust and the error message more informative. I've provided a suggestion to address this.

vllm/attention/selector.py

MatthewBonanni · 2025-10-27T15:18:38Z

vllm/platforms/interface.py

+        mamba_type: str = "",
+    ) -> str:
+        """Get mamba attention backend class of a device."""
+        mamba_type_to_backend_map = {


Could you add a _MambaBackend enum and MAMBA_BACKEND_MAP in registry.py and use these instead?

@MatthewBonanni Hello, could you please help ask someone of the maintainers if anything else need updated or if this PR can be merged?

tdoublep

Sorry for the delay in reviewing. I think this change looks fine - have some minor questions + suggestions.

vllm/model_executor/models/qwen3_next.py

tdoublep · 2025-11-06T15:21:06Z

vllm/model_executor/models/registry.py

+    MAMBA1 = "vllm.v1.attention.backends.mamba1_attn.Mamba1AttentionBackend"
+    MAMBA2 = "vllm.v1.attention.backends.mamba2_attn.Mamba2AttentionBackend"
+    LINEAR = "vllm.v1.attention.backends.linear_attn.LinearAttentionBackend"
+    GDN = "vllm.v1.attention.backends.gdn_attn.GDNAttentionBackend"


I think we also need to handle the Kimi Linear case here.

I think we also need to handle the Kimi Linear case here.

Thanks for your suggestion! I will add it later.

mergify · 2025-11-10T08:27:53Z

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

shen-shanshan · 2025-11-10T08:34:08Z

CC @tdoublep @MatthewBonanni

I have updated this PR with updated descriptions in the purpose part of this PR.

LucasWilkinson · 2025-11-10T18:44:01Z

not super familiar with this area but is _MambaBackend the best name here? given the diverse set of implementation being added? would something like _LinearBackend, _LinearAttentionBackend, or _SSMAttentionBackend be better? cc @tdoublep ?

tdoublep · 2025-11-10T21:33:33Z

@LucasWilkinson I believe _MambaBackend is consistent with the rest of the codebase right now. We use "mamba" as a catch-all for mamba + linear attention mechnaisms.

vllm/attention/selector.py

vllm/model_executor/models/qwen3_next.py

MatthewBonanni · 2025-11-11T15:42:02Z

Now that #24794 has landed, could we refactor this PR to use the pattern from registry.py? i.e. _MambaBackendEnumMeta and MambaBackendEnum instead of _MambaBackend? You can make a separate _MAMBA_OVERRIDES too.

shen-shanshan · 2025-11-12T01:28:26Z

Now that #24794 has landed, could we refactor this PR to use the pattern from registry.py? i.e. _MambaBackendEnumMeta and MambaBackendEnum instead of _MambaBackend? You can make a separate _MAMBA_OVERRIDES too.

OK, I will update it soon.

wangxiyuan · 2025-11-17T01:26:00Z

vllm/platforms/interface.py

+        mamba_type: str,
+        linear_attn_type: str | None,
+    ) -> str:
+        """Get mamba attention backend class of a device."""


add more docstring here to describe the args usage. Thanks.

wangxiyuan · 2025-11-17T02:14:00Z

vllm/model_executor/layers/mamba/linear_attn.py

-        from vllm.v1.attention.backends.linear_attn import LinearAttentionBackend
-
-        return LinearAttentionBackend
+        return self.mamba_attn_backend


I notice all the get_attn_backend implementation is the same, why note implement it in MambaBase class?

wangxiyuan · 2025-11-17T02:15:27Z

vllm/model_executor/layers/mamba/linear_attn.py

            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self

+        self.mamba_attn_backend = get_mamba_attn_backend(self.mamba_type)


self.mamba_attn_backend is only used by get_attn_backend, why not let get_attn_backend call get_mamba_attn_backend directly? this self.mamba_attn_backend is looks unnecessary

Signed-off-by: shen-shanshan <467638484@qq.com>

shen-shanshan · 2025-11-17T09:23:30Z

CC @tdoublep @MatthewBonanni I have updated the code following #24794.

MatthewBonanni · 2025-11-17T20:06:55Z

LGTM!

jikunshang

thanks for refactor!

vllm/model_executor/models/qwen3_next.py

shen-shanshan · 2025-11-19T02:41:29Z

@jikunshang Hello, the CI are broken may due to something irrelevant to this PR. Could you please help retrigger it? Thanks.

jikunshang · 2025-11-19T02:49:05Z

failed case is irrelevant. I think it's ncessary to retrigger full test to avoid CI resource.
we can request a force merge. cc @DarkLight1337 Please also take a look, thanks!

gcanlin · 2025-11-19T05:39:08Z

The breaking of CI has been fixed by #28908. Please try merging main branch into this PR.

shen-shanshan · 2025-11-19T07:28:48Z

All the CI failures are due to the same error shown below:

RuntimeError: This flash attention build does not support headdim not being a multiple of 32.

gcanlin · 2025-11-19T08:04:12Z

All the CI failures are due to the same error shown below:
RuntimeError: This flash attention build does not support headdim not being a multiple of 32.

I also met this issue locally.

tdoublep

Great work - thank you

…luggable for other device (vllm-project#26487) Signed-off-by: shen-shanshan <467638484@qq.com>

…luggable for other device (vllm-project#26487) Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: LuminolT <lumischen01@gmail.com>

…luggable for other device (#26487) Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: jiang1.li <jiang1.li@intel.com>

…luggable for other device (vllm-project#26487) Signed-off-by: shen-shanshan <467638484@qq.com>

shen-shanshan requested review from LucasWilkinson, sighingnow and tdoublep as code owners October 9, 2025 12:40

shen-shanshan marked this pull request as draft October 9, 2025 12:40

mergify bot added the qwen Related to Qwen models label Oct 9, 2025

gemini-code-assist bot reviewed Oct 9, 2025

View reviewed changes

vllm/attention/selector.py Outdated Show resolved Hide resolved

shen-shanshan changed the title ~~[Refactor][Mamba] Add selector for mamba attention backend and make it pluggable for other device~~ [Mamba] Add selector for mamba attention backend and make it pluggable for other device Oct 10, 2025

shen-shanshan marked this pull request as ready for review October 10, 2025 07:47

shen-shanshan force-pushed the mamba branch from 693b288 to e30cdc7 Compare October 11, 2025 09:19

MatthewBonanni reviewed Oct 27, 2025

View reviewed changes

shen-shanshan force-pushed the mamba branch from e30cdc7 to 789c216 Compare October 28, 2025 11:58

mergify bot added the new-model Requests to new models label Oct 28, 2025

tdoublep reviewed Nov 6, 2025

View reviewed changes

shen-shanshan mentioned this pull request Nov 10, 2025

[RFC]: Remove VL Modeling Files vllm-project/vllm-ascend#4084

Open

9 tasks

shen-shanshan force-pushed the mamba branch from efa3468 to 0fa1282 Compare November 10, 2025 07:46

mergify bot added the documentation Improvements or additions to documentation label Nov 10, 2025

LucasWilkinson mentioned this pull request Nov 10, 2025

[Attention] Refactor CUDA attention backend selection logic #24794

Merged

5 tasks

shen-shanshan changed the title ~~[Mamba] Add selector for mamba attention backend and make it pluggable for other device~~ [Model][Mamba] Add selector for mamba attention backend and make it pluggable for other device Nov 11, 2025

tdoublep reviewed Nov 11, 2025

View reviewed changes

vllm/attention/selector.py Outdated Show resolved Hide resolved

vllm/model_executor/models/qwen3_next.py Outdated Show resolved Hide resolved

wangxiyuan reviewed Nov 17, 2025

View reviewed changes

shen-shanshan added 4 commits November 17, 2025 07:11

Add selector for mamba attention backend

8af7c1f

Signed-off-by: shen-shanshan <467638484@qq.com>

fix lint

59c6470

Signed-off-by: shen-shanshan <467638484@qq.com>

fix lint

d96aabc

Signed-off-by: shen-shanshan <467638484@qq.com>

add gdn_attention as a new mamba type

f65fc62

Signed-off-by: shen-shanshan <467638484@qq.com>

shen-shanshan added 6 commits November 17, 2025 07:11

add backend enum

caad557

Signed-off-by: shen-shanshan <467638484@qq.com>

update

f042764

Signed-off-by: shen-shanshan <467638484@qq.com>

update

1a6358d

Signed-off-by: shen-shanshan <467638484@qq.com>

add linear_attn_type param

0dcbf6a

Signed-off-by: shen-shanshan <467638484@qq.com>

update doc

4364d77

Signed-off-by: shen-shanshan <467638484@qq.com>

add MambaAttentionBackendEnum

386d9f9

Signed-off-by: shen-shanshan <467638484@qq.com>

shen-shanshan force-pushed the mamba branch from a7a1eb8 to 386d9f9 Compare November 17, 2025 09:07

remove redundant import

6fc1d98

Signed-off-by: shen-shanshan <467638484@qq.com>

Yikun added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 17, 2025

jikunshang approved these changes Nov 18, 2025

View reviewed changes

wangxiyuan approved these changes Nov 18, 2025

View reviewed changes

vllm/model_executor/models/qwen3_next.py Show resolved Hide resolved

Merge branch 'main' into mamba

a0da430

DarkLight1337 enabled auto-merge (squash) November 19, 2025 05:46

tdoublep self-requested a review November 19, 2025 13:57

tdoublep approved these changes Nov 19, 2025

View reviewed changes

Merge branch 'main' into mamba

301d2f7

DarkLight1337 merged commit d44e9df into vllm-project:main Nov 19, 2025
54 checks passed

Victor49152 pushed a commit to Victor49152/vllm that referenced this pull request Nov 20, 2025

[Model][Mamba] Add selector for mamba attention backend and make it p…

961b542

…luggable for other device (vllm-project#26487) Signed-off-by: shen-shanshan <467638484@qq.com>

bigPYJ1151 pushed a commit that referenced this pull request Nov 25, 2025

[Model][Mamba] Add selector for mamba attention backend and make it p…

f74c2fb

…luggable for other device (#26487) Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: jiang1.li <jiang1.li@intel.com>

bringlein pushed a commit to bringlein/vllm that referenced this pull request Nov 26, 2025

[Model][Mamba] Add selector for mamba attention backend and make it p…

2208ac6

…luggable for other device (vllm-project#26487) Signed-off-by: shen-shanshan <467638484@qq.com>

Uh oh!

[Model][Mamba] Add selector for mamba attention backend and make it pluggable for other device #26487

[Model][Mamba] Add selector for mamba attention backend and make it pluggable for other device #26487

Uh oh!

Conversation

shen-shanshan commented Oct 9, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tdoublep left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Nov 10, 2025

Uh oh!

shen-shanshan commented Nov 10, 2025

Uh oh!

LucasWilkinson commented Nov 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tdoublep commented Nov 10, 2025

Uh oh!

Uh oh!

Uh oh!

MatthewBonanni commented Nov 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shen-shanshan commented Nov 12, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shen-shanshan commented Nov 17, 2025

Uh oh!

MatthewBonanni commented Nov 17, 2025

Uh oh!

jikunshang left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

shen-shanshan commented Nov 19, 2025

Uh oh!

jikunshang commented Nov 19, 2025

Uh oh!

gcanlin commented Nov 19, 2025

Uh oh!

shen-shanshan commented Nov 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gcanlin commented Nov 19, 2025

Uh oh!

tdoublep left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

shen-shanshan commented Oct 9, 2025 •

edited by github-actions bot

Loading

LucasWilkinson commented Nov 10, 2025 •

edited

Loading

MatthewBonanni commented Nov 11, 2025 •

edited

Loading

shen-shanshan commented Nov 19, 2025 •

edited

Loading