from typing import Any, Callable, Optional, Union import torch from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput from vllm.sampling_params import RequestOutputKind from vllm.transformers_utils.tokenizer import AnyTokenizer from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason from vllm.v1.engine.detokenizer import IncrementalDetokenizer from vllm.v1.engine.logprobs import LogprobsProcessor from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor 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 from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) class OmniRequestState(RequestState): """Request state for omni models, tracking multimodal outputs. Extends the base RequestState with support for accumulating multimodal tensor outputs (e.g., images, audio, latents) that are produced incrementally during generation. """ def __init__( self, *args, **kwargs, ): super().__init__(*args, **kwargs) self.mm_type: Optional[str] = None self.mm_accumulated: Optional[torch.Tensor] = None @classmethod def from_new_request( cls, tokenizer: AnyTokenizer, request: EngineCoreRequest, prompt: Optional[str], parent_req: Optional[ParentRequest], request_index: int, queue: Optional[Any], log_stats: bool, ) -> "OmniRequestState": if sampling_params := request.sampling_params: if not sampling_params.detokenize: tokenizer = None output_kind = sampling_params.output_kind logprobs_processor = LogprobsProcessor.from_new_request( tokenizer=tokenizer, request=request, ) detokenizer = IncrementalDetokenizer.from_new_request( tokenizer=tokenizer, request=request, ) max_tokens_param = sampling_params.max_tokens top_p = sampling_params.top_p n = sampling_params.n temperature = sampling_params.temperature else: logprobs_processor = None detokenizer = None max_tokens_param = None top_p = None n = None temperature = None assert request.pooling_params is not None output_kind = request.pooling_params.output_kind return cls( request_id=request.request_id, parent_req=parent_req, request_index=request_index, lora_name=(request.lora_request.name if request.lora_request is not None else None), output_kind=output_kind, prompt=prompt, prompt_token_ids=request.prompt_token_ids, prompt_embeds=request.prompt_embeds, logprobs_processor=logprobs_processor, detokenizer=detokenizer, max_tokens_param=max_tokens_param, top_p=top_p, n=n, temperature=temperature, arrival_time=request.arrival_time, queue=queue, log_stats=log_stats, ) def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], mm_type: Optional[str]) -> None: """Add a multimodal tensor to the accumulated outputs. Accumulates multimodal tensors (e.g., image latents, audio waveforms) that are produced incrementally during generation. Tensors are concatenated along the first dimension. Args: tensor: Optional tensor to add. If None, this is a no-op. mm_type: Optional modality type identifier (e.g., "image", "audio"). Used to categorize the output type. """ if tensor is None: return try: if mm_type: self.mm_type = (mm_type or "").lower() t = tensor.detach() try: t = t.to("cpu") except Exception: # Best-effort CPU move; keep original device if conversion fails logger.debug("Failed to move multimodal tensor to CPU", exc_info=True) if self.mm_accumulated is None: self.mm_accumulated = t else: self.mm_accumulated = torch.cat([self.mm_accumulated, t], dim=0) except Exception: # Log and continue without crashing the output pipeline logger.exception("Error accumulating multimodal tensor") # Override: do not route to pooling-only path; always create completion # outputs, and attach pooling_result into the CompletionOutput. def make_request_output( self, new_token_ids: list[int], pooling_output: Optional[torch.Tensor], finish_reason: Optional[FinishReason], stop_reason: Optional[Union[int, str]], kv_transfer_params: Optional[dict[str, Any]] = None, ) -> Optional[Union[OmniRequestOutput, PoolingRequestOutput]]: """Create a request output from generation results. Creates a RequestOutput or PoolingRequestOutput from the generated tokens and accumulated multimodal outputs. Attaches multimodal tensors to the completion output if available. Args: new_token_ids: List of newly generated token IDs pooling_output: Optional pooling output tensor finish_reason: Optional finish reason indicating why generation stopped stop_reason: Optional stop reason (token ID or stop string) kv_transfer_params: Optional KV cache transfer parameters Returns: OmniRequestOutput or PoolingRequestOutput if output should be emitted (based on finish status and output kind), None otherwise """ finished = finish_reason is not None final_only = self.output_kind == RequestOutputKind.FINAL_ONLY if not finished and final_only: return None request_id = self.request_id output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) if self.parent_req is None: outputs = [output] else: request_id, outputs, finished = self.parent_req.get_outputs(request_id, output) if not outputs: return None return self._new_request_output(request_id, outputs, finished, kv_transfer_params) def _new_completion_output( self, token_ids: list[int], finish_reason: Optional[FinishReason], stop_reason: Optional[Union[int, str]], ) -> Any: # Reuse base text/logprobs logic, then annotate with pooling_result. base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason) try: if self.mm_accumulated is not None: tensor = self.mm_accumulated try: tensor = tensor.detach().to("cpu") except Exception: logger.debug( "Failed to move accumulated multimodal tensor to CPU", exc_info=True, ) # Attach on the completion output for downstream consumers. if not hasattr(base_output, "multimodal_output"): setattr(base_output, "multimodal_output", {}) setattr(base_output, "multimodal_output", {self.mm_type: tensor}) except Exception: logger.exception("Error in _new_completion_output") return base_output class MultimodalOutputProcessor(VLLMOutputProcessor): """Handles multimodal output processing by normalizing EngineCoreOutput before delegating to the base vLLM OutputProcessor. Strategy: - Route by EngineCoreOutput.output_type when present ("image", "text+image", "latents", "text"). - Fallback to pooling/text heuristics when output_type is absent. - Mutate EngineCoreOutput in-place to ensure vLLM's base processor can produce the correct RequestOutput/PoolingRequestOutput. - Allow custom per-modality handlers via register_handler(). """ def __init__( self, tokenizer: AnyTokenizer, log_stats: bool, engine_core_output_type: Optional[str] = None, ): """Initialize the multimodal output processor. Args: tokenizer: Tokenizer for detokenizing text outputs log_stats: Whether to log statistics engine_core_output_type: Optional output type specification (e.g., "image", "audio", "latents"). Used to route outputs to appropriate processors. If None, output type is inferred. """ 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: """Register a custom handler for a specific modality. Allows custom processing logic for specific output modalities. The handler is called before default processing for outputs matching the specified modality. Args: modality: Modality name (e.g., "image", "audio", "latents") handler: Callable that takes an EngineCoreOutput and processes it """ self.output_handlers[modality.lower()] = handler def add_request( self, request: EngineCoreRequest, prompt: Optional[str], parent_req: Optional[ParentRequest] = None, request_index: int = 0, queue: Optional[RequestOutputCollector] = None, ) -> None: """Add a new request to be processed. Creates an OmniRequestState for the request and registers it for output processing. Args: request: Engine core request to add prompt: Optional prompt string for the request parent_req: Optional parent request for parallel sampling request_index: Index of the request in the batch queue: Optional queue for collecting outputs Raises: 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 = OmniRequestState.from_new_request( tokenizer=self.tokenizer, request=request, prompt=prompt, parent_req=parent_req, request_index=request_index, queue=queue, log_stats=self.log_stats, ) self.request_states[request_id] = req_state self.lora_states.add_request(req_state) if parent_req: self.parent_requests[parent_req.request_id] = parent_req def process_outputs( self, engine_core_outputs: list[EngineCoreOutput], engine_core_timestamp: Optional[float] = None, iteration_stats: Optional[IterationStats] = 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: 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, ) # 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 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: mm_key = req_state.mm_type or "latents" if not hasattr(ro, "multimodal_output"): setattr(ro, "multimodal_output", {}) ro.multimodal_output[mm_key] = 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 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, ) # ---- routing helpers ---- def _route_and_normalize(self, eco: EngineCoreOutput) -> None: output_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower() # Custom handler first (if registered) if output_type in self.output_handlers: try: self.output_handlers[output_type](eco) # Fall through to default fixups in case the handler left gaps except Exception: logger.exception("Error in custom output handler for %s", output_type) if output_type == "image": self._process_image_output(eco) elif output_type in ("text+image", "text,image", "image+text"): self._process_text_image_output(eco) elif output_type in ("latents", "latent"): self._process_latents_output(eco) elif output_type in ("audio", "speech"): self._process_audio_output(eco) elif output_type == "text": self._process_text_output(eco) else: # Fallback heuristic if eco.pooling_output is not None: self._process_pooling_output(eco) else: self._process_text_output(eco) # ---- modality processors ---- def _process_image_output(self, eco: EngineCoreOutput) -> None: """Ensure image tensors are surfaced via pooling_output for vLLM.""" if eco.pooling_output is None: tensor = self._extract_from_multimodal_outputs(eco, keys=("image", "images", "pixel_values", "pixels")) if tensor is not None: eco.pooling_output = tensor def _process_text_image_output(self, eco: EngineCoreOutput) -> None: """Allow text+image outputs. Text path stays as new_token_ids; image/latents route via pooling_output.""" # Preserve text tokens as-is; ensure pooling_output carries image/latents if eco.pooling_output is None: tensor = self._extract_from_multimodal_outputs( eco, keys=( "image", "images", "pixel_values", "pixels", "latent", "latents", "z", ), ) if tensor is not None: eco.pooling_output = tensor def _process_latents_output(self, eco: EngineCoreOutput) -> None: """Ensure latent tensors are surfaced via pooling_output.""" if eco.pooling_output is None: tensor = self._extract_from_multimodal_outputs(eco, keys=("latent", "latents", "z", "posterior")) if tensor is not None: eco.pooling_output = tensor def _process_audio_output(self, eco: EngineCoreOutput) -> None: """Ensure audio tensors are surfaced via pooling_output.""" if eco.pooling_output is None: tensor = self._extract_from_multimodal_outputs( eco, keys=("audio", "audios", "wav", "waveform", "audio_pcm", "pcm") ) if tensor is not None: eco.pooling_output = tensor def _process_text_output(self, eco: EngineCoreOutput) -> None: """No-op; base processor will detokenize new_token_ids → text.""" return def _process_pooling_output(self, eco: EngineCoreOutput) -> None: """Optional sanity checks for pooling tensor.""" if eco.pooling_output is None: return if not isinstance(eco.pooling_output, torch.Tensor): # Best-effort: convert to tensor if it's a list/ndarray-like try: eco.pooling_output = torch.as_tensor(eco.pooling_output) except Exception: pass def _extract_from_multimodal_outputs(self, eco: EngineCoreOutput, keys: tuple[str, ...]) -> Optional[torch.Tensor]: mm = getattr(eco, "multimodal_outputs", None) if not isinstance(mm, dict): return None for k in keys: v = mm.get(k) if isinstance(v, torch.Tensor): return v # Try the first tensor in the dict as a fallback for v in mm.values(): if isinstance(v, torch.Tensor): return v return None