From 1c8fcf436f596699d6c96c1d293feb55a3e3dc0c Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Wed, 28 Jan 2026 10:41:33 +0800 Subject: [PATCH] [Perf] replace torch's conv with vLLM's conv to fix torch 2.9 performance regression (#982) Signed-off-by: zjy0516 Co-authored-by: Hongsheng Liu --- vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index deb626a97..e4ba6318b 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -12,6 +12,7 @@ from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbed from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import FP32LayerNorm from vllm.logger import init_logger +from vllm.model_executor.layers.conv import Conv3dLayer from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear from vllm.model_executor.model_loader.weight_utils import default_weight_loader @@ -605,7 +606,12 @@ class WanTransformer3DModel(nn.Module): # 1. Patch & position embedding self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) - self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) + self.patch_embedding = Conv3dLayer( + in_channels=in_channels, + out_channels=inner_dim, + kernel_size=patch_size, + stride=patch_size, + ) # 2. Condition embeddings self.condition_embedder = WanTimeTextImageEmbedding(