modify code based on review suggestions

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>

examples/offline_inference/qwen_2_5_omni/run.sh

@@ -6,4 +6,4 @@ python end2end.py --model Qwen/Qwen2.5-Omni-7B \
                                  --output-wav output_audio \
                                  --prompt_type text \
                                  --init-sleep-seconds 0 \
-                                 --pt-prompts top100.pt
+                                 --pt-prompts <Your_Prompt_File>.pt

vllm_omni/__init__.py

@@ -16,11 +16,11 @@ from . import patch  # noqa: F401
 from .config import OmniModelConfig
 
 # Main entry points
-from .entrypoints.omni_llm import BaseOmniLLM
+from .entrypoints.omni_llm import OmniLLM
 
 __all__ = [
     # Main components
-    "BaseOmniLLM",
+    "OmniLLM",
     # "AsyncOmniLLM", # TODO: add AsyncOmniLLM back when it is implemented
     # Configuration
     "OmniModelConfig",

vllm_omni/entrypoints/omni_llm.py

@@ -1,6 +1,7 @@
 from typing import Any, Optional, Sequence, Union
 import logging
 import multiprocessing as mp
+from concurrent.futures import ThreadPoolExecutor, as_completed
 import os
 import queue
 import sys
@@ -43,74 +44,7 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
-class BaseOmniLLM:
-    def __init__(
-        self, model: str, stage_configs=None, log_stats: bool = False, **kwargs
-    ):
-        if stage_configs is None:
-            self.initialize_stage_configs(model)
-        else:
-            self.stage_configs = stage_configs
-
-        self.stage_list = []
-        self.initialize_stages(model)
-
-    def initialize_stage_configs(self, model: str):
-        self.stage_configs = load_stage_configs_from_model(model)
-
-    def initialize_stages(self, model: str):
-        for stage_config in self.stage_configs:
-            stage = OmniStage(stage_config)
-            stage_llm = OmniStageLLM(model=model, **stage_config.engine_args)
-            stage.set_engine(stage_llm)
-            self.stage_list.append(stage)
-
-    def generate(
-        self,
-        prompts: Union[PromptType, Sequence[PromptType]],
-        sampling_params_list: Optional[
-            Union[SamplingParams, Sequence[SamplingParams]]
-        ] = None,
-    ) -> list[OmniRequestOutput]:
-        """Generate text outputs for the given prompts."""
-        final_outputs: list[OmniRequestOutput] = []
-        if len(sampling_params_list) != len(self.stage_list):
-            raise ValueError(
-                f"Expected {len(self.stage_list)} sampling params, "
-                f"got {len(sampling_params_list)}"
-            )
-        for stage_id, stage in enumerate(self.stage_list):
-            if stage_id > 0:
-                engine_inputs = stage.process_engine_inputs(self.stage_list, prompts)
-            else:
-                engine_inputs = prompts
-            engine_outputs = self._run_generation(
-                stage, sampling_params_list[stage_id], engine_inputs
-            )
-            stage.set_engine_outputs(engine_outputs)
-            if hasattr(stage, "final_output") and stage.final_output:
-                final_outputs.append(
-                    OmniRequestOutput(
-                        stage_id=stage_id,
-                        final_output_type=stage.final_output_type,
-                        request_output=engine_outputs,
-                    )
-                )
-        return final_outputs
-
-    def _run_generation(
-        self,
-        stage: OmniStage,
-        sampling_params: SamplingParams,
-        prompts: Union[PromptType, Sequence[PromptType]],
-    ):
-        engine_outputs = []
-        for ro in stage.engine.generate(prompts, sampling_params):
-            engine_outputs.append(ro)
-        return engine_outputs
-
-
-class OmniLLM(BaseOmniLLM):
+class OmniLLM:
     """Multi-process pipelined OmniLLM.
 
     - Per-stage process with copy-based IPC (Queues)
@@ -135,16 +69,7 @@ class OmniLLM(BaseOmniLLM):
         else:
             self.stage_configs = stage_configs
 
-        self.stage_list: list[OmniStage] = []
-        for stage_config in self.stage_configs:
-            # Only construct lightweight OmniStage without setting engine in parent
-            stage = OmniStage(stage_config)
-            self.stage_list.append(stage)
-        logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
-        self._ctx = mp.get_context("spawn")
-        self._stage_in_queues: list[mp.Queue] = []
-        self._stage_out_queues: list[mp.Queue] = []
         # Optional file handler for orchestrator
         self._log_file = log_file
         if self._log_file:
@@ -152,13 +77,34 @@ class OmniLLM(BaseOmniLLM):
             configure_orchestrator_logger(logger, self._log_file)
 
         self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
+        self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
 
+    def _initialize_stages(self, model: str, init_sleep_seconds: int, shm_threshold_bytes: int, init_timeout: int) -> None:
+        self.stage_list: list[OmniStage] = []
+        # Build OmniStage instances in parallel, preserve original order
+        def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
+            idx, cfg = idx_cfg
+            return idx, OmniStage(cfg)
+
+        with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor:
+            futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)]
+            results: list[tuple[int, OmniStage]] = []
+            for fut in as_completed(futures):
+                results.append(fut.result())
+        results.sort(key=lambda x: x[0])
+        self.stage_list = [st for _, st in results]
+        logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
+
+        self._ctx = mp.get_context("spawn")
+        self._stage_in_queues: list[mp.Queue] = []
+        self._stage_out_queues: list[mp.Queue] = []
         self._init_sleep_seconds = max(0, int(init_sleep_seconds))
         self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
         self._start_stage_processes(model)
         # Wait for all stages to report readiness before seeding
         self._stages_ready: set[int] = set()
         self._wait_for_stages_ready(timeout=init_timeout)
+        
 
     def _start_stage_processes(self, model: str) -> None:
         for stage_id, stage in enumerate(self.stage_list):
@@ -400,7 +346,7 @@ class OmniLLM(BaseOmniLLM):
             # Provide actionable suggestions before shutdown
             try:
                 suggestions = [
-                    "Verify GPU/device assignment in config (mrs.devices) is correct.",
+                    "Verify GPU/device assignment in config (runtime.devices) is correct.",
                     "Check GPU/host memory availability; reduce model or batch size if needed.",
                     "Check model weights path and network reachability (if loading remotely).",
                     "Increase initialization wait time (init_sleep_seconds or call-site timeout).",

vllm_omni/entrypoints/omni_stage.py

@@ -99,11 +99,11 @@ class OmniStage:
         ctx = ctx or mp.get_context("spawn")
         # Prepare lightweight dict config for worker
         engine_args = _to_dict(self.engine_args)
-        mrs_cfg = _to_dict(getattr(self.stage_config, "mrs", {}))
+        runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {}))
         stage_payload: Dict[str, Any] = {
             "stage_id": self.stage_id,
             "engine_args": engine_args,
-            "mrs": mrs_cfg,
+            "runtime": runtime_cfg,
             "shm_threshold_bytes": self._shm_threshold_bytes,
         }
         self._proc = ctx.Process(
@@ -197,7 +197,7 @@ def _stage_worker(
 
     stage_id = stage_payload["stage_id"]
     engine_args = stage_payload.get("engine_args", {})
-    mrs_cfg = stage_payload.get("mrs", {})
+    runtime_cfg = stage_payload.get("runtime", {})
     shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
 
     # Per-stage file logger (optional)
@@ -228,7 +228,7 @@ def _stage_worker(
 
     # Device mapping
     try:
-        set_stage_gpu_devices(stage_id, mrs_cfg.get("devices"))
+        set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
     except Exception as e:
         _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
 
@@ -250,7 +250,7 @@ def _stage_worker(
             _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id)
             break
 
-        max_batch_size = int(mrs_cfg.get("max_batch_size", 1) or 1)
+        max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
         batch_tasks: List[Dict[str, Any]] = [task]
         if max_batch_size > 1:
             while len(batch_tasks) < max_batch_size:

vllm_omni/model_executor/models/qwen2_5_omni.py

@@ -269,13 +269,13 @@ class Qwen2_5OmniForConditionalGeneration(
 
         # 2) Talker (if codec not provided)
         if self.model_stage == "talker":
-            # Mixed-mode support: in一个step内同时支持 Prefill*n 与 Decode*n。
-            # 规则：
-            # - 使用特殊token将 Prefill 片段包裹为 [BOS][PAD...][EOS]
-            # - Decode 片段则为“单一非特殊token”
-            # - 若提供了 additional_information（可为按请求分割的 list 或拼接后的 tensor + 形状列表），
-            #   则对 Prefill 片段按请求重建 thinker→talker 的输入嵌入；
-            # - 对 Decode 片段，若提供了 per-request 的 decode 辅助嵌入（可选），则进行叠加；否则保留原嵌入。
+            # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
+            # Rules:
+            # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
+            # - Decode segments consist of a single non-special token.
+            # - If additional_information is provided (can be a list split by request or a concatenated tensor plus a list of shapes),
+            #   then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments;
+            # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them; otherwise, keep the original embedding.
 
             if input_ids is None and additional_information is None:
                 input_ids = torch.zeros(

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -1,7 +1,7 @@
 # stage config for running qwen2.5-omni with architecture of OmniLLM.
 stage_args:
   - stage_id: 0
-    mrs:
+    runtime:
       process: true            # Run this stage in a separate process
       devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
       max_batch_size: 1
@@ -18,7 +18,7 @@ stage_args:
     final_output: true
     final_output_type: text
   - stage_id: 1
-    mrs:
+    runtime:
       process: true
       devices: "1"
       max_batch_size: 3
@@ -35,7 +35,7 @@ stage_args:
     engine_input_source: [0]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
   - stage_id: 2
-    mrs:
+    runtime:
       process: true
       devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
       max_batch_size: 1
@@ -53,8 +53,8 @@ stage_args:
     final_output: true
     final_output_type: audio
 
-# Top-level MRS config (concise): default windows and stage edges
-mrs:
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
   enabled: true
   defaults:
     window_size: -1             # Simplified: trigger downstream only after full upstream completion