Update Dockerfile for vllm wheel version and enhance SHM connector tests

- Updated Dockerfile to use a new vllm wheel version. - Modified test_put_get_inline to clarify that it always uses SHM and updated metadata assertions. - Added handling for speculative configuration in GPU model runners to delay clearing connector metadata. - Refactored payload decoding in OmniGPUModelRunner for better clarity and maintainability. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-24 05:45:48 +00:00
1 parent e22ae62396
commit 861698f88f
@@ -18,7 +18,7 @@ RUN apt-get update && \
 # Note: --prerelease=allow is required because nightly/commit wheels use
 # pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default,
 # causing it to fall back to the older stable release from PyPI.
-RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/944ffb59680c0210ec54ddb43a3c7ef015e1f842
+RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/ec85340531e83962877fa683e23f1d5d17cd786a


 RUN uv pip uninstall --system nvidia-cublas-cu12
@@ -64,14 +64,13 @@ def shm_connector():


 def test_put_get_inline(shm_connector):
-    """Test inline transfer for small data."""
+    """Test transfer for small data (currently always uses SHM)."""
    data = {"small": "data"}
-    # Ensure data is smaller than threshold (100 bytes)

    success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
    assert success is True
-    assert "inline_bytes" in metadata
-    assert "shm" not in metadata
+    assert "shm" in metadata
+    assert "size" in metadata

    # Retrieve
    retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata)
@@ -34,6 +34,7 @@ def _make_runner(multimodal_outputs):
    runner.use_async_scheduling = False
    runner.device = torch.device("cpu")
    runner.supports_mm_inputs = False
+    runner.speculative_config = None
    return runner


@@ -3,30 +3,26 @@ from dataclasses import dataclass, field
 from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
 from vllm.v1.request import Request

-from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
+from vllm_omni.engine import AdditionalInformationPayload


@dataclass
 class OmniNewRequestData(NewRequestData):
    """New request data for omni models with embeddings support.

-    Extends NewRequestData to include prompt embeddings and additional
-    information for direct transfer between pipeline stages.
+    Extends NewRequestData to include additional information for direct
+    transfer between pipeline stages.
+
+    Note: prompt_embeds is inherited from NewRequestData
+    (torch.Tensor | None).

    Args:
-        prompt_embeds: Optional serialized prompt embeddings payload
-            (overrides parent's torch.Tensor type with PromptEmbedsPayload
-            for cross-process serialization)
        external_req_id: Optional external request ID for tracking
        additional_information: Optional serialized additional information
            dictionary containing tensors or lists
    """

-    # Optional serialized prompt embeddings (override parent type for serialization)
-    prompt_embeds: PromptEmbedsPayload | None = None  # type: ignore[assignment]
-    # Optional external request ID for tracking
    external_req_id: str | None = None
-    # Optional serialized additional information
    additional_information: AdditionalInformationPayload | None = None

    @classmethod
@@ -48,6 +48,7 @@ class SharedMemoryConnector(OmniConnectorBase):
            payload = self.serialize_obj(data)
            size = len(payload)

+            # Currently, we always use SHM.
            if True:
                # Use Shared Memory
                lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock"
@@ -56,19 +56,19 @@ class AdditionalInformationPayload(msgspec.Struct):
 class OmniEngineCoreRequest(EngineCoreRequest):
    """Engine core request for omni models with embeddings support.

-    Extends the base EngineCoreRequest with support for prompt embeddings
-    and additional information payloads, enabling direct transfer of
-    pre-computed embeddings between pipeline stages.
+    Extends the base EngineCoreRequest with support for additional
+    information payloads, enabling direct transfer of pre-computed data
+    between pipeline stages.
+
+    Note: prompt_embeds is inherited from EngineCoreRequest
+    (torch.Tensor | None). PromptEmbedsPayload should be decoded to
+    torch.Tensor before constructing this request.

    Attributes:
-        prompt_embeds: Optional serialized prompt embeddings payload for
-            direct transfer between stages
        additional_information: Optional serialized additional information
            dictionary containing tensors or lists to pass along with the request
    """

