From ebfbcfe46aa895d428933244908bae08b1ca6397 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Sat, 20 Jun 2026 16:38:10 -0400 Subject: [PATCH] Stop setting CUDA_VISIBLE_DEVICES internally in vLLM, add device_ids arg (#45026) Signed-off-by: Tyler Michael Smith Co-authored-by: Claude Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: kourosh hakhamaneshi --- .../attention/mla/sm100_cutlass_mla_kernel.cu | 11 +- tests/engine/test_arg_utils.py | 193 ++++++++++++++++++ .../entrypoints/openai/test_dp_supervisor.py | 4 +- vllm/config/parallel.py | 9 + .../device_communicators/all2all.py | 9 +- .../device_communicators/all_reduce_utils.py | 33 ++- .../device_communicators/custom_all_reduce.py | 18 +- .../device_communicators/quick_all_reduce.py | 10 +- .../device_communicators/shm_broadcast.py | 8 +- .../v1/lmcache_integration/vllm_v1_adapter.py | 9 +- vllm/distributed/parallel_state.py | 28 ++- vllm/distributed/stateless_coordinator.py | 22 +- vllm/engine/arg_utils.py | 58 ++++++ vllm/entrypoints/openai/dp_supervisor.py | 37 ++-- vllm/platforms/cuda.py | 9 + vllm/platforms/interface.py | 103 +++++++++- vllm/v1/engine/core.py | 25 ++- vllm/v1/engine/utils.py | 109 ++++++---- vllm/v1/executor/multiproc_executor.py | 10 + vllm/v1/executor/ray_executor.py | 53 ++--- vllm/v1/executor/ray_executor_v2.py | 87 +++++--- vllm/v1/executor/ray_utils.py | 8 +- vllm/v1/worker/gpu_worker.py | 48 ++++- vllm/v1/worker/worker_base.py | 6 + 24 files changed, 722 insertions(+), 185 deletions(-) diff --git a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu b/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu index 55d75383476..de62052b4b0 100644 --- a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu +++ b/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu @@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba using TileShapeD = typename MlaSm100Type::TileShapeD; arguments.problem_shape = cute::make_tuple(TileShapeH{}, static_cast(max_seq_len), TileShapeD{}, static_cast(num_batches)); - // Assumes device 0 when getting sm_count. - arguments.hw_info.sm_count = - sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count; + if (sm_count <= 0) { + int current_device = 0; + cudaGetDevice(¤t_device); + arguments.hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device); + } else { + arguments.hw_info.sm_count = sm_count; + } arguments.split_kv = static_cast(num_kv_splits); MlaSm100Type::Fmha::set_split_kv(arguments); diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 9d34975032e..a35f4453027 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch): args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer") assert args.model == "s3://bucket/model" assert args.tokenizer == "s3://bucket/tokenizer" + + +class TestDeviceIds: + def test_device_ids_with_cvd_out_of_range(self, monkeypatch): + """--device-ids index beyond the CVD set raises ValueError.""" + from vllm.platforms import current_platform + + key = current_platform.device_control_env_var + monkeypatch.setenv(key, "4,5") + args = EngineArgs(model="m", device_ids=[0, 2]) + with pytest.raises(ValueError, match="out of range"): + args._resolve_device_ids() + + def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch): + """--device-ids are CVD-local indices resolved to physical ids.""" + from vllm.platforms import current_platform + + key = current_platform.device_control_env_var + monkeypatch.setenv(key, "4,5") + args = EngineArgs(model="m", device_ids=[0, 1]) + assert args._resolve_device_ids() == [4, 5] + + def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch): + """--device-ids support UUID CVD values resolved by the platform.""" + from vllm.platforms import current_platform + + key = current_platform.device_control_env_var + monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890") + monkeypatch.setattr( + type(current_platform), + "device_control_id_to_physical_device_id", + classmethod( + lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id] + ), + ) + + args = EngineArgs(model="m", device_ids=[0, 1]) + assert args._resolve_device_ids() == [4, 5] + + def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch): + """UUID --device-ids are resolved to physical IDs immediately.""" + from vllm.platforms import current_platform + + monkeypatch.setattr( + type(current_platform), + "device_control_id_to_physical_device_id", + classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]), + ) + + args = EngineArgs(model="m", device_ids=["GPU-abcd1234"]) + assert args._resolve_device_ids() == [4] + + def test_device_ids_reject_mixed_integer_and_uuid_args(self): + """--device-ids must not mix CVD indices and UUIDs.""" + args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"]) + with pytest.raises(ValueError, match="must not mix"): + args._resolve_device_ids() + + def test_no_device_ids(self): + """No --device-ids returns None.""" + args = EngineArgs(model="m") + assert args._resolve_device_ids() is None + + def test_cli_parsing(self): + """--device-ids parses comma-separated string from CLI.""" + parser = FlexibleArgumentParser() + EngineArgs.add_cli_args(parser) + parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"]) + assert parsed.device_ids == [0, 2, 4] + + def test_cli_parsing_uuid(self): + """--device-ids parses comma-separated UUID strings from CLI.""" + parser = FlexibleArgumentParser() + EngineArgs.add_cli_args(parser) + parsed = parser.parse_args( + ["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"] + ) + assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"] + + def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch): + """assigned_physical_gpu_ids are already physical and not composed with CVD.""" + import vllm.platforms.interface as platform_interface + from vllm.platforms import current_platform + + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5]) + monkeypatch.setenv(current_platform.device_control_env_var, "4,5") + + assert current_platform.device_id_to_physical_device_id(0) == 4 + assert current_platform.device_id_to_physical_device_id(1) == 5 + assert current_platform.logical_device_id_to_visible_device_id(0) == 0 + assert current_platform.logical_device_id_to_visible_device_id(1) == 1 + + def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch): + """Physical IDs map back to visible ordinals when CVD uses UUIDs.""" + import vllm.platforms.interface as platform_interface + from vllm.platforms import current_platform + + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5]) + monkeypatch.setenv( + current_platform.device_control_env_var, + "GPU-abcd1234,GPU-ef567890", + ) + monkeypatch.setattr( + type(current_platform), + "device_control_id_to_physical_device_id", + classmethod( + lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id] + ), + ) + + assert current_platform.logical_device_id_to_visible_device_id(0) == 1 + + def test_device_ids_reject_duplicates(self): + """--device-ids must not contain duplicate entries.""" + args = EngineArgs(model="m", device_ids=[2, 2]) + with pytest.raises(ValueError, match="duplicates"): + args._resolve_device_ids() + + def test_cli_parsing_strips_whitespace(self): + """--device-ids tolerates whitespace around commas.""" + parser = FlexibleArgumentParser() + EngineArgs.add_cli_args(parser) + parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"]) + assert parsed.device_ids == [0, 2, 4] + + def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch): + """visible_device_id_to_physical_device_id maps torch device ordinals, + independent of the logical-to-physical mapping. + + Regression test: CustomAllreduce passes device.index (a visible + ordinal) and must not index into assigned_physical_gpu_ids, which + raised IndexError for non-identity --device-ids like [2, 3]. + """ + import vllm.platforms.interface as platform_interface + from vllm.platforms import current_platform + + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3]) + monkeypatch.delenv(current_platform.device_control_env_var, raising=False) + + # CVD unset: visible ordinal == physical ID, even beyond the + # assigned list's length. + assert current_platform.visible_device_id_to_physical_device_id(2) == 2 + assert current_platform.visible_device_id_to_physical_device_id(3) == 3 + + monkeypatch.setenv(current_platform.device_control_env_var, "4,5") + assert current_platform.visible_device_id_to_physical_device_id(1) == 5 + with pytest.raises(IndexError, match="out of range"): + current_platform.visible_device_id_to_physical_device_id(2) + + +class TestDpDeviceIdSharding: + def test_dp_supervisor_device_ids_stay_env_relative(self): + """Regression test: the DP supervisor must pass env-relative indices, + not physical IDs, because each child re-resolves --device-ids + against its inherited device-control env var.""" + import argparse + + from vllm.entrypoints.openai.dp_supervisor import _build_device_ids + + args = argparse.Namespace( + tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None + ) + assert _build_device_ids(args, local_rank=0) == [0, 1] + assert _build_device_ids(args, local_rank=1) == [2, 3] + + def test_dp_supervisor_shards_user_device_ids(self): + """User-provided --device-ids are sharded across DP children.""" + import argparse + + from vllm.entrypoints.openai.dp_supervisor import _build_device_ids + + args = argparse.Namespace( + tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7] + ) + assert _build_device_ids(args, local_rank=0) == [4, 5] + assert _build_device_ids(args, local_rank=1) == [6, 7] + with pytest.raises(ValueError, match="needs devices"): + _build_device_ids(args, local_rank=2) + + def test_dp_rank_shards_user_assigned_gpu_ids(self): + """get_physical_gpu_ids_for_local_dp_rank slices the user-provided + --device-ids list instead of recomputing from the env var.""" + from vllm.platforms import current_platform + from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank + + evar = current_platform.device_control_env_var + assert get_physical_gpu_ids_for_local_dp_rank( + evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7] + ) == [6, 7] + with pytest.raises(ValueError, match="needs devices"): + get_physical_gpu_ids_for_local_dp_rank( + evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7] + ) diff --git a/tests/entrypoints/openai/test_dp_supervisor.py b/tests/entrypoints/openai/test_dp_supervisor.py index 9967e6d86d0..1dd6537f201 100644 --- a/tests/entrypoints/openai/test_dp_supervisor.py +++ b/tests/entrypoints/openai/test_dp_supervisor.py @@ -364,7 +364,7 @@ class MockVLLMServer: await self._serve_task -def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]): +def launch_mock_vllm(child_args: argparse.Namespace): logger.info("Launching mock vLLM on port %s", child_args.port) mock_vllm = MockVLLMServer( port=child_args.port, @@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str] def launch_mock_vllm_with_drain( - child_args: argparse.Namespace, env_updates: dict[str, str] + child_args: argparse.Namespace, ): logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port) mock_vllm = MockVLLMServer( diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index a194640f2ec..2ae773d79c7 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -302,6 +302,14 @@ class ParallelConfig: Each entry must use `numactl --physcpubind` CPU-list syntax, for example `"0-3"` or `"0,2,4-7"`. """ + assigned_physical_gpu_ids: list[int] | None = None + """Mapping from vLLM-local logical GPU IDs to physical GPU IDs. + + For example, ``[2, 3]`` means logical GPU 0 maps to physical GPU 2, + and logical GPU 1 maps to physical GPU 3. Physical IDs are used only + at platform/topology boundaries such as NVML, NIC affinity, P2P + checks, and final CUDA device selection when needed. When None, + logical IDs map to visible device IDs in order.""" distributed_timeout_seconds: int | None = None """Timeout in seconds for distributed operations (e.g., init_process_group). @@ -772,6 +780,7 @@ class ParallelConfig: "numa_bind", "numa_bind_nodes", "numa_bind_cpus", + "assigned_physical_gpu_ids", } from vllm.config.utils import get_hash_factors, hash_factors diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 967ce5d75c3..0066a60dd02 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -704,7 +704,14 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): self.num_experts = num_experts self.cleanup() - gpus_per_node = torch.accelerator.device_count() + from vllm.platforms.interface import get_assigned_physical_gpu_ids + + assigned_physical_gpu_ids = get_assigned_physical_gpu_ids() + gpus_per_node = ( + len(assigned_physical_gpu_ids) + if assigned_physical_gpu_ids is not None + else torch.accelerator.device_count() + ) logger.debug( "Making One-sided NVLink mapping: rank=%d, world size=%d", self.rank, diff --git a/vllm/distributed/device_communicators/all_reduce_utils.py b/vllm/distributed/device_communicators/all_reduce_utils.py index cebf2c49b44..d50d84fa5ca 100644 --- a/vllm/distributed/device_communicators/all_reduce_utils.py +++ b/vllm/distributed/device_communicators/all_reduce_utils.py @@ -320,13 +320,21 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool: is_distributed = dist.is_initialized() - num_dev = current_platform.device_count() - cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES - if cuda_visible_devices is None: - cuda_visible_devices = ",".join(str(i) for i in range(num_dev)) + from vllm.platforms.interface import get_assigned_physical_gpu_ids + + assigned_physical_gpu_ids = get_assigned_physical_gpu_ids() + if assigned_physical_gpu_ids is not None: + # Key by the ordered list: the cache stores directed local-index + # pairs, so permutations of the same set are distinct mappings. + cache_key = ",".join(str(i) for i in assigned_physical_gpu_ids) + num_dev = len(assigned_physical_gpu_ids) + else: + num_dev = current_platform.device_count() + cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES + cache_key = cuda_visible_devices or ",".join(str(i) for i in range(num_dev)) path = os.path.join( - envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json" + envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cache_key}.json" ) os.makedirs(os.path.dirname(path), exist_ok=True) from vllm.distributed.parallel_state import get_world_group @@ -338,7 +346,15 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool: # enter this block to calculate the cache logger.info("generating GPU P2P access cache in %s", path) cache: dict[str, bool] = {} - ids = list(range(num_dev)) + # The probe subprocesses inherit this process's device-control env + # var, so they must be given visible ordinals, not physical IDs. + if assigned_physical_gpu_ids is not None: + ids = [ + current_platform.logical_device_id_to_visible_device_id(local) + for local in range(num_dev) + ] + else: + ids = list(range(num_dev)) # batch of all pairs of GPUs batch_src, batch_tgt = zip(*list(product(ids, ids))) # NOTE: we use `subprocess` rather than `multiprocessing` here @@ -368,8 +384,11 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool: ) from e with open(output_file.name, "rb") as f: result = pickle.load(f) + # Cache entries must be keyed by local indices (0..N-1) because + # gpu_p2p_access_check() is called with local ranks. + id_to_local = {device_id: local for local, device_id in enumerate(ids)} for _i, _j, r in zip(batch_src, batch_tgt, result): - cache[f"{_i}->{_j}"] = r + cache[f"{id_to_local[_i]}->{id_to_local[_j]}"] = r with open(path, "w") as f: json.dump(cache, f, indent=4) if is_distributed: diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index c57cc74fc06..95db6cc9245 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -34,7 +34,12 @@ def _can_p2p(rank: int, world_size: int) -> bool: continue if envs.VLLM_SKIP_P2P_CHECK: logger.debug("Skipping P2P check and trusting the driver's P2P report.") - return torch.cuda.can_device_access_peer(rank, i) + # can_device_access_peer takes visible device ordinals, while + # rank and i are logical local IDs. + return torch.cuda.can_device_access_peer( + current_platform.logical_device_id_to_visible_device_id(rank), + current_platform.logical_device_id_to_visible_device_id(i), + ) if not gpu_p2p_access_check(rank, i): return False return True @@ -126,13 +131,10 @@ class CustomAllreduce: CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size], max_size, ) - cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES - if cuda_visible_devices: - device_ids = list(map(int, cuda_visible_devices.split(","))) - else: - device_ids = list(range(current_platform.device_count())) - - physical_device_id = device_ids[device.index] + # device.index is a visible ordinal, not a logical local ID. + physical_device_id = current_platform.visible_device_id_to_physical_device_id( + device.index + ) tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") gather_list = [ torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size) diff --git a/vllm/distributed/device_communicators/quick_all_reduce.py b/vllm/distributed/device_communicators/quick_all_reduce.py index 8c7ee7452f1..c54eaf7555d 100644 --- a/vllm/distributed/device_communicators/quick_all_reduce.py +++ b/vllm/distributed/device_communicators/quick_all_reduce.py @@ -129,12 +129,10 @@ class QuickAllReduce: assert isinstance(device, torch.device) self.device = device - cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES - if cuda_visible_devices: - device_ids = list(map(int, cuda_visible_devices.split(","))) - else: - device_ids = list(range(current_platform.device_count())) - physical_device_id = device_ids[device.index] + # device.index is a visible ordinal, not a logical local ID. + physical_device_id = current_platform.visible_device_id_to_physical_device_id( + device.index + ) tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") gather_list = [ torch.tensor([0], dtype=torch.int, device="cpu") diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 9482568461c..43e066c44b0 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -840,7 +840,13 @@ class MessageQueue: The MessageQueue instance for the calling process, and a list of handles (only non-empty for the reader process). """ - local_size = current_platform.device_count() + from vllm.platforms.interface import get_assigned_physical_gpu_ids + + assigned_physical_gpu_ids = get_assigned_physical_gpu_ids() + if assigned_physical_gpu_ids is not None: + local_size = len(assigned_physical_gpu_ids) + else: + local_size = current_platform.device_count() rank = dist.get_rank() same_node = rank // local_size == reader_rank // local_size buffer_io = MessageQueue( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py index d16fbee585a..d72ebb5cd1e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py @@ -482,10 +482,11 @@ def _init_lmcache_engine( ) # Change current device. - num_gpus = torch.accelerator.device_count() - local_rank = parallel_config.rank % num_gpus - torch.accelerator.set_device_index(local_rank) - device = torch.device(f"cuda:{local_rank}") + from vllm.distributed.parallel_state import get_world_group + + device_index = get_world_group().device_index + torch.accelerator.set_device_index(device_index) + device = torch.device(f"cuda:{device_index}") metadata = LMCacheEngineMetadata( model_config.model, parallel_config.world_size, diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 8bd6e92157a..11b9e24e864 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -392,6 +392,14 @@ class GroupCoordinator: self.rank = torch.distributed.get_rank() self.local_rank = local_rank + self.device_index: int + if _WORLD is not None: + self.device_index = _WORLD.device_index + else: + assert local_rank >= 0, ( + "local_rank must be provided when creating the world group" + ) + self.device_index = local_rank self_device_group = None self_cpu_group = None @@ -442,11 +450,18 @@ class GroupCoordinator: from vllm.platforms import current_platform if current_platform.is_cuda_alike(): - self.device = torch.device(f"cuda:{local_rank}") + visible_device_index = ( + current_platform.logical_device_id_to_visible_device_id( + self.device_index + ) + ) + self.device = torch.device(f"cuda:{visible_device_index}") elif current_platform.is_xpu(): - self.device = torch.device(f"xpu:{local_rank}") + self.device = torch.device(f"xpu:{self.device_index}") elif current_platform.is_out_of_tree(): - self.device = torch.device(f"{current_platform.device_name}:{local_rank}") + self.device = torch.device( + f"{current_platform.device_name}:{self.device_index}" + ) else: self.device = torch.device("cpu") @@ -1438,7 +1453,12 @@ def _init_process_group_for_split_group( """ if torch.accelerator.is_available() and backend != "gloo": init_backend = "cpu:gloo,cuda:nccl" - device_id: torch.device | None = torch.device(f"cuda:{local_rank}") + from vllm.platforms import current_platform + + visible_device_index = current_platform.logical_device_id_to_visible_device_id( + local_rank + ) + device_id: torch.device | None = torch.device(f"cuda:{visible_device_index}") else: init_backend = "gloo" device_id = None diff --git a/vllm/distributed/stateless_coordinator.py b/vllm/distributed/stateless_coordinator.py index 549284df32d..5f4597d07cb 100644 --- a/vllm/distributed/stateless_coordinator.py +++ b/vllm/distributed/stateless_coordinator.py @@ -86,6 +86,15 @@ class StatelessGroupCoordinator(GroupCoordinator): self.rank = global_rank self.local_rank = local_rank + from vllm.distributed.parallel_state import _WORLD + + if _WORLD is not None: + self.device_index = _WORLD.device_index + else: + assert local_rank >= 0, ( + "local_rank must be provided when creating the world group" + ) + self.device_index = local_rank self_device_group = None self_cpu_group = None @@ -152,11 +161,18 @@ class StatelessGroupCoordinator(GroupCoordinator): self.tcp_store_group = self_tcp_store_group if current_platform.is_cuda_alike(): - self.device = torch.device(f"cuda:{local_rank}") + visible_device_index = ( + current_platform.logical_device_id_to_visible_device_id( + self.device_index + ) + ) + self.device = torch.device(f"cuda:{visible_device_index}") elif current_platform.is_xpu(): - self.device = torch.device(f"xpu:{local_rank}") + self.device = torch.device(f"xpu:{self.device_index}") elif current_platform.is_out_of_tree(): - self.device = torch.device(f"{current_platform.device_name}:{local_rank}") + self.device = torch.device( + f"{current_platform.device_name}:{self.device_index}" + ) else: self.device = torch.device("cpu") diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 9172a8728a0..921f31466b3 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -6,6 +6,7 @@ import copy import dataclasses import functools import json +import os import sys from collections.abc import Callable from dataclasses import MISSING, asdict, dataclass, fields, is_dataclass @@ -465,6 +466,7 @@ class EngineArgs: numa_bind: bool = ParallelConfig.numa_bind numa_bind_nodes: list[int] | None = ParallelConfig.numa_bind_nodes numa_bind_cpus: list[str] | None = ParallelConfig.numa_bind_cpus + device_ids: list[int | str] | None = None tensor_parallel_size: int = ParallelConfig.tensor_parallel_size prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size @@ -979,6 +981,20 @@ class EngineArgs: parallel_group.add_argument( "--numa-bind-cpus", **parallel_kwargs["numa_bind_cpus"] ) + parallel_group.add_argument( + "--device-ids", + type=lambda s: [ + int(device_id) if device_id.isdigit() else device_id + for device_id in (part.strip() for part in s.split(",")) + ], + default=None, + help="Comma-separated physical GPU device IDs or UUIDs to use " + '(e.g. --device-ids "2,3,5,7"). Avoids setting ' + "CUDA_VISIBLE_DEVICES, preserving full GPU topology " + "visibility for GPU-NIC affinity and DeepGEMM. " + "Note: has no effect with Ray executors; use Ray " + "placement groups for GPU selection instead.", + ) parallel_group.add_argument( "--tensor-parallel-size", "-tp", **parallel_kwargs["tensor_parallel_size"] ) @@ -1716,6 +1732,47 @@ class EngineArgs: ) return SpeculativeConfig(**self.speculative_config) + def _resolve_device_ids(self) -> list[int] | None: + if not self.device_ids: + return None + if self.distributed_executor_backend == "ray": + logger.warning( + "--device-ids has no effect when using the Ray executor. " + "Use Ray placement groups for GPU selection instead." + ) + ids = self.device_ids + if len(set(ids)) != len(ids): + raise ValueError(f"--device-ids must not contain duplicates: {ids}") + if all(isinstance(i, str) for i in ids): + return [ + current_platform.device_control_id_to_physical_device_id(i) + for i in cast(list[str], ids) + ] + if any(isinstance(i, str) for i in ids): + raise ValueError("--device-ids must not mix integer IDs and UUIDs") + int_ids = cast(list[int], ids) + # Compose with CUDA_VISIBLE_DEVICES: if CVD is set, treat + # --device-ids values as indices into the CVD-visible set. + cvd = getattr( + envs, + current_platform.device_control_env_var, + os.environ.get(current_platform.device_control_env_var), + ) + if cvd: + cvd_ids = [ + current_platform.device_control_id_to_physical_device_id(x) + for x in cvd.split(",") + ] + for i in int_ids: + if i >= len(cvd_ids): + raise ValueError( + f"--device-ids index {i} is out of range for " + f"{current_platform.device_control_env_var}" + f"={cvd} ({len(cvd_ids)} devices visible)" + ) + return [cvd_ids[i] for i in int_ids] + return int_ids + def create_diffusion_config(self) -> DiffusionConfig | None: if self.diffusion_config is None: return None @@ -2029,6 +2086,7 @@ class EngineArgs: cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, _api_process_count=self._api_process_count, _api_process_rank=self._api_process_rank, + assigned_physical_gpu_ids=self._resolve_device_ids(), numa_bind=self.numa_bind, numa_bind_nodes=self.numa_bind_nodes, numa_bind_cpus=self.numa_bind_cpus, diff --git a/vllm/entrypoints/openai/dp_supervisor.py b/vllm/entrypoints/openai/dp_supervisor.py index 13444015ecc..73b10a04ea5 100644 --- a/vllm/entrypoints/openai/dp_supervisor.py +++ b/vllm/entrypoints/openai/dp_supervisor.py @@ -23,12 +23,10 @@ import uvloop from fastapi import FastAPI, Response from vllm.logger import init_logger -from vllm.platforms import current_platform from vllm.utils.system_utils import ( decorate_logs, kill_process_tree, set_process_title, - update_environment_variables, ) logger = init_logger(__name__) @@ -127,22 +125,29 @@ def _build_vllm_dp_server_args( child_args.data_parallel_multi_port_external_lb = False child_args.data_parallel_supervisor_port = None child_args.api_server_count = 1 + child_args.device_ids = _build_device_ids(args, local_rank) return child_args -def _build_vllm_dp_server_env( - args: argparse.Namespace, local_rank: int -) -> dict[str, str]: - # set visible devices for the child process +def _build_device_ids(args: argparse.Namespace, local_rank: int) -> list[int | str]: + """Build the --device-ids value for a DP child process. + + The child resolves these against its own inherited device-control env + var (e.g. CUDA_VISIBLE_DEVICES), so integer IDs must stay env-relative + here rather than being translated to physical IDs. + """ devices_per_rank = args.tensor_parallel_size * args.pipeline_parallel_size start = local_rank * devices_per_rank stop = start + devices_per_rank - device_env = current_platform.device_control_env_var - visible_devices = ",".join( - str(current_platform.device_id_to_physical_device_id(idx)) - for idx in range(start, stop) - ) - return {device_env: visible_devices} + device_ids = getattr(args, "device_ids", None) + if device_ids is not None: + if stop > len(device_ids): + raise ValueError( + f"--device-ids has {len(device_ids)} entries, but DP rank " + f"{local_rank} needs devices [{start}, {stop})" + ) + return device_ids[start:stop] + return list(range(start, stop)) def _child_base_url(args: argparse.Namespace, port: int) -> str: @@ -228,9 +233,7 @@ def _build_dp_supervisor_app(supervisor: DPSupervisor) -> FastAPI: return app -def _run_vllm_dp_server( - child_args: argparse.Namespace, env_updates: dict[str, str] -) -> None: +def _run_vllm_dp_server(child_args: argparse.Namespace) -> None: """ Entrypoint function for the vLLM DP Server. """ @@ -241,7 +244,6 @@ def _run_vllm_dp_server( os.setpgrp() name = f"APIServer_DP{child_args.data_parallel_rank}" - update_environment_variables(env_updates) set_process_title(name) decorate_logs(name) uvloop.run(run_server(child_args)) @@ -345,11 +347,10 @@ class DPSupervisor: context = multiprocessing.get_context("spawn") for local_rank in range(self.args.data_parallel_size_local): child_args = _build_vllm_dp_server_args(self.args, local_rank) - child_env = _build_vllm_dp_server_env(self.args, local_rank) process = context.Process( target=_run_vllm_dp_server, name=f"APIServer_DPRank_{child_args.data_parallel_rank}", - args=(child_args, child_env), + args=(child_args,), ) process.start() self._processes.append(process) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 30a16e27469..6bf1793eefd 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -685,6 +685,15 @@ class CudaPlatformBase(Platform): # all the related functions work on real physical device ids. # the major benefit of using NVML is that it will not initialize CUDA class NvmlCudaPlatform(CudaPlatformBase): + @classmethod + @with_nvml_context + def device_control_id_to_physical_device_id(cls, device_id: str) -> int: + try: + return int(device_id) + except ValueError: + handle = pynvml.nvmlDeviceGetHandleByUUID(device_id) + return pynvml.nvmlDeviceGetIndex(handle) + @classmethod @cache @with_nvml_context diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 7fed06950bd..82c87416093 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -30,6 +30,33 @@ else: logger = init_logger(__name__) +_assigned_physical_gpu_ids: list[int] | None = None + + +def set_assigned_physical_gpu_ids(ids: list[int]) -> None: + """Set the physical GPU IDs assigned to this worker process. + Called during worker init so that device_id_to_physical_device_id() + can map local_rank to the correct physical device without relying + on CUDA_VISIBLE_DEVICES. + + Idempotent: a second call with the same value is a no-op. + Raises RuntimeError if called again with a different value. + + This is expected to run during single-threaded worker initialization.""" + global _assigned_physical_gpu_ids + if _assigned_physical_gpu_ids is not None: + if _assigned_physical_gpu_ids != ids: + raise RuntimeError( + f"set_assigned_physical_gpu_ids called with conflicting values: " + f"existing={_assigned_physical_gpu_ids}, new={ids}" + ) + return + _assigned_physical_gpu_ids = ids + + +def get_assigned_physical_gpu_ids() -> list[int] | None: + return _assigned_physical_gpu_ids + @functools.cache def in_wsl() -> bool: @@ -233,8 +260,34 @@ class Platform: """ import vllm.kernels # noqa: F401 + @classmethod + def device_control_id_to_physical_device_id(cls, device_id: str) -> int: + """Map one device-control env entry to an integer physical device ID.""" + try: + return int(device_id) + except ValueError as e: + raise ValueError( + f"Non-integer device ID {device_id!r} is not supported by " + f"{cls.device_name}." + ) from e + @classmethod def device_id_to_physical_device_id(cls, device_id: int): + """Map a vLLM-local logical device ID to a physical device ID. + + The input is a logical local ID (e.g. a local rank), NOT a visible + device ordinal; for the latter use + visible_device_id_to_physical_device_id(). The two coincide only + when no logical-to-physical mapping is in effect. + """ + if _assigned_physical_gpu_ids is not None: + if device_id >= len(_assigned_physical_gpu_ids): + raise IndexError( + f"device_id {device_id} is out of range for " + f"assigned_physical_gpu_ids {_assigned_physical_gpu_ids} " + f"({len(_assigned_physical_gpu_ids)} devices assigned)" + ) + return _assigned_physical_gpu_ids[device_id] # Treat empty device control env var as unset. This is a valid # configuration in Ray setups where the engine is launched in # a CPU-only placement group located on a GPU node. @@ -244,10 +297,58 @@ class Platform: ): device_ids = os.environ[cls.device_control_env_var].split(",") physical_device_id = device_ids[device_id] - return int(physical_device_id) + return cls.device_control_id_to_physical_device_id(physical_device_id) else: return device_id + @classmethod + def logical_device_id_to_visible_device_id(cls, device_id: int) -> int: + """Map a vLLM-local logical device ID to the current process's + visible accelerator ordinal. + + vLLM internals use logical local IDs. Physical IDs are used only + at platform/topology boundaries. This helper performs the final + translation needed by APIs such as ``torch.device("cuda:N")``. + """ + physical_device_id = cls.device_id_to_physical_device_id(device_id) + device_control_env = os.environ.get(cls.device_control_env_var, "") + if not device_control_env: + return physical_device_id + + visible_physical_device_ids = [ + cls.device_control_id_to_physical_device_id(physical_id) + for physical_id in device_control_env.split(",") + ] + if physical_device_id not in visible_physical_device_ids: + raise RuntimeError( + f"Physical device {physical_device_id} for logical device " + f"{device_id} is not visible in {cls.device_control_env_var}=" + f"{device_control_env}" + ) + return visible_physical_device_ids.index(physical_device_id) + + @classmethod + def visible_device_id_to_physical_device_id(cls, device_id: int) -> int: + """Map a visible accelerator ordinal (e.g. ``torch.device.index``) + to a physical device ID. + + This is the inverse of the env-var translation performed by + logical_device_id_to_visible_device_id() and is independent of any + logical-to-physical mapping set via set_assigned_physical_gpu_ids(). + """ + device_control_env = os.environ.get(cls.device_control_env_var, "") + if not device_control_env: + return device_id + visible_device_ids = device_control_env.split(",") + if device_id >= len(visible_device_ids): + raise IndexError( + f"visible device ordinal {device_id} is out of range for " + f"{cls.device_control_env_var}={device_control_env}" + ) + return cls.device_control_id_to_physical_device_id( + visible_device_ids[device_id] + ) + @classmethod def import_kernels(cls) -> None: """Import any platform-specific C kernels.""" diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index ac7037800a0..8f6baa46936 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -74,7 +74,7 @@ from vllm.v1.engine.utils import ( EngineHandshakeMetadata, EngineZmqAddresses, SignalCallback, - get_device_indices, + get_physical_gpu_ids_for_local_dp_rank, ) from vllm.v1.executor import Executor from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind @@ -2175,23 +2175,30 @@ class EngineCoreActorMixin: pass else: device_control_env_var = current_platform.device_control_env_var - self._set_cuda_visible_devices( + self._set_assigned_physical_gpu_ids( vllm_config, local_dp_rank, device_control_env_var ) - def _set_cuda_visible_devices( - self, vllm_config: VllmConfig, local_dp_rank: int, device_control_env_var: str + def _set_assigned_physical_gpu_ids( + self, + vllm_config: VllmConfig, + local_dp_rank: int, + device_control_env_var: str, ): world_size = vllm_config.parallel_config.world_size - # Set CUDA_VISIBLE_DEVICES or equivalent. try: - value = get_device_indices( - device_control_env_var, local_dp_rank, world_size + physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank( + device_control_env_var, + local_dp_rank, + world_size, + user_assigned_gpu_ids=( + vllm_config.parallel_config.assigned_physical_gpu_ids + ), ) - os.environ[device_control_env_var] = value + vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids except IndexError as e: raise Exception( - f"Error setting {device_control_env_var}: " + f"Error computing assigned_physical_gpu_ids: " f"local range: [{local_dp_rank * world_size}, " f"{(local_dp_rank + 1) * world_size}) " f'base value: "{os.getenv(device_control_env_var)}"' diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index e13301f03c6..093f065475a 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -12,7 +12,6 @@ from multiprocessing import Process, connection from multiprocessing.process import BaseProcess from multiprocessing.queues import Queue from typing import TYPE_CHECKING, cast -from unittest.mock import patch import msgspec import zmq @@ -175,38 +174,38 @@ class CoreEngineProcManager: self.manager_stopped = threading.Event() self.failed_proc_name: str | None = None + # All ranks share this config object: capture the user-provided + # --device-ids list before the per-rank shard overwrites it. Mutating + # the config before each proc.start() works because the spawn method + # pickles process args at start() time, sequentially per rank. + user_assigned_gpu_ids = vllm_config.parallel_config.assigned_physical_gpu_ids try: for proc, local_dp_rank in zip(self.processes, local_dp_ranks): - # Adjust device control in DP for platforms that cannot rely - # on torch.accelerator.set_device_index(), and for Ray launchers. - device_control_context: contextlib.AbstractContextManager[None] = ( - contextlib.nullcontext() - ) + # Populate the logical-to-physical GPU mapping in DP for + # platforms that cannot rely on + # torch.accelerator.set_device_index(), and for Ray. needs_device_env_isolation = not ( current_platform.is_cuda_alike() or current_platform.is_xpu() ) if is_dp and ( needs_device_env_isolation or vllm_config.parallel_config.use_ray ): - device_control_context = set_device_control_env_var( - vllm_config, local_dp_rank + set_assigned_physical_gpu_ids_for_dp_rank( + vllm_config, local_dp_rank, user_assigned_gpu_ids ) - with ( - device_control_context, - numa_utils.configure_subprocess( - # EngineCore itself does not have a TP/PP-local rank. - # When DP is enabled, set_device_control_env_var() - # narrows visible devices to this DP shard first, so - # local_rank=0 means "the first local GPU in this - # shard". The actual TP/PP worker processes spawned by - # the executor are bound separately with their own - # local_rank values. - vllm_config, - local_rank=0, - dp_local_rank=local_dp_rank, - process_kind="EngineCore", - ), + with numa_utils.configure_subprocess( + # EngineCore itself does not have a TP/PP-local rank. + # When DP is enabled, set_assigned_physical_gpu_ids_for_dp_rank() + # populates the logical-to-physical mapping for this DP + # shard, so local_rank=0 means "the first local GPU in + # this shard". The actual TP/PP worker processes spawned + # by the executor are bound separately with their own + # local_rank values. + vllm_config, + local_rank=0, + dp_local_rank=local_dp_rank, + process_kind="EngineCore", ): proc.start() finally: @@ -281,55 +280,79 @@ class SignalCallback: self._event.set() -@contextlib.contextmanager -def set_device_control_env_var( - vllm_config: VllmConfig, local_dp_rank: int -) -> Iterator[None]: +def set_assigned_physical_gpu_ids_for_dp_rank( + vllm_config: VllmConfig, + local_dp_rank: int, + user_assigned_gpu_ids: list[int] | None = None, +) -> None: """ - Temporarily set CUDA_VISIBLE_DEVICES or equivalent - for engine subprocess. + Populate assigned_physical_gpu_ids on the config for the given DP rank. + + user_assigned_gpu_ids is the full (un-sharded) --device-ids list, if the + user provided one; this DP rank's shard is sliced from it. It is passed + explicitly rather than read from the config because callers may reuse + one config object across DP ranks, overwriting the field each time. """ world_size = vllm_config.parallel_config.world_size local_world_size = vllm_config.parallel_config.local_world_size evar = current_platform.device_control_env_var - value = get_device_indices(evar, local_dp_rank, world_size, local_world_size) - with patch.dict(os.environ, values=((evar, value),)): - yield + physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank( + evar, + local_dp_rank, + world_size, + local_world_size, + user_assigned_gpu_ids=user_assigned_gpu_ids, + ) + vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids -def get_device_indices( +def get_physical_gpu_ids_for_local_dp_rank( device_control_env_var: str, local_dp_rank: int, world_size: int, local_world_size: int | None = None, -): + user_assigned_gpu_ids: list[int] | None = None, +) -> list[int]: """ - Returns a comma-separated string of device indices for the specified + Returns list of physical GPU IDs for the specified data parallel rank. For example, if world_size=2 and local_dp_rank=1, and there are 4 devices, - this will select devices 2 and 3 for local_dp_rank=1. + this will return [2, 3] for local_dp_rank=1. + + If user_assigned_gpu_ids is provided (e.g. from --device-ids), this DP + rank's shard is sliced from it instead of being derived from the + device-control env var. """ if local_world_size is None: local_world_size = world_size + if user_assigned_gpu_ids is not None: + start = local_dp_rank * world_size + stop = start + local_world_size + if stop > len(user_assigned_gpu_ids): + raise ValueError( + f"--device-ids provides {len(user_assigned_gpu_ids)} devices, " + f"but DP rank {local_dp_rank} needs devices [{start}, {stop})" + ) + return user_assigned_gpu_ids[start:stop] try: - value = ",".join( - str(current_platform.device_id_to_physical_device_id(i)) + return [ + current_platform.device_id_to_physical_device_id(i) for i in range( local_dp_rank * world_size, local_dp_rank * world_size + local_world_size, ) - ) + ] except IndexError as e: raise Exception( - f"Error setting {device_control_env_var}: " + f"Error computing device indices for " + f"{device_control_env_var}: " f"local range: [{local_dp_rank * world_size}, " f"{(local_dp_rank + 1) * world_size}) " "base value: " f'"{os.getenv(device_control_env_var)}"' ) from e - return value def _apply_dp_identity_suffix(dp_vllm_config, dp_rank: int) -> None: @@ -453,11 +476,11 @@ class CoreEngineActorManager: # https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/intel_gpu.py#L56 # noqa: E501 if current_platform.is_xpu(): device_evar = current_platform.device_control_env_var - device_indices = get_device_indices( + physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank( device_evar, local_index, world_size ) actor_env_vars = self.env_vars_dict.copy() - actor_env_vars[device_evar] = device_indices + actor_env_vars[device_evar] = ",".join(str(d) for d in physical_gpu_ids) runtime_env = RuntimeEnv(env_vars=actor_env_vars) actor = ( diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 7bc81118e6b..9b7581311e8 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -826,6 +826,16 @@ class WorkerProc: signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGINT, signal_handler) + # Publish the logical-to-physical mapping early so topology helpers + # work before init_device (needed by set_worker_net_device below). + assigned_physical_gpu_ids = kwargs[ + "vllm_config" + ].parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + # Set net device env vars for the worker if VLLM_GPU_NIC_PCIE_MAPPING is set set_worker_net_device(kwargs.get("local_rank", 0), kwargs["vllm_config"]) diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index 749e59e04c2..39749ffc257 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -258,30 +258,35 @@ class RayDistributedExecutor(Executor): } self.collective_rpc("adjust_rank", args=(rerank_mapping,)) - # Get the set of GPU IDs used on each node. - worker_node_and_gpu_ids = [] + # Get the set of physical GPU IDs used on each node. + worker_node_and_physical_gpu_ids = [] for worker in [self.driver_dummy_worker] + self.workers: if worker is None: # driver_dummy_worker can be None when using ray spmd worker. continue - worker_node_and_gpu_ids.append( - ray.get(worker.get_node_and_gpu_ids.remote()) # type: ignore[attr-defined] + worker_node_and_physical_gpu_ids.append( + ray.get(worker.get_node_and_physical_gpu_ids.remote()) # type: ignore[attr-defined] ) node_workers = defaultdict(list) # node id -> list of worker ranks - node_gpus = defaultdict(list) # node id -> list of gpu ids + node_physical_gpu_ids = defaultdict(list) # node id -> physical GPU IDs - for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids): + for i, (node_id, physical_gpu_ids) in enumerate( + worker_node_and_physical_gpu_ids + ): node_workers[node_id].append(i) - # `gpu_ids` can be a list of strings or integers. + # `physical_gpu_ids` can be a list of strings or integers. # convert them to integers for consistency. - # NOTE: gpu_ids can be larger than 9 (e.g. 16 GPUs), + # NOTE: physical GPU IDs can be larger than 9 (e.g. 16 GPUs), # string sorting is not sufficient. # see https://github.com/vllm-project/vllm/issues/5590 - gpu_ids = [int(x) for x in gpu_ids] - node_gpus[node_id].extend(gpu_ids) - for node_id, gpu_ids in node_gpus.items(): - node_gpus[node_id] = sorted(gpu_ids) + physical_gpu_ids = [ + current_platform.device_control_id_to_physical_device_id(str(x)) + for x in physical_gpu_ids + ] + node_physical_gpu_ids[node_id].extend(physical_gpu_ids) + for node_id, physical_gpu_ids in node_physical_gpu_ids.items(): + node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids) all_ips = set(worker_ips + [driver_ip]) n_ips = len(all_ips) @@ -297,23 +302,8 @@ class RayDistributedExecutor(Executor): " each node." ) - # Set environment variables for the driver and workers. - # We set CUDA_VISIBLE_DEVICES to ALL GPUs on the node for each worker. - # This is needed because: - # 1. Ray's compiled DAG needs to find the allocated GPU in - # CUDA_VISIBLE_DEVICES. - # 2. vLLM's communication layer (NCCL, CustomAllreduce) needs to see - # all GPUs for P2P checks and communication setup. Though if it was - # just this reason, we could have also just kept the visible devices - # unset. - # Each worker will use local_rank to index into the visible devices. - all_args_to_update_environment_variables = [ - { - current_platform.device_control_env_var: ",".join( - map(str, node_gpus[node_id]) - ), - } - for (node_id, _) in worker_node_and_gpu_ids + all_args_to_update_environment_variables: list[dict[str, str]] = [ + {} for _ in worker_node_and_physical_gpu_ids ] # Environment variables to copy from driver to workers @@ -336,7 +326,7 @@ class RayDistributedExecutor(Executor): "update_environment_variables", args=(self._get_env_vars_to_be_updated(),) ) - if len(node_gpus) == 1: + if len(node_physical_gpu_ids) == 1: # in single node case, we don't need to get the IP address. # the loopback address is sufficient # NOTE: a node may have several IP addresses, one for each @@ -352,10 +342,11 @@ class RayDistributedExecutor(Executor): # Initialize the actual workers inside worker wrapper. all_kwargs = [] - for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids): + for rank, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids): local_rank = node_workers[node_id].index(rank) kwargs = dict( vllm_config=self.vllm_config, + assigned_physical_gpu_ids=sorted(node_physical_gpu_ids[node_id]), local_rank=local_rank, rank=rank, distributed_init_method=distributed_init_method, diff --git a/vllm/v1/executor/ray_executor_v2.py b/vllm/v1/executor/ray_executor_v2.py index 0665b5fc1b8..d50f06cc620 100644 --- a/vllm/v1/executor/ray_executor_v2.py +++ b/vllm/v1/executor/ray_executor_v2.py @@ -79,24 +79,25 @@ class RayWorkerProc(WorkerProc): 1. __init__: lightweight setup, stores init args (no device/model init) 2. initialize_worker: called after GPU IDs are discovered, completes the full WorkerProc initialization with the correct local_rank and - CUDA_VISIBLE_DEVICES. + logical-to-physical GPU mapping. - CUDA_VISIBLE_DEVICES setup flow: + GPU assignment flow: 1. RayExecutorV2 enables RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES so Ray does not set CUDA_VISIBLE_DEVICES on RayWorkerProc actors at creation time. 2. Each actor is scheduled with a placement group and bundle index; Ray resolves the physical GPU ID for that bundle at placement time. - 3. After placement, the worker discovers that GPU ID and sets - CUDA_VISIBLE_DEVICES before finishing WorkerProc initialization. + 3. After placement, the executor discovers each worker's GPU ID and passes the + node's logical-to-physical mapping (assigned_physical_gpu_ids) to + initialize_worker(); CUDA_VISIBLE_DEVICES is never modified. - There is no workaround for this unset-and-reset sequence when the placement group - is externally managed: scheduling must complete before CUDA_VISIBLE_DEVICES can - match the GPU tied to the worker's bundle. + Scheduling must complete before the mapping is known when the placement + group is externally managed: only then is the GPU tied to the worker's + bundle resolved. This sequence allows multiple vLLM instances to coexist on the same node: each instance is unaware which physical devices others hold, and the - externally managed placement group avoids CUDA_VISIBLE_DEVICES conflicts + externally managed placement group avoids device assignment conflicts by binding workers to specific placement group bundles. """ @@ -120,28 +121,33 @@ class RayWorkerProc(WorkerProc): is_driver_worker=is_driver_worker, ) - def get_node_and_gpu_ids(self) -> tuple[str, list[int]]: - """Return (node_id, gpu_ids) assigned to this actor by Ray.""" + def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]: + """Return (node_id, physical_gpu_ids) assigned to this actor by Ray.""" node_id = ray.get_runtime_context().get_node_id() device_key = current_platform.ray_device_key if not device_key: raise RuntimeError( f"current platform {current_platform.device_name} does not support ray." ) - gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key] - return node_id, [int(x) for x in gpu_ids] + physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key] + return node_id, [ + current_platform.device_control_id_to_physical_device_id(str(x)) + for x in physical_gpu_ids + ] def initialize_worker( self, local_rank: int, env_vars: dict[str, str], driver_env_vars: dict[str, str] | None = None, + assigned_physical_gpu_ids: list[int] | None = None, ) -> None: """Complete initialization after GPU assignment is known. *driver_env_vars* are applied with ``setdefault`` — they fill in missing vars but never overwrite node-local values. - *env_vars* (e.g. CUDA_VISIBLE_DEVICES) always overwrite. + *env_vars* always overwrite. + *assigned_physical_gpu_ids* maps local_rank to physical CUDA device ID. """ if driver_env_vars: for key, value in driver_env_vars.items(): @@ -149,6 +155,13 @@ class RayWorkerProc(WorkerProc): for key, value in env_vars.items(): os.environ[key] = value + if assigned_physical_gpu_ids is not None: + vllm_config = self._init_kwargs["vllm_config"] + assert isinstance(vllm_config, VllmConfig) + vllm_config.parallel_config.assigned_physical_gpu_ids = ( + assigned_physical_gpu_ids + ) + self.local_rank = local_rank super().__init__( local_rank=local_rank, @@ -365,36 +378,48 @@ class RayExecutorV2(MultiprocExecutor): ) self.ray_worker_handles.append(handle) - # Step 6: Discover GPU IDs assigned to each worker via Ray runtime context. - worker_node_and_gpu_ids = ray.get( - [h.actor.get_node_and_gpu_ids.remote() for h in self.ray_worker_handles] + # Step 6: Discover physical GPU IDs assigned to each worker via Ray + # runtime context. + worker_node_and_physical_gpu_ids = ray.get( + [ + h.actor.get_node_and_physical_gpu_ids.remote() + for h in self.ray_worker_handles + ] ) node_workers: dict[str, list[int]] = defaultdict(list) - node_gpus: dict[str, list[int]] = defaultdict(list) - for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids): + node_physical_gpu_ids: dict[str, list[int]] = defaultdict(list) + for i, (node_id, physical_gpu_ids) in enumerate( + worker_node_and_physical_gpu_ids + ): node_workers[node_id].append(i) - node_gpus[node_id].extend(gpu_ids) - for node_id, gpu_ids in node_gpus.items(): - node_gpus[node_id] = sorted(gpu_ids) + node_physical_gpu_ids[node_id].extend(physical_gpu_ids) + for node_id, physical_gpu_ids in node_physical_gpu_ids.items(): + node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids) - # Step 7: Initialize workers with correct local_rank and - # CUDA_VISIBLE_DEVICES. Each worker sees all GPUs assigned to - # this executor on its node; local_rank indexes into that set. + # Step 7: Initialize workers with local logical ranks and the + # logical-to-physical GPU mapping discovered from Ray placement. init_worker_refs = [] - for i, (node_id, _) in enumerate(worker_node_and_gpu_ids): + for i, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids): local_rank = node_workers[node_id].index(i) - worker_env_vars = { - current_platform.device_control_env_var: ",".join( - map(str, node_gpus[node_id]) - ), - } + assigned_physical_gpu_ids = sorted(node_physical_gpu_ids[node_id]) + worker_env_vars: dict[str, str] = {} self.ray_worker_handles[i].local_rank = local_rank init_worker_refs.append( self.ray_worker_handles[i].actor.initialize_worker.remote( - local_rank, worker_env_vars, self.driver_env_vars + local_rank, + worker_env_vars, + self.driver_env_vars, + assigned_physical_gpu_ids=assigned_physical_gpu_ids, ) ) + # Also set on the executor-side config for consistency. The mapping + # is per-node, so only do this when all workers share one node. + if len(node_physical_gpu_ids) == 1: + node_id_0 = worker_node_and_physical_gpu_ids[0][0] + self.vllm_config.parallel_config.assigned_physical_gpu_ids = sorted( + node_physical_gpu_ids[node_id_0] + ) ray.get(init_worker_refs) # Step 8: Collect response MQ handles diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py index 9083b919591..cc17c39e35f 100644 --- a/vllm/v1/executor/ray_utils.py +++ b/vllm/v1/executor/ray_utils.py @@ -93,7 +93,7 @@ try: def get_node_ip(self) -> str: return get_ip() - def get_node_and_gpu_ids(self) -> tuple[str, list[int]]: + def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]: node_id = ray.get_runtime_context().get_node_id() device_key = vllm.platforms.current_platform.ray_device_key if not device_key: @@ -101,8 +101,10 @@ try: "current platform %s does not support ray.", vllm.platforms.current_platform.device_name, ) - gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key] - return node_id, gpu_ids + physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[ + device_key + ] + return node_id, physical_gpu_ids def setup_device_if_necessary(self): # TODO(swang): This is needed right now because Ray CG executes diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 0291faf1afc..5e266a31354 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -270,19 +270,47 @@ class Worker(WorkerBase): # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK self.local_rank += dp_local_rank * tp_pp_world_size + + # Publish the logical-to-physical mapping for topology queries + # such as NIC affinity and P2P checks. + assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + assert self.local_rank < len(assigned_physical_gpu_ids), ( + f"local_rank {self.local_rank} is out of bounds for " + f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}" + ) + # NOTE(patch pr45026): local_world_size is derived from + # parallel_config.nnodes, which is only set for the "mp" + # multi-node backend. With the "ray"/"external_launcher" + # backends nnodes stays 1, so local_world_size collapses to + # the full world_size and this check wrongly fires on + # cross-node deployments. assigned_physical_gpu_ids is already + # per-node and the local_rank bound above fully validates the + # mapping for these backends, so skip the check for them. + if parallel_config.distributed_executor_backend not in ( + "ray", + "external_launcher", + ): + assert self.parallel_config.local_world_size <= len( + assigned_physical_gpu_ids + ), ( + f"local_world_size ({self.parallel_config.local_world_size})" + " exceeds assigned_physical_gpu_ids count " + f"({len(assigned_physical_gpu_ids)})" + ) + else: assert self.local_rank < torch.accelerator.device_count(), ( - f"DP adjusted local rank {self.local_rank} is out of bounds. " - ) - visible_device_count = ( - torch.accelerator.device_count() if torch.cuda.is_available() else 0 - ) - assert self.parallel_config.local_world_size <= visible_device_count, ( - f"local_world_size ({self.parallel_config.local_world_size}) must " - f"be less than or equal to the number of visible devices " - f"({visible_device_count})." + f"DP adjusted local rank {self.local_rank} is out of " + f"bounds for {torch.accelerator.device_count()} devices." ) - self.device = torch.device(f"cuda:{self.local_rank}") + visible_device_index = ( + current_platform.logical_device_id_to_visible_device_id(self.local_rank) + ) + self.device = torch.device(f"cuda:{visible_device_index}") torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 19bb18bd39f..9381d71913d 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -286,6 +286,12 @@ class WorkerWrapperBase: extended_calls, ) + assigned_physical_gpu_ids = kwargs.pop("assigned_physical_gpu_ids", None) + if assigned_physical_gpu_ids is not None: + vllm_config.parallel_config.assigned_physical_gpu_ids = ( + assigned_physical_gpu_ids + ) + shared_worker_lock = kwargs.pop("shared_worker_lock", None) if shared_worker_lock is None: msg = (