@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ebfbcfe46a

Stop setting CUDA_VISIBLE_DEVICES internally in vLLM, add device_ids arg (#45026)

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: kourosh hakhamaneshi <kouroshHakha@users.noreply.github.com>

Tyler Michael Smith committed Jun 20, 2026 at 16:38 UTC ebfbcfe46aa895d428933244908bae08b1ca6397

24 files changed +722 -185

csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu

+8 -3

@@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba

268 using TileShapeD = typename MlaSm100Type::TileShapeD;

269 arguments.problem_shape =

270 cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));

271 - // Assumes device 0 when getting sm_count.

272 - arguments.hw_info.sm_count =

273 - sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;

271 + if (sm_count <= 0) {

272 + int current_device = 0;

273 + cudaGetDevice(&current_device);

274 + arguments.hw_info.sm_count =

275 + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);

276 + } else {

277 + arguments.hw_info.sm_count = sm_count;

278 + }

279 arguments.split_kv = static_cast<int>(num_kv_splits);

280 MlaSm100Type::Fmha::set_split_kv(arguments);

281

tests/engine/test_arg_utils.py

+193

@@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch):

649 args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer")

650 assert args.model == "s3://bucket/model"

651 assert args.tokenizer == "s3://bucket/tokenizer"

652 +

653 +

654 +class TestDeviceIds:

655 + def test_device_ids_with_cvd_out_of_range(self, monkeypatch):

656 + """--device-ids index beyond the CVD set raises ValueError."""

657 + from vllm.platforms import current_platform

658 +

659 + key = current_platform.device_control_env_var

660 + monkeypatch.setenv(key, "4,5")

661 + args = EngineArgs(model="m", device_ids=[0, 2])

662 + with pytest.raises(ValueError, match="out of range"):

663 + args._resolve_device_ids()

664 +

665 + def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch):

666 + """--device-ids are CVD-local indices resolved to physical ids."""

667 + from vllm.platforms import current_platform

668 +

669 + key = current_platform.device_control_env_var

670 + monkeypatch.setenv(key, "4,5")

671 + args = EngineArgs(model="m", device_ids=[0, 1])

672 + assert args._resolve_device_ids() == [4, 5]

673 +

674 + def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch):

675 + """--device-ids support UUID CVD values resolved by the platform."""

676 + from vllm.platforms import current_platform

677 +

678 + key = current_platform.device_control_env_var

679 + monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890")

680 + monkeypatch.setattr(

681 + type(current_platform),

682 + "device_control_id_to_physical_device_id",

683 + classmethod(

684 + lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]

685 + ),

686 + )

687 +

688 + args = EngineArgs(model="m", device_ids=[0, 1])

689 + assert args._resolve_device_ids() == [4, 5]

690 +

691 + def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch):

692 + """UUID --device-ids are resolved to physical IDs immediately."""

693 + from vllm.platforms import current_platform

694 +

695 + monkeypatch.setattr(

696 + type(current_platform),

697 + "device_control_id_to_physical_device_id",

698 + classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]),

699 + )

700 +

701 + args = EngineArgs(model="m", device_ids=["GPU-abcd1234"])

702 + assert args._resolve_device_ids() == [4]

703 +

704 + def test_device_ids_reject_mixed_integer_and_uuid_args(self):

705 + """--device-ids must not mix CVD indices and UUIDs."""

706 + args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"])

707 + with pytest.raises(ValueError, match="must not mix"):

708 + args._resolve_device_ids()

709 +

710 + def test_no_device_ids(self):

711 + """No --device-ids returns None."""

712 + args = EngineArgs(model="m")

713 + assert args._resolve_device_ids() is None

714 +

715 + def test_cli_parsing(self):

716 + """--device-ids parses comma-separated string from CLI."""

717 + parser = FlexibleArgumentParser()

718 + EngineArgs.add_cli_args(parser)

719 + parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"])

720 + assert parsed.device_ids == [0, 2, 4]

721 +

722 + def test_cli_parsing_uuid(self):

723 + """--device-ids parses comma-separated UUID strings from CLI."""

724 + parser = FlexibleArgumentParser()

725 + EngineArgs.add_cli_args(parser)

726 + parsed = parser.parse_args(

727 + ["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"]

728 + )

729 + assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"]

730 +

731 + def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch):

732 + """assigned_physical_gpu_ids are already physical and not composed with CVD."""

733 + import vllm.platforms.interface as platform_interface

734 + from vllm.platforms import current_platform

735 +

736 + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5])

737 + monkeypatch.setenv(current_platform.device_control_env_var, "4,5")

738 +

739 + assert current_platform.device_id_to_physical_device_id(0) == 4

740 + assert current_platform.device_id_to_physical_device_id(1) == 5

741 + assert current_platform.logical_device_id_to_visible_device_id(0) == 0

742 + assert current_platform.logical_device_id_to_visible_device_id(1) == 1

743 +

744 + def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch):

745 + """Physical IDs map back to visible ordinals when CVD uses UUIDs."""

746 + import vllm.platforms.interface as platform_interface

747 + from vllm.platforms import current_platform

748 +

749 + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5])

750 + monkeypatch.setenv(

751 + current_platform.device_control_env_var,

752 + "GPU-abcd1234,GPU-ef567890",

753 + )

754 + monkeypatch.setattr(

755 + type(current_platform),

756 + "device_control_id_to_physical_device_id",

757 + classmethod(

758 + lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]

759 + ),

760 + )

761 +

762 + assert current_platform.logical_device_id_to_visible_device_id(0) == 1

763 +

764 + def test_device_ids_reject_duplicates(self):

765 + """--device-ids must not contain duplicate entries."""

766 + args = EngineArgs(model="m", device_ids=[2, 2])

767 + with pytest.raises(ValueError, match="duplicates"):

768 + args._resolve_device_ids()

769 +

770 + def test_cli_parsing_strips_whitespace(self):

771 + """--device-ids tolerates whitespace around commas."""

772 + parser = FlexibleArgumentParser()

773 + EngineArgs.add_cli_args(parser)

774 + parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"])

775 + assert parsed.device_ids == [0, 2, 4]

776 +

777 + def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch):

778 + """visible_device_id_to_physical_device_id maps torch device ordinals,

779 + independent of the logical-to-physical mapping.

780 +

781 + Regression test: CustomAllreduce passes device.index (a visible

782 + ordinal) and must not index into assigned_physical_gpu_ids, which

783 + raised IndexError for non-identity --device-ids like [2, 3].

784 + """

785 + import vllm.platforms.interface as platform_interface

786 + from vllm.platforms import current_platform

787 +

788 + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3])

789 + monkeypatch.delenv(current_platform.device_control_env_var, raising=False)

790 +

791 + # CVD unset: visible ordinal == physical ID, even beyond the

792 + # assigned list's length.

793 + assert current_platform.visible_device_id_to_physical_device_id(2) == 2

794 + assert current_platform.visible_device_id_to_physical_device_id(3) == 3

795 +

796 + monkeypatch.setenv(current_platform.device_control_env_var, "4,5")

797 + assert current_platform.visible_device_id_to_physical_device_id(1) == 5

798 + with pytest.raises(IndexError, match="out of range"):

799 + current_platform.visible_device_id_to_physical_device_id(2)

800 +

801 +

802 +class TestDpDeviceIdSharding:

803 + def test_dp_supervisor_device_ids_stay_env_relative(self):

804 + """Regression test: the DP supervisor must pass env-relative indices,

805 + not physical IDs, because each child re-resolves --device-ids

806 + against its inherited device-control env var."""

807 + import argparse

808 +

809 + from vllm.entrypoints.openai.dp_supervisor import _build_device_ids

810 +

811 + args = argparse.Namespace(

812 + tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None

813 + )

