[Attention] Mamba attention module refactor - Final part (#44857)
wangxiyuan committed
Aug 8, 2026 at 02:27 UTC
4eccf906ca8778bc02cb7549bab38e8b53cd9a90
1 file changed
+3
-3
vllm/model_executor/layers/mamba/short_conv.py
+3
-3
@@ -7,7 +7,7 @@ import torch
7
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
8
from vllm.distributed import get_tensor_model_parallel_world_size
9
from vllm.forward_context import ForwardContext, get_forward_context
10
-from vllm.model_executor.custom_op import CustomOp
10
+from vllm.model_executor.custom_op import PluggableLayer
11
from vllm.model_executor.layers.linear import (
12
ColumnParallelLinear,
13
MergedColumnParallelLinear,
@@ -32,8 +32,8 @@ from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadat
32
33
34
# --8<-- [start:short_conv]
35
-@CustomOp.register("short_conv")
36
-class ShortConv(MambaBase, CustomOp):
35
+@PluggableLayer.register("short_conv")
36
+class ShortConv(MambaBase, PluggableLayer):
37
# --8<-- [end:short_conv]
38
39
def __init__(