tzhouam
bf69cd94bb
Update Dockerfile for vllm wheel version and enhance input processing in Omni models. Added decoding for byte messages in async requests, included new metrics in benchmarks, and improved pause state handling in the scheduler. Refined input and output processing methods for better multimodal support and error handling. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
477 lines
20 KiB
Python
477 lines
20 KiB
Python
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
|
|
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
|
|
|
|
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)
|
|
# Omni-specific: multimodal output accumulation
|
|
self.mm_type: str | None = None
|
|
self.mm_accumulated: dict[str, Any] | None = None
|
|
|
|
def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None:
|
|
if payload is None:
|
|
return
|
|
try:
|
|
if mm_type:
|
|
self.mm_type = (mm_type or "").lower()
|
|
|
|
# Normalize incoming payload to dict on CPU
|
|
def _to_cpu(x):
|
|
if isinstance(x, torch.Tensor):
|
|
try:
|
|
return x.detach().to("cpu", non_blocking=True).contiguous()
|
|
except Exception:
|
|
return x
|
|
return x
|
|
|
|
if isinstance(payload, dict):
|
|
incoming: dict[str, Any] = {}
|
|
target_key = self.mm_type or "hidden"
|
|
|
|
# Iterate directly without unnecessary dict copy
|
|
for k, v in payload.items():
|
|
# Optional remap: if producer used "model_outputs" or "hidden", rename to mm_type
|
|
if k == "model_outputs":
|
|
k = target_key
|
|
elif k == "hidden" and target_key != "hidden":
|
|
k = target_key
|
|
|
|
if isinstance(v, dict):
|
|
incoming[k] = {str(sk): _to_cpu(sv) for sk, sv in v.items()}
|
|
else:
|
|
incoming[k] = _to_cpu(v)
|
|
else:
|
|
key = self.mm_type or "hidden"
|
|
incoming = {key: _to_cpu(payload)}
|
|
|
|
if self.mm_accumulated is None:
|
|
self.mm_accumulated = incoming
|
|
else:
|
|
# Merge keys; accumulate tensors in lists for deferred concatenation
|
|
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):
|
|
# Use list accumulation to avoid O(n²) repeated concatenation
|
|
self.mm_accumulated[k] = [existing, v]
|
|
elif isinstance(v, torch.Tensor) and isinstance(existing, list):
|
|
# Append to existing list
|
|
existing.append(v)
|
|
elif isinstance(v, dict) and isinstance(existing, dict):
|
|
# Merge nested dicts with list accumulation for tensors
|
|
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."""
|
|
if self.mm_accumulated is None:
|
|
return
|
|
try:
|
|
for k, v in self.mm_accumulated.items():
|
|
if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
|
|
try:
|
|
if k == "audio":
|
|
# When the audio tensor shape is inconsistent, torch.cat will fail.
|
|
# We need to use torch.cat in -1 dimension.
|
|
continue
|
|
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")
|
|
|
|
# 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,
|
|
routed_experts: np.ndarray | 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,
|
|
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
|
|
|
|
# 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 self.output_kind == RequestOutputKind.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, routed_experts)
|
|
|
|
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,
|
|
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, routed_experts)
|
|
try:
|
|
if self.mm_accumulated is not None:
|
|
# Attach accumulated multimodal dict on the completion output
|
|
if not hasattr(base_output, "multimodal_output"):
|
|
setattr(base_output, "multimodal_output", {})
|
|
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)
|
|
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: 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", "latents"). Used to route outputs
|
|
to appropriate processors. If None, output type is inferred.
|
|
"""
|
|
super().__init__(
|
|
tokenizer=tokenizer,
|
|
log_stats=log_stats,
|
|
stream_interval=stream_interval,
|
|
tracing_enabled=tracing_enabled,
|
|
)
|
|
self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
|
|
self._reqid_to_mm_type: dict[str, str] = {}
|
|
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: 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,
|
|
)
|
|
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,
|
|
engine_core_outputs: list[EngineCoreOutput],
|
|
engine_core_timestamp: float | None = None,
|
|
iteration_stats: IterationStats | None = None,
|
|
) -> OutputProcessorOutput:
|
|
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)
|
|
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:
|
|
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
|
|
|
|
return super().process_outputs(
|
|
engine_core_outputs,
|
|
engine_core_timestamp=engine_core_timestamp,
|
|
iteration_stats=iteration_stats,
|
|
)
|
|
|
|
# ---- 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, ...]) -> torch.Tensor | None:
|
|
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
|