From 034d579d770bb39fa2fb8d2bd1caca27f1affae3 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Thu, 19 Mar 2026 05:14:34 +0000 Subject: [PATCH] Update vLLM precompiled wheel version and refactor various components for improved functionality and error handling. Changes include modifying the `get_name` method in `DiffusionQuantizationConfig` to use class-level attributes, enhancing metadata handling in `StageEngineCoreClient`, and adjusting the `AsyncOmni` class constructor for better parameter management. Additionally, improved error handling in `omni_snapshot_download` and optimized initial chunk size computation in `talker2code2wav_async_chunk` for dynamic load adaptation. Signed-off-by: tzhouam --- docker/Dockerfile.ci | 2 +- vllm_omni/diffusion/quantization/base.py | 11 ++++--- vllm_omni/engine/stage_engine_core_client.py | 33 +++++++++++-------- vllm_omni/entrypoints/async_omni.py | 21 +++++++++--- vllm_omni/entrypoints/omni_base.py | 5 +++ .../qwen3_omni/qwen3_omni_moe_thinker.py | 9 +++-- .../stage_input_processors/qwen3_tts.py | 22 +++++-------- 7 files changed, 61 insertions(+), 42 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index c56e9e83f..834636873 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -16,7 +16,7 @@ RUN uv pip uninstall --system -y vllm || true # Install vLLM from precompiled wheel at the selected commit. # Must use direct URL because the wheel has a PEP 440 local version identifier # (e.g. +g0a0a1a198) which pip/uv refuse to install from a PEP 503 package index. -ENV VLLM_PRECOMPILED_WHEEL_COMMIT=0a0a1a198be88e1782b52fa31738896468200a76 +ENV VLLM_PRECOMPILED_WHEEL_COMMIT=89138b21cc246ae944c741d5c399c148e2b770ab RUN VLLM_WHEEL_URL=$(python3 -c "import urllib.request,re; \ html=urllib.request.urlopen('https://wheels.vllm.ai/${VLLM_PRECOMPILED_WHEEL_COMMIT}/vllm/').read().decode(); \ m=re.search(r'>(\S+x86_64\.whl)<',html); \ diff --git a/vllm_omni/diffusion/quantization/base.py b/vllm_omni/diffusion/quantization/base.py index 17e6d32ea..8a35547f9 100644 --- a/vllm_omni/diffusion/quantization/base.py +++ b/vllm_omni/diffusion/quantization/base.py @@ -31,14 +31,15 @@ class DiffusionQuantizationConfig(ABC): # The underlying vLLM config instance _vllm_config: "QuantizationConfig | None" = None - def get_name(self) -> str: + @classmethod + def get_name(cls) -> str: """Return the quantization method name (e.g., 'fp8', 'int8'). - By default, delegates to the underlying vLLM config instance. + Delegates to the underlying vLLM config class's get_name(). """ - if self._vllm_config is not None: - return self._vllm_config.get_name() - raise NotImplementedError("Subclass must initialize _vllm_config or override get_name().") + if cls.quant_config_cls is not None: + return cls.quant_config_cls.get_name() + raise NotImplementedError("Subclass must set quant_config_cls or override get_name().") def get_vllm_quant_config(self) -> "QuantizationConfig | None": """Return the underlying vLLM QuantizationConfig for linear layers.""" diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index a4e8bc0e9..284cc2d31 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -38,8 +38,12 @@ class StageEngineCoreClient(AsyncMPClient): self, vllm_config: Any, executor_class: type, - metadata: StageMetadata, + log_stats: bool = False, client_addresses: dict[str, str] | None = None, + client_count: int = 1, + client_index: int = 0, + *, + metadata: StageMetadata | None = None, engine_manager: Any = None, coordinator: Any = None, ): @@ -51,17 +55,18 @@ class StageEngineCoreClient(AsyncMPClient): and calls super().__init__(). """ # -------- Stage metadata (public fields used at runtime) -------- - self.stage_id = metadata.stage_id - self.stage_type = metadata.stage_type - self.engine_output_type = metadata.engine_output_type - self.is_comprehension = metadata.is_comprehension - self.requires_multimodal_data = metadata.requires_multimodal_data - self.engine_input_source = metadata.engine_input_source - self.final_output = metadata.final_output - self.final_output_type = metadata.final_output_type - self.default_sampling_params = metadata.default_sampling_params - self.custom_process_input_func = metadata.custom_process_input_func - self.model_stage = metadata.model_stage + if metadata is not None: + self.stage_id = metadata.stage_id + self.stage_type = metadata.stage_type + self.engine_output_type = metadata.engine_output_type + self.is_comprehension = metadata.is_comprehension + self.requires_multimodal_data = metadata.requires_multimodal_data + self.engine_input_source = metadata.engine_input_source + self.final_output = metadata.final_output + self.final_output_type = metadata.final_output_type + self.default_sampling_params = metadata.default_sampling_params + self.custom_process_input_func = metadata.custom_process_input_func + self.model_stage = metadata.model_stage self.engine_outputs: Any = None @@ -73,8 +78,10 @@ class StageEngineCoreClient(AsyncMPClient): super().__init__( vllm_config, executor_class, - log_stats=False, + log_stats=log_stats, client_addresses=client_addresses, + client_count=client_count, + client_index=client_index, ) if engine_manager is not None: self.resources.engine_manager = engine_manager diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 3ba2c4a7e..57c432ee1 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -29,6 +29,7 @@ from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: from vllm.inputs.preprocess import InputPreprocessor from vllm.tokenizers import TokenizerLike + from vllm.v1.engine import PauseMode from vllm_omni.inputs.data import OmniPromptType, OmniSamplingParams @@ -66,7 +67,7 @@ class AsyncOmni(EngineClient, OmniBase): ... print(output) """ - def __init__(self, model: str, **kwargs: Any) -> None: + def __init__(self, *args: Any, model: str = "", **kwargs: Any) -> None: OmniBase.__init__(self, model=model, **kwargs) self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False @@ -130,9 +131,17 @@ class AsyncOmni(EngineClient, OmniBase): async def generate( self, prompt: OmniPromptType, - request_id: str, - sampling_params_list: Sequence[OmniSamplingParams] | None = None, + sampling_params: Any = None, + request_id: str = "", *, + prompt_text: str | None = None, + lora_request: Any = None, + tokenization_kwargs: dict[str, Any] | None = None, + trace_headers: Any = None, + priority: int = 0, + data_parallel_rank: int | None = None, + reasoning_ended: bool | None = None, + sampling_params_list: Sequence[OmniSamplingParams] | None = None, output_modalities: list[str] | None = None, ) -> AsyncGenerator[OmniRequestOutput, None]: """Generate outputs for the given prompt asynchronously. @@ -226,6 +235,7 @@ class AsyncOmni(EngineClient, OmniBase): trace_headers: dict[str, str] | None = None, priority: int = 0, tokenization_kwargs: dict[str, Any] | None = None, + reasoning_ended: bool | None = None, ) -> AsyncGenerator[PoolingRequestOutput, None]: """EngineClient.encode() stub. @@ -396,6 +406,7 @@ class AsyncOmni(EngineClient, OmniBase): async def pause_generation( self, *, + mode: PauseMode = "abort", wait_for_inflight_requests: bool = False, clear_cache: bool = True, ) -> None: @@ -467,7 +478,7 @@ class AsyncOmni(EngineClient, OmniBase): logger.warning("[AsyncOmni] reset_prefix_cache not yet supported with Orchestrator process") return True - async def sleep(self, level: int = 1) -> None: + async def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None: """Sleep all stages. Best-effort: unsupported stages will emit a TODO result. @@ -584,7 +595,7 @@ class AsyncOmni(EngineClient, OmniBase): # ==================== Shutdown ==================== - def shutdown(self) -> None: + def shutdown(self, timeout: float | None = None) -> None: """Shutdown the engine.""" if self.final_output_task is not None: self.final_output_task.cancel() diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index fb5f3788d..af34a4221 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -53,6 +53,11 @@ def omni_snapshot_download(model_id: str) -> str: ) except huggingface_hub.errors.RepositoryNotFoundError: logger.warning("Repository not found for '%s'.", model_id) + except PermissionError: + logger.warning( + "Permission denied when downloading '%s'. Assuming the model is already cached locally.", + model_id, + ) return model_id 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 f6140f594..4693a2314 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 @@ -155,8 +155,7 @@ class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer): cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) except RuntimeError: logger.warning( - "torch.repeat_interleave not executable, " - "switching to vectorized searchsorted implementation." + "torch.repeat_interleave not executable, switching to vectorized searchsorted implementation." ) repeat_counts = grid_thw_tensor[:, 0] values = grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2] @@ -179,9 +178,9 @@ class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer): max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens) grid_thw_np = grid_thw_tensor.cpu().numpy().astype(np.int32) - cu_seqlens_np = np.repeat( - grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0] - ).cumsum(axis=0, dtype=np.int32) + cu_seqlens_np = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum( + axis=0, dtype=np.int32 + ) cu_seqlens_np = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens_np]) sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens( self.attn_backend, diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 69724dfc0..4db14f57a 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -116,18 +116,12 @@ def talker2code2wav_async_chunk( initial_chunk_size = int(entry.list_data[0]) per_request_override = True - # Dynamic IC: cache per request so boundaries stay stable for its lifetime. + # Dynamic IC: recalculate each call so IC adapts to current load. if not per_request_override: - _ic_cache = getattr(transfer_manager, "_cached_ic", None) - if _ic_cache is None: - _ic_cache = {} - transfer_manager._cached_ic = _ic_cache - if request_id not in _ic_cache: - max_ic = max_ic_for_chunk_size(chunk_size) - active = sum(1 for v in transfer_manager.code_prompt_token_ids.values() if len(v) > 0) - capacity = getattr(transfer_manager, "scheduler_max_num_seqs", 1) - _ic_cache[request_id] = compute_dynamic_initial_chunk_size(active, capacity, max_ic) - initial_chunk_size = _ic_cache[request_id] + max_ic = max_ic_for_chunk_size(chunk_size) + active = sum(1 for v in transfer_manager.code_prompt_token_ids.values() if len(v) > 0) + capacity = getattr(transfer_manager, "scheduler_max_num_seqs", 1) + initial_chunk_size = compute_dynamic_initial_chunk_size(active, capacity, max_ic) if chunk_size <= 0 or left_context_size_config < 0 or initial_chunk_size < 0: raise ValueError( @@ -145,11 +139,13 @@ def talker2code2wav_async_chunk( initial_chunk_size = chunk_size length = len(transfer_manager.code_prompt_token_ids[request_id]) + finished_tensor = torch.tensor(finished, dtype=torch.bool) + if length <= 0: if finished: return { "code_predictor_codes": [], - "finished": True, + "finished": finished_tensor, } return None @@ -197,5 +193,5 @@ def talker2code2wav_async_chunk( return { "code_predictor_codes": code_predictor_codes, "left_context_size": left_context_size, - "finished": finished, + "finished": finished_tensor, }