[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"]
-                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":

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(

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)