814 + assert _build_device_ids(args, local_rank=0) == [0, 1]

815 + assert _build_device_ids(args, local_rank=1) == [2, 3]

816 +

817 + def test_dp_supervisor_shards_user_device_ids(self):

818 + """User-provided --device-ids are sharded across DP children."""

819 + import argparse

820 +

821 + from vllm.entrypoints.openai.dp_supervisor import _build_device_ids

822 +

823 + args = argparse.Namespace(

824 + tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7]

825 + )

826 + assert _build_device_ids(args, local_rank=0) == [4, 5]

827 + assert _build_device_ids(args, local_rank=1) == [6, 7]

828 + with pytest.raises(ValueError, match="needs devices"):

829 + _build_device_ids(args, local_rank=2)

830 +

831 + def test_dp_rank_shards_user_assigned_gpu_ids(self):

832 + """get_physical_gpu_ids_for_local_dp_rank slices the user-provided

833 + --device-ids list instead of recomputing from the env var."""

834 + from vllm.platforms import current_platform

835 + from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank

836 +

837 + evar = current_platform.device_control_env_var

838 + assert get_physical_gpu_ids_for_local_dp_rank(

839 + evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]

840 + ) == [6, 7]

841 + with pytest.raises(ValueError, match="needs devices"):

842 + get_physical_gpu_ids_for_local_dp_rank(

843 + evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]

844 + )

tests/entrypoints/openai/test_dp_supervisor.py

+2 -2

@@ -364,7 +364,7 @@ class MockVLLMServer:

364 await self._serve_task

365

366

367 -def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]):

367 +def launch_mock_vllm(child_args: argparse.Namespace):

368 logger.info("Launching mock vLLM on port %s", child_args.port)

369 mock_vllm = MockVLLMServer(

370 port=child_args.port,

@@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]

375

376

377 def launch_mock_vllm_with_drain(

378 - child_args: argparse.Namespace, env_updates: dict[str, str]

378 + child_args: argparse.Namespace,

379 ):

380 logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)

381 mock_vllm = MockVLLMServer(

vllm/config/parallel.py

@@ -302,6 +302,14 @@ class ParallelConfig:

302 Each entry must use `numactl --physcpubind` CPU-list syntax, for example

303 `"0-3"` or `"0,2,4-7"`.

304 """

305 + assigned_physical_gpu_ids: list[int] | None = None

306 + """Mapping from vLLM-local logical GPU IDs to physical GPU IDs.

307 +

308 + For example, ``[2, 3]`` means logical GPU 0 maps to physical GPU 2,

309 + and logical GPU 1 maps to physical GPU 3. Physical IDs are used only

310 + at platform/topology boundaries such as NVML, NIC affinity, P2P

311 + checks, and final CUDA device selection when needed. When None,

312 + logical IDs map to visible device IDs in order."""

313

314 distributed_timeout_seconds: int | None = None

315 """Timeout in seconds for distributed operations (e.g., init_process_group).

@@ -772,6 +780,7 @@ class ParallelConfig:

780 "numa_bind",

781 "numa_bind_nodes",

782 "numa_bind_cpus",

783 + "assigned_physical_gpu_ids",

784 }

785

786 from vllm.config.utils import get_hash_factors, hash_factors

vllm/distributed/device_communicators/all2all.py

+8 -1

@@ -704,7 +704,14 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):

704 self.num_experts = num_experts

705

706 self.cleanup()

707 - gpus_per_node = torch.accelerator.device_count()

707 + from vllm.platforms.interface import get_assigned_physical_gpu_ids

708 +

709 + assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()

710 + gpus_per_node = (

711 + len(assigned_physical_gpu_ids)

712 + if assigned_physical_gpu_ids is not None

713 + else torch.accelerator.device_count()

714 + )

715 logger.debug(

716 "Making One-sided NVLink mapping: rank=%d, world size=%d",

717 self.rank,

vllm/distributed/device_communicators/all_reduce_utils.py

+26 -7

@@ -320,13 +320,21 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:

320

321 is_distributed = dist.is_initialized()

322

323 - num_dev = current_platform.device_count()

324 - cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES

325 - if cuda_visible_devices is None:

326 - cuda_visible_devices = ",".join(str(i) for i in range(num_dev))

323 + from vllm.platforms.interface import get_assigned_physical_gpu_ids

324 +

325 + assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()

326 + if assigned_physical_gpu_ids is not None:

327 + # Key by the ordered list: the cache stores directed local-index

328 + # pairs, so permutations of the same set are distinct mappings.

329 + cache_key = ",".join(str(i) for i in assigned_physical_gpu_ids)

330 + num_dev = len(assigned_physical_gpu_ids)

331 + else:

332 + num_dev = current_platform.device_count()

333 + cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES

334 + cache_key = cuda_visible_devices or ",".join(str(i) for i in range(num_dev))

335

336 path = os.path.join(

329 - envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json"

337 + envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cache_key}.json"

338 )

339 os.makedirs(os.path.dirname(path), exist_ok=True)

340 from vllm.distributed.parallel_state import get_world_group

@@ -338,7 +346,15 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:

346 # enter this block to calculate the cache

347 logger.info("generating GPU P2P access cache in %s", path)

348 cache: dict[str, bool] = {}

341 - ids = list(range(num_dev))

349 + # The probe subprocesses inherit this process's device-control env

350 + # var, so they must be given visible ordinals, not physical IDs.

351 + if assigned_physical_gpu_ids is not None:

352 + ids = [

353 + current_platform.logical_device_id_to_visible_device_id(local)

354 + for local in range(num_dev)

355 + ]

356 + else:

357 + ids = list(range(num_dev))

358 # batch of all pairs of GPUs

359 batch_src, batch_tgt = zip(*list(product(ids, ids)))

360 # NOTE: we use `subprocess` rather than `multiprocessing` here

@@ -368,8 +384,11 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:

384 ) from e

385 with open(output_file.name, "rb") as f:

386 result = pickle.load(f)

387 + # Cache entries must be keyed by local indices (0..N-1) because

388 + # gpu_p2p_access_check() is called with local ranks.

389 + id_to_local = {device_id: local for local, device_id in enumerate(ids)}

390 for _i, _j, r in zip(batch_src, batch_tgt, result):

372 - cache[f"{_i}->{_j}"] = r

391 + cache[f"{id_to_local[_i]}->{id_to_local[_j]}"] = r

392 with open(path, "w") as f:

393 json.dump(cache, f, indent=4)

394 if is_distributed:

vllm/distributed/device_communicators/custom_all_reduce.py

+10 -8

@@ -34,7 +34,12 @@ def _can_p2p(rank: int, world_size: int) -> bool:

34 continue

35 if envs.VLLM_SKIP_P2P_CHECK:

36 logger.debug("Skipping P2P check and trusting the driver's P2P report.")

37 - return torch.cuda.can_device_access_peer(rank, i)

37 + # can_device_access_peer takes visible device ordinals, while

38 + # rank and i are logical local IDs.

39 + return torch.cuda.can_device_access_peer(

40 + current_platform.logical_device_id_to_visible_device_id(rank),

41 + current_platform.logical_device_id_to_visible_device_id(i),

42 + )

43 if not gpu_p2p_access_check(rank, i):

44 return False

45 return True

@@ -126,13 +131,10 @@ class CustomAllreduce:

131 CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size],

132 max_size,

133 )

129 - cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES

130 - if cuda_visible_devices:

131 - device_ids = list(map(int, cuda_visible_devices.split(",")))

132 - else:

133 - device_ids = list(range(current_platform.device_count()))

134 -

135 - physical_device_id = device_ids[device.index]

134 + # device.index is a visible ordinal, not a logical local ID.

