[Enhancement] Refactor OmniInputProcessor and related classes for improved prompt handling

- Updated OmniInputProcessor to streamline input processing, including enhanced handling of prompt embeddings and additional information.
- Introduced OmniTextPrompt and OmniEmbedsInputs to support new input structures with optional metadata.
- Modified input preprocessing methods to accommodate new prompt types and ensure compatibility with multimodal inputs.
- Improved error handling and type checking for data parallel configurations.

Signed-off-by: [Your Name] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/engine/input_processor.py

@@ -1,7 +1,9 @@
+import os
 import time
 from collections.abc import Mapping
 from typing import Any, cast
 
+import numpy as np
 import torch
 from vllm.config import VllmConfig
 from vllm.inputs import ProcessorInputs, PromptType
@@ -9,12 +11,13 @@ from vllm.inputs.parse import split_enc_dec_inputs
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
 from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalUUIDDict
+from vllm.multimodal.processing.context import set_request_id
 from vllm.multimodal.utils import argsort_mm_positions
 from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams
-from vllm.tokenizers import TokenizerLike
 from vllm.utils import length_from_prompt_token_ids_or_embeds
+from vllm.utils.torch_utils import set_default_torch_num_threads
 from vllm.v1.engine.input_processor import InputProcessor
 
 from vllm_omni.engine import (
@@ -75,13 +78,12 @@ class OmniInputProcessor(InputProcessor):
     def __init__(
         self,
         vllm_config: VllmConfig,
-        tokenizer: TokenizerLike,
         mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
     ):
-        super().__init__(vllm_config, tokenizer, mm_registry)
+        super().__init__(vllm_config, mm_registry)
         self.input_preprocessor = OmniInputPreprocessor(
             self.model_config,
-            self.tokenizer,
+            vllm_config.observability_config,
             mm_registry,
             mm_processor_cache=self.mm_processor_cache,
         )
@@ -97,7 +99,8 @@ class OmniInputProcessor(InputProcessor):
         trace_headers: Mapping[str, str] | None = None,
         priority: int = 0,
         data_parallel_rank: int | None = None,
-    ) -> tuple[str | None, OmniEngineCoreRequest]:
+        resumable: bool = False,
+    ) -> OmniEngineCoreRequest:
         """Process input prompt into an engine core request.
 
         Converts a prompt (text, tokens, or multimodal) into an
@@ -130,9 +133,12 @@ class OmniInputProcessor(InputProcessor):
         self._validate_lora(lora_request)
         self._validate_params(params)
 
-        data_parallel_size = self.vllm_config.parallel_config.data_parallel_size
-        if data_parallel_rank is not None and not (0 <= data_parallel_rank < data_parallel_size):
-            raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {data_parallel_size}).")
+        parallel_config = self.vllm_config.parallel_config
+        dp_size = parallel_config.data_parallel_size
+        dp_local_size = parallel_config.data_parallel_size_local
+        num_ranks = dp_local_size if parallel_config.local_engines_only else dp_size
+        if data_parallel_rank is not None and not (0 <= data_parallel_rank < num_ranks):
+            raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {num_ranks}).")
 
         if arrival_time is None:
             arrival_time = time.time()
@@ -154,7 +160,7 @@ class OmniInputProcessor(InputProcessor):
         else:
             # Otherwise, use user-provided uuids as multimodal hash overrides
             # if provided.
-            self._validate_multi_modal_uuids(prompt)
+            self._validate_mm_uuids(prompt)
             if isinstance(prompt, dict):
                 mm_uuids = cast(MultiModalUUIDDict | None, prompt.get("multi_modal_uuids"))
             else:
@@ -164,11 +170,19 @@ class OmniInputProcessor(InputProcessor):
         # 1. Tokenize text prompt, with LoRA request if one exists.
         # 2. For multimodal models with a merged preprocessor, preprocess
         #   multimodal data and expand prompt token ids accordingly.
-        processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
-            prompt,
-            tokenization_kwargs=tokenization_kwargs,
-            mm_uuids=mm_uuids,
-        )
+        num_threads = int(os.environ.get("OMP_NUM_THREADS", "1"))
+        if "OMP_NUM_THREADS" not in os.environ:
+            logger.debug_once(
+                "OMP_NUM_THREADS is not set; defaulting Torch threads to %d for input preprocessing.",
+                num_threads,
+            )
+
+        with set_request_id(request_id), set_default_torch_num_threads(num_threads):
+            processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
+                prompt,
+                tokenization_kwargs=tokenization_kwargs,
+                mm_uuids=mm_uuids,
+            )
 
         current_platform.validate_request(
             prompt=prompt,
@@ -181,13 +195,13 @@ class OmniInputProcessor(InputProcessor):
         encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)
         self._validate_model_inputs(encoder_inputs, decoder_inputs)
 
-        # Mypy does not always properly infer the types of some elements of
-        # discriminated unions of TypedDicts, because of how it handles
-        # inheritance of TypedDict. If we explicitly extract the items we want
-        # we can avoid type errors from using `dict.get` later in the method.
-        _prompt_str: str | None = None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt")
-        prompt_token_ids = decoder_inputs["prompt_token_ids"] if decoder_inputs["type"] != "embeds" else None
-        prompt_embeds = decoder_inputs["prompt_embeds"] if decoder_inputs["type"] == "embeds" else None
+        # Normalize decoder prompt access across TypedDict variants.
+        if decoder_inputs["type"] == "embeds":
+            prompt_token_ids = None
+            prompt_embeds = decoder_inputs["prompt_embeds"]
+        else:
+            prompt_token_ids = decoder_inputs["prompt_token_ids"]
+            prompt_embeds = decoder_inputs.get("prompt_embeds")
 
         sampling_params = None
         pooling_params = None
@@ -219,35 +233,26 @@ class OmniInputProcessor(InputProcessor):
 
             mm_features = []
             for modality, idx in sorted_mm_idxs:
+                base_mm_hash = decoder_mm_hashes[modality][idx]
                 mm_features.append(
                     MultiModalFeatureSpec(
                         data=decoder_mm_inputs[modality][idx],
                         modality=modality,
-                        identifier=decoder_mm_hashes[modality][idx],
+                        identifier=self._get_mm_identifier(base_mm_hash, lora_request),
                         mm_position=decoder_mm_positions[modality][idx],
+                        mm_hash=base_mm_hash,
                     )
                 )
 
-        # Serialize prompt_embeds and additional_information if provided
-        # (direct-transfer path)
-        prompt_embeds_payload: PromptEmbedsPayload | None = None
+        # Compatibility: decode serialized prompt embeds if provided.
+        if isinstance(prompt_embeds, PromptEmbedsPayload):
+            prompt_embeds = self._decode_prompt_embeds(prompt_embeds)
+
         additional_information_payload: AdditionalInformationPayload | None = None
-        pe: torch.Tensor | None = decoder_inputs.get("prompt_embeds")  # type: ignore[operator]
-        if pe is not None:
-            if pe.ndim != 2:
-                raise ValueError("prompt_embeds must be of shape (seq_len, hidden_size)")
-            # Move to CPU and ensure contiguous memory for stable serialization
-            pe_cpu = pe.detach().to("cpu").contiguous()
-            seq_len, hidden_size = pe_cpu.shape
-            dtype_str = self._dtype_to_name(pe_cpu.dtype)
-            data_bytes = pe_cpu.numpy().tobytes()
-            prompt_embeds_payload = PromptEmbedsPayload(
-                data=data_bytes,
-                shape=[int(seq_len), int(hidden_size)],
-                dtype=dtype_str,
-            )
-        raw_info: dict[str, Any] | None = decoder_inputs.get("additional_information")  # type: ignore[operator]
-        if raw_info is not None:
+        raw_info: dict[str, Any] | AdditionalInformationPayload | None = decoder_inputs.get("additional_information")
+        if isinstance(raw_info, AdditionalInformationPayload):
+            additional_information_payload = raw_info
+        elif raw_info is not None:
             entries: dict[str, AdditionalInformationEntry] = {}
             for key, value in raw_info.items():
                 if isinstance(value, torch.Tensor):
@@ -279,6 +284,14 @@ class OmniInputProcessor(InputProcessor):
             priority=priority,
             data_parallel_rank=data_parallel_rank,
             trace_headers=trace_headers,
-            prompt_embeds=prompt_embeds_payload,
+            prompt_embeds=prompt_embeds,
             additional_information=additional_information_payload,
+            resumable=resumable,
         )
+
+    @staticmethod
+    def _decode_prompt_embeds(payload: PromptEmbedsPayload) -> torch.Tensor:
+        dtype = getattr(np, payload.dtype)
+        arr = np.frombuffer(payload.data, dtype=dtype)
+        arr = arr.reshape(payload.shape)
+        return torch.from_numpy(arr)

vllm_omni/engine/output_processor.py

@@ -2,6 +2,7 @@ from ast import Dict
 from collections.abc import Callable
 from typing import Any
 
+import numpy as np
 import torch
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput
@@ -9,7 +10,11 @@ from vllm.sampling_params import RequestOutputKind
 from vllm.tokenizers import TokenizerLike
 from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
 from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
-from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestOutputCollector, RequestState
+from vllm.v1.engine.output_processor import (
+    OutputProcessorOutput,
+    RequestOutputCollector,
+    RequestState,
+)
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.metrics.stats import IterationStats
 
@@ -140,6 +145,7 @@ class OmniRequestState(RequestState):
         finish_reason: FinishReason | None,
         stop_reason: int | str | None,
         kv_transfer_params: dict[str, Any] | None = None,
+        routed_experts: np.ndarray | None = None,
     ) -> OmniRequestOutput | PoolingRequestOutput | None:
         """Create a request output from generation results.
 
