| 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 Conv2dLayer, Conv3dLayer |
| 9 | |
| 10 | |
| 11 | @pytest.mark.parametrize( |
| 12 | "input_shape,bias", |
| 13 | [ |
| 14 | pytest.param((256, 3, 14, 14), False, id="packed_patches"), |
| 15 | pytest.param((2, 3, 28, 28), True, id="multiple_spatial_patches"), |
| 16 | ], |
| 17 | ) |
| 18 | def test_conv2d_patch_embedding_correctness(default_vllm_config, input_shape, bias): |
| 19 | device = "cuda" if torch.cuda.is_available() else "cpu" |
| 20 | dtype = torch.bfloat16 if device == "cuda" else torch.float32 |
| 21 | reference = nn.Conv2d(3, 1024, 14, stride=14, bias=bias).to( |
| 22 | device=device, dtype=dtype |
| 23 | ) |
| 24 | layer = Conv2dLayer(3, 1024, 14, stride=14, bias=bias, params_dtype=dtype).to( |
| 25 | device |
| 26 | ) |
| 27 | layer.load_state_dict(reference.state_dict(), strict=True) |
| 28 | assert layer.enable_linear |
| 29 | |
| 30 | x = torch.randn(input_shape, device=device, dtype=dtype) |
| 31 | expected = reference(x) |
| 32 | atol = rtol = 0.02 if dtype == torch.bfloat16 else 1e-5 |
| 33 | torch.testing.assert_close(layer.forward_native(x), expected, atol=atol, rtol=rtol) |
| 34 | |
| 35 | |
| 36 | @pytest.mark.parametrize( |
| 37 | "input_shape,kernel_size,out_channels,bias", |
| 38 | [ |
| 39 | pytest.param((4, 3, 2, 14, 14), (2, 14, 14), 1280, False, id="minimax_m3"), |
| 40 | pytest.param((4, 3, 2, 14, 14), (2, 14, 14), 1152, False, id="qwen2_vl"), |
| 41 | pytest.param((4, 3, 2, 14, 14), (2, 14, 14), 1152, True, id="qwen3_vl"), |
| 42 | pytest.param((4, 3, 1, 14, 14), (1, 14, 14), 1536, True, id="glm_vl"), |
| 43 | pytest.param((1, 3, 4, 28, 28), (2, 14, 14), 32, True, id="multiple_patches"), |
| 44 | ], |
| 45 | ) |
| 46 | def test_conv3d_patch_embedding_correctness( |
| 47 | default_vllm_config, input_shape, kernel_size, out_channels, bias |
| 48 | ): |
| 49 | device = "cuda" if torch.cuda.is_available() else "cpu" |
| 50 | dtype = torch.bfloat16 if device == "cuda" else torch.float32 |
| 51 | reference = nn.Conv3d( |
| 52 | input_shape[1], out_channels, kernel_size, stride=kernel_size, bias=bias |
| 53 | ).to(device=device, dtype=dtype) |
| 54 | layer = Conv3dLayer( |
| 55 | input_shape[1], |
| 56 | out_channels, |
| 57 | kernel_size, |
| 58 | stride=kernel_size, |
| 59 | bias=bias, |
| 60 | params_dtype=dtype, |
| 61 | ).to(device) |
| 62 | layer.load_state_dict(reference.state_dict(), strict=True) |
| 63 | assert layer.enable_linear |
| 64 | |
| 65 | x = torch.randn(input_shape, device=device, dtype=dtype) |
| 66 | expected = reference(x) |
| 67 | atol = rtol = 0.02 if dtype == torch.bfloat16 else 1e-5 |
| 68 | torch.testing.assert_close(layer(x), expected, atol=atol, rtol=rtol) |