[Entrypoint][Refactor]Stage CLI Refactor (#2020)

Signed-off-by: wuhang <wuhang6@huawei.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

examples/online_serving/bagel/run_server_stage_cli.sh

@@ -1,34 +1,164 @@
 #!/bin/bash
-# Bagel multi-stage online serving startup script
-# Starts stage 0 as master with API server, and stage 1 in headless mode
+# Bagel multi-stage online serving startup script.
+#
+# Usage:
+#   ./run_server_stage_cli.sh --stage 0
+#   ./run_server_stage_cli.sh --stage 1
+#   ./run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2
+#   ./run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9
+#
+# By default, `--stage all` keeps the old behavior and launches both stages in
+# one session. Use `--stage 0` / `--stage 1` to launch each stage separately in
+# different terminal sessions, with stage-specific extra CLI arguments passed
+# after `--`.
+
+set -euo pipefail
 
 MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}"
 PORT="${PORT:-8091}"
 MASTER_ADDRESS="${MASTER_ADDRESS:-127.0.0.1}"
 MASTER_PORT="${MASTER_PORT:-8092}"
-STAGE_CONFIGS_PATH="$(dirname "$0")/../../../vllm_omni/model_executor/stage_configs/bagel.yaml"
+STAGE="all"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+STAGE_CONFIGS_PATH="${STAGE_CONFIGS_PATH:-$SCRIPT_DIR/../../../vllm_omni/model_executor/stage_configs/bagel.yaml}"
+EXTRA_ARGS=()
+
+usage() {
+    cat <<EOF
+Usage: $0 [OPTIONS] [-- EXTRA_VLLM_ARGS...]
+
+Options:
+  --stage {0|1|all}          Stage to launch (default: all)
+  --model MODEL              Model name/path (default: $MODEL)
+  --port PORT                API port for stage 0 (default: $PORT)
+  --master-address ADDRESS   Master/orchestrator address (default: $MASTER_ADDRESS)
+  --master-port PORT         Master/orchestrator port (default: $MASTER_PORT)
+  --stage-configs-path PATH  Stage config YAML path (default: $STAGE_CONFIGS_PATH)
+  --help                     Show this help message
+
+Examples:
+  $0 --stage 0
+  $0 --stage 1
+  $0 --stage 0 -- --tensor-parallel-size 2
+  $0 --stage 1 -- --gpu-memory-utilization 0.9
+
+Notes:
+  - Use different terminal sessions to launch stage 0 and stage 1 separately.
+  - Extra args after '--' are forwarded only to the selected stage.
+  - When using '--stage all', the extra args are forwarded to both stages.
+EOF
+}
+
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --stage)
+            STAGE="$2"
+            shift 2
+            ;;
+        --model)
+            MODEL="$2"
+            shift 2
+            ;;
+        --port)
+            PORT="$2"
+            shift 2
+            ;;
+        --master-address)
+            MASTER_ADDRESS="$2"
+            shift 2
+            ;;
+        --master-port)
+            MASTER_PORT="$2"
+            shift 2
+            ;;
+        --stage-configs-path)
+            STAGE_CONFIGS_PATH="$2"
+            shift 2
+            ;;
+        --help|-h)
+            usage
+            exit 0
+            ;;
+        --)
+            shift
+            EXTRA_ARGS=("$@")
+            break
+            ;;
+        *)
+            echo "Unknown option: $1" >&2
+            usage
+            exit 1
+            ;;
+    esac
+done
+
+if [[ "$STAGE" != "0" && "$STAGE" != "1" && "$STAGE" != "all" ]]; then
+    echo "Invalid --stage value: $STAGE" >&2
+    usage
+    exit 1
+fi
+
+print_config() {
+    echo "Model: $MODEL"
+    echo "API Port: $PORT"
+    echo "Master Address: $MASTER_ADDRESS"
+    echo "Master Port: $MASTER_PORT"
+    echo "Stage Configs: $STAGE_CONFIGS_PATH"
+    echo "Selected Stage: $STAGE"
+    if [[ ${#EXTRA_ARGS[@]} -gt 0 ]]; then
+        echo "Extra Args: ${EXTRA_ARGS[*]}"
+    fi
+}
+
+run_stage_0() {
+    echo "Starting Stage 0 (Thinker) as master..."
+    vllm serve "$MODEL" --omni \
+        --port "$PORT" \
+        --stage-configs-path "$STAGE_CONFIGS_PATH" \
+        --stage-id 0 \
+        -oma "$MASTER_ADDRESS" \
+        -omp "$MASTER_PORT" \
+        "${EXTRA_ARGS[@]}"
+}
+
+run_stage_1() {
+    echo "Starting Stage 1 (DiT) in headless mode..."
+    vllm serve "$MODEL" --omni \
+        --stage-configs-path "$STAGE_CONFIGS_PATH" \
+        --stage-id 1 \
+        --headless \
+        -oma "$MASTER_ADDRESS" \
+        -omp "$MASTER_PORT" \
+        "${EXTRA_ARGS[@]}"
+}
 
 echo "Starting Bagel multi-stage server..."
-echo "Model: $MODEL"
-echo "API Port: $PORT"
-echo "Master Address: $MASTER_ADDRESS"
-echo "Master Port: $MASTER_PORT"
-echo "Stage Configs: $STAGE_CONFIGS_PATH"
+print_config
 
-# Start stage 1 (DiT) in headless mode first
-echo "Starting Stage 1 (DiT) in headless mode..."
-vllm serve "$MODEL" --omni \
-    --stage-configs-path "$STAGE_CONFIGS_PATH" \
-    --stage-id 1 \
-    --headless \
-    -oma "$MASTER_ADDRESS" \
-    -omp "$MASTER_PORT" &
+case "$STAGE" in
+    0)
+        run_stage_0
+        ;;
+    1)
+        run_stage_1
+        ;;
+    all)
+        echo "Launching both stages in one session (legacy mode)..."
+        echo "Starting Stage 0 (Thinker) in background first..."
+        run_stage_0 &
+        STAGE_0_PID=$!
 
-# Start stage 0 (Thinker) as master with API server
-echo "Starting Stage 0 (Thinker) as master..."
-vllm serve "$MODEL" --omni \
-    --port "$PORT" \
-    --stage-configs-path "$STAGE_CONFIGS_PATH" \
-    --stage-id 0 \
-    -oma "$MASTER_ADDRESS" \
-    -omp "$MASTER_PORT"
+        cleanup() {
+            if [[ -n "${STAGE_0_PID:-}" ]]; then
+                kill "$STAGE_0_PID" 2>/dev/null || true
+                wait "$STAGE_0_PID" 2>/dev/null || true
+            fi
+        }
+
+        trap cleanup EXIT INT TERM
+
+        echo "Waiting briefly for Stage 0 to initialize..."
+        sleep 2
+        run_stage_1
+        ;;
+esac

tests/conftest.py

@@ -73,6 +73,8 @@ class OmniServerParams(NamedTuple):
     server_args: list[str] | None = None
     env_dict: dict[str, str] | None = None
     use_omni: bool = True
+    use_stage_cli: bool = False
+    init_timeout: int | None = None
 
 
 def assert_image_diffusion_response(
@@ -1546,6 +1548,183 @@ class OmniServer:
         cleanup_dist_env_and_memory()
 
 
+class OmniServerStageCli(OmniServer):
+    """Omni server harness that exercises the stage CLI flow."""
+
+    def __init__(
+        self,
+        model: str,
+        stage_config_path: str,
+        serve_args: list[str] | None = None,
+        *,
+        stage_ids: list[int] | None = None,
+        port: int | None = None,
+        env_dict: dict[str, str] | None = None,
+    ) -> None:
+        super().__init__(model, serve_args or [], port=port, env_dict=env_dict, use_omni=True)
+        self.stage_config_path = stage_config_path
+        self.master_port = get_open_port()
+        self.visible_device_list = self._load_visible_device_list(env_dict)
+        self.stage_runtime_devices = self._load_stage_runtime_devices(stage_config_path)
+        self.stage_ids = stage_ids or self._load_stage_ids(stage_config_path)
+        if 0 not in self.stage_ids:
+            raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}")
+        self.stage_procs: dict[int, subprocess.Popen] = {}
+        self.proc = None
+
+    @staticmethod
+    def _load_stage_ids(stage_config_path: str) -> list[int]:
+        with open(stage_config_path, encoding="utf-8") as f:
+            cfg = yaml.safe_load(f) or {}
+
+        stage_ids = [stage["stage_id"] for stage in cfg.get("stage_args", []) if "stage_id" in stage]
+        if not stage_ids:
+            raise ValueError(f"No stage IDs found in config: {stage_config_path}")
+        return stage_ids
+
+    @staticmethod
+    def _load_stage_runtime_devices(stage_config_path: str) -> dict[int, str]:
+        with open(stage_config_path, encoding="utf-8") as f:
+            cfg = yaml.safe_load(f) or {}
+
+        runtime_devices: dict[int, str] = {}
+        for stage in cfg.get("stage_args", []):
+            stage_id = stage.get("stage_id")
+            devices = stage.get("runtime", {}).get("devices")
+            if stage_id is not None and devices:
+                runtime_devices[int(stage_id)] = str(devices)
+        return runtime_devices
+
+    @classmethod
+    def _parse_device_list(cls, devices: str | int) -> list[str]:
+        if isinstance(devices, int):
+            if devices < 0:
+                raise ValueError("Device IDs must be non-negative integers")
+            return [str(devices)]
+        return [token.strip() for token in str(devices).split(",") if token.strip()]
+
+    @classmethod
+    def _load_visible_device_list(cls, env_dict: dict[str, str] | None) -> list[str] | None:
+        env = os.environ.copy()
+        if env_dict is not None:
+            env.update(env_dict)
+
+        env_var = getattr(current_omni_platform, "device_control_env_var", None)
+        if env_var and env_var in env:
+            return [token.strip() for token in env[env_var].split(",") if token.strip()]
+        return None
+
+    @classmethod
+    def _map_stage_devices(cls, stage_id: int, visible_device_list: list[str] | None, devices: str) -> str:
+        device_list = cls._parse_device_list(devices)
+
+        if visible_device_list is None:
+            return ",".join(device_list)
+
+        if not all(device.isdigit() for device in device_list):
+            raise ValueError("Logical devices must be non-negative integers")
+
+        logical_ids = [int(device) for device in device_list]
+        if logical_ids and max(logical_ids) >= len(visible_device_list):
+            raise ValueError(
+                f"Stage {stage_id} has logical IDs {device_list}, one or more of which exceed the number of visible devices"
+            )
+
+        return ",".join(visible_device_list[idx] for idx in logical_ids)
+
+    def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str) -> None:
+        mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, devices)
+        env_var = getattr(current_omni_platform, "device_control_env_var", None)
+        if env_var:
+            env[env_var] = mapped_devices
+
+    def _build_stage_cmd(self, stage_id: int, *, headless: bool) -> list[str]:
+        cmd = [
+            sys.executable,
+            "-m",
+            "vllm_omni.entrypoints.cli.main",
+            "serve",
+            self.model,
+            "--omni",
+            "--stage-configs-path",
+            self.stage_config_path,
+            "--stage-id",
+            str(stage_id),
+            "--omni-master-address",
+            self.host,
+            "--omni-master-port",
+            str(self.master_port),
+        ]
+
+        if headless:
+            cmd.append("--headless")
+        else:
+            cmd += ["--host", self.host, "--port", str(self.port)]
+
+        cmd += self.serve_args
+        return cmd
+
+    def _launch_stage(self, stage_id: int, *, headless: bool) -> None:
+        env = os.environ.copy()
+        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        if self.env_dict is not None:
+            env.update(self.env_dict)
+
+        devices = self.stage_runtime_devices.get(stage_id)
+        if devices:
+            self._set_stage_device_env(stage_id, env, devices)
+
+        cmd = self._build_stage_cmd(stage_id, headless=headless)
+        print(f"Launching OmniServerStageCli stage {stage_id}: {' '.join(cmd)}")
+        proc = subprocess.Popen(
+            cmd,
+            env=env,
+            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        )
+        self.stage_procs[stage_id] = proc
+        if stage_id == 0:
+            self.proc = proc
+
+    def _ensure_stage_processes_alive(self) -> None:
+        for stage_id, proc in self.stage_procs.items():
+            ret = proc.poll()
+            if ret is not None:
+                raise RuntimeError(f"Stage {stage_id} exited with code {ret} before API server became ready.")
+
+    def _start_server(self) -> None:
+        ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]]
+
+        self._launch_stage(0, headless=False)
+        time.sleep(2)
+        self._ensure_stage_processes_alive()
+
+        for stage_id in ordered_stage_ids[1:]:
+            self._launch_stage(stage_id, headless=True)
+
+        max_wait = 1200
+        start_time = time.time()
+        while time.time() - start_time < max_wait:
+            self._ensure_stage_processes_alive()
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+                sock.settimeout(1)
+                result = sock.connect_ex((self.host, self.port))
+                if result == 0:
+                    print(f"OmniServerStageCli ready on {self.host}:{self.port}")
+                    return
+            time.sleep(2)
+
+        raise RuntimeError(f"OmniServerStageCli failed to start within {max_wait} seconds")
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        for stage_id in sorted(self.stage_procs, reverse=True):
+            proc = self.stage_procs[stage_id]
+            if proc.poll() is None:
+                self._kill_process_tree(proc.pid)
+        _run_pre_test_cleanup(enable_force=True)
+        _run_post_test_cleanup(enable_force=True)
+        cleanup_dist_env_and_memory()
+
+
 def pytest_addoption(parser):
     parser.addoption(
         "--run-level",
@@ -1568,9 +1747,11 @@ _omni_server_lock = threading.Lock()
 
 @pytest.fixture(scope="module")
 def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]:
