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>
This commit is contained in:
co-authored by
Claude Opus 4.7
ishandhanani
parent
f6a5a1b59c
commit
86afa21ca7
@@ -324,6 +324,8 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
image_data: Optional[MultimodalDataInputFormat] = None,
|
image_data: Optional[MultimodalDataInputFormat] = None,
|
||||||
audio_data: Optional[MultimodalDataInputFormat] = None,
|
audio_data: Optional[MultimodalDataInputFormat] = None,
|
||||||
video_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,
|
return_logprob: Optional[Union[List[bool], bool]] = False,
|
||||||
logprob_start_len: Optional[Union[List[int], int]] = None,
|
logprob_start_len: Optional[Union[List[int], int]] = None,
|
||||||
top_logprobs_num: 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,
|
image_data=image_data,
|
||||||
audio_data=audio_data,
|
audio_data=audio_data,
|
||||||
video_data=video_data,
|
video_data=video_data,
|
||||||
|
mm_hashes=mm_hashes,
|
||||||
return_logprob=return_logprob,
|
return_logprob=return_logprob,
|
||||||
logprob_start_len=logprob_start_len,
|
logprob_start_len=logprob_start_len,
|
||||||
top_logprobs_num=top_logprobs_num,
|
top_logprobs_num=top_logprobs_num,
|
||||||
@@ -416,6 +419,13 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
image_data: Optional[MultimodalDataInputFormat] = None,
|
image_data: Optional[MultimodalDataInputFormat] = None,
|
||||||
audio_data: Optional[MultimodalDataInputFormat] = None,
|
audio_data: Optional[MultimodalDataInputFormat] = None,
|
||||||
video_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,
|
return_logprob: Optional[Union[List[bool], bool]] = False,
|
||||||
logprob_start_len: Optional[Union[List[int], int]] = None,
|
logprob_start_len: Optional[Union[List[int], int]] = None,
|
||||||
top_logprobs_num: 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,
|
image_data=image_data,
|
||||||
audio_data=audio_data,
|
audio_data=audio_data,
|
||||||
video_data=video_data,
|
video_data=video_data,
|
||||||
|
mm_hashes=mm_hashes,
|
||||||
return_logprob=return_logprob,
|
return_logprob=return_logprob,
|
||||||
logprob_start_len=logprob_start_len,
|
logprob_start_len=logprob_start_len,
|
||||||
top_logprobs_num=top_logprobs_num,
|
top_logprobs_num=top_logprobs_num,
|
||||||
|
|||||||
@@ -158,6 +158,15 @@ class GenerateReqInput(BaseReq):
|
|||||||
video_data: Optional[MultimodalDataInputFormat] = None
|
video_data: Optional[MultimodalDataInputFormat] = None
|
||||||
# The audio input. Like image data, it can be a file name, a url, or base64 encoded string.
|
# The audio input. Like image data, it can be a file name, a url, or base64 encoded string.
|
||||||
audio_data: Optional[MultimodalDataInputFormat] = None
|
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.
|
# Whether to extract and process audio from video inputs.
|
||||||
use_audio_in_video: bool = False
|
use_audio_in_video: bool = False
|
||||||
# The sampling_params. See descriptions below.
|
# The sampling_params. See descriptions below.
|
||||||
|
|||||||
@@ -840,6 +840,37 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
token_type_ids = mm_inputs.token_type_ids
|
token_type_ids = mm_inputs.token_type_ids
|
||||||
if not isinstance(token_type_ids, list):
|
if not isinstance(token_type_ids, list):
|
||||||
token_type_ids = token_type_ids.flatten().tolist()
|
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 (
|
if (
|
||||||
envs.SGLANG_MM_PRECOMPUTE_HASH.get()
|
envs.SGLANG_MM_PRECOMPUTE_HASH.get()
|
||||||
and mm_inputs
|
and mm_inputs
|
||||||
|
|||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user