from typing import Any import torch from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput 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.parallel_sampling import ParentRequest from vllm.v1.metrics.stats import IterationStats, RequestStateStats from vllm_omni.data_entry_keys import unflatten_payload from vllm_omni.engine.output_modality import DRAINABLE_MODALITIES 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, ): arrival_time = kwargs.get("arrival_time") if arrival_time is None and len(args) > 12: arrival_time = args[12] super().__init__(*args, **kwargs) self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0)) # Omni-specific: multimodal output accumulation # NOTE: Keys in mm_accumulated matter, because they dictate which # outputs are drained (i.e., only drain modality keys, don't drain # hidden states). self.mm_accumulated: dict[str, Any] = {} def apply_streaming_update(self, update) -> None: super().apply_streaming_update(update) self.native_text_stats.arrival_time = update.arrival_time @staticmethod def _to_cpu(x): """Try to convert to CPU tensor if needed.""" # TODO: Make this more robust and unify with other payload # building utils, we do this in multiple places. if isinstance(x, torch.Tensor): try: return x.detach().to("cpu", non_blocking=True).contiguous() except Exception: return x return x def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None: if payload is None: return mm_type = (mm_type or "").lower() try: if isinstance(payload, dict): # Keep payload flat (dotted keys like "hidden_states.layer_0") # during accumulation so that all values are tensors/scalars and # the merge logic below works correctly. Unflatten happens # later in _consolidate_multimodal_tensors after concatenation. incoming: dict[str, Any] = {} # TODO (Alex): Clean up and simplify key management target_key = mm_type or "hidden" for k, v in payload.items(): if k == "model_outputs": k = target_key elif k == "hidden" and target_key != "hidden": k = target_key incoming[k] = self._to_cpu(v) else: key = mm_type or "hidden" incoming = {key: self._to_cpu(payload)} if not self.mm_accumulated: self.mm_accumulated = incoming else: for k, v in incoming.items(): if k not in self.mm_accumulated: self.mm_accumulated[k] = v else: existing = self.mm_accumulated[k] if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor): self.mm_accumulated[k] = [existing, v] elif isinstance(v, torch.Tensor) and isinstance(existing, list): existing.append(v) elif isinstance(v, dict) and isinstance(existing, dict): for sk, sv in v.items(): if sk not in existing: existing[sk] = sv elif isinstance(sv, torch.Tensor) and isinstance(existing[sk], torch.Tensor): existing[sk] = [existing[sk], sv] elif isinstance(sv, torch.Tensor) and isinstance(existing[sk], list): existing[sk].append(sv) else: existing[sk] = sv else: self.mm_accumulated[k] = v except Exception: # Log and continue without crashing the output pipeline logger.exception("Error accumulating multimodal tensor") def _consolidate_multimodal_tensors(self) -> None: """Consolidate accumulated tensor lists into single tensors via concatenation. Only DELTA drains modality keys per-step, so they will never be lists here and can be skipped. For CUMULATIVE and FINAL_ONLY, modality keys accumulate across steps and need consolidation. """ if not self.mm_accumulated: return skip_modality = self.output_kind == RequestOutputKind.DELTA try: for k, v in self.mm_accumulated.items(): if skip_modality and k in DRAINABLE_MODALITIES: continue if isinstance(v, list) and v and isinstance(v[0], torch.Tensor): try: if k == "audio": # Preserve channel dimension when chunks are compatible; # fall back to 1-D waveform concatenation for uneven chunks. try: self.mm_accumulated[k] = torch.cat(v, dim=-1) except RuntimeError: self.mm_accumulated[k] = torch.cat([t.reshape(-1) for t in v], dim=0) elif k == "sr": # Sample rate is a constant scalar, keep last value. self.mm_accumulated[k] = v[-1] else: self.mm_accumulated[k] = torch.cat(v, dim=0) except Exception: # Keep last tensor on failure logger.warning(f"Error concatenating tensor for key {k}; keeping last tensor") self.mm_accumulated[k] = v[-1] elif isinstance(v, dict): for sk, sv in v.items(): if isinstance(sv, list) and sv and isinstance(sv[0], torch.Tensor): try: v[sk] = torch.cat(sv, dim=0) except Exception: v[sk] = sv[-1] except Exception: logger.exception("Error consolidating multimodal tensors") # Restore nested structure from flat dotted keys now that all tensor # lists have been concatenated into single tensors. try: self.mm_accumulated = unflatten_payload(self.mm_accumulated) except Exception: logger.exception("Error unflattening consolidated multimodal tensors") # 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: torch.Tensor | None, finish_reason: FinishReason | None, stop_reason: int | str | None, kv_transfer_params: dict[str, Any] | None = None, ) -> OmniRequestOutput | PoolingRequestOutput | None: """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 """ # 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, ) finished = finish_reason is not None is_final_only = self.output_kind == RequestOutputKind.FINAL_ONLY is_delta = self.output_kind == RequestOutputKind.DELTA if not finished and is_final_only: return None if finished or not is_delta: self._consolidate_multimodal_tensors() if self.stream_interval > 1: assert self.detokenizer is not None # Send output request only when # 1. It has finished, or # 2. It is the first token, or # 3. It has reached the stream interval number of tokens if not ( finished or self.sent_tokens_offset == 0 or self.detokenizer.num_output_tokens() - self.sent_tokens_offset >= self.stream_interval ): return None if is_delta: # Send tokens from the offset in DELTA mode, otherwise all # tokens are sent. new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :] self.sent_tokens_offset = self.detokenizer.num_output_tokens() external_req_id = self.external_req_id output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) if self.parent_req is None: outputs = [output] else: outputs, finished = self.parent_req.get_outputs(self.request_id, output) if not outputs: return None external_req_id = self.parent_req.external_req_id return self._new_request_output( external_req_id, outputs, finished, kv_transfer_params, ) def _new_completion_output( self, token_ids: list[int], finish_reason: FinishReason | None, stop_reason: int | str | None, ) -> Any: # Reuse base text/logprobs logic, then annotate with pooling_result. # Note: upstream _new_completion_output no longer accepts routed_experts # as a parameter; it reads from ``self.routed_experts_chunks`` internally. base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason) # Inter-stage processors need the full cumulative token sequence. # In DELTA mode, base_output.token_ids only has the latest step's # tokens, so we always store a cumulative copy here. base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids) if not hasattr(base_output, "multimodal_output"): setattr(base_output, "multimodal_output", {}) if self.mm_accumulated: mm_out = getattr(base_output, "multimodal_output") if isinstance(mm_out, dict): for k, v in self.mm_accumulated.items(): mm_out[k] = v else: setattr(base_output, "multimodal_output", self.mm_accumulated) if self.output_kind == RequestOutputKind.DELTA: for modality_key in DRAINABLE_MODALITIES: self.mm_accumulated.pop(modality_key, None) return base_output class MultimodalOutputProcessor(VLLMOutputProcessor): """Handles multimodal output processing by capturing pooling_output from EngineCoreOutput and accumulating it as multimodal tensors, before delegating to the base vLLM OutputProcessor for text handling. The actual data flow is: 1. For each EngineCoreOutput with pooling_output and a detokenizer: - Capture pooling_output into OmniRequestState.add_multimodal_tensor() - Clear eco.pooling_output to force text path in base processor 2. Base vLLM OutputProcessor handles text detokenization 3. On finish, _consolidate_multimodal_tensors() concatenates accumulated tensors 4. _new_completion_output() attaches mm_accumulated to CompletionOutput """ def __init__( self, tokenizer: TokenizerLike | None, *, log_stats: bool, stream_interval: int = 1, tracing_enabled: bool = False, engine_core_output_type: str | None = None, ): """Initialize the multimodal output processor. Args: tokenizer: Tokenizer for detokenizing text outputs log_stats: Whether to log statistics stream_interval: Stream interval for output generation engine_core_output_type: Optional output type specification (e.g., "image", "audio", "latent"). Used to tag multimodal outputs with the correct modality key. """ super().__init__( tokenizer=tokenizer, log_stats=log_stats, stream_interval=stream_interval, tracing_enabled=tracing_enabled, ) self.engine_core_output_type = engine_core_output_type self._native_text_metrics_by_request: dict[str, dict[str, Any]] = {} def _native_text_metric_record(self, request_id: str) -> dict[str, Any]: return self._native_text_metrics_by_request.setdefault( request_id, { "vllm_ttft_ms": 0.0, "vllm_tpot_ms": 0.0, "vllm_itl_ms": 0.0, "vllm_itls_ms": [], }, ) def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]: return self._native_text_metrics_by_request.pop(request_id, {}) def abort_requests(self, request_ids, internal: bool) -> list[str]: request_ids = list(request_ids) for request_id in request_ids: if internal: req_state = self.request_states.get(request_id) if req_state is not None: self._native_text_metrics_by_request.pop(req_state.external_req_id, None) else: self._native_text_metrics_by_request.pop(request_id, None) return super().abort_requests(request_ids, internal) def add_request( self, request: EngineCoreRequest, prompt: str | None, parent_req: ParentRequest | None = None, request_index: int = 0, queue: RequestOutputCollector | None = 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 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, request=request, prompt=prompt, parent_req=parent_req, request_index=request_index, queue=queue, log_stats=self.log_stats, stream_interval=self.stream_interval, ) 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 remove_request(self, request_id: str) -> None: """Rollback one previously registered request if it was never submitted.""" req_state = self.request_states.pop(request_id, None) if req_state is None: return external_req_id = getattr(req_state, "external_req_id", None) if external_req_id is not None: self._native_text_metrics_by_request.pop(external_req_id, None) request_ids = self.external_req_ids.get(external_req_id) if request_ids is not None: self.external_req_ids[external_req_id] = [rid for rid in request_ids if rid != request_id] if not self.external_req_ids[external_req_id]: self.external_req_ids.pop(external_req_id, None) parent_req = getattr(req_state, "parent_req", None) if parent_req is not None: self.parent_requests.pop(parent_req.request_id, None) def process_outputs( self, engine_core_outputs: list[EngineCoreOutput], engine_core_timestamp: float | None = None, iteration_stats: IterationStats | None = None, ) -> OutputProcessorOutput: for eco in engine_core_outputs: req_state = self.request_states.get(eco.request_id) if req_state is None or not isinstance(req_state, OmniRequestState): continue if eco.pooling_output is not None and req_state.detokenizer is not None: mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower() req_state.add_multimodal_tensor(eco.pooling_output, mm_type) # Force text path in base processor for multimodal outputs. eco.pooling_output = None return super().process_outputs( engine_core_outputs, engine_core_timestamp=engine_core_timestamp, iteration_stats=iteration_stats, ) def _update_stats_from_output( self, req_state: RequestState, engine_core_output: EngineCoreOutput, engine_core_timestamp: float | None, iteration_stats: IterationStats | None, ): was_prefilling = req_state.is_prefilling native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0 super()._update_stats_from_output( req_state, engine_core_output, engine_core_timestamp, None, ) if iteration_stats is None or engine_core_timestamp is None or native_stats is None: return iteration_stats.update_from_output( engine_core_output, engine_core_timestamp, was_prefilling, native_stats, self.lora_states, req_state.lora_name, ) record = self._native_text_metric_record(req_state.external_req_id) if was_prefilling: record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0) return if previous_last_token_ts > 0: itl_ms = max((engine_core_timestamp - previous_last_token_ts) * 1000.0, 0.0) itls_ms = record.setdefault("vllm_itls_ms", []) itls_ms.append(itl_ms) record["vllm_itl_ms"] = sum(itls_ms) / float(len(itls_ms)) def _update_stats_from_finished( self, req_state: RequestState, finish_reason: FinishReason | None, iteration_stats: IterationStats | None, ): super()._update_stats_from_finished(req_state, finish_reason, None) native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None if iteration_stats is None or native_stats is None or finish_reason is None: return iteration_stats.update_from_finished_request( finish_reason=finish_reason, request_id=req_state.external_req_id, num_prompt_tokens=req_state.prompt_len, max_tokens_param=req_state.max_tokens_param, req_stats=native_stats, num_cached_tokens=req_state.num_cached_tokens, ) if not iteration_stats.finished_requests: return finished_request = iteration_stats.finished_requests[-1] if finished_request.request_id != req_state.external_req_id: return self._native_text_metric_record(req_state.external_req_id)["vllm_tpot_ms"] = ( float(finished_request.mean_time_per_output_token) * 1000.0 )