Chendi.XueandHongsheng Liu
4e41b7bc23
Signed-off-by: Chendi Xue <chendi.xue@intel.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
425 lines
15 KiB
Python
425 lines
15 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import torch
|
|
from PIL import Image
|
|
from vllm.outputs import 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
|
|
|
|
|
|
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 dictionary mapping modality names to
|
|
output tensors (e.g., {"image": tensor, "audio": tensor})
|
|
"""
|
|
|
|
multimodal_outputs: dict[str, torch.Tensor] | 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
|
|
|
|
|
|
@dataclass
|
|
class OmniRequestOutput:
|
|
"""Unified request output for both pipeline stages and diffusion models.
|
|
|
|
This class handles outputs from:
|
|
1. Multi-stage LLM pipelines (with stage_id, final_output_type, request_output)
|
|
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)
|
|
final_output_type: Type of output ("text", "image", "audio", "latents")
|
|
request_output: The underlying RequestOutput from the stage (pipeline mode)
|
|
images: List of generated PIL images (diffusion mode)
|
|
prompt: The prompt used for generation (diffusion mode)
|
|
latents: Optional tensor of latent representations (diffusion mode)
|
|
metrics: Optional dictionary of generation metrics
|
|
"""
|
|
|
|
request_id: str = ""
|
|
finished: bool = True
|
|
|
|
# Pipeline stage fields
|
|
stage_id: int | None = None
|
|
final_output_type: str = "text"
|
|
request_output: RequestOutput | None = None
|
|
|
|
# Diffusion model fields
|
|
images: list[Image.Image] = field(default_factory=list)
|
|
prompt: OmniPromptType | 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
|
|
metrics: dict[str, Any] = field(default_factory=dict)
|
|
_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
|
|
|
|
@classmethod
|
|
def from_error(
|
|
cls,
|
|
request_id: str,
|
|
error_message: str,
|
|
) -> "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,
|
|
)
|
|
|
|
@classmethod
|
|
def from_pipeline(
|
|
cls,
|
|
stage_id: int,
|
|
final_output_type: str,
|
|
request_output: RequestOutput,
|
|
) -> "OmniRequestOutput":
|
|
"""Create output from pipeline stage.
|
|
|
|
Args:
|
|
stage_id: Stage identifier
|
|
final_output_type: Type of output
|
|
request_output: The stage's output
|
|
|
|
Returns:
|
|
OmniRequestOutput configured for pipeline mode
|
|
"""
|
|
return cls(
|
|
request_id=getattr(request_output, "request_id", ""),
|
|
stage_id=stage_id,
|
|
final_output_type=final_output_type,
|
|
request_output=request_output,
|
|
finished=True,
|
|
)
|
|
|
|
@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,
|
|
) -> "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=True,
|
|
)
|
|
|
|
@property
|
|
def multimodal_output(self) -> dict[str, Any]:
|
|
"""Return multimodal output from the underlying request output or local field.
|
|
|
|
For pipeline outputs, this checks completion outputs first, then request_output.
|
|
For diffusion outputs, this returns the local _multimodal_output field.
|
|
"""
|
|
if self.request_output is None:
|
|
return self._multimodal_output
|
|
|
|
# Check completion outputs first (where multimodal_output is attached)
|
|
for output in getattr(self.request_output, "outputs", []):
|
|
if mm := getattr(output, "multimodal_output", None):
|
|
return mm
|
|
if mm := getattr(self.request_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.
|
|
|
|
For diffusion outputs, returns the local _custom_output field.
|
|
For pipeline outputs with an inner OmniRequestOutput, forwards
|
|
the custom_output from the inner request output.
|
|
"""
|
|
if self.request_output is not None:
|
|
if isinstance(self.request_output, OmniRequestOutput):
|
|
return self.request_output._custom_output
|
|
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)
|
|
|
|
# Pass-through properties keep vLLM serving codepaths compatible with
|
|
# OmniRequestOutput for pipeline outputs (Issue #345).
|
|
@property
|
|
def prompt_token_ids(self) -> list[int] | None:
|
|
"""Return prompt token IDs from the underlying request output.
|
|
|
|
This property is required for compatibility with vLLM's streaming
|
|
chat completion generator which checks res.prompt_token_ids.
|
|
"""
|
|
if self.request_output is not None:
|
|
return getattr(self.request_output, "prompt_token_ids", None)
|
|
return None
|
|
|
|
@property
|
|
def outputs(self) -> list[Any]:
|
|
"""Return outputs from the underlying request output.
|
|
|
|
This property is required for compatibility with vLLM's streaming
|
|
and non-streaming chat completion generators.
|
|
"""
|
|
if self.request_output is not None:
|
|
return getattr(self.request_output, "outputs", [])
|
|
return []
|
|
|
|
@property
|
|
def encoder_prompt_token_ids(self) -> list[int] | None:
|
|
"""Return encoder prompt token IDs from the underlying request output."""
|
|
if self.request_output is not None:
|
|
return getattr(self.request_output, "encoder_prompt_token_ids", None)
|
|
return None
|
|
|
|
@property
|
|
def prompt_logprobs(self) -> Any:
|
|
"""Return prompt logprobs from the underlying request output."""
|
|
if self.request_output is not None:
|
|
return getattr(self.request_output, "prompt_logprobs", None)
|
|
return None
|
|
|
|
@property
|
|
def num_cached_tokens(self) -> int | None:
|
|
"""Return number of cached tokens from the underlying request output."""
|
|
if self.request_output is not None:
|
|
return getattr(self.request_output, "num_cached_tokens", None)
|
|
return None
|
|
|
|
@property
|
|
def kv_transfer_params(self) -> Any:
|
|
"""Return KV transfer params from the underlying request output."""
|
|
if self.request_output is not None:
|
|
return getattr(self.request_output, "kv_transfer_params", None)
|
|
return None
|
|
|
|
@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 and self.request_output is not None
|
|
|
|
def unwrap(self) -> "OmniRequestOutput":
|
|
"""Unwrap nested OmniRequestOutput to get the innermost result.
|
|
|
|
This helper handles the common pattern where pipeline outputs may wrap
|
|
other OmniRequestOutput instances. It recursively unwraps until it reaches
|
|
the final output with actual content (images, text, etc.).
|
|
|
|
Returns:
|
|
The innermost OmniRequestOutput containing the actual generation results.
|
|
|
|
Example:
|
|
```python
|
|
result = omni.generate(...)
|
|
output = OmniRequestOutput.unwrap_result(result)
|
|
if output.images:
|
|
# Access images directly
|
|
video_frames = output.images
|
|
```
|
|
"""
|
|
current = self
|
|
# Unwrap nested pipeline outputs
|
|
while current.is_pipeline_output and current.request_output is not None:
|
|
if isinstance(current.request_output, OmniRequestOutput):
|
|
current = current.request_output
|
|
else:
|
|
break
|
|
return current
|
|
|
|
@staticmethod
|
|
def unwrap_result(result: Any) -> "OmniRequestOutput":
|
|
"""Unwrap result from omni.generate() to get the final OmniRequestOutput.
|
|
|
|
This static helper handles the full unwrapping pattern including:
|
|
1. Extracting from list if needed
|
|
2. Type validation
|
|
3. Recursive unwrapping of nested pipeline outputs
|
|
|
|
Args:
|
|
result: The result from omni.generate() - may be a list or OmniRequestOutput
|
|
|
|
Returns:
|
|
The innermost OmniRequestOutput with actual content
|
|
|
|
Raises:
|
|
ValueError: If result is not an OmniRequestOutput or list containing one
|
|
|
|
Example:
|
|
```python
|
|
result = omni.generate(...)
|
|
output = OmniRequestOutput.unwrap_result(result)
|
|
# output is guaranteed to be the final OmniRequestOutput
|
|
```
|
|
"""
|
|
# Handle list wrapper
|
|
if isinstance(result, list):
|
|
if not result:
|
|
raise ValueError("Result list is empty")
|
|
result = result[0]
|
|
|
|
# Validate type
|
|
if not isinstance(result, OmniRequestOutput):
|
|
raise ValueError(f"Expected OmniRequestOutput, got {type(result)}")
|
|
|
|
# Unwrap nested outputs
|
|
return result.unwrap()
|
|
|
|
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"request_output={self.request_output}",
|
|
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)})"
|