-    """Start vLLM-Omni server as a subprocess with actual model weights.
-    Uses session scope so the server starts only once for the entire test session.
-    Multi-stage initialization can take 10-20+ minutes.
+    """Start vLLM-Omni through the standard or stage-CLI launcher.
+
+    The fixture stays module-scoped because multi-stage initialization is costly.
+    The ``use_stage_cli`` flag on ``OmniServerParams`` routes the setup through the
+    stage-CLI harness while still reusing the same fixture grouping semantics.
     """
     with _omni_server_lock:
         params: OmniServerParams = request.param
@@ -1589,28 +1770,47 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st
         server_args = params.server_args or []
         if params.use_omni:
             server_args = ["--stage-init-timeout", "120", *server_args]
-        if stage_config_path is not None:
-            server_args += ["--stage-configs-path", stage_config_path]
+        if params.init_timeout is not None:
+            server_args = [*server_args, "--init-timeout", str(params.init_timeout)]
+        if params.use_stage_cli:
+            if not params.use_omni:
+                raise ValueError("omni_server with use_stage_cli=True requires use_omni=True")
+            if stage_config_path is None:
+                raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path")
 
-        with (
-            OmniServer(
+            with OmniServerStageCli(
                 model,
+                stage_config_path,
                 server_args,
                 port=port,
                 env_dict=params.env_dict,
-                use_omni=params.use_omni,
-            )
-            if port
-            else OmniServer(
-                model,
-                server_args,
-                env_dict=params.env_dict,
-                use_omni=params.use_omni,
-            )
-        ) as server:
-            print("OmniServer started successfully")
-            yield server
-            print("OmniServer stopping...")
+            ) as server:
+                print("OmniServer started successfully")
+                yield server
+                print("OmniServer stopping...")
+        else:
+            if stage_config_path is not None:
+                server_args += ["--stage-configs-path", stage_config_path]
+
+            with (
+                OmniServer(
+                    model,
+                    server_args,
+                    port=port,
+                    env_dict=params.env_dict,
+                    use_omni=params.use_omni,
+                )
+                if port
+                else OmniServer(
+                    model,
+                    server_args,
+                    env_dict=params.env_dict,
+                    use_omni=params.use_omni,
+                )
+            ) as server:
+                print("OmniServer started successfully")
+                yield server
+                print("OmniServer stopping...")
 
         print("OmniServer stopped")
 
@@ -2653,10 +2853,11 @@ class OpenAIClientHandler:
 
 
 @pytest.fixture
-def openai_client(omni_server: OmniServer, run_level: str):
+def openai_client(request: pytest.FixtureRequest, run_level: str):
     """Create OpenAIClientHandler fixture to facilitate communication with OmniServer
     with encapsulated request sending, concurrent requests, response handling, and validation."""
-    return OpenAIClientHandler(host=omni_server.host, port=omni_server.port, api_key="EMPTY", run_level=run_level)
+    server = request.getfixturevalue("omni_server")
+    return OpenAIClientHandler(host=server.host, port=server.port, api_key="EMPTY", run_level=run_level)
 
 
 class OmniRunner:

tests/dfx/perf/scripts/run_benchmark.py

@@ -43,7 +43,7 @@ def omni_server(request):
 
         print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-        server_args = ["--stage-init-timeout", "120"]
+        server_args = ["--stage-init-timeout", "120", "--init-timeout", "900"]
         if stage_config_path:
             server_args = ["--stage-configs-path", stage_config_path] + server_args
         with OmniServer(model, server_args) as server:

tests/e2e/online_serving/test_flux_2_dev_expansion.py

@@ -2,13 +2,11 @@
 End-to-end diffusion coverage for FLUX.2-dev in online serving mode.
 
 Coverage:
-- Cache-DiT cache acceleration backend
 - CPU offload
 
-This test verifies that FLUX.2-dev can be launched with the Cache-DiT backend
-and CPU offload enabled, accepts text-to-image requests through the
-OpenAI-compatible API, and returns valid generated images with the requested
-resolution.
+This test verifies that FLUX.2-dev can be launched with CPU offload enabled,
+accepts text-to-image requests through the OpenAI-compatible API, and returns
+valid generated images with the requested resolution.
 
 assert_diffusion_response validates successful generation and the expected
 image resolution.
@@ -33,19 +31,17 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2)
 
 
 def _get_flux_2_dev_feature_cases(model: str):
-    """Return FLUX.2-dev diffusion feature cases for Cache-DiT + CPU offload."""
+    """Return FLUX.2-dev diffusion feature cases for CPU offload."""
 
     return [
         pytest.param(
             OmniServerParams(
                 model=model,
                 server_args=[
-                    "--cache-backend",
-                    "cache_dit",
                     "--enable-cpu-offload",
                 ],
             ),
-            id="cache_dit_cpu_offload",
+            id="cpu_offload",
             marks=SINGLE_CARD_FEATURE_MARKS,
         ),
         pytest.param(
@@ -76,7 +72,7 @@ def test_flux_2_dev(
     omni_server: OmniServer,
     openai_client: OpenAIClientHandler,
 ):
-    """Validate FLUX.2-dev online serving with Cache-DiT and CPU offload."""
+    """Validate FLUX.2-dev online serving with CPU offload."""
 
     messages = dummy_messages_from_mix_data(content_text=PROMPT)
 

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -67,13 +67,17 @@ if current_omni_platform.is_xpu():
 
 # Create parameter combinations for model and stage config
 test_params = [
-    pytest.param(OmniServerParams(model=model, stage_config_path=default_path), id="default"),
-    pytest.param(OmniServerParams(model=model, stage_config_path=get_chunk_config(default_path)), id="async_chunk"),
+    pytest.param(OmniServerParams(model=model, stage_config_path=default_path, use_stage_cli=True), id="default"),
+    pytest.param(
+        OmniServerParams(model=model, stage_config_path=get_chunk_config(default_path), use_stage_cli=True),
+        id="async_chunk",
+    ),
 ]
 
 test_token_params = [
     pytest.param(
-        OmniServerParams(model=model, stage_config_path=get_batch_token_config(default_path)), id="batch_token_64"
+        OmniServerParams(model=model, stage_config_path=get_batch_token_config(default_path), use_stage_cli=True),
+        id="batch_token_64",
     )
 ]
 

tests/engine/test_arg_utils.py

@@ -24,6 +24,18 @@ def test_sync_config_is_omni():
     assert isinstance(cfg, OmniModelConfig)
 
 
+def test_default_stage_id_is_concrete_int():
+    """Ensure `stage_id` stays safe for downstream arithmetic/indexing."""
+    engine_args = OmniEngineArgs()
+
+    assert engine_args.stage_id == 0
+    assert isinstance(engine_args.stage_id, int)
+    assert engine_args.log_stats is False
+
+    cfg = engine_args.create_model_config()
+    assert cfg.stage_id == 0
+
+
 def test_multimodal_kwarg_overrides():
     """Ensure that overrides in the multimodal config are preserved."""
     # Get a different value than the default for a multimodal field

tests/engine/test_async_omni_engine_stage_init.py

@@ -1,3 +1,4 @@
+import importlib
 import os
 import types
 
@@ -8,6 +9,17 @@ from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
+    """Regression test for forward references in make_async_mp_client."""
+    import vllm_omni.engine.stage_engine_core_client as client_mod
+
+    importlib.reload(client_mod)
+
+    assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
+        "StageEngineCoreClient | DPLBStageEngineCoreClient"
+    )
+
+
 def test_initialize_stages_restores_device_visibility_after_diffusion_init(monkeypatch):
     """Regression test for stage device env leakage across stage init.
 
@@ -23,6 +35,9 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke
     engine.num_stages = 1
     engine.async_chunk = False
     engine.diffusion_batch_size = 1
+    engine.single_stage_mode = False
+    engine._single_stage_id_filter = None
+    engine._omni_master_server = None
     engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion")]
 
     env_var = current_omni_platform.device_control_env_var
@@ -49,7 +64,7 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke
         current_omni_platform.set_device_control_env_var("1")
 
     monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices)
-    monkeypatch.setattr(engine_mod, "_inject_kv_stage_info", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
     monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: diffusion_client)
     monkeypatch.setattr(
         engine_mod,
@@ -101,7 +116,11 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
             self.vllm_config = vllm_config
             self.renderer = renderer
 
-    monkeypatch.setattr(engine_mod, "StageEngineCoreClient", DummyStageEngineCoreClient)
+    monkeypatch.setattr(
+        engine_mod.StageEngineCoreClientBase,
+        "make_async_mp_client",
+        staticmethod(lambda **kwargs: DummyStageEngineCoreClient(**kwargs)),
+    )
     monkeypatch.setattr(engine_mod, "MultimodalOutputProcessor", DummyOutputProcessor)
     monkeypatch.setattr(engine_mod, "InputProcessor", DummyInputProcessor)
     monkeypatch.setattr(engine_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor)

tests/entrypoints/test_serve.py

@@ -0,0 +1,195 @@
+"""Unit tests for the Omni serve CLI helpers."""
+
+from __future__ import annotations
+
+import argparse
+from unittest.mock import Mock, patch
+
+import pytest
+
+from vllm_omni.entrypoints.cli.serve import run_headless
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_headless_args() -> argparse.Namespace:
+    return argparse.Namespace(
+        model="fake-model",
+        stage_id=3,
+        omni_master_address="127.0.0.1",
+        omni_master_port=26000,
+        api_server_count=0,
+        worker_backend="multi_process",
+        stage_configs_path=None,
+        log_stats=False,
+        disable_log_stats=False,
+    )
+
+
+def test_run_headless_registers_stage_once_and_launches_all_local_engines() -> None:
+    args = _make_headless_args()
+    stage_cfg = Mock(stage_id=3)
+    stage_cfgs = [stage_cfg]
+    parallel_config = Mock(
+        data_parallel_size_local=2,
+        data_parallel_rank=4,
+        data_parallel_rank_local=1,
+        node_rank_within_dp=0,
+    )
+    vllm_config = Mock(parallel_config=parallel_config)
+    executor_class = Mock()
+    engine_manager = Mock()
+
+    with (
+        patch(
+            "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+            return_value=("/fake/stages.yaml", stage_cfgs),
+        ),
+        patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"),
+        patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()),
+        patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}),
+        patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}),
+        patch(
+            "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+            return_value=(None, None, None),
+        ),
+        patch(
+            "vllm_omni.engine.stage_init_utils.build_vllm_config",
+            return_value=(vllm_config, executor_class),
+        ) as mock_build_vllm_config,
+        patch(
+            "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+            return_value="tcp://127.0.0.1:26001",
+        ) as mock_register,
+        patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) as mock_manager_cls,
+        patch("signal.signal"),
+    ):
+        run_headless(args)
+
+    mock_build_vllm_config.assert_called_once_with(
+        stage_cfg,
+        "fake-model",
+        stage_connector_spec={},
+        engine_args_dict={},
+        headless=True,
+    )
+    mock_register.assert_called_once_with(
+        omni_master_address="127.0.0.1",
+        omni_master_port=26000,
+        omni_stage_id=3,
+        omni_stage_config=stage_cfg,
+        coordinator=None,
+    )
+    mock_manager_cls.assert_called_once()
+    manager_kwargs = mock_manager_cls.call_args.kwargs
+    assert manager_kwargs["local_engine_count"] == 2
+    assert manager_kwargs["start_index"] == 4
+    assert manager_kwargs["local_start_index"] == 0
+    assert manager_kwargs["local_client"] is False
+    assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
+    assert manager_kwargs["log_stats"] is False
+    engine_manager.join_first.assert_called_once_with()
+    engine_manager.shutdown.assert_called_once_with()
+
+
+def test_run_headless_honors_explicit_log_stats_flag() -> None:
+    args = _make_headless_args()
+    args.log_stats = True
+    stage_cfg = Mock(stage_id=3)
+    stage_cfgs = [stage_cfg]
+    parallel_config = Mock(
+        data_parallel_size_local=2,
+        data_parallel_rank=4,
+        data_parallel_rank_local=1,
+        node_rank_within_dp=0,
+    )
+    vllm_config = Mock(parallel_config=parallel_config)
+    executor_class = Mock()
+    engine_manager = Mock()
+
+    with (
+        patch(
+            "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+            return_value=("/fake/stages.yaml", stage_cfgs),
+        ),
+        patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"),
+        patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()),
+        patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}),
+        patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}),
+        patch(
+            "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+            return_value=(None, None, None),
+        ),
+        patch(
+            "vllm_omni.engine.stage_init_utils.build_vllm_config",
+            return_value=(vllm_config, executor_class),
+        ),
+        patch(
+            "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+            return_value="tcp://127.0.0.1:26001",
+        ),
+        patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) as mock_manager_cls,
+        patch("signal.signal"),
+    ):
+        run_headless(args)
+
+    manager_kwargs = mock_manager_cls.call_args.kwargs
+    assert manager_kwargs["log_stats"] is True
+
+
+def test_run_headless_launches_diffusion_stage_via_omni_master() -> None:
+    args = _make_headless_args()
+    stage_cfg = Mock(stage_id=3, stage_type="diffusion")
+    stage_cfg.engine_args = Mock()
+    stage_cfg.engine_input_source = []
+    stage_cfgs = [stage_cfg]
+    metadata = Mock(stage_id=3)
+    od_config = Mock()
+    proc = Mock()
+    proc.exitcode = 0
+    proc.is_alive.return_value = False
+
+    with (
+        patch(
+            "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+            return_value=("/fake/stages.yaml", stage_cfgs),
+        ),
+        patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"),
+        patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()),
+        patch(
+            "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+            return_value=(None, None, None),
+        ),
+        patch("vllm_omni.engine.stage_init_utils.extract_stage_metadata", return_value=metadata),
+        patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info") as mock_inject_stage_info,
+        patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=od_config),
+        patch(
+            "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+            return_value=("tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"),
+        ) as mock_register,
+        patch(
+            "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc",
+            return_value=(proc, "tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"),
+        ) as mock_spawn,
+        patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") as mock_handshake,
+        patch("signal.signal"),
+    ):
+        run_headless(args)
+
+    mock_inject_stage_info.assert_called_once_with(stage_cfg, 3)
+    mock_register.assert_called_once_with(
+        omni_master_address="127.0.0.1",
+        omni_master_port=26000,
+        omni_stage_id=3,
+        omni_stage_config=stage_cfg,
+        return_addresses=True,
+    )
+    mock_spawn.assert_called_once_with(
+        "fake-model",
+        od_config,
+        handshake_address="tcp://127.0.0.1:26001",
+        request_address="tcp://127.0.0.1:26002",
+        response_address="tcp://127.0.0.1:26003",
+    )
+    mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001")
+    proc.join.assert_called_once_with()

