diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 8f84a18fe..f95eae089 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -55,6 +55,7 @@ from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.managers.embed_types import PositionalEmbeds from sglang.srt.managers.schedule_batch import ( Modality, + MultimodalProcessorOutput, ReturnHiddenStatesMode, get_return_hidden_states_mode, ) @@ -62,6 +63,7 @@ from sglang.srt.multimodal.mm_utils import has_valid_data from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.utils import ImageData, VideoData from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d +from sglang.srt.utils.msgpack_utils import dec_hook, enc_hook, ext_hook from sglang.srt.utils.msgspec_utils import ( Base64Bytes, msgspec_struct_pydantic_core_schema, @@ -945,7 +947,7 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): # The input embeds input_embeds: Optional[List[List[float]]] # The multimodal inputs - mm_inputs: Optional[PickleWrapper] # Pickled Optional[MultimodalProcessorOutput] + mm_inputs: Optional[MultimodalProcessorOutput] token_type_ids: Optional[List[int]] # The sampling parameters sampling_params: SamplingParams @@ -1042,12 +1044,10 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): cache_salt: Optional[str] = None 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) @@ -1307,7 +1307,7 @@ class TokenizedEmbeddingReqInput(BaseReq, kw_only=True): # The input token ids input_ids: Optional[array] # array[int] # The multimodal inputs - mm_inputs: Optional[PickleWrapper] # Pickled Optional[MultimodalProcessorOutput] + mm_inputs: Optional[MultimodalProcessorOutput] # The token type ids token_type_ids: Optional[List[int]] # Dummy sampling params for compatibility @@ -1332,11 +1332,9 @@ class TokenizedEmbeddingReqInput(BaseReq, kw_only=True): time_stats: Optional[PickleWrapper] = None def wrap_pickle_fields(self): - self.mm_inputs = wrap_as_pickle(self.mm_inputs) self.time_stats = wrap_as_pickle(self.time_stats) def unwrap_pickle_fields(self): - self.mm_inputs = unwrap_from_pickle(self.mm_inputs) self.time_stats = unwrap_from_pickle(self.time_stats) @@ -2329,6 +2327,8 @@ def _check_all_req_types(): for class_type in all_classes: # check its name name = class_type[0] + if class_type[1].__module__ != __name__: + continue if name in _IGNORE_REQ_TYPES_CHECK: continue is_io_struct = ( @@ -2368,53 +2368,6 @@ def unwrap_from_pickle(obj: Optional[object]) -> Optional[object]: return pickle.loads(obj.data) -def enc_hook(obj: Any) -> Any: - if isinstance(obj, array): - return (obj.typecode, obj.tobytes()) - elif isinstance(obj, torch.Tensor): - tensor_dtype = str(obj.dtype).removeprefix("torch.") - raw_data = ( - obj.cpu().contiguous().reshape(-1).view(torch.uint8).numpy().tobytes() - ) - return (obj.shape, tensor_dtype, raw_data) - elif isinstance(obj, np.ndarray): - raw_data = np.ascontiguousarray(obj).reshape(-1).view(np.uint8).data - return (obj.shape, obj.dtype.str, raw_data) - elif isinstance(obj, np.floating): - return float(obj) - else: - raise TypeError( - f"Cannot msgpack encode object of type {type(obj)} with enc_hook. " - "Use an explicit PickleWrapper field via wrap_as_pickle(...) for " - "arbitrary payloads, or add a dedicated enc_hook/dec_hook branch " - "for this transport type." - ) - - -def dec_hook(tp: Type, obj: Any) -> Any: - if tp is array: - typecode, raw_data = obj - res = array(typecode) - res.frombytes(raw_data) - return res - elif tp is torch.Tensor: - shape, dtype, data = obj - tensor_dtype = getattr(torch, dtype) - if len(data) == 0: - return torch.empty(shape, dtype=tensor_dtype) - return torch.frombuffer(bytearray(data), dtype=tensor_dtype).reshape(shape) - elif tp is np.ndarray: - shape, dtype, data = obj - return np.frombuffer(data, dtype=np.dtype(dtype)).copy().reshape(shape) - else: - raise TypeError( - f"Cannot msgpack decode object of type {type(obj)} as {tp} with " - "dec_hook. Use an explicit PickleWrapper field via wrap_as_pickle(...) " - "and unwrap_from_pickle(...) for arbitrary payloads, or add a " - "dedicated enc_hook/dec_hook branch for this transport type." - ) - - _struct_types = tuple( cls for cls in BaseReq.__subclasses__() @@ -2429,14 +2382,18 @@ _primitive_types = (int, float, bool, bytes) _all_types = _struct_types + _primitive_types _msgpack_encoder = msgspec.msgpack.Encoder(enc_hook=enc_hook) -_msgpack_decoder = msgspec.msgpack.Decoder(Union[_all_types], dec_hook=dec_hook) +_msgpack_decoder = msgspec.msgpack.Decoder( + Union[_all_types], dec_hook=dec_hook, ext_hook=ext_hook +) _USE_PICKLE_IPC = envs.SGLANG_USE_PICKLE_IPC.get() def hook_custom_types(*new_types: Type): global _msgpack_decoder, _all_types _all_types = tuple(dict.fromkeys(_all_types + new_types)) - _msgpack_decoder = msgspec.msgpack.Decoder(Union[_all_types], dec_hook=dec_hook) + _msgpack_decoder = msgspec.msgpack.Decoder( + Union[_all_types], dec_hook=dec_hook, ext_hook=ext_hook + ) def _maybe_wrap_pickle(obj: Any) -> Any: diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 97d913c0c..b6f504088 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -67,6 +67,7 @@ from typing import ( Optional, Set, Tuple, + TypeAlias, Union, ) @@ -145,6 +146,7 @@ INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5 # Constant used as the base offset for MM (multimodal) pad values. # This ensures pad_values don't overlap with valid text token IDs. MM_PAD_SHIFT_VALUE = 1_000_000 +_MM_HASH_MASK = (1 << 64) - 1 logger = logging.getLogger(__name__) @@ -314,8 +316,12 @@ class MultimodalInputFormat(Enum): PRECOMPUTED_EMBEDDING = auto() -@dataclasses.dataclass -class MultimodalDataItem: +# Msgpack-native containers and Ext-decoded tensor/transport leaves. Tuple +# containers intentionally decode as lists, matching msgpack's native model. +MultimodalDataValue: TypeAlias = object + + +class MultimodalDataItem(msgspec.Struct, kw_only=True, dict=True, array_like=True): """ One MultimodalDataItem represents a single multimodal input (one image, one video, or one audio). For example, if there are 3 images and 1 audio, there will be 4 MultimodalDataItems. @@ -326,39 +332,49 @@ class MultimodalDataItem: """ modality: Modality - hash: int = None - pad_value: int = None - offsets: Optional[list] = None + hash: Optional[int] = None + pad_value: Optional[int] = None + offsets: Optional[List[Tuple[int, int]]] = None format: MultimodalInputFormat = MultimodalInputFormat.NORMAL # the raw features returned by processor, e.g. pixel_values or audio_features - feature: Union[torch.Tensor, np.ndarray] = None + feature: Optional[MultimodalDataValue] = None # the precomputed embeddings, passed as final encoder embeddings # One and only one of the feature and precomputed_embeddings will be empty - precomputed_embeddings: Optional[Union[torch.Tensor, np.ndarray]] = None + precomputed_embeddings: Optional[MultimodalDataValue] = None - # Model-specific data stored in a dictionary - model_specific_data: dict[str, Any] = dataclasses.field(default_factory=dict) + # Processor-owned tensors/arrays/scalars/transports. msgspec rejects a + # precise union with multiple custom types, but accepts Ext-decoded values + # under object. + model_specific_data: Dict[str, MultimodalDataValue] = msgspec.field( + default_factory=dict + ) - def __getattr__(self, name: str): - if ( - "model_specific_data" in self.__dict__ - and name in self.__dict__["model_specific_data"] - ): - return self.__dict__["model_specific_data"][name] + def __post_init__(self) -> None: + if self.hash is not None: + msgspec.Struct.__setattr__(self, "hash", self.hash & _MM_HASH_MASK) + + def __getattr__(self, name: str) -> MultimodalDataValue: + if name in self.model_specific_data: + return self.model_specific_data[name] + + raise AttributeError( + f"'{self.__class__.__name__}' object has no attribute '{name}'" + ) + + def __setattr__(self, name: str, value: MultimodalDataValue) -> None: + if name in self.__struct_fields__: + if name == "hash" and isinstance(value, int): + value &= _MM_HASH_MASK + msgspec.Struct.__setattr__(self, name, value) else: - raise AttributeError( - f"'{self.__class__.__name__}' object has no attribute '{name}'" - ) + self.model_specific_data[name] = value - def __setitem__(self, key: str, value: Any): - if key in self.__dict__: - self.__dict__[key] = value - else: - self.model_specific_data[key] = value + def __setitem__(self, key: str, value: MultimodalDataValue) -> None: + setattr(self, key, value) - def set(self, key: str, value: Any): + def set(self, key: str, value: MultimodalDataValue) -> None: self.__setitem__(key, value) def set_hash(self, hash_value: int) -> None: @@ -491,8 +507,9 @@ class MultimodalDataItem: return min(requested_count, proxy_count) -@dataclasses.dataclass -class MultimodalProcessorOutput: +class MultimodalProcessorOutput( + msgspec.Struct, kw_only=True, dict=True, array_like=True, weakref=True +): """Raw output from multimodal processors before scheduler-side preparation (pad, hash). This is the typed replacement for the dict previously returned by diff --git a/python/sglang/srt/models/nano_nemotron_vl.py b/python/sglang/srt/models/nano_nemotron_vl.py index bfd6cc128..021d33ad8 100644 --- a/python/sglang/srt/models/nano_nemotron_vl.py +++ b/python/sglang/srt/models/nano_nemotron_vl.py @@ -41,7 +41,6 @@ from sglang.srt.models.parakeet import ProjectedParakeet from sglang.srt.models.radio import RadioModel from sglang.srt.models.utils import WeightsMapper from sglang.srt.multimodal.evs import EVS, EVSConfig -from sglang.srt.multimodal.evs.evs_module import VideoEVSDataItem from sglang.srt.utils import add_prefix logger = logging.getLogger(__name__) @@ -148,8 +147,11 @@ class NemotronH_Nano_VL_V2(EVS): if audio_items: for item in visual_items: - if isinstance(item, VideoEVSDataItem): - item.pre_chunked_input_ids = input_ids + if ( + item.is_video() + and "pre_chunked_input_ids" in item.model_specific_data + ): + item.set("pre_chunked_input_ids", input_ids) return input_ids diff --git a/python/sglang/srt/multimodal/evs/evs_module.py b/python/sglang/srt/multimodal/evs/evs_module.py index a7ab6296e..b62ce08de 100644 --- a/python/sglang/srt/multimodal/evs/evs_module.py +++ b/python/sglang/srt/multimodal/evs/evs_module.py @@ -13,7 +13,6 @@ # ============================================================================== -import dataclasses import typing from abc import ABC, abstractmethod from dataclasses import dataclass @@ -29,19 +28,6 @@ from sglang.utils import logger from .evs_core import compute_retention_mask, replace_offsets_with_tokens_per_frame -@dataclasses.dataclass(kw_only=True) -class EVSDataItem(MultimodalDataItem): - thw_grids: list[tuple[int, int, int]] - - -@dataclasses.dataclass(kw_only=True) -class VideoEVSDataItem(EVSDataItem): - pre_chunked_input_ids: torch.Tensor - - def __post_init__(self): - assert self.is_video() - - @dataclass(kw_only=True) class EVSEmbeddingResult(EmbeddingResult): """ @@ -65,14 +51,12 @@ class EVSEmbeddingResult(EmbeddingResult): input_ids: torch.Tensor, offsets: list[tuple[int, int]], *, - item: VideoEVSDataItem, + item: MultimodalDataItem, extend_prefix_len: int, extend_seq_len: int, ) -> tuple[torch.Tensor, list[tuple[int, int]]]: assert len(input_ids) == extend_seq_len - assert isinstance( - item, VideoEVSDataItem - ), f"Expected VideoEVSDataItem, got {type(item)}" + assert item.is_video() and "pre_chunked_input_ids" in item.model_specific_data pre_chunked_input_ids = item.pre_chunked_input_ids filler_token_id = item.pad_value input_ids_list = replace_offsets_with_tokens_per_frame( @@ -151,7 +135,7 @@ class EVS(torch.nn.Module, ABC): Apply EVS pruning to video embeddings. Args: - items: List containing a single VideoEVSDataItem with video features. + items: List containing a single EVS video item with video features. Returns: EVSEmbeddingResult with pruned embeddings and actual token counts per frame. @@ -161,9 +145,7 @@ class EVS(torch.nn.Module, ABC): ) assert len(items) == 1, f"Expected 1 item, got {len(items)}" item = items[0] - assert isinstance( - item, VideoEVSDataItem - ), f"Expected VideoEVSDataItem with modality VIDEO, got {item}" + assert item.is_video() and "thw_grids" in item.model_specific_data q = self.evs_config.video_pruning_rate merge = self.evs_config.spatial_merge_size diff --git a/python/sglang/srt/multimodal/evs/evs_processor.py b/python/sglang/srt/multimodal/evs/evs_processor.py index 251f17bca..98e731afa 100644 --- a/python/sglang/srt/multimodal/evs/evs_processor.py +++ b/python/sglang/srt/multimodal/evs/evs_processor.py @@ -20,7 +20,7 @@ from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.utils import logger from .evs_core import tokens_per_frame -from .evs_module import EVS, EVSConfig, EVSDataItem, VideoEVSDataItem +from .evs_module import EVS, EVSConfig def _non_evs_data_items( @@ -99,23 +99,25 @@ class EVSProcessor: items = [] if image is not None: image_thw_grids = [(1, rows, cols)] * num_images - item = EVSDataItem( + item = MultimodalDataItem( modality=Modality.IMAGE, feature=image, offsets=image_offsets, - thw_grids=image_thw_grids, + model_specific_data={"thw_grids": image_thw_grids}, ) items.append(item) if video is not None: video_thw_grids = [ (num_frames, rows, cols) for num_frames in frames_per_video ] - item = VideoEVSDataItem( + item = MultimodalDataItem( modality=Modality.VIDEO, feature=video, offsets=video_offsets, - thw_grids=video_thw_grids, - pre_chunked_input_ids=input_ids_list, + model_specific_data={ + "thw_grids": video_thw_grids, + "pre_chunked_input_ids": input_ids_list, + }, ) items.append(item) return items diff --git a/python/sglang/srt/multimodal/processors/transformers_auto.py b/python/sglang/srt/multimodal/processors/transformers_auto.py index b9dddb320..19c075b20 100644 --- a/python/sglang/srt/multimodal/processors/transformers_auto.py +++ b/python/sglang/srt/multimodal/processors/transformers_auto.py @@ -238,7 +238,7 @@ class TransformersAutoMultimodalProcessor(BaseMultimodalProcessor): else "token_type_ids" ) if token_type_key in processor_output: - ret.token_type_ids = processor_output[token_type_key].flatten().tolist() + ret.token_type_ids = processor_output[token_type_key].flatten() if self.mm_tokens.image_token_id is not None: ret.im_token_id = self.mm_tokens.image_token_id diff --git a/python/sglang/srt/utils/msgpack_utils.py b/python/sglang/srt/utils/msgpack_utils.py new file mode 100644 index 000000000..1383a3cd1 --- /dev/null +++ b/python/sglang/srt/utils/msgpack_utils.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +import struct +from array import array +from typing import Sequence, Union + +import msgspec +import numpy as np +import torch + +from sglang.srt.utils.cuda_ipc_transport_utils import CudaIpcTensorTransportProxy + +# Stable wire IDs. Changing these requires updating the golden-wire test. +_MSGPACK_EXT_ARRAY = 1 +_MSGPACK_EXT_TORCH_TENSOR = 2 +_MSGPACK_EXT_NP_ARRAY = 3 +_MSGPACK_EXT_SHM_POINTER_MM_DATA = 4 +_MSGPACK_EXT_CUDA_IPC_TENSOR_PROXY = 5 +_MSGPACK_BUFFER_METADATA_SIZE = struct.Struct(">I") + + +def _pack_ext(code: int, obj: object) -> msgspec.msgpack.Ext: + return msgspec.msgpack.Ext(code, msgspec.msgpack.encode(obj, enc_hook=enc_hook)) + + +def _unpack_ext(data: memoryview) -> object: + return msgspec.msgpack.decode(data, ext_hook=ext_hook) + + +def _pack_buffer_ext( + code: int, metadata: object, raw_data: memoryview +) -> msgspec.msgpack.Ext: + metadata_bytes = msgspec.msgpack.encode(metadata) + payload = bytearray(_MSGPACK_BUFFER_METADATA_SIZE.pack(len(metadata_bytes))) + payload.extend(metadata_bytes) + payload.extend(raw_data) + return msgspec.msgpack.Ext(code, payload) + + +def _unpack_buffer_ext(data: memoryview) -> tuple[object, memoryview]: + if len(data) < _MSGPACK_BUFFER_METADATA_SIZE.size: + raise msgspec.DecodeError("MessagePack buffer extension is missing metadata") + + (metadata_size,) = _MSGPACK_BUFFER_METADATA_SIZE.unpack_from(data) + raw_data_offset = _MSGPACK_BUFFER_METADATA_SIZE.size + metadata_size + if raw_data_offset > len(data): + raise msgspec.DecodeError("MessagePack buffer extension has invalid metadata") + + metadata = msgspec.msgpack.decode( + data[_MSGPACK_BUFFER_METADATA_SIZE.size : raw_data_offset] + ) + return metadata, data[raw_data_offset:] + + +def _torch_dtype_name(dtype: torch.dtype) -> str: + return str(dtype).removeprefix("torch.") + + +def _torch_dtype_from_name(name: str) -> torch.dtype: + return getattr(torch, name) + + +def _restore_torch_tensor( + shape: Sequence[int], + dtype: str, + data: Union[bytes, memoryview], + device: str = "cpu", +) -> torch.Tensor: + tensor_dtype = _torch_dtype_from_name(dtype) + if len(data) == 0: + tensor = torch.empty(shape, dtype=tensor_dtype, device="cpu") + else: + tensor = torch.frombuffer(bytearray(data), dtype=tensor_dtype).reshape(shape) + if device != "cpu": + tensor = tensor.to(device) + return tensor + + +def _to_msgpack_state(obj: object) -> object: + if isinstance(obj, torch.dtype): + return {"__torch_dtype__": _torch_dtype_name(obj)} + if isinstance(obj, torch.device): + return {"__torch_device__": str(obj)} + if isinstance(obj, np.dtype): + return {"__np_dtype__": obj.str} + if isinstance(obj, dict): + return {key: _to_msgpack_state(value) for key, value in obj.items()} + if isinstance(obj, torch.Size): + return {"__torch_size__": list(obj)} + if isinstance(obj, tuple): + return {"__tuple__": [_to_msgpack_state(value) for value in obj]} + if isinstance(obj, list): + return [_to_msgpack_state(value) for value in obj] + return obj + + +def _from_msgpack_state(obj: object) -> object: + if isinstance(obj, dict): + if "__torch_dtype__" in obj: + return _torch_dtype_from_name(obj["__torch_dtype__"]) + if "__torch_device__" in obj: + return torch.device(obj["__torch_device__"]) + if "__np_dtype__" in obj: + return np.dtype(obj["__np_dtype__"]) + if "__torch_size__" in obj: + return torch.Size(obj["__torch_size__"]) + if "__tuple__" in obj: + return tuple(_from_msgpack_state(value) for value in obj["__tuple__"]) + return {key: _from_msgpack_state(value) for key, value in obj.items()} + if isinstance(obj, list): + return [_from_msgpack_state(value) for value in obj] + return obj + + +def _is_shm_pointer_mm_data(obj: object) -> bool: + cls = type(obj) + return cls.__name__ == "ShmPointerMMData" and cls.__module__.endswith( + ".managers.mm_utils" + ) + + +def _encode_shm_pointer_mm_data(obj: object) -> object: + return _to_msgpack_state(obj.__getstate__()) + + +def _decode_shm_pointer_mm_data(state: dict[str, object]) -> object: + from sglang.srt.managers.mm_utils import ShmPointerMMData + + obj = ShmPointerMMData.__new__(ShmPointerMMData) + obj.__setstate__(_from_msgpack_state(state)) + return obj + + +def _encode_cuda_ipc_tensor_proxy(obj: CudaIpcTensorTransportProxy) -> object: + return { + "proxy_state": _to_msgpack_state(obj.proxy_state), + "sync_data_meta": _to_msgpack_state(obj.sync_data_meta), + } + + +def _decode_cuda_ipc_tensor_proxy( + state: dict[str, object], +) -> CudaIpcTensorTransportProxy: + obj = CudaIpcTensorTransportProxy.__new__(CudaIpcTensorTransportProxy) + obj.proxy_state = _from_msgpack_state(state["proxy_state"]) + obj.reconstruct_tensor = None + obj.sync_data_meta = _from_msgpack_state(state["sync_data_meta"]) + obj.sync_buffer = None + obj._consumer_acknowledged = False + return obj + + +def enc_hook(obj: object) -> object: + if isinstance(obj, array): + return _pack_buffer_ext( + _MSGPACK_EXT_ARRAY, obj.typecode, memoryview(obj).cast("B") + ) + if isinstance(obj, torch.Tensor): + tensor_dtype = _torch_dtype_name(obj.dtype) + tensor = obj.cpu().contiguous() + raw_data = tensor.reshape(-1).view(torch.uint8).numpy().data + return _pack_buffer_ext( + _MSGPACK_EXT_TORCH_TENSOR, + (tuple(obj.shape), tensor_dtype, str(obj.device)), + raw_data, + ) + if isinstance(obj, np.ndarray): + arr = np.ascontiguousarray(obj) + raw_data = arr.reshape(-1).view(np.uint8).data + return _pack_buffer_ext( + _MSGPACK_EXT_NP_ARRAY, + (arr.shape, arr.dtype.str), + raw_data, + ) + if isinstance(obj, np.floating): + return float(obj) + if isinstance(obj, np.integer): + return int(obj) + if isinstance(obj, np.bool_): + return bool(obj) + if isinstance(obj, CudaIpcTensorTransportProxy): + return _pack_ext( + _MSGPACK_EXT_CUDA_IPC_TENSOR_PROXY, + _encode_cuda_ipc_tensor_proxy(obj), + ) + if _is_shm_pointer_mm_data(obj): + return _pack_ext( + _MSGPACK_EXT_SHM_POINTER_MM_DATA, + _encode_shm_pointer_mm_data(obj), + ) + raise TypeError( + f"Cannot msgpack encode object of type {type(obj)} with enc_hook. " + "Use an explicit PickleWrapper field via wrap_as_pickle(...) for " + "arbitrary payloads, or add a dedicated enc_hook/dec_hook branch " + "for this transport type." + ) + + +def dec_hook(tp: type, obj: object) -> object: + if isinstance(obj, tp): + return obj + if tp is array: + typecode, raw_data = obj + res = array(typecode) + res.frombytes(raw_data) + return res + if tp is torch.Tensor: + shape, dtype, data, *device = obj + return _restore_torch_tensor(shape, dtype, data, device[0] if device else "cpu") + if tp is np.ndarray: + shape, dtype, data = obj + return np.frombuffer(data, dtype=np.dtype(dtype)).copy().reshape(shape) + raise TypeError( + f"Cannot msgpack decode object of type {type(obj)} as {tp} with " + "dec_hook. Use an explicit PickleWrapper field via wrap_as_pickle(...) " + "and unwrap_from_pickle(...) for arbitrary payloads, or add a " + "dedicated enc_hook/dec_hook branch for this transport type." + ) + + +def ext_hook(code: int, data: memoryview) -> object: + if code not in ( + _MSGPACK_EXT_ARRAY, + _MSGPACK_EXT_TORCH_TENSOR, + _MSGPACK_EXT_NP_ARRAY, + _MSGPACK_EXT_SHM_POINTER_MM_DATA, + _MSGPACK_EXT_CUDA_IPC_TENSOR_PROXY, + ): + return msgspec.msgpack.Ext(code, bytes(data)) + + if code == _MSGPACK_EXT_ARRAY: + typecode, raw_data = _unpack_buffer_ext(data) + res = array(typecode) + res.frombytes(raw_data) + return res + if code == _MSGPACK_EXT_TORCH_TENSOR: + metadata, raw_data = _unpack_buffer_ext(data) + shape, dtype, device = metadata + return _restore_torch_tensor(shape, dtype, raw_data, device) + if code == _MSGPACK_EXT_NP_ARRAY: + metadata, raw_data = _unpack_buffer_ext(data) + shape, dtype = metadata + return np.frombuffer(raw_data, dtype=np.dtype(dtype)).copy().reshape(shape) + if code == _MSGPACK_EXT_SHM_POINTER_MM_DATA: + return _decode_shm_pointer_mm_data(_unpack_ext(data)) + if code == _MSGPACK_EXT_CUDA_IPC_TENSOR_PROXY: + return _decode_cuda_ipc_tensor_proxy(_unpack_ext(data)) + raise AssertionError(f"Unhandled known MessagePack extension code: {code}") diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index 9c9896204..2f145220a 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -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.""" diff --git a/test/registered/unit/multimodal/test_evs.py b/test/registered/unit/multimodal/test_evs.py index aee3f42f1..acbd79b48 100644 --- a/test/registered/unit/multimodal/test_evs.py +++ b/test/registered/unit/multimodal/test_evs.py @@ -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