[Feat] Support scalar types in AdditionalInformationEntry (#2105)

Signed-off-by: Nick Cao <ncao@redhat.com>

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 = {}

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.

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 {}

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

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)