iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix] Share PCP-aware local GPU indexing for NUMA binding (#61237)

Browse Source
This commit is contained in:
Sherif Waly
2026-10-14 10:54:03 +02:00
committed by GitHub
parent 937d6c19df
commit be4ce2e505
6 changed files with 134 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/utils_/test_numa_utils.py
+86 -3
View File
@@ -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)
tests/v1/executor/test_vllm_net_devices.py
+25
View File
@@ -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]
vllm/config/parallel.py
+16
View File
@@ -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")
vllm/utils/numa_utils.py
+3 -35
View File
@@ -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))
vllm/v1/executor/vllm_net_devices.py
+3 -28
View File
@@ -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)
vllm/v1/worker/gpu_worker.py
+1 -19
View File
@@ -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.