diff --git a/vllm_omni/engine/input_processor.py b/vllm_omni/engine/input_processor.py index 91e12f6a6..37cd4143e 100644 --- a/vllm_omni/engine/input_processor.py +++ b/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) diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index ab5cf878c..4d17d4f8d 100644 --- a/vllm_omni/engine/output_processor.py +++ b/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 ---- diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index ec291e43e..f8c0f5724 100644 --- a/vllm_omni/inputs/data.py +++ b/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 diff --git a/vllm_omni/inputs/parse.py b/vllm_omni/inputs/parse.py index a3700e61f..f21af5d2c 100644 --- a/vllm_omni/inputs/parse.py +++ b/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 diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 4eeadbff4..db7918770 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/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