@@ -158,12 +164,27 @@ class OmniRequestState(RequestState):
             OmniRequestOutput or PoolingRequestOutput if output should be
             emitted (based on finish status and output kind), None otherwise
         """
+        # Pooling-only requests should follow base behavior.
+        if self.detokenizer is None and pooling_output is not None:
+            return super().make_request_output(
+                new_token_ids,
+                pooling_output,
+                finish_reason,
+                stop_reason,
+                kv_transfer_params,
+                routed_experts,
+            )
+
         finished = finish_reason is not None
         final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
 
         if not finished and final_only:
             return None
 
+        # Consolidate accumulated tensors when finishing.
+        if finished:
+            self._consolidate_multimodal_tensors()
+
         if self.stream_interval > 1:
             assert self.detokenizer is not None
 
@@ -185,7 +206,7 @@ class OmniRequestState(RequestState):
                 self.sent_tokens_offset = len(self.detokenizer.output_token_ids)
 
         request_id = self.request_id
-        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)
+        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts)
 
         if self.parent_req is None:
             outputs = [output]
@@ -201,9 +222,10 @@ class OmniRequestState(RequestState):
         token_ids: list[int],
         finish_reason: FinishReason | None,
         stop_reason: int | str | None,
+        routed_experts: np.ndarray | None = None,
     ) -> Any:
         # Reuse base text/logprobs logic, then annotate with pooling_result.
-        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
+        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason, routed_experts)
         try:
             if self.mm_accumulated is not None:
                 # Attach accumulated multimodal dict on the completion output
@@ -251,7 +273,6 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         super().__init__(tokenizer=tokenizer, log_stats=log_stats)
         self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
         self._reqid_to_mm_type: dict[str, str] = {}
-        self.request_states: dict[str, OmniRequestState] = {}
         self.engine_core_output_type = engine_core_output_type
 
     def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None:
@@ -291,8 +312,10 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             ValueError: If the request ID is already registered
         """
         request_id = request.request_id
