diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index 7a11a4f5475..64fd14c3cfd 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, causal_conv1d_update, ) +from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backend import AttentionMetadata @@ -42,6 +43,7 @@ class ShortConv(MambaBase, CustomOp): layer_idx: int, model_config: ModelConfig | None = None, cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() @@ -67,12 +69,14 @@ class ShortConv(MambaBase, CustomOp): input_size=dim, output_sizes=[dim] * 3, bias=self.bias, + quant_config=quant_config, prefix=f"{prefix}.in_proj", ) self.out_proj = RowParallelLinear( input_size=dim, output_size=dim, bias=self.bias, + quant_config=quant_config, prefix=f"{prefix}.out_proj", ) diff --git a/vllm/model_executor/models/lfm2.py b/vllm/model_executor/models/lfm2.py index 601bd22b2fa..a8387a33641 100644 --- a/vllm/model_executor/models/lfm2.py +++ b/vllm/model_executor/models/lfm2.py @@ -257,6 +257,7 @@ class Lfm2ShortConvDecoderLayer(nn.Module): layer_idx=layer_idx, model_config=model_config, cache_config=cache_config, + quant_config=quant_config, prefix=f"{prefix}.conv", ) diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index 698f1fb72ef..abdf12c7651 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -351,6 +351,7 @@ class Lfm2MoeShortConvDecoderLayer(nn.Module): layer_idx=layer_idx, model_config=model_config, cache_config=cache_config, + quant_config=quant_config, prefix=f"{prefix}.conv", )