diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 7859fd8cd..3f58ca644 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -324,6 +324,8 @@ class Engine(EngineScoreMixin, EngineBase): image_data: Optional[MultimodalDataInputFormat] = None, audio_data: Optional[MultimodalDataInputFormat] = None, video_data: Optional[MultimodalDataInputFormat] = None, + # See GenerateReqInput.mm_hashes / async_generate for the contract. + mm_hashes: Optional[Union[List[str], List[List[str]]]] = None, return_logprob: Optional[Union[List[bool], bool]] = False, logprob_start_len: Optional[Union[List[int], int]] = None, top_logprobs_num: Optional[Union[List[int], int]] = None, @@ -361,6 +363,7 @@ class Engine(EngineScoreMixin, EngineBase): image_data=image_data, audio_data=audio_data, video_data=video_data, + mm_hashes=mm_hashes, return_logprob=return_logprob, logprob_start_len=logprob_start_len, top_logprobs_num=top_logprobs_num, @@ -416,6 +419,13 @@ class Engine(EngineScoreMixin, EngineBase): image_data: Optional[MultimodalDataInputFormat] = None, audio_data: Optional[MultimodalDataInputFormat] = None, video_data: Optional[MultimodalDataInputFormat] = None, + # Optional per-image hashes the caller has already computed (hex strings, + # one per image in `image_data`). When supplied, each MultimodalDataItem's + # `hash` is initialised from this list and `set_pad_value` skips the + # internal `hash_feature()` recompute. Intended for external KV routers + # that compute their own per-image hash for routing decisions and need + # sglang's prefix-cache key to align. See GenerateReqInput.mm_hashes. + mm_hashes: Optional[Union[List[str], List[List[str]]]] = None, return_logprob: Optional[Union[List[bool], bool]] = False, logprob_start_len: Optional[Union[List[int], int]] = None, top_logprobs_num: Optional[Union[List[int], int]] = None, @@ -453,6 +463,7 @@ class Engine(EngineScoreMixin, EngineBase): image_data=image_data, audio_data=audio_data, video_data=video_data, + mm_hashes=mm_hashes, return_logprob=return_logprob, logprob_start_len=logprob_start_len, top_logprobs_num=top_logprobs_num, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index d9d102202..8b6251171 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -158,6 +158,15 @@ class GenerateReqInput(BaseReq): video_data: Optional[MultimodalDataInputFormat] = None # The audio input. Like image data, it can be a file name, a url, or base64 encoded string. audio_data: Optional[MultimodalDataInputFormat] = None + # Optional per-image hashes the caller has already computed (hex strings, + # one per image in `image_data`). When supplied, each MultimodalDataItem's + # `hash` is initialised from this list and `set_pad_value` skips the + # internal `hash_feature()` recompute, so the resulting `pad_value` is + # deterministic from the caller's hash. Intended for external KV routers + # that compute their own per-image hash for routing decisions and need + # sglang's prefix-cache key to align. When unset, behavior is unchanged + # (sglang hashes the processor feature tensor). + mm_hashes: Optional[Union[List[str], List[List[str]]]] = None # Whether to extract and process audio from video inputs. use_audio_in_video: bool = False # The sampling_params. See descriptions below. diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 9ee9af094..2168f053c 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -840,6 +840,37 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): token_type_ids = mm_inputs.token_type_ids if not isinstance(token_type_ids, list): token_type_ids = token_type_ids.flatten().tolist() + # Caller-supplied per-image hashes (external KV routers, e.g. + # routing-aware orchestrators that compute a content-addressed + # hash before dispatch). Setting MultimodalDataItem.hash here + # short-circuits the internal hash_feature() recompute inside + # set_pad_value(), making the derived pad_value deterministic + # from the caller's hash. That alignment lets the router's + # routing decision agree with sglang's prefix-cache key for + # the same image. On any per-item parse error or list-length + # mismatch we fall back to the internal recompute so a + # malformed mm_hashes never blocks a request. + caller_mm_hashes = getattr(obj, "mm_hashes", None) + if caller_mm_hashes and mm_inputs and mm_inputs.mm_items: + if len(caller_mm_hashes) != len(mm_inputs.mm_items): + logger.warning( + "mm_hashes length (%d) != mm_items length (%d); " + "ignoring caller hashes for this request.", + len(caller_mm_hashes), + len(mm_inputs.mm_items), + ) + else: + for item, hex_hash in zip(mm_inputs.mm_items, caller_mm_hashes): + if not isinstance(item, MultimodalDataItem): + continue + try: + item.hash = int(hex_hash, 16) + except (TypeError, ValueError): + logger.warning( + "Ignoring malformed mm_hashes entry %r; " + "this item will fall back to hash_feature().", + hex_hash, + ) if ( envs.SGLANG_MM_PRECOMPUTE_HASH.get() and mm_inputs diff --git a/test/registered/unit/managers/test_mm_hashes.py b/test/registered/unit/managers/test_mm_hashes.py new file mode 100644 index 000000000..85fb500ec --- /dev/null +++ b/test/registered/unit/managers/test_mm_hashes.py @@ -0,0 +1,82 @@ +"""Tests for caller-supplied mm_hashes plumbing. + +Verifies the contract that: + 1. GenerateReqInput.mm_hashes is an optional list of hex strings. + 2. MultimodalDataItem.set_pad_value() honors a pre-set hash and does NOT + overwrite it via hash_feature(). + 3. The derived pad_value is deterministic across requests with identical + mm_hashes — the property external KV routers depend on. + +The wiring step that copies GenerateReqInput.mm_hashes into per-item +MultimodalDataItem.hash lives in tokenizer_manager.py and is exercised by +the e2e serve tests; this file pins the unit-level invariants the wiring +relies on. +""" + +import unittest +from unittest.mock import patch + +from sglang.srt.managers.io_struct import GenerateReqInput +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + _compute_pad_value, +) +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=2, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=2, suite="stage-b-test-1-gpu-small-amd") + + +class TestMmHashesContract(CustomTestCase): + def test_generate_req_input_accepts_mm_hashes(self): + """GenerateReqInput exposes mm_hashes as an optional field.""" + req = GenerateReqInput( + text="hi", + image_data=["http://example.com/img.png"], + mm_hashes=["deadbeefcafe1234"], + ) + self.assertEqual(req.mm_hashes, ["deadbeefcafe1234"]) + + def test_generate_req_input_defaults_mm_hashes_to_none(self): + """Absent mm_hashes preserves existing (None) behavior.""" + req = GenerateReqInput(text="hi") + self.assertIsNone(req.mm_hashes) + + def test_set_pad_value_honors_preset_hash(self): + """set_pad_value() must use a pre-set hash without recomputing.""" + item = MultimodalDataItem(modality=Modality.IMAGE, hash=0xDEADBEEF) + # If hash_feature is invoked, the test fails — we patch it to + # raise so any accidental recompute is loud. + with patch( + "sglang.srt.managers.mm_utils.hash_feature", + side_effect=AssertionError( + "hash_feature must NOT be called when hash is preset" + ), + ): + item.set_pad_value() + self.assertEqual(item.hash, 0xDEADBEEF) + self.assertEqual(item.pad_value, _compute_pad_value(0xDEADBEEF)) + + def test_set_pad_value_is_deterministic_across_items(self): + """Two items with the same preset hash must derive the same pad_value.""" + a = MultimodalDataItem(modality=Modality.IMAGE, hash=0x123456789ABCDEF0) + b = MultimodalDataItem(modality=Modality.IMAGE, hash=0x123456789ABCDEF0) + # No feature payload — set_pad_value uses the preset hash. + a.set_pad_value() + b.set_pad_value() + self.assertEqual(a.pad_value, b.pad_value) + self.assertEqual(a.hash, b.hash) + + def test_set_pad_value_distinguishes_different_preset_hashes(self): + """Distinct preset hashes must produce distinct pad_values.""" + a = MultimodalDataItem(modality=Modality.IMAGE, hash=0xAAAA) + b = MultimodalDataItem(modality=Modality.IMAGE, hash=0xBBBB) + a.set_pad_value() + b.set_pad_value() + self.assertNotEqual(a.pad_value, b.pad_value) + + +if __name__ == "__main__": + unittest.main()