From 9a7f7b7de57c6aecc6f915913cd94c353e23886f Mon Sep 17 00:00:00 2001
From: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Date: Tue, 9 Dec 2025 21:04:30 +0800
Subject: [PATCH] [Bugfix]multi-GPU tp scenarios, devices: "0,1" uses physical
 IDs instead of logical IDs (#253)

Signed-off-by: David Chen <530634352@qq.com>
---
 tests/entrypoints/test_stage_utils.py         | 58 +++++++++++++++++++
 vllm_omni/entrypoints/stage_utils.py          | 32 ++++++++--
 .../stage_configs/qwen3_omni_moe.yaml         |  2 +-
 3 files changed, 86 insertions(+), 6 deletions(-)
 create mode 100644 tests/entrypoints/test_stage_utils.py

diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py
new file mode 100644
index 000000000..d4e718cad
--- /dev/null
+++ b/tests/entrypoints/test_stage_utils.py
@@ -0,0 +1,58 @@
+import os
+import sys
+
+import pytest
+
+from vllm_omni.entrypoints.stage_utils import set_stage_devices
+
+
+def _make_dummy_torch(call_log):
+    class _Props:
+        def __init__(self, total):
+            self.total_memory = total
+
+    class _Cuda:
+        @staticmethod
+        def is_available():
+            return True
+
+        @staticmethod
+        def set_device(idx):
+            call_log.append(idx)
+
+        @staticmethod
+        def device_count():
+            return 2
+
+        @staticmethod
+        def get_device_properties(idx):
+            return _Props(total=16000)
+
+        @staticmethod
+        def mem_get_info(idx):
+            return (8000, 16000)
+
+        @staticmethod
+        def get_device_name(idx):
+            return f"gpu-{idx}"
+
+    class _Torch:
+        cuda = _Cuda
+
+    return _Torch
+
+
+@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
+def test_set_stage_devices_respects_logical_ids(monkeypatch):
+    # Preserve an existing logical mapping and ensure devices "0,1" map through it.
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7")
+    call_log: list[int] = []
+    dummy_torch = _make_dummy_torch(call_log)
+    monkeypatch.setitem(sys.modules, "torch", dummy_torch)
+    monkeypatch.setattr("vllm_omni.utils.detect_device_type", lambda: "cuda")
+    monkeypatch.setattr("vllm_omni.utils.get_device_control_env_var", lambda: "CUDA_VISIBLE_DEVICES")
+
+    set_stage_devices(stage_id=0, devices="0,1")
+
+    assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
+    assert call_log and call_log[0] == 0  # current device set after remap
diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py
index eda037b68..8dbbd8083 100644
--- a/vllm_omni/entrypoints/stage_utils.py
+++ b/vllm_omni/entrypoints/stage_utils.py
@@ -27,8 +27,11 @@ def set_stage_devices(
     Args:
         stage_id: Stage identifier for logging
         devices: Device specification:
-            - Comma-separated string (e.g. "2,5,7"): set device visibility env var
-              exactly to this list; logical index 0 is used as current device.
+            - Comma-separated string (e.g. "2,5,7"): interpreted as logical
+              indices against the current device visibility env var (e.g.
+              CUDA_VISIBLE_DEVICES/ASCEND_RT_VISIBLE_DEVICES) when present;
+              falls back to physical IDs if no mapping exists. Logical index 0
+              is used as current device.
             - Integer or digit-string: treat as logical index (0-based) into the
               current device visibility mapping; map to physical device, then set
               env var to this single device.
@@ -81,16 +84,35 @@ def set_stage_devices(
         logical_idx: int | None = None
 
         if isinstance(devices, str) and "," in devices:
-            os.environ[env_var] = devices
             toks = [t.strip() for t in devices.split(",") if t.strip() != ""]
+            vis = os.environ.get(env_var)
+            mapped_devices: list[str] = []
+            mapping: list[int] = []
+            if vis:
+                try:
+                    mapping = [int(x) for x in vis.split(",") if x.strip() != ""]
+                except Exception as e:
+                    logger.debug("[Stage-%s] Failed to parse existing %s: %s", stage_id, env_var, e)
+            for tok in toks:
+                try:
+                    idx = int(tok)
+                except Exception:
+                    mapped_devices.append(tok)
+                    continue
+                if mapping and 0 <= idx < len(mapping):
+                    mapped_devices.append(str(mapping[idx]))
+                else:
+                    mapped_devices.append(str(idx))
+            mapped_devices_str = ",".join(mapped_devices)
+            os.environ[env_var] = mapped_devices_str
             if toks:
                 try:
-                    selected_physical = int(toks[0])
+                    selected_physical = int(mapped_devices[0])
                     logger.debug(
                         "[Stage-%s] Set %s to %s; logical 0 -> physical %s",
                         stage_id,
                         env_var,
-                        devices,
+                        mapped_devices_str,
                         selected_physical,
                     )
                 except Exception as e:
diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
index 6480b7012..8178721dd 100644
--- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
@@ -44,7 +44,7 @@ stage_args:
        model_arch: Qwen3OmniMoeForConditionalGeneration
        worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.5
+       gpu_memory_utilization: 0.3
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav