diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 9630872a9..1cc190eb1 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -3276,8 +3276,8 @@ class KimiK3ForConditionalGeneration(nn.Module): "Kimi-K3 cannot mix deferred and preprocessed image features" ) first_config = deferred[0] - backend = first_config["backend"] - if any(config["backend"] != backend for config in deferred): + backend = first_config.backend + if any(config.backend != backend for config in deferred): raise ValueError( "Kimi-K3 cannot mix deferred preprocessing backends" ) @@ -3287,17 +3287,17 @@ class KimiK3ForConditionalGeneration(nn.Module): ) image_scale, image_bias = normalization_tensors( - first_config["image_mean"], first_config["image_std"], device + 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], + [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"] + x, first_config.transparent_bg_config ), ) elif backend == "cpu": diff --git a/python/sglang/srt/multimodal/kimi_k3_image_processing.py b/python/sglang/srt/multimodal/kimi_k3_image_processing.py index ab959ef87..4977c0ddc 100644 --- a/python/sglang/srt/multimodal/kimi_k3_image_processing.py +++ b/python/sglang/srt/multimodal/kimi_k3_image_processing.py @@ -1,6 +1,7 @@ import functools import math -from typing import Union +from dataclasses import dataclass +from typing import Literal, Optional, Union import numpy as np import torch @@ -9,6 +10,22 @@ from PIL import Image DEFERRED_PREPROCESSING_KEY = "kimi_k3_deferred_preprocessing" +@dataclass(frozen=True) +class KimiK3DeferredPreprocessing: + """Parameters the vision-DP owner needs to finish one deferred image. + + ``backend`` is the producer's decision and cannot be recovered from the + item; the feature's own layout stays observable on ``item.feature``, so it + is not mirrored here. + """ + + backend: Literal["gpu", "cpu"] + image_mean: list[float] + image_std: list[float] + transparent_bg_config: Optional[dict] + resize_config: dict + + def prepare_kimi_k3_encoder_inputs( images, image_processor, *, use_gpu_preprocessing=False ): @@ -61,12 +78,13 @@ def prepare_kimi_k3_encoder_inputs( patch_size = int(media_proc_cfg["patch_size"]) merge_kernel_size = int(media_proc_cfg["merge_kernel_size"]) - common_deferred_config = { - "backend": "gpu" if use_gpu_preprocessing else "cpu", - "image_mean": list(media_proc_cfg["image_mean"]), - "image_std": list(media_proc_cfg["image_std"]), - "transparent_bg_config": media_proc_cfg.get("transparent_bg_config"), - } + deferred_preprocessing = functools.partial( + KimiK3DeferredPreprocessing, + backend="gpu" if use_gpu_preprocessing else "cpu", + image_mean=list(media_proc_cfg["image_mean"]), + image_std=list(media_proc_cfg["image_std"]), + transparent_bg_config=media_proc_cfg.get("transparent_bg_config"), + ) items = [] grids = [] @@ -93,11 +111,9 @@ def prepare_kimi_k3_encoder_inputs( feature=to_chw_uint8(image) if use_gpu_preprocessing else image, model_specific_data={ "grid_thws": grid_tensor, - DEFERRED_PREPROCESSING_KEY: { - **common_deferred_config, - "feature_layout": "chw" if use_gpu_preprocessing else "raw", - "resize_config": resize_config, - }, + DEFERRED_PREPROCESSING_KEY: deferred_preprocessing( + resize_config=resize_config + ), }, ) if not use_gpu_preprocessing: @@ -133,9 +149,6 @@ def materialize_kimi_k3_cpu_features(items, image_processor) -> torch.Tensor: medias = [] for item in items: image = item.feature - config = item.model_specific_data[DEFERRED_PREPROCESSING_KEY] - if config["feature_layout"] != "raw": - raise ValueError("Kimi-K3 deferred CPU preprocessing expects raw inputs") if not isinstance(image, Image.Image): if not isinstance(image, torch.Tensor) or image.dtype != torch.uint8: raise TypeError( diff --git a/python/sglang/srt/multimodal/processors/kimi_k3.py b/python/sglang/srt/multimodal/processors/kimi_k3.py index a5686f906..c3aec5cc3 100644 --- a/python/sglang/srt/multimodal/processors/kimi_k3.py +++ b/python/sglang/srt/multimodal/processors/kimi_k3.py @@ -8,6 +8,7 @@ images onto the checkpoint-configured background at load time. """ +import functools import re from typing import Dict, List, Union @@ -23,6 +24,7 @@ from sglang.srt.managers.schedule_batch import ( from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration from sglang.srt.multimodal.kimi_k3_image_processing import ( DEFERRED_PREPROCESSING_KEY, + KimiK3DeferredPreprocessing, ) from sglang.srt.multimodal.kimi_k3_image_processing import ( fill_transparent_bg as _fill_transparent_bg, @@ -272,12 +274,16 @@ class KimiK3GPUProcessorWrapper(KimiGPUProcessorWrapper): 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 + # This path only ever defers GPU preprocessing: the caller gates on + # `_should_defer_gpu_preprocessing` and stages CHW uint8 features. + deferred_preprocessing = functools.partial( + KimiK3DeferredPreprocessing, + backend="gpu", + 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_preprocessing class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): @@ -365,7 +371,11 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): return raw_bytes <= processed_bytes def _build_deferred_output(self, base_output): - input_ids, resize_configs, deferred_config = self._processor.prepare_deferred( + ( + input_ids, + resize_configs, + deferred_preprocessing, + ) = self._processor.prepare_deferred( base_output.input_text, base_output.images, base_output.input_ids, @@ -389,10 +399,9 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): offsets=[offset], model_specific_data={ "image_grid_thw": torch.tensor([grid_thw], dtype=torch.int64), - DEFERRED_PREPROCESSING_KEY: { - **deferred_config, - "resize_config": resize_config, - }, + DEFERRED_PREPROCESSING_KEY: deferred_preprocessing( + resize_config=resize_config + ), }, ) items.append(item) diff --git a/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py b/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py index b25362fba..cba64f49b 100644 --- a/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py +++ b/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py @@ -188,8 +188,8 @@ def test_kimi_k3_epd_preprocess_preserves_raw_per_image_items(): assert item.hash is not None assert item.pad_value is not None deferred = item.model_specific_data[DEFERRED_PREPROCESSING_KEY] - assert deferred["image_mean"] == [0.5, 0.5, 0.5] - assert deferred["image_std"] == [0.5, 0.5, 0.5] + assert deferred.image_mean == [0.5, 0.5, 0.5] + assert deferred.image_std == [0.5, 0.5, 0.5] def test_kimi_k3_epd_model_preprocessor_receives_image_processor(): @@ -289,7 +289,7 @@ def test_kimi_k3_epd_default_cpu_materialization_is_owner_only_and_exact(): assert len(processor.calls[0]) == 1 assert processor.calls[0][0]["image"].getpixel((0, 0)) == (11, 0, 0) assert torch.all(materialized == 11) - assert items[0].model_specific_data[DEFERRED_PREPROCESSING_KEY]["backend"] == "cpu" + assert items[0].model_specific_data[DEFERRED_PREPROCESSING_KEY].backend == "cpu" def test_encoder_preprocess_materializes_only_local_size_balanced_items(): diff --git a/test/registered/unit/models/test_kimi_k25.py b/test/registered/unit/models/test_kimi_k25.py index c8d2694de..3cdbf0f52 100644 --- a/test/registered/unit/models/test_kimi_k25.py +++ b/test/registered/unit/models/test_kimi_k25.py @@ -1,6 +1,7 @@ """CPU coverage for Kimi-K2.5/K2.7 encoder-DP wiring.""" import asyncio +import functools from types import SimpleNamespace from unittest.mock import AsyncMock, Mock, patch @@ -690,6 +691,7 @@ 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, + KimiK3DeferredPreprocessing, ) processor = object.__new__(KimiK3ImageProcessor) @@ -717,11 +719,13 @@ def test_kimi_k3_cpu_transport_defers_gpu_preprocessing(): "pad_height": 2, }, ], - { - "image_mean": [0.5, 0.5, 0.5], - "image_std": [0.5, 0.5, 0.5], - "transparent_bg_config": None, - }, + functools.partial( + KimiK3DeferredPreprocessing, + backend="gpu", + image_mean=[0.5, 0.5, 0.5], + image_std=[0.5, 0.5, 0.5], + transparent_bg_config=None, + ), ) ), ) @@ -751,10 +755,13 @@ def test_kimi_k3_cpu_transport_defers_gpu_preprocessing(): ] 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 - ) + deferred = [ + item.model_specific_data[DEFERRED_PREPROCESSING_KEY] for item in output.mm_items + ] + # The staged features are CHW uint8, so the config has to route them to the + # GPU arm of `materialize_item_features`. + assert [config.backend for config in deferred] == ["gpu", "gpu"] + assert [config.resize_config["new_width"] for config in deferred] == [4, 2] @pytest.mark.parametrize( diff --git a/test/registered/unit/models/test_kimi_k3_vision.py b/test/registered/unit/models/test_kimi_k3_vision.py index ddc3b6026..d934a52ca 100644 --- a/test/registered/unit/models/test_kimi_k3_vision.py +++ b/test/registered/unit/models/test_kimi_k3_vision.py @@ -509,6 +509,7 @@ def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch): from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration from sglang.srt.multimodal.kimi_k3_image_processing import ( DEFERRED_PREPROCESSING_KEY, + KimiK3DeferredPreprocessing, ) model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration) @@ -517,20 +518,19 @@ def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch): model.vision_tower = _K3TowerStub() model.mm_projector = lambda image_embeds: image_embeds - deferred_config = { - "backend": "gpu", - "feature_layout": "chw", - "image_mean": [0.5, 0.5, 0.5], - "image_std": [0.5, 0.5, 0.5], - "transparent_bg_config": None, - "resize_config": { + deferred_config = KimiK3DeferredPreprocessing( + backend="gpu", + 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,