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 10, 21:30 UTC
Replace `nn.ConvNd` with vLLM's `ConvNdLayer` for Transformers modeling backend (#31498)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Dec 29, 2025 at 16:20 UTC
e37e7349e693e50995b6df5b9b6be52729f44cc6
2 files changed
+43
vllm/model_executor/models/transformers/base.py
+3
@@ -47,6 +47,7 @@ from vllm.model_executor.models.transformers.utils import (
47
get_feature_request_tip,
48
init_on_device_without_buffers,
49
log_replacement,
50
+ replace_conv_class,
51
replace_linear_class,
52
replace_rms_norm_class,
53
)
@@ -314,6 +315,8 @@ class Base(
315
new_module = replace_linear_class(
316
child_module, style, self.quant_config, prefix=qual_name
317
)
318
+ elif isinstance(child_module, (nn.Conv2d, nn.Conv3d)):
319
+ new_module = replace_conv_class(child_module)
320
elif child_module.__class__.__name__.endswith("RMSNorm"):
321
new_module = replace_rms_norm_class(
322
child_module, self.text_config.hidden_size
vllm/model_executor/models/transformers/utils.py
+40
@@ -25,6 +25,7 @@ from torch import nn
25
26
from vllm.config.utils import getattr_iter
27
from vllm.logger import init_logger
28
+from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
29
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
30
from vllm.model_executor.layers.linear import (
31
ColumnParallelLinear,
@@ -136,6 +137,45 @@ def replace_linear_class(
137
)
138
139
140
+TorchConv = nn.Conv2d | nn.Conv3d
141
+VllmConv = Conv2dLayer | Conv3dLayer
142
+
143
+
144
+def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv:
145
+ """Replace a Transformers Conv2d/Conv3d with vLLM's Conv2d/Conv3d.
146
+
147
+ Args:
148
+ conv: `nn.Conv2d` or `nn.Conv3d` to be replaced.
149
+ Returns:
150
+ The new `Conv2dLayer` or `Conv3dLayer`. If the conv module is not supported,
151
+ returns the original conv module.
152
+ """
153
+ # vLLM does not handle non-zero padding modes
154
+ if conv.padding_mode != "zeros":
155
+ return conv
156
+
157
+ vllm_conv_cls = {
158
+ nn.Conv2d: Conv2dLayer,
159
+ nn.Conv3d: Conv3dLayer,
160
+ }.get(type(conv))
161
+
162
+ if vllm_conv_cls is None:
163
+ return conv
164
+
165
+ return vllm_conv_cls(
166
+ in_channels=conv.in_channels,
167
+ out_channels=conv.out_channels,
168
+ kernel_size=conv.kernel_size,
169
+ stride=conv.stride,
170
+ padding=conv.padding,
171
+ dilation=conv.dilation,
172
+ groups=conv.groups,
173
+ bias=conv.bias is not None,
174
+ padding_mode=conv.padding_mode,
175
+ params_dtype=conv.weight.dtype,
176
+ )
177
+
178
+
179
def replace_rms_norm_class(rms_norm: nn.Module, hidden_size: int) -> RMSNorm:
180
"""Replace a Transformers RMSNorm with vLLM's RMSNorm.
181