tests/entrypoints/test_utils.py

@@ -8,6 +8,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 
+from vllm_omni.config.yaml_util import create_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -16,6 +17,7 @@ from vllm_omni.entrypoints.utils import (
     _filter_dict_like_object,
     filter_dataclass_kwargs,
     load_and_resolve_stage_configs,
+    load_stage_configs_from_yaml,
     resolve_model_config_path,
 )
 
@@ -322,3 +324,70 @@ class TestLoadAndResolveStageConfigs:
         assert config_path is None
         assert len(stage_configs) == 1
         assert "dtype" in stage_configs[0]["engine_args"]
+
+
+class TestLoadStageConfigsFromYaml:
+    """Regression tests for stage-config loading and merging."""
+
+    def test_deep_merges_stage_engine_args(self, mocker: MockerFixture):
+        yaml_config = create_config(
+            {
+                "async_chunk": True,
+                "stage_args": [
+                    {
+                        "stage_id": 0,
+                        "runtime": {"device": 0},
+                        "engine_args": {
+                            "parallel_config": {"tensor_parallel_size": 4},
+                        },
+                    }
+                ],
+            }
+        )
+        mocker.patch(
+            "vllm_omni.entrypoints.utils.load_yaml_config",
+            return_value=yaml_config,
+        )
+
+        stages = load_stage_configs_from_yaml(
+            "fake.yaml",
+            base_engine_args={
+                "parallel_config": {
+                    "tensor_parallel_size": 1,
+                    "pipeline_parallel_size": 2,
+                },
+                "model": "base-model",
+            },
+        )
+
+        merged_engine_args = stages[0]["engine_args"]
+        assert merged_engine_args["parallel_config"]["tensor_parallel_size"] == 4
+        assert merged_engine_args["parallel_config"]["pipeline_parallel_size"] == 2
+        assert merged_engine_args["model"] == "base-model"
+        assert merged_engine_args["async_chunk"] is True
+
+    def test_merges_nested_stage_engine_args(self, mocker: MockerFixture):
+        yaml_config = create_config(
+            {
+                "stage_args": [
+                    {
+                        "stage_id": 0,
+                        "engine_args": {
+                            "nested": {"override": 2},
+                        },
+                    }
+                ],
+            }
+        )
+        mocker.patch(
+            "vllm_omni.entrypoints.utils.load_yaml_config",
+            return_value=yaml_config,
+        )
+
+        stages = load_stage_configs_from_yaml(
+            "fake.yaml",
+            base_engine_args={"nested": {"base": 1}},
+        )
+
+        assert stages[0]["engine_args"]["nested"]["base"] == 1
+        assert stages[0]["engine_args"]["nested"]["override"] == 2

vllm_omni/diffusion/stage_diffusion_client.py

@@ -51,6 +51,41 @@ class StageDiffusionClient:
         od_config: OmniDiffusionConfig,
         metadata: StageMetadata,
         batch_size: int = 1,
+    ) -> None:
+        # Spawn StageDiffusionProc subprocess and wait for READY.
+        proc, handshake_address, request_address, response_address = spawn_diffusion_proc(model, od_config)
+        complete_diffusion_handshake(proc, handshake_address)
+        self._initialize_client(metadata, request_address, response_address, proc=proc, batch_size=batch_size)
+
+    @classmethod
+    def from_addresses(
+        cls,
+        metadata: StageMetadata,
+        request_address: str,
+        response_address: str,
+        *,
+        proc: Any = None,
+        batch_size: int = 1,
+    ) -> StageDiffusionClient:
+        """Create a client for an already-running diffusion subprocess."""
+        client = cls.__new__(cls)
+        client._initialize_client(
+            metadata,
+            request_address,
+            response_address,
+            proc=proc,
+            batch_size=batch_size,
+        )
+        return client
+
+    def _initialize_client(
+        self,
+        metadata: StageMetadata,
+        request_address: str,
+        response_address: str,
+        *,
+        proc: Any,
+        batch_size: int,
     ) -> None:
         self.stage_id = metadata.stage_id
         self.final_output = metadata.final_output
@@ -58,13 +93,9 @@ class StageDiffusionClient:
         self.default_sampling_params = metadata.default_sampling_params
         self.custom_process_input_func = metadata.custom_process_input_func
         self.engine_input_source = metadata.engine_input_source
-
-        # Spawn StageDiffusionProc subprocess and wait for READY.
-        proc, handshake_address, request_address, response_address = spawn_diffusion_proc(model, od_config)
-        complete_diffusion_handshake(proc, handshake_address)
         self._proc = proc
+        self._owns_process = proc is not None
 
-        # ZMQ sockets (sync) for communicating with the subprocess.
         self._zmq_ctx = zmq.Context()
         self._request_socket = self._zmq_ctx.socket(zmq.PUSH)
         self._request_socket.connect(request_address)
@@ -74,14 +105,18 @@ class StageDiffusionClient:
         self._encoder = OmniMsgpackEncoder()
         self._decoder = OmniMsgpackDecoder()
 
-        # Buffers for demultiplexing response messages.
         self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
         self._rpc_results: dict[str, Any] = {}
         self._pending_rpcs: set[str] = set()
         self._tasks: dict[str, asyncio.Task] = {}
         self._shutting_down = False
 
-        logger.info("[StageDiffusionClient] Stage-%s initialized (batch_size=%d)", self.stage_id, batch_size)
+        logger.info(
+            "[StageDiffusionClient] Stage-%s initialized (owns_process=%s, batch_size=%d)",
+            self.stage_id,
+            self._owns_process,
+            batch_size,
+        )
 
     # ------------------------------------------------------------------
     # Internal helpers
@@ -253,7 +288,7 @@ class StageDiffusionClient:
         try:
             return self._output_queue.get_nowait()
         except asyncio.QueueEmpty:
-            if not self._shutting_down and self._proc is not None and not self._proc.is_alive():
+            if not self._shutting_down and self._owns_process and self._proc is not None and not self._proc.is_alive():
                 exitcode = self._proc.exitcode
                 # One final drain – the last ZMQ frame may have arrived
                 # between the first drain and the is_alive() check.
@@ -325,7 +360,7 @@ class StageDiffusionClient:
                 self._drain_responses()
                 if rpc_id in self._rpc_results:
                     return self._rpc_results.pop(rpc_id)
-                if self._proc is not None and not self._proc.is_alive():
+                if self._owns_process and self._proc is not None and not self._proc.is_alive():
                     raise RuntimeError(
                         f"StageDiffusionProc died while waiting for "
                         f"collective_rpc '{method}' (exit code {self._proc.exitcode})"
@@ -343,7 +378,7 @@ class StageDiffusionClient:
         except Exception:
             pass
 
-        if self._proc is not None and self._proc.is_alive():
+        if self._owns_process and self._proc is not None and self._proc.is_alive():
             self._proc.join(timeout=10)
             terminate_alive_proc(self._proc)
 

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -580,14 +580,17 @@ class StageDiffusionProc:
 def spawn_diffusion_proc(
     model: str,
     od_config: OmniDiffusionConfig,
+    handshake_address: str | None = None,
+    request_address: str | None = None,
+    response_address: str | None = None,
 ) -> tuple[BaseProcess, str, str, str]:
     """Spawn a StageDiffusionProc subprocess.
 
     Returns ``(proc, handshake_address, request_address, response_address)``.
     """
-    handshake_address = get_open_zmq_ipc_path()
-    request_address = get_open_zmq_ipc_path()
-    response_address = get_open_zmq_ipc_path()
+    handshake_address = handshake_address or get_open_zmq_ipc_path()
+    request_address = request_address or get_open_zmq_ipc_path()
+    response_address = response_address or get_open_zmq_ipc_path()
 
     ctx = get_mp_context()
     proc = ctx.Process(

vllm_omni/engine/arg_utils.py

@@ -86,7 +86,11 @@ class OmniEngineArgs(EngineArgs):
     Adds omni-specific configuration fields for multi-stage pipeline
     processing and output type specification.
     Args:
-        stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
+        stage_id: Identifier for the stage in a multi-stage pipeline.
+            Defaults to 0 for per-stage engine construction. The CLI-level
+            single-stage selector remains optional on the parsed argparse
+            namespace and should not be forwarded as a nullable per-stage
+            engine argument.
         model_stage: Stage type identifier, e.g., "thinker" or "talker"
             (default: "thinker")
         model_arch: Model architecture name
@@ -105,6 +109,18 @@ class OmniEngineArgs(EngineArgs):
         worker_type: Model Type, e.g., "ar" or "generation"
         task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base).
             If not specified, will be inferred from model path.
+        omni_master_address: TCP address that the OmniMasterServer (running
+            inside AsyncOmniEngine) listens on for engine core registrations.
+            Required when single-stage mode is active.
+        omni_master_port: TCP port for the OmniMasterServer registration
+            socket.  Required when single-stage mode is active.
+        stage_configs_path: Optional path to a JSON/YAML file containing
+            stage configurations for the multi-stage pipeline. If None,
+            stage configs are resolved from the model's default configuration.
+        output_modalities: Optional list of output modality names to enable
+            (e.g. ["text", "audio"]). If None, all modalities supported by
+            the model are used.
+        log_stats: Whether to log engine statistics. Defaults to False.
     """
 
     stage_id: int = 0
@@ -119,6 +135,11 @@ class OmniEngineArgs(EngineArgs):
     quantization_config: Any | None = None
     worker_type: str | None = None
     task_type: str | None = None
+    omni_master_address: str | None = None
+    omni_master_port: int | None = None
+    stage_configs_path: str | None = None
+    output_modalities: list[str] | None = None
+    log_stats: bool = False
 
     def __post_init__(self) -> None:
         load_omni_general_plugins()

vllm_omni/engine/async_omni_engine.py

@@ -18,12 +18,10 @@ import time
 import uuid
 import weakref
 from collections.abc import Mapping, Sequence
+from contextlib import ExitStack
 from dataclasses import asdict
 from typing import TYPE_CHECKING, Any
 
-if TYPE_CHECKING:
-    from vllm_omni.engine.arg_utils import OmniEngineArgs
-
 import janus
 import torch
 from omegaconf import OmegaConf
@@ -34,26 +32,36 @@ from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.input_processor import InputProcessor
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+from vllm_omni.diffusion.stage_diffusion_proc import (
+    complete_diffusion_handshake,
+    spawn_diffusion_proc,
+)
 from vllm_omni.distributed.omni_connectors.utils.initialization import (
     resolve_omni_kv_config_for_stage,
 )
-from vllm_omni.engine import (
-    OmniEngineCoreRequest,
-)
+from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.orchestrator import Orchestrator
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.serialization import (
     deserialize_additional_information,
     serialize_additional_information,
 )
-from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
 from vllm_omni.engine.stage_engine_core_proc import (
     complete_stage_handshake,
     spawn_stage_core,
 )
+from vllm_omni.engine.stage_engine_startup import (
+    OmniMasterServer,
+    connect_remote_engine_cores,
+    launch_omni_core_engines,
+    register_stage_with_omni_master,
+)
 from vllm_omni.engine.stage_init_utils import (
     StartedLlmStage,
     acquire_device_locks,
+    build_diffusion_config,
     build_engine_args_dict,
     build_vllm_config,
     cleanup_failed_stage_initialization,
@@ -62,17 +70,20 @@ from vllm_omni.engine.stage_init_utils import (
     finalize_initialized_stages,
     get_stage_connector_spec,
     initialize_diffusion_stage,
+    inject_kv_stage_info,
     load_omni_transfer_config_for_model,
     prepare_engine_environment,
     release_device_locks,
     setup_stage_devices,
+    terminate_alive_proc,
 )
-from vllm_omni.entrypoints.utils import (
-    load_and_resolve_stage_configs,
-)
+from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs
 from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 from vllm_omni.platforms import current_omni_platform
 
+if TYPE_CHECKING:
+    from vllm_omni.engine.arg_utils import OmniEngineArgs
+
 logger = init_logger(__name__)
 
 
@@ -86,39 +97,6 @@ def _patch_generation_config_if_needed(model_config: Any) -> None:
         model_config.try_get_generation_config = lambda: {}
 
 
-def _inject_kv_stage_info(stage_cfg: Any, stage_id: int) -> None:
-    """Inject stage_id and engine_input_source into omni_kv_config.
-
-    OmniKVTransferManager needs stage_id to compute recv_stages for the
-    receiving side. In the old Omni architecture, OmniDiffusion.__init__
-    performed this injection; replicate it here for AsyncOmniEngine.
-    """
-    try:
-        engine_args = stage_cfg.engine_args
-        if hasattr(engine_args, "get"):
-            omni_kv = engine_args.get("omni_kv_config", None)
-        else:
-            omni_kv = getattr(engine_args, "omni_kv_config", None)
-
-        if omni_kv is None:
-            return
-
-        if hasattr(omni_kv, "setdefault"):
-            omni_kv.setdefault("stage_id", stage_id)
-        elif hasattr(omni_kv, "__setitem__"):
-            if "stage_id" not in omni_kv:
-                omni_kv["stage_id"] = stage_id
-
-        engine_input_source = getattr(stage_cfg, "engine_input_source", None)
-        if engine_input_source is not None:
-            if hasattr(omni_kv, "setdefault"):
-                omni_kv.setdefault("engine_input_source", list(engine_input_source))
-            elif hasattr(omni_kv, "__setitem__") and "engine_input_source" not in omni_kv:
-                omni_kv["engine_input_source"] = list(engine_input_source)
-    except Exception as e:
-        logger.debug("Failed to inject stage info into omni_kv_config: %s", e)
-
-
 def _inject_global_id(target: Any, request_id: str) -> None:
     """Inject global_request_id into a prompt dict's additional_information."""
     if isinstance(target, dict):
@@ -255,6 +233,7 @@ class AsyncOmniEngine:
         stage_init_timeout: int = 300,
         init_timeout: int = 600,
         diffusion_batch_size: int = 1,
+        single_stage_mode: bool = False,
         **kwargs: Any,
     ) -> None:
         self.model = model
