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

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,