feat: make mm_inputs msgpack-native (#29656)
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
co-authored by
Alex Nails
parent
5a100d9086
commit
0f744b6848
@@ -55,6 +55,7 @@ from sglang.srt.lora.lora_registry import LoRARef
|
|||||||
from sglang.srt.managers.embed_types import PositionalEmbeds
|
from sglang.srt.managers.embed_types import PositionalEmbeds
|
||||||
from sglang.srt.managers.schedule_batch import (
|
from sglang.srt.managers.schedule_batch import (
|
||||||
Modality,
|
Modality,
|
||||||
|
MultimodalProcessorOutput,
|
||||||
ReturnHiddenStatesMode,
|
ReturnHiddenStatesMode,
|
||||||
get_return_hidden_states_mode,
|
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.sampling.sampling_params import SamplingParams
|
||||||
from sglang.srt.utils import ImageData, VideoData
|
from sglang.srt.utils import ImageData, VideoData
|
||||||
from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d
|
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 (
|
from sglang.srt.utils.msgspec_utils import (
|
||||||
Base64Bytes,
|
Base64Bytes,
|
||||||
msgspec_struct_pydantic_core_schema,
|
msgspec_struct_pydantic_core_schema,
|
||||||
@@ -945,7 +947,7 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True):
|
|||||||
# The input embeds
|
# The input embeds
|
||||||
input_embeds: Optional[List[List[float]]]
|
input_embeds: Optional[List[List[float]]]
|
||||||
# The multimodal inputs
|
# The multimodal inputs
|
||||||
mm_inputs: Optional[PickleWrapper] # Pickled Optional[MultimodalProcessorOutput]
|
mm_inputs: Optional[MultimodalProcessorOutput]
|
||||||
token_type_ids: Optional[List[int]]
|
token_type_ids: Optional[List[int]]
|
||||||
# The sampling parameters
|
# The sampling parameters
|
||||||
sampling_params: SamplingParams
|
sampling_params: SamplingParams
|
||||||
@@ -1042,12 +1044,10 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True):
|
|||||||
cache_salt: Optional[str] = None
|
cache_salt: Optional[str] = None
|
||||||
|
|
||||||
def wrap_pickle_fields(self):
|
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.mm_data_mooncake = wrap_as_pickle(self.mm_data_mooncake)
|
||||||
self.time_stats = wrap_as_pickle(self.time_stats)
|
self.time_stats = wrap_as_pickle(self.time_stats)
|
||||||
|
|
||||||
def unwrap_pickle_fields(self):
|
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.mm_data_mooncake = unwrap_from_pickle(self.mm_data_mooncake)
|
||||||
self.time_stats = unwrap_from_pickle(self.time_stats)
|
self.time_stats = unwrap_from_pickle(self.time_stats)
|
||||||
|
|
||||||
@@ -1307,7 +1307,7 @@ class TokenizedEmbeddingReqInput(BaseReq, kw_only=True):
|
|||||||
# The input token ids
|
# The input token ids
|
||||||
input_ids: Optional[array] # array[int]
|
input_ids: Optional[array] # array[int]
|
||||||
# The multimodal inputs
|
# The multimodal inputs
|
||||||
mm_inputs: Optional[PickleWrapper] # Pickled Optional[MultimodalProcessorOutput]
|
mm_inputs: Optional[MultimodalProcessorOutput]
|
||||||
# The token type ids
|
# The token type ids
|
||||||
token_type_ids: Optional[List[int]]
|
token_type_ids: Optional[List[int]]
|
||||||
# Dummy sampling params for compatibility
|
# Dummy sampling params for compatibility
|
||||||
@@ -1332,11 +1332,9 @@ class TokenizedEmbeddingReqInput(BaseReq, kw_only=True):
|
|||||||
time_stats: Optional[PickleWrapper] = None
|
time_stats: Optional[PickleWrapper] = None
|
||||||
|
|
||||||
def wrap_pickle_fields(self):
|
def wrap_pickle_fields(self):
|
||||||
self.mm_inputs = wrap_as_pickle(self.mm_inputs)
|
|
||||||
self.time_stats = wrap_as_pickle(self.time_stats)
|
self.time_stats = wrap_as_pickle(self.time_stats)
|
||||||
|
|
||||||
def unwrap_pickle_fields(self):
|
def unwrap_pickle_fields(self):
|
||||||
self.mm_inputs = unwrap_from_pickle(self.mm_inputs)
|
|
||||||
self.time_stats = unwrap_from_pickle(self.time_stats)
|
self.time_stats = unwrap_from_pickle(self.time_stats)
|
||||||
|
|
||||||
|
|
||||||
@@ -2329,6 +2327,8 @@ def _check_all_req_types():
|
|||||||
for class_type in all_classes:
|
for class_type in all_classes:
|
||||||
# check its name
|
# check its name
|
||||||
name = class_type[0]
|
name = class_type[0]
|
||||||
|
if class_type[1].__module__ != __name__:
|
||||||
|
continue
|
||||||
if name in _IGNORE_REQ_TYPES_CHECK:
|
if name in _IGNORE_REQ_TYPES_CHECK:
|
||||||
continue
|
continue
|
||||||
is_io_struct = (
|
is_io_struct = (
|
||||||
@@ -2368,53 +2368,6 @@ def unwrap_from_pickle(obj: Optional[object]) -> Optional[object]:
|
|||||||
return pickle.loads(obj.data)
|
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(
|
_struct_types = tuple(
|
||||||
cls
|
cls
|
||||||
for cls in BaseReq.__subclasses__()
|
for cls in BaseReq.__subclasses__()
|
||||||
@@ -2429,14 +2382,18 @@ _primitive_types = (int, float, bool, bytes)
|
|||||||
_all_types = _struct_types + _primitive_types
|
_all_types = _struct_types + _primitive_types
|
||||||
|
|
||||||
_msgpack_encoder = msgspec.msgpack.Encoder(enc_hook=enc_hook)
|
_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()
|
_USE_PICKLE_IPC = envs.SGLANG_USE_PICKLE_IPC.get()
|
||||||
|
|
||||||
|
|
||||||
def hook_custom_types(*new_types: Type):
|
def hook_custom_types(*new_types: Type):
|
||||||
global _msgpack_decoder, _all_types
|
global _msgpack_decoder, _all_types
|
||||||
_all_types = tuple(dict.fromkeys(_all_types + new_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:
|
def _maybe_wrap_pickle(obj: Any) -> Any:
|
||||||
|
|||||||
@@ -67,6 +67,7 @@ from typing import (
|
|||||||
Optional,
|
Optional,
|
||||||
Set,
|
Set,
|
||||||
Tuple,
|
Tuple,
|
||||||
|
TypeAlias,
|
||||||
Union,
|
Union,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -145,6 +146,7 @@ INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
|
|||||||
# Constant used as the base offset for MM (multimodal) pad values.
|
# Constant used as the base offset for MM (multimodal) pad values.
|
||||||
# This ensures pad_values don't overlap with valid text token IDs.
|
# This ensures pad_values don't overlap with valid text token IDs.
|
||||||
MM_PAD_SHIFT_VALUE = 1_000_000
|
MM_PAD_SHIFT_VALUE = 1_000_000
|
||||||
|
_MM_HASH_MASK = (1 << 64) - 1
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -314,8 +316,12 @@ class MultimodalInputFormat(Enum):
|
|||||||
PRECOMPUTED_EMBEDDING = auto()
|
PRECOMPUTED_EMBEDDING = auto()
|
||||||
|
|
||||||
|
|
||||||
@dataclasses.dataclass
|
# Msgpack-native containers and Ext-decoded tensor/transport leaves. Tuple
|
||||||
class MultimodalDataItem:
|
# 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).
|
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.
|
For example, if there are 3 images and 1 audio, there will be 4 MultimodalDataItems.
|
||||||
@@ -326,39 +332,49 @@ class MultimodalDataItem:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
modality: Modality
|
modality: Modality
|
||||||
hash: int = None
|
hash: Optional[int] = None
|
||||||
pad_value: int = None
|
pad_value: Optional[int] = None
|
||||||
offsets: Optional[list] = None
|
offsets: Optional[List[Tuple[int, int]]] = None
|
||||||
|
|
||||||
format: MultimodalInputFormat = MultimodalInputFormat.NORMAL
|
format: MultimodalInputFormat = MultimodalInputFormat.NORMAL
|
||||||
|
|
||||||
# the raw features returned by processor, e.g. pixel_values or audio_features
|
# 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
|
# the precomputed embeddings, passed as final encoder embeddings
|
||||||
# One and only one of the feature and precomputed_embeddings will be empty
|
# 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
|
# Processor-owned tensors/arrays/scalars/transports. msgspec rejects a
|
||||||
model_specific_data: dict[str, Any] = dataclasses.field(default_factory=dict)
|
# 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):
|
def __post_init__(self) -> None:
|
||||||
if (
|
if self.hash is not None:
|
||||||
"model_specific_data" in self.__dict__
|
msgspec.Struct.__setattr__(self, "hash", self.hash & _MM_HASH_MASK)
|
||||||
and name in self.__dict__["model_specific_data"]
|
|
||||||
):
|
def __getattr__(self, name: str) -> MultimodalDataValue:
|
||||||
return self.__dict__["model_specific_data"][name]
|
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:
|
else:
|
||||||
raise AttributeError(
|
self.model_specific_data[name] = value
|
||||||
f"'{self.__class__.__name__}' object has no attribute '{name}'"
|
|
||||||
)
|
|
||||||
|
|
||||||
def __setitem__(self, key: str, value: Any):
|
def __setitem__(self, key: str, value: MultimodalDataValue) -> None:
|
||||||
if key in self.__dict__:
|
setattr(self, key, value)
|
||||||
self.__dict__[key] = value
|
|
||||||
else:
|
|
||||||
self.model_specific_data[key] = value
|
|
||||||
|
|
||||||
def set(self, key: str, value: Any):
|
def set(self, key: str, value: MultimodalDataValue) -> None:
|
||||||
self.__setitem__(key, value)
|
self.__setitem__(key, value)
|
||||||
|
|
||||||
def set_hash(self, hash_value: int) -> None:
|
def set_hash(self, hash_value: int) -> None:
|
||||||
@@ -491,8 +507,9 @@ class MultimodalDataItem:
|
|||||||
return min(requested_count, proxy_count)
|
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).
|
"""Raw output from multimodal processors before scheduler-side preparation (pad, hash).
|
||||||
|
|
||||||
This is the typed replacement for the dict previously returned by
|
This is the typed replacement for the dict previously returned by
|
||||||
|
|||||||
@@ -41,7 +41,6 @@ from sglang.srt.models.parakeet import ProjectedParakeet
|
|||||||
from sglang.srt.models.radio import RadioModel
|
from sglang.srt.models.radio import RadioModel
|
||||||
from sglang.srt.models.utils import WeightsMapper
|
from sglang.srt.models.utils import WeightsMapper
|
||||||
from sglang.srt.multimodal.evs import EVS, EVSConfig
|
from sglang.srt.multimodal.evs import EVS, EVSConfig
|
||||||
from sglang.srt.multimodal.evs.evs_module import VideoEVSDataItem
|
|
||||||
from sglang.srt.utils import add_prefix
|
from sglang.srt.utils import add_prefix
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -148,8 +147,11 @@ class NemotronH_Nano_VL_V2(EVS):
|
|||||||
|
|
||||||
if audio_items:
|
if audio_items:
|
||||||
for item in visual_items:
|
for item in visual_items:
|
||||||
if isinstance(item, VideoEVSDataItem):
|
if (
|
||||||
item.pre_chunked_input_ids = input_ids
|
item.is_video()
|
||||||
|
and "pre_chunked_input_ids" in item.model_specific_data
|
||||||
|
):
|
||||||
|
item.set("pre_chunked_input_ids", input_ids)
|
||||||
|
|
||||||
return input_ids
|
return input_ids
|
||||||
|
|
||||||
|
|||||||
@@ -13,7 +13,6 @@
|
|||||||
# ==============================================================================
|
# ==============================================================================
|
||||||
|
|
||||||
|
|
||||||
import dataclasses
|
|
||||||
import typing
|
import typing
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from dataclasses import dataclass
|
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
|
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)
|
@dataclass(kw_only=True)
|
||||||
class EVSEmbeddingResult(EmbeddingResult):
|
class EVSEmbeddingResult(EmbeddingResult):
|
||||||
"""
|
"""
|
||||||
@@ -65,14 +51,12 @@ class EVSEmbeddingResult(EmbeddingResult):
|
|||||||
input_ids: torch.Tensor,
|
input_ids: torch.Tensor,
|
||||||
offsets: list[tuple[int, int]],
|
offsets: list[tuple[int, int]],
|
||||||
*,
|
*,
|
||||||
item: VideoEVSDataItem,
|
item: MultimodalDataItem,
|
||||||
extend_prefix_len: int,
|
extend_prefix_len: int,
|
||||||
extend_seq_len: int,
|
extend_seq_len: int,
|
||||||
) -> tuple[torch.Tensor, list[tuple[int, int]]]:
|
) -> tuple[torch.Tensor, list[tuple[int, int]]]:
|
||||||
assert len(input_ids) == extend_seq_len
|
assert len(input_ids) == extend_seq_len
|
||||||
assert isinstance(
|
assert item.is_video() and "pre_chunked_input_ids" in item.model_specific_data
|
||||||
item, VideoEVSDataItem
|
|
||||||
), f"Expected VideoEVSDataItem, got {type(item)}"
|
|
||||||
pre_chunked_input_ids = item.pre_chunked_input_ids
|
pre_chunked_input_ids = item.pre_chunked_input_ids
|
||||||
filler_token_id = item.pad_value
|
filler_token_id = item.pad_value
|
||||||
input_ids_list = replace_offsets_with_tokens_per_frame(
|
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.
|
Apply EVS pruning to video embeddings.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
items: List containing a single VideoEVSDataItem with video features.
|
items: List containing a single EVS video item with video features.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
EVSEmbeddingResult with pruned embeddings and actual token counts per frame.
|
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)}"
|
assert len(items) == 1, f"Expected 1 item, got {len(items)}"
|
||||||
item = items[0]
|
item = items[0]
|
||||||
assert isinstance(
|
assert item.is_video() and "thw_grids" in item.model_specific_data
|
||||||
item, VideoEVSDataItem
|
|
||||||
), f"Expected VideoEVSDataItem with modality VIDEO, got {item}"
|
|
||||||
|
|
||||||
q = self.evs_config.video_pruning_rate
|
q = self.evs_config.video_pruning_rate
|
||||||
merge = self.evs_config.spatial_merge_size
|
merge = self.evs_config.spatial_merge_size
|
||||||
|
|||||||
@@ -20,7 +20,7 @@ from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
|
|||||||
from sglang.utils import logger
|
from sglang.utils import logger
|
||||||
|
|
||||||
from .evs_core import tokens_per_frame
|
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(
|
def _non_evs_data_items(
|
||||||
@@ -99,23 +99,25 @@ class EVSProcessor:
|
|||||||
items = []
|
items = []
|
||||||
if image is not None:
|
if image is not None:
|
||||||
image_thw_grids = [(1, rows, cols)] * num_images
|
image_thw_grids = [(1, rows, cols)] * num_images
|
||||||
item = EVSDataItem(
|
item = MultimodalDataItem(
|
||||||
modality=Modality.IMAGE,
|
modality=Modality.IMAGE,
|
||||||
feature=image,
|
feature=image,
|
||||||
offsets=image_offsets,
|
offsets=image_offsets,
|
||||||
thw_grids=image_thw_grids,
|
model_specific_data={"thw_grids": image_thw_grids},
|
||||||
)
|
)
|
||||||
items.append(item)
|
items.append(item)
|
||||||
if video is not None:
|
if video is not None:
|
||||||
video_thw_grids = [
|
video_thw_grids = [
|
||||||
(num_frames, rows, cols) for num_frames in frames_per_video
|
(num_frames, rows, cols) for num_frames in frames_per_video
|
||||||
]
|
]
|
||||||
item = VideoEVSDataItem(
|
item = MultimodalDataItem(
|
||||||
modality=Modality.VIDEO,
|
modality=Modality.VIDEO,
|
||||||
feature=video,
|
feature=video,
|
||||||
offsets=video_offsets,
|
offsets=video_offsets,
|
||||||
thw_grids=video_thw_grids,
|
model_specific_data={
|
||||||
pre_chunked_input_ids=input_ids_list,
|
"thw_grids": video_thw_grids,
|
||||||
|
"pre_chunked_input_ids": input_ids_list,
|
||||||
|
},
|
||||||
)
|
)
|
||||||
items.append(item)
|
items.append(item)
|
||||||
return items
|
return items
|
||||||
|
|||||||
@@ -238,7 +238,7 @@ class TransformersAutoMultimodalProcessor(BaseMultimodalProcessor):
|
|||||||
else "token_type_ids"
|
else "token_type_ids"
|
||||||
)
|
)
|
||||||
if token_type_key in processor_output:
|
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:
|
if self.mm_tokens.image_token_id is not None:
|
||||||
ret.im_token_id = self.mm_tokens.image_token_id
|
ret.im_token_id = self.mm_tokens.image_token_id
|
||||||
|
|||||||
@@ -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}")
|
||||||
@@ -1,7 +1,29 @@
|
|||||||
import copy
|
import copy
|
||||||
import unittest
|
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 (
|
from sglang.test.ci.ci_register import (
|
||||||
register_amd_ci,
|
register_amd_ci,
|
||||||
register_cpu_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")
|
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):
|
class TestGenerateReqInputNormalization(CustomTestCase):
|
||||||
"""Test the normalization of GenerateReqInput for batch processing and different input formats."""
|
"""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
|
from types import SimpleNamespace
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
import torch
|
||||||
|
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
from sglang.test.test_utils import run_doctests
|
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)
|
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__":
|
if __name__ == "__main__":
|
||||||
import sys
|
import sys
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user