@@ -274,6 +253,31 @@ class AsyncOmniEngine:
             ea_dict.pop("model", None)
             kwargs = {**ea_dict, **kwargs}
 
+        # ------------------------------------------------------------------ #
+        # Single-stage mode detection                                        #
+        # ------------------------------------------------------------------ #
+        # Single-stage mode is enabled when the caller explicitly passes      #
+        # single_stage_mode=True, or when a stage_id is provided in the args. #
+        _stage_id_kwarg = kwargs.get("stage_id")
+        if isinstance(_stage_id_kwarg, int) and not single_stage_mode:
+            single_stage_mode = True
+
+        self.single_stage_mode: bool = single_stage_mode
+        self._single_stage_id_filter: int | None = (
+            int(_stage_id_kwarg) if single_stage_mode and isinstance(_stage_id_kwarg, int) else None
+        )
+        self._omni_master_address: str | None = kwargs.get("omni_master_address")
+        self._omni_master_port: int | None = kwargs.get("omni_master_port")
+        self._omni_master_server: OmniMasterServer | None = None
+
+        if single_stage_mode:
+            logger.info(
+                "[AsyncOmniEngine] Single-stage mode enabled (stage_id_filter=%s, master=%s:%s)",
+                self._single_stage_id_filter,
+                self._omni_master_address,
+                self._omni_master_port,
+            )
+
         self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs)
 
         self.num_stages = len(self.stage_configs)
@@ -350,61 +354,89 @@ class AsyncOmniEngine:
         started_stage: StartedLlmStage | None = None
         lock_fds: list[int] = []
         device_control_env = current_omni_platform.device_control_env_var
-
         try:
-            with llm_stage_launch_lock:
-                previous_visible_devices = os.environ.get(device_control_env)
-                try:
-                    setup_stage_devices(metadata.stage_id, metadata.runtime_cfg)
-                    engine_args_dict = build_engine_args_dict(
-                        stage_cfg,
-                        self.model,
-                        stage_connector_spec=stage_connector_spec,
-                    )
-                    omni_conn_cfg, omni_from, omni_to = omni_kv_connector
-                    if omni_conn_cfg:
-                        omni_kv = engine_args_dict.get("omni_kv_config") or {}
-                        if not isinstance(omni_kv, dict):
-                            omni_kv = dict(omni_kv)
-                        omni_kv["connector_config"] = omni_conn_cfg
-                        omni_kv["omni_from_stage"] = omni_from
-                        omni_kv["omni_to_stage"] = omni_to
-                        omni_kv.setdefault("stage_id", metadata.stage_id)
-                        engine_args_dict["omni_kv_config"] = omni_kv
-                    vllm_config, executor_class = build_vllm_config(
-                        stage_cfg,
-                        self.model,
-                        stage_connector_spec=stage_connector_spec,
-                        engine_args_dict=engine_args_dict,
-                    )
-                    lock_fds = acquire_device_locks(
-                        metadata.stage_id,
-                        engine_args_dict,
-                        stage_init_timeout,
-                    )
-                    addresses, proc, handshake_address = spawn_stage_core(
-                        vllm_config=vllm_config,
-                        executor_class=executor_class,
-                        log_stats=False,
-                    )
-                    started_stage = StartedLlmStage(
-                        stage_id=metadata.stage_id,
-                        metadata=metadata,
-                        vllm_config=vllm_config,
-                        executor_class=executor_class,
-                        proc=proc,
-                        addresses=addresses,
-                    )
-                    logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id)
-                    # Keep the stage-specific device visibility until vLLM
-                    # finishes starting all child processes.
-                    complete_stage_handshake(proc, handshake_address, addresses, vllm_config)
-                    logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id)
-                finally:
-                    if previous_visible_devices is None:
-                        current_omni_platform.unset_device_control_env_var()
-                    else:
-                        current_omni_platform.set_device_control_env_var(previous_visible_devices)
+            proc = None
+            handshake_address = None
+            with ExitStack() as launch_stack:
+                with llm_stage_launch_lock:
+                    previous_visible_devices = os.environ.get(device_control_env)
+                    try:
+                        setup_stage_devices(metadata.stage_id, metadata.runtime_cfg)
+                        engine_args_dict = build_engine_args_dict(
+                            stage_cfg,
+                            self.model,
+                            stage_connector_spec=stage_connector_spec,
+                        )
+                        omni_conn_cfg, omni_from, omni_to = omni_kv_connector
+                        if omni_conn_cfg:
+                            omni_kv = engine_args_dict.get("omni_kv_config") or {}
+                            if not isinstance(omni_kv, dict):
+                                omni_kv = dict(omni_kv)
+                            omni_kv["connector_config"] = omni_conn_cfg
+                            omni_kv["omni_from_stage"] = omni_from
+                            omni_kv["omni_to_stage"] = omni_to
+                            omni_kv.setdefault("stage_id", metadata.stage_id)
+                            engine_args_dict["omni_kv_config"] = omni_kv
+                        vllm_config, executor_class = build_vllm_config(
+                            stage_cfg,
+                            self.model,
+                            stage_connector_spec=stage_connector_spec,
+                            engine_args_dict=engine_args_dict,
+                        )
+                        lock_fds = acquire_device_locks(
+                            metadata.stage_id,
+                            engine_args_dict,
+                            stage_init_timeout,
+                        )
+                        if self.single_stage_mode and self._omni_master_server is not None:
+                            engine_manager, coordinator, addresses = launch_stack.enter_context(
+                                launch_omni_core_engines(
+                                    vllm_config=vllm_config,
+                                    executor_class=executor_class,
+                                    log_stats=False,
+                                    omni_master_server=self._omni_master_server,
+                                    stage_id=metadata.stage_id,
+                                    stage_config=stage_cfg,
+                                )
+                            )
+                            started_stage = StartedLlmStage(
+                                stage_id=metadata.stage_id,
+                                metadata=metadata,
+                                vllm_config=vllm_config,
+                                executor_class=executor_class,
+                                addresses=addresses,
+                                engine_manager=engine_manager,
+                                coordinator=coordinator,
+                            )
+                        else:
+                            addresses, proc, handshake_address = spawn_stage_core(
+                                vllm_config=vllm_config,
+                                executor_class=executor_class,
+                                log_stats=False,
+                            )
+                            started_stage = StartedLlmStage(
+                                stage_id=metadata.stage_id,
+                                metadata=metadata,
+                                vllm_config=vllm_config,
+                                executor_class=executor_class,
+                                addresses=addresses,
+                                proc=proc,
+                            )
+                        logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id)
+                        # Keep the stage-specific device visibility until vLLM
+                        # finishes starting all child processes.
+                        if self.single_stage_mode and self._omni_master_server is not None:
+                            launch_stack.close()
+                        else:
+                            assert proc is not None
+                            assert handshake_address is not None
+                            complete_stage_handshake(proc, handshake_address, addresses, vllm_config)
+                        logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id)
+                    finally:
+                        if previous_visible_devices is None:
+                            current_omni_platform.unset_device_control_env_var()
+                        else:
+                            current_omni_platform.set_device_control_env_var(previous_visible_devices)
 
             assert started_stage is not None
             return started_stage
@@ -416,13 +448,138 @@ class AsyncOmniEngine:
             if lock_fds:
                 release_device_locks(lock_fds)
 
+    def _create_remote_llm_stage(
+        self,
+        stage_cfg: Any,
+        metadata: Any,
+        stage_connector_spec: dict[str, Any],
+        stage_init_timeout: int,
+        omni_master_server: OmniMasterServer,
+    ) -> StartedLlmStage:
+        """Attach to a remote engine core and wait for its startup handshake."""
+        started_stage: StartedLlmStage | None = None
+        try:
+            raw_stage_cfg = omni_master_server.get_stage_config(
+                metadata.stage_id,
+                timeout_s=stage_init_timeout,
+            )
+            if raw_stage_cfg is None:
+                raise ValueError(f"Remote stage {metadata.stage_id} registered without stage config")
+            stage_cfg = OmegaConf.create(raw_stage_cfg)
+            engine_args_dict = build_engine_args_dict(
+                stage_cfg,
+                self.model,
+                stage_connector_spec=stage_connector_spec,
+            )
+            vllm_config, executor_class = build_vllm_config(
+                stage_cfg,
+                self.model,
+                stage_connector_spec=stage_connector_spec,
+                engine_args_dict=engine_args_dict,
+            )
+            vllm_config.parallel_config.data_parallel_size_local = 0
+            launch_cm = connect_remote_engine_cores(
+                vllm_config=vllm_config,
+                omni_master_server=omni_master_server,
+                stage_id=metadata.stage_id,
+            )
+            logger.info("[AsyncOmniEngine] Stage %s remote engine handshake started", metadata.stage_id)
+            with launch_cm as (engine_manager, coordinator, addresses):
+                started_stage = StartedLlmStage(
+                    stage_id=metadata.stage_id,
+                    metadata=metadata,
+                    vllm_config=vllm_config,
+                    executor_class=executor_class,
+                    engine_manager=engine_manager,
+                    coordinator=coordinator,
+                    addresses=addresses,
+                )
+            logger.info("[AsyncOmniEngine] Stage %s remote engine startup completed", metadata.stage_id)
+            assert started_stage is not None
+            return started_stage
+        except Exception:
+            if started_stage is not None:
+                close_started_llm_stage(started_stage)
+            raise
+
+    def _launch_diffusion_stage(
+        self,
+        stage_cfg: Any,
+        metadata: Any,
+        omni_master_server: OmniMasterServer,
+    ) -> StageDiffusionClient:
+        """Launch a local diffusion stage on OmniMasterServer-allocated sockets."""
+        proc = None
+        try:
+            od_config = build_diffusion_config(self.model, stage_cfg, metadata)
+            handshake_address, request_address, response_address = register_stage_with_omni_master(
+                omni_master_address=omni_master_server.address,
+                omni_master_port=omni_master_server.port,
+                omni_stage_id=metadata.stage_id,
+                omni_stage_config=stage_cfg,
+                return_addresses=True,
+            )
+            logger.info(
+                "[AsyncOmniEngine] Stage %s diffusion registration completed",
+                metadata.stage_id,
+            )
+            proc, _, _, _ = spawn_diffusion_proc(
+                self.model,
+                od_config,
+                handshake_address=handshake_address,
+                request_address=request_address,
+                response_address=response_address,
+            )
+            complete_diffusion_handshake(proc, handshake_address)
+            logger.info(
+                "[AsyncOmniEngine] Stage %s diffusion startup completed",
+                metadata.stage_id,
+            )
+            return StageDiffusionClient.from_addresses(
+                metadata,
+                request_address=request_address,
+                response_address=response_address,
+                proc=proc,
+                batch_size=self.diffusion_batch_size,
+            )
+        except Exception:
+            if proc is not None:
+                terminate_alive_proc(proc)
+            raise
+
+    def _create_remote_diffusion_stage(
+        self,
+        metadata: Any,
+        stage_init_timeout: int,
+        omni_master_server: OmniMasterServer,
+    ) -> StageDiffusionClient:
+        """Attach to a remote diffusion stage registered with OmniMasterServer."""
+        remote_stage_cfg = OmegaConf.create(
+            omni_master_server.get_stage_config(
+                metadata.stage_id,
+                timeout_s=stage_init_timeout,
+            )
+        )
+        remote_metadata = extract_stage_metadata(remote_stage_cfg)
+        addresses = omni_master_server.get_zmq_addresses(metadata.stage_id)
+        logger.info(
+            "[AsyncOmniEngine] Stage %s remote diffusion startup completed",
+            metadata.stage_id,
+        )
+        return StageDiffusionClient.from_addresses(
+            remote_metadata,
+            request_address=addresses.inputs[0],
+            response_address=addresses.outputs[0],
+            batch_size=self.diffusion_batch_size,
+        )
+
     def _attach_llm_stage(
         self,
         started: StartedLlmStage,
     ) -> tuple[Any, Any, Any, InputProcessor | None]:
         """Attach a READY LLM stage to the orchestrator event loop."""
 