-    # Optional prompt embeddings (direct-transfer version)
-    prompt_embeds: PromptEmbedsPayload | None = None
    # Optional additional information dictionary (serialized)
    additional_information: AdditionalInformationPayload | None = None

@@ -24,6 +24,7 @@ from vllm_ascend.utils import enable_sp, lmhead_tp_enable
 from vllm_ascend.worker.model_runner_v1 import NPUModelRunner

 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner

 if TYPE_CHECKING:
    from vllm.v1.core.sched.output import SchedulerOutput
@@ -194,54 +195,16 @@ class OmniNPUModelRunner(NPUModelRunner):
            )
            self.requests[req_id] = req_state

-            #  -------------------------------------- Omni-new -------------------------------------------------
-            # If prompt embeddings are provided, decode and attach to inter_data
-            try:
-                if getattr(new_req_data, "prompt_embeds", None) is not None:
-                    payload = new_req_data.prompt_embeds
-                    dtype = getattr(np, payload.dtype)
-                    arr = np.frombuffer(payload.data, dtype=dtype)
-                    arr = arr.reshape(payload.shape)
-                    pe_cpu = torch.from_numpy(arr)
-                    # Store temporarily on CPU; later moved to device in builder
-                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
-                    # Also replace payload with Tensor for user visibility in
-                    # scheduler_output
-                    try:
-                        new_req_data.prompt_embeds = pe_cpu  # type: ignore[assignment]
-                    except Exception:
-                        pass
-            except Exception as e:
-                logger.error(f"Error decoding prompt embeds: {e}")
-            # Decode additional_information payloads (dictionary)
-            try:
-                if getattr(new_req_data, "additional_information", None) is not None:
-                    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
+            # Decode and store omni-specific payloads on CPU.
+            pe_cpu = OmniGPUModelRunner._resolve_prompt_embeds_cpu(new_req_data.prompt_embeds)
+            if pe_cpu is not None:
+                setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)

