From 22e003580bc0032918f1ebd1111de4850e33ec0f Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 9 Aug 2026 16:13:20 +0800 Subject: [PATCH] [Kimi K3] optimize: preprocess cpu-transport images on the vision owner (#33921) --- python/sglang/srt/managers/mm_schedule.py | 1 + python/sglang/srt/models/kimi_k3.py | 67 +++++-- .../multimodal/kimi_k3_image_processing.py | 81 ++++++++ .../srt/multimodal/processors/kimi_k3.py | 185 ++++++++++++------ test/registered/unit/models/test_kimi_k25.py | 134 +++++++++++++ .../unit/models/test_kimi_k3_vision.py | 84 +++++++- 6 files changed, 470 insertions(+), 82 deletions(-) create mode 100644 python/sglang/srt/multimodal/kimi_k3_image_processing.py diff --git a/python/sglang/srt/managers/mm_schedule.py b/python/sglang/srt/managers/mm_schedule.py index f08b4b1fe..9e91f6366 100644 --- a/python/sglang/srt/managers/mm_schedule.py +++ b/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", } diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 4c21d790e..3adb1831a 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/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. - - 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 - ) + """Materialize only the images assigned to this vision-DP rank.""" + ipc_consumer_count = max(get_parallel().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 - if not isinstance(feature, torch.Tensor): + selected_items.append(item) + + 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 ) diff --git a/python/sglang/srt/multimodal/kimi_k3_image_processing.py b/python/sglang/srt/multimodal/kimi_k3_image_processing.py new file mode 100644 index 000000000..e4d959f79 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/multimodal/processors/kimi_k3.py b/python/sglang/srt/multimodal/processors/kimi_k3.py index 0e6c67631..e3735b622 100644 --- a/python/sglang/srt/multimodal/processors/kimi_k3.py +++ b/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 - # 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() + return to_chw_uint8(image, device="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, diff --git a/test/registered/unit/models/test_kimi_k25.py b/test/registered/unit/models/test_kimi_k25.py index 2bc9dfeca..425b2c13c 100644 --- a/test/registered/unit/models/test_kimi_k25.py +++ b/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] diff --git a/test/registered/unit/models/test_kimi_k3_vision.py b/test/registered/unit/models/test_kimi_k3_vision.py index d1bb4cbda..4505baf62 100644 --- a/test/registered/unit/models/test_kimi_k3_vision.py +++ b/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