| 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 | from typing import Literal |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | import torch.nn.functional as F |
| 11 | |
| 12 | from vllm.model_executor.custom_op import CustomOp |
| 13 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 14 | |
| 15 | |
| 16 | class ConvLayerBase(CustomOp): |
| 17 | """Conv layer base class.""" |
| 18 | |
| 19 | num_dim: int |
| 20 | |
| 21 | def __init__( |
| 22 | self, |
| 23 | in_channels: int, |
| 24 | out_channels: int, |
| 25 | kernel_size: int | tuple[int, ...], |
| 26 | stride: int | tuple[int, ...] = 1, |
| 27 | padding: int | tuple[int, ...] | Literal["same", "valid"] = 0, |
| 28 | dilation: int | tuple[int, ...] = 1, |
| 29 | groups: int = 1, |
| 30 | bias: bool = True, |
| 31 | padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros", |
| 32 | *, |
| 33 | params_dtype: torch.dtype | None = None, |
| 34 | ) -> None: |
| 35 | super().__init__() |
| 36 | |
| 37 | if params_dtype is None: |
| 38 | params_dtype = torch.get_default_dtype() |
| 39 | |
| 40 | valid_padding_strings = {"same", "valid"} |
| 41 | if isinstance(padding, str) and padding not in valid_padding_strings: |
| 42 | raise ValueError( |
| 43 | f"Invalid padding string '{padding}'. " |
| 44 | f"Expected one of {valid_padding_strings}." |
| 45 | ) |
| 46 | |
| 47 | if padding == "same": |
| 48 | padding = ( |
| 49 | kernel_size // 2 |
| 50 | if isinstance(kernel_size, int) |
| 51 | else tuple(k // 2 for k in kernel_size) |
| 52 | ) |
| 53 | elif padding == "valid": |
| 54 | padding = 0 |
| 55 | |
| 56 | kernel_size = ( |
| 57 | (kernel_size,) * self.num_dim |
| 58 | if isinstance(kernel_size, int) |
| 59 | else kernel_size |
| 60 | ) |
| 61 | stride = (stride,) * self.num_dim if isinstance(stride, int) else stride |
| 62 | padding = (padding,) * self.num_dim if isinstance(padding, int) else padding |
| 63 | dilation = (dilation,) * self.num_dim if isinstance(dilation, int) else dilation |
| 64 | |
| 65 | if padding == "same" and any(s != 1 for s in stride): |
| 66 | raise ValueError("padding='same' is not supported for strided convolutions") |
| 67 | |
| 68 | self.in_channels = in_channels |
| 69 | self.out_channels = out_channels |
| 70 | self.kernel_size = kernel_size |
| 71 | self.stride = stride |
| 72 | self.padding = padding |
| 73 | self.dilation = dilation |
| 74 | self.groups = groups |
| 75 | self.padding_mode = padding_mode |
| 76 | |
| 77 | self.enable_linear = ( |
| 78 | (self.kernel_size == self.stride) |
| 79 | and not any(self.padding) |
| 80 | and self.groups == 1 |
| 81 | ) |
| 82 | self.input_size = in_channels * math.prod(self.kernel_size) |
| 83 | |
| 84 | self.weight = nn.Parameter( |
| 85 | torch.empty( |
| 86 | out_channels, |
| 87 | in_channels // groups, |
| 88 | *kernel_size, |
| 89 | dtype=params_dtype, |
| 90 | ), |
| 91 | ) |
| 92 | |
| 93 | if bias: |
| 94 | self.bias = nn.Parameter(torch.empty(self.out_channels, dtype=params_dtype)) |
| 95 | else: |
| 96 | self.register_parameter("bias", None) |
| 97 | |
| 98 | def extra_repr(self) -> str: |
| 99 | s = f"in_channels={self.in_channels}, " |
| 100 | s += f"out_channels={self.out_channels}, " |
| 101 | s += f"kernel_size={self.kernel_size}, " |
| 102 | s += f"stride={self.stride}, " |
| 103 | s += f"padding={self.padding}, " |
| 104 | s += f"bias={self.bias is not None}" |
| 105 | return s |
| 106 | |
| 107 | |
| 108 | # --8<-- [start:conv2d] |
| 109 | @CustomOp.register("conv2d") |
| 110 | class Conv2dLayer(ConvLayerBase): |
| 111 | """Conv layer with Conv2d.""" |
| 112 | |
| 113 | # --8<-- [end:conv2d] |
| 114 | |
| 115 | num_dim = 2 |
| 116 | |
| 117 | def _forward_mulmat(self, x: torch.Tensor) -> torch.Tensor: |
| 118 | assert x.dim() == 4 |
| 119 | B, C, H, W = x.shape |
| 120 | K1, K2 = self.kernel_size |
| 121 | H, W = H // K1, W // K2 |
| 122 | x = x.unfold(2, K1, K1).unfold(3, K2, K2) |
| 123 | x = x.permute(0, 2, 3, 1, 4, 5).reshape(-1, self.input_size) |
| 124 | x = F.linear( |
| 125 | x, |
| 126 | self.weight.view(self.out_channels, self.input_size), |
| 127 | self.bias, |
| 128 | ) |
| 129 | x = x.view(B, H, W, self.out_channels).permute(0, 3, 1, 2) |
| 130 | return x |
| 131 | |
| 132 | def _forward_conv(self, x: torch.Tensor) -> torch.Tensor: |
| 133 | assert x.dim() == 4 |
| 134 | x = F.conv2d( |
| 135 | x, |
| 136 | self.weight, |
| 137 | self.bias, |
| 138 | stride=self.stride, |
| 139 | padding=self.padding, |
| 140 | dilation=self.dilation, |
| 141 | groups=self.groups, |
| 142 | ) |
| 143 | return x |
| 144 | |
| 145 | def forward_native(self, x: torch.Tensor) -> torch.Tensor: |
| 146 | """Expected input shape: (batch_size, in_channels, height, width).""" |
| 147 | assert x.dim() == 4 |
| 148 | if self.enable_linear: |
| 149 | return self._forward_mulmat(x) |
| 150 | else: |
| 151 | return self._forward_conv(x) |
| 152 | |
| 153 | def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: |
| 154 | # By default, we use CUDNN's convolution ops with optimization. |
| 155 | return self._forward_conv(x) |
| 156 | |
| 157 | |
| 158 | # --8<-- [start:conv3d] |
| 159 | @CustomOp.register("conv3d") |
| 160 | class Conv3dLayer(ConvLayerBase): |
| 161 | """Conv layer with Conv3d.""" |
| 162 | |
| 163 | # --8<-- [end:conv3d] |
| 164 | |
| 165 | num_dim = 3 |
| 166 | |
| 167 | def _forward_mulmat(self, x: torch.Tensor) -> torch.Tensor: |
| 168 | assert x.dim() == 5 |
| 169 | B, C, T, H, W = x.shape |
| 170 | K1, K2, K3 = self.kernel_size |
| 171 | T, H, W = T // K1, H // K2, W // K3 |
| 172 | x = x.unfold(2, K1, K1).unfold(3, K2, K2).unfold(4, K3, K3) |
| 173 | x = x.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(-1, self.input_size) |
| 174 | x = F.linear( |
| 175 | x, |
| 176 | self.weight.view(self.out_channels, self.input_size), |
| 177 | self.bias, |
| 178 | ) |
| 179 | x = x.view(B, T, H, W, self.out_channels).permute(0, 4, 1, 2, 3) |
| 180 | return x |
| 181 | |
| 182 | def _forward_conv(self, x: torch.Tensor) -> torch.Tensor: |
| 183 | assert x.dim() == 5 |
| 184 | x = F.conv3d( |
| 185 | x, |
| 186 | self.weight, |
| 187 | self.bias, |
| 188 | stride=self.stride, |
| 189 | padding=self.padding, |
| 190 | dilation=self.dilation, |
| 191 | groups=self.groups, |
| 192 | ) |
| 193 | return x |
| 194 | |
| 195 | def forward_native(self, x: torch.Tensor) -> torch.Tensor: |
| 196 | """Expected input shape: (batch_size, in_channels, time, height, width).""" |
| 197 | if self.enable_linear: |
| 198 | return self._forward_mulmat(x) |
| 199 | else: |
| 200 | return self._forward_conv(x) |
| 201 | |
| 202 | def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: |
| 203 | # PyTorch 2.9.0+ disabled CUDNN's Conv3D, which caused a |
| 204 | # significant performance regression. |
| 205 | # See: https://github.com/vllm-project/vllm/issues/27406 |
| 206 | # and https://github.com/pytorch/pytorch/issues/166122 |
| 207 | # and https://github.com/huggingface/transformers/pull/45041 |
| 208 | # By default, we use CUDNN's convolution ops with optimization. |
| 209 | if self.enable_linear and is_torch_equal_or_newer("2.9.0"): |
| 210 | return self._forward_mulmat(x) |
| 211 | return self._forward_conv(x) |