-                        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
-                    if info_dict:
-                        setattr(
-                            self.requests[req_id],
-                            "additional_information_cpu",
-                            info_dict,
-                        )
-            except Exception as e:
-                logger.error(f"Error decoding additional information: {e}")
-                pass
-            #  -------------------------------------- Omni-new -------------------------------------------------
+            info_dict = OmniGPUModelRunner._resolve_additional_information(
+                getattr(new_req_data, "additional_information", None)
+            )
+            if info_dict:
+                setattr(self.requests[req_id], "additional_information_cpu", info_dict)

            if sampling_params and sampling_params.prompt_logprobs is not None:
                self.num_prompt_logprobs[req_id] = (
@@ -627,59 +590,24 @@ class OmniNPUModelRunner(NPUModelRunner):
                self.eplb_updator.forward_end()
            return hidden_states, hidden_states

-    def _decode_and_store_request_payloads(self, scheduler_output: "SchedulerOutput") -> None:
-        """Decode per-request prompt_embeds and additional_information for newly
-        scheduled requests and store them to CPU in the request state.
-        This version avoids hard dependency on payload classes by duck-typing."""
-        try:
-            new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
-            if not new_reqs:
-                return
-            for nr in new_reqs:
-                req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
-                if req_id is None:
-                    continue
-                # prompt_embeds
-                payload_pe = getattr(nr, "prompt_embeds", None)
-                pe_cpu = None
-                if payload_pe is not None:
-                    if isinstance(payload_pe, torch.Tensor):
-                        pe_cpu = payload_pe.detach().to("cpu").contiguous()
-                    else:
-                        # Try duck-typing a payload with data/shape/dtype
-                        data = getattr(payload_pe, "data", None)
-                        shape = getattr(payload_pe, "shape", None)
-                        if data is not None and shape is not None:
-                            dt = np.dtype(getattr(payload_pe, "dtype", "float32"))
-                            arr = np.frombuffer(data, dtype=dt)
-                            arr = arr.reshape(shape)
-                            pe_cpu = torch.from_numpy(arr.copy())
-                if pe_cpu is not None and req_id in self.requests:
-                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
-                # additional_information
-                payload_info = getattr(nr, "additional_information", None)
-                if payload_info is not None:
-                    info_dict = {}
-                    if isinstance(payload_info, dict):
-                        info_dict = payload_info
-                    else:
-                        # Try duck-typing a payload with entries, each entry may have
-                        # tensor_data/tensor_dtype/tensor_shape or list_data
-                        entries = getattr(payload_info, "entries", None)
-                        if isinstance(entries, dict):
-                            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_dict[k] = torch.from_numpy(arr.copy())
-                                else:
-                                    info_dict[k] = getattr(entry, "list_data", None)
-                    if info_dict and req_id in self.requests:
-                        setattr(self.requests[req_id], "additional_information_cpu", info_dict)
-        except Exception as e:
-            logger.error(f"Error decoding prompt_embeds / additional_information: {e}")
+    def _decode_and_store_request_payloads(
+        self,
+        scheduler_output: "SchedulerOutput",
+    ) -> None:
+        """Decode per-request prompt_embeds and additional_information for
+        newly scheduled requests and store them on CPU in the request state.
+        """
+        new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
+        for nr in new_reqs:
+            req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
+            if req_id is None or req_id not in self.requests:
+                continue
+            pe_cpu = OmniGPUModelRunner._resolve_prompt_embeds_cpu(getattr(nr, "prompt_embeds", None))
+            if pe_cpu is not None:
+                setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
+            info_dict = OmniGPUModelRunner._resolve_additional_information(getattr(nr, "additional_information", None))
+            if info_dict:
+                setattr(self.requests[req_id], "additional_information_cpu", info_dict)

    def _gather_runtime_additional_information(self) -> list[dict]:
        """Gather per-request additional_information stored in request state in batch order."""
@@ -273,6 +273,9 @@ class GPUARModelRunner(OmniGPUModelRunner):

        # Run the model.
        # Use persistent buffers for CUDA graphs.
+        # When spec decode is enabled, delay clearing connector metadata
+        # until after draft model runs in sample_tokens.
+        clear_kv_metadata = self.speculative_config is None
        with (
            set_forward_context(
                attn_metadata,
@@ -285,7 +288,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
                slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
            ),
            record_function_or_nullcontext("gpu_model_runner: forward"),
-            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+            self.maybe_get_kv_connector_output(
+                scheduler_output, clear_metadata=clear_kv_metadata
+            ) as kv_connector_output,
        ):
            model_output = self._model_forward(
                input_ids=input_ids,
@@ -519,6 +524,12 @@ class GPUARModelRunner(OmniGPUModelRunner):
            # tokens on the CPU, so they are run after bookkeeping.
            propose_draft_token_ids(valid_sampled_token_ids)

+        # Clear KV connector metadata after draft model runs (if spec decode).
+        # This was deferred from target model forward to allow draft model
+        # to also save its KV cache.
+        if self.speculative_config is not None:
+            self.clear_kv_connector_metadata()
+
        with record_function_or_nullcontext("gpu_model_runner: eplb"):
            self.eplb_step()

@@ -264,6 +264,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):

        # Run the model.
        # Use persistent buffers for CUDA graphs.
+        # When spec decode is enabled, delay clearing connector metadata
+        # until after draft model runs in sample_tokens.
+        clear_kv_metadata = self.speculative_config is None
        with (
            set_forward_context(
                attn_metadata,
@@ -276,7 +279,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
            ),
            record_function_or_nullcontext("Forward"),
-            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+            self.maybe_get_kv_connector_output(
+                scheduler_output, clear_metadata=clear_kv_metadata
+            ) as kv_connector_output,
        ):
            outputs = self._run_generation_model(
                input_ids=input_ids,
@@ -345,6 +350,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
        ) = self.execute_model_state
        self.execute_model_state = None

+        # Clear KV connector metadata after draft model runs (if spec decode).
+        if self.speculative_config is not None:
+            self.clear_kv_connector_metadata()
+
        pooler_output: list[object] = []
        if isinstance(multimodal_outputs, torch.Tensor):
            assert multimodal_outputs.shape[0] == 1, (
@@ -466,6 +475,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
        remove_lora: bool = True,
        is_graph_capturing: bool = False,
        num_active_loras: int = 0,
+        activate_lora: bool | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Run a dummy forward pass to warm up/profile run or capture the
@@ -489,7 +499,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                (1 token) and prefill (multiple tokens) requests.
            remove_lora: If False, dummy LoRAs are not destroyed after the run
            num_active_loras: Number of active LoRAs to capture for.
+            activate_lora: Backward-compatible override for LoRA activation.
        """
+        if activate_lora is None:
+            activate_lora = num_active_loras > 0
+
        mm_config = self.vllm_config.model_config.multimodal_config
        if mm_config and mm_config.mm_encoder_only:
            # The current dummy run only covers LM execution, so we can skip it.
@@ -566,7 +580,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                # `force_has_lora` is used for cudagraph capture; because LoRA is
                # activated later in the context manager, but we need to know the
                # LoRA state when determining the batch descriptor for capture
-                force_has_lora=num_active_loras > 0,
+                force_has_lora=activate_lora,
                # Capture shape specialization for specific active LoRA counts.
                force_num_active_loras=num_active_loras,
            )
@@ -639,8 +653,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
            self.lora_config,
            num_scheduled_tokens,
            num_sampled_tokens,
+            activate_lora,
            remove_lora,
-            num_active_loras,
        ):
            # Make sure padding doesn't exceed max_num_tokens
            assert num_tokens_padded <= self.max_num_tokens
@@ -727,7 +741,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                # lora cases when cudagraph_specialize_lora is enabled. This is a
                # short term mitigation for issue mentioned in
                # https://github.com/vllm-project/vllm/issues/28334
-                if self.compilation_config.cudagraph_specialize_lora and num_active_loras > 0:
+                if self.compilation_config.cudagraph_specialize_lora and activate_lora:
                    use_cudagraphs = False

                self.drafter.dummy_run(
@@ -240,52 +240,14 @@ class OmniGPUModelRunner(GPUModelRunner):
            )
            self.requests[req_id] = req_state

-            # If prompt embeddings are provided, decode and attach to inter_data
-            try:
-                if getattr(new_req_data, "prompt_embeds", None) is not None:
-                    payload = new_req_data.prompt_embeds
-                    dtype = getattr(np, payload.dtype)
-                    arr = np.frombuffer(payload.data, dtype=dtype)
-                    arr = arr.reshape(payload.shape)
-                    pe_cpu = torch.from_numpy(arr)
-                    # Store temporarily on CPU; later moved to device in builder
-                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
-                    # Also replace payload with Tensor for user visibility in
-                    # scheduler_output
-                    try:
-                        new_req_data.prompt_embeds = pe_cpu  # type: ignore[assignment]
-                    except Exception:
-                        pass
-            except Exception as e:
-                logger.error(f"Error decoding prompt embeds: {e}")
-            # Decode additional_information payloads (dictionary)
-            try:
-                if getattr(new_req_data, "additional_information", None) is not None:
-                    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
+            # Decode and store omni-specific payloads on CPU.
+            pe_cpu = self._resolve_prompt_embeds_cpu(new_req_data.prompt_embeds)
+            if pe_cpu is not None:
+                setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)

-                        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
-                    if info_dict:
-                        setattr(
-                            self.requests[req_id],
-                            "additional_information_cpu",
-                            info_dict,
-                        )
-            except Exception as e:
-                logger.error(f"Error decoding additional information: {e}")
-                pass
+            info_dict = self._resolve_additional_information(getattr(new_req_data, "additional_information", None))
+            if info_dict:
+                setattr(self.requests[req_id], "additional_information_cpu", info_dict)

            if sampling_params and sampling_params.prompt_logprobs is not None:
                self.num_prompt_logprobs[req_id] = (
@@ -776,59 +738,92 @@ class OmniGPUModelRunner(GPUModelRunner):
        logit_indices_device = torch.from_numpy(logit_indices).to(self.device, non_blocking=True)
        return hidden_states, hidden_states[logit_indices_device]

-    def _decode_and_store_request_payloads(self, scheduler_output: "SchedulerOutput") -> None:
-        """Decode per-request prompt_embeds and additional_information for newly
-        scheduled requests and store them to CPU in the request state.
-        This version avoids hard dependency on payload classes by duck-typing."""
+    # ------------------------------------------------------------------
+    # Payload decoding helpers (torch.Tensor passthrough + legacy
+    # PromptEmbedsPayload / AdditionalInformationPayload support)
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _resolve_prompt_embeds_cpu(
+        pe: "torch.Tensor | object | None",
+    ) -> torch.Tensor | None:
+        """Convert *prompt_embeds* to a contiguous CPU tensor.
+
+        Accepts:
+        - ``torch.Tensor`` – moved to CPU as-is (the normal path after
+          upstream added ``prompt_embeds`` to ``EngineCoreRequest``).
+        - Legacy ``PromptEmbedsPayload`` (or any duck-typed object with
+          ``.data``, ``.shape``, ``.dtype``) – decoded via numpy.
+        - ``None`` – returns ``None``.
+        """
+        if pe is None:
+            return None
        try:
-            new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
-            if not new_reqs:
-                return
-            for nr in new_reqs:
-                req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
-                if req_id is None:
-                    continue
-                # prompt_embeds
-                payload_pe = getattr(nr, "prompt_embeds", None)
-                pe_cpu = None
-                if payload_pe is not None:
-                    if isinstance(payload_pe, torch.Tensor):
-                        pe_cpu = payload_pe.detach().to("cpu").contiguous()
-                    else:
-                        # Try duck-typing a payload with data/shape/dtype
-                        data = getattr(payload_pe, "data", None)
-                        shape = getattr(payload_pe, "shape", None)
-                        if data is not None and shape is not None:
-                            dt = np.dtype(getattr(payload_pe, "dtype", "float32"))
-                            arr = np.frombuffer(data, dtype=dt)
-                            arr = arr.reshape(shape)
-                            pe_cpu = torch.from_numpy(arr.copy())
-                if pe_cpu is not None and req_id in self.requests:
-                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
-                # additional_information
-                payload_info = getattr(nr, "additional_information", None)
-                if payload_info is not None:
-                    info_dict = {}
-                    if isinstance(payload_info, dict):
-                        info_dict = payload_info
-                    else:
-                        # Try duck-typing a payload with entries, each entry may have
-                        # tensor_data/tensor_dtype/tensor_shape or list_data
-                        entries = getattr(payload_info, "entries", None)
-                        if isinstance(entries, dict):
-                            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_dict[k] = torch.from_numpy(arr.copy())
-                                else:
-                                    info_dict[k] = getattr(entry, "list_data", None)
-                    if info_dict and req_id in self.requests:
-                        setattr(self.requests[req_id], "additional_information_cpu", info_dict)
-        except Exception as e:
-            logger.error(f"Error decoding prompt_embeds / additional_information: {e}")
+            if isinstance(pe, torch.Tensor):
+                return pe.detach().cpu().contiguous()
+            data = getattr(pe, "data", None)
+            shape = getattr(pe, "shape", None)
+            if data is not None and shape is not None:
+                dt = np.dtype(getattr(pe, "dtype", "float32"))
+                arr = np.frombuffer(data, dtype=dt).reshape(shape)
+                return torch.from_numpy(arr.copy())
+        except Exception:
+            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",
+    ) -> None:
+        """Decode per-request prompt_embeds and additional_information for
+        newly scheduled requests and store them on CPU in the request state.
+        """
+        new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
+        for nr in new_reqs:
+            req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
+            if req_id is None or req_id not in self.requests:
+                continue
+            pe_cpu = self._resolve_prompt_embeds_cpu(getattr(nr, "prompt_embeds", None))
+            if pe_cpu is not None:
+                setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
+            info_dict = self._resolve_additional_information(getattr(nr, "additional_information", None))
+            if info_dict:
+                setattr(self.requests[req_id], "additional_information_cpu", info_dict)

    def _gather_runtime_additional_information(self) -> list[dict]:
        """Gather per-request additional_information stored in request state in batch order."""