135 + physical_device_id = current_platform.visible_device_id_to_physical_device_id(

136 + device.index

137 + )

138 tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")

139 gather_list = [

140 torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)

vllm/distributed/device_communicators/quick_all_reduce.py

+4 -6

@@ -129,12 +129,10 @@ class QuickAllReduce:

129 assert isinstance(device, torch.device)

130 self.device = device

131

132 - cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES

133 - if cuda_visible_devices:

134 - device_ids = list(map(int, cuda_visible_devices.split(",")))

135 - else:

136 - device_ids = list(range(current_platform.device_count()))

137 - physical_device_id = device_ids[device.index]

132 + # device.index is a visible ordinal, not a logical local ID.

133 + physical_device_id = current_platform.visible_device_id_to_physical_device_id(

134 + device.index

135 + )

136 tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")

137 gather_list = [

138 torch.tensor([0], dtype=torch.int, device="cpu")

vllm/distributed/device_communicators/shm_broadcast.py

+7 -1

@@ -840,7 +840,13 @@ class MessageQueue:

840 The MessageQueue instance for the calling process,

841 and a list of handles (only non-empty for the reader process).

842 """

843 - local_size = current_platform.device_count()

843 + from vllm.platforms.interface import get_assigned_physical_gpu_ids

844 +

845 + assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()

846 + if assigned_physical_gpu_ids is not None:

847 + local_size = len(assigned_physical_gpu_ids)

848 + else:

849 + local_size = current_platform.device_count()

850 rank = dist.get_rank()

851 same_node = rank // local_size == reader_rank // local_size

852 buffer_io = MessageQueue(

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py

+5 -4

@@ -482,10 +482,11 @@ def _init_lmcache_engine(

482 )

483

484 # Change current device.

485 - num_gpus = torch.accelerator.device_count()

486 - local_rank = parallel_config.rank % num_gpus

487 - torch.accelerator.set_device_index(local_rank)

488 - device = torch.device(f"cuda:{local_rank}")

485 + from vllm.distributed.parallel_state import get_world_group

486 +

487 + device_index = get_world_group().device_index

488 + torch.accelerator.set_device_index(device_index)

489 + device = torch.device(f"cuda:{device_index}")

490 metadata = LMCacheEngineMetadata(

491 model_config.model,

492 parallel_config.world_size,

vllm/distributed/parallel_state.py

+24 -4

@@ -392,6 +392,14 @@ class GroupCoordinator:

392

393 self.rank = torch.distributed.get_rank()

394 self.local_rank = local_rank

395 + self.device_index: int

396 + if _WORLD is not None:

397 + self.device_index = _WORLD.device_index

398 + else:

399 + assert local_rank >= 0, (

400 + "local_rank must be provided when creating the world group"

401 + )

402 + self.device_index = local_rank

403

404 self_device_group = None

405 self_cpu_group = None

@@ -442,11 +450,18 @@ class GroupCoordinator:

450 from vllm.platforms import current_platform

451

452 if current_platform.is_cuda_alike():

445 - self.device = torch.device(f"cuda:{local_rank}")

453 + visible_device_index = (

454 + current_platform.logical_device_id_to_visible_device_id(

455 + self.device_index

456 + )

457 + )

458 + self.device = torch.device(f"cuda:{visible_device_index}")

459 elif current_platform.is_xpu():

447 - self.device = torch.device(f"xpu:{local_rank}")

460 + self.device = torch.device(f"xpu:{self.device_index}")

461 elif current_platform.is_out_of_tree():

449 - self.device = torch.device(f"{current_platform.device_name}:{local_rank}")

462 + self.device = torch.device(

463 + f"{current_platform.device_name}:{self.device_index}"

464 + )

465 else:

466 self.device = torch.device("cpu")

467

@@ -1438,7 +1453,12 @@ def _init_process_group_for_split_group(

1453 """

1454 if torch.accelerator.is_available() and backend != "gloo":

1455 init_backend = "cpu:gloo,cuda:nccl"

1441 - device_id: torch.device | None = torch.device(f"cuda:{local_rank}")

1456 + from vllm.platforms import current_platform

1457 +

1458 + visible_device_index = current_platform.logical_device_id_to_visible_device_id(

1459 + local_rank

1460 + )

1461 + device_id: torch.device | None = torch.device(f"cuda:{visible_device_index}")

1462 else:

1463 init_backend = "gloo"

1464 device_id = None

vllm/distributed/stateless_coordinator.py

+19 -3

@@ -86,6 +86,15 @@ class StatelessGroupCoordinator(GroupCoordinator):

87 self.rank = global_rank

88 self.local_rank = local_rank

89 + from vllm.distributed.parallel_state import _WORLD

90 +

91 + if _WORLD is not None:

92 + self.device_index = _WORLD.device_index

93 + else:

94 + assert local_rank >= 0, (

95 + "local_rank must be provided when creating the world group"

96 + )

97 + self.device_index = local_rank

99 self_device_group = None

100 self_cpu_group = None

@@ -152,11 +161,18 @@ class StatelessGroupCoordinator(GroupCoordinator):

161 self.tcp_store_group = self_tcp_store_group

162

163 if current_platform.is_cuda_alike():

155 - self.device = torch.device(f"cuda:{local_rank}")

164 + visible_device_index = (

165 + current_platform.logical_device_id_to_visible_device_id(

166 + self.device_index

167 + )

168 + )

169 + self.device = torch.device(f"cuda:{visible_device_index}")

170 elif current_platform.is_xpu():

157 - self.device = torch.device(f"xpu:{local_rank}")

171 + self.device = torch.device(f"xpu:{self.device_index}")

172 elif current_platform.is_out_of_tree():

159 - self.device = torch.device(f"{current_platform.device_name}:{local_rank}")

173 + self.device = torch.device(

174 + f"{current_platform.device_name}:{self.device_index}"

175 + )

176 else:

177 self.device = torch.device("cpu")

178

vllm/engine/arg_utils.py

+58

@@ -6,6 +6,7 @@ import copy

6 import dataclasses

7 import functools

8 import json

9 +import os

10 import sys

11 from collections.abc import Callable

12 from dataclasses import MISSING, asdict, dataclass, fields, is_dataclass

@@ -465,6 +466,7 @@ class EngineArgs:

466 numa_bind: bool = ParallelConfig.numa_bind

467 numa_bind_nodes: list[int] | None = ParallelConfig.numa_bind_nodes

468 numa_bind_cpus: list[str] | None = ParallelConfig.numa_bind_cpus

469 + device_ids: list[int | str] | None = None

470 tensor_parallel_size: int = ParallelConfig.tensor_parallel_size

471 prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size

472 decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size

@@ -979,6 +981,20 @@ class EngineArgs:

981 parallel_group.add_argument(

982 "--numa-bind-cpus", **parallel_kwargs["numa_bind_cpus"]

983 )

984 + parallel_group.add_argument(

985 + "--device-ids",

986 + type=lambda s: [

987 + int(device_id) if device_id.isdigit() else device_id

988 + for device_id in (part.strip() for part in s.split(","))

989 + ],

990 + default=None,

991 + help="Comma-separated physical GPU device IDs or UUIDs to use "

992 + '(e.g. --device-ids "2,3,5,7"). Avoids setting '

993 + "CUDA_VISIBLE_DEVICES, preserving full GPU topology "

994 + "visibility for GPU-NIC affinity and DeepGEMM. "

995 + "Note: has no effect with Ray executors; use Ray "

996 + "placement groups for GPU selection instead.",

997 + )

998 parallel_group.add_argument(

999 "--tensor-parallel-size", "-tp", **parallel_kwargs["tensor_parallel_size"]

1000 )

@@ -1716,6 +1732,47 @@ class EngineArgs:

1732 )

1733 return SpeculativeConfig(**self.speculative_config)

1734

1735 + def _resolve_device_ids(self) -> list[int] | None:

1736 + if not self.device_ids:

1737 + return None

1738 + if self.distributed_executor_backend == "ray":

1739 + logger.warning(

1740 + "--device-ids has no effect when using the Ray executor. "

1741 + "Use Ray placement groups for GPU selection instead."

1742 + )

1743 + ids = self.device_ids

1744 + if len(set(ids)) != len(ids):

1745 + raise ValueError(f"--device-ids must not contain duplicates: {ids}")

1746 + if all(isinstance(i, str) for i in ids):

1747 + return [

1748 + current_platform.device_control_id_to_physical_device_id(i)

1749 + for i in cast(list[str], ids)

1750 + ]

1751 + if any(isinstance(i, str) for i in ids):

1752 + raise ValueError("--device-ids must not mix integer IDs and UUIDs")

1753 + int_ids = cast(list[int], ids)

1754 + # Compose with CUDA_VISIBLE_DEVICES: if CVD is set, treat

1755 + # --device-ids values as indices into the CVD-visible set.

1756 + cvd = getattr(

1757 + envs,

1758 + current_platform.device_control_env_var,

1759 + os.environ.get(current_platform.device_control_env_var),

1760 + )

1761 + if cvd:

1762 + cvd_ids = [

1763 + current_platform.device_control_id_to_physical_device_id(x)

1764 + for x in cvd.split(",")

1765 + ]

1766 + for i in int_ids:

1767 + if i >= len(cvd_ids):

1768 + raise ValueError(

1769 + f"--device-ids index {i} is out of range for "

1770 + f"{current_platform.device_control_env_var}"

1771 + f"={cvd} ({len(cvd_ids)} devices visible)"

1772 + )

1773 + return [cvd_ids[i] for i in int_ids]

1774 + return int_ids

1775 +

1776 def create_diffusion_config(self) -> DiffusionConfig | None:

1777 if self.diffusion_config is None:

1778 return None

@@ -2029,6 +2086,7 @@ class EngineArgs:

2086 cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,

2087 _api_process_count=self._api_process_count,

2088 _api_process_rank=self._api_process_rank,

2089 + assigned_physical_gpu_ids=self._resolve_device_ids(),

2090 numa_bind=self.numa_bind,

2091 numa_bind_nodes=self.numa_bind_nodes,

2092 numa_bind_cpus=self.numa_bind_cpus,

vllm/entrypoints/openai/dp_supervisor.py

+19 -18

@@ -23,12 +23,10 @@ import uvloop

23 from fastapi import FastAPI, Response

25 from vllm.logger import init_logger

26 -from vllm.platforms import current_platform

26 from vllm.utils.system_utils import (

27 decorate_logs,

28 kill_process_tree,

29 set_process_title,

31 - update_environment_variables,

30 )

32 logger = init_logger(__name__)

