[Bugfix][Core] Fix replica device split when a stage omits devices (#5445)

Signed-off-by: Yash Jain <yashjainjain1704@gmail.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -15,6 +15,7 @@ from vllm_omni.engine.stage_init_utils import (
     ReplicaInitPlan,
     build_stage0_input_processor,
     compute_replica_layout,
+    split_devices_for_replicas,
 )
 from vllm_omni.engine.stage_runtime import StageRuntime
 
@@ -232,6 +233,76 @@ def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
     assert replica_devices_map == {0: ["0,1", "2,3"]}
 
 
+@pytest.mark.parametrize("num_replicas", [1, 2, 4])
+def test_split_devices_for_replicas_returns_one_slot_per_replica_when_devices_unset(num_replicas):
+    """``devices`` is optional, and the result is indexed by replica id.
+
+    A stage that declares ``num_replicas`` without ``devices`` lets every
+    replica inherit the launcher's CUDA_VISIBLE_DEVICES, so each one still
+    needs its own (empty) slot.
+    """
+    assert split_devices_for_replicas(None, num_replicas, 1, 0) == [None] * num_replicas
+
+
+def test_compute_replica_layout_covers_every_replica_when_devices_unset():
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={},
+        runtime=types.SimpleNamespace(num_replicas=3),
+    )
+
+    replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg])
+
+    assert replicas_per_stage == [3]
+    assert replica_devices_map == {0: [None, None, None]}
+
+
+def test_build_logical_stage_init_plans_handles_stage_without_devices(monkeypatch):
+    """Regression: a multi-replica stage with no ``devices`` must still plan.
+
+    ``_build_logical_stage_init_plans`` indexes ``replica_devices_map`` by
+    replica id, so a short list raised ``IndexError`` for every replica after
+    the first.
+    """
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={},
+        runtime=types.SimpleNamespace(num_replicas=3),
+    )
+    runtime = StageRuntime(
+        stage_configs=[stage_cfg],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+    )
+
+    monkeypatch.setattr(
+        runtime_mod,
+        "extract_legacy_stage_metadata",
+        lambda cfg: _make_llm_metadata(cfg.stage_id),
+    )
+    monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {})
+    monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
+    monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {})
+    monkeypatch.setattr(runtime_mod, "build_vllm_config", lambda *_args, **_kwargs: (types.SimpleNamespace(), object))
+
+    replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg])
+    stage_plans = runtime._build_logical_stage_init_plans(
+        omni_transfer_config=None,
+        replicas_per_stage=replicas_per_stage,
+        replica_devices_map=replica_devices_map,
+    )
+
+    assert [replica.replica_id for replica in stage_plans[0].replicas] == [0, 1, 2]
+    assert [replica.stage_cfg.runtime.devices for replica in stage_plans[0].replicas] == [None, None, None]
+
+
 def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
     shared = types.SimpleNamespace(name="shared")
     extra = types.SimpleNamespace(name="extra")

vllm_omni/engine/stage_init_utils.py

@@ -567,9 +567,14 @@ def split_devices_for_replicas(
     num_replicas: int,
     tp_size: int,
     stage_id: int,
-) -> list[str]:
+) -> list[str | None]:
     """Split a devices string into per-replica subsets.
 
+    The result always has one entry per replica, because callers index it by
+    replica id. When ``devices_str`` is ``None`` the stage declares no explicit
+    placement, so every replica gets ``None`` and inherits the launcher's
+    ``CUDA_VISIBLE_DEVICES``.
+
     When ``num_replicas`` is 1, returns ``[devices_str]`` unchanged.
     Otherwise, two YAML shapes are accepted:
 
@@ -597,8 +602,14 @@ def split_devices_for_replicas(
     Any other length raises ``ValueError`` (the two modes are
     length-disjoint for ``num_replicas > 1``).
     """
-    if num_replicas <= 1 or devices_str is None:
-        return [devices_str] if devices_str is not None else [devices_str]
+    if devices_str is None:
+        # No explicit placement: hand back one empty slot per replica, matching
+        # ``get_headless_replica_devices``. Returning a single-element list here
+        # made callers index past the end for every replica after the first.
+        return [None] * max(1, num_replicas)
+
+    if num_replicas <= 1:
+        return [devices_str]
 
     device_list = [d.strip() for d in devices_str.split(",") if d.strip()]
 
@@ -688,7 +699,7 @@ def compute_replica_layout(
             raise ValueError(f"num_replicas must be >= 0, got {num_replicas}")
         replicas_per_stage.append(num_replicas if allow_zero else max(1, num_replicas))
 
-    replica_devices_map: dict[int, list[str]] = {}
+    replica_devices_map: dict[int, list[str | None]] = {}
     for stage_id, stage_cfg in enumerate(stage_configs):
         num_replicas = replicas_per_stage[stage_id]
         if num_replicas <= 1:

vllm_omni/engine/stage_runtime.py

@@ -337,7 +337,7 @@ class StageRuntime:
         self,
         omni_transfer_config: Any,
         replicas_per_stage: Sequence[int],
-        replica_devices_map: Mapping[int, Sequence[str]],
+        replica_devices_map: Mapping[int, Sequence[str | None]],
     ) -> list[LogicalStageInitPlan]:
         """Build startup plans for every logical stage and replica."""
         stage_plans: list[LogicalStageInitPlan] = []