feat: make mm_inputs msgpack-native (#29656)
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
ishandhanani
co-authored by
Alex Nails
Alex Nails
parent
5a100d9086
commit
0f744b6848
@@ -1,7 +1,29 @@
|
||||
import copy
|
||||
import unittest
|
||||
import weakref
|
||||
from array import array
|
||||
|
||||
from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput
|
||||
import msgspec
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.managers.io_struct import (
|
||||
EmbeddingReqInput,
|
||||
GenerateReqInput,
|
||||
TokenizedEmbeddingReqInput,
|
||||
TokenizedGenerateReqInput,
|
||||
msgpack_decode,
|
||||
msgpack_encode,
|
||||
)
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
MultimodalInputFormat,
|
||||
MultimodalProcessorOutput,
|
||||
)
|
||||
from sglang.srt.sampling.sampling_params import SamplingParams
|
||||
from sglang.srt.utils.cuda_ipc_transport_utils import CudaIpcTensorTransportProxy
|
||||
from sglang.srt.utils.msgpack_utils import _restore_torch_tensor, enc_hook, ext_hook
|
||||
from sglang.test.ci.ci_register import (
|
||||
register_amd_ci,
|
||||
register_cpu_ci,
|
||||
@@ -18,6 +40,328 @@ 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 TestTokenizedReqInputMsgpack(unittest.TestCase):
|
||||
def _make_mm_inputs(self, device="cpu"):
|
||||
return MultimodalProcessorOutput(
|
||||
mm_items=[
|
||||
MultimodalDataItem(
|
||||
modality=Modality.IMAGE,
|
||||
offsets=[(0, 1)],
|
||||
format=MultimodalInputFormat.NORMAL,
|
||||
feature=torch.tensor(
|
||||
[[1.0, 2.0]], dtype=torch.float32, device=device
|
||||
),
|
||||
model_specific_data={
|
||||
"image_grid_thw": torch.tensor(
|
||||
[[1, 1, 2]], dtype=torch.int64, device=device
|
||||
),
|
||||
"patch_counts": np.array([2], dtype=np.int32),
|
||||
"names": ["image0"],
|
||||
"count": np.int64(2),
|
||||
"enabled": np.bool_(True),
|
||||
"size": (336, 336),
|
||||
},
|
||||
)
|
||||
],
|
||||
input_ids=[1, 2],
|
||||
padded_input_ids=[10, 10],
|
||||
im_token_id=10,
|
||||
mrope_positions=torch.tensor([[0, 1]], dtype=torch.int64, device=device),
|
||||
token_type_ids=torch.tensor([0, 0], dtype=torch.int64, device=device),
|
||||
)
|
||||
|
||||
def _round_trip(self, req):
|
||||
req.wrap_pickle_fields()
|
||||
decoded = msgpack_decode(msgpack_encode(req))
|
||||
decoded.unwrap_pickle_fields()
|
||||
return decoded
|
||||
|
||||
def _round_trip_mm_inputs(self, mm_inputs):
|
||||
decoded = self._round_trip(
|
||||
TokenizedGenerateReqInput(
|
||||
input_text="",
|
||||
input_ids=array("q", [1, 2]),
|
||||
input_embeds=None,
|
||||
mm_inputs=mm_inputs,
|
||||
token_type_ids=[0, 0],
|
||||
sampling_params=SamplingParams(),
|
||||
return_logprob=False,
|
||||
logprob_start_len=0,
|
||||
top_logprobs_num=0,
|
||||
token_ids_logprob=None,
|
||||
stream=False,
|
||||
)
|
||||
)
|
||||
return decoded.mm_inputs
|
||||
|
||||
def test_generate_mm_inputs_round_trip_without_pickle_wrapper(self):
|
||||
decoded = self._round_trip(
|
||||
TokenizedGenerateReqInput(
|
||||
input_text="",
|
||||
input_ids=array("q", [1, 2]),
|
||||
input_embeds=None,
|
||||
mm_inputs=self._make_mm_inputs(),
|
||||
token_type_ids=[0, 0],
|
||||
sampling_params=SamplingParams(),
|
||||
return_logprob=False,
|
||||
logprob_start_len=0,
|
||||
top_logprobs_num=0,
|
||||
token_ids_logprob=None,
|
||||
stream=False,
|
||||
)
|
||||
)
|
||||
|
||||
self.assertIsInstance(decoded.mm_inputs, MultimodalProcessorOutput)
|
||||
item = decoded.mm_inputs.mm_items[0]
|
||||
self.assertIsInstance(item, MultimodalDataItem)
|
||||
self.assertEqual(item.modality, Modality.IMAGE)
|
||||
self.assertEqual(item.offsets, [(0, 1)])
|
||||
self.assertTrue(
|
||||
torch.equal(item.feature, torch.tensor([[1.0, 2.0]], device="cpu"))
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
item.model_specific_data["image_grid_thw"],
|
||||
torch.tensor([[1, 1, 2]], dtype=torch.int64, device="cpu"),
|
||||
)
|
||||
)
|
||||
np.testing.assert_array_equal(
|
||||
item.model_specific_data["patch_counts"],
|
||||
np.array([2], dtype=np.int32),
|
||||
)
|
||||
self.assertEqual(item.model_specific_data["count"], 2)
|
||||
self.assertIs(item.model_specific_data["enabled"], True)
|
||||
self.assertEqual(item.model_specific_data["size"], [336, 336])
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
decoded.mm_inputs.mrope_positions,
|
||||
torch.tensor([[0, 1]], dtype=torch.int64, device="cpu"),
|
||||
)
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.equal(decoded.mm_inputs.token_type_ids, torch.tensor([0, 0]))
|
||||
)
|
||||
|
||||
def test_dynamic_model_specific_attribute_round_trip(self):
|
||||
mm_inputs = self._make_mm_inputs()
|
||||
mm_inputs.mm_items[0].audio_feature_lens = torch.tensor([2])
|
||||
|
||||
decoded = self._round_trip_mm_inputs(mm_inputs)
|
||||
|
||||
self.assertTrue(
|
||||
torch.equal(decoded.mm_items[0].audio_feature_lens, torch.tensor([2]))
|
||||
)
|
||||
self.assertIn("audio_feature_lens", decoded.mm_items[0].model_specific_data)
|
||||
|
||||
def test_multimodal_hash_is_normalized_to_uint64(self):
|
||||
mm_inputs = self._make_mm_inputs()
|
||||
mm_inputs.mm_items[0].hash = (1 << 256) - 1
|
||||
constructed = MultimodalDataItem(modality=Modality.IMAGE, hash=(1 << 128) - 1)
|
||||
|
||||
decoded = self._round_trip_mm_inputs(mm_inputs)
|
||||
|
||||
self.assertEqual(decoded.mm_items[0].hash, (1 << 64) - 1)
|
||||
self.assertEqual(constructed.hash, (1 << 64) - 1)
|
||||
|
||||
def test_multimodal_processor_output_supports_weakrefs(self):
|
||||
mm_inputs = self._make_mm_inputs()
|
||||
|
||||
ref = weakref.ref(mm_inputs)
|
||||
|
||||
self.assertIs(ref(), mm_inputs)
|
||||
|
||||
def test_unknown_ext_payload_is_preserved_without_decoding(self):
|
||||
ext = msgspec.msgpack.Ext(99, b"not msgpack")
|
||||
|
||||
decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(ext), ext_hook=ext_hook)
|
||||
|
||||
self.assertEqual(decoded, ext)
|
||||
|
||||
def test_malformed_known_buffer_ext_is_rejected(self):
|
||||
with self.assertRaisesRegex(msgspec.DecodeError, "missing metadata"):
|
||||
ext_hook(3, memoryview(b"bad"))
|
||||
|
||||
def test_embedding_mm_inputs_round_trip_without_pickle_wrapper(self):
|
||||
decoded = self._round_trip(
|
||||
TokenizedEmbeddingReqInput(
|
||||
input_text="",
|
||||
input_ids=array("q", [1, 2]),
|
||||
mm_inputs=self._make_mm_inputs(),
|
||||
token_type_ids=[0, 0],
|
||||
sampling_params=SamplingParams(),
|
||||
)
|
||||
)
|
||||
|
||||
self.assertIsInstance(decoded.mm_inputs, MultimodalProcessorOutput)
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
decoded.mm_inputs.mm_items[0].feature,
|
||||
torch.tensor([[1.0, 2.0]], device="cpu"),
|
||||
)
|
||||
)
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is not available")
|
||||
def test_generate_mm_inputs_round_trip_preserves_cuda_tensor_device(self):
|
||||
decoded = self._round_trip(
|
||||
TokenizedGenerateReqInput(
|
||||
input_text="",
|
||||
input_ids=array("q", [1, 2]),
|
||||
input_embeds=None,
|
||||
mm_inputs=self._make_mm_inputs(device="cuda:0"),
|
||||
token_type_ids=[0, 0],
|
||||
sampling_params=SamplingParams(),
|
||||
return_logprob=False,
|
||||
logprob_start_len=0,
|
||||
top_logprobs_num=0,
|
||||
token_ids_logprob=None,
|
||||
stream=False,
|
||||
)
|
||||
)
|
||||
|
||||
item = decoded.mm_inputs.mm_items[0]
|
||||
self.assertEqual(item.feature.device.type, "cuda")
|
||||
self.assertEqual(item.model_specific_data["image_grid_thw"].device.type, "cuda")
|
||||
self.assertEqual(decoded.mm_inputs.mrope_positions.device.type, "cuda")
|
||||
|
||||
def test_cuda_ipc_proxy_state_round_trip_preserves_tuple_types(self):
|
||||
proxy = CudaIpcTensorTransportProxy.__new__(CudaIpcTensorTransportProxy)
|
||||
proxy.proxy_state = {
|
||||
"ipc_extra": {
|
||||
"shape": torch.Size([2, 3]),
|
||||
"stride": (3, 1),
|
||||
"dtype": torch.float16,
|
||||
"nested": [(1, 2), torch.Size([4])],
|
||||
},
|
||||
"tensor_data": None,
|
||||
}
|
||||
proxy.reconstruct_tensor = None
|
||||
proxy.sync_data_meta = {
|
||||
"handle": "dummy",
|
||||
"shape": torch.Size([1]),
|
||||
"dtype": np.dtype("float32"),
|
||||
}
|
||||
proxy.sync_buffer = None
|
||||
|
||||
mm_inputs = self._make_mm_inputs()
|
||||
mm_inputs.mm_items[0].model_specific_data["ipc_proxy"] = proxy
|
||||
decoded = self._round_trip(
|
||||
TokenizedGenerateReqInput(
|
||||
input_text="",
|
||||
input_ids=array("q", [1, 2]),
|
||||
input_embeds=None,
|
||||
mm_inputs=mm_inputs,
|
||||
token_type_ids=[0, 0],
|
||||
sampling_params=SamplingParams(),
|
||||
return_logprob=False,
|
||||
logprob_start_len=0,
|
||||
top_logprobs_num=0,
|
||||
token_ids_logprob=None,
|
||||
stream=False,
|
||||
)
|
||||
)
|
||||
|
||||
decoded_proxy = decoded.mm_inputs.mm_items[0].model_specific_data["ipc_proxy"]
|
||||
ipc_extra = decoded_proxy.proxy_state["ipc_extra"]
|
||||
self.assertIsInstance(ipc_extra["shape"], torch.Size)
|
||||
self.assertEqual(ipc_extra["shape"], torch.Size([2, 3]))
|
||||
self.assertIsInstance(ipc_extra["stride"], tuple)
|
||||
self.assertEqual(ipc_extra["stride"], (3, 1))
|
||||
self.assertIsInstance(ipc_extra["nested"][0], tuple)
|
||||
self.assertIsInstance(ipc_extra["nested"][1], torch.Size)
|
||||
self.assertIsInstance(decoded_proxy.sync_data_meta["shape"], torch.Size)
|
||||
self.assertIsInstance(decoded_proxy.sync_data_meta["dtype"], np.dtype)
|
||||
self.assertFalse(decoded_proxy._consumer_acknowledged)
|
||||
|
||||
def test_cuda_ipc_proxy_tensor_fallback_round_trip(self):
|
||||
proxy = CudaIpcTensorTransportProxy.__new__(CudaIpcTensorTransportProxy)
|
||||
proxy.proxy_state = {
|
||||
"ipc_extra": None,
|
||||
"tensor_data": torch.tensor([1.0, 2.0], device="cpu"),
|
||||
}
|
||||
proxy.reconstruct_tensor = None
|
||||
proxy.sync_data_meta = {
|
||||
"handle": "dummy",
|
||||
"shape": (1,),
|
||||
"dtype": np.dtype("uint8"),
|
||||
}
|
||||
proxy.sync_buffer = None
|
||||
|
||||
mm_inputs = self._make_mm_inputs()
|
||||
mm_inputs.mm_items[0].model_specific_data["ipc_proxy"] = proxy
|
||||
decoded = self._round_trip(
|
||||
TokenizedGenerateReqInput(
|
||||
input_text="",
|
||||
input_ids=array("q", [1, 2]),
|
||||
input_embeds=None,
|
||||
mm_inputs=mm_inputs,
|
||||
token_type_ids=[0, 0],
|
||||
sampling_params=SamplingParams(),
|
||||
return_logprob=False,
|
||||
logprob_start_len=0,
|
||||
top_logprobs_num=0,
|
||||
token_ids_logprob=None,
|
||||
stream=False,
|
||||
)
|
||||
)
|
||||
|
||||
decoded_proxy = decoded.mm_inputs.mm_items[0].model_specific_data["ipc_proxy"]
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
decoded_proxy.proxy_state["tensor_data"],
|
||||
torch.tensor([1.0, 2.0], device="cpu"),
|
||||
)
|
||||
)
|
||||
|
||||
def test_evs_model_specific_data_round_trip(self):
|
||||
mm_inputs = self._make_mm_inputs()
|
||||
item = mm_inputs.mm_items[0]
|
||||
item.modality = Modality.VIDEO
|
||||
item.model_specific_data.update(
|
||||
{
|
||||
"thw_grids": [(2, 3, 4)],
|
||||
"pre_chunked_input_ids": [1, 2, 3],
|
||||
}
|
||||
)
|
||||
decoded = self._round_trip(
|
||||
TokenizedGenerateReqInput(
|
||||
input_text="",
|
||||
input_ids=array("q", [1, 2]),
|
||||
input_embeds=None,
|
||||
mm_inputs=mm_inputs,
|
||||
token_type_ids=[0, 0],
|
||||
sampling_params=SamplingParams(),
|
||||
return_logprob=False,
|
||||
logprob_start_len=0,
|
||||
top_logprobs_num=0,
|
||||
token_ids_logprob=None,
|
||||
stream=False,
|
||||
)
|
||||
)
|
||||
|
||||
decoded_item = decoded.mm_inputs.mm_items[0]
|
||||
self.assertEqual(decoded_item.thw_grids, [[2, 3, 4]])
|
||||
self.assertEqual(decoded_item.pre_chunked_input_ids, [1, 2, 3])
|
||||
|
||||
def test_torch_tensor_ext_wire_format(self):
|
||||
ext = enc_hook(torch.tensor([1, 2], dtype=torch.int16, device="cpu"))
|
||||
self.assertIsInstance(ext, msgspec.msgpack.Ext)
|
||||
self.assertEqual(ext.code, 2)
|
||||
self.assertEqual(
|
||||
bytes(ext.data).hex(),
|
||||
"0000000d939102a5696e743136a363707501000200",
|
||||
)
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is not available")
|
||||
def test_empty_cpu_tensor_restore_ignores_default_device(self):
|
||||
previous_device = torch.get_default_device()
|
||||
try:
|
||||
torch.set_default_device("cuda")
|
||||
tensor = _restore_torch_tensor((0,), "float32", b"", "cpu")
|
||||
self.assertEqual(tensor.device.type, "cpu")
|
||||
finally:
|
||||
torch.set_default_device(previous_device)
|
||||
|
||||
|
||||
class TestGenerateReqInputNormalization(CustomTestCase):
|
||||
"""Test the normalization of GenerateReqInput for batch processing and different input formats."""
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@ from dataclasses import asdict, dataclass
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import run_doctests
|
||||
@@ -53,6 +54,29 @@ def test_replace_offsets_with_tokens_per_frame():
|
||||
run_doctests(replace_offsets_with_tokens_per_frame)
|
||||
|
||||
|
||||
def test_evs_items_store_wire_data_in_model_specific_data():
|
||||
from sglang.srt.managers.schedule_batch import MultimodalDataItem
|
||||
from sglang.srt.multimodal.evs import EVSConfig, EVSProcessor
|
||||
|
||||
processor = EVSProcessor.__new__(EVSProcessor)
|
||||
processor.evs_config = EVSConfig(video_pruning_rate=0.1)
|
||||
make_items, _ = processor.static_size_data_items(
|
||||
frames_per_video=[2], num_images=1, rows=2, cols=3
|
||||
)
|
||||
items = make_items(
|
||||
input_ids_list=[1, 2, 3],
|
||||
image=torch.zeros(1),
|
||||
image_offsets=[(0, 0)],
|
||||
video=torch.zeros(1),
|
||||
video_offsets=[(1, 2)],
|
||||
)
|
||||
|
||||
assert all(type(item) is MultimodalDataItem for item in items)
|
||||
assert items[0].thw_grids == [(1, 2, 3)]
|
||||
assert items[1].thw_grids == [(2, 2, 3)]
|
||||
assert items[1].pre_chunked_input_ids == [1, 2, 3]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
|
||||
Reference in New Issue
Block a user