@@ -127,22 +125,29 @@ def _build_vllm_dp_server_args(

125 child_args.data_parallel_multi_port_external_lb = False

126 child_args.data_parallel_supervisor_port = None

127 child_args.api_server_count = 1

128 + child_args.device_ids = _build_device_ids(args, local_rank)

129 return child_args

130

131

133 -def _build_vllm_dp_server_env(

134 - args: argparse.Namespace, local_rank: int

135 -) -> dict[str, str]:

136 - # set visible devices for the child process

132 +def _build_device_ids(args: argparse.Namespace, local_rank: int) -> list[int | str]:

133 + """Build the --device-ids value for a DP child process.

134 +

135 + The child resolves these against its own inherited device-control env

136 + var (e.g. CUDA_VISIBLE_DEVICES), so integer IDs must stay env-relative

137 + here rather than being translated to physical IDs.

138 + """

139 devices_per_rank = args.tensor_parallel_size * args.pipeline_parallel_size

140 start = local_rank * devices_per_rank

141 stop = start + devices_per_rank

140 - device_env = current_platform.device_control_env_var

141 - visible_devices = ",".join(

142 - str(current_platform.device_id_to_physical_device_id(idx))

143 - for idx in range(start, stop)

144 - )

145 - return {device_env: visible_devices}

142 + device_ids = getattr(args, "device_ids", None)

143 + if device_ids is not None:

144 + if stop > len(device_ids):

145 + raise ValueError(

146 + f"--device-ids has {len(device_ids)} entries, but DP rank "

147 + f"{local_rank} needs devices [{start}, {stop})"

148 + )

149 + return device_ids[start:stop]

150 + return list(range(start, stop))

151

152

153 def _child_base_url(args: argparse.Namespace, port: int) -> str:

@@ -228,9 +233,7 @@ def _build_dp_supervisor_app(supervisor: DPSupervisor) -> FastAPI:

233 return app

234

235

231 -def _run_vllm_dp_server(

232 - child_args: argparse.Namespace, env_updates: dict[str, str]

233 -) -> None:

236 +def _run_vllm_dp_server(child_args: argparse.Namespace) -> None:

237 """

238 Entrypoint function for the vLLM DP Server.

239 """

@@ -241,7 +244,6 @@ def _run_vllm_dp_server(

244 os.setpgrp()

245

246 name = f"APIServer_DP{child_args.data_parallel_rank}"

244 - update_environment_variables(env_updates)

247 set_process_title(name)

248 decorate_logs(name)

249 uvloop.run(run_server(child_args))

@@ -345,11 +347,10 @@ class DPSupervisor:

347 context = multiprocessing.get_context("spawn")

348 for local_rank in range(self.args.data_parallel_size_local):

349 child_args = _build_vllm_dp_server_args(self.args, local_rank)

348 - child_env = _build_vllm_dp_server_env(self.args, local_rank)

350 process = context.Process(

351 target=_run_vllm_dp_server,

352 name=f"APIServer_DPRank_{child_args.data_parallel_rank}",

352 - args=(child_args, child_env),

353 + args=(child_args,),

354 )

355 process.start()

356 self._processes.append(process)

vllm/platforms/cuda.py

@@ -685,6 +685,15 @@ class CudaPlatformBase(Platform):

685 # all the related functions work on real physical device ids.

686 # the major benefit of using NVML is that it will not initialize CUDA

687 class NvmlCudaPlatform(CudaPlatformBase):

688 + @classmethod

689 + @with_nvml_context

690 + def device_control_id_to_physical_device_id(cls, device_id: str) -> int:

691 + try:

692 + return int(device_id)

693 + except ValueError:

694 + handle = pynvml.nvmlDeviceGetHandleByUUID(device_id)

695 + return pynvml.nvmlDeviceGetIndex(handle)

696 +

697 @classmethod

698 @cache

699 @with_nvml_context

vllm/platforms/interface.py

+102 -1

@@ -30,6 +30,33 @@ else:

31 logger = init_logger(__name__)

33 +_assigned_physical_gpu_ids: list[int] | None = None

34 +

35 +

36 +def set_assigned_physical_gpu_ids(ids: list[int]) -> None:

37 + """Set the physical GPU IDs assigned to this worker process.

38 + Called during worker init so that device_id_to_physical_device_id()

39 + can map local_rank to the correct physical device without relying

40 + on CUDA_VISIBLE_DEVICES.

41 +

42 + Idempotent: a second call with the same value is a no-op.

43 + Raises RuntimeError if called again with a different value.

44 +

45 + This is expected to run during single-threaded worker initialization."""

46 + global _assigned_physical_gpu_ids

47 + if _assigned_physical_gpu_ids is not None:

48 + if _assigned_physical_gpu_ids != ids:

49 + raise RuntimeError(

50 + f"set_assigned_physical_gpu_ids called with conflicting values: "

51 + f"existing={_assigned_physical_gpu_ids}, new={ids}"

52 + )

53 + return

54 + _assigned_physical_gpu_ids = ids

55 +

56 +

57 +def get_assigned_physical_gpu_ids() -> list[int] | None:

58 + return _assigned_physical_gpu_ids

59 +

61 @functools.cache

62 def in_wsl() -> bool:

@@ -233,8 +260,34 @@ class Platform:

260 """

261 import vllm.kernels # noqa: F401

262

263 + @classmethod

264 + def device_control_id_to_physical_device_id(cls, device_id: str) -> int:

265 + """Map one device-control env entry to an integer physical device ID."""

266 + try:

267 + return int(device_id)

268 + except ValueError as e:

269 + raise ValueError(

270 + f"Non-integer device ID {device_id!r} is not supported by "

271 + f"{cls.device_name}."

272 + ) from e

273 +

274 @classmethod

275 def device_id_to_physical_device_id(cls, device_id: int):

276 + """Map a vLLM-local logical device ID to a physical device ID.

277 +

278 + The input is a logical local ID (e.g. a local rank), NOT a visible

279 + device ordinal; for the latter use

280 + visible_device_id_to_physical_device_id(). The two coincide only

281 + when no logical-to-physical mapping is in effect.

282 + """

283 + if _assigned_physical_gpu_ids is not None:

284 + if device_id >= len(_assigned_physical_gpu_ids):

285 + raise IndexError(

286 + f"device_id {device_id} is out of range for "

287 + f"assigned_physical_gpu_ids {_assigned_physical_gpu_ids} "

288 + f"({len(_assigned_physical_gpu_ids)} devices assigned)"

289 + )

290 + return _assigned_physical_gpu_ids[device_id]

291 # Treat empty device control env var as unset. This is a valid

292 # configuration in Ray setups where the engine is launched in

293 # a CPU-only placement group located on a GPU node.

@@ -244,10 +297,58 @@ class Platform:

297 ):

298 device_ids = os.environ[cls.device_control_env_var].split(",")

299 physical_device_id = device_ids[device_id]

247 - return int(physical_device_id)

300 + return cls.device_control_id_to_physical_device_id(physical_device_id)

301 else:

302 return device_id

303

304 + @classmethod

305 + def logical_device_id_to_visible_device_id(cls, device_id: int) -> int:

306 + """Map a vLLM-local logical device ID to the current process's

307 + visible accelerator ordinal.

308 +

309 + vLLM internals use logical local IDs. Physical IDs are used only

310 + at platform/topology boundaries. This helper performs the final

311 + translation needed by APIs such as ``torch.device("cuda:N")``.

312 + """

313 + physical_device_id = cls.device_id_to_physical_device_id(device_id)

314 + device_control_env = os.environ.get(cls.device_control_env_var, "")

315 + if not device_control_env:

316 + return physical_device_id

317 +

318 + visible_physical_device_ids = [

319 + cls.device_control_id_to_physical_device_id(physical_id)

320 + for physical_id in device_control_env.split(",")

321 + ]

322 + if physical_device_id not in visible_physical_device_ids:

323 + raise RuntimeError(

324 + f"Physical device {physical_device_id} for logical device "

325 + f"{device_id} is not visible in {cls.device_control_env_var}="

326 + f"{device_control_env}"

327 + )

328 + return visible_physical_device_ids.index(physical_device_id)

329 +

330 + @classmethod

331 + def visible_device_id_to_physical_device_id(cls, device_id: int) -> int:

332 + """Map a visible accelerator ordinal (e.g. ``torch.device.index``)

333 + to a physical device ID.

334 +

335 + This is the inverse of the env-var translation performed by

336 + logical_device_id_to_visible_device_id() and is independent of any

337 + logical-to-physical mapping set via set_assigned_physical_gpu_ids().

338 + """

339 + device_control_env = os.environ.get(cls.device_control_env_var, "")

340 + if not device_control_env:

341 + return device_id

342 + visible_device_ids = device_control_env.split(",")

343 + if device_id >= len(visible_device_ids):

344 + raise IndexError(

345 + f"visible device ordinal {device_id} is out of range for "

346 + f"{cls.device_control_env_var}={device_control_env}"

347 + )

348 + return cls.device_control_id_to_physical_device_id(

349 + visible_device_ids[device_id]

350 + )

351 +

352 @classmethod

353 def import_kernels(cls) -> None:

354 """Import any platform-specific C kernels."""

vllm/v1/engine/core.py

+16 -9

@@ -74,7 +74,7 @@ from vllm.v1.engine.utils import (

74 EngineHandshakeMetadata,

75 EngineZmqAddresses,

76 SignalCallback,

77 - get_device_indices,

77 + get_physical_gpu_ids_for_local_dp_rank,

78 )

79 from vllm.v1.executor import Executor

80 from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind

@@ -2175,23 +2175,30 @@ class EngineCoreActorMixin:

2175 pass

2176 else:

2177 device_control_env_var = current_platform.device_control_env_var

2178 - self._set_cuda_visible_devices(

2178 + self._set_assigned_physical_gpu_ids(

2179 vllm_config, local_dp_rank, device_control_env_var

2180 )

2181

2182 - def _set_cuda_visible_devices(

2183 - self, vllm_config: VllmConfig, local_dp_rank: int, device_control_env_var: str

2182 + def _set_assigned_physical_gpu_ids(

2183 + self,

2184 + vllm_config: VllmConfig,

2185 + local_dp_rank: int,

2186 + device_control_env_var: str,

2187 ):

2188 world_size = vllm_config.parallel_config.world_size

2186 - # Set CUDA_VISIBLE_DEVICES or equivalent.

2189 try:

2188 - value = get_device_indices(

2189 - device_control_env_var, local_dp_rank, world_size

2190 + physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(

2191 + device_control_env_var,

2192 + local_dp_rank,

2193 + world_size,

2194 + user_assigned_gpu_ids=(

2195 + vllm_config.parallel_config.assigned_physical_gpu_ids

2196 + ),

2197 )

2191 - os.environ[device_control_env_var] = value

2198 + vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids

2199 except IndexError as e:

2200 raise Exception(

2194 - f"Error setting {device_control_env_var}: "

2201 + f"Error computing assigned_physical_gpu_ids: "

2202 f"local range: [{local_dp_rank * world_size}, "

2203 f"{(local_dp_rank + 1) * world_size}) "

2204 f'base value: "{os.getenv(device_control_env_var)}"'

vllm/v1/engine/utils.py

+66 -43

@@ -12,7 +12,6 @@ from multiprocessing import Process, connection

12 from multiprocessing.process import BaseProcess

13 from multiprocessing.queues import Queue

14 from typing import TYPE_CHECKING, cast

15 -from unittest.mock import patch

16 import msgspec

17 import zmq

@@ -175,38 +174,38 @@ class CoreEngineProcManager:

174 self.manager_stopped = threading.Event()

175 self.failed_proc_name: str | None = None

176

177 + # All ranks share this config object: capture the user-provided

178 + # --device-ids list before the per-rank shard overwrites it. Mutating

179 + # the config before each proc.start() works because the spawn method

180 + # pickles process args at start() time, sequentially per rank.

181 + user_assigned_gpu_ids = vllm_config.parallel_config.assigned_physical_gpu_ids

182 try:

183 for proc, local_dp_rank in zip(self.processes, local_dp_ranks):

180 - # Adjust device control in DP for platforms that cannot rely

181 - # on torch.accelerator.set_device_index(), and for Ray launchers.

182 - device_control_context: contextlib.AbstractContextManager[None] = (

183 - contextlib.nullcontext()

184 - )

184 + # Populate the logical-to-physical GPU mapping in DP for

185 + # platforms that cannot rely on

186 + # torch.accelerator.set_device_index(), and for Ray.

187 needs_device_env_isolation = not (

188 current_platform.is_cuda_alike() or current_platform.is_xpu()

189 )

190 if is_dp and (

191 needs_device_env_isolation or vllm_config.parallel_config.use_ray

192 ):

191 - device_control_context = set_device_control_env_var(

192 - vllm_config, local_dp_rank

193 + set_assigned_physical_gpu_ids_for_dp_rank(

194 + vllm_config, local_dp_rank, user_assigned_gpu_ids

195 )

196

195 - with (

196 - device_control_context,

197 - numa_utils.configure_subprocess(

198 - # EngineCore itself does not have a TP/PP-local rank.

199 - # When DP is enabled, set_device_control_env_var()

200 - # narrows visible devices to this DP shard first, so

201 - # local_rank=0 means "the first local GPU in this

202 - # shard". The actual TP/PP worker processes spawned by

203 - # the executor are bound separately with their own

204 - # local_rank values.

205 - vllm_config,

206 - local_rank=0,

207 - dp_local_rank=local_dp_rank,

208 - process_kind="EngineCore",

209 - ),

197 + with numa_utils.configure_subprocess(

198 + # EngineCore itself does not have a TP/PP-local rank.

199 + # When DP is enabled, set_assigned_physical_gpu_ids_for_dp_rank()

200 + # populates the logical-to-physical mapping for this DP

201 + # shard, so local_rank=0 means "the first local GPU in

202 + # this shard". The actual TP/PP worker processes spawned

203 + # by the executor are bound separately with their own

204 + # local_rank values.

205 + vllm_config,

206 + local_rank=0,

207 + dp_local_rank=local_dp_rank,

208 + process_kind="EngineCore",

209 ):

210 proc.start()

211 finally:

@@ -281,55 +280,79 @@ class SignalCallback:

280 self._event.set()

281

282

284 -@contextlib.contextmanager

285 -def set_device_control_env_var(

286 - vllm_config: VllmConfig, local_dp_rank: int

287 -) -> Iterator[None]:

283 +def set_assigned_physical_gpu_ids_for_dp_rank(

284 + vllm_config: VllmConfig,

285 + local_dp_rank: int,

286 + user_assigned_gpu_ids: list[int] | None = None,

287 +) -> None:

288 """

289 - Temporarily set CUDA_VISIBLE_DEVICES or equivalent

290 - for engine subprocess.

289 + Populate assigned_physical_gpu_ids on the config for the given DP rank.

290 +

291 + user_assigned_gpu_ids is the full (un-sharded) --device-ids list, if the

292 + user provided one; this DP rank's shard is sliced from it. It is passed

293 + explicitly rather than read from the config because callers may reuse

294 + one config object across DP ranks, overwriting the field each time.

295 """

296 world_size = vllm_config.parallel_config.world_size

297 local_world_size = vllm_config.parallel_config.local_world_size

298 evar = current_platform.device_control_env_var

299

296 - value = get_device_indices(evar, local_dp_rank, world_size, local_world_size)

297 - with patch.dict(os.environ, values=((evar, value),)):

298 - yield

300 + physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(

301 + evar,

302 + local_dp_rank,

303 + world_size,

304 + local_world_size,

305 + user_assigned_gpu_ids=user_assigned_gpu_ids,

306 + )

307 + vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids

308

309

301 -def get_device_indices(

310 +def get_physical_gpu_ids_for_local_dp_rank(

311 device_control_env_var: str,

312 local_dp_rank: int,

313 world_size: int,

314 local_world_size: int | None = None,

306 -):

315 + user_assigned_gpu_ids: list[int] | None = None,

316 +) -> list[int]:

317 """

308 - Returns a comma-separated string of device indices for the specified

318 + Returns list of physical GPU IDs for the specified

319 data parallel rank.

320

321 For example, if world_size=2 and local_dp_rank=1, and there are 4 devices,

312 - this will select devices 2 and 3 for local_dp_rank=1.

322 + this will return [2, 3] for local_dp_rank=1.

323 +

324 + If user_assigned_gpu_ids is provided (e.g. from --device-ids), this DP

325 + rank's shard is sliced from it instead of being derived from the

326 + device-control env var.

327 """

328 if local_world_size is None:

329 local_world_size = world_size

330 + if user_assigned_gpu_ids is not None:

331 + start = local_dp_rank * world_size

332 + stop = start + local_world_size

333 + if stop > len(user_assigned_gpu_ids):

334 + raise ValueError(

335 + f"--device-ids provides {len(user_assigned_gpu_ids)} devices, "

336 + f"but DP rank {local_dp_rank} needs devices [{start}, {stop})"

337 + )

338 + return user_assigned_gpu_ids[start:stop]

339 try:

317 - value = ",".join(

318 - str(current_platform.device_id_to_physical_device_id(i))

340 + return [

341 + current_platform.device_id_to_physical_device_id(i)

342 for i in range(

343 local_dp_rank * world_size,

344 local_dp_rank * world_size + local_world_size,

345 )

323 - )

346 + ]

347 except IndexError as e:

348 raise Exception(

326 - f"Error setting {device_control_env_var}: "

349 + f"Error computing device indices for "

350 + f"{device_control_env_var}: "

351 f"local range: [{local_dp_rank * world_size}, "

352 f"{(local_dp_rank + 1) * world_size}) "

353 "base value: "

354 f'"{os.getenv(device_control_env_var)}"'

355 ) from e

332 - return value

356

357

358 def _apply_dp_identity_suffix(dp_vllm_config, dp_rank: int) -> None:

@@ -453,11 +476,11 @@ class CoreEngineActorManager:

476 # https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/intel_gpu.py#L56 # noqa: E501

477 if current_platform.is_xpu():

478 device_evar = current_platform.device_control_env_var

456 - device_indices = get_device_indices(

479 + physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(

480 device_evar, local_index, world_size

481 )

482 actor_env_vars = self.env_vars_dict.copy()

460 - actor_env_vars[device_evar] = device_indices

483 + actor_env_vars[device_evar] = ",".join(str(d) for d in physical_gpu_ids)

484 runtime_env = RuntimeEnv(env_vars=actor_env_vars)

485

486 actor = (

vllm/v1/executor/multiproc_executor.py

+10

@@ -826,6 +826,16 @@ class WorkerProc:

826 signal.signal(signal.SIGTERM, signal_handler)

827 signal.signal(signal.SIGINT, signal_handler)

828

829 + # Publish the logical-to-physical mapping early so topology helpers

830 + # work before init_device (needed by set_worker_net_device below).

831 + assigned_physical_gpu_ids = kwargs[

832 + "vllm_config"

833 + ].parallel_config.assigned_physical_gpu_ids

834 + if assigned_physical_gpu_ids is not None:

835 + from vllm.platforms.interface import set_assigned_physical_gpu_ids

836 +

837 + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)

838 +

839 # Set net device env vars for the worker if VLLM_GPU_NIC_PCIE_MAPPING is set

840 set_worker_net_device(kwargs.get("local_rank", 0), kwargs["vllm_config"])

841

vllm/v1/executor/ray_executor.py

+22 -31

@@ -258,30 +258,35 @@ class RayDistributedExecutor(Executor):

258 }

259 self.collective_rpc("adjust_rank", args=(rerank_mapping,))

260

261 - # Get the set of GPU IDs used on each node.

262 - worker_node_and_gpu_ids = []

261 + # Get the set of physical GPU IDs used on each node.

262 + worker_node_and_physical_gpu_ids = []

263 for worker in [self.driver_dummy_worker] + self.workers:

264 if worker is None:

265 # driver_dummy_worker can be None when using ray spmd worker.

266 continue

267 - worker_node_and_gpu_ids.append(

268 - ray.get(worker.get_node_and_gpu_ids.remote()) # type: ignore[attr-defined]

267 + worker_node_and_physical_gpu_ids.append(

268 + ray.get(worker.get_node_and_physical_gpu_ids.remote()) # type: ignore[attr-defined]

269 )

270

271 node_workers = defaultdict(list) # node id -> list of worker ranks

272 - node_gpus = defaultdict(list) # node id -> list of gpu ids

272 + node_physical_gpu_ids = defaultdict(list) # node id -> physical GPU IDs

273

274 - for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids):

274 + for i, (node_id, physical_gpu_ids) in enumerate(

275 + worker_node_and_physical_gpu_ids

276 + ):

277 node_workers[node_id].append(i)

276 - # `gpu_ids` can be a list of strings or integers.

278 + # `physical_gpu_ids` can be a list of strings or integers.

279 # convert them to integers for consistency.

278 - # NOTE: gpu_ids can be larger than 9 (e.g. 16 GPUs),

280 + # NOTE: physical GPU IDs can be larger than 9 (e.g. 16 GPUs),

281 # string sorting is not sufficient.

282 # see https://github.com/vllm-project/vllm/issues/5590

281 - gpu_ids = [int(x) for x in gpu_ids]

282 - node_gpus[node_id].extend(gpu_ids)

283 - for node_id, gpu_ids in node_gpus.items():

284 - node_gpus[node_id] = sorted(gpu_ids)

283 + physical_gpu_ids = [

284 + current_platform.device_control_id_to_physical_device_id(str(x))

285 + for x in physical_gpu_ids

286 + ]

287 + node_physical_gpu_ids[node_id].extend(physical_gpu_ids)

288 + for node_id, physical_gpu_ids in node_physical_gpu_ids.items():

289 + node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids)

290

291 all_ips = set(worker_ips + [driver_ip])

292 n_ips = len(all_ips)

@@ -297,23 +302,8 @@ class RayDistributedExecutor(Executor):

302 " each node."

303 )

304

300 - # Set environment variables for the driver and workers.

301 - # We set CUDA_VISIBLE_DEVICES to ALL GPUs on the node for each worker.

302 - # This is needed because:

303 - # 1. Ray's compiled DAG needs to find the allocated GPU in

304 - # CUDA_VISIBLE_DEVICES.

305 - # 2. vLLM's communication layer (NCCL, CustomAllreduce) needs to see

306 - # all GPUs for P2P checks and communication setup. Though if it was

307 - # just this reason, we could have also just kept the visible devices

308 - # unset.

309 - # Each worker will use local_rank to index into the visible devices.

310 - all_args_to_update_environment_variables = [

311 - {

312 - current_platform.device_control_env_var: ",".join(

313 - map(str, node_gpus[node_id])

314 - ),

315 - }

316 - for (node_id, _) in worker_node_and_gpu_ids

305 + all_args_to_update_environment_variables: list[dict[str, str]] = [

306 + {} for _ in worker_node_and_physical_gpu_ids

307 ]

308

309 # Environment variables to copy from driver to workers

@@ -336,7 +326,7 @@ class RayDistributedExecutor(Executor):

326 "update_environment_variables", args=(self._get_env_vars_to_be_updated(),)

327 )

328

339 - if len(node_gpus) == 1:

329 + if len(node_physical_gpu_ids) == 1:

330 # in single node case, we don't need to get the IP address.

331 # the loopback address is sufficient

332 # NOTE: a node may have several IP addresses, one for each

@@ -352,10 +342,11 @@ class RayDistributedExecutor(Executor):

342

343 # Initialize the actual workers inside worker wrapper.

344 all_kwargs = []

355 - for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids):

345 + for rank, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids):

346 local_rank = node_workers[node_id].index(rank)

347 kwargs = dict(

348 vllm_config=self.vllm_config,

349 + assigned_physical_gpu_ids=sorted(node_physical_gpu_ids[node_id]),

350 local_rank=local_rank,

351 rank=rank,

352 distributed_init_method=distributed_init_method,

vllm/v1/executor/ray_executor_v2.py

+56 -31

@@ -79,24 +79,25 @@ class RayWorkerProc(WorkerProc):

79 1. __init__: lightweight setup, stores init args (no device/model init)

80 2. initialize_worker: called after GPU IDs are discovered, completes

81 the full WorkerProc initialization with the correct local_rank and

82 - CUDA_VISIBLE_DEVICES.

82 + logical-to-physical GPU mapping.

84 - CUDA_VISIBLE_DEVICES setup flow:

84 + GPU assignment flow:

86 1. RayExecutorV2 enables RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES so Ray does

87 not set CUDA_VISIBLE_DEVICES on RayWorkerProc actors at creation time.

88 2. Each actor is scheduled with a placement group and bundle index; Ray resolves

89 the physical GPU ID for that bundle at placement time.

90 - 3. After placement, the worker discovers that GPU ID and sets

91 - CUDA_VISIBLE_DEVICES before finishing WorkerProc initialization.

90 + 3. After placement, the executor discovers each worker's GPU ID and passes the

91 + node's logical-to-physical mapping (assigned_physical_gpu_ids) to

92 + initialize_worker(); CUDA_VISIBLE_DEVICES is never modified.

93 - There is no workaround for this unset-and-reset sequence when the placement group

94 - is externally managed: scheduling must complete before CUDA_VISIBLE_DEVICES can

95 - match the GPU tied to the worker's bundle.

94 + Scheduling must complete before the mapping is known when the placement

95 + group is externally managed: only then is the GPU tied to the worker's

96 + bundle resolved.

98 This sequence allows multiple vLLM instances to coexist on the same node:

99 each instance is unaware which physical devices others hold, and the

99 - externally managed placement group avoids CUDA_VISIBLE_DEVICES conflicts

100 + externally managed placement group avoids device assignment conflicts

101 by binding workers to specific placement group bundles.

102 """

103

@@ -120,28 +121,33 @@ class RayWorkerProc(WorkerProc):

121 is_driver_worker=is_driver_worker,

122 )

123

123 - def get_node_and_gpu_ids(self) -> tuple[str, list[int]]:

124 - """Return (node_id, gpu_ids) assigned to this actor by Ray."""

124 + def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]:

