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 <tzhouam@connect.ust.hk>

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); \

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:
+        if cls.quant_config_cls is not None:
-            return self._vllm_config.get_name()
+            return cls.quant_config_cls.get_name()
-        raise NotImplementedError("Subclass must initialize _vllm_config or override 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."""

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
+        if metadata is not None:
-        self.stage_type = metadata.stage_type
+            self.stage_id = metadata.stage_id
-        self.engine_output_type = metadata.engine_output_type
+            self.stage_type = metadata.stage_type
-        self.is_comprehension = metadata.is_comprehension
+            self.engine_output_type = metadata.engine_output_type
-        self.requires_multimodal_data = metadata.requires_multimodal_data
+            self.is_comprehension = metadata.is_comprehension
-        self.engine_input_source = metadata.engine_input_source
+            self.requires_multimodal_data = metadata.requires_multimodal_data
-        self.final_output = metadata.final_output
+            self.engine_input_source = metadata.engine_input_source
-        self.final_output_type = metadata.final_output_type
+            self.final_output = metadata.final_output
-        self.default_sampling_params = metadata.default_sampling_params
+            self.final_output_type = metadata.final_output_type
-        self.custom_process_input_func = metadata.custom_process_input_func
+            self.default_sampling_params = metadata.default_sampling_params
-        self.model_stage = metadata.model_stage
+            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

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: Any = None,
-        sampling_params_list: Sequence[OmniSamplingParams] | None = 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()

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
 

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, "
+                "torch.repeat_interleave not executable, switching to vectorized searchsorted implementation."
-                "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(
+        cu_seqlens_np = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum(
-            grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]
+            axis=0, dtype=np.int32
-        ).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,

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)
+        max_ic = max_ic_for_chunk_size(chunk_size)
-        if _ic_cache is None:
+        active = sum(1 for v in transfer_manager.code_prompt_token_ids.values() if len(v) > 0)
-            _ic_cache = {}
+        capacity = getattr(transfer_manager, "scheduler_max_num_seqs", 1)
-            transfer_manager._cached_ic = _ic_cache
+        initial_chunk_size = compute_dynamic_initial_chunk_size(active, capacity, max_ic)
-        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]
 
     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,
     }