minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat: make mm_inputs msgpack-native (#29656)

Browse Source 
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
ishandhanani
2026-08-20 14:30:07 -07:00
committed by GitHub
co-authored by Alex Nails
parent 5a100d9086
commit 0f744b6848
9 changed files with 690 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/io_struct.py
+12 -55
View File
@@ -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:
python/sglang/srt/managers/schedule_batch.py
+43 -26
View File
@@ -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
python/sglang/srt/models/nano_nemotron_vl.py
+5 -3
View File
@@ -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
python/sglang/srt/multimodal/evs/evs_module.py
+4 -22
View File
@@ -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
python/sglang/srt/multimodal/evs/evs_processor.py
+8 -6
View File
@@ -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
python/sglang/srt/multimodal/processors/transformers_auto.py
+1 -1
View File
@@ -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
python/sglang/srt/utils/msgpack_utils.py
+248
View File
@@ -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}")