[Kimi K3] optimize: preprocess cpu-transport images on the vision owner (#33921)

python/sglang/srt/managers/mm_schedule.py

@@ -178,6 +178,7 @@ def _can_skip_pre_embed_feature_move(data_embedding_func: DataEmbeddingFunc) ->
         "Qwen3_5ForConditionalGeneration",
         "Qwen3_5MoeForConditionalGeneration",
         "KimiK25ForConditionalGeneration",
+        "KimiK3ForConditionalGeneration",
     }
 
 

python/sglang/srt/models/kimi_k3.py

@@ -104,6 +104,12 @@ from sglang.srt.models.kimi_k3_vl import (
 )
 from sglang.srt.models.transformers import maybe_prefix
 from sglang.srt.models.utils import WeightsMapper
+from sglang.srt.multimodal.kimi_k3_image_processing import (
+    DEFERRED_PREPROCESSING_KEY,
+    fill_transparent_bg,
+    normalization_tensors,
+    to_chw_uint8,
+)
 from sglang.srt.multimodal.mm_utils import materialize_multimodal_features
 from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args
 from sglang.srt.utils import is_blackwell_supported, is_hip, make_layers
@@ -3075,40 +3081,59 @@ class KimiK3ForConditionalGeneration(nn.Module):
         grid_thw_list = grid_thws_host.tolist()
 
         def materialize_item_features(image_indices: List[int]) -> torch.Tensor:
-            """Materialize features for the images assigned to this rank.
+            """Materialize only the images assigned to this vision-DP rank."""
-
+            ipc_consumer_count = max(get_parallel().tp_size, 1)
-            K3 vision is image-wise data-parallel, so each image is consumed
-            by exactly one TP rank. Deferred CUDA-IPC proxies are
-            reconstructed here, after the assignment is known, so an image
-            crosses the tokenizer/scheduler boundary once instead of once
-            per rank; CPU-transport features likewise only pay their H2D
-            copy on the owner rank. The consumer count matches
-            MmItemMemoryPool.try_to_recycle(), which waits for the server TP
-            size rather than the attention subgroup size.
-            """
-            parallel = get_parallel()
-            server_args = get_server_args()
-            ipc_consumer_count = max(
-                getattr(server_args, "tp_size", parallel.attn_tp_size), 1
-            )
             device_index = device.index
             if device.type == "cuda" and device_index is None:
                 device_index = torch.cuda.current_device()
 
-            features = []
+            selected_items = []
             for image_index in image_indices:
                 item = items[image_index]
                 if device.type == "cuda":
                     item.reconstruct(
                         device_index, ipc_consumer_count=ipc_consumer_count
                     )
-                feature = item.feature
+                selected_items.append(item)
-                if not isinstance(feature, torch.Tensor):
+
+            deferred = [
+                item.model_specific_data.get(DEFERRED_PREPROCESSING_KEY)
+                for item in selected_items
+            ]
+            if any(config is not None for config in deferred):
+                if not all(config is not None for config in deferred):
+                    raise ValueError(
+                        "Kimi-K3 cannot mix deferred and preprocessed image features"
+                    )
+                from sglang.srt.multimodal.processors.kimi_k25 import (
+                    _gpu_preprocess_images,
+                )
+
+                first_config = deferred[0]
+                image_scale, image_bias = normalization_tensors(
+                    first_config["image_mean"], first_config["image_std"], device
+                )
+                pixel_values, _ = _gpu_preprocess_images(
+                    [item.feature for item in selected_items],
+                    [config["resize_config"] for config in deferred],
+                    image_scale,
+                    image_bias,
+                    self.vision_tower.patch_size,
+                    to_chw=lambda image: to_chw_uint8(image, device=device),
+                    post_resize=lambda x: fill_transparent_bg(
+                        x, first_config["transparent_bg_config"]
+                    ),
+                )
+                return pixel_values.to(dtype=target_dtype)
+
+            features = []
+            for item in selected_items:
+                if not isinstance(item.feature, torch.Tensor):
                     raise TypeError(
                         "Kimi-K3 image feature must be a torch.Tensor, "
-                        f"got {type(feature)}"
+                        f"got {type(item.feature)}"
                     )
