A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 16:30 UTC
[PluggableLayer][3/N] Apply PluggableLayer to mamba layers. (#33660)
Signed-off-by: whx-sjtu <2952154980@qq.com>
whx committed
Feb 7, 2026 at 21:26 UTC
ce9b3cd3e91c2f6c901f12b038ee7797aaa5c829
3 files changed
+13
-31
vllm/model_executor/layers/mamba/mamba_mixer.py
+5
-8
@@ -13,7 +13,7 @@ from vllm.distributed.parallel_state import (
13
get_tensor_model_parallel_world_size,
14
)
15
from vllm.forward_context import ForwardContext, get_forward_context
16
-from vllm.model_executor.custom_op import CustomOp
16
+from vllm.model_executor.custom_op import PluggableLayer
17
from vllm.model_executor.layers.layernorm import RMSNorm
18
from vllm.model_executor.layers.linear import (
19
ColumnParallelLinear,
@@ -41,8 +41,8 @@ from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata
41
42
# Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
43
# --8<-- [start:mamba_mixer]
44
-@CustomOp.register("mamba_mixer")
45
-class MambaMixer(MambaBase, CustomOp):
44
+@PluggableLayer.register("mamba_mixer")
45
+class MambaMixer(MambaBase, PluggableLayer):
46
"""
47
Compute ∆, A, B, C, and D the state space parameters and compute
48
the `contextualized_states`. A, D are input independent
@@ -230,10 +230,7 @@ class MambaMixer(MambaBase, CustomOp):
230
self.prefix,
231
)
232
233
- def forward_native(self, hidden_states: torch.Tensor, output: torch.Tensor):
234
- pass
235
-
236
- def forward_cuda(self, hidden_states: torch.Tensor, output: torch.Tensor):
233
+ def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor):
234
"""
235
Run the Mamba-1 SSM pipeline.
236
@@ -528,7 +525,7 @@ def mamba_mixer(
525
) -> None:
526
forward_context: ForwardContext = get_forward_context()
527
self = forward_context.no_compile_layers[layer_name]
531
- self.forward_cuda(hidden_states=hidden_states, output=output)
528
+ self.forward_impl(hidden_states=hidden_states, output=output)
529
530
531
def mamba_mixer_fake(
vllm/model_executor/layers/mamba/mamba_mixer2.py
+3
-10
@@ -14,7 +14,7 @@ from vllm.distributed import (
14
tensor_model_parallel_all_reduce,
15
)
16
from vllm.forward_context import ForwardContext, get_forward_context
17
-from vllm.model_executor.custom_op import CustomOp
17
+from vllm.model_executor.custom_op import CustomOp, PluggableLayer
18
from vllm.model_executor.layers.linear import (
19
ColumnParallelLinear,
20
RowParallelLinear,
@@ -219,8 +219,8 @@ def mamba_v2_sharded_weight_loader(
219
220
# Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
221
# --8<-- [start:mamba_mixer2]
222
-@CustomOp.register("mamba_mixer2")
223
-class MambaMixer2(MambaBase, CustomOp):
222
+@PluggableLayer.register("mamba_mixer2")
223
+class MambaMixer2(MambaBase, PluggableLayer):
224
"""
225
Compute ∆, A, B, C, and D the state space parameters and compute
226
the `contextualized_states`. A, D are input independent
@@ -472,13 +472,6 @@ class MambaMixer2(MambaBase, CustomOp):
472
# Check if running on Blackwell (SM100+) for kernel tuning
473
self.is_blackwell = current_platform.is_device_capability_family(100)
474
475
- def forward_native(
476
- self,
477
- hidden_states: torch.Tensor,
478
- mup_vector: torch.Tensor | None = None,
479
- ):
480
- pass
481
-
475
def forward(
476
self,
477
hidden_states: torch.Tensor,
vllm/model_executor/models/plamo2.py
+5
-13
@@ -14,7 +14,7 @@ from vllm.config import VllmConfig, get_current_vllm_config
14
from vllm.distributed import divide, get_tensor_model_parallel_world_size
15
from vllm.distributed.parallel_state import get_pp_group
16
from vllm.forward_context import ForwardContext, get_forward_context
17
-from vllm.model_executor.custom_op import CustomOp
17
+from vllm.model_executor.custom_op import PluggableLayer
18
from vllm.model_executor.layers.activation import SiluAndMul
19
from vllm.model_executor.layers.attention import Attention
20
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -107,8 +107,8 @@ def is_mamba(config: Plamo2Config, i: int) -> bool:
107
# vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2
108
# transformers.models.mamba.modeling_mamba.MambaMixer
109
# --8<-- [start:plamo2_mamba_mixer]
110
-@CustomOp.register("plamo2_mamba_mixer")
111
-class Plamo2MambaMixer(MambaBase, CustomOp):
110
+@PluggableLayer.register("plamo2_mamba_mixer")
111
+class Plamo2MambaMixer(MambaBase, PluggableLayer):
112
# --8<-- [end:plamo2_mamba_mixer]
113
114
def __init__(self, vllm_config: VllmConfig, *, prefix: str = "", **kwargs) -> None:
@@ -233,14 +233,6 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
233
dt = self.dt_proj(time_step)
234
return B, C, dt
235
236
- def forward_native(
237
- self,
238
- hidden_states: torch.Tensor,
239
- output: torch.Tensor,
240
- **kwargs,
241
- ):
242
- pass
243
-
236
def forward(
237
self,
238
hidden_states: torch.Tensor,
@@ -253,7 +245,7 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
245
self.prefix,
246
)
247
256
- def forward_cuda(
248
+ def forward_impl(
249
self,
250
hidden_states: torch.Tensor,
251
output: torch.Tensor,
@@ -494,7 +486,7 @@ def plamo2_mamba_mixer(
486
) -> None:
487
forward_context: ForwardContext = get_forward_context()
488
self = forward_context.no_compile_layers[layer_name]
497
- self.forward_cuda(hidden_states=hidden_states, output=output)
489
+ self.forward_impl(hidden_states=hidden_states, output=output)
490
491
492
def plamo2_mamba_mixer_fake(