From 846167b23766f6294970798bd18e7bbb8f14105c Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 6 Mar 2026 14:45:53 +0800 Subject: [PATCH] [1/N][Refactor] Clean up dead code in output processor (#1579) Signed-off-by: gcanlin --- vllm_omni/engine/output_processor.py | 158 ++---------------- .../entrypoints/openai/serving_speech.py | 2 +- 2 files changed, 19 insertions(+), 141 deletions(-) diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 697622f12..a37e84097 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -1,4 +1,3 @@ -from collections.abc import Callable from typing import Any import numpy as np @@ -60,9 +59,11 @@ class OmniRequestState(RequestState): 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 + # Normalize producer keys to the modality name. + # AR runners produce {"hidden": ...} and generation + # runners produce {"model_outputs": ...}; remap both + # to the semantic modality key (e.g. "audio", "latent"). if k == "model_outputs": k = target_key elif k == "hidden" and target_key != "hidden": @@ -247,16 +248,17 @@ class OmniRequestState(RequestState): class MultimodalOutputProcessor(VLLMOutputProcessor): - """Handles multimodal output processing by normalizing EngineCoreOutput - before delegating to the base vLLM OutputProcessor. + """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. - 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(). + 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__( @@ -273,27 +275,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): 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. + (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) - 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, @@ -346,20 +333,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): 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(), - ) + 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 @@ -368,105 +348,3 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): 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 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 9ff7bc767..8c292cc7a 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -418,7 +418,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): mm = getattr(ro, "multimodal_output", None) if ro else None if not mm: return None, None - key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None) + key = "audio" if "audio" in mm else None return mm, key def _build_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: