[Fix] Carry the backend on Kimi-K3 deferred preprocessing configs (#34766)

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"]
+                backend = first_config.backend
-                if any(config["backend"] != backend for config in deferred):
+                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":

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 = {
+    deferred_preprocessing = functools.partial(
-        "backend": "gpu" if use_gpu_preprocessing else "cpu",
+        KimiK3DeferredPreprocessing,
-        "image_mean": list(media_proc_cfg["image_mean"]),
+        backend="gpu" if use_gpu_preprocessing else "cpu",
-        "image_std": list(media_proc_cfg["image_std"]),
+        image_mean=list(media_proc_cfg["image_mean"]),
-        "transparent_bg_config": media_proc_cfg.get("transparent_bg_config"),
+        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: {
+                DEFERRED_PREPROCESSING_KEY: deferred_preprocessing(
-                    **common_deferred_config,
+                    resize_config=resize_config
-                    "feature_layout": "chw" if use_gpu_preprocessing else "raw",
+                ),
-                    "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(

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 = {
+        # This path only ever defers GPU preprocessing: the caller gates on
-            "image_mean": list(self._image_mean),
+        # `_should_defer_gpu_preprocessing` and stages CHW uint8 features.
-            "image_std": list(self._image_std),
+        deferred_preprocessing = functools.partial(
-            "transparent_bg_config": self._transparent_bg_config,
+            KimiK3DeferredPreprocessing,
-        }
+            backend="gpu",
-        return input_ids, resize_configs, 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_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_PREPROCESSING_KEY: deferred_preprocessing(
-                        **deferred_config,
+                        resize_config=resize_config
-                        "resize_config": resize_config,
+                    ),
-                    },
                 },
             )
             items.append(item)

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_mean == [0.5, 0.5, 0.5]
-        assert deferred["image_std"] == [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():

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,
                     },
                 ],
-                {
+                functools.partial(
-                    "image_mean": [0.5, 0.5, 0.5],
+                    KimiK3DeferredPreprocessing,
-                    "image_std": [0.5, 0.5, 0.5],
+                    backend="gpu",
-                    "transparent_bg_config": None,
+                    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 = [
-        DEFERRED_PREPROCESSING_KEY in item.model_specific_data
+        item.model_specific_data[DEFERRED_PREPROCESSING_KEY] for item in output.mm_items
-        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(

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 = {
+    deferred_config = KimiK3DeferredPreprocessing(
-        "backend": "gpu",
+        backend="gpu",
-        "feature_layout": "chw",
+        image_mean=[0.5, 0.5, 0.5],
-        "image_mean": [0.5, 0.5, 0.5],
+        image_std=[0.5, 0.5, 0.5],
-        "image_std": [0.5, 0.5, 0.5],
+        transparent_bg_config=None,
-        "transparent_bg_config": None,
+        resize_config={
-        "resize_config": {
             "num_tokens": 1,
             "new_width": 2,
             "new_height": 2,
             "pad_width": 0,
             "pad_height": 0,
         },
-    }
+    )
     items = [
         MultimodalDataItem(
             modality=Modality.IMAGE,