From def8cb0f02bc3d2f309f0d99bbd517e356e2f8ed Mon Sep 17 00:00:00 2001 From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Date: Sun, 28 Jun 2026 00:08:54 -0700 Subject: [PATCH 1/2] refactor: make mm_data_mooncake msgpack-native (drop PickleWrapper) --- .../srt/disaggregation/encode_receiver.py | 44 +++++----- python/sglang/srt/managers/io_struct.py | 20 +++-- .../unit/managers/test_io_struct_msgpack.py | 86 +++++++++++++++++++ 3 files changed, 123 insertions(+), 27 deletions(-) create mode 100644 test/registered/unit/managers/test_io_struct_msgpack.py diff --git a/python/sglang/srt/disaggregation/encode_receiver.py b/python/sglang/srt/disaggregation/encode_receiver.py index bbe120c9a1cc..273368a8621f 100644 --- a/python/sglang/srt/disaggregation/encode_receiver.py +++ b/python/sglang/srt/disaggregation/encode_receiver.py @@ -30,7 +30,11 @@ get_mooncake_transfer_engine, ) from sglang.srt.environ import envs -from sglang.srt.managers.io_struct import GenerateReqInput, TokenizedGenerateReqInput +from sglang.srt.managers.io_struct import ( + GenerateReqInput, + MooncakeMMUrlItem, + TokenizedGenerateReqInput, +) from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors from sglang.srt.managers.schedule_batch import Modality, Req from sglang.srt.server_args import ServerArgs @@ -945,7 +949,7 @@ async def _send_encode_and_rdma_request(self): for modality in modalities: assigned_nums = self.num_items_assigned[modality] num_parts = modality_num_parts[modality] - mm_data_modality = [d for d in mm_data_all if d["modality"] == modality] + mm_data_modality = [d for d in mm_data_all if d.modality == modality] cum_num_items = 0 cum_idx = 0 for idx, assigned_num in enumerate(assigned_nums): @@ -957,7 +961,7 @@ async def _send_encode_and_rdma_request(self): { "encoder_idx": idx, "mm_items": [ - d["url"] + d.url for d in mm_data_modality[ cum_num_items : cum_num_items + assigned_num ] @@ -1568,7 +1572,7 @@ async def recv_mm_data( self.context, zmq.PULL, host=self.host ) mm_data = self._extract_url_data(request_obj) - modalities = [m.get("modality") for m in mm_data] + modalities = [m.modality for m in mm_data] logger.info( f"[{req_id}] Sending encode request to E, " f"modalities={modalities}, num_items={len(mm_data)}" @@ -1916,7 +1920,7 @@ def _assign_items_by_modality( Assign multimodal items across encoders by modality with cross-modality load balancing. Args: - mm_data: List of multimodal data items, each with a "modality" key + mm_data: List of multimodal data items, each with a modality encoder_num: Number of encoders random_shuffle: Whether to shuffle the encoder indices @@ -1928,14 +1932,14 @@ def _assign_items_by_modality( if random_shuffle: random.shuffle(encode_idx) # Get unique modalities with order preserved - modalities = list(dict.fromkeys(mm_item.get("modality") for mm_item in mm_data)) + modalities = list(dict.fromkeys(mm_item.modality for mm_item in mm_data)) # Use OrderedDict to explicitly maintain modality order num_items_assigned = OrderedDict() current_offset = 0 for modality in modalities: mm_data_modality = [ - mm_item for mm_item in mm_data if mm_item.get("modality") == modality + mm_item for mm_item in mm_data if mm_item.modality == modality ] num_items = len(mm_data_modality) if num_items == 0: @@ -1955,7 +1959,7 @@ def _assign_items_by_modality( return num_items_assigned - def _extract_url_data(self, request_obj) -> List[Dict]: + def _extract_url_data(self, request_obj) -> List[MooncakeMMUrlItem]: def flatten_mm_items(items): if not isinstance(items, list): return [items] @@ -1987,10 +1991,10 @@ def to_raw_url(mm_item): mm_items = flatten_mm_items(mm_items) for mm_item in mm_items: mm_data.append( - { - "url": to_raw_url(mm_item), - "modality": modality, - } + MooncakeMMUrlItem( + url=to_raw_url(mm_item), + modality=modality, + ) ) return mm_data @@ -2080,7 +2084,7 @@ async def encode( effective_urls = encode_urls if encode_urls is not None else self.encode_urls # get unique modalities with order preserved - modalities = [mm_item.get("modality") for mm_item in mm_data] + modalities = [mm_item.modality for mm_item in mm_data] modalities = list(dict.fromkeys(modalities)) encode_requests = [] @@ -2098,7 +2102,7 @@ async def encode( for modality in modalities: num_items_assigned_modality = num_items_assigned.get(modality) mm_data_modality = [ - mm_item for mm_item in mm_data if mm_item.get("modality") == modality + mm_item for mm_item in mm_data if mm_item.modality == modality ] num_parts = modality_num_parts[modality] @@ -2114,7 +2118,7 @@ async def encode( "encoder_idx": idx, "encoder_url": effective_urls[idx], "mm_items": [ - mm_item.get("url") + mm_item.url for mm_item in mm_data_modality[ cum_num_items : cum_num_items + assigned_num ] @@ -2241,18 +2245,16 @@ async def encode( effective_urls = encode_urls if encode_urls is not None else self.encode_urls - # gRPC currently only supports image; flatten new dict formats to simple lists - if mm_data and isinstance(mm_data[0], dict): + # gRPC currently only supports image; flatten typed multimodal items to simple lists + if mm_data and isinstance(mm_data[0], MooncakeMMUrlItem): non_image = [ - item.get("modality") - for item in mm_data - if item.get("modality") != Modality.IMAGE + item.modality for item in mm_data if item.modality != Modality.IMAGE ] if non_image: raise NotImplementedError( f"gRPC encode only supports IMAGE modality, got: {non_image}" ) - img_data = [item.get("url") for item in mm_data] + img_data = [item.url for item in mm_data] else: img_data = mm_data if isinstance(num_items_assigned, dict): diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 3772068100b7..64ac8e7270fc 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -148,6 +148,17 @@ class SessionParams(msgspec.Struct, kw_only=True, array_like=True): ] +class MooncakeMMUrlItem(msgspec.Struct, array_like=True): + """One multimodal item captured at tokenizer dispatch for mooncake encode routing. + + ``url`` is the raw URL string from MMReceiverBase._extract_url_data + (ImageData.url when present); ``modality`` tags which encoder pool handles it. + """ + + url: str + modality: Modality + + @dataclass class GenerateReqInput: # Request ID(s). If omitted, generated during normalization. For batch @@ -845,10 +856,9 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): need_wait_for_mm_inputs: Optional[bool] = None num_items_assigned: Optional[Dict[Modality, List[int]]] = None - # Pickled Optional[List[{"url": MultimodalDataInputItem, "modality": Modality}]] - # from MMReceiverBase._extract_url_data. "url" is ImageData.url, - # dict["url"] when present, or the original raw multimodal item. - mm_data_mooncake: Optional[PickleWrapper] = None + # Multimodal URL snapshot from MMReceiverBase._extract_url_data for mooncake. + # Each item carries the raw URL/payload and modality captured at tokenizer dispatch. + mm_data_mooncake: Optional[List[MooncakeMMUrlItem]] = None # Encoder URL snapshot frozen at tokenizer-side dispatch time so that # encoder_idx assignments stay consistent in the scheduler subprocess. # Internal IPC only. @@ -863,12 +873,10 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): def wrap_pickle_fields(self): self.mm_inputs = wrap_as_pickle(self.mm_inputs) - self.mm_data_mooncake = wrap_as_pickle(self.mm_data_mooncake) self.time_stats = wrap_as_pickle(self.time_stats) def unwrap_pickle_fields(self): self.mm_inputs = unwrap_from_pickle(self.mm_inputs) - self.mm_data_mooncake = unwrap_from_pickle(self.mm_data_mooncake) self.time_stats = unwrap_from_pickle(self.time_stats) diff --git a/test/registered/unit/managers/test_io_struct_msgpack.py b/test/registered/unit/managers/test_io_struct_msgpack.py new file mode 100644 index 000000000000..90ca6c90c4b8 --- /dev/null +++ b/test/registered/unit/managers/test_io_struct_msgpack.py @@ -0,0 +1,86 @@ +import unittest + +from sglang.srt.managers.io_struct import ( + MooncakeMMUrlItem, + TokenizedGenerateReqInput, + msgpack_decode, + msgpack_encode, +) +from sglang.srt.managers.schedule_batch import Modality +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.test.ci.ci_register import ( + register_amd_ci, + register_cpu_ci, + register_cuda_ci, +) + +register_cuda_ci(est_time=8, stage="base-b", runner_config="1-gpu-large") +register_amd_ci(est_time=8, suite="stage-b-test-1-gpu-small-amd") +register_cpu_ci(est_time=8, suite="base-c-test-cpu") + + +class TestTokenizedGenerateReqInputMsgpack(unittest.TestCase): + def _make_request(self, mm_data_mooncake): + return TokenizedGenerateReqInput( + input_text="", + input_ids=None, + input_embeds=None, + mm_inputs=None, + token_type_ids=None, + sampling_params=SamplingParams(), + return_logprob=False, + logprob_start_len=0, + top_logprobs_num=0, + token_ids_logprob=None, + stream=False, + mm_data_mooncake=mm_data_mooncake, + ) + + def _round_trip(self, req): + req.wrap_pickle_fields() + decoded = msgpack_decode(msgpack_encode(req)) + decoded.unwrap_pickle_fields() + return decoded + + def test_mm_data_mooncake_round_trips_populated(self): + items = [ + MooncakeMMUrlItem(url="http://x/a.jpg", modality=Modality.IMAGE), + MooncakeMMUrlItem(url="http://x/b.mp4", modality=Modality.VIDEO), + ] + + decoded = self._round_trip(self._make_request(items)) + + self.assertEqual(decoded.mm_data_mooncake, items) + self.assertEqual(decoded.mm_data_mooncake[0].url, "http://x/a.jpg") + self.assertEqual(decoded.mm_data_mooncake[0].modality, Modality.IMAGE) + self.assertEqual(decoded.mm_data_mooncake[1].url, "http://x/b.mp4") + self.assertEqual(decoded.mm_data_mooncake[1].modality, Modality.VIDEO) + + def test_mm_data_mooncake_round_trips_none(self): + decoded = self._round_trip(self._make_request(None)) + + self.assertIsNone(decoded.mm_data_mooncake) + + def test_mm_data_mooncake_round_trips_empty_list(self): + decoded = self._round_trip(self._make_request([])) + + self.assertEqual(decoded.mm_data_mooncake, []) + + def test_mm_data_mooncake_round_trips_all_modalities(self): + for modality in (Modality.IMAGE, Modality.VIDEO, Modality.AUDIO): + with self.subTest(modality=modality): + items = [ + MooncakeMMUrlItem( + url=f"http://x/{modality.name.lower()}", + modality=modality, + ) + ] + + decoded = self._round_trip(self._make_request(items)) + + self.assertEqual(decoded.mm_data_mooncake[0].url, items[0].url) + self.assertEqual(decoded.mm_data_mooncake[0].modality, modality) + + +if __name__ == "__main__": + unittest.main() From 95b01f7aa01defc97e3dfe30dbdcf8216fe8f09a Mon Sep 17 00:00:00 2001 From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Date: Sun, 28 Jun 2026 00:15:49 -0700 Subject: [PATCH 2/2] refactor: widen mm_data_mooncake url to Any for non-string payloads --- python/sglang/srt/managers/io_struct.py | 11 ++++++++--- .../unit/managers/test_io_struct_msgpack.py | 18 ++++++++++++++++++ 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 64ac8e7270fc..cb82b795df98 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -151,11 +151,16 @@ class SessionParams(msgspec.Struct, kw_only=True, array_like=True): class MooncakeMMUrlItem(msgspec.Struct, array_like=True): """One multimodal item captured at tokenizer dispatch for mooncake encode routing. - ``url`` is the raw URL string from MMReceiverBase._extract_url_data - (ImageData.url when present); ``modality`` tags which encoder pool handles it. + ``url`` is the raw payload from MMReceiverBase._extract_url_data: the URL + string (ImageData.url, or dict["url"] when present) in the common case, or + the original raw multimodal item (e.g. a dict for processor/precomputed + formats, or bytes) when no URL is available. Typed as Any because msgspec + cannot express a union mixing str and bytes; the field still round-trips + over msgpack since every passthrough payload is a native msgpack type. + ``modality`` tags which encoder pool handles it. """ - url: str + url: Any modality: Modality diff --git a/test/registered/unit/managers/test_io_struct_msgpack.py b/test/registered/unit/managers/test_io_struct_msgpack.py index 90ca6c90c4b8..638d34a3210e 100644 --- a/test/registered/unit/managers/test_io_struct_msgpack.py +++ b/test/registered/unit/managers/test_io_struct_msgpack.py @@ -56,6 +56,24 @@ def test_mm_data_mooncake_round_trips_populated(self): self.assertEqual(decoded.mm_data_mooncake[1].url, "http://x/b.mp4") self.assertEqual(decoded.mm_data_mooncake[1].modality, Modality.VIDEO) + def test_mm_data_mooncake_round_trips_dict_payload(self): + # _extract_url_data passes raw dict multimodal items through unchanged + # when they carry no "url" key, so the field must survive non-string + # payloads over msgpack. + items = [ + MooncakeMMUrlItem( + url={"image": "data:base64", "detail": "high"}, + modality=Modality.IMAGE, + ), + ] + + decoded = self._round_trip(self._make_request(items)) + + self.assertEqual( + decoded.mm_data_mooncake[0].url, {"image": "data:base64", "detail": "high"} + ) + self.assertEqual(decoded.mm_data_mooncake[0].modality, Modality.IMAGE) + def test_mm_data_mooncake_round_trips_none(self): decoded = self._round_trip(self._make_request(None))