125 + """Return (node_id, physical_gpu_ids) assigned to this actor by Ray."""

126 node_id = ray.get_runtime_context().get_node_id()

127 device_key = current_platform.ray_device_key

128 if not device_key:

129 raise RuntimeError(

130 f"current platform {current_platform.device_name} does not support ray."

131 )

131 - gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]

132 - return node_id, [int(x) for x in gpu_ids]

132 + physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]

133 + return node_id, [

134 + current_platform.device_control_id_to_physical_device_id(str(x))

135 + for x in physical_gpu_ids

136 + ]

137

138 def initialize_worker(

139 self,

140 local_rank: int,

141 env_vars: dict[str, str],

142 driver_env_vars: dict[str, str] | None = None,

143 + assigned_physical_gpu_ids: list[int] | None = None,

144 ) -> None:

145 """Complete initialization after GPU assignment is known.

146

147 *driver_env_vars* are applied with ``setdefault`` — they fill

148 in missing vars but never overwrite node-local values.

144 - *env_vars* (e.g. CUDA_VISIBLE_DEVICES) always overwrite.

149 + *env_vars* always overwrite.

150 + *assigned_physical_gpu_ids* maps local_rank to physical CUDA device ID.

151 """

152 if driver_env_vars:

153 for key, value in driver_env_vars.items():

