From bf69cd94bbc1ee755ae82c075002c75ec9d296d3 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 23 Feb 2026 03:36:01 +0000 Subject: [PATCH] Rebase 2026/02/23 Update Dockerfile for vllm wheel version and enhance input processing in Omni models. Added decoding for byte messages in async requests, included new metrics in benchmarks, and improved pause state handling in the scheduler. Refined input and output processing methods for better multimodal support and error handling. Signed-off-by: tzhouam --- docker/Dockerfile.ci | 2 +- vllm_omni/benchmarks/patch/patch.py | 3 ++ .../core/sched/omni_generation_scheduler.py | 12 ++++- vllm_omni/engine/input_processor.py | 32 +++++------- vllm_omni/engine/output_processor.py | 4 +- vllm_omni/inputs/preprocess.py | 8 +-- .../qwen3_omni/qwen3_omni_moe_thinker.py | 2 +- .../worker/gpu_generation_model_runner.py | 52 ++++++++++--------- vllm_omni/worker/gpu_model_runner.py | 52 ++++++++++--------- 9 files changed, 90 insertions(+), 77 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index e8c23f18f..e7c8d2df4 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -18,7 +18,7 @@ RUN apt-get update && \ # Note: --prerelease=allow is required because nightly/commit wheels use # pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default, # causing it to fall back to the older stable release from PyPI. -RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/16f24e87975ef4cd2c12879425062913ef62f6fd +RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/944ffb59680c0210ec54ddb43a3c7ef015e1f842 RUN uv pip uninstall --system nvidia-cublas-cu12 diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 5a03a8069..f89fbe035 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -128,6 +128,8 @@ async def async_request_openai_chat_omni_completions( messages = handler.add_chunk(chunk_bytes) for message in messages: + if type(message) is bytes: + message = message.decode("utf-8") # NOTE: SSE comments (often used as pings) start with # a colon. These are not JSON data payload and should # be skipped. @@ -486,6 +488,7 @@ async def benchmark( "errors": [output.error for output in outputs], "max_output_tokens_per_s": metrics.max_output_tokens_per_s, "max_concurrent_requests": metrics.max_concurrent_requests, + "rtfx": metrics.rtfx, } else: result = { diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index ef1c4c7c9..fff0730f8 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -6,6 +6,7 @@ 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.interface import PauseState from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.request_queue import create_request_queue from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler @@ -39,8 +40,12 @@ class OmniGenerationScheduler(VLLMScheduler): """ token_budget = self.max_num_scheduled_tokens + if self._pause_state == PauseState.PAUSED_ALL: + token_budget = 0 scheduled_timestamp = time.monotonic() + self.kv_cache_manager.new_step_starts() + scheduled_new_reqs: list[Request] = [] req_to_new_blocks: dict[str, KVCacheBlocks] = {} @@ -116,7 +121,12 @@ class OmniGenerationScheduler(VLLMScheduler): # 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: + while ( + self.waiting + and token_budget > 0 + and len(self.running) < self.max_num_running_reqs + and self._pause_state == PauseState.UNPAUSED + ): request = self.waiting.peek_request() # OMNI: Skip requests that are not in self.requests if request.request_id not in self.requests or ( diff --git a/vllm_omni/engine/input_processor.py b/vllm_omni/engine/input_processor.py index 37aadc7b0..e5c45ba6e 100644 --- a/vllm_omni/engine/input_processor.py +++ b/vllm_omni/engine/input_processor.py @@ -10,25 +10,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 - -try: - from vllm.multimodal.processing import set_request_id -except ImportError: # vllm without set_request_id (older releases) - from contextlib import contextmanager - - @contextmanager - def set_request_id(_request_id: str): - yield - - from vllm.multimodal.utils import argsort_mm_positions from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer -from vllm.renderers.inputs import DictPrompt, TokPrompt from vllm.sampling_params import SamplingParams from vllm.tasks import SupportedTask from vllm.utils import length_from_prompt_token_ids_or_embeds -from vllm.utils.torch_utils import set_default_torch_num_threads +from vllm.utils.jsontree import json_iter_leaves from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.engine import ( @@ -131,7 +119,7 @@ class OmniInputProcessor(InputProcessor): def process_inputs( self, request_id: str, - prompt: PromptType | DictPrompt | TokPrompt, + prompt: PromptType | ProcessorInputs, params: SamplingParams | PoolingParams, arrival_time: float | None = None, lora_request: LoRARequest | None = None, @@ -192,11 +180,10 @@ class OmniInputProcessor(InputProcessor): if arrival_time is None: arrival_time = time.time() - with set_request_id(request_id), set_default_torch_num_threads(): - processed_inputs = self.input_preprocessor.preprocess( - prompt, - tokenization_kwargs=tokenization_kwargs, - ) + processed_inputs = self.input_preprocessor.preprocess( + prompt, + tokenization_kwargs=tokenization_kwargs, + ) self._platform_validate_request(processed_inputs, params) @@ -236,6 +223,13 @@ class OmniInputProcessor(InputProcessor): decoder_mm_positions = decoder_inputs["mm_placeholders"] decoder_mm_hashes = decoder_inputs["mm_hashes"] + if not all(isinstance(leaf, str) for leaf in json_iter_leaves(decoder_mm_hashes)): + raise ValueError( + f"mm_hashes must contain only strings, got: {decoder_mm_hashes}. " + "This is likely due to an incorrect custom implementation of " + "MultiModalProcessor.apply method." + ) + # Merge and flatten multimodal placeholders, hashes and inputs # from dictionaries to lists, and sort them by each item's position # in the input sequence. diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index f87dc70ef..8aa8ad9f4 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -195,7 +195,7 @@ class OmniRequestState(RequestState): if not ( finished or self.sent_tokens_offset == 0 - or len(self.detokenizer.output_token_ids) - self.sent_tokens_offset >= self.stream_interval + or self.detokenizer.num_output_tokens() - self.sent_tokens_offset >= self.stream_interval ): return None @@ -203,7 +203,7 @@ class OmniRequestState(RequestState): # Send tokens from the offset in DELTA mode, otherwise all # tokens are sent. new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :] - self.sent_tokens_offset = len(self.detokenizer.output_token_ids) + self.sent_tokens_offset = self.detokenizer.num_output_tokens() external_req_id = self.external_req_id output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts) diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 53e09f08a..4c16295c5 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -92,6 +92,8 @@ class OmniInputPreprocessor(InputPreprocessor): inputs["prompt_embeds"] = prompt_embeds if additional_information is not None: inputs["additional_information"] = additional_information + if prompt_text := parsed_content.get("prompt"): + inputs["prompt"] = prompt_text if cache_salt := parsed_content.get("cache_salt"): inputs["cache_salt"] = cache_salt @@ -132,14 +134,14 @@ class OmniInputPreprocessor(InputPreprocessor): * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance """ + if "prompt_embeds" in prompt: + return self._process_embeds(prompt) # type: ignore[arg-type] + if "prompt_token_ids" in prompt: return self._process_tokens( prompt, # type: ignore[arg-type] ) - if "prompt_embeds" in prompt: - return self._process_embeds(prompt) # type: ignore[arg-type] - if "prompt" in prompt: return self._process_text( prompt, # type: ignore[arg-type] 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 a0920c069..6e1b4b9b4 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 @@ -253,7 +253,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( return x # NOTE: WhisperFeatureExtractor cannot handle empty list of audios - feature_extractor = self.info.get_feature_extractor() + feature_extractor = self.info.get_feature_extractor(**mm_kwargs) hop_length = feature_extractor.hop_length if audios: # NOTE: Qwen3-Omni processor accept "audio" diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index aa75d201c..192fadd7e 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -605,33 +605,35 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): 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: - if create_mixed_batch: - # In the mixed batch mode (used for FI warmup), we use - # shorter sequence lengths to run faster. - # TODO(luka) better system for describing dummy batches - seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] - else: - seq_lens = max_query_len # type: ignore[assignment] - self.seq_lens.np[:num_reqs] = seq_lens - self.seq_lens.np[num_reqs:] = 0 - self.seq_lens.copy_to_gpu() + with self.synchronize_input_prep(): + # 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: + if create_mixed_batch: + # In the mixed batch mode (used for FI warmup), we use + # shorter sequence lengths to run faster. + # TODO(luka) better system for describing dummy batches + seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] + else: + seq_lens = max_query_len # type: ignore[assignment] + self.seq_lens.np[:num_reqs] = seq_lens + self.seq_lens.np[num_reqs:] = 0 + self.seq_lens.copy_to_gpu() - cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) - self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens - self.query_start_loc.copy_to_gpu() + cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() - pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL - attn_metadata, _ = self._build_attention_metadata( - num_tokens=num_tokens_unpadded, - num_reqs=num_reqs_padded, - 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, - ) + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded if pad_attn else None, + num_reqs=num_reqs_padded, + 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( self.lora_config, diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index a01ad113a..713da73fd 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -603,33 +603,35 @@ class OmniGPUModelRunner(GPUModelRunner): 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: - if create_mixed_batch: - # In the mixed batch mode (used for FI warmup), we use - # shorter sequence lengths to run faster. - # TODO(luka) better system for describing dummy batches - seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] - else: - seq_lens = max_query_len # type: ignore[assignment] - self.seq_lens.np[:num_reqs] = seq_lens - self.seq_lens.np[num_reqs:] = 0 - self.seq_lens.copy_to_gpu() + with self.synchronize_input_prep(): + # 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: + if create_mixed_batch: + # In the mixed batch mode (used for FI warmup), we use + # shorter sequence lengths to run faster. + # TODO(luka) better system for describing dummy batches + seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] + else: + seq_lens = max_query_len # type: ignore[assignment] + self.seq_lens.np[:num_reqs] = seq_lens + self.seq_lens.np[num_reqs:] = 0 + self.seq_lens.copy_to_gpu() - cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) - self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens - self.query_start_loc.copy_to_gpu() + cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() - pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL - attn_metadata, _ = self._build_attention_metadata( - num_tokens=num_tokens_unpadded, - num_reqs=num_reqs_padded, - 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, - ) + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded if pad_attn else None, + num_reqs=num_reqs_padded, + 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( self.lora_config,