A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 12:15 UTC
[Feat] Add support for per GPU worker RDMA NIC selection (#42083)
Signed-off-by: Raj Joshi <rajjoshi@redhat.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Raj Joshi committed
May 28, 2026 at 15:45 UTC
9090368b650896bf5fc990c921df7eb4c20355a5
7 files changed
+365
tests/v1/executor/test_vllm_net_devices.py
new
+76
@@ -0,0 +1,76 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.v1.executor.vllm_net_devices import normalize_pci
7
+
8
+
9
+@pytest.mark.parametrize(
10
+ "addr, expected",
11
+ [
12
+ ("0000:3f:00.0", (0, 0x3F, 0, 0)),
13
+ ("0001:00:00.0", (1, 0, 0, 0)),
14
+ ("00000001:00:00.0", (1, 0, 0, 0)),
15
+ ("0009:00:00.0", (9, 0, 0, 0)),
16
+ ("0105:00:00.0", (0x105, 0, 0, 0)),
17
+ ("0000:0a:1f.7", (0, 0x0A, 0x1F, 7)),
18
+ ],
19
+)
20
+def test_normalize_pci_full_domain(addr, expected):
21
+ assert normalize_pci(addr) == expected
22
+
23
+
24
+@pytest.mark.parametrize(
25
+ "addr, expected",
26
+ [
27
+ ("01:00.0", (0, 1, 0, 0)),
28
+ ("3f:00.0", (0, 0x3F, 0, 0)),
29
+ ("40:00.0", (0, 0x40, 0, 0)),
30
+ ("ff:1f.7", (0, 0xFF, 0x1F, 7)),
31
+ ],
32
+)
33
+def test_normalize_pci_short_form(addr, expected):
34
+ assert normalize_pci(addr) == expected
35
+
36
+
37
+def test_normalize_pci_case_insensitive():
38
+ assert normalize_pci("0A:1F.7") == normalize_pci("0a:1f.7")
39
+
40
+
41
+def test_normalize_pci_strips_whitespace():
42
+ assert normalize_pci(" 0001:00:00.0 ") == (1, 0, 0, 0)
43
+
44
+
45
+def test_normalize_pci_strips_0x_prefix():
46
+ assert normalize_pci("0x0001:00:00.0") == (1, 0, 0, 0)
47
+
48
+
49
+def test_normalize_pci_missing_function_raises():
50
+ with pytest.raises(ValueError, match="missing function suffix"):
51
+ normalize_pci("0001:00:00")
52
+
53
+
54
+def test_normalize_pci_invalid_function_char_raises():
55
+ with pytest.raises(ValueError, match="invalid PCI function"):
56
+ normalize_pci("0001:00:00.z")
57
+
58
+
59
+def test_normalize_pci_too_many_segments_raises():
60
+ with pytest.raises(ValueError, match="invalid PCI BDF"):
61
+ normalize_pci("a:b:c:d.0")
62
+
63
+
64
+def test_normalize_pci_bus_out_of_range_raises():
65
+ with pytest.raises(ValueError, match="out of range"):
66
+ normalize_pci("0000:1ff:00.0")
67
+
68
+
69
+def test_normalize_pci_device_out_of_range_raises():
70
+ with pytest.raises(ValueError, match="out of range"):
71
+ normalize_pci("0000:00:20.0")
72
+
73
+
74
+def test_normalize_pci_empty_string_raises():
75
+ with pytest.raises(ValueError):
76
+ normalize_pci("")
vllm/envs.py
+9
@@ -278,6 +278,8 @@ if TYPE_CHECKING:
278
VLLM_XPU_ENABLE_XPU_GRAPH: bool = False
279
VLLM_XPU_USE_SAMPLER_KERNEL: bool = True
280
VLLM_LORA_ENABLE_DUAL_STREAM: bool = False
281
+ VLLM_GPU_NIC_PCIE_MAPPING: str = ""
282
+ VLLM_NIC_SELECTION_VARS: str = ""
283
284
285
def get_default_cache_root():
@@ -1973,6 +1975,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
1975
# If set to 1, use Python spinloop extension to poll in a more efficient
1976
# way when using the mp backend.
1977
"VLLM_USE_SPINLOOP_EXT": lambda: bool(int(os.getenv("VLLM_USE_SPINLOOP_EXT", "0"))),
1978
+ # Comma-separated GPU_BDF=NIC_BDF pairs for RDMA NIC selection.
1979
+ # Must be set together with VLLM_NIC_SELECTION_VARS.
1980
+ "VLLM_GPU_NIC_PCIE_MAPPING": lambda: os.getenv("VLLM_GPU_NIC_PCIE_MAPPING", ""),
1981
+ # Comma-separated list of env vars to set from the GPU-NIC mapping.
1982
+ # Each entry is VAR_NAME or VAR_NAME:<suffix> (suffix appended to
1983
+ # RDMA device name). Must be set together with VLLM_GPU_NIC_PCIE_MAPPING.
1984
+ "VLLM_NIC_SELECTION_VARS": lambda: os.getenv("VLLM_NIC_SELECTION_VARS", ""),
1985
}
1986
1987
vllm/platforms/cuda.py
+16
@@ -809,6 +809,22 @@ class NvmlCudaPlatform(CudaPlatformBase):
809
logger.warning("Failed to get NUMA nodes for GPUs: %s", e)
810
return None
811
812
+ @classmethod
813
+ @with_nvml_context
814
+ def get_all_gpu_pci_bus_ids(cls) -> dict[int, str]:
815
+ """Query NVML for GPU index -> PCI bus ID mapping."""
816
+ out: dict[int, str] = {}
817
+ for idx in range(pynvml.nvmlDeviceGetCount()):
818
+ handle = pynvml.nvmlDeviceGetHandleByIndex(idx)
819
+ pci_info = pynvml.nvmlDeviceGetPciInfo(handle)
820
+ bus_id = pci_info.busId
821
+ if isinstance(bus_id, bytes):
822
+ bus_id = bus_id.decode("utf-8")
823
+ out[idx] = bus_id.rstrip("\x00")
824
+ if not out:
825
+ raise RuntimeError("NVML returned no GPU PCI bus ID rows")
826
+ return out
827
+
828
@classmethod
829
@with_nvml_context
830
def log_warnings(cls):
vllm/platforms/interface.py
+13
@@ -381,6 +381,19 @@ class Platform:
381
"""Get the total memory of a device in bytes."""
382
raise NotImplementedError
383
384
+ @classmethod
385
+ def get_all_gpu_pci_bus_ids(cls) -> dict[int, str]:
386
+ """Return a mapping of device index to PCI bus ID string.
387
+
388
+ Used by ``VLLM_GPU_NIC_PCIE_MAPPING`` for RDMA NIC selection.
389
+ Subclasses should override with platform-specific discovery
390
+ (e.g. pynvml for CUDA).
391
+ """
392
+ raise NotImplementedError(
393
+ "VLLM_GPU_NIC_PCIE_MAPPING is not supported on the "
394
+ f"current platform ({cls.device_name})"
395
+ )
396
+
397
@classmethod
398
def inference_mode(cls):
399
"""A device-specific wrapper of `torch.inference_mode`.
vllm/v1/executor/multiproc_executor.py
+4
@@ -60,6 +60,7 @@ from vllm.utils.system_utils import (
60
)
61
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
62
from vllm.v1.executor.abstract import Executor, FailureCallback
63
+from vllm.v1.executor.vllm_net_devices import set_worker_net_device
64
from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput
65
from vllm.v1.worker.worker_base import WorkerWrapperBase
66
@@ -811,6 +812,9 @@ class WorkerProc:
812
signal.signal(signal.SIGTERM, signal_handler)
813
signal.signal(signal.SIGINT, signal_handler)
814
815
+ # Set net device env vars for the worker if VLLM_GPU_NIC_PCIE_MAPPING is set
816
+ set_worker_net_device(kwargs.get("local_rank", 0), kwargs["vllm_config"])
817
+
818
worker = None
819
ready_writer = kwargs.pop("ready_pipe")
820
death_pipe = kwargs.pop("death_pipe", None)
vllm/v1/executor/uniproc_executor.py
+4
@@ -16,6 +16,7 @@ from vllm.platforms import current_platform
16
from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
17
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
18
from vllm.v1.executor.abstract import Executor
19
+from vllm.v1.executor.vllm_net_devices import set_worker_net_device
20
from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput
21
from vllm.v1.serial_utils import run_method
22
from vllm.v1.worker.worker_base import WorkerWrapperBase
@@ -56,6 +57,9 @@ class UniProcExecutor(Executor):
57
shared_worker_lock=Lock(),
58
)
59
60
+ # Set net device env vars for the worker if VLLM_GPU_NIC_PCIE_MAPPING is set
61
+ set_worker_net_device(local_rank, self.vllm_config)
62
+
63
self.driver_worker.init_worker(all_kwargs=[kwargs])
64
self.driver_worker.init_device()
65
vllm/v1/executor/vllm_net_devices.py
new
+243
@@ -0,0 +1,243 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""GPU-to-NIC net-device mapping for RDMA transports (UCX, NVSHMEM, ...).
4
+
5
+Shared by both UniProcExecutor (TP=1) and MultiprocExecutor (TP>1).
6
+All transport-specific env vars and sysfs lookups live here so executor
7
+files only need to call ``set_worker_net_device(local_rank, vllm_config)``.
8
+
9
+Requires two env vars set together:
10
+- ``VLLM_GPU_NIC_PCIE_MAPPING`` -- comma-separated GPU_BDF=NIC_BDF pairs.
11
+- ``VLLM_NIC_SELECTION_VARS`` -- comma-separated list of env vars to set,
12
+ each optionally suffixed (e.g. ``UCX_NET_DEVICES:1,NCCL_IB_HCA:1``).
13
+"""
14
+
15
+import os
16
+from pathlib import Path
17
+
18
+import vllm.envs as envs
19
+from vllm.config import VllmConfig
20
+from vllm.logger import init_logger
21
+from vllm.platforms import current_platform
22
+
23
+logger = init_logger(__name__)
24
+
25
+
26
+def normalize_pci(addr: str) -> tuple[int, int, int, int]:
27
+ """Parse PCI BDF/domain-bus-device-function into comparable ints (all hex).
28
+
29
+ Supported shapes:
30
+ - ``domain:bus:dev.fn`` -- domain width varies (e.g. ``00000001:00:00.0``,
31
+ ``0001:00:00.0``, ``0000:3f:00.0``).
32
+ - ``bus:dev.fn`` -- domain **0** (e.g. ``01:00.0``, ``40:00.0``).
33
+
34
+ Function suffix is hex (typically ``0``--``7``). Raises ``ValueError`` if malformed.
35
+ """
36
+ s = addr.strip().lower().replace(" ", "")
37
+ if s.startswith("0x"):
38
+ s = s[2:]
39
+ if "." not in s:
40
+ raise ValueError(f"invalid PCI BDF (missing function suffix): {addr!r}")
41
+ body, fn_s = s.rsplit(".", 1)
42
+ if not fn_s or any(c not in "0123456789abcdef" for c in fn_s):
43
+ raise ValueError(f"invalid PCI function in BDF: {addr!r}")
44
+ fn = int(fn_s, 16)
45
+ if fn > 0xFF:
46
+ raise ValueError(f"PCI function out of range: {addr!r}")
47
+
48
+ parts = body.split(":")
49
+ if len(parts) == 2:
50
+ domain = 0
51
+ bus = int(parts[0], 16)
52
+ device = int(parts[1], 16)
53
+ elif len(parts) == 3:
54
+ domain = int(parts[0], 16)
55
+ bus = int(parts[1], 16)
56
+ device = int(parts[2], 16)
57
+ else:
58
+ raise ValueError(
59
+ f"invalid PCI BDF (want domain:bus:dev.fn or bus:dev.fn): {addr!r}"
60
+ )
61
+
62
+ if bus > 0xFF or device > 0x1F:
63
+ raise ValueError(f"PCI bus or device out of range: {addr!r}")
64
+ return (domain, bus, device, fn)
65
+
66
+
67
+def parse_gpu_nic_mapping(
68
+ raw: str,
69
+) -> dict[tuple[int, int, int, int], tuple[int, int, int, int]]:
70
+ out: dict[tuple[int, int, int, int], tuple[int, int, int, int]] = {}
71
+ for segment in raw.split(","):
72
+ segment = segment.strip()
73
+ if not segment:
74
+ continue
75
+ if "=" not in segment:
76
+ raise ValueError(
77
+ "VLLM_GPU_NIC_PCIE_MAPPING: expected comma-separated"
78
+ f" gpu_bdf=nic_bdf pairs; ambiguous segment: {segment!r}"
79
+ )
80
+ gpu_s, nic_s = segment.split("=", 1)
81
+ gpu_key = normalize_pci(gpu_s.strip())
82
+ nic_val = normalize_pci(nic_s.strip())
83
+ out[gpu_key] = nic_val
84
+ return out
85
+
86
+
87
+def rdma_name_for_nic_pci(nic_pci: tuple[int, int, int, int]) -> str:
88
+ """Map NIC PCI BDF to sysfs RDMA name (mlx5_*, ibp*, ...).
89
+
90
+ Under ``/sys/class/infiniband/<name>/``, ``device`` is a **symlink** to the PCI
91
+ device directory (e.g. ``.../0101:00:00.0``). We take ``Path(...).resolve().name``
92
+ as the BDF string.
93
+
94
+ ``VLLM_GPU_NIC_PCIE_MAPPING`` NIC keys must **normalize** (via ``normalize_pci``)
95
+ to the same tuple as this basename.
96
+ """
97
+ ib = Path("/sys/class/infiniband")
98
+ if not ib.is_dir():
99
+ raise RuntimeError("/sys/class/infiniband not found or not a directory")
100
+ names = sorted(p.name for p in ib.iterdir() if p.is_dir())
101
+ for name in names:
102
+ dev_link = ib / name / "device"
103
+ if not dev_link.exists():
104
+ continue
105
+ try:
106
+ resolved = dev_link.resolve()
107
+ except OSError:
108
+ continue
109
+ pci_name = resolved.name
110
+ try:
111
+ if normalize_pci(pci_name) == nic_pci:
112
+ return name
113
+ except ValueError:
114
+ continue
115
+ raise RuntimeError(
116
+ f"No /sys/class/infiniband device for NIC PCI {nic_pci}; have entries: {names}"
117
+ )
118
+
119
+
120
+def parse_nic_selection_vars(raw: str) -> list[tuple[str, str]]:
121
+ """Parse ``VLLM_NIC_SELECTION_VARS`` into ``(env_var_name, suffix)`` pairs.
122
+
123
+ Each entry is ``VAR_NAME`` or ``VAR_NAME:<suffix>``. The colon and
124
+ everything after it is appended verbatim to the RDMA device name.
125
+ """
126
+ result: list[tuple[str, str]] = []
127
+ for entry in raw.split(","):
128
+ entry = entry.strip()
129
+ if not entry:
130
+ continue
131
+ if ":" in entry:
132
+ var_name, suffix = entry.split(":", 1)
133
+ result.append((var_name, ":" + suffix))
134
+ else:
135
+ result.append((entry, ""))
136
+ return result
137
+
138
+
139
+def set_worker_gpu_nic_mapping(local_rank: int) -> None:
140
+ """Set NIC selection env vars from VLLM_GPU_NIC_PCIE_MAPPING for a worker.
141
+
142
+ Which env vars are set is controlled by ``VLLM_NIC_SELECTION_VARS``.
143
+ """
144
+ raw = envs.VLLM_GPU_NIC_PCIE_MAPPING.strip()
145
+ if not raw:
146
+ return
147
+ selection_raw = envs.VLLM_NIC_SELECTION_VARS.strip()
148
+ selection_vars = parse_nic_selection_vars(selection_raw)
149
+ mapping = parse_gpu_nic_mapping(raw)
150
+ pci_by_index = current_platform.get_all_gpu_pci_bus_ids()
151
+ # Translate CUDA-relative local_rank to the physical device index,
152
+ # which accounts for CUDA_VISIBLE_DEVICES narrowing (e.g. DP sharding).
153
+ physical_id = current_platform.device_id_to_physical_device_id(local_rank)
154
+ if physical_id not in pci_by_index:
155
+ raise RuntimeError(
156
+ f"No GPU PCI for physical device index {physical_id} "
157
+ f"(local_rank={local_rank}) in map "
158
+ f"(have indices {sorted(pci_by_index.keys())})"
159
+ )
160
+ gpu_bdf = pci_by_index[physical_id]
161
+ gpu_key = normalize_pci(gpu_bdf)
162
+ if gpu_key not in mapping:
163
+ keys_fmt = ", ".join(
164
+ f"{d:04x}:{b:02x}:{dev:02x}.{fn}"
165
+ for d, b, dev, fn in sorted(mapping.keys())
166
+ )
167
+ raise RuntimeError(
168
+ f"No VLLM_GPU_NIC_PCIE_MAPPING entry for GPU PCI {gpu_bdf} "
169
+ f"(worker local_rank={local_rank}); mapped GPUs: {keys_fmt}"
170
+ )
171
+ nic_pci = mapping[gpu_key]
172
+ rdma_dev = rdma_name_for_nic_pci(nic_pci)
173
+
174
+ set_vars: list[str] = []
175
+ for var_name, suffix in selection_vars:
176
+ value = f"{rdma_dev}{suffix}"
177
+ existing = os.environ.get(var_name, "").strip()
178
+ if existing:
179
+ value = f"{value},{existing}"
180
+ os.environ[var_name] = value
181
+ set_vars.append(f"{var_name}={value}")
182
+
183
+ nic_fmt = f"{nic_pci[0]:04x}:{nic_pci[1]:02x}:{nic_pci[2]:02x}.{nic_pci[3]}"
184
+ logger.info(
185
+ "GPU rank %s (PCIe addr %s) mapped to NIC %s (PCIe addr %s) via env vars: %s",
186
+ local_rank,
187
+ gpu_bdf,
188
+ rdma_dev,
189
+ nic_fmt,
190
+ ", ".join(set_vars),
191
+ )
192
+
193
+
194
+def _dp_adjusted_local_rank(tp_local_rank: int, vllm_config: VllmConfig) -> int:
195
+ """Compute the node-wide GPU index accounting for data parallelism.
196
+
197
+ On CUDA-alike platforms without env-var device isolation (the common
198
+ MP-backend path), the worker sees *all* GPUs on the node and selects
199
+ its device via ``torch.accelerator.set_device_index()`` using::
200
+
201
+ dp_local_rank * tp_pp_world_size + tp_local_rank
202
+
203
+ This mirrors the adjustment in ``Worker.init_device()`` so we resolve
204
+ the correct GPU PCI address *before* the CUDA device is initialised.
205
+ """
206
+ pc = vllm_config.parallel_config
207
+ if (
208
+ pc.distributed_executor_backend not in ("ray", "external_launcher")
209
+ and pc.data_parallel_backend != "ray"
210
+ and pc.nnodes_within_dp == 1
211
+ ):
212
+ dp_local_rank = pc.data_parallel_rank_local
213
+ if dp_local_rank is None:
214
+ dp_local_rank = pc.data_parallel_index
215
+ tp_pp_world_size = pc.pipeline_parallel_size * pc.tensor_parallel_size
216
+ return dp_local_rank * tp_pp_world_size + tp_local_rank
217
+ return tp_local_rank
218
+
219
+
220
+def set_worker_net_device(local_rank: int, vllm_config: VllmConfig) -> None:
221
+ """Top-level entry point for both UniProcExecutor and MultiprocExecutor.
222
+
223
+ Sets NIC selection env vars from ``VLLM_GPU_NIC_PCIE_MAPPING`` and
224
+ ``VLLM_NIC_SELECTION_VARS`` if present; no-op otherwise.
225
+ """
226
+ has_pcie_mapping = bool(envs.VLLM_GPU_NIC_PCIE_MAPPING.strip())
227
+ has_selection_vars = bool(envs.VLLM_NIC_SELECTION_VARS.strip())
228
+ if has_pcie_mapping and not has_selection_vars:
229
+ raise RuntimeError(
230
+ "VLLM_GPU_NIC_PCIE_MAPPING is set but VLLM_NIC_SELECTION_VARS "
231
+ "is not; both must be set together."
232
+ )
233
+ if has_selection_vars and not has_pcie_mapping:
234
+ raise RuntimeError(
235
+ "VLLM_NIC_SELECTION_VARS is set but VLLM_GPU_NIC_PCIE_MAPPING "
236
+ "is not; both must be set together."
237
+ )
238
+ # No-op when neither env var is present.
239
+ if not has_pcie_mapping and not has_selection_vars:
240
+ return
241
+ # Both env vars are present, so set the NIC selection env vars.
242
+ adjusted_rank = _dp_adjusted_local_rank(local_rank, vllm_config)
243
+ set_worker_gpu_nic_mapping(adjusted_rank)