[Perf] Use Conv3dLayer for MiniMax M3 patch embedding (#58512)
Signed-off-by: OpenAI Codex <codex@openai.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Thien Tran committed
Sep 24, 2026 at 15:41 UTC
00b7847c8036b667742b4efb21aab1de51fd4721
2 files changed
+45
-1
tests/model_executor/layers/test_conv.py
new
+43
@@ -0,0 +1,43 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+from torch import nn
7
+
8
+from vllm.model_executor.layers.conv import Conv3dLayer
9
+
10
+
11
+@pytest.mark.parametrize(
12
+ "input_shape,kernel_size,out_channels,bias",
13
+ [
14
+ pytest.param((4, 3, 2, 14, 14), (2, 14, 14), 1280, False, id="minimax_m3"),
15
+ pytest.param((4, 3, 2, 14, 14), (2, 14, 14), 1152, False, id="qwen2_vl"),
16
+ pytest.param((4, 3, 2, 14, 14), (2, 14, 14), 1152, True, id="qwen3_vl"),
17
+ pytest.param((4, 3, 1, 14, 14), (1, 14, 14), 1536, True, id="glm_vl"),
18
+ pytest.param((1, 3, 4, 28, 28), (2, 14, 14), 32, True, id="multiple_patches"),
19
+ ],
20
+)
21
+def test_conv3d_patch_embedding_correctness(
22
+ default_vllm_config, input_shape, kernel_size, out_channels, bias
23
+):
24
+ device = "cuda" if torch.cuda.is_available() else "cpu"
25
+ dtype = torch.bfloat16 if device == "cuda" else torch.float32
26
+ reference = nn.Conv3d(
27
+ input_shape[1], out_channels, kernel_size, stride=kernel_size, bias=bias
28
+ ).to(device=device, dtype=dtype)
29
+ layer = Conv3dLayer(
30
+ input_shape[1],
31
+ out_channels,
32
+ kernel_size,
33
+ stride=kernel_size,
34
+ bias=bias,
35
+ params_dtype=dtype,
36
+ ).to(device)
37
+ layer.load_state_dict(reference.state_dict(), strict=True)
38
+ assert layer.enable_linear
39
+
40
+ x = torch.randn(input_shape, device=device, dtype=dtype)
41
+ expected = reference(x)
42
+ atol = rtol = 0.02 if dtype == torch.bfloat16 else 1e-5
43
+ torch.testing.assert_close(layer(x), expected, atol=atol, rtol=rtol)
vllm/models/minimax_m3/common/vision_tower.py
+2
-1
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.activation import get_act_fn
14
from vllm.model_executor.layers.attention.mm_encoder_attention import (
15
MMEncoderAttention,
16
)
17
+from vllm.model_executor.layers.conv import Conv3dLayer
18
from vllm.model_executor.layers.linear import (
19
ColumnParallelLinear,
20
QKVParallelLinear,
@@ -55,7 +56,7 @@ class MiniMaxVLPatchEmbed(nn.Module):
56
self.num_channels = num_channels
57
self.hidden_size = config.hidden_size
58
58
- self.patch_embedding = nn.Conv3d(
59
+ self.patch_embedding = Conv3dLayer(
60
in_channels=num_channels,
61
out_channels=config.hidden_size,
62
kernel_size=(temporal_patch_size, patch_size, patch_size),