From 4eccf906ca8778bc02cb7549bab38e8b53cd9a90 Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Sat, 8 Aug 2026 02:27:10 +0800 Subject: [PATCH] [Attention] Mamba attention module refactor - Final part (#44857) --- vllm/model_executor/layers/mamba/short_conv.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index 64fd14c3cf..f5548cc7de 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -7,7 +7,7 @@ import torch from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import ForwardContext, get_forward_context -from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, @@ -32,8 +32,8 @@ from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadat # --8<-- [start:short_conv] -@CustomOp.register("short_conv") -class ShortConv(MambaBase, CustomOp): +@PluggableLayer.register("short_conv") +class ShortConv(MambaBase, PluggableLayer): # --8<-- [end:short_conv] def __init__(