-        if request_id in self.request_states:
-            raise ValueError(f"Request id {request_id} already running.")
+        req_state = self.request_states.get(request_id)
+        if req_state is not None:
+            self._update_streaming_request_state(req_state, request, prompt)
+            return
 
         req_state = OmniRequestState.from_new_request(
             tokenizer=self.tokenizer,
@@ -304,9 +327,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             log_stats=self.log_stats,
             stream_interval=self.stream_interval,
         )
+        if self._requests_drained.is_set():
+            self._requests_drained.clear()
         self.request_states[request_id] = req_state
         if parent_req:
             self.parent_requests[parent_req.request_id] = parent_req
+        self.external_req_ids[req_state.external_req_id].append(request_id)
 
     def process_outputs(
         self,
@@ -314,121 +340,27 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         engine_core_timestamp: float | None = None,
         iteration_stats: IterationStats | None = None,
     ) -> OutputProcessorOutput:
-        """Process engine core outputs into request outputs.
-
-        Converts EngineCoreOutput objects into RequestOutput objects,
-        handling multimodal outputs by routing them through appropriate
-        processors and accumulating tensors in request states.
-
-        Args:
-            engine_core_outputs: List of engine core outputs to process
-            engine_core_timestamp: Optional timestamp for the outputs
-            iteration_stats: Optional iteration statistics
-
-        Returns:
-            OutputProcessorOutput containing processed request outputs
-            and list of request IDs to abort
-        """
         self._reqid_to_mm_type.clear()
         for eco in engine_core_outputs:
             mm_type = (self.engine_core_output_type or "").lower()
             if mm_type:
                 self._reqid_to_mm_type[eco.request_id] = mm_type
             self._route_and_normalize(eco)