-        client_addresses = {
+        client_addresses: dict[str, str] = {
             "input_address": started.addresses.inputs[0],
             "output_address": started.addresses.outputs[0],
         }
@@ -430,14 +587,18 @@ class AsyncOmniEngine:
             client_addresses["stats_update_address"] = started.addresses.frontend_stats_publish_address
 
         try:
-            stage_client = StageEngineCoreClient(
+            stage_client = StageEngineCoreClientBase.make_async_mp_client(
                 vllm_config=started.vllm_config,
                 executor_class=started.executor_class,
                 metadata=started.metadata,
                 client_addresses=client_addresses,
                 proc=started.proc,
+                engine_manager=started.engine_manager,
+                coordinator=started.coordinator,
             )
             started.proc = None
+            started.engine_manager = None
+            started.coordinator = None
         except Exception:
             close_started_llm_stage(started)
             raise
@@ -493,7 +654,7 @@ class AsyncOmniEngine:
         output_processors: list[Any | None] = [None] * num_stages
         stage_vllm_configs: list[Any | None] = [None] * num_stages
         input_processor: InputProcessor | None = None
-        llm_stage_ids: list[int] = []
+        llm_stage_positions: list[int] = []
         llm_launch_futures: dict[int, concurrent.futures.Future[StartedLlmStage]] = {}
         started_llm_stages: dict[int, StartedLlmStage] = {}
         llm_stage_launch_lock = threading.Lock()
@@ -507,45 +668,102 @@ class AsyncOmniEngine:
         prepare_engine_environment()
         omni_transfer_config = load_omni_transfer_config_for_model(self.model, self.config_path)
 
+        # ------------------------------------------------------------------ #
+        # Single-stage mode: start OmniMasterServer before launching stages.  #
+        # ------------------------------------------------------------------ #
+        if self.single_stage_mode:
+            if not self._omni_master_address or not self._omni_master_port:
+                raise ValueError(
+                    "AsyncOmniEngine single_stage_mode requires both "
+                    "omni_master_address and omni_master_port to be set."
+                )
+            # Collect all configured stage IDs for pre-allocation.
+            all_stage_ids: list[int] = []
+            seen_stage_ids: set[int] = set()
+            for i, sc in enumerate(self.stage_configs):
+                stage_id = int(getattr(sc, "stage_id", i))
+                if stage_id in seen_stage_ids:
+                    raise ValueError(
+                        f"Duplicate stage_id {stage_id!r} detected among configured stages; stage_ids must be unique."
+                    )
+                seen_stage_ids.add(stage_id)
+                all_stage_ids.append(stage_id)
+            self._omni_master_server = OmniMasterServer(
+                master_address=self._omni_master_address,
+                master_port=self._omni_master_port,
+                stage_ids=all_stage_ids,
+            )
+            self._omni_master_server.start()
+            logger.info(
+                "[AsyncOmniEngine] OmniMasterServer started for stages %s",
+                all_stage_ids,
+            )
+
         try:
             with concurrent.futures.ThreadPoolExecutor(
                 max_workers=max(1, llm_stage_count),
                 thread_name_prefix="llm-stage-launch",
             ) as launch_executor:
-                for stage_id, stage_cfg in enumerate(self.stage_configs):
-                    logger.info("[AsyncOmniEngine] Initializing stage %s", stage_id)
+                for stage_idx, stage_cfg in enumerate(self.stage_configs):
                     metadata = extract_stage_metadata(stage_cfg)
+                    configured_stage_id = metadata.stage_id
+                    logger.info("[AsyncOmniEngine] Initializing stage %s", configured_stage_id)
                     if metadata.prompt_expand_func is not None:
                         prompt_expand_func = metadata.prompt_expand_func
 
+                    if self.single_stage_mode:
+                        metadata.runtime_cfg = None
+
                     stage_connector_spec = get_stage_connector_spec(
                         omni_transfer_config=omni_transfer_config,
-                        stage_id=stage_id,
+                        stage_id=configured_stage_id,
                         async_chunk=async_chunk,
                     )
 
-                    omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id)
+                    omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, configured_stage_id)
 
                     if metadata.stage_type == "diffusion":
+                        is_remote_diffusion_stage = (
+                            self.single_stage_mode
+                            and self._single_stage_id_filter is not None
+                            and configured_stage_id != self._single_stage_id_filter
+                        )
+                        if is_remote_diffusion_stage:
+                            assert self._omni_master_server is not None
+                            stage_clients[stage_idx] = self._create_remote_diffusion_stage(
+                                metadata,
+                                stage_init_timeout,
+                                self._omni_master_server,
+                            )
+                            continue
+
                         with llm_stage_launch_lock:
                             previous_visible_devices = os.environ.get(device_control_env)
                             try:
-                                setup_stage_devices(stage_id, metadata.runtime_cfg)
+                                setup_stage_devices(configured_stage_id, metadata.runtime_cfg)
                                 omni_conn_cfg, omni_from, omni_to = omni_kv_connector
                                 if omni_conn_cfg:
                                     from vllm_omni.entrypoints.utils import inject_omni_kv_config
 
                                     inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
-                                _inject_kv_stage_info(stage_cfg, stage_id)
-                                stage_clients[stage_id] = initialize_diffusion_stage(
-                                    self.model,
-                                    stage_cfg,
-                                    metadata,
-                                    batch_size=self.diffusion_batch_size,
-                                )
+                                inject_kv_stage_info(stage_cfg, configured_stage_id)
+                                if self.single_stage_mode:
+                                    assert self._omni_master_server is not None
+                                    stage_clients[stage_idx] = self._launch_diffusion_stage(
+                                        stage_cfg,
+                                        metadata,
+                                        self._omni_master_server,
+                                    )
+                                else:
+                                    stage_clients[stage_idx] = initialize_diffusion_stage(
+                                        self.model,
+                                        stage_cfg,
+                                        metadata,
+                                        batch_size=self.diffusion_batch_size,
+                                    )
                                 logger.info(
                                     "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)",
-                                    stage_id,
+                                    configured_stage_id,
                                     self.diffusion_batch_size,
                                 )
                             finally:
@@ -555,30 +773,58 @@ class AsyncOmniEngine:
                                     current_omni_platform.set_device_control_env_var(previous_visible_devices)
                         continue
 
-                    llm_stage_ids.append(stage_id)
-                    llm_launch_futures[stage_id] = launch_executor.submit(
-                        self._launch_llm_stage,
-                        stage_cfg,
-                        metadata,
-                        stage_connector_spec,
-                        stage_init_timeout,
-                        llm_stage_launch_lock,
-                        omni_kv_connector,
-                    )
+                    llm_stage_positions.append(stage_idx)
+
+                    # In single-stage mode, stages that don't match the local
+                    # stage_id filter are skipped.
+                    if (
+                        self.single_stage_mode
+                        and self._single_stage_id_filter is not None
+                        and configured_stage_id != self._single_stage_id_filter
+                    ):
+                        assert self._omni_master_server is not None
+                        llm_launch_futures[stage_idx] = launch_executor.submit(
+                            self._create_remote_llm_stage,
+                            stage_cfg,
+                            metadata,
+                            stage_connector_spec,
+                            stage_init_timeout,
+                            self._omni_master_server,
+                        )
+                    else:
+                        llm_launch_futures[stage_idx] = launch_executor.submit(
+                            self._launch_llm_stage,
+                            stage_cfg,
+                            metadata,
+                            stage_connector_spec,
+                            stage_init_timeout,
+                            llm_stage_launch_lock,
+                            omni_kv_connector,
+                        )
 
                 concurrent.futures.wait(list(llm_launch_futures.values()))
 
-                for stage_id in llm_stage_ids:
-                    started_llm_stages[stage_id] = llm_launch_futures[stage_id].result()
+                for stage_idx in llm_stage_positions:
+                    started_llm_stages[stage_idx] = llm_launch_futures[stage_idx].result()
 
-            for stage_id in llm_stage_ids:
-                started = started_llm_stages[stage_id]
-                stage_client, output_processor, vllm_config, stage0_input_processor = self._attach_llm_stage(started)
-                stage_clients[stage_id] = stage_client
-                output_processors[stage_id] = output_processor
-                stage_vllm_configs[stage_id] = vllm_config
-                if stage0_input_processor is not None:
-                    input_processor = stage0_input_processor
+            attach_futures: dict[concurrent.futures.Future[tuple[Any, Any, Any, InputProcessor | None]], int] = {}
+            with concurrent.futures.ThreadPoolExecutor(
+                max_workers=max(1, len(llm_stage_positions)),
+                thread_name_prefix="llm-stage-attach",
+            ) as attach_executor:
+                for stage_idx in llm_stage_positions:
+                    attach_futures[attach_executor.submit(self._attach_llm_stage, started_llm_stages[stage_idx])] = (
+                        stage_idx
+                    )
+
+                for future in concurrent.futures.as_completed(attach_futures):
+                    stage_idx = attach_futures[future]
+                    stage_client, output_processor, vllm_config, stage0_input_processor = future.result()
+                    stage_clients[stage_idx] = stage_client
+                    output_processors[stage_idx] = output_processor
+                    stage_vllm_configs[stage_idx] = vllm_config
+                    if stage0_input_processor is not None:
+                        input_processor = stage0_input_processor
 
             initialized_stage_clients, default_sampling_params_list, stage_metadata = finalize_initialized_stages(
                 stage_clients,
@@ -595,8 +841,13 @@ class AsyncOmniEngine:
             )
             cleanup_failed_stage_initialization(
                 stage_clients,
-                [started_llm_stages[stage_id] for stage_id in llm_stage_ids if stage_id in started_llm_stages],
+                [started_llm_stages[stage_idx] for stage_idx in llm_stage_positions if stage_idx in started_llm_stages],
             )
+            if self._omni_master_server is not None:
+                try:
+                    self._omni_master_server.stop()
+                except Exception:
+                    logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during stage-init cleanup")
             raise
 
         self.stage_clients = initialized_stage_clients
@@ -1310,3 +1561,10 @@ class AsyncOmniEngine:
                 q.close()
             except Exception:
                 pass
+
+        if self._omni_master_server is not None:
+            try:
+                self._omni_master_server.stop()
+            except Exception:
+                logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during shutdown")
+            self._omni_master_server = None

vllm_omni/engine/stage_engine_core_client.py

@@ -12,7 +12,7 @@ from urllib.parse import urlparse
 
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
-from vllm.v1.engine.core_client import AsyncMPClient
+from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
 
 from vllm_omni.distributed.omni_connectors.utils.initialization import KV_TRANSFER_PORT_OFFSET
 from vllm_omni.engine.stage_init_utils import StageMetadata
@@ -25,18 +25,54 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
-class StageEngineCoreClient(AsyncMPClient):
-    """Stage async client that inherits from vLLM's AsyncMPClient.
+class StageEngineCoreClientBase:
+    """Shared stage-aware behavior for async EngineCore clients.
 
-    Fully reuses AsyncMPClient for:
+    The concrete transport/load-balancing behavior is supplied by the
+    multiprocessing client subclass in the MRO.
+
+    Fully reuses the underlying vLLM async MP client ``__init__`` for:
     - ZMQ setup, sockets
     - outputs_queue, output_queue_task
     - All utility methods (get_output_async, abort_requests_async, etc.)
 
     The subprocess is spawned externally via ``spawn_stage_core`` /
     ``complete_stage_handshake`` from *stage_engine_core_proc.py*.
+    In single-stage CLI mode, the client may instead attach to an
+    ``engine_manager`` / ``coordinator`` pair created elsewhere.
     """
 
+    @staticmethod
+    def make_async_mp_client(
+        vllm_config: Any,
+        executor_class: type,
+        metadata: StageMetadata,
+        client_addresses: dict[str, str] | None = None,
+        proc: Any = None,
+        engine_manager: Any = None,
+        coordinator: Any = None,
+        client_count: int = 1,
+        client_index: int = 0,
+    ) -> StageEngineCoreClient | DPLBStageEngineCoreClient:
+        """Create the appropriate stage async client for the DP mode."""
+        parallel_config = vllm_config.parallel_config
+        client_args = dict(
+            vllm_config=vllm_config,
+            executor_class=executor_class,
+            metadata=metadata,
+            client_addresses=client_addresses,
+            proc=proc,
+            engine_manager=engine_manager,
+            coordinator=coordinator,
+            client_count=client_count,
+            client_index=client_index,
+        )
+
+        if parallel_config.data_parallel_size > 1 and not parallel_config.data_parallel_external_lb:
+            return DPLBStageEngineCoreClient(**client_args)
+
+        return StageEngineCoreClient(**client_args)
+
     def __init__(
         self,
         vllm_config: Any,
@@ -85,8 +121,10 @@ class StageEngineCoreClient(AsyncMPClient):
         self._kv_sender_info: dict[str, Any] | None = None
         self._kv_sender_initialized = False
 
+        client_name = self.__class__.__name__
         logger.info(
-            "[StageEngineCoreClient] Stage-%s initializing EngineCore",
+            "[%s] Stage-%s initializing EngineCore",
+            client_name,
             self.stage_id,
         )
         try:
@@ -98,23 +136,30 @@ class StageEngineCoreClient(AsyncMPClient):
                 client_count=client_count,
                 client_index=client_index,
             )
+            if engine_manager is not None:
+                self.resources.engine_manager = engine_manager
+            if coordinator is not None:
+                self.resources.coordinator = coordinator
         except Exception:
             logger.exception(
-                "[StageEngineCoreClient] Stage-%s EngineCore init failed",
+                "[%s] Stage-%s EngineCore init failed",
+                client_name,
                 self.stage_id,
             )
             try:
                 self.shutdown()
             except Exception as shutdown_error:
                 logger.warning(
-                    "[StageEngineCoreClient] Stage-%s cleanup after init failure failed: %s",
+                    "[%s] Stage-%s cleanup after init failure failed: %s",
+                    client_name,
                     self.stage_id,
                     shutdown_error,
                 )
             raise
         self._initialize_kv_sender_endpoint()
         logger.info(
-            "[StageEngineCoreClient] Stage-%s EngineCore running",
+            "[%s] Stage-%s EngineCore running",
+            client_name,
             self.stage_id,
         )
 
@@ -122,7 +167,12 @@ class StageEngineCoreClient(AsyncMPClient):
 
     async def add_request_async(self, request: EngineCoreRequest) -> None:
         """Add request to the stage engine core."""
-        logger.info(f"[StageEngineCoreClient] Stage-{self.stage_id} adding request: {request.request_id}")
+        logger.info(
+            "[%s] Stage-%s adding request: %s",
+            self.__class__.__name__,
+            self.stage_id,
+            request.request_id,
+        )
         await super().add_request_async(request)
 
     # ==================== Stage Methods ====================
@@ -287,9 +337,9 @@ class StageEngineCoreClient(AsyncMPClient):
     ) -> Any:
         """Forward control RPCs to the underlying AsyncMPClient stage engine.
 