@@ -149,6 +155,13 @@ class RayWorkerProc(WorkerProc):

155 for key, value in env_vars.items():

156 os.environ[key] = value

157

158 + if assigned_physical_gpu_ids is not None:

159 + vllm_config = self._init_kwargs["vllm_config"]

160 + assert isinstance(vllm_config, VllmConfig)

161 + vllm_config.parallel_config.assigned_physical_gpu_ids = (

162 + assigned_physical_gpu_ids

163 + )

164 +

165 self.local_rank = local_rank

166 super().__init__(

167 local_rank=local_rank,

@@ -365,36 +378,48 @@ class RayExecutorV2(MultiprocExecutor):

378 )

379 self.ray_worker_handles.append(handle)

380

368 - # Step 6: Discover GPU IDs assigned to each worker via Ray runtime context.

369 - worker_node_and_gpu_ids = ray.get(

370 - [h.actor.get_node_and_gpu_ids.remote() for h in self.ray_worker_handles]

381 + # Step 6: Discover physical GPU IDs assigned to each worker via Ray

382 + # runtime context.

383 + worker_node_and_physical_gpu_ids = ray.get(

384 + [

385 + h.actor.get_node_and_physical_gpu_ids.remote()

386 + for h in self.ray_worker_handles

387 + ]

388 )

389

390 node_workers: dict[str, list[int]] = defaultdict(list)

374 - node_gpus: dict[str, list[int]] = defaultdict(list)

375 - for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids):

