[Kimi-K3][Perf] Dispatch GEMM for vision patch embedder (#58527)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

tests/model_executor/layers/test_conv.py

@@ -5,7 +5,32 @@ import pytest
 import torch
 from torch import nn
 
-from vllm.model_executor.layers.conv import Conv3dLayer
+from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
+
+
+@pytest.mark.parametrize(
+    "input_shape,bias",
+    [
+        pytest.param((256, 3, 14, 14), False, id="packed_patches"),
+        pytest.param((2, 3, 28, 28), True, id="multiple_spatial_patches"),
+    ],
+)
+def test_conv2d_patch_embedding_correctness(default_vllm_config, input_shape, bias):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
+    reference = nn.Conv2d(3, 1024, 14, stride=14, bias=bias).to(
+        device=device, dtype=dtype
+    )
+    layer = Conv2dLayer(3, 1024, 14, stride=14, bias=bias, params_dtype=dtype).to(
+        device
+    )
+    layer.load_state_dict(reference.state_dict(), strict=True)
+    assert layer.enable_linear
+
+    x = torch.randn(input_shape, device=device, dtype=dtype)
+    expected = reference(x)
+    atol = rtol = 0.02 if dtype == torch.bfloat16 else 1e-5
+    torch.testing.assert_close(layer.forward_native(x), expected, atol=atol, rtol=rtol)
 
 
 @pytest.mark.parametrize(

vllm/model_executor/models/kimi_k25_vit.py

@@ -21,6 +21,7 @@ from vllm.distributed import divide, get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
+from vllm.model_executor.layers.conv import Conv2dLayer
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
@@ -194,7 +195,7 @@ class MoonVision3dPatchEmbed(nn.Module):
         )
         self.patch_size = patch_size
 
-        self.proj = nn.Conv2d(
+        self.proj = Conv2dLayer(
             in_dim,
             out_dim,
             kernel_size=patch_size,
@@ -220,26 +221,13 @@ class MoonVision3dPatchEmbed(nn.Module):
         *,
         pos_embeds: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        x = self._proj(x).view(x.size(0), -1)
+        # forward_native dispatches this non-overlapping patch projection to GEMM.
+        x = self.proj.forward_native(x).view(x.size(0), self.proj.out_channels)
         if pos_embeds is not None:
             return x + pos_embeds
         assert grid_thws is not None
         return self.pos_emb(x, grid_thws)
 
-    def _proj(self, x: torch.Tensor) -> torch.Tensor:
-        # MIOpen conv2d intermittently fails under load on ROCm; use aiter Triton.
-        if current_platform.is_rocm() and x.dtype in (torch.float16, torch.bfloat16):
-            from aiter.ops.triton.conv.conv2d import conv2d
-
-            return conv2d(
-                x,
-                self.proj.weight,
-                self.proj.bias,
-                stride=self.patch_size,
-                layout="nchw",
-            )
-        return self.proj(x)
-
 
 class Rope2DPosEmbRepeated(nn.Module):
     """2D rotary position embedding with multi-resolution support."""