-        Each ``StageEngineCoreClient`` already represents one logical stage, so
-        stage-scoped control operations should be executed here and then fanned
-        in-core across the workers managed by this EngineCore client.
+        Each stage client already represents one logical stage, so stage-scoped
+        control operations should be executed here and then fanned in-core
+        across the workers managed by this EngineCore client.
         """
         return await super().collective_rpc_async(
             method=method,
@@ -299,10 +349,19 @@ class StageEngineCoreClient(AsyncMPClient):
         )
 
     def shutdown(self) -> None:
-        """Shutdown ZMQ connections and the subprocess."""
+        """Shutdown managed resources and any externally spawned subprocess."""
         super().shutdown()
         if self._proc is not None and self._proc.is_alive():
             self._proc.terminate()
             self._proc.join(timeout=5)
             if self._proc.is_alive():
                 self._proc.kill()
+        self._proc = None
+
+
+class StageEngineCoreClient(StageEngineCoreClientBase, AsyncMPClient):
+    """Stage async client backed by vLLM's ``AsyncMPClient``."""
+
+
+class DPLBStageEngineCoreClient(StageEngineCoreClientBase, DPLBAsyncMPClient):
+    """Stage async client backed by vLLM's ``DPLBAsyncMPClient``."""

vllm_omni/engine/stage_engine_startup.py