-
-        # Build RequestOutputs without delegating to base, so we can keep ids
-        request_outputs: list[Any] = []
-        reqs_to_abort: list[str] = []
-        for eco in engine_core_outputs:
-            req_id = eco.request_id
-            req_state = self.request_states.get(req_id)
-            if req_state is None:
+            req_state = self.request_states.get(eco.request_id)
+            if req_state is None or not isinstance(req_state, OmniRequestState):
                 continue
-
-            # 1) Stats
-            self._update_stats_from_output(req_state, eco, engine_core_timestamp, iteration_stats)
-
-            new_token_ids = eco.new_token_ids
-            pooling_output = eco.pooling_output
-            finish_reason = eco.finish_reason
-            stop_reason = eco.stop_reason
-            kv_transfer_params = eco.kv_transfer_params
-            req_state.num_cached_tokens = eco.num_cached_tokens
-            req_state.is_prefilling = False
-
-            # 2) Detokenize and logprobs when text path
-            assert req_state.detokenizer is not None
-            assert req_state.logprobs_processor is not None
-            stop_string = req_state.detokenizer.update(new_token_ids, finish_reason == FinishReason.STOP)
-            if stop_string:
-                finish_reason = FinishReason.STOP
-                stop_reason = stop_string
-            req_state.logprobs_processor.update_from_output(eco)
-
-            # 2.5) Accumulate multimodal tensors in RequestState
-            try:
-                mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
-                if pooling_output is not None and isinstance(req_state, OmniRequestState):
-                    req_state.add_multimodal_tensor(pooling_output, mm_type)
-            except Exception:
-                logger.debug(
-                    "Failed to accumulate multimodal tensor for request %s",
-                    req_id,
-                    exc_info=True,
+            if eco.pooling_output is not None and req_state.detokenizer is not None:
+                req_state.add_multimodal_tensor(
+                    eco.pooling_output,
+                    (getattr(eco, "output_type", self.engine_core_output_type) or "").lower(),
                 )
+                # Force text path in base processor for multimodal outputs.
+                eco.pooling_output = None
 
-            # 3) Create RequestOutput objects, forcing combined mode to keep ids
-            pooling_for_make = pooling_output
-            if pooling_output is not None and new_token_ids:
-                # Do not consume pooling path now; keep ids and attach mm later
-                pooling_for_make = None
-
-            # Consolidate any accumulated tensor lists before creating output
-            if isinstance(req_state, OmniRequestState) and finish_reason is not None:
-                req_state._consolidate_multimodal_tensors()
-
-            ro = req_state.make_request_output(
-                new_token_ids,
-                pooling_for_make,
-                finish_reason,
-                stop_reason,
-                kv_transfer_params,
-            )
-            if ro:
-                # Attach accumulated multimodal payload if any
-                try:
-                    if isinstance(req_state, OmniRequestState) and req_state.mm_accumulated is not None:
-                        if not hasattr(ro, "multimodal_output"):
-                            setattr(ro, "multimodal_output", {})
-                        ro.multimodal_output = req_state.mm_accumulated
-                except Exception:
-                    logger.exception("Error attaching multimodal payload in process_outputs")
-                if req_state.queue is not None:
-                    req_state.queue.put(ro)
-                else:
-                    request_outputs.append(ro)
-
-            # 4) Free completed
-            if finish_reason is not None:
-                self.request_states.pop(req_id)
-                parent_req = req_state.parent_req
-                if parent_req and not parent_req.child_requests:
-                    self.parent_requests.pop(parent_req.request_id, None)
-                if not self.request_states:
-                    self._requests_drained.set()
-                if not eco.finished:
-                    reqs_to_abort.append(req_id)
-                self._update_stats_from_finished(req_state, finish_reason, iteration_stats)
-                if self.tracer:
-                    self.do_tracing(eco, req_state, iteration_stats)
-                # Cleanup per-request mm state
-                if isinstance(req_state, OmniRequestState):
-                    req_state.mm_accumulated = None
-                    req_state.mm_type = None
-
-        return OutputProcessorOutput(
-            request_outputs=request_outputs,
-            reqs_to_abort=reqs_to_abort,
+        return super().process_outputs(
+            engine_core_outputs,
+            engine_core_timestamp=engine_core_timestamp,
+            iteration_stats=iteration_stats,
         )
 
     # ---- routing helpers ----

vllm_omni/inputs/data.py

@@ -6,8 +6,23 @@ except ImportError:
     # Python < 3.11: use typing_extensions
     from typing_extensions import NotRequired
 
+from typing import TypeAlias
+
 import torch
-from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt
+from vllm.inputs.data import (
+    EmbedsInputs,
+    EmbedsPrompt,
+    TextPrompt,
+    TokenInputs,
+    TokensPrompt,
+)
+
+
+class OmniTextPrompt(TextPrompt):
+    """Text prompt with optional embeddings and additional information."""
+
+    prompt_embeds: NotRequired[torch.Tensor]
+    additional_information: NotRequired[dict[str, Any]]
 
 
 class OmniTokensPrompt(TokensPrompt):
@@ -52,25 +67,20 @@ class OmniTokenInputs(TokenInputs):
 
 
 class OmniEmbedsPrompt(EmbedsPrompt):
-    """Embeddings prompt with optional additional information.
+    """Embeddings prompt with optional additional information."""
 
-    Extends EmbedsPrompt to support additional information payloads
-    for direct transfer between pipeline stages.
-
-    Attributes:
-        prompt_embeds: Optional tensor containing prompt embeddings
-        additional_information: Optional dictionary containing additional
-            information (tensors or lists) to pass along with the prompt
-    """
-
-    # New: optional prompt embeddings aligned with token ids
-    prompt_embeds: NotRequired[torch.Tensor]
-
-    # New: optional additional information dictionary
-    # Values may be torch.Tensor or list
     additional_information: NotRequired[dict[str, Any]]
 
 
+class OmniEmbedsInputs(EmbedsInputs):
+    """Embeddings inputs with optional additional information."""
+
+    additional_information: NotRequired[dict[str, Any]]
+
+
+OmniSingletonPrompt: TypeAlias = str | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt
+
+
 def token_inputs_omni(
     prompt_token_ids: list[int],
     prompt: str | None = None,
@@ -106,3 +116,17 @@ def token_inputs_omni(
         inputs["additional_information"] = additional_information
 
     return inputs
+
+
+def embeds_inputs_omni(
+    prompt_embeds: torch.Tensor,
+    cache_salt: str | None = None,
+    additional_information: dict[str, Any] | None = None,
+) -> OmniEmbedsInputs:
+    """Construct embeddings inputs with optional metadata."""
+    inputs: OmniEmbedsInputs = OmniEmbedsInputs(type="embeds", prompt_embeds=prompt_embeds)
+    if cache_salt is not None:
+        inputs["cache_salt"] = cache_salt
+    if additional_information is not None:
+        inputs["additional_information"] = additional_information
+    return inputs

vllm_omni/inputs/parse.py

@@ -1,4 +1,3 @@
-from vllm.inputs.data import SingletonPrompt
 from vllm.inputs.parse import (
     ParsedEmbedsPrompt,
     ParsedSingletonPrompt,
@@ -7,8 +6,12 @@ from vllm.inputs.parse import (
     ParsedTokensPrompt,
 )
 
+from vllm_omni.inputs.data import OmniSingletonPrompt
 
-def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPrompt:
+
+def parse_singleton_prompt_omni(
+    prompt: OmniSingletonPrompt,
+) -> ParsedSingletonPrompt:
     """Parse a singleton prompt into a typed parsed prompt.
 
     Handles omni-specific prompt types including tokens prompts with

vllm_omni/inputs/preprocess.py

@@ -1,12 +1,20 @@
 from typing import Any
 
 from typing_extensions import assert_never
-from vllm.inputs.data import SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt
+from vllm.inputs.data import SingletonInputs, SingletonPrompt
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
 from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict
 
-from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni
+from vllm_omni.inputs.data import (
+    OmniEmbedsInputs,
+    OmniEmbedsPrompt,
+    OmniTextPrompt,
+    OmniTokenInputs,
+    OmniTokensPrompt,
+    embeds_inputs_omni,
+    token_inputs_omni,
+)
 from vllm_omni.inputs.parse import parse_singleton_prompt_omni
 
 logger = init_logger(__name__)
@@ -22,7 +30,7 @@ class OmniInputPreprocessor(InputPreprocessor):
 
     def _process_text(
         self,
-        parsed_content: TextPrompt,
+        parsed_content: OmniTextPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
         *,
         mm_uuids: MultiModalUUIDDict | None = None,
@@ -62,7 +70,7 @@ class OmniInputPreprocessor(InputPreprocessor):
 
     def _process_tokens(
         self,
-        parsed_content: TokensPrompt,
+        parsed_content: OmniTokensPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
         *,
         mm_uuids: MultiModalUUIDDict | None = None,
@@ -80,6 +88,10 @@ class OmniInputPreprocessor(InputPreprocessor):
                 tokenization_kwargs=tokenization_kwargs,
                 mm_uuids=mm_uuids,
             )
+            if prompt_embeds is not None:
+                inputs["prompt_embeds"] = prompt_embeds
+            if additional_information is not None:
+                inputs["additional_information"] = additional_information
         else:
             inputs = token_inputs_omni(
                 prompt_token_ids=prompt_token_ids,
@@ -129,9 +141,20 @@ class OmniInputPreprocessor(InputPreprocessor):
             )
         if parsed["type"] == "str":
             return self._process_text(
-                TextPrompt(prompt=parsed["content"]),
+                OmniTextPrompt(prompt=parsed["content"]),
                 tokenization_kwargs=tokenization_kwargs,
                 mm_uuids=mm_uuids,
             )
 
         assert_never(parsed)
+
+    def _process_embeds(self, parsed_content: OmniEmbedsPrompt) -> OmniEmbedsInputs:
+        inputs = super()._process_embeds(parsed_content)
+        additional_information = parsed_content.get("additional_information")
+        if additional_information is not None:
+            inputs = embeds_inputs_omni(
+                prompt_embeds=inputs["prompt_embeds"],
+                cache_salt=inputs.get("cache_salt"),
+                additional_information=additional_information,
+            )
+        return inputs