391 + node_physical_gpu_ids: dict[str, list[int]] = defaultdict(list)

392 + for i, (node_id, physical_gpu_ids) in enumerate(

393 + worker_node_and_physical_gpu_ids

394 + ):

395 node_workers[node_id].append(i)

377 - node_gpus[node_id].extend(gpu_ids)

378 - for node_id, gpu_ids in node_gpus.items():

379 - node_gpus[node_id] = sorted(gpu_ids)

396 + node_physical_gpu_ids[node_id].extend(physical_gpu_ids)

397 + for node_id, physical_gpu_ids in node_physical_gpu_ids.items():

398 + node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids)

399

381 - # Step 7: Initialize workers with correct local_rank and

382 - # CUDA_VISIBLE_DEVICES. Each worker sees all GPUs assigned to

383 - # this executor on its node; local_rank indexes into that set.

400 + # Step 7: Initialize workers with local logical ranks and the

401 + # logical-to-physical GPU mapping discovered from Ray placement.

402 init_worker_refs = []

385 - for i, (node_id, _) in enumerate(worker_node_and_gpu_ids):

403 + for i, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids):

404 local_rank = node_workers[node_id].index(i)

387 - worker_env_vars = {

388 - current_platform.device_control_env_var: ",".join(

389 - map(str, node_gpus[node_id])

390 - ),

391 - }

405 + assigned_physical_gpu_ids = sorted(node_physical_gpu_ids[node_id])

406 + worker_env_vars: dict[str, str] = {}

407 self.ray_worker_handles[i].local_rank = local_rank

408 init_worker_refs.append(

409 self.ray_worker_handles[i].actor.initialize_worker.remote(

395 - local_rank, worker_env_vars, self.driver_env_vars

410 + local_rank,

411 + worker_env_vars,

412 + self.driver_env_vars,

413 + assigned_physical_gpu_ids=assigned_physical_gpu_ids,

414 )

415 )

416 + # Also set on the executor-side config for consistency. The mapping

417 + # is per-node, so only do this when all workers share one node.

418 + if len(node_physical_gpu_ids) == 1:

419 + node_id_0 = worker_node_and_physical_gpu_ids[0][0]

420 + self.vllm_config.parallel_config.assigned_physical_gpu_ids = sorted(

421 + node_physical_gpu_ids[node_id_0]

422 + )

423 ray.get(init_worker_refs)

424

425 # Step 8: Collect response MQ handles

vllm/v1/executor/ray_utils.py

+5 -3

@@ -93,7 +93,7 @@ try:

93 def get_node_ip(self) -> str:

94 return get_ip()

96 - def get_node_and_gpu_ids(self) -> tuple[str, list[int]]:

96 + def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]:

97 node_id = ray.get_runtime_context().get_node_id()

98 device_key = vllm.platforms.current_platform.ray_device_key

99 if not device_key:

@@ -101,8 +101,10 @@ try:

101 "current platform %s does not support ray.",

102 vllm.platforms.current_platform.device_name,

103 )

104 - gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]

105 - return node_id, gpu_ids

104 + physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[

105 + device_key

106 + ]

107 + return node_id, physical_gpu_ids

108

109 def setup_device_if_necessary(self):

110 # TODO(swang): This is needed right now because Ray CG executes

vllm/v1/worker/gpu_worker.py

+38 -10

@@ -270,19 +270,47 @@ class Worker(WorkerBase):

270

271 # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK

272 self.local_rank += dp_local_rank * tp_pp_world_size

273 - assert self.local_rank < torch.accelerator.device_count(), (

274 - f"DP adjusted local rank {self.local_rank} is out of bounds. "

275 - )

276 - visible_device_count = (

277 - torch.accelerator.device_count() if torch.cuda.is_available() else 0

273 +

274 + # Publish the logical-to-physical mapping for topology queries

275 + # such as NIC affinity and P2P checks.

276 + assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids

277 + if assigned_physical_gpu_ids is not None:

278 + from vllm.platforms.interface import set_assigned_physical_gpu_ids

279 +

280 + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)

281 + assert self.local_rank < len(assigned_physical_gpu_ids), (

282 + f"local_rank {self.local_rank} is out of bounds for "

283 + f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}"

284 )

279 - assert self.parallel_config.local_world_size <= visible_device_count, (

280 - f"local_world_size ({self.parallel_config.local_world_size}) must "

281 - f"be less than or equal to the number of visible devices "

282 - f"({visible_device_count})."

285 + # NOTE(patch pr45026): local_world_size is derived from

286 + # parallel_config.nnodes, which is only set for the "mp"

287 + # multi-node backend. With the "ray"/"external_launcher"

288 + # backends nnodes stays 1, so local_world_size collapses to

289 + # the full world_size and this check wrongly fires on

290 + # cross-node deployments. assigned_physical_gpu_ids is already

291 + # per-node and the local_rank bound above fully validates the

292 + # mapping for these backends, so skip the check for them.

293 + if parallel_config.distributed_executor_backend not in (

294 + "ray",

295 + "external_launcher",

296 + ):

297 + assert self.parallel_config.local_world_size <= len(

298 + assigned_physical_gpu_ids

299 + ), (

300 + f"local_world_size ({self.parallel_config.local_world_size})"

301 + " exceeds assigned_physical_gpu_ids count "

302 + f"({len(assigned_physical_gpu_ids)})"

303 + )

304 + else:

305 + assert self.local_rank < torch.accelerator.device_count(), (

306 + f"DP adjusted local rank {self.local_rank} is out of "

307 + f"bounds for {torch.accelerator.device_count()} devices."

308 )

309

285 - self.device = torch.device(f"cuda:{self.local_rank}")

310 + visible_device_index = (

311 + current_platform.logical_device_id_to_visible_device_id(self.local_rank)

312 + )

313 + self.device = torch.device(f"cuda:{visible_device_index}")

314 torch.accelerator.set_device_index(self.device)

315

316 current_platform.check_if_supports_dtype(self.model_config.dtype)

vllm/v1/worker/worker_base.py

@@ -286,6 +286,12 @@ class WorkerWrapperBase:

286 extended_calls,

287 )

288

289 + assigned_physical_gpu_ids = kwargs.pop("assigned_physical_gpu_ids", None)

290 + if assigned_physical_gpu_ids is not None:

291 + vllm_config.parallel_config.assigned_physical_gpu_ids = (

292 + assigned_physical_gpu_ids

293 + )

294 +

295 shared_worker_lock = kwargs.pop("shared_worker_lock", None)

296 if shared_worker_lock is None:

297 msg = (