@@ -0,0 +1,599 @@
+"""Helpers for launching and handshaking omni engine cores."""
+
+from __future__ import annotations
+
+import contextlib
+import dataclasses
+import threading
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Any
+
+import msgspec
+import zmq
+from omegaconf import OmegaConf
+from vllm.config import CacheConfig, VllmConfig
+from vllm.logger import init_logger
+from vllm.utils.network_utils import get_open_port, zmq_socket_ctx
+from vllm.v1.engine.coordinator import DPCoordinator
+from vllm.v1.engine.utils import (
+    STARTUP_POLL_PERIOD_MS,
+    CoreEngine,
+    CoreEngineProcManager,
+    CoreEngineState,
+    EngineHandshakeMetadata,
+    EngineZmqAddresses,
+    wait_for_engine_startup,
+)
+from vllm.v1.executor import Executor
+
+logger = init_logger(__name__)
+
+# Poll period (ms) used by the registration/handshake loop.
+_POLL_PERIOD_MS = 5_000
+# Default timeout (s) for a stage to send READY.
+_DEFAULT_STARTUP_TIMEOUT_S = 300
+
+
+def _serialize_stage_config(stage_config: Any) -> Any:
+    """Convert a stage config to msgpack-friendly builtins."""
+    if stage_config is None or isinstance(stage_config, (str, bytes, int, float, bool)):
+        return stage_config
+
+    if OmegaConf.is_config(stage_config):
+        return _serialize_stage_config(OmegaConf.to_container(stage_config, resolve=True))
+
+    if dataclasses.is_dataclass(stage_config):
+        return _serialize_stage_config(dataclasses.asdict(stage_config))
+
+    if isinstance(stage_config, dict):
+        return {key: _serialize_stage_config(value) for key, value in stage_config.items() if not callable(value)}
+
+    if isinstance(stage_config, (list, tuple, set)):
+        return [_serialize_stage_config(item) for item in stage_config if not callable(item)]
+
+    if hasattr(stage_config, "items"):
+        return {key: _serialize_stage_config(value) for key, value in stage_config.items() if not callable(value)}
+
+    if hasattr(stage_config, "__dict__"):
+        return {
+            key: _serialize_stage_config(value)
+            for key, value in vars(stage_config).items()
+            if not key.startswith("_") and not callable(value)
+        }
+
+    return stage_config
+
+
+# ---------------------------------------------------------------------------
+# Per-stage address allocation
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class StageAllocation:
+    """ZMQ addresses reserved for a single stage."""
+
+    # Per-stage handshake socket (OmniMasterServer binds, engine connects)
+    handshake_bind_address: str
+    handshake_connect_address: str
+    # Input channel: client binds ROUTER, engine connects DEALER
+    input_bind_address: str
+    input_connect_address: str
+    # Output channel: client binds PULL, engine connects PUSH
+    output_bind_address: str
+    output_connect_address: str
+
+
+@dataclass(frozen=True)
+class StageCoordinatorAddresses:
+    """Optional DP coordinator addresses registered for a stage."""
+
+    coordinator_input: str | None = None
+    coordinator_output: str | None = None
+    frontend_stats_publish_address: str | None = None
+
+
+# ---------------------------------------------------------------------------
+# OmniMasterServer
+# ---------------------------------------------------------------------------
+
+
+class OmniMasterServer:
+    """Registration server for single-stage engine startup."""
+
+    def __init__(
+        self,
+        master_address: str,
+        master_port: int,
+        stage_ids: list[int],
+    ) -> None:
+        self._address = master_address
+        self._port = master_port
+        self._allocations: dict[int, StageAllocation] = {}
+        self._stage_configs: dict[int, Any] = {}
+        self._stage_coordinator_addresses: dict[int, StageCoordinatorAddresses] = {}
+        self._stage_config_events: dict[int, threading.Event] = {}
+        self._thread: threading.Thread | None = None
+        self._stop_event = threading.Event()
+
+        for sid in stage_ids:
+            self._stage_config_events[sid] = threading.Event()
+            self._stage_coordinator_addresses[sid] = StageCoordinatorAddresses()
+            hs_port = get_open_port()
+            inp_port = get_open_port()
+            out_port = get_open_port()
+            self._allocations[sid] = StageAllocation(
+                handshake_bind_address=f"tcp://{master_address}:{hs_port}",
+                handshake_connect_address=f"tcp://{master_address}:{hs_port}",
+                input_bind_address=f"tcp://{master_address}:{inp_port}",
+                input_connect_address=f"tcp://{master_address}:{inp_port}",
+                output_bind_address=f"tcp://{master_address}:{out_port}",
+                output_connect_address=f"tcp://{master_address}:{out_port}",
+            )
+
+        logger.info(
+            "[OmniMasterServer] Pre-allocated addresses for stages %s (master=%s:%d)",
+            list(stage_ids),
+            master_address,
+            master_port,
+        )
+
+    # ------------------------------------------------------------------
+    # Public helpers
+    # ------------------------------------------------------------------
+    @property
+    def address(self) -> str:
+        """Return the registration address exposed to stage launchers."""
+        return self._address
+
+    @property
+    def port(self) -> int:
+        """Return the registration port exposed to stage launchers."""
+        return self._port
+
+    def get_allocation(self, stage_id: int) -> StageAllocation:
+        """Return the full address allocation for *stage_id*."""
+        return self._allocations[stage_id]
+
+    def register_stage_config(
+        self,
+        stage_id: int,
+        stage_config: Any,
+        coordinator_addresses: StageCoordinatorAddresses | None = None,
+    ) -> None:
+        """Store the latest stage registration payload for *stage_id*."""
+        if stage_id not in self._allocations:
+            raise KeyError(stage_id)
+        self._stage_configs[stage_id] = stage_config
+        if coordinator_addresses is not None:
+            self._stage_coordinator_addresses[stage_id] = coordinator_addresses
+        self._stage_config_events[stage_id].set()
+
+    def get_stage_config(self, stage_id: int, timeout_s: float | None = None) -> Any:
+        """Return the stage config for *stage_id*, waiting if necessary."""
+        if stage_id not in self._allocations:
+            raise KeyError(stage_id)
+
+        if stage_id in self._stage_configs:
+            return self._stage_configs[stage_id]
+
+        if not self._stage_config_events[stage_id].wait(timeout=timeout_s):
+            raise TimeoutError(f"Timed out waiting for stage config for stage {stage_id}.")
+
+        return self._stage_configs[stage_id]
+
+    def get_stage_coordinator_addresses(
+        self,
+        stage_id: int,
+        timeout_s: float | None = None,
+    ) -> StageCoordinatorAddresses:
+        """Return the registered coordinator addresses for *stage_id*."""
+        if stage_id not in self._allocations:
+            raise KeyError(stage_id)
+
+        if not self._stage_config_events[stage_id].is_set():
+            if not self._stage_config_events[stage_id].wait(timeout=timeout_s):
+                raise TimeoutError(f"Timed out waiting for stage registration for stage {stage_id}.")
+
+        return self._stage_coordinator_addresses[stage_id]
+
+    def get_client_addresses(self, stage_id: int) -> dict[str, str]:
+        """Return the addresses the client-side sockets should *bind* to."""
+        alloc = self._allocations[stage_id]
+        return {
+            "input_address": alloc.input_bind_address,
+            "output_address": alloc.output_bind_address,
+        }
+
+    def get_zmq_addresses(self, stage_id: int) -> EngineZmqAddresses:
+        """Return EngineZmqAddresses using the *bind* (client) side addresses."""
+        alloc = self._allocations[stage_id]
+        return EngineZmqAddresses(
+            inputs=[alloc.input_bind_address],
+            outputs=[alloc.output_bind_address],
+        )
+
+    def get_engine_zmq_addresses(self, stage_id: int) -> EngineZmqAddresses:
+        """Return EngineZmqAddresses using the *connect* (engine) addresses."""
+        alloc = self._allocations[stage_id]
+        return EngineZmqAddresses(
+            inputs=[alloc.input_connect_address],
+            outputs=[alloc.output_connect_address],
+        )
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def start(self) -> None:
+        """Start the background server thread."""
+        self._thread = threading.Thread(
+            target=self._run,
+            name="OmniMasterServer",
+            daemon=True,
+        )
+        self._thread.start()
+        logger.info(
+            "[OmniMasterServer] Listening on tcp://%s:%d",
+            self.address,
+            self.port,
+        )
+
+    def stop(self) -> None:
+        """Signal stop and join the background thread."""
+        self._stop_event.set()
+        if self._thread is not None:
+            self._thread.join(timeout=10)
+
+    # ------------------------------------------------------------------
+    # Internal server logic
+    # ------------------------------------------------------------------
+
+    def _run(self) -> None:
+        ctx = zmq.Context()
+        try:
+            self._serve(ctx)
+        except Exception:
+            logger.exception("[OmniMasterServer] Server thread crashed")
+        finally:
+            ctx.term()
+
+    def _serve(self, ctx: zmq.Context) -> None:  # type: ignore[type-arg]
+        # Registration socket for the initial stage registration.
+        # Per-stage handshake sockets are bound by the launch helpers.
+        reg_socket: zmq.Socket = ctx.socket(zmq.ROUTER)  # type: ignore[attr-defined]
+        reg_socket.bind(f"tcp://{self.address}:{self.port}")
+
+        poller = zmq.Poller()
+        poller.register(reg_socket, zmq.POLLIN)
+
+        pending: set[int] = set(self._allocations.keys())
+
+        while pending and not self._stop_event.is_set():
+            events: list[tuple[zmq.Socket, int]] = poller.poll(_POLL_PERIOD_MS)  # type: ignore[assignment]
+            if not events:
+                logger.debug("[OmniMasterServer] Still waiting for registration from stages: %s", pending)
+                continue
+
+            for sock, _ in events:
+                if sock is reg_socket:
+                    sid = self._handle_registration(reg_socket)
+                    if sid is not None:
+                        pending.discard(sid)
+
+        # Cleanup
+        reg_socket.close(linger=0)
+        logger.info("[OmniMasterServer] All stages registered; server thread exiting.")
+
+    def _handle_registration(self, reg_socket: zmq.Socket) -> int | None:  # type: ignore[type-arg]
+        """Receive a stage registration and reply with the handshake address.
+
+        Returns the registered stage_id on success, or None on failure.
+        """
+        frames = reg_socket.recv_multipart()
+        if len(frames) < 2:
+            logger.warning(
+                "[OmniMasterServer] Unexpected registration frame count: %d",
+                len(frames),
+            )
+            return None
+        identity = frames[0]
+        msg_bytes = frames[-1]
+        try:
+            msg = msgspec.msgpack.decode(msg_bytes)
+        except Exception as exc:
+            logger.warning("[OmniMasterServer] Failed to decode registration message: %s", exc)
+            return None
+
+        stage_id: int | None = msg.get("stage_id")
+        if stage_id not in self._allocations:
+            logger.warning(
+                "[OmniMasterServer] Received registration for unknown stage_id=%s",
+                stage_id,
+            )
+            return None
+
+        self.register_stage_config(
+            stage_id,
+            msg.get("stage_config"),
+            coordinator_addresses=StageCoordinatorAddresses(
+                coordinator_input=msg.get("coordinator_input"),
+                coordinator_output=msg.get("coordinator_output"),
+                frontend_stats_publish_address=msg.get("frontend_stats_publish_address"),
+            ),
+        )
+
+        alloc = self._allocations[stage_id]
+        response = msgspec.msgpack.encode(
+            {
+                "handshake_address": alloc.handshake_connect_address,
+                "input_address": alloc.input_bind_address,
+                "output_address": alloc.output_bind_address,
+            }
+        )
+        # ROUTER-DEALER: reply is [identity, payload] (no empty delimiter).
+        reg_socket.send_multipart([identity, response])
+        logger.info(
+            "[OmniMasterServer] Stage %d registered; assigned handshake=%s",
+            stage_id,
+            alloc.handshake_connect_address,
+        )
+        return stage_id
+
+
+def register_stage_with_omni_master(
+    *,
+    omni_master_address: str,
+    omni_master_port: int,
+    omni_stage_id: int,
+    omni_stage_config: Any = None,
+    coordinator: DPCoordinator | None = None,
+    return_addresses: bool = False,
+) -> str | tuple[str, str, str]:
+    """Register a stage with the omni master server.
+
+    Returns the per-stage handshake address by default. When
+    ``return_addresses`` is true, also returns the stage input/output
+    addresses allocated by the master.
+    """
+
+    reg_ctx = zmq.Context()
+    try:
+        reg_sock: zmq.Socket = reg_ctx.socket(zmq.DEALER)  # type: ignore[attr-defined]
+        try:
+            reg_sock.connect(f"tcp://{omni_master_address}:{omni_master_port}")
+            payload = {
+                "stage_id": omni_stage_id,
+                "stage_config": _serialize_stage_config(omni_stage_config),
+            }
+            if coordinator is not None:
+                coordinator_input, coordinator_output = coordinator.get_engine_socket_addresses()
+                payload["coordinator_input"] = coordinator_input
+                payload["coordinator_output"] = coordinator_output
+                payload["frontend_stats_publish_address"] = coordinator.get_stats_publish_address()
+
+            reg_sock.send(msgspec.msgpack.encode(payload))
+            timeout_ms = _DEFAULT_STARTUP_TIMEOUT_S * 1_000
+            if not reg_sock.poll(timeout=timeout_ms):
+                raise RuntimeError(
+                    f"Timed out waiting for registration "
+                    f"response from OmniMasterServer "
+                    f"({omni_master_address}:{omni_master_port}) "
+                    f"for stage {omni_stage_id}."
+                )
+            response_bytes = reg_sock.recv()
+            response = msgspec.msgpack.decode(response_bytes)
+            handshake_address: str = response["handshake_address"]
+            input_address: str = response["input_address"]
+            output_address: str = response["output_address"]
+            logger.info(
+                "Stage %d registered; handshake_address=%s",
+                omni_stage_id,
+                handshake_address,
+            )
+        finally:
+            reg_sock.close(linger=0)
+    finally:
+        reg_ctx.term()
+
+    if return_addresses:
+        return handshake_address, input_address, output_address
+    return handshake_address
+
+
+def _wait_for_omni_engine_startup(
+    handshake_socket: zmq.Socket,
+    engine_addresses: EngineZmqAddresses,
+    engines: list[CoreEngine],
+    cache_config: CacheConfig,
+) -> None:
+    """Wait for omni-managed engines to finish the HELLO/READY handshake."""
+    conn_pending = len(engines)
+    start_pending = 0
+
+    poller = zmq.Poller()
+    poller.register(handshake_socket, zmq.POLLIN)
+
+    while conn_pending or start_pending:
+        events = poller.poll(STARTUP_POLL_PERIOD_MS)
+        if not events:
+            logger.debug(
+                "[omni] Waiting for %d engine(s) to connect, %d to start.",
+                conn_pending,
+                start_pending,
+            )
+            continue
+
+        eng_identity, msg_bytes = handshake_socket.recv_multipart()
+        eng_index = int.from_bytes(eng_identity, "little")
+        engine = next((e for e in engines if e.identity == eng_identity), None)
+        if engine is None:
+            raise RuntimeError(f"[omni] Handshake message from unexpected engine rank: {eng_index}")
+
+        msg = msgspec.msgpack.decode(msg_bytes)
+        status: str = msg["status"]
+
+        if status == "HELLO" and engine.state == CoreEngineState.NEW:
+            init_message = msgspec.msgpack.encode(
+                EngineHandshakeMetadata(addresses=engine_addresses, parallel_config={})
+            )
+            handshake_socket.send_multipart((eng_identity, init_message), copy=False)
+            conn_pending -= 1
+            start_pending += 1
+            engine.state = CoreEngineState.CONNECTED
+            logger.debug("[omni] HELLO from engine %d", eng_index)
+
+        elif status == "READY" and engine.state == CoreEngineState.CONNECTED:
+            num_gpu_blocks = (cache_config.num_gpu_blocks or 0) + msg["num_gpu_blocks"]
+            cache_config.num_gpu_blocks = num_gpu_blocks
+            if engine_addresses.frontend_stats_publish_address is None:
+                engine_addresses.frontend_stats_publish_address = msg.get("dp_stats_address")
+            start_pending -= 1
+            engine.state = CoreEngineState.READY
+            logger.debug("[omni] READY from engine %d (num_gpu_blocks=%d)", eng_index, msg["num_gpu_blocks"])
+
+        else:
+            raise RuntimeError(f"[omni] Unexpected status '{status}' from engine {eng_index} in state {engine.state}.")
+
+
+@contextlib.contextmanager
+def connect_remote_engine_cores(
+    vllm_config: VllmConfig,
+    omni_master_server: OmniMasterServer,
+    stage_id: int,
+) -> Iterator[tuple[None, DPCoordinator | None, EngineZmqAddresses]]:
+    """Wait for remote engine cores to connect through the omni handshake."""
+    addresses = omni_master_server.get_zmq_addresses(stage_id)
+    parallel_config = vllm_config.parallel_config
+    # Mirror the engine-count logic from launch_omni_core_engines.
+    remote_engine_count = (
+        parallel_config.data_parallel_size_local
+        if parallel_config.data_parallel_size_local is not None and parallel_config.data_parallel_size_local > 0
+        else max(1, parallel_config.data_parallel_size)
+    )
+    start_index = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0
+    coordinator = None
+
+    registered_coordinator_addresses = omni_master_server.get_stage_coordinator_addresses(stage_id)
+    addresses.coordinator_input = registered_coordinator_addresses.coordinator_input
+    addresses.coordinator_output = registered_coordinator_addresses.coordinator_output
+    addresses.frontend_stats_publish_address = registered_coordinator_addresses.frontend_stats_publish_address
+
+    engines_to_handshake = [CoreEngine(index=start_index + i, local=False) for i in range(remote_engine_count)]
+
+    logger.info(
+        "Waiting for %d remote engine(s) for stage %d",
+        remote_engine_count,
+        stage_id,
+    )
+
+    handshake_bind_address = omni_master_server.get_allocation(stage_id).handshake_bind_address
+
+    with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
+        yield None, coordinator, addresses
+
+        _wait_for_omni_engine_startup(
+            handshake_socket,
+            addresses,
+            engines_to_handshake,
+            vllm_config.cache_config,
+        )
+
+
+@contextlib.contextmanager
+def launch_omni_core_engines(
+    vllm_config: VllmConfig,
+    executor_class: type[Executor],
+    log_stats: bool,
+    omni_master_server: OmniMasterServer,
+    stage_id: int,
+    stage_config: Any = None,
+) -> Iterator[tuple[CoreEngineProcManager, DPCoordinator | None, EngineZmqAddresses]]:
+    """Launch local engine cores using the omni registration flow."""
+    addresses = omni_master_server.get_zmq_addresses(stage_id)
+    parallel_config = vllm_config.parallel_config
+    # Determine the number of local engines and their ranks.
+    local_engine_count = (
+        parallel_config.data_parallel_size_local
+        if parallel_config.data_parallel_size_local is not None and parallel_config.data_parallel_size_local > 0
+        else max(1, parallel_config.data_parallel_size)
+    )
+    dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0
+    local_start_index = 0
+    start_index = dp_rank
+
+    # Run the DP Coordinator process with rank 0 when in online DP mode.
+    # The coordinator is needed for:
+    # 1. Internal/hybrid LB: collecting and publishing queue stats
+    # 2. MoE models: wave coordination in addition to stats
+    run_coordinator = vllm_config.needs_dp_coordinator and dp_rank == 0
+
+    if run_coordinator:
+        coordinator = DPCoordinator(
+            parallel_config,
+            enable_wave_coordination=vllm_config.model_config.is_moe,
+        )
+
+        addresses.coordinator_input, addresses.coordinator_output = coordinator.get_engine_socket_addresses()
+        addresses.frontend_stats_publish_address = coordinator.get_stats_publish_address()
+
+        logger.info(
+            "[omni] Started DP Coordinator process for stage %d (PID: %d)",
+            stage_id,
+            coordinator.proc.pid,
+        )
+    else:
+        coordinator = None
+
+    logger.info(
+        "Starting %d local engine(s) for stage %d (dp_rank=%d)",
+        local_engine_count,
+        stage_id,
+        dp_rank,
+    )
+
+    # Register the stage once and reuse the returned per-stage handshake
+    # address for all local engine-core processes.
+    handshake_address = register_stage_with_omni_master(
+        omni_master_address=omni_master_server.address,
+        omni_master_port=omni_master_server.port,
+        omni_stage_id=stage_id,
+        omni_stage_config=stage_config,
+        coordinator=coordinator,
+    )
+
+    # One CoreEngine entry per local engine so wait_for_engine_startup can
+    # track the HELLO/READY handshake for each of them.
+    engines_to_handshake = [CoreEngine(index=start_index + i, local=True) for i in range(local_engine_count)]
+
+    # Bind the pre-allocated handshake socket for this stage.
+    handshake_bind_address = omni_master_server.get_allocation(stage_id).handshake_bind_address
+
+    with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
+        local_engine_manager = CoreEngineProcManager(
+            local_engine_count=local_engine_count,
+            start_index=start_index,
+            local_start_index=local_start_index,
+            vllm_config=vllm_config,
+            local_client=True,
+            handshake_address=handshake_address,
+            executor_class=executor_class,
+            log_stats=log_stats,
+        )
+
+        yield local_engine_manager, coordinator, addresses
+
+        # Wait for all local engine-core processes to complete the
+        # standard HELLO/READY handshake — mirrors launch_core_engines.
+        coordinated_dp = parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe
+        wait_for_engine_startup(
+            handshake_socket,
+            addresses,
+            engines_to_handshake,
+            parallel_config,
+            coordinated_dp,
+            vllm_config.cache_config,
+            local_engine_manager,
+            coordinator.proc if coordinator else None,
+        )

vllm_omni/engine/stage_init_utils.py

@@ -101,6 +101,34 @@ def resolve_worker_cls(engine_args: dict[str, Any]) -> None:
         raise ValueError(f"Unknown worker_type: {worker_type}")
 
 
+def inject_kv_stage_info(stage_cfg: Any, stage_id: int) -> None:
+    """Inject stage metadata into omni_kv_config when present."""
+    try:
+        engine_args = stage_cfg.engine_args
+        if hasattr(engine_args, "get"):
+            omni_kv = engine_args.get("omni_kv_config", None)
+        else:
+            omni_kv = getattr(engine_args, "omni_kv_config", None)
+
+        if omni_kv is None:
+            return
+
+        if hasattr(omni_kv, "setdefault"):
+            omni_kv.setdefault("stage_id", stage_id)
+        elif hasattr(omni_kv, "__setitem__"):
+            if "stage_id" not in omni_kv:
+                omni_kv["stage_id"] = stage_id
+
+        engine_input_source = getattr(stage_cfg, "engine_input_source", None)
+        if engine_input_source is not None:
+            if hasattr(omni_kv, "setdefault"):
+                omni_kv.setdefault("engine_input_source", list(engine_input_source))
+            elif hasattr(omni_kv, "__setitem__") and "engine_input_source" not in omni_kv:
+                omni_kv["engine_input_source"] = list(engine_input_source)
+    except Exception as e:
+        logger.debug("Failed to inject stage info into omni_kv_config: %s", e)
+
+
 @dataclass
 class StageMetadata:
     """Lightweight stage attributes extracted from stage_config."""
@@ -129,8 +157,10 @@ class StartedLlmStage:
     metadata: Any
     vllm_config: Any
     executor_class: type
-    proc: Any
     addresses: Any
+    proc: Any = None
+    engine_manager: Any = None
+    coordinator: Any = None
 
 
 def extract_stage_metadata(stage_config: Any) -> StageMetadata:
