From 70ef548448430d1452c2eed9482faeb755fc0841 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Tue, 3 Feb 2026 12:24:45 +0800 Subject: [PATCH] Dev/rebase v0.15.0 (#1159) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Taichang Zhou Signed-off-by: tzhouam Signed-off-by: princepride Co-authored-by: 汪志鹏 --- docker/Dockerfile.ci | 19 +- docker/Dockerfile.rocm | 2 +- .../offline_inference/qwen2_5_omni/end2end.py | 2 +- .../offline_inference/qwen3_omni/end2end.py | 2 +- .../offline_inference/qwen3_tts/end2end.py | 4 +- .../qwen3_tts/openai_speech_client.py | 4 + .../offline_inference/test_qwen2_5_omni.py | 2 +- .../e2e/offline_inference/test_qwen3_omni.py | 2 +- .../e2e/online_serving/test_image_gen_edit.py | 2 + vllm_omni/config/model.py | 11 ++ vllm_omni/core/sched/omni_ar_scheduler.py | 67 +++---- .../core/sched/omni_generation_scheduler.py | 64 ++++++- vllm_omni/core/sched/output.py | 13 +- vllm_omni/engine/arg_utils.py | 5 +- vllm_omni/engine/input_processor.py | 98 +++++----- vllm_omni/engine/output_processor.py | 156 +++++----------- vllm_omni/entrypoints/async_omni.py | 12 +- vllm_omni/entrypoints/async_omni_llm.py | 1 - vllm_omni/entrypoints/chat_utils.py | 41 +++-- vllm_omni/entrypoints/omni_llm.py | 4 +- vllm_omni/entrypoints/omni_stage.py | 57 +++--- vllm_omni/entrypoints/openai/api_server.py | 119 ++++++++----- .../openai/protocol/chat_completion.py | 4 +- vllm_omni/entrypoints/openai/serving_chat.py | 89 ++++++---- .../entrypoints/openai/serving_speech.py | 17 +- vllm_omni/inputs/data.py | 3 + vllm_omni/inputs/parse.py | 2 +- vllm_omni/inputs/preprocess.py | 21 ++- .../layers/rotary_embedding/__init__.py | 6 + .../layers/{ => rotary_embedding}/mrope.py | 168 ++---------------- .../models/qwen2_5_omni/qwen2_5_omni.py | 4 +- .../qwen2_5_omni/qwen2_5_omni_thinker.py | 1 - .../qwen3_omni/qwen3_omni_moe_thinker.py | 6 +- .../qwen3_tts/configuration_qwen3_tts.py | 11 +- .../models/qwen3_tts/modeling_qwen3_tts.py | 25 ++- vllm_omni/model_executor/models/vision.py | 23 --- vllm_omni/outputs.py | 17 +- vllm_omni/patch.py | 4 +- vllm_omni/request.py | 25 ++- vllm_omni/worker/gpu_ar_model_runner.py | 26 ++- vllm_omni/worker/gpu_ar_worker.py | 5 + .../worker/gpu_generation_model_runner.py | 38 +++- vllm_omni/worker/gpu_generation_worker.py | 5 + vllm_omni/worker/gpu_model_runner.py | 15 +- 44 files changed, 648 insertions(+), 554 deletions(-) create mode 100644 vllm_omni/model_executor/layers/rotary_embedding/__init__.py rename vllm_omni/model_executor/layers/{ => rotary_embedding}/mrope.py (82%) delete mode 100644 vllm_omni/model_executor/models/vision.py diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index f0f585fba..4cbcb4d08 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,20 +1,21 @@ ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.14.0 +ARG VLLM_BASE_TAG=v0.15.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} - COPY . . - # Install system dependencies RUN apt-get update && \ - apt-get install -y ffmpeg && \ - apt-get clean && \ - rm -rf /var/lib/apt/lists/* + apt-get install -y ffmpeg && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +# Install vLLM wheel for the specific commit +# URL format: https://wheels.vllm.ai// +# Current vLLM commit: f176443446f659dbab5315e056e605d8984fd976 (v0.15.0) +RUN uv pip install --system --upgrade --force-reinstall vllm --extra-index-url https://wheels.vllm.ai/f176443446f659dbab5315e056e605d8984fd976 # Install vllm-omni into the same uv-managed Python environment used by the base image. -RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" - +RUN uv pip install --system --no-cache-dir ".[dev]" RUN ln -sf /usr/bin/python3 /usr/bin/python - ENTRYPOINT [] diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 0f5923342..bbb75a196 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.14.0 +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.15.0 FROM ${BASE_IMAGE} AS final ARG COMMON_WORKDIR=/app diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 22f2161c2..63b08202a 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -410,7 +410,7 @@ def main(args): elif stage_outputs.final_output_type == "audio": for output in stage_outputs.request_output: request_id = output.request_id - audio_tensor = output.multimodal_output["audio"] + audio_tensor = output.outputs[0].multimodal_output["audio"] output_wav = os.path.join(output_dir, f"output_{request_id}.wav") sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000) print(f"Request ID: {request_id}, Saved audio to {output_wav}") diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 3cd891820..0b484c507 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -418,7 +418,7 @@ def main(args): elif stage_outputs.final_output_type == "audio": for output in stage_outputs.request_output: request_id = output.request_id - audio_tensor = output.multimodal_output["audio"] + audio_tensor = output.outputs[0].multimodal_output["audio"] output_wav = os.path.join(output_dir, f"output_{request_id}.wav") # Convert to numpy array and ensure correct format diff --git a/examples/offline_inference/qwen3_tts/end2end.py b/examples/offline_inference/qwen3_tts/end2end.py index 7bdd0ae7d..18bb9dedc 100644 --- a/examples/offline_inference/qwen3_tts/end2end.py +++ b/examples/offline_inference/qwen3_tts/end2end.py @@ -244,9 +244,9 @@ def main(args): for stage_outputs in omni_generator: for output in stage_outputs.request_output: request_id = output.request_id - audio_tensor = output.multimodal_output["audio"] + audio_tensor = output.outputs[0].multimodal_output["audio"] output_wav = os.path.join(output_dir, f"output_{request_id}.wav") - audio_samplerate = output.multimodal_output["sr"].item() + audio_samplerate = output.outputs[0].multimodal_output["sr"].item() # Convert to numpy array and ensure correct format audio_numpy = audio_tensor.float().detach().cpu().numpy() diff --git a/examples/online_serving/qwen3_tts/openai_speech_client.py b/examples/online_serving/qwen3_tts/openai_speech_client.py index 71ba5b70a..4d7046094 100644 --- a/examples/online_serving/qwen3_tts/openai_speech_client.py +++ b/examples/online_serving/qwen3_tts/openai_speech_client.py @@ -111,6 +111,10 @@ def run_tts_generation(args) -> None: print(response.text) return + if response.content.decode("utf-8").startswith('{"error"'): + print(f"Error: {response.content.decode('utf-8')}") + return + # Save audio response output_path = args.output or "tts_output.wav" with open(output_path, "wb") as f: diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni.py b/tests/e2e/offline_inference/test_qwen2_5_omni.py index 3537ea127..af9d793c1 100644 --- a/tests/e2e/offline_inference/test_qwen2_5_omni.py +++ b/tests/e2e/offline_inference/test_qwen2_5_omni.py @@ -88,7 +88,7 @@ def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: t assert len(audio_output.request_output) > 0 # Verify audio tensor exists and has content - audio_tensor = audio_output.request_output[0].multimodal_output["audio"] + audio_tensor = audio_output.request_output[0].outputs[0].multimodal_output["audio"] assert audio_tensor is not None assert audio_tensor.numel() > 0 diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index 470cdb73b..2c5c66fe3 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -75,6 +75,6 @@ def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None: assert len(audio_output.request_output) > 0 # Verify audio tensor exists and has content - audio_tensor = audio_output.request_output[0].multimodal_output["audio"] + audio_tensor = audio_output.request_output[0].outputs[0].multimodal_output["audio"] assert audio_tensor is not None assert audio_tensor.numel() > 0 diff --git a/tests/e2e/online_serving/test_image_gen_edit.py b/tests/e2e/online_serving/test_image_gen_edit.py index e796fe959..8db0d50fb 100644 --- a/tests/e2e/online_serving/test_image_gen_edit.py +++ b/tests/e2e/online_serving/test_image_gen_edit.py @@ -23,6 +23,8 @@ from vllm.assets.image import ImageAsset from vllm.utils.network_utils import get_open_port os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +# Increase timeout for downloading assets from S3 (default 5s is too short for CI) +os.environ.setdefault("VLLM_IMAGE_FETCH_TIMEOUT", "60") models = ["Qwen/Qwen-Image-Edit-2509"] test_params = models diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 751a487c3..dad07dc22 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -110,6 +110,7 @@ class OmniModelConfig(ModelConfig): mm_processor_cache_gb: float | None, mm_processor_cache_type: MMCacheType | None, mm_shm_cache_max_object_size_mb: int | None, + mm_encoder_only: bool | None, mm_encoder_tp_mode: MMEncoderTPMode | None, mm_encoder_attn_backend: AttentionBackendEnum | str | None, interleave_mm_strings: bool | None, @@ -181,6 +182,15 @@ class OmniModelConfig(ModelConfig): ) self.model_arch_config = self.get_model_arch_config() + if self.convert == "mm_encoder_only": + logger.warning_once( + "`--convert mm_encoder_only` is deprecated and " + "will be removed in v0.15. " + "Please use --mm-encoder-only` instead." + ) + mm_encoder_only = True + self.convert = "none" + architectures = self.architectures registry = self.registry is_generative_model = registry.is_text_generation_model(architectures, self) @@ -262,6 +272,7 @@ class OmniModelConfig(ModelConfig): mm_processor_cache_gb=mm_processor_cache_gb, mm_processor_cache_type=mm_processor_cache_type, mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb, + mm_encoder_only=mm_encoder_only, mm_encoder_tp_mode=mm_encoder_tp_mode, mm_encoder_attn_backend=mm_encoder_attn_backend, interleave_mm_strings=interleave_mm_strings, diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 647eb4867..2af96e718 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -6,20 +6,20 @@ from dataclasses import asdict, dataclass from time import time from typing import Any -import numpy as np from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.distributed.kv_events import KVEventBatch from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.logger import init_logger from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler -from vllm.v1.core.sched.utils import check_stop, remove_all +from vllm.v1.core.sched.utils import remove_all from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs from vllm.v1.metrics.perf import PerfStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm_omni.core.sched.output import OmniSchedulerOutput from vllm_omni.distributed.omni_connectors.adapter import get_chunk, put_chunk from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec @@ -185,13 +185,19 @@ class OmniARScheduler(VLLMScheduler): get_chunk(self.omni_connector, scheduler_output) # Add information about requests needing KV cache transfer - scheduler_output.finished_requests_needing_kv_transfer = self.get_finished_requests_needing_kv_transfer() + finished_reqs = self.get_finished_requests_needing_kv_transfer() except Exception: # If anything goes wrong, leave the original output unchanged init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData") - scheduler_output.finished_requests_needing_kv_transfer = {} + finished_reqs = {} - return scheduler_output + # Wrap in omni scheduler output to carry transfer metadata. + base_fields = SchedulerOutput.__dataclass_fields__.keys() + base_data = {name: getattr(scheduler_output, name) for name in base_fields} + return OmniSchedulerOutput( + **base_data, + finished_requests_needing_kv_transfer=finished_reqs, + ) def update_from_output( self, @@ -278,47 +284,32 @@ class OmniARScheduler(VLLMScheduler): pooler_output = pooler_outputs[req_index] if pooler_outputs else None kv_transfer_params = None status_before_stop = request.status + finish_reason = None + routed_experts = None # Check for stop and update request status. if new_token_ids: new_token_ids, stopped = self._update_request_with_output(request, new_token_ids) + elif request.pooling_params and pooler_output is not None: + # Pooling stops as soon as there is output. + request.status = RequestStatus.FINISHED_STOPPED + stopped = True - # If criteria returns True, it means we must STOP the request (e.g. legacy behavior, not used now) - # If criteria returns False, it might have triggered a background transfer - # (e.g. prefill finished / special token) but continues decoding + # If criteria returns True, it means we must STOP the request. + # If criteria returns False, it might have triggered a background + # transfer (e.g. prefill finished / special token) but continues decoding. if not stopped and self._process_kv_transfer_trigger(request, new_token_ids): stopped = True - # Stop checking for pooler models. - pooler_output = None - if pooler_outputs: - pooler_output = pooler_outputs[req_index] - if request.output_token_ids: - stopped = check_stop(request, self.max_model_len) - routed_experts = None if stopped: - # [Omni] Handle routed experts if enabled - if self.vllm_config.model_config.enable_return_routed_experts: - kv_blocks = self.kv_cache_manager.get_blocks(request.request_id) - block_ids = kv_blocks.get_block_ids()[0] - num_tokens = request.num_tokens - 1 + routed_experts = self._get_routed_experts(request) - # compute slot mapping - block_ids_array = np.array(block_ids, dtype=np.int32) - num_blocks = len(block_ids) - block_size = self.block_size - - # generate block offsets - block_offsets = np.arange(0, block_size) - - # compute slot mapping: slot = block_id * block_size + offset - slot_mapping = ( - block_offsets.reshape((1, block_size)) + block_ids_array.reshape((num_blocks, 1)) * block_size - ).flatten()[:num_tokens] - - routed_experts = self.routed_experts_reader.get_routed_experts(indices=slot_mapping) - - kv_transfer_params = self._free_request(request) + # Capture finish_reason BEFORE _handle_stopped_request, which may + # reset the status to WAITING for streaming requests that continue. + finish_reason = request.get_finished_reason() + finished = self._handle_stopped_request(request) + if finished: + kv_transfer_params = self._free_request(request) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) else: @@ -345,13 +336,13 @@ class OmniARScheduler(VLLMScheduler): # Get prompt logprobs for this request. prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) - if new_token_ids or pooler_output is not None or kv_transfer_params: + if new_token_ids or pooler_output is not None or kv_transfer_params or stopped: # Add EngineCoreOutput for this Request. outputs[request.client_index].append( EngineCoreOutput( request_id=req_id, new_token_ids=new_token_ids, - finish_reason=request.get_finished_reason(), + finish_reason=finish_reason, new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, pooling_output=pooler_output, diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 5974adb52..50e7b8dbc 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -1,7 +1,9 @@ import time from collections import defaultdict +from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.distributed.kv_events import KVEventBatch +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.logger import init_logger from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput @@ -9,6 +11,7 @@ from vllm.v1.core.sched.request_queue import create_request_queue from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.core.sched.utils import remove_all from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs +from vllm.v1.metrics.perf import PerfStats from vllm.v1.request import Request, RequestStatus from vllm.v1.spec_decode.metrics import SpecDecodingStats @@ -56,10 +59,21 @@ class OmniGenerationScheduler(VLLMScheduler): # Temporary queue: preserve waiting order, do not disturb non-diffusion requests skipped_waiting_requests = create_request_queue(self.policy) req_index = 0 + # OMNI: Track requests that are already finished (e.g., marked by connector) + # These should be removed from running and not scheduled + already_finished_reqs: set[Request] = set() while req_index < len(self.running) and token_budget > 0: request = self.running[req_index] if self.omni_connector is not None: get_chunk_for_generation(self.omni_connector, request) + + # OMNI: Skip requests that are already finished or not in self.requests + # This can happen when connector marks request as finished + if request.status == RequestStatus.FINISHED_STOPPED or request.request_id not in self.requests: + already_finished_reqs.add(request) + req_index += 1 + continue + num_computed_tokens = request.num_computed_tokens required_tokens = max(len(request.prompt_token_ids) - num_computed_tokens, 1) num_new_tokens = min(required_tokens, token_budget) @@ -83,12 +97,24 @@ class OmniGenerationScheduler(VLLMScheduler): scheduled_running_reqs.append(request) req_index += 1 + # OMNI: Remove already finished requests from running queue + if already_finished_reqs: + self.running = remove_all(self.running, already_finished_reqs) + # Fast path selection and scheduling (treat all as diffusion requests, # independent of pooling_params) while self.waiting and token_budget > 0 and len(self.running) < self.max_num_running_reqs: request = self.waiting.peek_request() if self.omni_connector is not None: get_chunk_for_generation(self.omni_connector, request) + + # OMNI: Skip requests that are already finished or not in self.requests + # This can happen when connector marks request as finished + if request.status == RequestStatus.FINISHED_STOPPED or request.request_id not in self.requests: + # Pop the finished request from waiting queue and don't schedule it + self.waiting.pop_request() + continue + # Uniformly treat as diffusion. A feature flag can be added later # via config or request tag. @@ -257,9 +283,16 @@ class OmniGenerationScheduler(VLLMScheduler): num_nans_in_logits = model_runner_output.num_nans_in_logits kv_connector_output = model_runner_output.kv_connector_output + cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats + perf_stats: PerfStats | None = None + if self.perf_metrics and self.perf_metrics.is_enabled(): + perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output) + outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None - kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) # Merge connector-side stats (align with v0.14.0) if kv_connector_stats and self.connector: kv_stats = self.connector.get_kv_connector_stats() @@ -308,12 +341,14 @@ class OmniGenerationScheduler(VLLMScheduler): num_accepted_tokens=num_accepted, ) + stopped = False new_logprobs = None new_token_ids = generated_token_ids kv_transfer_params = None - pooler_output = None - if pooler_outputs: - pooler_output = pooler_outputs[req_index] + pooler_output = pooler_outputs[req_index] if pooler_outputs else None + status_before_stop = request.status + finish_reason = None + routed_experts = None # Diffusion request: completes in one step; mark finished and free resources if request.status == RequestStatus.FINISHED_STOPPED or ( @@ -323,8 +358,18 @@ class OmniGenerationScheduler(VLLMScheduler): # Optional: set a stop_reason for front-end clarity # (does not affect protocol) request.stop_reason = request.stop_reason # or "generation_done" - kv_transfer_params = self._free_request(request) - stopped_running_reqs.add(request) + stopped = True + + if stopped: + routed_experts = self._get_routed_experts(request) + finish_reason = request.get_finished_reason() + finished = self._handle_stopped_request(request) + if finished: + kv_transfer_params = self._free_request(request) + if status_before_stop == RequestStatus.RUNNING: + stopped_running_reqs.add(request) + else: + stopped_preempted_reqs.add(request) # Extract sample logprobs if needed. if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs: @@ -344,13 +389,13 @@ class OmniGenerationScheduler(VLLMScheduler): # Get prompt logprobs for this request. prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) - if new_token_ids or pooler_output is not None or kv_transfer_params: + if new_token_ids or pooler_output is not None or kv_transfer_params or stopped: # Add EngineCoreOutput for this Request. outputs[request.client_index].append( EngineCoreOutput( request_id=req_id, new_token_ids=new_token_ids, - finish_reason=request.get_finished_reason(), + finish_reason=finish_reason, new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, pooling_output=pooler_output, @@ -359,6 +404,7 @@ class OmniGenerationScheduler(VLLMScheduler): kv_transfer_params=kv_transfer_params, trace_headers=request.trace_headers, num_cached_tokens=request.num_cached_tokens, + routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, ) ) @@ -406,7 +452,7 @@ class OmniGenerationScheduler(VLLMScheduler): engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set) finished_req_ids.clear() - if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None: + if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats)) is not None: # Return stats to only one of the front-ends. if (eco := next(iter(engine_core_outputs.values()), None)) is None: # We must return the stats even if there are no request diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index f4acb034d..933ae7629 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -1,6 +1,6 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field -from vllm.v1.core.sched.output import CachedRequestData, NewRequestData +from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput from vllm.v1.request import Request from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload @@ -66,4 +66,11 @@ class OmniCachedRequestData(CachedRequestData): prompt_token_ids: Mapping from request ID to list of prompt token IDs """ - prompt_token_ids: dict[int, list[int]] + prompt_token_ids: dict[str, list[int]] + + +@dataclass +class OmniSchedulerOutput(SchedulerOutput): + """Scheduler output with omni-specific transfer metadata.""" + + finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 21bb761d2..d550cf555 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -1,4 +1,5 @@ from dataclasses import dataclass, field +from typing import Any from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeTextConfig from vllm.engine.arg_utils import EngineArgs @@ -70,7 +71,7 @@ class OmniEngineArgs(EngineArgs): engine_output_type: str | None = None hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None - stage_connector_spec: dict[str, any] = field(default_factory=dict) + stage_connector_spec: dict[str, Any] = field(default_factory=dict) async_chunk: bool = False omni_kv_config: dict | None = None @@ -172,7 +173,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): engine_output_type: str | None = None hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None - stage_connector_spec: dict[str, any] = field(default_factory=dict) + stage_connector_spec: dict[str, Any] = field(default_factory=dict) async_chunk: bool = False omni_kv_config: dict | None = None diff --git a/vllm_omni/engine/input_processor.py b/vllm_omni/engine/input_processor.py index 91e12f6a6..a1e467a88 100644 --- a/vllm_omni/engine/input_processor.py +++ b/vllm_omni/engine/input_processor.py @@ -1,7 +1,9 @@ +import os import time from collections.abc import Mapping from typing import Any, cast +import numpy as np import torch from vllm.config import VllmConfig from vllm.inputs import ProcessorInputs, PromptType @@ -9,12 +11,13 @@ from vllm.inputs.parse import split_enc_dec_inputs from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalUUIDDict +from vllm.multimodal.processing.context import set_request_id from vllm.multimodal.utils import argsort_mm_positions from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams -from vllm.tokenizers import TokenizerLike from vllm.utils import length_from_prompt_token_ids_or_embeds +from vllm.utils.torch_utils import set_default_torch_num_threads from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.engine import ( @@ -38,7 +41,6 @@ class OmniInputProcessor(InputProcessor): Args: vllm_config: Global vLLM configuration - tokenizer: Tokenizer instance for text processing mm_registry: Multi-modal registry for processing multimodal inputs """ @@ -75,13 +77,12 @@ class OmniInputProcessor(InputProcessor): def __init__( self, vllm_config: VllmConfig, - tokenizer: TokenizerLike, mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, ): - super().__init__(vllm_config, tokenizer, mm_registry) + super().__init__(vllm_config, mm_registry) self.input_preprocessor = OmniInputPreprocessor( self.model_config, - self.tokenizer, + vllm_config.observability_config, mm_registry, mm_processor_cache=self.mm_processor_cache, ) @@ -97,7 +98,8 @@ class OmniInputProcessor(InputProcessor): trace_headers: Mapping[str, str] | None = None, priority: int = 0, data_parallel_rank: int | None = None, - ) -> tuple[str | None, OmniEngineCoreRequest]: + resumable: bool = False, + ) -> OmniEngineCoreRequest: """Process input prompt into an engine core request. Converts a prompt (text, tokens, or multimodal) into an @@ -130,9 +132,12 @@ class OmniInputProcessor(InputProcessor): self._validate_lora(lora_request) self._validate_params(params) - data_parallel_size = self.vllm_config.parallel_config.data_parallel_size - if data_parallel_rank is not None and not (0 <= data_parallel_rank < data_parallel_size): - raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {data_parallel_size}).") + parallel_config = self.vllm_config.parallel_config + dp_size = parallel_config.data_parallel_size + dp_local_size = parallel_config.data_parallel_size_local + num_ranks = dp_local_size if parallel_config.local_engines_only else dp_size + if data_parallel_rank is not None and not (0 <= data_parallel_rank < num_ranks): + raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {num_ranks}).") if arrival_time is None: arrival_time = time.time() @@ -154,7 +159,7 @@ class OmniInputProcessor(InputProcessor): else: # Otherwise, use user-provided uuids as multimodal hash overrides # if provided. - self._validate_multi_modal_uuids(prompt) + self._validate_mm_uuids(prompt) if isinstance(prompt, dict): mm_uuids = cast(MultiModalUUIDDict | None, prompt.get("multi_modal_uuids")) else: @@ -164,11 +169,19 @@ class OmniInputProcessor(InputProcessor): # 1. Tokenize text prompt, with LoRA request if one exists. # 2. For multimodal models with a merged preprocessor, preprocess # multimodal data and expand prompt token ids accordingly. - processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess( - prompt, - tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, - ) + num_threads = int(os.environ.get("OMP_NUM_THREADS", "1")) + if "OMP_NUM_THREADS" not in os.environ: + logger.debug_once( + "OMP_NUM_THREADS is not set; defaulting Torch threads to %d for input preprocessing.", + num_threads, + ) + + with set_request_id(request_id), set_default_torch_num_threads(num_threads): + processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess( + prompt, + tokenization_kwargs=tokenization_kwargs, + mm_uuids=mm_uuids, + ) current_platform.validate_request( prompt=prompt, @@ -181,13 +194,13 @@ class OmniInputProcessor(InputProcessor): encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs) self._validate_model_inputs(encoder_inputs, decoder_inputs) - # Mypy does not always properly infer the types of some elements of - # discriminated unions of TypedDicts, because of how it handles - # inheritance of TypedDict. If we explicitly extract the items we want - # we can avoid type errors from using `dict.get` later in the method. - _prompt_str: str | None = None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt") - prompt_token_ids = decoder_inputs["prompt_token_ids"] if decoder_inputs["type"] != "embeds" else None - prompt_embeds = decoder_inputs["prompt_embeds"] if decoder_inputs["type"] == "embeds" else None + # Normalize decoder prompt access across TypedDict variants. + if decoder_inputs["type"] == "embeds": + prompt_token_ids = None + prompt_embeds = decoder_inputs["prompt_embeds"] + else: + prompt_token_ids = decoder_inputs["prompt_token_ids"] + prompt_embeds = decoder_inputs.get("prompt_embeds") sampling_params = None pooling_params = None @@ -219,35 +232,26 @@ class OmniInputProcessor(InputProcessor): mm_features = [] for modality, idx in sorted_mm_idxs: + base_mm_hash = decoder_mm_hashes[modality][idx] mm_features.append( MultiModalFeatureSpec( data=decoder_mm_inputs[modality][idx], modality=modality, - identifier=decoder_mm_hashes[modality][idx], + identifier=self._get_mm_identifier(base_mm_hash, lora_request), mm_position=decoder_mm_positions[modality][idx], + mm_hash=base_mm_hash, ) ) - # Serialize prompt_embeds and additional_information if provided - # (direct-transfer path) - prompt_embeds_payload: PromptEmbedsPayload | None = None + # Compatibility: decode serialized prompt embeds if provided. + if isinstance(prompt_embeds, PromptEmbedsPayload): + prompt_embeds = self._decode_prompt_embeds(prompt_embeds) + additional_information_payload: AdditionalInformationPayload | None = None - pe: torch.Tensor | None = decoder_inputs.get("prompt_embeds") # type: ignore[operator] - if pe is not None: - if pe.ndim != 2: - raise ValueError("prompt_embeds must be of shape (seq_len, hidden_size)") - # Move to CPU and ensure contiguous memory for stable serialization - pe_cpu = pe.detach().to("cpu").contiguous() - seq_len, hidden_size = pe_cpu.shape - dtype_str = self._dtype_to_name(pe_cpu.dtype) - data_bytes = pe_cpu.numpy().tobytes() - prompt_embeds_payload = PromptEmbedsPayload( - data=data_bytes, - shape=[int(seq_len), int(hidden_size)], - dtype=dtype_str, - ) - raw_info: dict[str, Any] | None = decoder_inputs.get("additional_information") # type: ignore[operator] - if raw_info is not None: + raw_info: dict[str, Any] | AdditionalInformationPayload | None = decoder_inputs.get("additional_information") + if isinstance(raw_info, AdditionalInformationPayload): + additional_information_payload = raw_info + elif raw_info is not None: entries: dict[str, AdditionalInformationEntry] = {} for key, value in raw_info.items(): if isinstance(value, torch.Tensor): @@ -279,6 +283,14 @@ class OmniInputProcessor(InputProcessor): priority=priority, data_parallel_rank=data_parallel_rank, trace_headers=trace_headers, - prompt_embeds=prompt_embeds_payload, + prompt_embeds=prompt_embeds, additional_information=additional_information_payload, + resumable=resumable, ) + + @staticmethod + def _decode_prompt_embeds(payload: PromptEmbedsPayload) -> torch.Tensor: + dtype = getattr(np, payload.dtype) + arr = np.frombuffer(payload.data, dtype=dtype) + arr = arr.reshape(payload.shape) + return torch.from_numpy(arr) diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 65a69e0d4..e39dc9fec 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -2,6 +2,7 @@ from ast import Dict 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 @@ -9,7 +10,11 @@ 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.output_processor import ( + OutputProcessorOutput, + RequestOutputCollector, + RequestState, +) from vllm.v1.engine.parallel_sampling import ParentRequest from vllm.v1.metrics.stats import IterationStats @@ -140,6 +145,7 @@ class OmniRequestState(RequestState): 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. @@ -158,12 +164,27 @@ class OmniRequestState(RequestState): 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 @@ -185,7 +206,7 @@ class OmniRequestState(RequestState): self.sent_tokens_offset = len(self.detokenizer.output_token_ids) request_id = self.request_id - output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) + output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts) if self.parent_req is None: outputs = [output] @@ -201,9 +222,10 @@ class OmniRequestState(RequestState): 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) + 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 @@ -251,7 +273,6 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): super().__init__(tokenizer=tokenizer, log_stats=log_stats) self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {} self._reqid_to_mm_type: dict[str, str] = {} - self.request_states: dict[str, OmniRequestState] = {} self.engine_core_output_type = engine_core_output_type def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None: @@ -291,8 +312,10 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): ValueError: If the request ID is already registered """ request_id = request.request_id - if request_id in self.request_states: - raise ValueError(f"Request id {request_id} already running.") + 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, @@ -304,9 +327,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): 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, @@ -314,121 +340,27 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): engine_core_timestamp: float | None = None, iteration_stats: IterationStats | None = None, ) -> OutputProcessorOutput: - """Process engine core outputs into request outputs. - - Converts EngineCoreOutput objects into RequestOutput objects, - handling multimodal outputs by routing them through appropriate - processors and accumulating tensors in request states. - - Args: - engine_core_outputs: List of engine core outputs to process - engine_core_timestamp: Optional timestamp for the outputs - iteration_stats: Optional iteration statistics - - Returns: - OutputProcessorOutput containing processed request outputs - and list of request IDs to abort - """ 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) - - # Build RequestOutputs without delegating to base, so we can keep ids - request_outputs: list[Any] = [] - reqs_to_abort: list[str] = [] - for eco in engine_core_outputs: - req_id = eco.request_id - req_state = self.request_states.get(req_id) - if req_state is None: + req_state = self.request_states.get(eco.request_id) + if req_state is None or not isinstance(req_state, OmniRequestState): continue - - # 1) Stats - self._update_stats_from_output(req_state, eco, engine_core_timestamp, iteration_stats) - - new_token_ids = eco.new_token_ids - pooling_output = eco.pooling_output - finish_reason = eco.finish_reason - stop_reason = eco.stop_reason - kv_transfer_params = eco.kv_transfer_params - req_state.num_cached_tokens = eco.num_cached_tokens - req_state.is_prefilling = False - - # 2) Detokenize and logprobs when text path - assert req_state.detokenizer is not None - assert req_state.logprobs_processor is not None - stop_string = req_state.detokenizer.update(new_token_ids, finish_reason == FinishReason.STOP) - if stop_string: - finish_reason = FinishReason.STOP - stop_reason = stop_string - req_state.logprobs_processor.update_from_output(eco) - - # 2.5) Accumulate multimodal tensors in RequestState - try: - mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower() - if pooling_output is not None and isinstance(req_state, OmniRequestState): - req_state.add_multimodal_tensor(pooling_output, mm_type) - except Exception: - logger.debug( - "Failed to accumulate multimodal tensor for request %s", - req_id, - exc_info=True, + 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 - # 3) Create RequestOutput objects, forcing combined mode to keep ids - pooling_for_make = pooling_output - if pooling_output is not None and new_token_ids: - # Do not consume pooling path now; keep ids and attach mm later - pooling_for_make = None - - # Consolidate any accumulated tensor lists before creating output - if isinstance(req_state, OmniRequestState) and finish_reason is not None: - req_state._consolidate_multimodal_tensors() - - ro = req_state.make_request_output( - new_token_ids, - pooling_for_make, - finish_reason, - stop_reason, - kv_transfer_params, - ) - if ro: - # Attach accumulated multimodal payload if any - try: - if isinstance(req_state, OmniRequestState) and req_state.mm_accumulated is not None: - if not hasattr(ro, "multimodal_output"): - setattr(ro, "multimodal_output", {}) - ro.multimodal_output = req_state.mm_accumulated - except Exception: - logger.exception("Error attaching multimodal payload in process_outputs") - if req_state.queue is not None: - req_state.queue.put(ro) - else: - request_outputs.append(ro) - - # 4) Free completed - if finish_reason is not None: - self.request_states.pop(req_id) - parent_req = req_state.parent_req - if parent_req and not parent_req.child_requests: - self.parent_requests.pop(parent_req.request_id, None) - if not self.request_states: - self._requests_drained.set() - if not eco.finished: - reqs_to_abort.append(req_id) - self._update_stats_from_finished(req_state, finish_reason, iteration_stats) - if self.tracer: - self.do_tracing(eco, req_state, iteration_stats) - # Cleanup per-request mm state - if isinstance(req_state, OmniRequestState): - req_state.mm_accumulated = None - req_state.mm_type = None - - return OutputProcessorOutput( - request_outputs=request_outputs, - reqs_to_abort=reqs_to_abort, + return super().process_outputs( + engine_core_outputs, + engine_core_timestamp=engine_core_timestamp, + iteration_stats=iteration_stats, ) # ---- routing helpers ---- diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index ab829beaa..144c73f3a 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -196,11 +196,10 @@ class AsyncOmni(OmniBase): if stage.vllm_config is not None and stage.tokenizer is not None: try: vllm_config = stage.vllm_config - tokenizer = stage.tokenizer # Initialize input_processor + # OMNI: OmniInputProcessor creates tokenizer internally from vllm_config self.input_processor = OmniInputProcessor( vllm_config=vllm_config, - tokenizer=tokenizer, ) # Initialize model_config self.model_config = vllm_config.model_config @@ -675,6 +674,15 @@ class AsyncOmni(OmniBase): return stage.is_tracing_enabled return False + @property + def renderer(self): + """Return the renderer from input_processor if available. + + OMNI: Required by upstream OpenAIServingModels.__init__ which + accesses engine_client.renderer. + """ + return self.input_processor.renderer + async def do_log_stats(self) -> None: pass diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py index 287f12b9e..b557c07dd 100644 --- a/vllm_omni/entrypoints/async_omni_llm.py +++ b/vllm_omni/entrypoints/async_omni_llm.py @@ -116,7 +116,6 @@ class AsyncOmniLLM(AsyncLLM): # InputProcessor (converts Inputs --> EngineCoreRequests). self.input_processor = OmniInputProcessor( vllm_config=vllm_config, - tokenizer=tokenizer, mm_registry=mm_registry, ) diff --git a/vllm_omni/entrypoints/chat_utils.py b/vllm_omni/entrypoints/chat_utils.py index 0fdef5edb..ca15305c5 100644 --- a/vllm_omni/entrypoints/chat_utils.py +++ b/vllm_omni/entrypoints/chat_utils.py @@ -11,11 +11,11 @@ from vllm.entrypoints.chat_utils import ( BaseMultiModalItemTracker, ChatCompletionContentPartParam, ChatCompletionMessageParam, + ChatTemplateContentFormat, ConversationMessage, MultiModalDataDict, MultiModalUUIDDict, _AssistantParser, - _ChatTemplateContentFormat, _ContentPart, _get_full_multimodal_text_prompt, _parse_chat_message_content_part, @@ -39,17 +39,27 @@ class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser): self._mm_processor_kwargs = mm_processor_kwargs def parse_video(self, video_url: str | None, uuid: str | None = None) -> None: - video = self._connector.fetch_video_async(video_url=video_url) if video_url else None - - placeholder = self._tracker.add("video", video, uuid) + # OMNI: Follow upstream async pattern - create coroutine that resolves to (data, uuid) + coro = self._video_with_uuid_async(video_url, uuid) + placeholder = self._tracker.add("video", coro) self._add_placeholder("video", placeholder) # Extract audio from video if use_audio_in_video is True if video_url and self._mm_processor_kwargs and self._mm_processor_kwargs.get("use_audio_in_video", False): - audio_coro = self._extract_audio_from_video_async(video_url) - audio_placeholder = self._tracker.add("audio", audio_coro, uuid) + audio_coro = self._audio_from_video_with_uuid_async(video_url, uuid) + audio_placeholder = self._tracker.add("audio", audio_coro) self._add_placeholder("audio", audio_placeholder) + async def _video_with_uuid_async(self, video_url: str | None, uuid: str | None): + """Fetch video and return (video, uuid) tuple.""" + video = await self._connector.fetch_video_async(video_url=video_url) if video_url else None + return video, uuid + + async def _audio_from_video_with_uuid_async(self, video_url: str, uuid: str | None): + """Extract audio from video and return (audio, uuid) tuple.""" + audio = await self._extract_audio_from_video_async(video_url) + return audio, uuid + async def _extract_audio_from_video_async(self, video_url: str) -> tuple[np.ndarray, int | float]: """ Extract audio from video URL using librosa. @@ -128,13 +138,21 @@ class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser): def parse_chat_messages_futures( messages: list[ChatCompletionMessageParam], model_config: ModelConfig, - content_format: _ChatTemplateContentFormat, + content_format: ChatTemplateContentFormat, mm_processor_kwargs: dict[str, Any] | None = None, ) -> tuple[ list[ConversationMessage], - Awaitable[MultiModalDataDict | None], - MultiModalUUIDDict | None, + Awaitable[tuple[MultiModalDataDict | None, MultiModalUUIDDict | None]], ]: + """Parse chat messages and return conversation with multimodal data future. + + OMNI: Updated to use upstream vLLM v0.15.0 API where resolve_items() + returns both mm_data and mm_uuids together as a tuple. + + Returns: + Tuple of (conversation, mm_future) where mm_future resolves to + (mm_data, mm_uuids) when awaited. + """ conversation: list[ConversationMessage] = [] mm_tracker = OmniAsyncMultiModalItemTracker(model_config) @@ -155,13 +173,14 @@ def parse_chat_messages_futures( _postprocess_messages(conversation) - return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids() + # OMNI: Use upstream resolve_items() which returns (mm_data, mm_uuids) tuple + return conversation, mm_tracker.resolve_items() def _parse_chat_message_content( message: ChatCompletionMessageParam, mm_tracker: BaseMultiModalItemTracker, - content_format: _ChatTemplateContentFormat, + content_format: ChatTemplateContentFormat, interleave_strings: bool, mm_processor_kwargs: dict[str, Any] | None = None, ) -> list[ConversationMessage]: diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 6bb8ac366..0356ea231 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -161,9 +161,7 @@ class OmniLLM(LLM): log_stats=self.llm_engine.log_stats, engine_core_output_type=engine_args.engine_output_type, ) - self.llm_engine.input_processor = OmniInputProcessor( - vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer - ) + self.llm_engine.input_processor = OmniInputProcessor(vllm_config=self.llm_engine.vllm_config) self.engine_class = type(self.llm_engine) self.request_counter = Counter() diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index fc01c147c..133e9ec84 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -670,12 +670,22 @@ def _stage_worker( break lock_files = acquired_lock_fds + + # Set FD_CLOEXEC on all lock file descriptors to prevent child processes + # (e.g., EngineCore) from inheriting them, which would cause deadlock + for lock_fd in acquired_lock_fds: + try: + flags = fcntl.fcntl(lock_fd, fcntl.F_GETFD) + fcntl.fcntl(lock_fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC) + except (OSError, ValueError): + pass except Exception as e: logger.debug( "[Stage-%s] Failed to set up sequential initialization lock: %s", stage_id, e, ) + # Init engine based on stage_type logger.debug("[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys())) if engine_args.get("async_chunk", False): @@ -686,30 +696,29 @@ def _stage_worker( break engine_args["stage_connector_spec"] = stage_connector_spec engine_args["stage_id"] = stage_id - try: - if stage_type == "diffusion": - engine_args.pop("model_stage", None) - engine_args.pop("model", None) - stage_engine = OmniDiffusion( - model=model, - stage_id=stage_id, - engine_input_source=stage_payload.get("engine_input_source", []), - **engine_args, - ) - else: - # Default to LLM engine - stage_engine = OmniLLM(model=model, **engine_args) - finally: - # Release all locks by closing file descriptors - # Locks are automatically released when file descriptors are closed - # or when process dies - for lock_fd in lock_files: - try: - fcntl.flock(lock_fd, fcntl.LOCK_UN) - _os.close(lock_fd) - logger.debug("Released initialization lock (fd=%s)", lock_fd) - except (OSError, ValueError): - pass + if stage_type == "diffusion": + engine_args.pop("model_stage", None) + engine_args.pop("model", None) + stage_engine = OmniDiffusion( + model=model, + stage_id=stage_id, + engine_input_source=stage_payload.get("engine_input_source", []), + **engine_args, + ) + else: + # Default to LLM engine + stage_engine = OmniLLM(model=model, **engine_args) + + # Release all locks AFTER engine initialization completes + for lock_fd in lock_files: + try: + fcntl.flock(lock_fd, fcntl.LOCK_UN) + _os.close(lock_fd) + logger.debug("Released initialization lock (fd=%s)", lock_fd) + except (OSError, ValueError): + pass + lock_files = [] # Clear after release + logger.debug("Engine initialized") # Initialize OmniConnectors if configured connectors: dict[tuple[str, str], OmniConnectorBase] | None = {} diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 2e057901a..fb52c7e46 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -19,40 +19,41 @@ from typing import Annotated, Any, cast import httpx import vllm.envs as envs -from fastapi import Depends, File, Form, HTTPException, Request, UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile from fastapi.responses import JSONResponse, StreamingResponse from PIL import Image from starlette.datastructures import State from starlette.routing import Route from vllm import SamplingParams from vllm.engine.protocol import EngineClient -from vllm.entrypoints.anthropic.serving_messages import AnthropicServingMessages +from vllm.entrypoints.anthropic.serving import AnthropicServingMessages +from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launcher import serve_http from vllm.entrypoints.logger import RequestLogger -from vllm.entrypoints.openai.api_server import ( - base, - build_app, - load_log_config, - router, - setup_server, -) -from vllm.entrypoints.openai.orca_metrics import metrics_header -from vllm.entrypoints.openai.protocol import ( +from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer +from vllm.entrypoints.openai.api_server import base, load_log_config +from vllm.entrypoints.openai.api_server import build_app as build_openai_app +from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server +from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, - ErrorResponse, - ModelCard, - ModelList, - ModelPermission, ) # yapf conflicts with isort for this block # yapf: disable # yapf: enable -from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion -from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels -from vllm.entrypoints.openai.serving_responses import OpenAIServingResponses -from vllm.entrypoints.openai.serving_transcription import ( +from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion +from vllm.entrypoints.openai.engine.protocol import ( + ErrorResponse, + ModelCard, + ModelList, + ModelPermission, +) +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.openai.orca_metrics import metrics_header +from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses +from vllm.entrypoints.openai.translations.serving import ( OpenAIServingTranscription, OpenAIServingTranslation, ) @@ -63,10 +64,8 @@ from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling from vllm.entrypoints.pooling.score.serving import ServingScores from vllm.entrypoints.serve.disagg.serving import ServingTokens from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization -from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer from vllm.entrypoints.utils import ( load_aware_call, - process_chat_template, process_lora_modules, with_cancellation, ) @@ -93,29 +92,41 @@ from vllm_omni.lora.request import LoRARequest from vllm_omni.lora.utils import stable_lora_int_id logger = init_logger(__name__) +router = APIRouter() + + +def _remove_route_from_router( + router: APIRouter, + path: str, + methods: set[str] | None = None, +) -> None: + methods_set = {method.upper() for method in methods} if methods else None + for route in list(router.routes): + if getattr(route, "path", None) != path: + continue + if methods_set is not None: + route_methods = {method.upper() for method in (getattr(route, "methods", None) or set())} + if not (route_methods & methods_set): + continue + router.routes.remove(route) + ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format" -def _remove_route_from_router(router_obj, path: str, methods: set[str] | None = None): - """Remove a route from the router by path and optionally by methods. +def _remove_route_from_app(app, path: str, methods: set[str] | None = None): + """Remove a route from the app by path and optionally by methods. - This is needed because vllm's api_server registers routes when imported, - and we need to override some routes (like /v1/chat/completions) with - omni-specific implementations. + OMNI: used to override upstream /v1/chat/completions with omni behavior. """ routes_to_remove = [] - for route in router_obj.routes: + for route in app.routes: if isinstance(route, Route) and route.path == path: if methods is None or (hasattr(route, "methods") and route.methods & methods): routes_to_remove.append(route) for route in routes_to_remove: - router_obj.routes.remove(route) - - -# Remove vllm's /v1/chat/completions route so we can register our own omni version -_remove_route_from_router(router, "/v1/chat/completions", {"POST"}) + app.routes.remove(route) class _DiffusionServingModels: @@ -162,7 +173,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None: # Add process-specific prefix to stdout and stderr. decorate_logs("APIServer") - listen_address, sock = setup_server(args) + listen_address, sock = setup_openai_server(args) # Unified use of omni_run_server_worker, AsyncOmni automatically handles LLM and Diffusion models await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs) @@ -179,7 +190,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) # Load logging config for uvicorn if specified - log_config = load_log_config(args.log_config_file) + log_config = load_log_config(getattr(args, "log_config_file", None)) if log_config is not None: uvicorn_kwargs["log_config"] = log_config @@ -187,7 +198,19 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, args, client_config=client_config, ) as engine_client: - app = build_app(args) + supported_tasks: tuple[str, ...] + if hasattr(engine_client, "get_supported_tasks"): + supported_tasks = tuple(await engine_client.get_supported_tasks()) + else: + supported_tasks = ("generate",) + if not supported_tasks: + supported_tasks = ("generate",) + + app = build_openai_app(args) + # OMNI: Remove upstream routes that we override with omni-specific handlers + _remove_route_from_app(app, "/v1/chat/completions", {"POST"}) + _remove_route_from_app(app, "/v1/models", {"GET"}) # Remove upstream /v1/models to use omni's handler + app.include_router(router) await omni_init_app_state(engine_client, app.state, args) @@ -374,6 +397,8 @@ async def omni_init_app_state( state.vllm_config = None state.diffusion_engine = engine_client state.openai_serving_models = _DiffusionServingModels(base_model_paths) + # OMNI: tokenization endpoints are not supported in pure diffusion mode. + state.openai_serving_tokenization = None # Use for_diffusion method to create chat handler state.openai_serving_chat = OmniOpenAIServingChat.for_diffusion( @@ -401,11 +426,7 @@ async def omni_init_app_state( supported_tasks = set(await engine_client.get_supported_tasks()) logger.info("Supported tasks: %s", supported_tasks) - resolved_chat_template = await process_chat_template( - args.chat_template, - engine_client, - vllm_config.model_config if vllm_config is not None else None, - ) + resolved_chat_template = load_chat_template(args.chat_template) if args.tool_server == "demo": tool_server: ToolServer | None = DemoToolServer() @@ -444,10 +465,10 @@ async def omni_init_app_state( tokenizer = await engine_client.get_tokenizer() if tokenizer is not None: # Initialize input_processor + # OMNI: OmniInputProcessor creates tokenizer internally from vllm_config if not hasattr(engine_client, "input_processor") or engine_client.input_processor is None: engine_client.input_processor = OmniInputProcessor( vllm_config=vllm_config, - tokenizer=tokenizer, ) logger.info("Initialized input_processor for AsyncOmni") @@ -691,7 +712,13 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re metrics_header_format = raw_request.headers.get(ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, "") handler = Omnichat(raw_request) if handler is None: - return base(raw_request).create_error_response(message="The model does not support Chat Completions API") + base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None) + if base_server is None: + raise HTTPException( + status_code=HTTPStatus.NOT_FOUND.value, + detail="The model does not support Chat Completions API", + ) + return base_server.create_error_response(message="The model does not support Chat Completions API") try: generator = await handler.create_chat_completion(request, raw_request) except Exception as e: @@ -760,7 +787,13 @@ _remove_route_from_router(router, "/v1/audio/speech", {"POST"}) async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request): handler = Omnispeech(raw_request) if handler is None: - return base(raw_request).create_error_response(message="The model does not support Speech API") + base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None) + if base_server is None: + raise HTTPException( + status_code=HTTPStatus.NOT_FOUND.value, + detail="The model does not support Speech API", + ) + return base_server.create_error_response(message="The model does not support Speech API") try: return await handler.create_speech(request, raw_request) except Exception as e: diff --git a/vllm_omni/entrypoints/openai/protocol/chat_completion.py b/vllm_omni/entrypoints/openai/protocol/chat_completion.py index d0c83f56f..cdc93f672 100644 --- a/vllm_omni/entrypoints/openai/protocol/chat_completion.py +++ b/vllm_omni/entrypoints/openai/protocol/chat_completion.py @@ -1,6 +1,4 @@ -from vllm.entrypoints.openai.protocol import ( - ChatCompletionStreamResponse, -) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionStreamResponse class OmniChatCompletionStreamResponse(ChatCompletionStreamResponse): diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index ae2fb6552..cfee79157 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -13,6 +13,7 @@ import torch from fastapi import Request from PIL import Image from pydantic import TypeAdapter +from vllm.renderers import RendererLike from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt @@ -28,17 +29,10 @@ from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, ChatTemplateContentFormatOption, ConversationMessage, - apply_hf_chat_template, - apply_mistral_chat_template, get_history_tool_calls_cnt, make_tool_call_id, - resolve_chat_template_content_format, ) -from vllm.entrypoints.openai.parser.harmony_utils import ( - get_streamable_parser_for_assistant, - parse_chat_output, -) -from vllm.entrypoints.openai.protocol import ( +from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionResponse, @@ -46,6 +40,9 @@ from vllm.entrypoints.openai.protocol import ( ChatCompletionResponseStreamChoice, ChatCompletionStreamResponse, ChatMessage, +) +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat +from vllm.entrypoints.openai.engine.protocol import ( DeltaFunctionCall, DeltaMessage, DeltaToolCall, @@ -55,20 +52,29 @@ from vllm.entrypoints.openai.protocol import ( FunctionDefinition, PromptTokenUsageInfo, RequestResponseMetadata, - ResponsesRequest, ToolCall, UsageInfo, ) -from vllm.entrypoints.openai.serving_chat import OpenAIServingChat -from vllm.entrypoints.openai.serving_engine import ( - ChatLikeRequest, - clamp_prompt_logprobs, +from vllm.entrypoints.openai.engine.serving import ChatLikeRequest, clamp_prompt_logprobs +from vllm.entrypoints.openai.parser.harmony_utils import ( + get_streamable_parser_for_assistant, + parse_chat_output, ) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls from vllm.entrypoints.utils import should_include_usage from vllm.inputs.data import PromptType, TokensPrompt from vllm.logger import init_logger from vllm.outputs import RequestOutput +from vllm.renderers.hf import ( + resolve_chat_template_content_format, +) +from vllm.renderers.hf import ( + safe_apply_chat_template as apply_hf_chat_template, +) +from vllm.renderers.mistral import ( + safe_apply_chat_template as apply_mistral_chat_template, +) from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike from vllm.tokenizers.mistral import ( @@ -171,7 +177,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): model_name = self.models.model_name(lora_request) - tokenizer = await self.engine_client.get_tokenizer() + renderer = self.renderer + tokenizer = renderer.get_tokenizer() + if tokenizer is None: + tokenizer = await self.engine_client.get_tokenizer() tool_parser = self.tool_parser @@ -228,7 +237,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): engine_prompts, ) = await self._preprocess_chat( request, - tokenizer, + renderer, request.messages, chat_template=request.chat_template or self.chat_template, chat_template_content_format=self.chat_template_content_format, @@ -320,7 +329,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): async def _preprocess_chat( self, request: ChatLikeRequest | ResponsesRequest, - tokenizer: TokenizerLike, + renderer: RendererLike, messages: list[ChatCompletionMessageParam], chat_template: str | None, chat_template_content_format: ChatTemplateContentFormatOption, @@ -338,15 +347,22 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): list[TokensPrompt], ]: model_config = self.model_config + tokenizer = renderer.get_tokenizer() if renderer is not None else None - resolved_content_format = resolve_chat_template_content_format( - chat_template, - tool_dicts, - chat_template_content_format, - tokenizer, - model_config=model_config, - ) - conversation, mm_data_future, mm_uuids = parse_chat_messages_futures( + if tokenizer is None or isinstance(tokenizer, MistralTokenizer): + resolved_content_format = ( + chat_template_content_format if chat_template_content_format != "auto" else "string" + ) + else: + resolved_content_format = resolve_chat_template_content_format( + chat_template, + tool_dicts, + chat_template_content_format, + tokenizer, + model_config=model_config, + ) + # OMNI: Updated for vLLM v0.15.0 API - resolve_items() returns (mm_data, mm_uuids) tuple + conversation, mm_future = parse_chat_messages_futures( messages, model_config, content_format=resolved_content_format, @@ -380,14 +396,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): **_chat_template_kwargs, ) else: + hf_chat_template_kwargs = dict(_chat_template_kwargs) + hf_chat_template_kwargs.pop("tools", None) + hf_chat_template_kwargs.pop("chat_template", None) request_prompt = apply_hf_chat_template( + model_config=model_config, tokenizer=tokenizer, conversation=conversation, - model_config=model_config, - **_chat_template_kwargs, + tools=tool_dicts, + chat_template=chat_template, + tokenize=False, + **hf_chat_template_kwargs, ) - mm_data = await mm_data_future + # OMNI: Await the combined future to get both mm_data and mm_uuids + mm_data, mm_uuids = await mm_future # tool parsing is done only if a tool_parser has been set and if # tool_choice is not "none" (if tool_choice is "none" but a tool_parser @@ -1653,7 +1676,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ): choices: list[ChatCompletionResponseChoice] = [] final_res = omni_outputs.request_output - audio_data = final_res.multimodal_output.get("audio") + # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput + # Reference: examples/offline_inference/qwen3_omni/end2end.py line 421 + audio_data = final_res.outputs[0].multimodal_output.get("audio") if stream: audio_tensor = audio_data[-1].float().detach().cpu().numpy() else: @@ -1739,9 +1764,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if omni_outputs.images: images = omni_outputs.images # Fall back to request_output for pipeline mode - elif final_res is not None: - if hasattr(final_res, "multimodal_output") and final_res.multimodal_output: - image_data = final_res.multimodal_output.get("image") + # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput + elif final_res is not None and final_res.outputs: + completion_output = final_res.outputs[0] + if hasattr(completion_output, "multimodal_output") and completion_output.multimodal_output: + image_data = completion_output.multimodal_output.get("image") if image_data is not None: if isinstance(image_data, Image.Image): images.append(image_data) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 5e1767700..a8bae9e99 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -3,7 +3,7 @@ from typing import Any from fastapi import Request from fastapi.responses import Response -from vllm.entrypoints.openai.serving_engine import OpenAIServing +from vllm.entrypoints.openai.engine.serving import OpenAIServing from vllm.logger import init_logger from vllm.utils import random_uuid @@ -258,17 +258,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Extract audio from output # Audio can be in final_output.multimodal_output or final_output.request_output.multimodal_output + # Support both "audio" and "model_outputs" keys for compatibility with different models audio_output = None if hasattr(final_output, "multimodal_output") and final_output.multimodal_output: audio_output = final_output.multimodal_output - if (not audio_output or "audio" not in audio_output) and hasattr(final_output, "request_output"): + if not audio_output and hasattr(final_output, "request_output"): if final_output.request_output and hasattr(final_output.request_output, "multimodal_output"): audio_output = final_output.request_output.multimodal_output - if not audio_output or "audio" not in audio_output: + # Check for audio data using either "audio" or "model_outputs" key + audio_key = None + if audio_output: + if "audio" in audio_output: + audio_key = "audio" + elif "model_outputs" in audio_output: + audio_key = "model_outputs" + + if not audio_output or audio_key is None: return self.create_error_response("TTS model did not produce audio output.") - audio_tensor = audio_output["audio"] + audio_tensor = audio_output[audio_key] sample_rate = audio_output.get("sr", 24000) if hasattr(sample_rate, "item"): sample_rate = sample_rate.item() diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 01de61b96..3684744ba 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -104,6 +104,9 @@ class OmniEmbedsPrompt(EmbedsPrompt): # Must ensure that all additional prompt types are inherited from vLLM prompt types # Because TypedDict doesn't support isinstance and are dict. Cannot distinguish them in runtime. # Inheritance ensure that there are only additional fields but not removing fields--safe to route to LLM.generate() +OmniSingletonPrompt: TypeAlias = str | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt +"""Omni singleton prompt type extending vLLM's SingletonPrompt with additional fields.""" + OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt diff --git a/vllm_omni/inputs/parse.py b/vllm_omni/inputs/parse.py index a3700e61f..3fe46ff53 100644 --- a/vllm_omni/inputs/parse.py +++ b/vllm_omni/inputs/parse.py @@ -1,10 +1,10 @@ -from vllm.inputs.data import SingletonPrompt from vllm.inputs.parse import ( ParsedEmbedsPrompt, ParsedSingletonPrompt, ParsedStrPrompt, ParsedTextPrompt, ParsedTokensPrompt, + SingletonPrompt, ) diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 4eeadbff4..d04c66331 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -1,12 +1,17 @@ from typing import Any from typing_extensions import assert_never -from vllm.inputs.data import SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt +from vllm.inputs.data import SingletonInputs, SingletonPrompt from vllm.inputs.preprocess import InputPreprocessor from vllm.logger import init_logger from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict -from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni +from vllm_omni.inputs.data import ( + OmniTextPrompt, + OmniTokenInputs, + OmniTokensPrompt, + token_inputs_omni, +) from vllm_omni.inputs.parse import parse_singleton_prompt_omni logger = init_logger(__name__) @@ -22,7 +27,7 @@ class OmniInputPreprocessor(InputPreprocessor): def _process_text( self, - parsed_content: TextPrompt, + parsed_content: OmniTextPrompt, tokenization_kwargs: dict[str, Any] | None = None, *, mm_uuids: MultiModalUUIDDict | None = None, @@ -62,7 +67,7 @@ class OmniInputPreprocessor(InputPreprocessor): def _process_tokens( self, - parsed_content: TokensPrompt, + parsed_content: OmniTokensPrompt, tokenization_kwargs: dict[str, Any] | None = None, *, mm_uuids: MultiModalUUIDDict | None = None, @@ -80,6 +85,10 @@ class OmniInputPreprocessor(InputPreprocessor): tokenization_kwargs=tokenization_kwargs, mm_uuids=mm_uuids, ) + if prompt_embeds is not None: + inputs["prompt_embeds"] = prompt_embeds + if additional_information is not None: + inputs["additional_information"] = additional_information else: inputs = token_inputs_omni( prompt_token_ids=prompt_token_ids, @@ -114,8 +123,6 @@ class OmniInputPreprocessor(InputPreprocessor): """ parsed = parse_singleton_prompt_omni(prompt) - if parsed["type"] == "embeds": - return self._process_embeds(parsed["content"]) if parsed["type"] == "tokens": return self._process_tokens( parsed["content"], @@ -129,7 +136,7 @@ class OmniInputPreprocessor(InputPreprocessor): ) if parsed["type"] == "str": return self._process_text( - TextPrompt(prompt=parsed["content"]), + OmniTextPrompt(prompt=parsed["content"]), tokenization_kwargs=tokenization_kwargs, mm_uuids=mm_uuids, ) diff --git a/vllm_omni/model_executor/layers/rotary_embedding/__init__.py b/vllm_omni/model_executor/layers/rotary_embedding/__init__.py new file mode 100644 index 000000000..02ebeb194 --- /dev/null +++ b/vllm_omni/model_executor/layers/rotary_embedding/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .mrope import OmniMRotaryEmbedding + +__all__ = ["OmniMRotaryEmbedding"] diff --git a/vllm_omni/model_executor/layers/mrope.py b/vllm_omni/model_executor/layers/rotary_embedding/mrope.py similarity index 82% rename from vllm_omni/model_executor/layers/mrope.py rename to vllm_omni/model_executor/layers/rotary_embedding/mrope.py index 9ca6a36e2..463e55507 100644 --- a/vllm_omni/model_executor/layers/mrope.py +++ b/vllm_omni/model_executor/layers/rotary_embedding/mrope.py @@ -1,136 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Omni-extended MRotaryEmbedding with multimodal position computation methods. + +This module extends the upstream vLLM MRotaryEmbedding with additional methods +for computing input positions for various multimodal scenarios including: +- Image/Video inputs (Qwen2.5-VL style) +- Audio inputs (Qwen2.5-Omni style) +- Audio-in-video interleaved inputs +- GLM4V style inputs +""" + import itertools -import numpy as np import torch from transformers import PretrainedConfig from vllm.logger import init_logger -from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding +from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding logger = init_logger(__name__) -def _apply_rotary_emb( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - is_neox_style: bool, -) -> torch.Tensor: +class OmniMRotaryEmbedding(MRotaryEmbedding): + """Omni-extended MRotaryEmbedding with multimodal position computation. + + Extends the upstream MRotaryEmbedding with additional class methods for + computing input positions for various multimodal scenarios. """ - Args: - x: [num_tokens, num_heads, head_size] - cos: [num_tokens, head_size // 2] - sin: [num_tokens, head_size // 2] - is_neox_style: Whether to use the Neox-style or GPT-J-style rotary - positional embeddings. - """ - cos = cos.unsqueeze(-2).to(x.dtype) - sin = sin.unsqueeze(-2).to(x.dtype) - if is_neox_style: - x1, x2 = torch.chunk(x, 2, dim=-1) - else: - x1 = x[..., ::2] - x2 = x[..., 1::2] - o1 = x1 * cos - x2 * sin - o2 = x2 * cos + x1 * sin - if is_neox_style: - return torch.cat((o1, o2), dim=-1) - else: - return torch.stack((o1, o2), dim=-1).flatten(-2) - - -class MRotaryEmbedding(RotaryEmbedding): - """Rotary Embedding with Multimodal Sections.""" - - def __init__( - self, - head_size: int, - rotary_dim: int, - max_position_embeddings: int, - base: float, - is_neox_style: bool, - dtype: torch.dtype, - mrope_section: list[int] | None = None, - mrope_interleaved: bool = False, - # YaRN parameters. - *, - scaling_factor: float | None = None, - extrapolation_factor: float = 1, - attn_factor: float = 1, - beta_fast: int = 32, - beta_slow: int = 1, - ) -> None: - # In Qwen2.5-VL, the maximum index value is related to the duration of - # the input video. We enlarge max_position_embeddings to 4 times to get - # a larger the cos and sin cache. - self.cache_max_position_num = max_position_embeddings * 4 - super().__init__( - head_size, - rotary_dim, - self.cache_max_position_num, - base, - is_neox_style, - dtype, - ) - - self.mrope_section = mrope_section - if self.mrope_section: - print( - "Warning: mrope_section check is disabled in Qwen2.5-Omni, " - "this may cause errors, and should be restored in the future." - ) - # assert sum(self.mrope_section) == rotary_dim // 2 - - def forward( - self, - positions: torch.Tensor, - query: torch.Tensor, - key: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """PyTorch-native implementation equivalent to forward(). - - Args: - positions: - [num_tokens,] (text only) or - [3, num_tokens] (T/H/W positions with multimodal inputs) - query: [num_tokens, num_heads * head_size] - key: [num_tokens, num_kv_heads * head_size] - """ - assert positions.ndim == 1 or positions.ndim == 2 - assert key is not None - - num_tokens = positions.shape[-1] - cos_sin = self.cos_sin_cache[positions] - cos, sin = cos_sin.chunk(2, dim=-1) - if positions.ndim == 2: - assert self.mrope_section - - cos = torch.cat( - [m[i] for i, m in enumerate(cos.split(self.mrope_section, dim=-1))], - dim=-1, - ) - sin = torch.cat( - [m[i] for i, m in enumerate(sin.split(self.mrope_section, dim=-1))], - dim=-1, - ) - - query_shape = query.shape - query = query.view(num_tokens, -1, self.head_size) - query_rot = query[..., : self.rotary_dim] - query_pass = query[..., self.rotary_dim :] - # query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin, - # self.is_neox_style) - query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style) - query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) - - key_shape = key.shape - key = key.view(num_tokens, -1, self.head_size) - key_rot = key[..., : self.rotary_dim] - key_pass = key[..., self.rotary_dim :] - # key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin, - # self.is_neox_style) - key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style) - key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) - return query, key @classmethod def get_input_positions( @@ -440,9 +336,6 @@ class MRotaryEmbedding(RotaryEmbedding): |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |... """ - # TODO(fyabc): refactor and share more code with - # _vl_get_input_positions_tensor. - thinker_config = hf_config.thinker_config try: audio_token_id = thinker_config.audio_token_index @@ -613,29 +506,6 @@ class MRotaryEmbedding(RotaryEmbedding): ranges[index].append(num) return ranges - @staticmethod - def get_next_input_positions( - mrope_position_delta: int, - context_len: int, - seq_len: int, - ) -> list[list[int]]: - return [list(range(context_len + mrope_position_delta, seq_len + mrope_position_delta)) for _ in range(3)] - - @staticmethod - def get_next_input_positions_tensor( - out: np.ndarray, - out_offset: int, - mrope_position_delta: int, - context_len: int, - num_new_tokens: int, - ): - values = np.arange( - mrope_position_delta + context_len, - mrope_position_delta + context_len + num_new_tokens, - dtype=out.dtype, - ) - out[:, out_offset : out_offset + num_new_tokens] = values - @classmethod def omni_get_updates_use_audio_in_video( cls, diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index f4afe675a..936a18427 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -22,6 +22,9 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniThinkerProcessingInfo, ) from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix +from vllm.model_executor.models.vision import ( + get_llm_pos_ids_for_vision, +) # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav from vllm.multimodal import MULTIMODAL_REGISTRY @@ -35,7 +38,6 @@ from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges -from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision from vllm_omni.platforms import current_omni_platform TALKER_CODEC_EOS_TOKEN_ID = 8294 diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index a2250b82e..bb2de4f14 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -286,7 +286,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration( norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), quant_config=quant_config, prefix=maybe_prefix(prefix, "visual"), - multimodal_config=multimodal_config, ) else: self.visual = None diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 2d479062e..24f2aecc1 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -82,7 +82,7 @@ from vllm.model_executor.models.vision import ( from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItems from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems -from vllm.multimodal.processing import ( +from vllm.multimodal.processing.processor import ( MultiModalPromptUpdates, PlaceholderFeaturesInfo, PromptReplacement, @@ -609,7 +609,8 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix feature_lens=audio_feature_lengths, aftercnn_lens=audio_output_lengths, ) - audio_features = audio_outputs.last_hidden_state + # OMNI: audio_tower.forward() returns hidden_states tensor directly + audio_features = audio_outputs return audio_features.split(audio_output_lengths.tolist()) @@ -674,7 +675,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), quant_config=quant_config, prefix=maybe_prefix(prefix, "visual"), - multimodal_config=multimodal_config, ) self.quant_config = quant_config diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py index dcf9d54ac..dde690068 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py @@ -504,12 +504,19 @@ class Qwen3TTSConfig(PretrainedConfig): self.tts_bos_token_id = tts_bos_token_id self.tts_eos_token_id = tts_eos_token_id + # TODO: remove these dummy values after + self.image_token_id = 0 # dummy image token id + self.video_token_id = 0 # dummy video token id + self.vision_start_token_id = 0 # dummy vision start token id + self.vision_config = PretrainedConfig() # dummy vision config + self.vision_config.spatial_merge_size = 1 + def get_text_config(self, **kwargs): # vLLM expects text config to expose hidden_size/num_attention_heads. # For Qwen3 TTS, the talker config is the text model config. config = self.talker_config - if hasattr(config, "rope_parameters"): - delattr(config, "rope_parameters") + # if hasattr(config, "rope_parameters"): + # delattr(config, "rope_parameters") return config diff --git a/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py index 75fe5dbf4..1e759a8d2 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py @@ -430,6 +430,22 @@ def mel_spectrogram( return mel_spec +def _compute_default_rope_parameters( + config, + device, +): + base = config.rope_theta + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0) + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + dim = int(head_dim * partial_rotary_factor) + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)) + return inv_freq, attention_factor + + class Qwen3TTSPreTrainedModel(PreTrainedModel): config_class = Qwen3TTSConfig base_model_prefix = "model" @@ -502,7 +518,9 @@ class Qwen3TTSTalkerRotaryEmbedding(nn.Module): self.original_max_seq_len = config.max_position_embeddings self.config = config - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + self.rope_init_fn: Callable = _compute_default_rope_parameters + if self.rope_type != "default": + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) @@ -538,7 +556,9 @@ class Qwen3TTSRotaryEmbedding(nn.Module): self.original_max_seq_len = config.max_position_embeddings self.config = config - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + self.rope_init_fn: Callable = _compute_default_rope_parameters + if self.rope_type != "default": + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) @@ -1389,7 +1409,6 @@ class Qwen3TTSTalkerModel(Qwen3TTSTalkerTextPreTrainedModel): def __init__(self, config): super().__init__(config) - self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.layers = nn.ModuleList( [Qwen3TTSTalkerDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] diff --git a/vllm_omni/model_executor/models/vision.py b/vllm_omni/model_executor/models/vision.py deleted file mode 100644 index 850286a59..000000000 --- a/vllm_omni/model_executor/models/vision.py +++ /dev/null @@ -1,23 +0,0 @@ -import torch - - -def get_llm_pos_ids_for_vision( - start_idx: int, - vision_idx: int, - spatial_merge_size: int, - t_index: list[int], - grid_hs: torch.Tensor, - grid_ws: torch.Tensor, -) -> torch.Tensor: - llm_pos_ids_list = [] - llm_grid_h = grid_hs[vision_idx] // spatial_merge_size - llm_grid_w = grid_ws[vision_idx] // spatial_merge_size - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten() - t_index_tensor = ( - torch.Tensor(t_index).to(llm_grid_h.device).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten() - ) - _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index]) - llm_pos_ids_list.append(_llm_pos_ids + start_idx) - llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1) - return llm_pos_ids diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index d67922000..6cba9d96a 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -59,7 +59,7 @@ class OmniRequestOutput: prompt: OmniPromptType | None = None latents: torch.Tensor | None = None metrics: dict[str, Any] = field(default_factory=dict) - multimodal_output: dict[str, Any] = field(default_factory=dict) + _multimodal_output: dict[str, Any] = field(default_factory=dict) @classmethod def from_pipeline( @@ -116,10 +116,21 @@ class OmniRequestOutput: prompt=prompt, latents=latents, metrics=metrics or {}, - multimodal_output=multimodal_output or {}, + _multimodal_output=multimodal_output or {}, 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 proxies to request_output.multimodal_output. + For diffusion outputs, this returns the local _multimodal_output field. + """ + if self.request_output is not None: + return getattr(self.request_output, "multimodal_output", {}) + return self._multimodal_output + @property def num_images(self) -> int: """Return the number of generated images.""" @@ -229,7 +240,7 @@ class OmniRequestOutput: f"prompt={self.prompt!r}", f"latents={self.latents}", f"metrics={self.metrics}", - f"multimodal_output={self.multimodal_output}", + f"multimodal_output={self._multimodal_output}", ] return f"OmniRequestOutput({', '.join(parts)})" diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 687ff5186..90c718d40 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -12,7 +12,7 @@ from vllm.v1.request import Request as _OriginalRequest import vllm_omni.logger # noqa: F401 from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt -from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding +from vllm_omni.model_executor.layers.rotary_embedding import OmniMRotaryEmbedding from vllm_omni.request import OmniRequest for module_name, module in sys.modules.items(): @@ -26,7 +26,7 @@ for module_name, module in sys.modules.items(): if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt: module.TokensPrompt = OmniTokensPrompt if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding: - module.MRotaryEmbedding = MRotaryEmbedding + module.MRotaryEmbedding = OmniMRotaryEmbedding if hasattr(module, "Request") and module.Request == _OriginalRequest: module.Request = OmniRequest if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest: diff --git a/vllm_omni/request.py b/vllm_omni/request.py index 864d9dce4..c992816de 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -1,6 +1,8 @@ from collections.abc import Callable from typing import TYPE_CHECKING +import numpy as np +import torch from vllm.v1.request import Request if TYPE_CHECKING: @@ -25,21 +27,35 @@ class OmniRequest(Request): def __init__( self, - prompt_embeds: PromptEmbedsPayload | None = None, + prompt_embeds: PromptEmbedsPayload | torch.Tensor | None = None, # Optional external request ID for tracking external_req_id: str | None = None, additional_information: AdditionalInformationPayload | None = None, *args, **kwargs, ): - super().__init__(*args, **kwargs) - # Serialized prompt embeddings payload (optional) - self.prompt_embeds: PromptEmbedsPayload | None = prompt_embeds + prompt_embeds_tensor = self._maybe_decode_prompt_embeds(prompt_embeds) + super().__init__(prompt_embeds=prompt_embeds_tensor, *args, **kwargs) + # Preserve serialized prompt embeddings payload (optional) + self.prompt_embeds_payload: PromptEmbedsPayload | None = ( + prompt_embeds if isinstance(prompt_embeds, PromptEmbedsPayload) else None + ) # Optional external request ID for tracking self.external_req_id: str | None = external_req_id # Serialized additional information payload (optional) self.additional_information: AdditionalInformationPayload | None = additional_information + @staticmethod + def _maybe_decode_prompt_embeds( + prompt_embeds: PromptEmbedsPayload | torch.Tensor | None, + ) -> torch.Tensor | None: + if isinstance(prompt_embeds, PromptEmbedsPayload): + dtype = getattr(np, prompt_embeds.dtype) + arr = np.frombuffer(prompt_embeds.data, dtype=dtype) + arr = arr.reshape(prompt_embeds.shape) + return torch.from_numpy(arr) + return prompt_embeds + @classmethod def from_engine_core_request( cls, @@ -74,4 +90,5 @@ class OmniRequest(Request): trace_headers=request.trace_headers, block_hasher=block_hasher, additional_information=request.additional_information, + resumable=request.resumable, ) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 8440cb5b8..9332abeb1 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -13,7 +13,8 @@ import numpy as np import torch from vllm.config import CUDAGraphMode from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer -from vllm.distributed.kv_transfer import get_kv_transfer_group +from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group +from vllm.distributed.parallel_state import get_pp_group, get_tp_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( @@ -28,9 +29,6 @@ from vllm.v1.worker.gpu_model_runner import ( EMPTY_MODEL_RUNNER_OUTPUT, AsyncGPUModelRunnerOutput, IntermediateTensors, - get_pp_group, - get_tp_group, - has_kv_transfer_group, ) from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import is_residual_scattered_for_sp @@ -221,6 +219,24 @@ class GPUARModelRunner(OmniGPUModelRunner): use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + # True if any attention backend handles KV cache update separately + # from forward() (i.e., forward_includes_kv_cache_update=False). When true, + # slot_mappings must use padded dimensions to match the key/value tensors. + from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec + + has_separate_kv_update = not all( + all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id]) + for id, spec in enumerate(self.kv_cache_config.kv_cache_groups) + if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec) + ) + + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded, + num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs), + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata( num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded if pad_attn else None, @@ -232,6 +248,7 @@ class GPUARModelRunner(OmniGPUModelRunner): use_spec_decode=use_spec_decode, num_scheduled_tokens=scheduler_output.num_scheduled_tokens, cascade_attn_prefix_lens=cascade_attn_prefix_lens, + slot_mappings=slot_mappings_by_group, ) ( @@ -262,6 +279,7 @@ class GPUARModelRunner(OmniGPUModelRunner): cudagraph_runtime_mode=cudagraph_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, # OMNI: required for KV cache operations ), record_function_or_nullcontext("gpu_model_runner: forward"), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index a1d6b4665..98b5afdda 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -91,6 +91,11 @@ class GPUARWorker(OmniWorkerMixin, GPUWorker): num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 init_workspace_manager(self.device, num_ubatches) + if self.use_v2_model_runner: + # OMNI: v2 model runner does not yet include omni hooks. + logger.warning("OMNI GPUARWorker forces v1 model runner for omni hooks.") + self.use_v2_model_runner = False + # Construct the model runner self.model_runner = GPUARModelRunner(self.vllm_config, self.device) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 87e1419c8..9125cc52a 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -15,10 +15,11 @@ import torch from vllm.config import CUDAGraphMode from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group +from vllm.distributed.parallel_state import get_pp_group +from vllm.forward_context import set_forward_context from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( RoutedExpertsCapturer, ) -from vllm.model_executor.models.interfaces import supports_mm_encoder_only from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output @@ -29,8 +30,6 @@ from vllm.v1.worker.gpu_model_runner import ( AsyncGPUModelRunnerOutput, IntermediateTensors, PerLayerAttnMetadata, - get_pp_group, - set_forward_context, ) from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs @@ -190,6 +189,24 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + # OMNI: True if any attention backend handles KV cache update separately + # from forward() (i.e., forward_includes_kv_cache_update=False). When true, + # slot_mappings must use padded dimensions to match the key/value tensors. + from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec + + has_separate_kv_update = not all( + all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id]) + for id, spec in enumerate(self.kv_cache_config.kv_cache_groups) + if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec) + ) + + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded, + num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs), + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata( num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded if pad_attn else None, @@ -201,6 +218,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): use_spec_decode=use_spec_decode, num_scheduled_tokens=scheduler_output.num_scheduled_tokens, cascade_attn_prefix_lens=cascade_attn_prefix_lens, + slot_mappings=slot_mappings_by_group, ) ( @@ -235,6 +253,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): cudagraph_runtime_mode=cudagraph_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, # OMNI: required for KV cache operations ), record_function_or_nullcontext("Forward"), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, @@ -435,7 +454,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): remove_lora: If False, dummy LoRAs are not destroyed after the run activate_lora: If False, dummy_run is performed without LoRAs. """ - if supports_mm_encoder_only(self.model): + mm_config = self.vllm_config.model_config.multimodal_config + if mm_config and mm_config.mm_encoder_only: # The current dummy run only covers LM execution, so we can skip it. # mm encoder dummy run may need to add in the future. return torch.tensor([]), torch.tensor([]) @@ -539,6 +559,14 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): attn_metadata: PerLayerAttnMetadata | None = None + # OMNI: Get slot mappings before building attention metadata + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens, + num_reqs_padded=num_reqs_padded, + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + # If force_attention is True, we always capture attention. Otherwise, # it only happens for cudagraph_runtime_mode=FULL. if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: @@ -564,6 +592,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): max_query_len=max_query_len, ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, + slot_mappings=slot_mappings_by_group, ) with self.maybe_dummy_run_with_lora( @@ -628,6 +657,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): cudagraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, # OMNI: required for KV cache operations ), ): outputs = self.model( diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 0aa13ed76..91f90d6b4 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -91,6 +91,11 @@ class GPUGenerationWorker(OmniWorkerMixin, GPUWorker): num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 init_workspace_manager(self.device, num_ubatches) + if self.use_v2_model_runner: + # OMNI: v2 model runner does not yet include omni hooks. + logger.warning("OMNI GPUGenerationWorker forces v1 model runner for omni hooks.") + self.use_v2_model_runner = False + self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device) if self.rank == 0: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index b3d1833f2..cb14ed49a 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -8,7 +8,7 @@ from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding -from vllm.model_executor.models.interfaces import supports_mm_encoder_only, supports_mrope +from vllm.model_executor.models.interfaces import supports_mrope from vllm.model_executor.models.interfaces_base import VllmModelForPooling from vllm.sampling_params import SamplingType from vllm.utils.import_utils import LazyLoader @@ -448,7 +448,8 @@ class OmniGPUModelRunner(GPUModelRunner): remove_lora: If False, dummy LoRAs are not destroyed after the run activate_lora: If False, dummy_run is performed without LoRAs. """ - if supports_mm_encoder_only(self.model): + mm_config = self.vllm_config.model_config.multimodal_config + if mm_config and mm_config.mm_encoder_only: # The current dummy run only covers LM execution, so we can skip it. # mm encoder dummy run may need to add in the future. return torch.tensor([]), torch.tensor([]) @@ -552,6 +553,13 @@ class OmniGPUModelRunner(GPUModelRunner): attn_metadata: PerLayerAttnMetadata | None = None + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens, + num_reqs_padded=num_reqs_padded, + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + # If force_attention is True, we always capture attention. Otherwise, # it only happens for cudagraph_runtime_mode=FULL. if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: @@ -577,6 +585,7 @@ class OmniGPUModelRunner(GPUModelRunner): max_query_len=max_query_len, ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, + slot_mappings=slot_mappings_by_group, ) with self.maybe_dummy_run_with_lora( @@ -641,6 +650,7 @@ class OmniGPUModelRunner(GPUModelRunner): cudagraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, ), ): if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"): @@ -688,6 +698,7 @@ class OmniGPUModelRunner(GPUModelRunner): num_tokens, use_cudagraphs=use_cudagraphs, is_graph_capturing=is_graph_capturing, + slot_mappings=slot_mappings, ) # We register layerwise NVTX hooks here after the first dynamo tracing is