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 13, 05:45 UTC
[Model][MM] Extract conv layer as CustomOp (#28455)
Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Shanshan Shen committed
Nov 14, 2025 at 19:16 UTC
41b92f7d38d3f056004991c026f6a24846755ef4
8 files changed
+277
-66
vllm/model_executor/layers/conv.py
new
+236
@@ -0,0 +1,236 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Conv Layer Class."""
4
+
5
+import math
6
+
7
+import torch
8
+import torch.nn as nn
9
+import torch.nn.functional as F
10
+
11
+from vllm.model_executor.custom_op import CustomOp
12
+from vllm.utils.torch_utils import is_torch_equal
13
+
14
+
15
+class ConvLayerBase(CustomOp):
16
+ """Conv layer base class."""
17
+
18
+ num_dim: int
19
+
20
+ def __init__(
21
+ self,
22
+ in_channels: int,
23
+ out_channels: int,
24
+ kernel_size: int | tuple[int, ...],
25
+ stride: int | tuple[int, ...] = 1,
26
+ padding: int | tuple[int, ...] = 0,
27
+ dilation: int | tuple[int, ...] = 1,
28
+ groups: int = 1,
29
+ bias: bool = True,
30
+ padding_mode: str = "zeros",
31
+ *,
32
+ params_dtype: torch.dtype | None = None,
33
+ ) -> None:
34
+ super().__init__()
35
+
36
+ if params_dtype is None:
37
+ params_dtype = torch.get_default_dtype()
38
+
39
+ kernel_size = (
40
+ (kernel_size,) * self.num_dim
41
+ if isinstance(kernel_size, int)
42
+ else kernel_size
43
+ )
44
+ stride = (stride,) * self.num_dim if isinstance(stride, int) else stride
45
+ padding = (padding,) * self.num_dim if isinstance(padding, int) else padding
46
+ dilation = (dilation,) * self.num_dim if isinstance(dilation, int) else dilation
47
+
48
+ self.in_channels = in_channels
49
+ self.out_channels = out_channels
50
+ self.kernel_size = kernel_size
51
+ self.stride = stride
52
+ self.padding = padding
53
+ self.dilation = dilation
54
+ self.groups = groups
55
+ self.padding_mode = padding_mode
56
+
57
+ self.enable_linear = (
58
+ (self.kernel_size == self.stride)
59
+ and not any(self.padding)
60
+ and self.groups == 1
61
+ )
62
+ self.input_size = in_channels * math.prod(self.kernel_size)
63
+
64
+ self.weight = nn.Parameter(
65
+ torch.empty(
66
+ out_channels,
67
+ in_channels // groups,
68
+ *kernel_size,
69
+ dtype=params_dtype,
70
+ ),
71
+ )
72
+
73
+ if bias:
74
+ self.bias = nn.Parameter(torch.empty(self.out_channels, dtype=params_dtype))
75
+ else:
76
+ self.register_parameter("bias", None)
77
+
78
+ def extra_repr(self) -> str:
79
+ s = f"in_channels={self.in_channels}, "
80
+ s += f"out_channels={self.out_channels}, "
81
+ s += f"kernel_size={self.kernel_size}, "
82
+ s += f"stride={self.stride}, "
83
+ s += f"padding={self.padding}, "
84
+ s += f"bias={self.bias is not None}"
85
+ return s
86
+
87
+
88
+@CustomOp.register("conv2d")
89
+class Conv2dLayer(ConvLayerBase):
90
+ """Conv layer with Conv2d."""
91
+
92
+ num_dim = 2
93
+
94
+ def _forward_mulmat(self, x: torch.Tensor) -> torch.Tensor:
95
+ assert x.dim() == 4
96
+ B, C, H, W = x.shape
97
+ K1, K2 = self.kernel_size
98
+ H, W = H // K1, W // K2
99
+ x = x.unfold(2, K1, K1).unfold(3, K2, K2)
100
+ x = x.permute(0, 2, 3, 1, 4, 5).reshape(-1, self.input_size)
101
+ x = F.linear(
102
+ x,
103
+ self.weight.view(self.out_channels, self.input_size),
104
+ self.bias,
105
+ )
106
+ x = x.view(B, H, W, self.out_channels).permute(0, 3, 1, 2)
107
+ return x
108
+
109
+ def _forward_conv(self, x: torch.Tensor) -> torch.Tensor:
110
+ assert x.dim() == 4
111
+ x = F.conv2d(
112
+ x,
113
+ self.weight,
114
+ self.bias,
115
+ stride=self.stride,
116
+ padding=self.padding,
117
+ dilation=self.dilation,
118
+ groups=self.groups,
119
+ )
120
+ return x
121
+
122
+ def forward_native(self, x: torch.Tensor) -> torch.Tensor:
123
+ """Expected input shape: (batch_size, in_channels, height, width)"""
124
+ assert x.dim() == 4
125
+ if self.enable_linear:
126
+ return self._forward_mulmat(x)
127
+ else:
128
+ return self._forward_conv(x)
129
+
130
+ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
131
+ # By default, we use CUDNN's convolution ops with optimization.
132
+ return self._forward_conv(x)
133
+
134
+
135
+class CausalConv2dLayer(Conv2dLayer):
136
+ """
137
+ A causal version of nn.Conv2d where each location in the 2D matrix would
138
+ have no access to locations on its right or down
139
+ All arguments are the same as nn.Conv2d except padding which should be
140
+ set as None
141
+ """
142
+
143
+ def __init__(
144
+ self,
145
+ in_channels: int,
146
+ out_channels: int,
147
+ kernel_size: int,
148
+ stride: int,
149
+ padding: int = 0,
150
+ dilation: int = 1,
151
+ groups: int = 1,
152
+ bias: bool = True,
153
+ padding_mode: str = "zeros",
154
+ *,
155
+ params_dtype: torch.dtype | None = None,
156
+ ) -> None:
157
+ if padding is not None:
158
+ raise ValueError(
159
+ "Argument padding should be set to None for CausalConv2dLayer."
160
+ )
161
+ self._left_padding: int = kernel_size - 1
162
+ self._right_padding: int = stride - 1
163
+ padding = 0
164
+
165
+ super().__init__(
166
+ in_channels,
167
+ out_channels,
168
+ kernel_size,
169
+ stride,
170
+ padding,
171
+ dilation,
172
+ groups,
173
+ bias,
174
+ padding_mode,
175
+ params_dtype=params_dtype,
176
+ )
177
+
178
+ def forward(
179
+ self,
180
+ x: torch.Tensor,
181
+ ) -> torch.Tensor:
182
+ x = F.pad(x, pad=(self._left_padding, self._right_padding, 0, 0))
183
+ x = super().forward(x)
184
+ return x
185
+
186
+
187
+@CustomOp.register("conv3d")
188
+class Conv3dLayer(ConvLayerBase):
189
+ """Conv layer with Conv3d."""
190
+
191
+ num_dim = 3
192
+
193
+ def _forward_mulmat(self, x: torch.Tensor) -> torch.Tensor:
194
+ assert x.dim() == 5
195
+ B, C, T, H, W = x.shape
196
+ K1, K2, K3 = self.kernel_size
197
+ T, H, W = T // K1, H // K2, W // K3
198
+ x = x.unfold(2, K1, K1).unfold(3, K2, K2).unfold(4, K3, K3)
199
+ x = x.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(-1, self.input_size)
200
+ x = F.linear(
201
+ x,
202
+ self.weight.view(self.out_channels, self.input_size),
203
+ self.bias,
204
+ )
205
+ x = x.view(B, T, H, W, self.out_channels).permute(0, 4, 1, 2, 3)
206
+ return x
207
+
208
+ def _forward_conv(self, x: torch.Tensor) -> torch.Tensor:
209
+ assert x.dim() == 5
210
+ x = F.conv3d(
211
+ x,
212
+ self.weight,
213
+ self.bias,
214
+ stride=self.stride,
215
+ padding=self.padding,
216
+ dilation=self.dilation,
217
+ groups=self.groups,
218
+ )
219
+ return x
220
+
221
+ def forward_native(self, x: torch.Tensor) -> torch.Tensor:
222
+ """Expected input shape: (batch_size, in_channels, time, height, width)"""
223
+ if self.enable_linear:
224
+ return self._forward_mulmat(x)
225
+ else:
226
+ return self._forward_conv(x)
227
+
228
+ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
229
+ # PyTorch2.9.0 disabled CUDNN's Conv3D, which caused a
230
+ # significant performance regression.
231
+ # See: https://github.com/vllm-project/vllm/issues/27406
232
+ # and https://github.com/pytorch/pytorch/issues/166122
233
+ # By default, we use CUDNN's convolution ops with optimization.
234
+ if self.enable_linear and is_torch_equal("2.9.0"):
235
+ return self._forward_mulmat(x)
236
+ return self._forward_conv(x)
vllm/model_executor/models/clip.py
+2
-1
@@ -20,6 +20,7 @@ from vllm.config import VllmConfig
20
from vllm.config.multimodal import BaseDummyOptions
21
from vllm.distributed import divide, get_tensor_model_parallel_world_size
22
from vllm.model_executor.layers.activation import get_act_fn
23
+from vllm.model_executor.layers.conv import Conv2dLayer
24
from vllm.model_executor.layers.linear import (
25
ColumnParallelLinear,
26
QKVParallelLinear,
@@ -315,7 +316,7 @@ class CLIPVisionEmbeddings(nn.Module):
316
317
self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
318
318
- self.patch_embedding = nn.Conv2d(
319
+ self.patch_embedding = Conv2dLayer(
320
in_channels=config.num_channels,
321
out_channels=self.embed_dim,
322
kernel_size=self.patch_size,
vllm/model_executor/models/glm4_1v.py
+8
-9
@@ -56,12 +56,12 @@ from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
56
from vllm.distributed import get_tensor_model_parallel_world_size, parallel_state
57
from vllm.distributed import utils as dist_utils
58
from vllm.logger import init_logger
59
+from vllm.model_executor.layers.conv import Conv3dLayer
60
from vllm.model_executor.layers.layernorm import RMSNorm
61
from vllm.model_executor.layers.linear import (
62
ColumnParallelLinear,
63
MergedColumnParallelLinear,
64
QKVParallelLinear,
64
- ReplicatedLinear,
65
RowParallelLinear,
66
)
67
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -103,7 +103,6 @@ from .utils import (
103
maybe_prefix,
104
)
105
from .vision import (
106
- conv3d_to_linear_weight,
106
get_vit_attn_backend,
107
run_dp_sharded_mrope_vision_model,
108
)
@@ -486,15 +485,18 @@ class Glm4vVisionPatchEmbed(nn.Module):
485
self.hidden_size = hidden_size
486
487
kernel_size = (temporal_patch_size, patch_size, patch_size)
489
- self.proj = ReplicatedLinear(
490
- in_channels * math.prod(kernel_size),
488
+ self.proj = Conv3dLayer(
489
+ in_channels,
490
hidden_size,
491
+ kernel_size=kernel_size,
492
+ stride=kernel_size,
493
bias=True,
493
- return_bias=False,
494
)
495
496
def forward(self, x: torch.Tensor) -> torch.Tensor:
497
- x = self.proj(x)
497
+ L, C = x.shape
498
+ x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size)
499
+ x = self.proj(x).view(L, self.hidden_size)
500
return x
501
502
@@ -893,9 +895,6 @@ class Glm4vVisionTransformer(nn.Module):
895
loaded_params: set[str] = set()
896
897
for name, loaded_weight in weights:
896
- if name.endswith("patch_embed.proj.weight"):
897
- loaded_weight = conv3d_to_linear_weight(loaded_weight)
898
-
898
for param_name, weight_name, shard_id in stacked_params_mapping:
899
if weight_name not in name:
900
continue
vllm/model_executor/models/qwen2_5_vl.py
+8
-10
@@ -26,7 +26,6 @@
26
# limitations under the License.
27
"""Inference-only Qwen2.5-VL model compatible with HuggingFace weights."""
28
29
-import math
29
from collections.abc import Callable, Iterable, Mapping, Sequence
30
from functools import lru_cache, partial
31
from typing import Annotated, Any, Literal, TypeAlias
@@ -56,12 +55,12 @@ from vllm.distributed import utils as dist_utils
55
from vllm.forward_context import set_forward_context
56
from vllm.logger import init_logger
57
from vllm.model_executor.layers.activation import get_act_and_mul_fn
58
+from vllm.model_executor.layers.conv import Conv3dLayer
59
from vllm.model_executor.layers.layernorm import RMSNorm
60
from vllm.model_executor.layers.linear import (
61
ColumnParallelLinear,
62
MergedColumnParallelLinear,
63
QKVParallelLinear,
64
- ReplicatedLinear,
64
RowParallelLinear,
65
)
66
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -110,7 +109,6 @@ from .utils import (
109
maybe_prefix,
110
)
111
from .vision import (
113
- conv3d_to_linear_weight,
112
get_vit_attn_backend,
113
run_dp_sharded_mrope_vision_model,
114
)
@@ -525,15 +523,18 @@ class Qwen2_5_VisionPatchEmbed(nn.Module):
523
self.hidden_size = hidden_size
524
525
kernel_size = (temporal_patch_size, patch_size, patch_size)
528
- self.proj = ReplicatedLinear(
529
- in_channels * math.prod(kernel_size),
526
+ self.proj = Conv3dLayer(
527
+ in_channels,
528
hidden_size,
529
+ kernel_size=kernel_size,
530
+ stride=kernel_size,
531
bias=False,
532
- return_bias=False,
532
)
533
534
def forward(self, x: torch.Tensor) -> torch.Tensor:
536
- x = self.proj(x)
535
+ L, C = x.shape
536
+ x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size)
537
+ x = self.proj(x).view(L, self.hidden_size)
538
return x
539
540
@@ -957,9 +958,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
958
loaded_params: set[str] = set()
959
960
for name, loaded_weight in weights:
960
- if name.endswith("patch_embed.proj.weight"):
961
- loaded_weight = conv3d_to_linear_weight(loaded_weight)
962
-
961
for param_name, weight_name, shard_id in stacked_params_mapping:
962
if weight_name not in name:
963
continue
vllm/model_executor/models/qwen2_vl.py
+8
-10
@@ -25,7 +25,6 @@
25
# limitations under the License.
26
"""Inference-only Qwen2-VL model compatible with HuggingFace weights."""
27
28
-import math
28
from collections.abc import Callable, Iterable, Mapping, Sequence
29
from functools import partial
30
from typing import Annotated, Any, Literal, TypeAlias
@@ -54,9 +53,9 @@ from vllm.distributed import parallel_state
53
from vllm.distributed import utils as dist_utils
54
from vllm.logger import init_logger
55
from vllm.model_executor.layers.activation import QuickGELU
56
+from vllm.model_executor.layers.conv import Conv3dLayer
57
from vllm.model_executor.layers.linear import (
58
ColumnParallelLinear,
59
- ReplicatedLinear,
59
RowParallelLinear,
60
)
61
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -107,7 +106,6 @@ from .utils import (
106
maybe_prefix,
107
)
108
from .vision import (
110
- conv3d_to_linear_weight,
109
get_vit_attn_backend,
110
run_dp_sharded_mrope_vision_model,
111
)
@@ -566,15 +564,18 @@ class Qwen2VisionPatchEmbed(nn.Module):
564
self.embed_dim = embed_dim
565
566
kernel_size = (temporal_patch_size, patch_size, patch_size)
569
- self.proj = ReplicatedLinear(
570
- in_channels * math.prod(kernel_size),
567
+ self.proj = Conv3dLayer(
568
+ in_channels,
569
embed_dim,
570
+ kernel_size=kernel_size,
571
+ stride=kernel_size,
572
bias=False,
573
- return_bias=False,
573
)
574
575
def forward(self, x: torch.Tensor) -> torch.Tensor:
577
- x = self.proj(x)
576
+ L, C = x.shape
577
+ x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size)
578
+ x = self.proj(x).view(L, self.embed_dim)
579
return x
580
581
@@ -844,9 +845,6 @@ class Qwen2VisionTransformer(nn.Module):
845
loaded_params: set[str] = set()
846
847
for name, loaded_weight in weights:
847
- if name.endswith("patch_embed.proj.weight"):
848
- loaded_weight = conv3d_to_linear_weight(loaded_weight)
849
-
848
for param_name, weight_name, shard_id in stacked_params_mapping:
849
if weight_name not in name:
850
continue
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+7
-10
@@ -22,7 +22,6 @@
22
# limitations under the License.
23
"""Inference-only Qwen3-Omni-Moe model (thinker part)."""
24
25
-import math
25
from collections.abc import Callable, Iterable, Mapping, Sequence
26
from functools import partial
27
from typing import Any
@@ -54,9 +53,9 @@ from vllm.config import VllmConfig
53
from vllm.distributed import get_pp_group
54
from vllm.logger import init_logger
55
from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
56
+from vllm.model_executor.layers.conv import Conv3dLayer
57
from vllm.model_executor.layers.linear import (
58
ColumnParallelLinear,
59
- ReplicatedLinear,
59
RowParallelLinear,
60
)
61
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -102,7 +101,6 @@ from .utils import (
101
maybe_prefix,
102
)
103
from .vision import (
105
- conv3d_to_linear_weight,
104
get_llm_pos_ids_for_vision,
105
get_vit_attn_backend,
106
)
@@ -138,16 +136,18 @@ class Qwen3_VisionPatchEmbed(nn.Module):
136
self.hidden_size = hidden_size
137
138
kernel_size = (temporal_patch_size, patch_size, patch_size)
141
- self.proj = ReplicatedLinear(
142
- in_channels * math.prod(kernel_size),
139
+ self.proj = Conv3dLayer(
140
+ in_channels,
141
hidden_size,
142
+ kernel_size=kernel_size,
143
+ stride=kernel_size,
144
bias=True,
145
- return_bias=False,
145
)
146
147
def forward(self, x: torch.Tensor) -> torch.Tensor:
148
L, C = x.shape
150
- x = self.proj(x)
149
+ x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size)
150
+ x = self.proj(x).view(L, self.hidden_size)
151
return x
152
153
@@ -566,9 +566,6 @@ class Qwen3Omni_VisionTransformer(nn.Module):
566
loaded_params: set[str] = set()
567
568
for name, loaded_weight in weights:
569
- if name.endswith("patch_embed.proj.weight"):
570
- loaded_weight = conv3d_to_linear_weight(loaded_weight)
571
-
569
for param_name, weight_name, shard_id in stacked_params_mapping:
570
if weight_name not in name:
571
continue
vllm/model_executor/models/qwen3_vl.py
+8
-10
@@ -24,7 +24,6 @@
24
# limitations under the License.
25
"""Inference-only Qwen3VL model compatible with HuggingFace weights."""
26
27
-import math
27
from collections.abc import Callable, Iterable, Mapping, Sequence
28
from functools import partial
29
from itertools import islice
@@ -57,9 +56,9 @@ from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
56
from vllm.distributed import get_pp_group
57
from vllm.logger import init_logger
58
from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
59
+from vllm.model_executor.layers.conv import Conv3dLayer
60
from vllm.model_executor.layers.linear import (
61
ColumnParallelLinear,
62
- ReplicatedLinear,
62
RowParallelLinear,
63
)
64
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -114,7 +113,6 @@ from .utils import (
113
maybe_prefix,
114
)
115
from .vision import (
117
- conv3d_to_linear_weight,
116
get_vit_attn_backend,
117
run_dp_sharded_mrope_vision_model,
118
)
@@ -139,15 +137,18 @@ class Qwen3_VisionPatchEmbed(nn.Module):
137
self.hidden_size = hidden_size
138
139
kernel_size = (temporal_patch_size, patch_size, patch_size)
142
- self.proj = ReplicatedLinear(
143
- in_channels * math.prod(kernel_size),
140
+ self.proj = Conv3dLayer(
141
+ in_channels,
142
hidden_size,
143
+ kernel_size=kernel_size,
144
+ stride=kernel_size,
145
bias=True,
146
- return_bias=False,
146
)
147
148
def forward(self, x: torch.Tensor) -> torch.Tensor:
150
- x = self.proj(x)
149
+ L, C = x.shape
150
+ x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size)
151
+ x = self.proj(x).view(L, self.hidden_size)
152
return x
153
154
@@ -579,9 +580,6 @@ class Qwen3_VisionTransformer(nn.Module):
580
loaded_params: set[str] = set()
581
582
for name, loaded_weight in weights:
582
- if name.endswith("patch_embed.proj.weight"):
583
- loaded_weight = conv3d_to_linear_weight(loaded_weight)
584
-
583
for param_name, weight_name, shard_id in stacked_params_mapping:
584
if weight_name not in name:
585
continue
vllm/model_executor/models/vision.py
-16
@@ -550,19 +550,3 @@ def get_llm_pos_ids_for_vision(
550
llm_pos_ids_list.append(_llm_pos_ids + start_idx)
551
llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
552
return llm_pos_ids
553
-
554
-
555
-# Due to a performance regression with Conv3D in PyTorch2.9, we reshape
556
-# Conv3D weights to Linear weights for better performance.
557
-# See: https://github.com/vllm-project/vllm/issues/27406
558
-# and https://github.com/pytorch/pytorch/issues/166122
559
-# FIXME(Isotr0py): Revert the PR introduces this workaround
560
-# (https://github.com/vllm-project/vllm/pull/27418),
561
-# once the performance issue is resolved in PyTorch.
562
-def conv3d_to_linear_weight(conv3d_weight: torch.Tensor) -> torch.Tensor:
563
- """
564
- Reshape Conv3D weight to Linear weight. Only work when kernel_size==stride.
565
- """
566
- out_channels, in_channels, kt, kh, kw = conv3d_weight.shape
567
- linear_weight = conv3d_weight.reshape(out_channels, in_channels * kt * kh * kw)
568
- return linear_weight