feat: optional caller-supplied mm_hashes on GenerateReqInput (#25300)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com>

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,

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.

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

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()