[Kimi-K3][Perf] Dispatch GEMM for vision patch embedder (#58527)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Thien Tran committed
Sep 25, 2026 at 09:30 UTC
ade1056ee5afda88c70ff3275288424ef4a9814a
2 files changed
+30
-17
tests/model_executor/layers/test_conv.py
+26
-1
@@ -5,7 +5,32 @@ import pytest
5
import torch
6
from torch import nn
7
8
-from vllm.model_executor.layers.conv import Conv3dLayer
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(
vllm/model_executor/models/kimi_k25_vit.py
+4
-16
@@ -21,6 +21,7 @@ from vllm.distributed import divide, get_tensor_model_parallel_world_size
21
from vllm.logger import init_logger
22
from vllm.model_executor.layers.activation import get_act_fn
23
from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
24
+from vllm.model_executor.layers.conv import Conv2dLayer
25
from vllm.model_executor.layers.linear import (
26
ColumnParallelLinear,
27
QKVParallelLinear,
@@ -194,7 +195,7 @@ class MoonVision3dPatchEmbed(nn.Module):
195
)
196
self.patch_size = patch_size
197
197
- self.proj = nn.Conv2d(
198
+ self.proj = Conv2dLayer(
199
in_dim,
200
out_dim,
201
kernel_size=patch_size,
@@ -220,26 +221,13 @@ class MoonVision3dPatchEmbed(nn.Module):
221
*,
222
pos_embeds: torch.Tensor | None = None,
223
) -> torch.Tensor:
223
- x = self._proj(x).view(x.size(0), -1)
224
+ # forward_native dispatches this non-overlapping patch projection to GEMM.
225
+ x = self.proj.forward_native(x).view(x.size(0), self.proj.out_channels)
226
if pos_embeds is not None:
227
return x + pos_embeds
228
assert grid_thws is not None
229
return self.pos_emb(x, grid_thws)
230
229
- def _proj(self, x: torch.Tensor) -> torch.Tensor:
230
- # MIOpen conv2d intermittently fails under load on ROCm; use aiter Triton.
231
- if current_platform.is_rocm() and x.dtype in (torch.float16, torch.bfloat16):
232
- from aiter.ops.triton.conv.conv2d import conv2d
233
-
234
- return conv2d(
235
- x,
236
- self.proj.weight,
237
- self.proj.bias,
238
- stride=self.patch_size,
239
- layout="nchw",
240
- )
241
- return self.proj(x)
242
-
231
232
class Rope2DPosEmbRepeated(nn.Module):
233
"""2D rotary position embedding with multi-resolution support."""