@@ -263,6 +293,7 @@ def build_vllm_config(
     model: str,
     stage_connector_spec: dict[str, Any] | None = None,
     engine_args_dict: dict[str, Any] | None = None,
+    headless: bool = False,
 ) -> tuple[Any, type]:
     """Build engine args, then create VllmConfig and executor_class.
 
@@ -278,7 +309,10 @@ def build_vllm_config(
 
     filtered_engine_args_dict = filter_dataclass_kwargs(OmniEngineArgs, engine_args_dict)
     omni_engine_args = OmniEngineArgs(**filtered_engine_args_dict)
-    vllm_config = omni_engine_args.create_engine_config(usage_context=UsageContext.LLM_CLASS)
+    vllm_config = omni_engine_args.create_engine_config(
+        usage_context=UsageContext.LLM_CLASS,
+        headless=headless,
+    )
     executor_class = Executor.get_class(vllm_config)
 
     return vllm_config, executor_class
@@ -445,6 +479,37 @@ def get_stage_connector_spec(
     return {}
 
 
+def build_diffusion_config(
+    model: str,
+    stage_cfg: Any,
+    metadata: StageMetadata,
+) -> Any:
+    """Build diffusion config for a stage."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    engine_args_dict = build_engine_args_dict(stage_cfg, model)
+    od_config = OmniDiffusionConfig.from_kwargs(**engine_args_dict)
+
+    num_devices_per_stage = od_config.parallel_config.world_size
+    device_control_env = current_omni_platform.device_control_env_var
+    visible_devices_str = os.environ.get(device_control_env) if device_control_env else None
+    if visible_devices_str:
+        physical_devices = [device.strip() for device in visible_devices_str.split(",") if device.strip()]
+    else:
+        physical_devices = list(range(current_omni_platform.get_device_count()))
+
+    if len(physical_devices) < num_devices_per_stage:
+        raise ValueError(
+            f"Stage {metadata.stage_id} requires {num_devices_per_stage} device(s) based on parallel_config, "
+            f"but {len(physical_devices)} device(s) are available: {physical_devices}"
+        )
+
+    od_config.num_gpus = num_devices_per_stage
+    if metadata.cfg_kv_collect_func is not None:
+        od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func
+    return od_config
+
+
 def initialize_diffusion_stage(
     model: str,
     stage_cfg: Any,
@@ -461,30 +526,9 @@ def initialize_diffusion_stage(
             diffusion engine.  Passed through to ``StageDiffusionClient``
             and ultimately to ``AsyncOmni``.
     """
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
     from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
 
-    od_config = OmniDiffusionConfig.from_kwargs(
-        model=model,
-        **_to_dict(stage_cfg.engine_args),
-    )
-    num_devices_per_stage = od_config.parallel_config.world_size
-    device_control_env = current_omni_platform.device_control_env_var
-    visible_devices_str = os.environ.get(device_control_env)
-    if visible_devices_str:
-        physical_devices = [device.strip() for device in visible_devices_str.split(",") if device.strip()]
-    else:
-        physical_devices = list(range(current_omni_platform.get_device_count()))
-
-    if len(physical_devices) < num_devices_per_stage:
-        raise ValueError(
-            f"Stage {metadata.stage_id} requires {num_devices_per_stage} device(s) based on parallel_config, "
-            f"but {len(physical_devices)} device(s) are available: {physical_devices}"
-        )
-
-    od_config.num_gpus = num_devices_per_stage
-    if metadata.cfg_kv_collect_func is not None:
-        od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func
+    od_config = build_diffusion_config(model, stage_cfg, metadata)
     return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
 
 
@@ -518,17 +562,18 @@ def _shutdown_or_close_resource(resource: Any, resource_name: str, stage_id: int
 
 
 def close_started_llm_stage(started: StartedLlmStage) -> None:
-    """Terminate the subprocess owned by a launched stage that never attached."""
-    if started.proc is None:
-        return
-    try:
-        terminate_alive_proc(started.proc)
-    except Exception as cleanup_error:
-        logger.warning(
-            "[stage_init] Failed to terminate process for stage %s: %s",
-            started.stage_id,
-            cleanup_error,
-        )
+    """Release resources owned by a launched stage that never attached."""
+    if started.proc is not None:
+        try:
+            terminate_alive_proc(started.proc)
+        except Exception as cleanup_error:
+            logger.warning(
+                "[stage_init] Failed to terminate process for stage %s: %s",
+                started.stage_id,
+                cleanup_error,
+            )
+    _shutdown_or_close_resource(started.engine_manager, "engine manager", started.stage_id)
+    _shutdown_or_close_resource(started.coordinator, "coordinator", started.stage_id)
 
 
 def finalize_initialized_stages(

vllm_omni/entrypoints/cli/serve.py

@@ -8,6 +8,8 @@ diffusion models (e.g., Qwen-Image) through the same CLI interface.
 import argparse
 import json
 import os
+import signal
+from types import FrameType
 from typing import Any
 
 import uvloop
@@ -419,22 +421,228 @@ def _create_default_diffusion_stage_cfg(args: argparse.Namespace) -> list[dict[s
 
 
 def run_headless(args: argparse.Namespace) -> None:
-    """Run a single stage in headless mode.
+    """Run a single stage in headless mode."""
+    from vllm.v1.engine.coordinator import DPCoordinator
+    from vllm.v1.engine.utils import CoreEngineProcManager
+    from vllm.v1.executor.multiproc_executor import MultiprocExecutor
+    from vllm.version import __version__ as VLLM_VERSION
 
-    .. deprecated:: 0.x.x
-        Headless mode is deprecated and will be removed in a future version.
-        It is only compatible with the old OmniStage-based runtime.
-        The current AsyncOmniEngine-based runtime does not support headless mode.
-
-    Raises:
-        RuntimeError: Always raises an error indicating headless mode is deprecated.
-    """
-    raise RuntimeError(
-        "Headless mode is deprecated and not supported in the current runtime. "
-        "Please use the standard orchestrator mode (without --headless flag). "
-        "If you need distributed deployment, consider using Ray backend or "
-        "other distributed serving solutions."
+    from vllm_omni.diffusion.stage_diffusion_proc import (
+        complete_diffusion_handshake,
+        spawn_diffusion_proc,
     )
+    from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_omni_kv_config_for_stage
+    from vllm_omni.engine.stage_engine_startup import register_stage_with_omni_master
+    from vllm_omni.engine.stage_init_utils import (
+        build_diffusion_config,
+        build_engine_args_dict,
+        build_vllm_config,
+        extract_stage_metadata,
+        get_stage_connector_spec,
+        inject_kv_stage_info,
+        load_omni_transfer_config_for_model,
+        prepare_engine_environment,
+        terminate_alive_proc,
+    )
+    from vllm_omni.entrypoints.utils import inject_omni_kv_config, load_and_resolve_stage_configs
+
+    model = args.model
+    stage_id: int | None = args.stage_id
+    omni_master_address: str | None = args.omni_master_address
+    omni_master_port: int | None = args.omni_master_port
+
+    if stage_id is None:
+        raise ValueError("--stage-id is required in headless mode")
+    if omni_master_address is None or omni_master_port is None:
+        raise ValueError("--omni-master-address and --omni-master-port are required in headless mode")
+    if getattr(args, "api_server_count", 0) and args.api_server_count > 1:
+        raise ValueError("api_server_count can't be set in headless mode")
+    if args.worker_backend != "multi_process":
+        raise ValueError("headless mode requires worker_backend=multi_process")
+
+    args_dict = vars(args).copy()
+    config_path, stage_configs = load_and_resolve_stage_configs(
+        model,
+        args_dict.get("stage_configs_path"),
+        args_dict,
+    )
+
+    # Locate the stage config that matches stage_id.
+    stage_cfg = None
+    for cfg in stage_configs:
+        if getattr(cfg, "stage_id", None) == stage_id:
+            stage_cfg = cfg
+            break
+    if stage_cfg is None:
+        raise ValueError(
+            f"No stage config found for stage_id={stage_id}. "
+            f"Available stage ids: {[getattr(c, 'stage_id', None) for c in stage_configs]}"
+        )
+
+    prepare_engine_environment()
+    omni_transfer_config = load_omni_transfer_config_for_model(model, config_path)
+    omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id)
+
+    if getattr(stage_cfg, "stage_type", "llm") == "diffusion":
+        metadata = extract_stage_metadata(stage_cfg)
+        if omni_conn_cfg:
+            inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
+        inject_kv_stage_info(stage_cfg, stage_id)
+        od_config = build_diffusion_config(model, stage_cfg, metadata)
+
+        logger.info(
+            "[Headless] Launching diffusion stage %d via OmniMasterServer at %s:%d",
+            stage_id,
+            omni_master_address,
+            omni_master_port,
+        )
+
+        proc = None
+        try:
+            handshake_address, request_address, response_address = register_stage_with_omni_master(
+                omni_master_address=omni_master_address,
+                omni_master_port=omni_master_port,
+                omni_stage_id=stage_id,
+                omni_stage_config=stage_cfg,
+                return_addresses=True,
+            )
+            proc, _, _, _ = spawn_diffusion_proc(
+                model,
+                od_config,
+                handshake_address=handshake_address,
+                request_address=request_address,
+                response_address=response_address,
+            )
+            complete_diffusion_handshake(proc, handshake_address)
+            proc.join()
+            if proc.exitcode not in (None, 0):
+                raise RuntimeError(f"Diffusion stage {stage_id} exited with code {proc.exitcode}")
+            return
+        finally:
+            logger.info("[Headless] Shutting down stage %d.", stage_id)
+            if proc is not None and proc.is_alive():
+                terminate_alive_proc(proc)
+
+    stage_connector_spec = get_stage_connector_spec(
+        omni_transfer_config=omni_transfer_config,
+        stage_id=stage_id,
+        async_chunk=False,
+    )
+
+    # Device assignment is managed externally (e.g. CUDA_VISIBLE_DEVICES);
+    # runtime_cfg is intentionally ignored in headless mode.
+    engine_args_dict = build_engine_args_dict(
+        stage_cfg,
+        model,
+        stage_connector_spec=stage_connector_spec,
+    )
+
+    # Inject omni KV connector config so the engine runner can initialize the
+    # correct connector (sender/receiver role, type, addresses, etc.).
+    if omni_conn_cfg:
+        omni_kv = engine_args_dict.get("omni_kv_config") or {}
+        if not isinstance(omni_kv, dict):
+            omni_kv = dict(omni_kv)
+        omni_kv["connector_config"] = omni_conn_cfg
+        omni_kv["omni_from_stage"] = omni_from
+        omni_kv["omni_to_stage"] = omni_to
+        omni_kv.setdefault("stage_id", stage_id)
+        engine_args_dict["omni_kv_config"] = omni_kv
+
+    vllm_config, executor_class = build_vllm_config(
+        stage_cfg,
+        model,
+        stage_connector_spec=stage_connector_spec,
+        engine_args_dict=engine_args_dict,
+        headless=True,
+    )
+    parallel_config = vllm_config.parallel_config
+    local_engine_count = parallel_config.data_parallel_size_local
+
+    if local_engine_count <= 0:
+        raise ValueError("data_parallel_size_local must be > 0 in headless mode")
+
+    shutdown_requested = False
+
+    def signal_handler(signum: int, frame: FrameType | None) -> None:
+        nonlocal shutdown_requested
+        logger.debug("Received %d signal.", signum)
+        if not shutdown_requested:
+            shutdown_requested = True
+            raise SystemExit
+
+    signal.signal(signal.SIGTERM, signal_handler)
+    signal.signal(signal.SIGINT, signal_handler)
+
+    if parallel_config.node_rank_within_dp > 0:
+        head_node_address = f"{parallel_config.master_addr}:{parallel_config.master_port}"
+        logger.info(
+            "Launching vLLM-Omni (v%s) headless multiproc executor, "
+            "with head node address %s for torch.distributed process group.",
+            VLLM_VERSION,
+            head_node_address,
+        )
+
+        executor = MultiprocExecutor(vllm_config, monitor_workers=False)
+        executor.start_worker_monitor(inline=True)
+        return
+
+    dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0
+    coordinator = None
+    if vllm_config.needs_dp_coordinator and dp_rank == 0:
+        coordinator = DPCoordinator(
+            parallel_config,
+            enable_wave_coordination=vllm_config.model_config.is_moe,
+        )
+        logger.info(
+            "[Headless] Started DP Coordinator process for stage %d (PID: %d)",
+            stage_id,
+            coordinator.proc.pid,
+        )
+
+    logger.info(
+        "[Headless] Launching %d engine core(s) for stage %d via OmniMasterServer at %s:%d",
+        local_engine_count,
+        stage_id,
+        omni_master_address,
+        omni_master_port,
+    )
+
+    # Headless mode launches all local engine cores for a single stage.
+    # The OmniMasterServer allocates one handshake endpoint per stage, so we
+    # register the stage once here and let every local engine core reuse the
+    # returned handshake address directly.
+    handshake_address = register_stage_with_omni_master(
+        omni_master_address=omni_master_address,
+        omni_master_port=omni_master_port,
+        omni_stage_id=stage_id,
+        omni_stage_config=stage_cfg,
+        coordinator=coordinator,
+    )
+
+    engine_manager = None
+    log_stats = bool(getattr(args, "log_stats", False))
+    if getattr(args, "disable_log_stats", False):
+        log_stats = False
+
+    try:
+        engine_manager = CoreEngineProcManager(
+            local_engine_count=local_engine_count,
+            start_index=dp_rank,
+            local_start_index=0,
+            vllm_config=vllm_config,
+            local_client=False,
+            handshake_address=handshake_address,
+            executor_class=executor_class,
+            log_stats=log_stats,
+        )
+        engine_manager.join_first()
+    finally:
+        logger.info("[Headless] Shutting down stage %d.", stage_id)
+        if engine_manager is not None:
+            engine_manager.shutdown()
+        if coordinator is not None:
+            coordinator.shutdown()
 
 
 def cmd_init() -> list[CLISubcommand]:

vllm_omni/entrypoints/openai/api_server.py

@@ -353,7 +353,8 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
     try:
         await shutdown_task
     finally:
-        serving_speech = getattr(getattr(app, "state", None), "openai_serving_speech", None)
+        state = getattr(app, "state", None)
+        serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None
         if serving_speech is not None:
             serving_speech.shutdown()
         sock.close()