From be4ce2e50545c6c1bf13f224022aeb2bcf0c0070 Mon Sep 17 00:00:00 2001 From: Sherif Waly Date: Wed, 14 Oct 2026 09:54:03 +0100 Subject: [PATCH] [Bugfix] Share PCP-aware local GPU indexing for NUMA binding (#61237) --- tests/utils_/test_numa_utils.py | 89 +++++++++++++++++++++- tests/v1/executor/test_vllm_net_devices.py | 25 ++++++ vllm/config/parallel.py | 16 ++++ vllm/utils/numa_utils.py | 38 +-------- vllm/v1/executor/vllm_net_devices.py | 31 +------- vllm/v1/worker/gpu_worker.py | 20 +---- 6 files changed, 134 insertions(+), 85 deletions(-) diff --git a/tests/utils_/test_numa_utils.py b/tests/utils_/test_numa_utils.py index 9f718703a7..3c09722f82 100644 --- a/tests/utils_/test_numa_utils.py +++ b/tests/utils_/test_numa_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os +from functools import partial from types import SimpleNamespace import pytest @@ -49,9 +50,18 @@ def _make_config(**parallel_kwargs): data_parallel_index=0, pipeline_parallel_size=1, tensor_parallel_size=1, + prefill_context_parallel_size=1, ) parallel_defaults.update(parallel_kwargs) parallel_config = SimpleNamespace(**parallel_defaults) + parallel_config.world_size = ( + parallel_config.tensor_parallel_size + * parallel_config.prefill_context_parallel_size + * parallel_config.pipeline_parallel_size + ) + parallel_config.get_local_gpu_index = partial( + ParallelConfig.get_local_gpu_index, parallel_config + ) return SimpleNamespace(parallel_config=parallel_config) @@ -368,13 +378,86 @@ def test_get_numactl_args_user_cpus_override_pct(monkeypatch): ) -def test_get_numactl_args_uses_dp_offset(): +@pytest.mark.parametrize("pcp,pp", [(1, 1), (2, 1), (2, 2)]) +def test_get_numactl_args_uses_dp_offset(pcp, pp): + workers_per_dp = 2 * pcp * pp vllm_config = _make_config( numa_bind=True, - numa_bind_nodes=[0, 0, 1, 1], + numa_bind_nodes=[0] * workers_per_dp + [1] * workers_per_dp, data_parallel_rank_local=1, - pipeline_parallel_size=1, + pipeline_parallel_size=pp, tensor_parallel_size=2, + prefill_context_parallel_size=pcp, + ) + assert ( + numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1) + == "--cpunodebind=1 --membind=1" + ) + + +@pytest.mark.parametrize( + "local_dp_rank,global_dp_rank,launch_dp_rank", + [(1, 3, None), (None, 1, None), (0, 0, 1)], + ids=["local-rank", "global-fallback", "launch-override"], +) +def test_pcp_worker_cpu_binding_uses_dp_gpu_index( + local_dp_rank, global_dp_rank, launch_dp_rank +): + """CPU overrides use the same PCP-aware GPU index as device selection.""" + vllm_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0] * 4 + [1] * 4, + numa_bind_cpus=[f"{4 * i}-{4 * i + 3}" for i in range(8)], + tensor_parallel_size=2, + prefill_context_parallel_size=2, + data_parallel_rank_local=local_dp_rank, + data_parallel_index=global_dp_rank, + ) + assert ( + numa_utils._get_numactl_worker_args( + vllm_config.parallel_config, local_rank=1, dp_local_rank=launch_dp_rank + ) + == "--physcpubind=20-23 --membind=1" + ) + + +@pytest.mark.parametrize("dp_rank,expected_nodes", [(0, "0,1"), (1, "2,3")]) +def test_pcp_engine_core_covers_all_shard_workers(dp_rank, expected_nodes): + """EngineCore includes every PCP worker's NUMA node in its DP shard.""" + vllm_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0, 0, 1, 1, 2, 2, 3, 3], + tensor_parallel_size=2, + prefill_context_parallel_size=2, + data_parallel_rank_local=dp_rank, + ) + assert ( + numa_utils._get_numactl_enginecore_args( + vllm_config.parallel_config, local_rank=0 + ) + == f"--cpunodebind={expected_nodes} --membind={expected_nodes}" + ) + + +@pytest.mark.parametrize( + "parallel_kwargs", + [ + {"distributed_executor_backend": "ray"}, + {"distributed_executor_backend": "external_launcher"}, + {"data_parallel_backend": "ray"}, + {"nnodes_within_dp": 2}, + ], + ids=["ray-executor", "external-launcher", "ray-dp", "multi-node-dp"], +) +def test_pcp_worker_keeps_local_gpu_index_without_dp_offset(parallel_kwargs): + """Node-local launcher ranks do not acquire a colocated DP offset.""" + vllm_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0, 1], + tensor_parallel_size=2, + prefill_context_parallel_size=2, + data_parallel_rank_local=1, + **parallel_kwargs, ) assert ( numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1) diff --git a/tests/v1/executor/test_vllm_net_devices.py b/tests/v1/executor/test_vllm_net_devices.py index 366dd001f2..214ff1af80 100644 --- a/tests/v1/executor/test_vllm_net_devices.py +++ b/tests/v1/executor/test_vllm_net_devices.py @@ -1,8 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest +from vllm.config import ParallelConfig +from vllm.v1.executor import vllm_net_devices from vllm.v1.executor.vllm_net_devices import normalize_pci @@ -74,3 +78,24 @@ def test_normalize_pci_device_out_of_range_raises(): def test_normalize_pci_empty_string_raises(): with pytest.raises(ValueError): normalize_pci("") + + +@pytest.mark.parametrize("pcp,pp", [(1, 1), (2, 1), (2, 2)]) +def test_set_worker_net_device_uses_pcp_aware_gpu_index(monkeypatch, pcp, pp): + """NIC lookup uses the same DP-offset GPU index as device selection.""" + monkeypatch.setenv("VLLM_GPU_NIC_PCIE_MAPPING", "unused") + monkeypatch.setenv("VLLM_NIC_SELECTION_VARS", "unused") + seen: list[int] = [] + monkeypatch.setattr(vllm_net_devices, "set_worker_gpu_nic_mapping", seen.append) + parallel_config = ParallelConfig( + tensor_parallel_size=2, + prefill_context_parallel_size=pcp, + pipeline_parallel_size=pp, + data_parallel_size=2, + data_parallel_rank_local=1, + distributed_executor_backend="mp", + ) + vllm_net_devices.set_worker_net_device( + 1, SimpleNamespace(parallel_config=parallel_config) + ) + assert seen == [1 + 2 * pcp * pp] diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 73d83442ba..6716d5a21b 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -838,6 +838,22 @@ class ParallelConfig: def local_world_size(self) -> int: return self.world_size // self.nnodes_within_dp + def get_local_gpu_index( + self, local_rank: int, dp_local_rank: int | None = None + ) -> int: + """Map a worker's local rank to its host-local logical GPU index.""" + if ( + self.distributed_executor_backend not in ("ray", "external_launcher") + and self.data_parallel_backend != "ray" + and self.nnodes_within_dp == 1 + ): + if dp_local_rank is None: + dp_local_rank = self.data_parallel_rank_local + if dp_local_rank is None: + dp_local_rank = self.data_parallel_index + return local_rank + dp_local_rank * self.world_size + return local_rank + @staticmethod def has_unfinished_dp(dp_group: ProcessGroup, has_unfinished: bool) -> bool: tensor = torch.tensor([has_unfinished], dtype=torch.int32, device="cpu") diff --git a/vllm/utils/numa_utils.py b/vllm/utils/numa_utils.py index 2e52935ea6..755d46e711 100644 --- a/vllm/utils/numa_utils.py +++ b/vllm/utils/numa_utils.py @@ -223,29 +223,6 @@ def _pct_sku_config() -> _PctSku | None: return sku -def _get_gpu_index( - parallel_config, local_rank: int, dp_local_rank: int | None = None -) -> int: - """Compute the physical GPU index used for NUMA lookup.""" - if ( - parallel_config.distributed_executor_backend not in ("ray", "external_launcher") - and parallel_config.data_parallel_backend != "ray" - and parallel_config.nnodes_within_dp == 1 - ): - if dp_local_rank is None: - dp_local_rank = parallel_config.data_parallel_rank_local - if dp_local_rank is None: - dp_local_rank = parallel_config.data_parallel_index - - tp_pp_world_size = ( - parallel_config.pipeline_parallel_size - * parallel_config.tensor_parallel_size - ) - return local_rank + dp_local_rank * tp_pp_world_size - - return local_rank - - def _get_numa_node(parallel_config, gpu_index: int) -> int: numa_nodes = parallel_config.numa_bind_nodes if numa_nodes is None: @@ -346,7 +323,7 @@ def _get_numactl_worker_args( parallel_config, local_rank: int, dp_local_rank: int | None = None ) -> str: """Compute the numactl args for a single TP/PP worker subprocess.""" - gpu_index = _get_gpu_index(parallel_config, local_rank, dp_local_rank) + gpu_index = parallel_config.get_local_gpu_index(local_rank, dp_local_rank) numa_node = _get_numa_node(parallel_config, gpu_index) cpu_binding = _get_cpu_binding(parallel_config, gpu_index, [numa_node]) @@ -384,17 +361,8 @@ def _get_enginecore_numa_nodes( and parallel_config.data_parallel_backend != "ray" and parallel_config.nnodes_within_dp == 1 ): - if dp_local_rank is None: - dp_local_rank = parallel_config.data_parallel_rank_local - if dp_local_rank is None: - dp_local_rank = parallel_config.data_parallel_index - - tp_pp_world_size = ( - parallel_config.pipeline_parallel_size - * parallel_config.tensor_parallel_size - ) - shard_start = dp_local_rank * tp_pp_world_size - shard_end = min(shard_start + tp_pp_world_size, len(numa_nodes)) + shard_start = parallel_config.get_local_gpu_index(0, dp_local_rank) + shard_end = min(shard_start + parallel_config.world_size, len(numa_nodes)) shard_indices: range | tuple[int, ...] = range(shard_start, shard_end) else: shard_indices = range(len(numa_nodes)) diff --git a/vllm/v1/executor/vllm_net_devices.py b/vllm/v1/executor/vllm_net_devices.py index be442f35ee..043e3dec9b 100644 --- a/vllm/v1/executor/vllm_net_devices.py +++ b/vllm/v1/executor/vllm_net_devices.py @@ -191,32 +191,6 @@ def set_worker_gpu_nic_mapping(local_rank: int) -> None: ) -def _dp_adjusted_local_rank(tp_local_rank: int, vllm_config: VllmConfig) -> int: - """Compute the node-wide GPU index accounting for data parallelism. - - On CUDA-alike platforms without env-var device isolation (the common - MP-backend path), the worker sees *all* GPUs on the node and selects - its device via ``torch.accelerator.set_device_index()`` using:: - - dp_local_rank * tp_pp_world_size + tp_local_rank - - This mirrors the adjustment in ``Worker.init_device()`` so we resolve - the correct GPU PCI address *before* the CUDA device is initialised. - """ - pc = vllm_config.parallel_config - if ( - pc.distributed_executor_backend not in ("ray", "external_launcher") - and pc.data_parallel_backend != "ray" - and pc.nnodes_within_dp == 1 - ): - dp_local_rank = pc.data_parallel_rank_local - if dp_local_rank is None: - dp_local_rank = pc.data_parallel_index - tp_pp_world_size = pc.pipeline_parallel_size * pc.tensor_parallel_size - return dp_local_rank * tp_pp_world_size + tp_local_rank - return tp_local_rank - - def set_worker_net_device(local_rank: int, vllm_config: VllmConfig) -> None: """Top-level entry point for both UniProcExecutor and MultiprocExecutor. @@ -239,5 +213,6 @@ def set_worker_net_device(local_rank: int, vllm_config: VllmConfig) -> None: if not has_pcie_mapping and not has_selection_vars: return # Both env vars are present, so set the NIC selection env vars. - adjusted_rank = _dp_adjusted_local_rank(local_rank, vllm_config) - set_worker_gpu_nic_mapping(adjusted_rank) + # Resolve the node-wide GPU index before the CUDA device is initialised. + gpu_index = vllm_config.parallel_config.get_local_gpu_index(local_rank) + set_worker_gpu_nic_mapping(gpu_index) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index aaab37a496..62a11d1eab 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -444,25 +444,7 @@ class Worker(WorkerBase): # This env var set by Ray causes exceptions with graph building. os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) parallel_config = self.parallel_config - if ( - parallel_config.distributed_executor_backend - not in ("ray", "external_launcher") - and parallel_config.data_parallel_backend != "ray" - and parallel_config.nnodes_within_dp == 1 - ): - # Use local DP rank if available, otherwise use global DP rank. - dp_local_rank = self.parallel_config.data_parallel_rank_local - if dp_local_rank is None: - dp_local_rank = self.parallel_config.data_parallel_index - - tp_pcp_pp_world_size = ( - self.parallel_config.pipeline_parallel_size - * self.parallel_config.prefill_context_parallel_size - * self.parallel_config.tensor_parallel_size - ) - - # DP_LOCAL_RANK * TP_PCP_PP_WORLD_SIZE + TP_LOCAL_RANK - self.local_rank += dp_local_rank * tp_pcp_pp_world_size + self.local_rank = parallel_config.get_local_gpu_index(self.local_rank) # Publish the logical-to-physical mapping for topology queries # such as NIC affinity and P2P checks.