[BugFix] Fix incorrect GPU device mapping in multi-replica stages (#4132)

Signed-off-by: ZhengWG <zwg0606@gmail.com>

tests/entrypoints/test_stage_utils.py

@@ -135,22 +135,6 @@ def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pyte
 # ---- _map_device_list unit tests ----
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_map_device_list_idempotency():
-    """Device IDs already in the visible set are returned as-is (idempotency)."""
-    result = _map_device_list(0, ["0", "1"], ["0", "1", "2", "3"])
-    assert result == ["0", "1"]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_map_device_list_idempotency_single():
-    """Single device ID in visible set passes through."""
-    result = _map_device_list(1, ["1"], ["0", "1"])
-    assert result == ["1"]
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_map_device_list_index_mapping():
@@ -211,6 +195,18 @@ def test_map_device_list_partial_mapping():
     assert result == ["5", "6"]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_multi_replica_offset_cvd():
+    """Logical indices from split_devices_for_replicas must map through an offset
+    CUDA_VISIBLE_DEVICES, not be treated as physical IDs.
+    """
+    visible = ["2", "3", "4", "5"]
+    assert _map_device_list(1, ["1"], visible) == ["3"]
+    assert _map_device_list(1, ["2"], visible) == ["4"]
+    assert _map_device_list(1, ["2", "3"], visible) == ["4", "5"]
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_map_device_list_raises_on_non_numeric():

vllm_omni/entrypoints/stage_utils.py

@@ -95,24 +95,6 @@ def _map_device_list(stage_id: int, device_list: list[str], visible_device_list:
     if not all(device.isdigit() for device in device_list):
         raise ValueError("Logical devices must be non-negative integers")
 
-    # Idempotency: if every requested id already names an entry in the visible
-    # pool (by value, not by index), treat the YAML field as physical and skip
-    # the remap. This avoids a double-mapping crash when a parent harness has
-    # already narrowed CUDA_VISIBLE_DEVICES to specific physical ids before
-    # spawning the subprocess (the test runtime's OmniServerStageCli does this
-    # — it sets each stage subprocess's env via ``_set_stage_device_env``).
-    # Without this guard, a YAML ``devices: "1"`` against a narrowed visible
-    # list ``["1"]`` would try ``visible[1]`` (out of range) and raise.
-    visible_set = set(visible_device_list)
-    if all(device in visible_set for device in device_list):
-        logger.info(
-            "Stage %s requested devices %s are already in the visible set %s; using as-is (no remap).",
-            stage_id,
-            device_list,
-            visible_device_list,
-        )
-        return list(device_list)
-
     logical_ids = [int(device) for device in device_list]
 
     # If ALL logical device IDs are >= num_visible, none can be 0-based indices