[EPD] feat: pipeline owner-only multimodal preprocessing (#34206)

python/sglang/srt/disaggregation/encode_server.py

@@ -58,6 +58,11 @@ from sglang.srt.model_executor.model_runner_components.load_model_utils import (
     maybe_precompile_model_kernels_after_loading,
 )
 from sglang.srt.model_loader import get_model
+from sglang.srt.multimodal.encoder_preprocessing import (
+    EncoderPreprocessOutput,
+    get_encoder_preprocessed_items,
+    invoke_encoder_preprocessor,
+)
 from sglang.srt.multimodal.processors.qwen_vl import preprocess_video
 from sglang.srt.observability.metrics_collector import EncoderMetricsCollector
 from sglang.srt.observability.req_time_stats import EncoderReqTimeStats
@@ -65,7 +70,13 @@ from sglang.srt.observability.trace import (
     process_tracing_init,
     trace_set_thread_info,
 )
-from sglang.srt.runtime_context import get_disagg, get_exec, get_mm, publish
+from sglang.srt.runtime_context import (
+    get_disagg,
+    get_exec,
+    get_mm,
+    get_parallel,
+    publish,
+)
 from sglang.srt.server_args import (
     PortArgs,
     ServerArgs,
@@ -862,9 +873,24 @@ class MMEncoder:
         return slices
 
     def _calculate_hashes_from_features(
-        self, mm_feature, grid_thw: List, modality: Modality
+        self, mm_feature, grid_thw: List, modality: Modality, mm_inputs=None
     ) -> List[int]:
         """CPU Task: Compute hashes based on processed feature patches."""
+        preprocessed_items = (
+            get_encoder_preprocessed_items(mm_inputs) if mm_inputs is not None else None
+        )
+        if preprocessed_items is not None:
+            if len(preprocessed_items) != len(grid_thw):
+                raise ValueError(
+                    "Encoder preprocess item/grid mismatch: "
+                    f"{len(preprocessed_items)} items != {len(grid_thw)} grids"
+                )
+            hashes = []
+            for item in preprocessed_items:
+                item.set_pad_value()
+                hashes.append(item.hash)
+            return hashes
+
         hashes = []
         if modality == Modality.AUDIO and isinstance(mm_feature, list):
             for feature in mm_feature:
@@ -884,21 +910,37 @@ class MMEncoder:
             offset += num_patches
         return hashes
 
-    def _encode_missing(
+    def _build_mm_data_items(
         self,
         mm_feature,
         mm_inputs: dict,
         indices: List[int],
-        modality: Modality = Modality.IMAGE,
+        modality: Modality,
-        get_feature_fn=None,
         grid_thw: Optional[List] = None,
-        keep_on_gpu: bool = False,
+    ) -> List[MultimodalDataItem]:
-    ) -> List[torch.Tensor]:
+        """Build the model-facing items selected for one encoder forward.
-        """
+
-        GPU Task: Run ViT inference ONLY on the subset of mm items missing from the cache.
+        A model preprocessor can preserve an item-wise representation with
+        ``EncoderPreprocessOutput``. This path avoids concatenating and then
+        re-slicing features before encoder-DP knows which rank owns each item.
+        Legacy Hugging Face processor outputs retain their existing aggregate
+        tensor behavior.
         """
         if grid_thw is None:
             grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type)
+
+        preprocessed_items = get_encoder_preprocessed_items(mm_inputs)
+        if preprocessed_items is not None:
+            if len(preprocessed_items) != len(grid_thw):
+                raise ValueError(
+                    "Encoder preprocess item/grid mismatch: "
+                    f"{len(preprocessed_items)} items != {len(grid_thw)} grids"
+                )
+            selected = [preprocessed_items[index] for index in indices]
+            if any(item.modality != modality for item in selected):
+                raise ValueError("Encoder preprocess output contains wrong modality")
+            return selected
+
         split_kimi_k3_images = (
             self.model_type == "kimi_k3" and modality == Modality.IMAGE
         )
@@ -916,11 +958,11 @@ class MMEncoder:
             sub_feature_list = []
             offsets = [0]
             curr = 0
-            for g in grid_thw:
+            for grid in grid_thw:
-                curr += self.get_num_patches(g, modality)
+                curr += self.get_num_patches(grid, modality)
                 offsets.append(curr)
-            for idx in indices:
+            for index in indices:
-                sub_feature_list.append(mm_feature[offsets[idx] : offsets[idx + 1]])
+                sub_feature_list.append(mm_feature[offsets[index] : offsets[index + 1]])
             if not split_kimi_k3_images:
                 sub_feature = torch.cat(sub_feature_list, dim=0)
 
@@ -948,19 +990,39 @@ class MMEncoder:
                 )
             ]
 
-        for k, v in mm_inputs.items():
+        for key, value in mm_inputs.items():
-            if k in _mm_feature_attrs.get(modality, []):
+            if key in _mm_feature_attrs.get(modality, []):
                 continue
-            val = _convert(v)
+            value = _convert(value)
-            if k in _mm_grid_attrs.get(modality, []):
+            if key in _mm_grid_attrs.get(modality, []):
                 if split_kimi_k3_images:
-                    for mm_item, idx in zip(mm_items, indices):
+                    for mm_item, index in zip(mm_items, indices):
-                        mm_item.set(k, val[idx : idx + 1])
+                        mm_item.set(key, value[index : index + 1])
                 else:
-                    mm_items[0].set(k, val[indices])
+                    mm_items[0].set(key, value[indices])
             else:
                 for mm_item in mm_items:
-                    mm_item.set(k, val)
+                    mm_item.set(key, value)
+        return mm_items
+
+    def _encode_missing(
+        self,
+        mm_feature,
+        mm_inputs: dict,
+        indices: List[int],
+        modality: Modality = Modality.IMAGE,
+        get_feature_fn=None,
+        grid_thw: Optional[List] = None,
+        keep_on_gpu: bool = False,
+    ) -> List[torch.Tensor]:
+        """
+        GPU Task: Run ViT inference ONLY on the subset of mm items missing from the cache.
+        """
+        if grid_thw is None:
+            grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type)
+        mm_items = self._build_mm_data_items(
+            mm_feature, mm_inputs, indices, modality, grid_thw
+        )
 
         forward_start = time.perf_counter()
         with torch.inference_mode():
@@ -1002,7 +1064,7 @@ class MMEncoder:
         if self.rank == 0:
             if hashes is None:
                 mm_hashes = self._calculate_hashes_from_features(
-                    mm_feature, grid_thw, modality
+                    mm_feature, grid_thw, modality, mm_inputs
                 )
             else:
                 mm_hashes = hashes
@@ -1651,12 +1713,31 @@ class MMEncoder:
         if not (self.image_processor or model_preprocessor):
             raise ValueError("No image processor available")
         images = await self._flatten_and_load_images(mm_items)
-        if model_preprocessor:
-            return model_preprocessor(images, Modality.IMAGE, self.vision_config)
-        image_config = self.vision_config.get("image", {})
-        original_image_sizes = [_get_original_image_size(item) for item in images]
         if self.model_type in ["kimi_k25", "kimi_k3", "kimi_vl"]:
             images = self._normalize_kimi_encoder_images(images)
+        original_image_sizes = [_get_original_image_size(item) for item in images]
+        if model_preprocessor:
+            processor_output = invoke_encoder_preprocessor(
+                model_preprocessor,
+                images,
+                Modality.IMAGE,
+                self.vision_config,
+                image_processor=self.image_processor,
+                use_gpu_preprocessing=self.use_image_processor_gpu,
+            )
+            if (
+                isinstance(processor_output, EncoderPreprocessOutput)
+                and processor_output.materialize_local_items is not None
+            ):
+                parallel = get_parallel()
+                await asyncio.get_running_loop().run_in_executor(
+                    self.preproc_executor,
+                    processor_output.materialize_for_rank,
+                    parallel.attn_tp_rank,
+                    parallel.attn_tp_size,
+                )
+            return processor_output
+        image_config = self.vision_config.get("image", {})
         processor_input = await asyncio.get_running_loop().run_in_executor(
             self.preproc_executor,
             functools.partial(self.image_processor, images=images, **image_config),
@@ -1775,25 +1856,25 @@ class MMEncoder:
             # support mm_cache
             mm_embedding = None
             mm_hash = None
-
+            mm_feature = _convert(_get_mm_feature(mm_inputs, modality))
-            mm_item = MultimodalDataItem.from_dict(
+            grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type)
-                {
+            model_mm_items = self._build_mm_data_items(
-                    "modality": modality,
+                mm_feature,
-                    "feature": _convert(_get_mm_feature(mm_inputs, modality)),
+                mm_inputs,
-                }
+                list(range(len(grid_thw))),
+                modality,
+                grid_thw,
             )
-            for k, v in mm_inputs.items():
-                if k in _mm_feature_attrs[modality]:
-                    continue
-                mm_item.set(k, _convert(v))
 
             cache_hit = False
             use_mm_cache = get_mm().enable_prefix_mm_cache and log_metrics
             if use_mm_cache:
-                mm_item.set_pad_value()
+                for item in model_mm_items:
-                mm_hash = MultiModalStaticCache.combine_hashes([mm_item.hash])
+                    item.set_pad_value()
+                item_hashes = [item.hash for item in model_mm_items]
+                mm_hash = MultiModalStaticCache.combine_hashes(item_hashes)
                 async with self.mm_cache_lock:
-                    mm_cache = self.mm_cache.get([mm_item.hash])
+                    mm_cache = self.mm_cache.get(item_hashes)
                     if mm_cache is not None:
                         mm_embedding = mm_cache.embedding
                         cache_hit = True
@@ -1801,7 +1882,7 @@ class MMEncoder:
             if mm_embedding is None:
                 forward_start = time.perf_counter()
                 with torch.inference_mode():
-                    mm_embedding: torch.Tensor = get_feature_fn([mm_item])
+                    mm_embedding: torch.Tensor = get_feature_fn(model_mm_items)
                     mm_embedding = mm_embedding.cpu()
                 if len(mm_embedding.shape) != 2:
                     mm_embedding = mm_embedding.reshape(-1, mm_embedding.shape[-1])
@@ -1810,7 +1891,8 @@ class MMEncoder:
                         time.perf_counter() - forward_start, modality=modality_str
                     )
 
-            # Per-request cache hit metrics: tokens = embedding rows, files = 1 item.
+            # Per-request cache hit metrics: tokens = embedding rows, files =
+            # logical multimodal items (not the legacy aggregate tensor count).
             if use_mm_cache and encoder_metrics_collector is not None:
                 total_tokens = int(mm_embedding.shape[0])
                 hit_tokens = total_tokens if cache_hit else 0
@@ -1818,7 +1900,9 @@ class MMEncoder:
                     hit_tokens, total_tokens, modality=modality_str
                 )
                 encoder_metrics_collector.record_cache_files(
-                    1 if cache_hit else 0, 1, modality=modality_str
+                    len(model_mm_items) if cache_hit else 0,
+                    len(model_mm_items),
+                    modality=modality_str,
                 )
 
             if use_mm_cache:
@@ -1865,7 +1949,7 @@ class MMEncoder:
                     )
 
             return (
-                _get_mm_grid_dim(mm_inputs, modality, self.model_type),
+                grid_thw,
                 mm_embedding,
                 aux_data,
             )
@@ -2223,23 +2307,21 @@ class MMEncoder:
                 )
             )
 
-            # Build mm_item (all ranks)
+            # Build model-facing items on all ranks. Owner-deferred processor
-            mm_item = MultimodalDataItem.from_dict(
+            # outputs stay per-item until encoder-DP assigns them.
-                {
+            mm_feature = _convert(_get_mm_feature(mm_inputs, modality))
-                    "modality": modality,
+            model_mm_items = self._build_mm_data_items(
-                    "feature": _convert(_get_mm_feature(mm_inputs, modality)),
+                mm_feature,
-                }
+                mm_inputs,
+                list(range(len(grid_thw))),
+                modality,
+                grid_thw,
             )
-            for k, v in mm_inputs.items():
-                if k in _mm_feature_attrs.get(modality, []):
-                    continue
-                val = _convert(v)
-                mm_item.set(k, val)
 
             async def _run_forward():
                 try:
                     with torch.inference_mode():
-                        emb = get_feature_fn([mm_item])
+                        emb = get_feature_fn(model_mm_items)
                         if len(emb.shape) != 2:
                             emb = emb.reshape(-1, emb.shape[-1])
                         # mooncake's transfer_sync is a host-side

python/sglang/srt/models/kimi_k3.py

@@ -3064,6 +3064,32 @@ class KimiK3ForConditionalGeneration(nn.Module):
             )
         self.language_model.set_dspark_layers_to_capture(layer_ids)
 
+    def preprocess_mm_for_encoder(
+        self,
+        mm_data,
+        modality,
+        config,
+        *,
+        image_processor=None,
+        use_gpu_preprocessing=False,
+    ):
+        """Prepare per-image raw inputs for owner-side EPD preprocessing."""
+        if modality != Modality.IMAGE:
+            raise ValueError("Kimi-K3 encoder mode supports image input only")
+        if image_processor is None:
+            raise ValueError("Kimi-K3 encoder preprocessing needs an image processor")
+
+        from sglang.srt.multimodal.kimi_k3_image_processing import (
+            prepare_kimi_k3_encoder_inputs,
+        )
+
+        self._encoder_image_processor = image_processor
+        return prepare_kimi_k3_encoder_inputs(
+            mm_data,
+            image_processor,
+            use_gpu_preprocessing=use_gpu_preprocessing,
+        )
+
     def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
         device = self.vision_tower.device
         target_dtype = self.vision_tower.patch_embed.proj.weight.dtype
@@ -3087,6 +3113,10 @@ class KimiK3ForConditionalGeneration(nn.Module):
 
         def materialize_item_features(image_indices: List[int]) -> torch.Tensor:
             """Materialize only the images assigned to this vision-DP rank."""
+            from sglang.srt.multimodal.encoder_preprocessing import (
+                LOCAL_PREPROCESSED_KEY,
+            )
+
             # Match the configured TP consumer count captured when the
             # tokenizer creates MmItemMemoryPool. A live attention subgroup
             # size could leave acknowledgements missing and strand the lease.
@@ -3104,6 +3134,21 @@ class KimiK3ForConditionalGeneration(nn.Module):
                     )
                 selected_items.append(item)
 
+            locally_preprocessed = [
+                item.model_specific_data.get(LOCAL_PREPROCESSED_KEY, False)
+                for item in selected_items
+            ]
+            if any(locally_preprocessed):
+                if not all(locally_preprocessed):
+                    raise ValueError(
+                        "Kimi-K3 cannot mix local preprocessed and deferred images"
+                    )
+                return materialize_multimodal_features(
+                    [item.feature for item in selected_items],
+                    device=device,
+                    dtype=target_dtype,
+                )
+
             deferred = [
                 item.model_specific_data.get(DEFERRED_PREPROCESSING_KEY)
                 for item in selected_items
@@ -3113,11 +3158,17 @@ class KimiK3ForConditionalGeneration(nn.Module):
                     raise ValueError(
                         "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):
+                    raise ValueError(
+                        "Kimi-K3 cannot mix deferred preprocessing backends"
+                    )
+                if backend == "gpu":
                     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
                     )
@@ -3132,6 +3183,19 @@ class KimiK3ForConditionalGeneration(nn.Module):
                             x, first_config["transparent_bg_config"]
                         ),
                     )
+                elif backend == "cpu":
+                    from sglang.srt.multimodal.kimi_k3_image_processing import (
+                        materialize_kimi_k3_cpu_features,
+                    )
+
+                    pixel_values = materialize_kimi_k3_cpu_features(
+                        selected_items, self._encoder_image_processor
+                    )
+                    pixel_values = pixel_values.to(device, non_blocking=True)
+                else:
+                    raise ValueError(
+                        f"Unsupported Kimi-K3 deferred preprocessing backend: {backend}"
+                    )
                 return pixel_values.to(dtype=target_dtype)
 
             features = []

python/sglang/srt/multimodal/encoder_preprocessing.py

@@ -0,0 +1,140 @@
+import hashlib
+import inspect
+from collections.abc import Mapping
+from typing import Any, Callable, Iterable, Sequence
+
+import numpy as np
+import torch
+
+from sglang.srt.managers.schedule_batch import MultimodalDataItem
+
+LOCAL_PREPROCESSED_KEY = "encoder_local_preprocessed"
+
+
+def hash_raw_encoder_item(value: Any) -> int:
+    """Hash raw CPU media including layout metadata, before owner materialization."""
+    if isinstance(value, torch.Tensor):
+        value = value.detach().cpu().contiguous().numpy()
+    elif not isinstance(value, np.ndarray):
+        from PIL import Image
+
+        if not isinstance(value, Image.Image):
+            raise TypeError(f"Unsupported raw encoder item: {type(value)}")
+        value = np.asarray(value)
+
+    value = np.ascontiguousarray(value)
+    hasher = hashlib.sha256()
+    hasher.update(value.dtype.str.encode())
+    hasher.update(repr(value.shape).encode())
+    hasher.update(memoryview(value))
+    return int.from_bytes(hasher.digest()[:8], byteorder="big", signed=False)
+
+
+class EncoderPreprocessOutput(dict):
+    """Processor output that preserves one feature object per multimodal item.
+
+    Most Hugging Face processors concatenate every image into one tensor before
+    the vision model decides its encoder-DP assignment.  A model-specific
+    ``preprocess_mm_for_encoder`` hook can return this mapping instead, allowing
+    the encoder to carry raw or partially processed items to the model.  The
+    model can then materialize only the items owned by its local vision rank.
+
+    The mapping remains compatible with existing encoder metadata helpers;
+    ``mm_items`` is an out-of-band, per-item representation used only by the
+    encoder forward/cache paths.
+    """
+
+    def __init__(
+        self,
+        values: Mapping[str, Any] | None = None,
+        *,
+        mm_items: Iterable[MultimodalDataItem],
+        item_sizes: Sequence[int] | None = None,
+        materialize_local_items: (
+            Callable[[list[MultimodalDataItem]], Sequence[torch.Tensor]] | None
+        ) = None,
+    ) -> None:
+        super().__init__(values or {})
+        self.mm_items = list(mm_items)
+        if not self.mm_items:
+            raise ValueError("EncoderPreprocessOutput requires at least one item")
+        self.item_sizes = list(item_sizes) if item_sizes is not None else None
+        if self.item_sizes is not None and len(self.item_sizes) != len(self.mm_items):
+            raise ValueError("Encoder preprocess item_sizes must match mm_items")
+        self.materialize_local_items = materialize_local_items
+
+    def local_item_indices(self, rank: int, world_size: int) -> list[int]:
+        """Return the same size-balanced owner assignment used by vision DP."""
+        if self.materialize_local_items is None:
+            return []
+        if world_size < 1 or not 0 <= rank < world_size:
+            raise ValueError(
+                f"Invalid encoder preprocess rank {rank} for world size {world_size}"
+            )
+        if world_size == 1:
+            return list(range(len(self.mm_items)))
+        if self.item_sizes is None:
+            raise ValueError(
+                "Owner-side encoder preprocessing requires per-item load sizes"
+            )
+
+        from sglang.srt.multimodal.mm_utils import get_dp_encoder_lb_assignment
+
+        shuffled, counts, _ = get_dp_encoder_lb_assignment(self.item_sizes, world_size)
+        start = sum(counts[:rank])
+        return shuffled[start : start + counts[rank]]
+
+    def materialize_for_rank(self, rank: int, world_size: int) -> None:
+        """Materialize only this vision-DP rank's items in-place."""
+        indices = self.local_item_indices(rank, world_size)
+        if not indices:
+            return
+        items = [self.mm_items[index] for index in indices]
+        materialize = self.materialize_local_items
+        assert materialize is not None
+        features = list(materialize(items))
+        if len(features) != len(items):
+            raise ValueError(
+                "Encoder local materializer must return one feature per item"
+            )
+        for item, feature in zip(items, features):
+            item.feature = feature
+            item.model_specific_data[LOCAL_PREPROCESSED_KEY] = True
+
+
+def get_encoder_preprocessed_items(
+    processor_output: Mapping[str, Any],
+) -> list[MultimodalDataItem] | None:
+    if isinstance(processor_output, EncoderPreprocessOutput):
+        return processor_output.mm_items
+    return None
+
+
+def invoke_encoder_preprocessor(
+    preprocessor,
+    mm_data,
+    modality,
+    config,
+    **available_context,
+):
+    """Call a model hook with only the optional context it declares.
+
+    Existing hooks keep their three-argument contract. New model integrations
+    can request shared processor state or backend policy by adding named
+    keyword-only parameters, without requiring encode-server model branches.
+    """
+    parameters = inspect.signature(preprocessor).parameters
+    accepts_kwargs = any(
+        parameter.kind == inspect.Parameter.VAR_KEYWORD
+        for parameter in parameters.values()
+    )
+    context = (
+        available_context
+        if accepts_kwargs
+        else {
+            name: value
+            for name, value in available_context.items()
+            if name in parameters
+        }
+    )
+    return preprocessor(mm_data, modality, config, **context)

python/sglang/srt/multimodal/kimi_k3_image_processing.py

@@ -1,3 +1,5 @@
+import functools
+import math
 from typing import Union
 
 import numpy as np
@@ -7,6 +9,195 @@ from PIL import Image
 DEFERRED_PREPROCESSING_KEY = "kimi_k3_deferred_preprocessing"
 
 
+def prepare_kimi_k3_encoder_inputs(
+    images, image_processor, *, use_gpu_preprocessing=False
+):
+    """Keep K3 EPD images raw until the vision-DP owner is known.
+
+    The lightweight NaViT shape calculation runs on every encoder rank.  The
+    expensive resize, normalization, and patchification remain deferred in each
+    item and are executed by ``KimiK3ForConditionalGeneration`` only for images
+    assigned to the local vision rank.
+    """
+    from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
+    from sglang.srt.multimodal.encoder_preprocessing import (
+        EncoderPreprocessOutput,
+        hash_raw_encoder_item,
+    )
+    from sglang.srt.multimodal.processors.kimi_k25 import (
+        _grid_thw_from_resize_config,
+        navit_resize_config,
+    )
+
+    media_proc_cfg = getattr(image_processor, "media_proc_cfg", None)
+    if not isinstance(media_proc_cfg, dict):
+        raise ValueError(
+            "Kimi-K3 EPD owner-side preprocessing requires "
+            "image_processor.media_proc_cfg"
+        )
+
+    required = (
+        "patch_size",
+        "merge_kernel_size",
+        "in_patch_limit",
+        "patch_limit_on_one_side",
+        "image_mean",
+        "image_std",
+    )
+    missing = [name for name in required if name not in media_proc_cfg]
+    if missing:
+        raise ValueError(
+            "Kimi-K3 image processor is missing deferred-preprocessing config: "
+            + ", ".join(missing)
+        )
+
+    concrete_images = []
+    for image in images:
+        if isinstance(image, dict):
+            if image.get("type") != "image" or "image" not in image:
+                raise ValueError(f"Unsupported Kimi-K3 encoder media item: {image}")
+            image = image["image"]
+        concrete_images.append(image)
+
+    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"),
+    }
+
+    items = []
+    grids = []
+    original_image_sizes = []
+    for image in concrete_images:
+        width, height = (
+            (int(image.shape[-1]), int(image.shape[-2]))
+            if isinstance(image, torch.Tensor)
+            else image.size
+        )
+        resize_config = navit_resize_config(
+            width,
+            height,
+            patch_size,
+            merge_kernel_size,
+            int(media_proc_cfg["in_patch_limit"]),
+            int(media_proc_cfg["patch_limit_on_one_side"]),
+            media_proc_cfg.get("fixed_output_tokens"),
+        )
+        grid_thw = _grid_thw_from_resize_config(resize_config, patch_size)
+        grid_tensor = torch.tensor([grid_thw], dtype=torch.int64)
+        item = MultimodalDataItem(
+            modality=Modality.IMAGE,
+            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,
+                },
+            },
+        )
+        if not use_gpu_preprocessing:
+            item.set_hash(hash_raw_encoder_item(image))
+        items.append(item)
+        grids.append(grid_thw)
+        original_image_sizes.append([width, height])
+
+    grid_thws = torch.tensor(grids, dtype=torch.int64)
+    return EncoderPreprocessOutput(
+        {
+            # Preserve the conventional feature key for cache/accounting code;
+            # EncoderPreprocessOutput.mm_items is the authoritative per-item view.
+            "pixel_values": [item.feature for item in items],
+            "grid_thws": grid_thws,
+            "original_image_sizes": original_image_sizes,
+        },
+        mm_items=items,
+        item_sizes=[math.prod(grid) for grid in grids],
+        materialize_local_items=(
+            None
+            if use_gpu_preprocessing
+            else functools.partial(
+                materialize_kimi_k3_cpu_item_features,
+                image_processor=image_processor,
+            )
+        ),
+    )
+
+
+def materialize_kimi_k3_cpu_features(items, image_processor) -> torch.Tensor:
+    """Run the checkpoint's exact processor only on locally owned images."""
+    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(
+                    "Kimi-K3 deferred CPU preprocessing expects PIL or uint8 tensors"
+                )
+            image = to_hwc_uint8(image).numpy()
+            channels = image.shape[-1]
+            if channels == 1:
+                image = Image.fromarray(image[..., 0], mode="L")
+            elif channels == 3:
+                image = Image.fromarray(image, mode="RGB")
+            elif channels == 4:
+                image = Image.fromarray(image, mode="RGBA")
+            else:
+                raise ValueError(f"Unsupported Kimi-K3 image channel count: {channels}")
+        medias.append({"type": "image", "image": image})
+
+    output = image_processor.preprocess(medias, return_tensors="pt")
+    expected_grids = torch.cat(
+        [item.model_specific_data["grid_thws"] for item in items], dim=0
+    )
+    if not torch.equal(output["grid_thws"].cpu(), expected_grids.cpu()):
+        raise ValueError("Kimi-K3 deferred CPU preprocessing produced wrong grids")
+    return output["pixel_values"]
+
+
+def materialize_kimi_k3_cpu_item_features(items, image_processor) -> list[torch.Tensor]:
+    """Return exact checkpoint-processor features split by logical image."""
+    pixel_values = materialize_kimi_k3_cpu_features(items, image_processor)
+    patch_counts = [
+        math.prod(item.model_specific_data["grid_thws"][0].tolist()) for item in items
+    ]
+    if sum(patch_counts) != pixel_values.shape[0]:
+        raise ValueError(
+            "Kimi-K3 processor feature length does not match image grids: "
+            f"{pixel_values.shape[0]} != {sum(patch_counts)}"
+        )
+    return list(pixel_values.split(patch_counts))
+
+
+def to_hwc_uint8(image: Union[torch.Tensor, Image.Image]) -> torch.Tensor:
+    """Stage an exact CPU image without doing resize/normalize work."""
+    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)
+        return torch.from_numpy(array)
+
+    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(-1)
+    elif image.dim() == 3 and image.shape[0] in (1, 3, 4):
+        image = image.permute(1, 2, 0)
+    if image.shape[-1] == 1:
+        image = image.repeat(1, 1, 3)
+    return image.cpu().contiguous()
+
+
 def to_chw_uint8(
     image: Union[torch.Tensor, Image.Image],
     device: torch.device | str | None = None,

test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py

@@ -22,11 +22,23 @@ from sglang.srt.disaggregation.encode_receiver import (
     _select_mm_processor_prompt,
 )
 from sglang.srt.disaggregation.encode_server import MMEncoder, _get_mm_grid_dim
-from sglang.srt.managers.schedule_batch import Modality
+from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
 from sglang.srt.managers.tokenizer_manager import (
     _reject_missing_dispatched_encoder_embedding,
 )
 from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
+from sglang.srt.multimodal.encoder_preprocessing import (
+    LOCAL_PREPROCESSED_KEY,
+    EncoderPreprocessOutput,
+    get_encoder_preprocessed_items,
+    hash_raw_encoder_item,
+    invoke_encoder_preprocessor,
+)
+from sglang.srt.multimodal.kimi_k3_image_processing import (
+    DEFERRED_PREPROCESSING_KEY,
+    materialize_kimi_k3_cpu_features,
+    prepare_kimi_k3_encoder_inputs,
+)
 from sglang.srt.runtime_context import get_context
 from sglang.srt.server_args import resolve_encoder_transfer_backend
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -138,6 +150,189 @@ def test_kimi_k3_encoder_passes_media_dicts_to_image_processor():
     assert kwargs == {"return_tensors": "pt"}
 
 
+def _kimi_k3_image_processor():
+    return SimpleNamespace(
+        media_proc_cfg={
+            "patch_size": 2,
+            "merge_kernel_size": 2,
+            "in_patch_limit": 1024,
+            "patch_limit_on_one_side": 64,
+            "fixed_output_tokens": None,
+            "image_mean": [0.5, 0.5, 0.5],
+            "image_std": [0.5, 0.5, 0.5],
+            "transparent_bg_config": {"type": "white"},
+        }
+    )
+
+
+def test_kimi_k3_epd_preprocess_preserves_raw_per_image_items():
+    first = Image.new("RGB", (8, 6), color=(1, 2, 3))
+    second = Image.new("RGB", (5, 9), color=(4, 5, 6))
+
+    output = prepare_kimi_k3_encoder_inputs(
+        [
+            {"type": "image", "image": first},
+            {"type": "image", "image": second},
+        ],
+        _kimi_k3_image_processor(),
+    )
+
+    items = get_encoder_preprocessed_items(output)
+    assert isinstance(output, EncoderPreprocessOutput)
+    assert len(items) == 2
+    assert output["original_image_sizes"] == [[8, 6], [5, 9]]
+    assert output["grid_thws"].tolist() == [[1, 4, 4], [1, 6, 4]]
+    for item, image in zip(items, (first, second)):
+        assert item.modality == Modality.IMAGE
+        assert item.feature is image
+        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]
+
+
+def test_kimi_k3_epd_model_preprocessor_receives_image_processor():
+    image = Image.new("RGB", (8, 6), color=(1, 2, 3))
+    image_processor = _kimi_k3_image_processor()
+    image_processor.preprocess = lambda medias, return_tensors: {
+        "pixel_values": torch.zeros(16, 12),
+        "grid_thws": torch.tensor([[1, 4, 4]]),
+    }
+    calls = []
+
+    def model_preprocessor(
+        mm_data,
+        modality,
+        config,
+        *,
+        image_processor=None,
+        use_gpu_preprocessing=False,
+    ):
+        calls.append(
+            (mm_data, modality, config, image_processor, use_gpu_preprocessing)
+        )
+        return prepare_kimi_k3_encoder_inputs(mm_data, image_processor)
+
+    encoder = _encoder()
+    encoder.image_processor = image_processor
+    encoder.use_image_processor_gpu = False
+    encoder.vision_config = {"image": {"return_tensors": "pt"}}
+    encoder._flatten_and_load_images = AsyncMock(return_value=[image])
+    encoder.preproc_executor = ThreadPoolExecutor(max_workers=1)
+    try:
+        with patch(
+            "sglang.srt.disaggregation.encode_server.get_parallel",
+            return_value=SimpleNamespace(attn_tp_rank=0, attn_tp_size=1),
+        ):
+            output = asyncio.run(
+                encoder._process_image_items([image], model_preprocessor)
+            )
+    finally:
+        encoder.preproc_executor.shutdown()
+
+    assert len(calls) == 1
+    assert calls[0][0][0] == {"type": "image", "image": image}
+    assert calls[0][1:] == (
+        Modality.IMAGE,
+        encoder.vision_config,
+        image_processor,
+        False,
+    )
+    assert len(get_encoder_preprocessed_items(output)) == 1
+
+
+def test_encoder_preprocessor_context_keeps_legacy_hooks_compatible():
+    calls = []
+
+    def legacy_hook(mm_data, modality, config):
+        calls.append((mm_data, modality, config))
+        return {"ok": True}
+
+    result = invoke_encoder_preprocessor(
+        legacy_hook,
+        ["image"],
+        Modality.IMAGE,
+        {"image": {}},
+        image_processor=object(),
+        use_gpu_preprocessing=True,
+    )
+
+    assert result == {"ok": True}
+    assert calls == [(["image"], Modality.IMAGE, {"image": {}})]
+
+
+def test_kimi_k3_epd_default_cpu_materialization_is_owner_only_and_exact():
+    class RecordingImageProcessor:
+        media_proc_cfg = _kimi_k3_image_processor().media_proc_cfg
+
+        def __init__(self):
+            self.calls = []
+
+        def preprocess(self, medias, return_tensors):
+            self.calls.append(medias)
+            features = [
+                torch.full((4, 3, 2, 2), media["image"].getpixel((0, 0))[0])
+                for media in medias
+            ]
+            grids = torch.tensor([[1, 2, 2]] * len(medias))
+            return {"pixel_values": torch.cat(features), "grid_thws": grids}
+
+    processor = RecordingImageProcessor()
+    images = [Image.new("RGB", (4, 4), color=(value, 0, 0)) for value in (7, 11)]
+    output = prepare_kimi_k3_encoder_inputs(images, processor)
+    items = get_encoder_preprocessed_items(output)
+
+    materialized = materialize_kimi_k3_cpu_features([items[1]], processor)
+
+    assert len(processor.calls) == 1
+    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"
+
+
+def test_encoder_preprocess_materializes_only_local_size_balanced_items():
+    items = [
+        MultimodalDataItem(
+            modality=Modality.IMAGE,
+            feature=torch.tensor([value], dtype=torch.uint8),
+        )
+        for value in (3, 5, 7)
+    ]
+    calls = []
+
+    def materialize(selected):
+        calls.append(selected)
+        return [item.feature.float() + 10 for item in selected]
+
+    output = EncoderPreprocessOutput(
+        {"pixel_values": [item.feature for item in items]},
+        mm_items=items,
+        item_sizes=[8, 5, 3],
+        materialize_local_items=materialize,
+    )
+
+    output.materialize_for_rank(rank=1, world_size=2)
+
+    assert calls == [[items[1], items[2]]]
+    assert items[0].feature.tolist() == [3]
+    assert items[1].feature.tolist() == [15.0]
+    assert items[2].feature.tolist() == [17.0]
+    assert LOCAL_PREPROCESSED_KEY not in items[0].model_specific_data
+    assert items[1].model_specific_data[LOCAL_PREPROCESSED_KEY]
+    assert items[2].model_specific_data[LOCAL_PREPROCESSED_KEY]
+
+
+def test_raw_encoder_hash_includes_shape_and_dtype():
+    flat = torch.arange(12, dtype=torch.uint8)
+
+    assert hash_raw_encoder_item(flat.reshape(2, 2, 3)) != hash_raw_encoder_item(
+        flat.reshape(3, 2, 2)
+    )
+    assert hash_raw_encoder_item(flat) != hash_raw_encoder_item(flat.to(torch.int16))
+
+
 @pytest.mark.parametrize(
     ("use_image_processor_gpu", "expected_decode_mode"),
     [(False, False), (True, "nvjpeg_fancy")],
@@ -233,6 +428,67 @@ def test_kimi_k3_encoder_splits_cross_request_batch_into_single_grid_items():
     torch.testing.assert_close(torch.cat(output), embeddings)
 
 
+def test_encoder_preprocessed_items_follow_dp_owner_selection_order():
+    encoder = _encoder()
+    grid_thws = torch.tensor([[1, 2, 2], [1, 2, 4], [1, 4, 2]])
+    items = [
+        MultimodalDataItem(
+            modality=Modality.IMAGE,
+            feature=torch.full((3, i + 2, i + 3), i, dtype=torch.uint8),
+            model_specific_data={"grid_thws": grid_thws[i : i + 1]},
+        )
+        for i in range(3)
+    ]
+    mm_inputs = EncoderPreprocessOutput(
+        {"pixel_values": [item.feature for item in items], "grid_thws": grid_thws},
+        mm_items=items,
+    )
+    embeddings = torch.arange(4, dtype=torch.float32).reshape(4, 1)
+    captured = {}
+
+    def get_feature_fn(selected_items):
+        captured["items"] = selected_items
+        return embeddings
+
+    output = encoder._encode_missing(
+        mm_inputs["pixel_values"],
+        mm_inputs,
+        indices=[2, 0],
+        modality=Modality.IMAGE,
+        get_feature_fn=get_feature_fn,
+        grid_thw=grid_thws,
+        keep_on_gpu=True,
+    )
+
+    assert captured["items"] == [items[2], items[0]]
+    assert [part.shape[0] for part in output] == [2, 1]
+    torch.testing.assert_close(torch.cat(output), embeddings[:3])
+
+
+def test_encoder_preprocessed_items_hash_individually():
+    encoder = _encoder()
+    grid_thws = torch.tensor([[1, 2, 2], [1, 2, 4]])
+    items = [
+        MultimodalDataItem(
+            modality=Modality.IMAGE,
+            feature=torch.full((3, 2, 2), value, dtype=torch.uint8),
+            model_specific_data={"grid_thws": grid_thws[i : i + 1]},
+        )
+        for i, value in enumerate((17, 29))
+    ]
+    mm_inputs = EncoderPreprocessOutput(
+        {"pixel_values": [item.feature for item in items], "grid_thws": grid_thws},
+        mm_items=items,
+    )
+
+    hashes = encoder._calculate_hashes_from_features(
+        mm_inputs["pixel_values"], grid_thws, Modality.IMAGE, mm_inputs
+    )
+
+    assert hashes == [item.hash for item in items]
+    assert hashes[0] != hashes[1]
+
+
 def test_kimi_k3_encoder_only_wrapper_guards_language_tower_hooks():
     model = SimpleNamespace(language_model=None)
 

test/registered/unit/models/test_kimi_k3_vision.py

@@ -518,6 +518,8 @@ def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch):
     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,