From bb3f5b5cb771ecd34ddc9c3c6bb291ac68905874 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Mon, 5 Oct 2026 05:12:14 +0800 Subject: [PATCH] [Core] Preserve native scalar encoding in stage tensor frames Signed-off-by: Sy03 <1370724210@qq.com> --- .../omni_connectors/test_shm_ring.py | 26 ++++- .../omni_connectors/utils/tensor_frame.py | 95 ++++++++----------- 2 files changed, 66 insertions(+), 55 deletions(-) diff --git a/tests/distributed/omni_connectors/test_shm_ring.py b/tests/distributed/omni_connectors/test_shm_ring.py index 4b24523de..2a7fc58a3 100644 --- a/tests/distributed/omni_connectors/test_shm_ring.py +++ b/tests/distributed/omni_connectors/test_shm_ring.py @@ -10,6 +10,7 @@ import uuid from dataclasses import asdict from multiprocessing import shared_memory +import msgspec import pytest import torch from vllm.outputs import CompletionOutput, RequestOutput @@ -200,11 +201,34 @@ def test_disabled_ring_and_unsupported_tree_keep_legacy_serializer(edge): _, _, metadata = legacy.put("0", "1", "legacy", {"codes": torch.arange(3)}) assert "shm" in metadata assert receiver.get("0", "1", "legacy", metadata)[0]["codes"].tolist() == [0, 1, 2] - assert prepare_tensor_frame({"codes": torch.arange(3), "opaque": {1, 2}}) is None + assert prepare_tensor_frame({"codes": torch.arange(3), "opaque": object()}) is None assert sender.put("0", "1", "falsey", False)[0] assert receiver.get("0", "1", "falsey")[0] is False +def test_tensor_free_lists_reuse_native_serialization(): + payload = {"prompt_token_ids": list(range(8192)), "condition": [0.25] * 8192} + with SharedMemoryConnector({}) as connector: + frame = prepare_tensor_frame(payload) + assert isinstance(frame, bytes) + assert frame == connector.serialize_obj(payload) + assert connector.put("0", "1", "native-lists", payload)[0] + assert connector.get("0", "1", "native-lists")[0] == payload + + +def test_scalar_lists_and_literal_extensions_preserve_wire_semantics(): + literal = msgspec.msgpack.Ext(42, bytes(20)) + payload = {"condition": [0.25] * 8192, "codes": torch.arange(28), "extension": literal} + with SharedMemoryConnector({}) as connector: + _, _, metadata = connector.put("0", "1", "mixed-lists", payload) + assert "host_ring" in metadata + result = connector.get("0", "1", "mixed-lists")[0] + assert result["condition"] == payload["condition"] + assert torch.equal(result["codes"], payload["codes"]) + assert result["extension"].code == literal.code + assert result["extension"].data == literal.data + + def _locked_writer(name, ready, release): fd = os.open(f"/dev/shm/{name}", os.O_RDWR) try: diff --git a/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py b/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py index 5bc7ea82f..42868907b 100644 --- a/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py +++ b/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Owned CPU tensor trees with a msgpack header and raw tensor regions. -The format uses tagged nodes rather than reserved dictionary keys. It never +Tensor references use message-scoped msgpack extensions, preserving literal +extensions and native scalar containers without reserved dictionary keys. It never exposes a view of transport storage to a caller. Unsupported objects retain the ordinary OmniSerializer path. """ @@ -12,6 +13,7 @@ from __future__ import annotations import ctypes import math import struct +import uuid from dataclasses import dataclass import msgspec @@ -19,8 +21,9 @@ import torch _LENGTH = struct.Struct(" int: @@ -51,61 +54,47 @@ class TensorFrame: ctypes.memmove(base + offset, tensor.data_ptr(), tensor.nbytes) -def prepare_tensor_frame(payload: object, *, max_bytes: int | None = None) -> TensorFrame | None: - """Prepare a plain tree containing dense CPU tensors; otherwise return None. +def prepare_tensor_frame(payload: object, *, max_bytes: int | None = None) -> TensorFrame | bytes | None: + """Encode native containers once, or prepare their dense CPU tensor regions. - Retaining tensors is enough because the caller writes before put() returns. - This is not an asynchronous snapshot or a CUDA source ownership protocol. + Scalar/token lists stay in msgpack's C encoder. Tensor-free messages reuse + those encoded bytes, rather than constructing and discarding a tagged tree. + Sources remain caller-owned until the synchronous put() returns. """ tensors: list[torch.Tensor] = [] offsets: list[int] = [] descriptors: list[tuple[str, tuple[int, ...], int]] = [] data_size = 0 + nonce = uuid.uuid4().bytes - def visit(value: object, depth: int) -> list: + def encode_tensor(value: object) -> msgspec.msgpack.Ext: nonlocal data_size - if depth > 64: + if type(value) is not torch.Tensor: raise _UnsupportedTreeError - if type(value) in (type(None), bool, int, float, str, bytes): - return [_SCALAR, value] - if type(value) is torch.Tensor: - if value.device.type != "cpu" or value.layout != torch.strided or value.is_quantized: - raise _UnsupportedTreeError - tensor = value - offset = _align(data_size) - index = len(tensors) - tensors.append(tensor) - offsets.append(offset) - descriptors.append((str(tensor.dtype).removeprefix("torch."), tuple(tensor.shape), offset)) - data_size = offset + tensor.nbytes - if max_bytes is not None and data_size > max_bytes: - raise _UnsupportedTreeError - return [_TENSOR, index] - if type(value) is dict and all(type(key) is str for key in value): - return [_DICT, {key: visit(item, depth + 1) for key, item in value.items()}] - if isinstance(value, (list, tuple)) and type(value) in (list, tuple): - # msgpack represents both as arrays; match OmniSerializer's wire type. - return [_LIST, [visit(item, depth + 1) for item in value]] - if isinstance(value, msgspec.Struct): - try: - fields = msgspec.to_builtins(value, builtin_types=(torch.Tensor, bytes)) - except TypeError: - raise _UnsupportedTreeError from None - return visit(fields, depth + 1) - raise _UnsupportedTreeError + if value.device.type != "cpu" or value.layout != torch.strided or value.is_quantized: + raise _UnsupportedTreeError + offset = _align(data_size) + index = len(tensors) + tensors.append(value) + offsets.append(offset) + descriptors.append((str(value.dtype).removeprefix("torch."), tuple(value.shape), offset)) + data_size = offset + value.nbytes + if max_bytes is not None and data_size > max_bytes: + raise _UnsupportedTreeError + return msgspec.msgpack.Ext(_REFERENCE_CODE, _REFERENCE.pack(nonce, index)) try: - tree = visit(payload, 0) + tree = msgspec.msgpack.Encoder(enc_hook=encode_tensor).encode(payload) if not tensors: - return None - header = _ENCODER.encode([descriptors, tree]) - except (_UnsupportedTreeError, OverflowError): + return tree + header = _ENCODER.encode([descriptors, tree, nonce]) + except (_UnsupportedTreeError, TypeError, OverflowError, RecursionError): return None size = _align(4 + len(header)) + data_size if max_bytes is not None and size > max_bytes: return None - # Only materialize non-contiguous/flagged storage after the entire tree - # is known to fit. Large and unsupported payloads avoid an unused copy. + # Only materialize non-contiguous/flagged storage after the entire message + # is known to fit. Unsupported and large frames avoid an unused copy. for index, tensor in enumerate(tensors): if tensor.is_conj(): tensor = tensor.resolve_conj() @@ -127,7 +116,7 @@ def read_tensor_frame(buffer: memoryview) -> object: raise ValueError("tensor frame header exceeds payload") header_view = buffer[4 : 4 + header_size] try: - descriptors, tree = _DECODER.decode(header_view) + descriptors, tree, nonce = _DECODER.decode(header_view) finally: header_view.release() # All leaves share one independently owned allocation. This amortizes CPU @@ -159,16 +148,14 @@ def read_tensor_frame(buffer: memoryview) -> object: ctypes.memmove(tensor.data_ptr(), ctypes.addressof(anchor) + start, size) tensors.append(tensor) - def restore(node: list) -> object: - tag, value = node - if tag == _SCALAR: - return value - if tag == _TENSOR: - return tensors[value] - if tag == _DICT: - return {key: restore(item) for key, item in value.items()} - if tag == _LIST: - return [restore(item) for item in value] - raise ValueError("unknown tensor frame node") + def restore_reference(code: int, data: memoryview) -> object: + if code == _REFERENCE_CODE and len(data) == _REFERENCE.size: + reference_nonce, index = _REFERENCE.unpack(data) + if reference_nonce == nonce: + if index >= len(tensors): + raise ValueError("invalid tensor frame reference") + return tensors[index] + # Literal user extensions keep their existing msgpack wire semantics. + return msgspec.msgpack.Ext(code, bytes(data)) - return restore(tree) + return msgspec.msgpack.decode(tree, ext_hook=restore_reference)