From e6d4cdd71a46ab97f5682d0cff6cd3e76db73400 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 9 Mar 2026 06:17:08 +0000 Subject: [PATCH 1/3] Refactor model mode initialization and enhance error handling in AsyncOmni. Update worker type retrieval in omni_stage and add audio token handling in Qwen3OmniMoe classes for improved robustness. Signed-off-by: tzhouam --- .../chunk_transfer_adapter.py | 2 +- vllm_omni/entrypoints/async_omni.py | 19 +++++++++++++++++-- vllm_omni/entrypoints/omni_stage.py | 2 +- .../qwen3_omni/qwen3_omni_moe_thinker.py | 8 ++++++-- 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index c7289d8b8..db8d5dbaa 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -39,7 +39,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs self.connector = self.create_connector(model_config) super().__init__(model_config) - self.model_mode = getattr(model_config, "worker_type", "ar") + self.model_mode = getattr(model_config, "worker_type", None) or "ar" # State specific to Chunk management self.custom_process_next_stage_input_func = None custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index b663430dc..463ecd7fb 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -401,7 +401,10 @@ class AsyncOmni(OmniBase): ) -> AsyncGenerator[OmniRequestOutput, None]: all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)} submit_flag = True + _loop_iter = 0 + _last_progress_ts = time.time() while not all(all_stages_finished.values()): + _loop_iter += 1 for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]): if all_stages_finished[stage_id]: continue @@ -410,6 +413,7 @@ class AsyncOmni(OmniBase): except asyncio.QueueEmpty: await asyncio.sleep(0.001) continue + _last_progress_ts = time.time() engine_outputs, finished, output_to_yield = self._process_single_result( result, stage, @@ -418,11 +422,20 @@ class AsyncOmni(OmniBase): ) if submit_flag and stage_id == 0: submit_flag = False - prompt_token_ids = engine_outputs.prompt_token_ids + prompt_token_ids = getattr(engine_outputs, "prompt_token_ids", None) + if prompt_token_ids is None: + prompt_token_ids = [] engine_input = copy.deepcopy(prompt) - next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) + try: + next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) + except Exception: + raise engine_input["prompt_token_ids"] = [0] * next_prompt_len engine_input["multi_modal_data"] = engine_input["mm_processor_kwargs"] = None + for _mm_key in ("mm_kwargs", "mm_hashes", "mm_placeholders", "multi_modal_uuids"): + engine_input.pop(_mm_key, None) + if engine_input.get("type") == "multimodal": + engine_input["type"] = "token" for i in range(1, len(self.stage_list)): task = { "request_id": request_id, @@ -709,6 +722,8 @@ class AsyncOmni(OmniBase): OMNI: Required by upstream OpenAIServingModels.__init__ which accesses engine_client.renderer. """ + if self.input_processor is None: + return None return self.input_processor.renderer async def do_log_stats(self) -> None: diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 4114ef3fb..c9356119e 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -207,7 +207,7 @@ def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int = def _resolve_worker_cls(engine_args: dict[str, Any]) -> None: - worker_type = engine_args.pop("worker_type", None) + worker_type = engine_args.get("worker_type", None) if not worker_type: return worker_cls = engine_args.get("worker_cls") 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 835dc6d6f..7dfb988eb 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 @@ -340,6 +340,11 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc seq_len=seq_len, mm_counts=mm_counts, ) + if audio_tokens is None: + feature_extractor = self.get_feature_extractor() + max_audio_samples = feature_extractor.chunk_length * feature_extractor.sampling_rate + max_audio_tokens = int(_get_feat_extract_output_lengths(torch.tensor([max_audio_samples])).item()) + audio_tokens = {"audio": max_audio_tokens} mm_max_tokens["audio"] = audio_tokens["audio"] return mm_max_tokens @@ -769,8 +774,7 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix feature_lens=audio_feature_lengths, aftercnn_lens=audio_output_lengths, ) - # OMNI: audio_tower.forward() returns hidden_states tensor directly - audio_features = audio_outputs + audio_features = audio_outputs if isinstance(audio_outputs, torch.Tensor) else audio_outputs.last_hidden_state return audio_features.split(audio_output_lengths.tolist()) From 79626586e1c21f9167ce2b5f371f17bc63b427c7 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 9 Mar 2026 08:02:25 +0000 Subject: [PATCH 2/3] Update GPU memory utilization in qwen3_omni.yaml from 0.6 to 0.45 to prevent out-of-memory errors during model execution. Signed-off-by: tzhouam --- tests/perf/stage_configs/qwen3_omni.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/perf/stage_configs/qwen3_omni.yaml b/tests/perf/stage_configs/qwen3_omni.yaml index 802f8dd24..fd5586b94 100644 --- a/tests/perf/stage_configs/qwen3_omni.yaml +++ b/tests/perf/stage_configs/qwen3_omni.yaml @@ -47,7 +47,7 @@ stage_args: model_arch: Qwen3OmniMoeForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.6 + gpu_memory_utilization: 0.45 enforce_eager: false trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav From a9b4367808a14ae03b5514b7f07b45402756e806 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 9 Mar 2026 08:22:26 +0000 Subject: [PATCH 3/3] Update GPU memory utilization and max number of batched tokens in qwen3_omni.yaml to optimize resource allocation during model execution. Signed-off-by: tzhouam --- tests/perf/stage_configs/qwen3_omni.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/perf/stage_configs/qwen3_omni.yaml b/tests/perf/stage_configs/qwen3_omni.yaml index fd5586b94..01ca7d9fb 100644 --- a/tests/perf/stage_configs/qwen3_omni.yaml +++ b/tests/perf/stage_configs/qwen3_omni.yaml @@ -47,7 +47,7 @@ stage_args: model_arch: Qwen3OmniMoeForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.45 + gpu_memory_utilization: 0.6 enforce_eager: false trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav @@ -85,7 +85,7 @@ stage_args: engine_output_type: audio # Final output: audio waveform gpu_memory_utilization: 0.1 distributed_executor_backend: "mp" - max_num_batched_tokens: 1000000 + max_num_batched_tokens: 100000 hf_config_name: thinker_config engine_input_source: [1] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav