[Model][MM] Extract conv layer as CustomOp

[shen-shanshan]

Purpose

Extract ConvLayer as CustomOp for better management and extensibility (especially for plugin devices).

Find more details at the comments below.

Test Plan

I have tested this PR on ascend platform and it worked well, even getting a better performance.

Test Result

See: vllm-project/vllm-ascend#4198.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

The pull request successfully extracts the VisionPatchEmbed layer as a CustomOp, improving modularity and extensibility across various vision models. The refactoring in ernie45_vl.py, glm4_1v.py, moonvit.py, qwen2_5_vl.py, qwen2_vl.py, qwen3_omni_moe_thinker.py, and qwen3_vl.py correctly utilizes the new factory function get_vision_patch_embed to instantiate the appropriate patch embedding layers. This change streamlines the codebase by centralizing the logic for creating vision patch embedding layers.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[Isotr0py]

BTW, I think we can finish the convolution layer implementation in this PR and only modify several typical model's implementation. Then migrate remained models in following PR.

Otherwise you have to modify all models' implementation in one time when refine layer implementation each time, which is quite inefficient and easy to make mistake.

[Isotr0py]

Suggested change

	self.proj = nn.Conv2d(
	in_channels=in_channels,
	out_channels=out_channels,
	kernel_size=kernel_size,
	stride=stride,
	padding=padding,
	dilation=dilation,
	groups=groups,
	bias=bias,
	padding_mode=padding_mode,
	)
	self.weight = nn.Parameter(
	torch.empty(
	(in_channels, out_channels // groups, *kernel_size),
	)
	)

This will create a nn.Conv2d as a submodule inside Conv2dLayer, let's initialize weight and bias in Conv2dLayer directly.

[Isotr0py]

Suggested change

	x = self.proj(x)
	return x
	return F.conv2d(
	input, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups
	)

Then we can control the dispatch at ops-level like F.conv2d etc.

[Isotr0py]

Suggested change

	class LinearConvLayer(ConvLayerBase):
	"""Conv layer with linear module."""
	def __init__(
	self,
	input_size: int,
	output_size: int,
	bias: bool = True,
	skip_bias_add: bool = False,
	params_dtype: torch.dtype | None = None,
	quant_config: QuantizationConfig | None = None,
	prefix: str = "",
	*,
	return_bias: bool = True,
	disable_tp: bool = False,
	) -> None:
	class LinearConvLayer(ConvLayerBase):
	"""Conv layer with linear module."""
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	kernel_size: int,
	stride: int = 1,
	padding: int | tuple | str = 0,
	dilation: int | tuple = 1,
	groups: int = 1,
	bias: bool = True,
	padding_mode: str = "zeros",
	) -> None:

Hmmm, we only use linear to replace convolution when kernel_size == stride, which is a special case for nn.Conv2d/nn.Conv3d. I think we can implement the linearized convolution inside Conv2dLayer, if not we should at least keep args consistent with ConvLayer. WDYT?

[shen-shanshan]

Hmmm, we only use linear to replace convolution when kernel_size == stride, which is a special case for nn.Conv2d/nn.Conv3d. I think we can implement the linearized convolution inside Conv2dLayer, if not we should at least keep args consistent with ConvLayer. WDYT?

Oh, that's a good suggestion, maybe we can merge these two types into Conv2dLayer and add some checks to dispatch the forward.

[Isotr0py]

We only replace conv3d with linear currently because of a performance regression in torch2.9 (see #27406).

BTW, I remember that linear is faster than convolution for a point-wise convolution computation in early pytorch, perhaps other platforms (like CPU platform) can benefit from linearized conv2d/3d because they have custom optimized gemm ops.

I think we can implement both conv layer like this to automatically linearize conv ops:

class Conv2dLayer(CustomOp):
    """Conv layer with Conv2d."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int | tuple,
        stride: int | tuple | None,
        padding: int | tuple | str | None,
        dilation: int | tuple | None,
        groups: int | None,
        bias: bool | None,
        padding_mode: str | None,
    ) -> None:
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        ...

        self.can_linearize = (self.kernel_size == self.stride and not self.padding)

        if self.can_linearize:
            self.weight = nn.Parameter(
                torch.empty(out_channels, in_channels * math.prod(kernel_size))
            )
        else:
            self.weight = nn.Parameter(
                torch.empty(out_channels, in_channels, *kernel_size)
            )
        ...
    
    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
        .... # <- do some reshape here
        if self.can_linearize:
            return F.linear(x, self.weight, self.bias)
        else:
            return F.conv2d(
                x,
                self.weight,
                bias=self.bias,
                stride=self.stride,
                padding=self.padding,
                dilation=self.dilation,
                groups=self.groups,
            )

[shen-shanshan]

@Isotr0py Currently, I find that vllm only convert Conv3D to Linear with conv3d_to_linear_weight() in vision.py and don't apply this optimization for Conv2D. Could we directly implement it as below?

• Conv2dLayer -> only use F.conv2d(). (because we don't convert the Conv2D weight when loading)
• Conv3dLayer -> directly use F.linear() when kernel_size==stride for better performance.

[Isotr0py]

(because we don't convert the Conv2D weight when loading)

In fact, conv3d_to_linear_weight is just a workaround for nn.Conv3d when fixing the performance regression. Since we decide to expose conv layer as CustomOP, we can optimize its usage with a better implementation.

We can implement weight_loader for conv layer like linear layer so that it can automatically convert weight when calling weight_loader(param, loaded_weight):

set_weight_attrs(weight, {"weight_loader": self.weight_loader})

def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
    if self.can_linearize:
        out_channels = loaded_weight.shape[0]
        loaded_weight = loaded_weight.reshape(out_channels, -1)
    param.data.copy_(loaded_weight)

[shen-shanshan]

OK, that's a good idea, I will follow this to implement it today.

[Isotr0py]

BTW, let me update the layer implementation for you tonight. :)

[shen-shanshan]

BTW, let me update the layer implementation for you tonight. :)

Oh sorry I didn't notice that and have just force pushed a new commit.. 😂

[Isotr0py]

/gemini review

[gemini-code-assist]

Code Review

The pull request successfully extracts ConvLayer as CustomOp, which improves modularity and extensibility. The changes correctly integrate the new Conv2dLayer and Conv3dLayer into existing models, replacing direct nn.Conv2d or ReplicatedLinear calls and removing the conv3d_to_linear_weight utility function. This is a positive step towards better management of convolution operations within the framework. However, there are a few areas in the new conv.py file that require attention to ensure robustness and correctness, particularly regarding the handling of kernel_size and stride in CausalConv2dLayer when they are provided as tuples.

[Isotr0py]

Let's see if the CI can pass now.

[shen-shanshan]

I have also tested this PR on ascend platform and it worked well, even getting a better performance.

see: vllm-project/vllm-ascend#4198.

[Isotr0py]

Also cc @jikunshang @bigPYJ1151. CPU/XPU platform may need to update forward_cpu/forward_xpu for convolution layer optimization as well.

[jikunshang]

@Isotr0py thanks for reminder! for XPU we will call forward_native by default. and for CPU, it will call forward_cuda. ideally both should work since it's just torch operator calls here. we will check whether this affect functionality/perf on XPU&CPU.