# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from dataclasses import dataclass, field from typing import Any import numpy as np import torch from PIL import Image from vllm.outputs import CompletionOutput, RequestOutput from vllm.v1.outputs import ModelRunnerOutput from vllm_omni.inputs.data import OmniPromptType @dataclass class OmniConnectorOutput: """Communication results from Model Runner to Scheduler. Carries transfer readiness signals so the Scheduler can make scheduling decisions without ever calling connector.put()/get() directly. Attributes: chunk_ready_req_ids: Request IDs with newly arrived chunks this cycle. chunk_finished_req_ids: Request IDs whose final chunk has arrived. request_metadata: Lightweight scheduling metadata keyed by request ID (e.g. next_stage_prompt_len, code_predictor_codes, left_context_size). Full payloads are owned by the Model Runner's local cache. kv_sent_req_ids: Request IDs whose KV cache was successfully sent. stage_recv_req_ids: Request IDs that received batch stage inputs. has_pending_kv_work: True if the mixin has pending, active, or completed KV transfers that the scheduler should account for. """ chunk_ready_req_ids: set[str] = field(default_factory=set) chunk_finished_req_ids: set[str] = field(default_factory=set) request_metadata: dict[str, dict[str, Any]] = field(default_factory=dict) kv_sent_req_ids: list[str] = field(default_factory=list) stage_recv_req_ids: set[str] = field(default_factory=set) has_pending_kv_work: bool = False @dataclass class OmniModelRunnerOutput(ModelRunnerOutput): """Model runner output for omni models. Extends the base ModelRunnerOutput with support for multimodal outputs that may be produced by non-autoregressive stages. Attributes: multimodal_outputs: Optional per-request list of client-facing multimodal output dicts, indexed by req_index. inter_stage_outputs: Optional per-request list of inter-stage payload dicts for connector transport (``save_async`` / full_payload). Not forwarded to the orchestrator output processor. """ multimodal_outputs: list[dict[str, object]] | None = None inter_stage_outputs: list[dict[str, Any] | None] | None = None # IDs of requests whose KV cache has been extracted from GPU/NPU to CPU. # The Scheduler can safely free the block tables for these requests. kv_extracted_req_ids: list[str] | None = None omni_connector_output: OmniConnectorOutput | None = None # True when sampled_token_ids has already been materialized on the host. # MRv2 consumers must treat the value as immutable and must not perform # another device-to-host conversion or rebuild it from sampler tensors. sampled_token_ids_materialized: bool = False @classmethod def with_kv_conn_output_only(cls, kv_connector_output: Any) -> "OmniModelRunnerOutput": return cls( req_ids=[], req_id_to_index={}, sampled_token_ids=[], logprobs=None, prompt_logprobs_dict={}, pooler_output=[], kv_connector_output=kv_connector_output, ) # RequestOutput attributes that hold generation content. When wrapping a # stage output, these are copied onto the OmniRequestOutput itself so that # the object IS the RequestOutput instead of holding one. _REQUEST_OUTPUT_CONTENT_ATTRS = ( "request_id", "prompt", "prompt_token_ids", "prompt_logprobs", "prompt_token_id_logprobs", "outputs", "finished", "lora_request", "encoder_prompt", "encoder_prompt_token_ids", "num_cached_tokens", "num_cache_creation_tokens", "kv_transfer_params", "ec_transfer_params", ) # Omni-specific content copied when the wrapped stage output is itself an # OmniRequestOutput (e.g. a diffusion stage inside a multi-stage pipeline). _OMNI_CONTENT_ATTRS = ( "images", "latents", "trajectory_latents", "trajectory_timesteps", "trajectory_log_probs", "trajectory_decoded", "_multimodal_output", "_custom_output", ) @dataclass class OmniRequestOutput(RequestOutput): """Unified request output for both pipeline stages and diffusion models. Extends vLLM's ``RequestOutput`` so that omni outputs can flow directly through vLLM serving codepaths (which expect ``prompt_token_ids``, ``outputs``, etc. as real attributes). The inherited fields store the LLM generation content; omni-specific fields store pipeline/diffusion extras. Note: ``RequestOutput`` is a plain class (not a dataclass), so all of its attributes are redeclared below as dataclass fields with defaults — the dataclass-generated ``__init__`` replaces ``RequestOutput.__init__`` and must set them itself. This class handles outputs from: 1. Multi-stage LLM pipelines (with stage_id, final_output_type, and the inherited RequestOutput fields carrying the stage's generation content) 2. Diffusion models (with images, prompt, metrics) Attributes: request_id: Unique identifier for this request finished: Whether generation is complete stage_id: Identifier of the stage that produced this output (pipeline mode) replica_id: Identifier of the stage replica that produced this output final_output_type: Type of output ("text", "image", "audio", "latents") images: List of generated PIL images (diffusion mode) prompt: The prompt used for generation latents: Optional tensor of latent representations (diffusion mode) metrics: Generation metrics. A plain dict for omni outputs; may carry vLLM's request stats object when copied from a raw RequestOutput. """ # --- Inherited RequestOutput attributes (redeclared as fields) --- request_id: str = "" prompt: OmniPromptType | None = None prompt_token_ids: list[int] | None = None prompt_logprobs: Any = None prompt_token_id_logprobs: np.ndarray | None = None outputs: list[CompletionOutput] = field(default_factory=list) finished: bool = True metrics: Any = field(default_factory=dict) lora_request: Any = None encoder_prompt: str | None = None encoder_prompt_token_ids: list[int] | None = None num_cached_tokens: int | None = None num_cache_creation_tokens: int | None = None kv_transfer_params: dict[str, Any] | None = None ec_transfer_params: dict[str, Any] | None = None # --- Pipeline stage fields --- stage_id: int | None = None replica_id: int | None = None final_output_type: str = "text" # --- Diffusion model fields --- images: list[Image.Image] = field(default_factory=list) latents: torch.Tensor | None = None trajectory_latents: torch.Tensor | None = None trajectory_timesteps: torch.Tensor | None = None trajectory_log_probs: torch.Tensor | None = None trajectory_decoded: list | None = None _multimodal_output: dict[str, Any] = field(default_factory=dict) _custom_output: dict[str, Any] = field(default_factory=dict) # profiling data stage_durations: dict[str, float] = field(default_factory=dict) # memory usage info peak_memory_mb: float = 0.0 # error handling error: str | None = None error_status_code: int | None = None error_type: str | None = None def _copy_content_from(self, source: RequestOutput) -> None: """Copy generation content from a stage output into this object. *source* should be a vLLM ``RequestOutput`` or a subclass such as ``OmniRequestOutput``. Control fields (``stage_id``, ``final_output_type``, ``stage_durations``, ``peak_memory_mb``, ``error``, ...) are intentionally not copied. """ # RequestOutput attributes — guaranteed on any RequestOutput. # Only copy attributes that actually exist on the source, so that # duck-typed mocks can omit optional fields without overwriting the # dataclass defaults (e.g. ``outputs=[]``, ``finished=True``). for name in _REQUEST_OUTPUT_CONTENT_ATTRS: if hasattr(source, name): setattr(self, name, getattr(source, name)) # Omni-specific fields — only present when the source is itself an # OmniRequestOutput (e.g. a diffusion stage inside a pipeline). if isinstance(source, OmniRequestOutput): for name in _OMNI_CONTENT_ATTRS: if hasattr(source, name): setattr(self, name, getattr(source, name)) # Propagate multimodal_output from duck-typed sources (e.g. test mocks # that set multimodal_output directly on a non-OmniRequestOutput). if not isinstance(source, OmniRequestOutput): src_mm = getattr(source, "multimodal_output", None) if src_mm: self._multimodal_output = src_mm @classmethod def from_stage_output(cls, source: RequestOutput, **kwargs: Any) -> "OmniRequestOutput": """Create an OmniRequestOutput from a stage's raw output. Copies generation content (``outputs``, ``prompt``, ``prompt_token_ids``, ``finished``, ``images``, ``latents``, etc.) from *source* onto the returned object. *source* may be a vLLM ``RequestOutput``, another ``OmniRequestOutput`` (which inherits from ``RequestOutput``). This is the **preferred** way to construct an ``OmniRequestOutput`` that wraps a stage result. Args: source: The stage output whose content is copied onto the new object. **kwargs: Passed through to the dataclass constructor (``request_id``, ``stage_id``, ``final_output_type``, ``metrics``, ``stage_durations``, ``peak_memory_mb``, ``finished``, etc.). Typed as ``Any`` because the exact set of valid keys is the dataclass field list, which is validated by ``cls(**kwargs)`` at call time. Returns: A new ``OmniRequestOutput`` with the stage's content flattened onto it. """ obj = cls(**kwargs) obj._copy_content_from(source) return obj @classmethod def from_error( cls, request_id: str, error_message: str, *, status_code: int | None = None, error_type: str | None = None, ) -> "OmniRequestOutput": """Create a terminal error output. Args: request_id: Request identifier error_message: Human-readable error description Returns: OmniRequestOutput with ``finished=True`` and the ``error`` field set. """ return cls( request_id=request_id, finished=True, error=error_message, error_status_code=status_code, error_type=error_type, ) @classmethod def from_diffusion( cls, request_id: str, images: list[Image.Image], prompt: OmniPromptType | None = None, metrics: dict[str, Any] | None = None, latents: torch.Tensor | None = None, trajectory_latents: torch.Tensor | None = None, trajectory_timesteps: torch.Tensor | None = None, trajectory_log_probs: torch.Tensor | None = None, trajectory_decoded: list | None = None, multimodal_output: dict[str, Any] | None = None, custom_output: dict[str, Any] | None = None, final_output_type: str = "image", stage_durations: dict[str, float] | None = None, peak_memory_mb: float = 0.0, finished: bool = True, ) -> "OmniRequestOutput": """Create output from diffusion model. Args: request_id: Request identifier images: Generated images prompt: The prompt used metrics: Generation metrics latents: Optional latent tensors trajectory_latents: Optional stacked trajectory latent tensors trajectory_timesteps: Optional stacked trajectory timestep tensors trajectory_log_probs: Optional stacked trajectory log-probability tensors trajectory_decoded: Optional list of decoded trajectory images multimodal_output: Optional multimodal output dict custom_output: Optional custom output dict (e.g. prompt embeds) stage_durations: Optional stage durations (execution time of each stage) dict peak_memory_mb: Peak memory usage in MB Returns: OmniRequestOutput configured for diffusion mode """ return cls( request_id=request_id, final_output_type=final_output_type, images=images, prompt=prompt, latents=latents, trajectory_latents=trajectory_latents, trajectory_timesteps=trajectory_timesteps, trajectory_log_probs=trajectory_log_probs, trajectory_decoded=trajectory_decoded, metrics=metrics or {}, _multimodal_output=multimodal_output or {}, _custom_output=custom_output or {}, stage_durations=stage_durations or {}, peak_memory_mb=peak_memory_mb, finished=finished, ) @property def multimodal_output(self) -> Any: """Return the multimodal output payload. Checks completion outputs first (where multimodal_output is attached by AR stages), then the local _multimodal_output field. Returns either a MultimodalPayload (Phase 3+) or a plain dict (legacy). """ for output in self.outputs: if mm := getattr(output, "multimodal_output", None): return mm return self._multimodal_output @property def custom_output(self) -> dict[str, Any]: """Return custom output data from diffusion pipelines.""" return self._custom_output @custom_output.setter def custom_output(self, value: dict[str, Any]) -> None: self._custom_output = value @property def num_images(self) -> int: """Return the number of generated images.""" return len(self.images) @property def is_diffusion_output(self) -> bool: """Check if this is a diffusion model output.""" return len(self.images) > 0 or self.final_output_type == "image" @property def is_pipeline_output(self) -> bool: """Check if this is a pipeline stage output.""" return self.stage_id is not None def to_dict(self) -> dict[str, Any]: """Convert to dictionary for JSON serialization.""" result = { "request_id": self.request_id, "finished": self.finished, "final_output_type": self.final_output_type, } if self.is_diffusion_output: result.update( { "num_images": self.num_images, "prompt": self.prompt, "metrics": self.metrics, } ) if self.is_pipeline_output: result.update( { "stage_id": self.stage_id, } ) return result def __repr__(self) -> str: """Custom repr to properly show image count instead of image objects.""" # For images, show count instead of full list images_repr = f"[{len(self.images)} PIL Images]" if self.images else "[]" # Build repr string parts = [ f"request_id={self.request_id!r}", f"finished={self.finished}", f"stage_id={self.stage_id}", f"final_output_type={self.final_output_type!r}", f"prompt_token_ids={self.prompt_token_ids}", f"outputs={self.outputs}", f"images={images_repr}", f"prompt={self.prompt!r}", f"latents={self.latents}", f"metrics={self.metrics}", f"multimodal_output={self._multimodal_output}", f"custom_output={self._custom_output}", f"stage_durations={self.stage_durations}", f"peak_memory_mb={self.peak_memory_mb}", ] return f"OmniRequestOutput({', '.join(parts)})"