From 47ba887b93d62446b4c9ab9709eecbe63071e356 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Tue, 24 Mar 2026 10:29:38 -0400 Subject: [PATCH] [Feat] Support scalar types in AdditionalInformationEntry (#2105) Signed-off-by: Nick Cao --- vllm_omni/core/sched/omni_ar_scheduler.py | 9 ++-- vllm_omni/engine/__init__.py | 8 +++- vllm_omni/engine/serialization.py | 58 ++++++++++++++++++----- vllm_omni/engine/test_serialization.py | 35 ++++++++++++++ vllm_omni/worker/gpu_model_runner.py | 54 ++------------------- 5 files changed, 92 insertions(+), 72 deletions(-) create mode 100644 vllm_omni/engine/test_serialization.py diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 3629ade63..ea3b0a8cd 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -22,6 +22,7 @@ from vllm_omni.core.sched.output import OmniSchedulerOutput from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, ) +from vllm_omni.engine.serialization import deserialize_additional_information logger = init_logger(__name__) @@ -521,17 +522,13 @@ class OmniARScheduler(VLLMScheduler): # Also update request.additional_information for good measure add_info = getattr(request, "additional_information", None) # If additional_information is an AdditionalInformationPayload-like object, - # unpack list_data into a plain dict. + # unpack it into a plain dict. if ( add_info is not None and hasattr(add_info, "entries") and isinstance(getattr(add_info, "entries"), dict) ): - request.additional_information = { - k: getattr(v, "list_data") - for k, v in getattr(add_info, "entries").items() - if getattr(v, "list_data", None) is not None - } + request.additional_information = deserialize_additional_information(add_info) add_info = request.additional_information if add_info is None: request.additional_information = {} diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index c4d42f95d..c8a96e6d2 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -29,10 +29,11 @@ class PromptEmbedsPayload(msgspec.Struct): class AdditionalInformationEntry(msgspec.Struct): """One entry of additional_information. - Two supported forms are encoded: + Three supported forms are encoded: - tensor: data/shape/dtype - list: a Python list (msgspec-serializable) - Exactly one of (tensor_data, list_data) should be non-None. + - scalar: a Python scalar (msgspec-serializable) + Exactly one of (tensor_data, list_data, scalar_data) should be non-None. """ # Tensor form @@ -43,6 +44,9 @@ class AdditionalInformationEntry(msgspec.Struct): # List form list_data: list[Any] | None = None + # Scalar form + scalar_data: Any | None = None + class AdditionalInformationPayload(msgspec.Struct): """Serialized dictionary payload for additional_information. diff --git a/vllm_omni/engine/serialization.py b/vllm_omni/engine/serialization.py index 41146a01b..386c97814 100644 --- a/vllm_omni/engine/serialization.py +++ b/vllm_omni/engine/serialization.py @@ -4,6 +4,7 @@ from __future__ import annotations from typing import Any +import numpy as np import torch from vllm.logger import init_logger @@ -64,18 +65,49 @@ def serialize_additional_information( entries[key] = AdditionalInformationEntry(list_data=value) continue - if log_prefix is None: - logger.warning( - "Dropping unsupported additional_information key=%s type=%s", - key, - type(value).__name__, - ) - else: - logger.warning( - "[%s] Dropping unsupported additional_information key=%s type=%s", - log_prefix, - key, - type(value).__name__, - ) + entries[key] = AdditionalInformationEntry(scalar_data=value) return AdditionalInformationPayload(entries=entries) if entries else None + + +def deserialize_additional_information( + payload: dict | AdditionalInformationPayload | object | None, +) -> dict: + """Deserialize an *additional_information* payload into a plain dict. + + Accepts: + - ``dict`` – returned as-is. + - ``AdditionalInformationPayload`` (or duck-typed with + ``.entries``) – decoded entry-by-entry. + - ``None`` – returns ``{}``. + """ + + if payload is None: + return {} + + if isinstance(payload, dict): + return payload + + try: + entries = getattr(payload, "entries", None) + if not isinstance(entries, dict): + logger.exception("Failed to decode additional_information payload, entries field not a dict") + return {} + info: dict[str, object] = {} + for k, entry in entries.items(): + if getattr(entry, "tensor_data", None) is not None: + dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) + arr = np.frombuffer(entry.tensor_data, dtype=dt) + arr = arr.reshape(getattr(entry, "tensor_shape", ())) + info[k] = torch.from_numpy(arr.copy()) + elif getattr(entry, "list_data", None) is not None: + info[k] = entry.list_data + elif getattr(entry, "scalar_data", None) is not None: + info[k] = entry.scalar_data + else: + info[k] = None + return info + except Exception: + logger.exception("Failed to decode additional_information payload") + + return {} diff --git a/vllm_omni/engine/test_serialization.py b/vllm_omni/engine/test_serialization.py new file mode 100644 index 000000000..69c7fb8a1 --- /dev/null +++ b/vllm_omni/engine/test_serialization.py @@ -0,0 +1,35 @@ +import torch + +from vllm_omni.engine.serialization import deserialize_additional_information, serialize_additional_information + + +def test_serialize_additional_information(): + info = { + "tensor_f32": torch.rand(3, dtype=torch.float32), + "tensor_f64": torch.rand(3, 3, dtype=torch.float64), + "list": [1, 2, 3], + "scalar": 1, + } + + payload = serialize_additional_information(info) + + assert payload.entries["tensor_f32"].tensor_dtype == "float32" + assert payload.entries["tensor_f32"].tensor_shape == [3] + assert payload.entries["tensor_f32"].tensor_data == info["tensor_f32"].numpy().tobytes() + + assert payload.entries["tensor_f64"].tensor_dtype == "float64" + assert payload.entries["tensor_f64"].tensor_shape == [3, 3] + assert payload.entries["tensor_f64"].tensor_data == info["tensor_f64"].numpy().tobytes() + + assert payload.entries["list"].list_data == info["list"] + assert payload.entries["scalar"].scalar_data == info["scalar"] + + deserialized = deserialize_additional_information(payload) + + # compare tensors separately first + assert torch.equal(deserialized["tensor_f32"], info["tensor_f32"]) + assert torch.equal(deserialized["tensor_f64"], info["tensor_f64"]) + del deserialized["tensor_f32"], info["tensor_f32"] + del deserialized["tensor_f64"], info["tensor_f64"] + + assert deserialized == info diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index a7441f674..8e5689986 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -21,6 +21,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices +from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding from vllm_omni.model_executor.models.output_templates import OmniOutput @@ -365,21 +366,7 @@ class OmniGPUModelRunner(GPUModelRunner): "additional_information on request data is deprecated, use model_intermediate_buffer" ) payload_info = new_req_data.additional_information - info_dict = {} - if isinstance(payload_info, dict): - info_dict = payload_info - else: - from vllm_omni.engine import AdditionalInformationPayload - - if isinstance(payload_info, AdditionalInformationPayload): - for k, entry in payload_info.entries.items(): - if entry.tensor_data is not None: - dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) - arr = np.frombuffer(entry.tensor_data, dtype=dt) - arr = arr.reshape(entry.tensor_shape) - info_dict[k] = torch.from_numpy(arr.copy()) - else: - info_dict[k] = entry.list_data + info_dict = deserialize_additional_information(payload_info) if info_dict: self.model_intermediate_buffer[req_id] = info_dict setattr( @@ -921,41 +908,6 @@ class OmniGPUModelRunner(GPUModelRunner): logger.exception("Failed to decode prompt_embeds payload") return None - @staticmethod - def _resolve_additional_information( - payload: "dict | object | None", - ) -> dict[str, object]: - """Convert an *additional_information* payload to a plain dict. - - Accepts: - - ``dict`` – returned as-is. - - ``AdditionalInformationPayload`` (or duck-typed with - ``.entries``) – decoded entry-by-entry. - - ``None`` – returns ``{}``. - """ - if payload is None: - return {} - if isinstance(payload, dict): - return payload - try: - entries = getattr(payload, "entries", None) - if not isinstance(entries, dict): - return {} - info: dict[str, object] = {} - for k, entry in entries.items(): - tensor_data = getattr(entry, "tensor_data", None) - if tensor_data is not None: - dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) - arr = np.frombuffer(tensor_data, dtype=dt) - arr = arr.reshape(getattr(entry, "tensor_shape", ())) - info[k] = torch.from_numpy(arr.copy()) - else: - info[k] = getattr(entry, "list_data", None) - return info - except Exception: - logger.exception("Failed to decode additional_information payload") - return {} - def _decode_and_store_request_payloads( self, scheduler_output: "SchedulerOutput", @@ -978,7 +930,7 @@ class OmniGPUModelRunner(GPUModelRunner): logger.warning_once( "additional_information on request data is deprecated, use model_intermediate_buffer" ) - info_dict = self._resolve_additional_information(info_payload) + info_dict = deserialize_additional_information(info_payload) if info_dict: self.model_intermediate_buffer[req_id] = info_dict setattr(self.requests[req_id], "additional_information_cpu", info_dict)