-                features.append(feature)
+                features.append(item.feature)
             return materialize_multimodal_features(
                 features, device=device, dtype=target_dtype
             )

python/sglang/srt/multimodal/kimi_k3_image_processing.py

@@ -0,0 +1,81 @@
+from typing import Union
+
+import numpy as np
+import torch
+from PIL import Image
+
+DEFERRED_PREPROCESSING_KEY = "kimi_k3_deferred_preprocessing"
+
+
+def to_chw_uint8(
+    image: Union[torch.Tensor, Image.Image],
+    device: torch.device | str | None = None,
+) -> torch.Tensor:
+    if isinstance(image, Image.Image):
+        has_alpha = image.mode != "RGB" and (
+            "A" in image.getbands() or "transparency" in image.info
+        )
+        array = np.array(image.convert("RGBA" if has_alpha else "RGB"), copy=True)
+        image = torch.from_numpy(array).permute(2, 0, 1)
+
+    if image.dtype != torch.uint8:
+        raise ValueError(
+            f"Kimi-K3 preprocessing expects raw uint8 pixels, got {image.dtype}"
+        )
+    if image.dim() == 2:
+        image = image.unsqueeze(0)
+    if image.shape[0] == 1:
+        image = image.repeat(3, 1, 1)
+    if device is not None:
+        image = image.to(device)
+    return image
+
+
+def fill_transparent_bg(x: torch.Tensor, bg_config: Union[dict, None]) -> torch.Tensor:
+    if x.shape[1] == 3:
+        return x
+    rgb = x[:, :3]
+    if bg_config is None:
+        return rgb
+
+    _, _, height, width = x.shape
+    pattern = bg_config.get("pattern", "black")
+    if pattern == "chessboard":
+        square = bg_config.get("chessboard_square_size", 16)
+        white = float(bg_config.get("chessboard_white_value", 255))
+        gray = float(bg_config.get("chessboard_gray_value", 200))
+        top_left = bg_config.get("chessboard_square_on_top_left", True)
+        ys = torch.arange(height, device=x.device) // square
+        xs = torch.arange(width, device=x.device) // square
+        parity = (ys.unsqueeze(1) + xs.unsqueeze(0)) % 2
+        background = torch.where(parity == (1 if top_left else 0), gray, white)
+        background = background.unsqueeze(0).expand(3, height, width)
+    elif pattern == "white":
+        background = torch.full((3, height, width), 255.0, device=x.device)
+    elif pattern == "black":
+        background = torch.zeros(3, height, width, device=x.device)
+    elif pattern == "gray":
+        background = torch.full((3, height, width), 128.0, device=x.device)
+    else:
+        raise ValueError(f"Invalid background pattern: {pattern}")
+
+    alpha = (x[:, 3:4] / 255.0).clamp(0.0, 1.0)
+    return (alpha * rgb + (1.0 - alpha) * background).clamp(0.0, 255.0).floor_()
+
+
+def normalization_tensors(
+    image_mean: list[float],
+    image_std: list[float],
+    device: torch.device | str,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    scale = torch.tensor(
+        [1.0 / (255.0 * std) for std in image_std],
+        device=device,
+        dtype=torch.float32,
+    ).view(1, 3, 1, 1)
+    bias = torch.tensor(
+        [-mean / std for mean, std in zip(image_mean, image_std)],
+        device=device,
+        dtype=torch.float32,
+    ).view(1, 3, 1, 1)
+    return scale, bias

python/sglang/srt/multimodal/processors/kimi_k3.py

@@ -15,8 +15,21 @@ import numpy as np
 import torch
 from PIL import Image
 
-from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput
+from sglang.srt.managers.schedule_batch import (
+    Modality,
+    MultimodalDataItem,
+    MultimodalProcessorOutput,
+)
 from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
+from sglang.srt.multimodal.kimi_k3_image_processing import (
+    DEFERRED_PREPROCESSING_KEY,
+)
+from sglang.srt.multimodal.kimi_k3_image_processing import (
+    fill_transparent_bg as _fill_transparent_bg,
+)
+from sglang.srt.multimodal.kimi_k3_image_processing import (
+    to_chw_uint8,
+)
 from sglang.srt.multimodal.processors.base_processor import (
     BaseMultimodalProcessor as SGLangBaseProcessor,
 )
@@ -28,8 +41,10 @@ from sglang.srt.multimodal.processors.kimi_k25 import (
     KimiGPUProcessorWrapper,
     _get_image_dimensions,
     _gpu_preprocess_images,
+    _grid_thw_from_resize_config,
     navit_resize_config,
 )
+from sglang.srt.utils import is_cuda
 from sglang.srt.utils.cuda_ipc_transport_utils import (
     DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
 )
@@ -124,15 +139,7 @@ def _expand_k3_image_prompt_text(
 
 def _k3_to_cuda_chw(image: Union[torch.Tensor, Image.Image]) -> torch.Tensor:
     if isinstance(image, Image.Image):
-        # The checkpoint's fill_transparent_bg_with() returns RGB-mode images
+        return to_chw_uint8(image, device="cuda")
-        # untouched before it ever inspects the alpha bands, so an RGB image
-        # carrying a stray "transparency" info key must NOT be promoted to
-        # RGBA here.
-        has_alpha = image.mode != "RGB" and (
-            "A" in image.getbands() or "transparency" in image.info
-        )
-        arr = np.asarray(image.convert("RGBA" if has_alpha else "RGB"))
-        return torch.from_numpy(arr).permute(2, 0, 1).cuda()
 
     image = image.cuda()
     if image.dim() == 2:
@@ -142,51 +149,6 @@ def _k3_to_cuda_chw(image: Union[torch.Tensor, Image.Image]) -> torch.Tensor:
     return image
 
 
-def _chessboard_background(
-    height: int, width: int, cfg: dict, device: torch.device
-) -> torch.Tensor:
-    square = cfg.get("chessboard_square_size", 16)
-    white = float(cfg.get("chessboard_white_value", 255))
-    gray = float(cfg.get("chessboard_gray_value", 200))
-    on_top_left = cfg.get("chessboard_square_on_top_left", True)
-
-    ys = torch.arange(height, device=device) // square
-    xs = torch.arange(width, device=device) // square
-    parity = (ys.unsqueeze(1) + xs.unsqueeze(0)) % 2
-    gray_parity = 1 if on_top_left else 0
-    bg = torch.where(parity == gray_parity, gray, white)
-    return bg.unsqueeze(0).expand(3, height, width)
-
-
-def _fill_transparent_bg(x: torch.Tensor, bg_cfg: Union[dict, None]) -> torch.Tensor:
-    """Composite a resized (1, 4, H, W) float image in [0, 255] onto the
-    configured background; 3-channel input passes through."""
-    if x.shape[1] == 3:
-        return x
-    rgb = x[:, :3]
-    if bg_cfg is None:
-        return rgb
-
-    _, _, height, width = x.shape
-    pattern = bg_cfg.get("pattern", "black")
-    if pattern == "chessboard":
-        bg = _chessboard_background(height, width, bg_cfg, x.device)
-    elif pattern == "white":
-        bg = torch.full((3, height, width), 255.0, device=x.device)
-    elif pattern == "black":
-        bg = torch.zeros(3, height, width, device=x.device)
-    elif pattern == "gray":
-        bg = torch.full((3, height, width), 128.0, device=x.device)
-    else:
-        raise ValueError(f"Invalid background pattern: {pattern}")
-
-    alpha = (x[:, 3:4] / 255.0).clamp(0.0, 1.0)
-    # The checkpoint processor casts the composited float result back with
-    # numpy's astype(np.uint8), which truncates; floor matches that exactly
-    # (a composite of [0, 255] inputs is always non-negative).
-    return (alpha * rgb + (1.0 - alpha) * bg).clamp(0.0, 255.0).floor_()
-
-
 class KimiK3GPUProcessorWrapper(KimiGPUProcessorWrapper):
     def __init__(self, *args, transparent_bg_config=None, **kwargs):
         super().__init__(*args, **kwargs)
@@ -292,6 +254,31 @@ class KimiK3GPUProcessorWrapper(KimiGPUProcessorWrapper):
             out["image_grid_thw"] = grid_thws
         return out
 
+    def prepare_deferred(self, text, images, original_input_ids=None):
+        input_text = text[0] if isinstance(text, list) else text
+        image_sizes = [_get_image_dimensions(image) for image in images]
+        resize_configs = [
+            navit_resize_config(
+                width,
+                height,
+                self._patch_size,
+                self._merge_kernel_size,
+                self._in_patch_limit,
+                self._patch_limit_on_one_side,
+                self._fixed_output_tokens,
+            )
+            for width, height in image_sizes
+        ]
+        input_ids = self._prepare_input_ids(
+            input_text, resize_configs, original_input_ids, image_sizes
+        )
+        deferred_config = {
+            "image_mean": list(self._image_mean),
+            "image_std": list(self._image_std),
+            "transparent_bg_config": self._transparent_bg_config,
+        }
+        return input_ids, resize_configs, deferred_config
+
 
 class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
     models = [KimiK3ForConditionalGeneration]
@@ -328,6 +315,91 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
         super().__init__(hf_config, server_args, processor, *args, **kwargs)
         self.mm_tokens = mm_tokens
 
+    def _should_defer_gpu_preprocessing(self, images) -> bool:
+        if (
+            not images
+            or self.mm_feature_transport != "cpu"
+            or not is_cuda()
+            or not all(
+                isinstance(image, Image.Image)
+                or (isinstance(image, torch.Tensor) and image.dtype == torch.uint8)
+                for image in images
+            )
+        ):
+            return False
+
+        raw_bytes = 0
+        processed_bytes = 0
+        patch_size = self._processor._patch_size
+        for image in images:
+            width, height = _get_image_dimensions(image)
+            resize_config = navit_resize_config(
+                width,
+                height,
+                patch_size,
+                self._processor._merge_kernel_size,
+                self._processor._in_patch_limit,
+                self._processor._patch_limit_on_one_side,
+                self._processor._fixed_output_tokens,
+            )
+            if isinstance(image, torch.Tensor):
+                channels = (
+                    3 if image.dim() == 2 or image.shape[0] == 1 else image.shape[0]
+                )
+            else:
+                channels = (
+                    4
+                    if image.mode != "RGB"
+                    and ("A" in image.getbands() or "transparency" in image.info)
+                    else 3
+                )
+            raw_bytes += channels * width * height
+            padded_width = resize_config["new_width"] + resize_config["pad_width"]
+            padded_height = resize_config["new_height"] + resize_config["pad_height"]
+            processed_bytes += 3 * padded_width * padded_height * torch.float32.itemsize
+
+        return raw_bytes <= processed_bytes
+
+    def _build_deferred_output(self, base_output):
+        input_ids, resize_configs, deferred_config = self._processor.prepare_deferred(
+            base_output.input_text,
+            base_output.images,
+            base_output.input_ids,
+        )
+        offsets = self.get_mm_items_offset(
+            input_ids.flatten(), self.mm_tokens.image_token_id
+        )
+        if len(offsets) != len(base_output.images):
+            raise ValueError("Expected one Kimi-K3 image span for each image")
+
+        items = []
+        for image, resize_config, offset in zip(
+            base_output.images, resize_configs, offsets
+        ):
+            grid_thw = _grid_thw_from_resize_config(
+                resize_config, self._processor._patch_size
+            )
+            item = MultimodalDataItem(
+                modality=Modality.IMAGE,
+                feature=to_chw_uint8(image),
+                offsets=[offset],
+                model_specific_data={
+                    "image_grid_thw": torch.tensor([grid_thw], dtype=torch.int64),
+                    DEFERRED_PREPROCESSING_KEY: {
+                        **deferred_config,
+                        "resize_config": resize_config,
+                    },
+                },
+            )
+            items.append(item)
+
+        self._precompute_hashes_before_cpu_transfer(items)
+        return MultimodalProcessorOutput(
+            input_ids=input_ids.flatten().tolist(),
+            mm_items=items,
+            im_token_id=self.mm_tokens.image_token_id,
+        )
+
     async def process_mm_data_async(
         self,
         image_data: List[Union[str, bytes, Dict]],
@@ -378,6 +450,9 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
                 f"expected {expected_image_count}, loaded {len(base_output.images)}"
             )
 
+        if self._should_defer_gpu_preprocessing(base_output.images):
+            return self._build_deferred_output(base_output)
+
         mm_items, input_ids, _ = await self.process_and_combine_mm_data_async(
             base_output,
             self.mm_tokens,

test/registered/unit/models/test_kimi_k25.py

@@ -680,6 +680,138 @@ def test_kimi_k3_epd_rebuild_uses_the_same_media_contract():
     )
 
 
+def test_kimi_k3_cpu_transport_defers_gpu_preprocessing():
+    from sglang.srt.multimodal.kimi_k3_image_processing import (
+        DEFERRED_PREPROCESSING_KEY,
+    )
+
+    processor = object.__new__(KimiK3ImageProcessor)
+    processor.mm_tokens = SimpleNamespace(image_token_id=99)
+    processor.mm_feature_transport = "cpu"
+    processor.use_cuda_ipc = False
+    processor._processor = SimpleNamespace(
+        _patch_size=2,
+        prepare_deferred=Mock(
+            return_value=(
+                torch.tensor([[1, 99, 99, 2, 99, 3]]),
+                [
+                    {
+                        "num_tokens": 2,
+                        "new_width": 4,
+                        "new_height": 2,
+                        "pad_width": 0,
+                        "pad_height": 2,
+                    },
+                    {
+                        "num_tokens": 1,
+                        "new_width": 2,
+                        "new_height": 2,
+                        "pad_width": 2,
+                        "pad_height": 2,
+                    },
+                ],
+                {
+                    "image_mean": [0.5, 0.5, 0.5],
+                    "image_std": [0.5, 0.5, 0.5],
+                    "transparent_bg_config": None,
+                },
+            )
+        ),
+    )
+    images = [
+        torch.arange(3 * 2 * 4, dtype=torch.uint8).reshape(3, 2, 4),
+        torch.arange(3 * 2 * 2, dtype=torch.uint8).reshape(3, 2, 2),
+    ]
+    base_output = SimpleNamespace(
+        input_text="prompt", images=images, input_ids=[1, 99, 2, 99, 3]
+    )
+
+    output = processor._build_deferred_output(base_output)
+
+    assert output.input_ids == [1, 99, 99, 2, 99, 3]
+    assert [item.offsets for item in output.mm_items] == [[(1, 2)], [(4, 4)]]
+    assert [item.feature.dtype for item in output.mm_items] == [
+        torch.uint8,
+        torch.uint8,
+    ]
+    assert [item.feature.shape for item in output.mm_items] == [
+        torch.Size([3, 2, 4]),
+        torch.Size([3, 2, 2]),
+    ]
+    assert [item.image_grid_thw.tolist() for item in output.mm_items] == [
+        [[1, 2, 2]],
+        [[1, 2, 2]],
+    ]
+    assert all(item.hash is not None for item in output.mm_items)
+    assert all(item.pad_value is not None for item in output.mm_items)
+    assert all(
+        DEFERRED_PREPROCESSING_KEY in item.model_specific_data
+        for item in output.mm_items
+    )
+
+
+@pytest.mark.parametrize(
+    ("image_shape", "in_patch_limit", "expected"),
+    [((3, 32, 32), 65536, True), ((3, 1024, 1024), 1, False)],
+)
+def test_kimi_k3_defers_only_when_raw_transport_is_smaller(
+    image_shape, in_patch_limit, expected
+):
+    processor = object.__new__(KimiK3ImageProcessor)
+    processor.mm_feature_transport = "cpu"
+    processor._processor = SimpleNamespace(
+        _patch_size=14,
+        _merge_kernel_size=2,
+        _in_patch_limit=in_patch_limit,
+        _patch_limit_on_one_side=512,
+        _fixed_output_tokens=None,
+    )
+    image = torch.zeros(image_shape, dtype=torch.uint8)
+
+    with patch("sglang.srt.multimodal.processors.kimi_k3.is_cuda", return_value=True):
+        assert processor._should_defer_gpu_preprocessing([image]) is expected
+
+
+def test_kimi_k3_does_not_defer_non_uint8_tensor_preprocessing():
+    processor = object.__new__(KimiK3ImageProcessor)
+    processor.mm_feature_transport = "cpu"
+
+    with patch("sglang.srt.multimodal.processors.kimi_k3.is_cuda", return_value=True):
+        assert not processor._should_defer_gpu_preprocessing(
+            [torch.zeros((3, 32, 32), dtype=torch.float32)]
+        )
+
+
+def test_kimi_k3_does_not_defer_empty_image_batch():
+    processor = object.__new__(KimiK3ImageProcessor)
+    processor.mm_feature_transport = "cpu"
+
+    with patch("sglang.srt.multimodal.processors.kimi_k3.is_cuda", return_value=True):
+        assert not processor._should_defer_gpu_preprocessing([])
+
+
+def test_kimi_k3_eager_preprocessing_preserves_float_tensor_support():
+    from sglang.srt.multimodal.processors.kimi_k3 import _k3_to_cuda_chw
+
+    image = torch.zeros((1, 4, 4), dtype=torch.float32)
+    with patch.object(torch.Tensor, "cuda", lambda self: self):
+        output = _k3_to_cuda_chw(image)
+
+    assert output.dtype == torch.float32
+    assert output.shape == (3, 4, 4)
+
+
+@pytest.mark.parametrize("transport", ["cuda_ipc", "fabric"])
+def test_kimi_k3_keeps_gpu_transport_preprocessing_eager(transport):
+    processor = object.__new__(KimiK3ImageProcessor)
+    processor.mm_feature_transport = transport
+
+    with patch("sglang.srt.multimodal.processors.kimi_k3.is_cuda", return_value=True):
+        assert not processor._should_defer_gpu_preprocessing(
+            [torch.zeros((3, 32, 32), dtype=torch.uint8)]
+        )
+
+
 def test_kimi_k3_rejects_silently_dropped_images():
     processor = object.__new__(KimiK3ImageProcessor)
     processor.mm_tokens = Mock()
@@ -699,6 +831,8 @@ def test_kimi_k3_uses_token_ids_to_preserve_media_boundaries():
     processor = object.__new__(KimiK3ImageProcessor)
     processor.mm_feature_transport = "cpu"
     processor.mm_tokens = SimpleNamespace(image_token_id=99)
+    processor.mm_feature_transport = "cuda_ipc"
+    processor.use_cuda_ipc = True
     processor.fast_load_mm_data = AsyncMock(
         return_value=SimpleNamespace(
             images=[object(), object()], input_ids=[1, 99, 2, 99, 3]

test/registered/unit/models/test_kimi_k3_vision.py

@@ -167,7 +167,7 @@ def test_kimi_k3_vision_tower_reuses_prepared_forward_metadata(monkeypatch):
     )
 
     assert len(actual) == len(reference) == 1
-    assert torch.equal(actual[0], reference[0])
+    torch.testing.assert_close(actual[0], reference[0], rtol=0, atol=0, equal_nan=True)
 
 
 def test_kimi_k3_dp_helper_passes_host_grid_list_to_capable_tower():
@@ -424,6 +424,7 @@ if __name__ == "__main__":
 class _K3TowerStub:
     device = torch.device("cpu")
     merge_kernel_size = (2, 2)
+    patch_size = 2
 
     def __init__(self):
         self.config = SimpleNamespace(hidden_size=2)
@@ -467,14 +468,11 @@ def test_kimi_k3_encoder_dp_defers_feature_materialization(monkeypatch):
         "sglang.srt.multimodal.mm_utils.run_dp_sharded_mrope_vision_model",
         return_value=sharded_embeddings,
     ) as run_dp, mock_patch(
-        "sglang.srt.models.kimi_k3.get_server_args",
-        return_value=SimpleNamespace(tp_size=1),
-    ), mock_patch(
         "sglang.srt.models.kimi_k3.get_parallel",
-        return_value=SimpleNamespace(attn_tp_size=1),
+        return_value=SimpleNamespace(tp_size=1, attn_tp_size=1),
     ):
         output = model.get_image_feature(items)
-        # exercise the loader inside the patch scope: it reads server args
+        # Exercise the loader while the runtime topology is patched.
         loader_in_scope = run_dp.call_args.kwargs["load_local_pixel_values"]
         local = loader_in_scope([1])
         both = loader_in_scope([0, 1])
@@ -502,6 +500,80 @@ def test_kimi_k3_encoder_dp_defers_feature_materialization(monkeypatch):
     )
 
 
+def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch):
+    from unittest.mock import patch as mock_patch
+
+    from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
+    from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
+    from sglang.srt.multimodal.kimi_k3_image_processing import (
+        DEFERRED_PREPROCESSING_KEY,
+    )
+
+    model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.use_data_parallel = True
+    model.vision_tower = _K3TowerStub()
+    model.mm_projector = lambda image_embeds: image_embeds
+
+    deferred_config = {
+        "image_mean": [0.5, 0.5, 0.5],
+        "image_std": [0.5, 0.5, 0.5],
+        "transparent_bg_config": None,
+        "resize_config": {
+            "num_tokens": 1,
+            "new_width": 2,
+            "new_height": 2,
+            "pad_width": 0,
+            "pad_height": 0,
+        },
+    }
+    items = [
+        MultimodalDataItem(
+            modality=Modality.IMAGE,
+            offsets=[(index, index)],
+            feature=torch.full((3, 2, 2), index, dtype=torch.uint8),
+            model_specific_data={
+                "image_grid_thw": torch.tensor([[1, 1, 1]]),
+                DEFERRED_PREPROCESSING_KEY: deferred_config,
+            },
+        )
+        for index in range(2)
+    ]
+    calls = []
+
+    def fake_preprocess(images, resize_configs, *args, **kwargs):
+        calls.append([int(image[0, 0, 0]) for image in images])
+        return torch.tensor([[float(calls[-1][0]), 0.0]]), torch.tensor([[1, 1, 1]])
+
+    with mock_patch(
+        "sglang.srt.multimodal.mm_utils.run_dp_sharded_mrope_vision_model",
+        return_value=torch.zeros(1, 2),
+    ) as run_dp, mock_patch(
+        "sglang.srt.models.kimi_k3.get_parallel",
+        return_value=SimpleNamespace(tp_size=1, attn_tp_size=1),
+    ), mock_patch(
+        "sglang.srt.multimodal.processors.kimi_k25._gpu_preprocess_images",
+        side_effect=fake_preprocess,
+    ):
+        model.get_image_feature(items)
+        loader = run_dp.call_args.kwargs["load_local_pixel_values"]
+        one = loader([1])
+
+    assert calls == [[1]]
+    assert one.dtype == torch.float32
+    assert one.tolist() == [[1.0, 0.0]]
+
+
+def test_kimi_k3_scheduler_leaves_feature_placement_to_dp_owner():
+    from sglang.srt.managers.mm_schedule import _can_skip_pre_embed_feature_move
+    from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
+
+    model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+
+    assert _can_skip_pre_embed_feature_move(model.get_image_feature)
+
+
 def test_kimi_k3_rejects_aggregated_items():
     """One item must carry exactly one logical image: the DP owner
     assignment and the bounded CUDA-IPC lease accounting are per-item, so