510 lines
21 KiB
Python
510 lines
21 KiB
Python
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
|
|
)
|