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 14, 05:30 UTC
[core] add nccl symmetric memory for all reduce (#24532)
Signed-off-by: Amir Samani <asamani@nvidia.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Amir Samani committed
Sep 23, 2025 at 11:33 UTC
8c1c81a3de344fba340cb2efc82f29d2c1563a2d
12 files changed
+489
-6
.buildkite/test-pipeline.yaml
+1
@@ -1039,3 +1039,4 @@ steps:
1039
num_gpus: 2
1040
commands:
1041
- pytest -v -s tests/distributed/test_context_parallel.py
1042
+ - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
benchmarks/kernels/benchmark_device_communicators.py
+24
-2
@@ -7,6 +7,10 @@ Benchmark script for device communicators:
7
CustomAllreduce (oneshot, twoshot), PyNcclCommunicator,
8
and SymmMemCommunicator (multimem, two-shot).
9
10
+for NCCL symmetric memory you need to set the environment variables
11
+NCCL_NVLS_ENABLE=1 NCCL_CUMEM_ENABLE=1 VLLM_USE_NCCL_SYMM_MEM=1, otherwise NCCL does
12
+not use fast NVLS implementation for all reduce.
13
+
14
Usage:
15
torchrun --nproc_per_node=<N> benchmark_device_communicators.py [options]
16
@@ -26,7 +30,13 @@ import torch.distributed as dist
30
from torch.distributed import ProcessGroup
31
32
from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
29
-from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
33
+from vllm.distributed.device_communicators.pynccl import (
34
+ PyNcclCommunicator,
35
+ register_nccl_symmetric_ops,
36
+)
37
+from vllm.distributed.device_communicators.pynccl_allocator import (
38
+ set_graph_pool_id,
39
+)
40
from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator
41
from vllm.logger import init_logger
42
from vllm.utils import FlexibleArgumentParser
@@ -98,6 +108,7 @@ class CommunicatorBenchmark:
108
)
109
if not self.pynccl_comm.disabled:
110
logger.info("Rank %s: PyNcclCommunicator initialized", self.rank)
111
+ register_nccl_symmetric_ops(self.pynccl_comm)
112
else:
113
logger.info("Rank %s: PyNcclCommunicator disabled", self.rank)
114
self.pynccl_comm = None
@@ -194,6 +205,15 @@ class CommunicatorBenchmark:
205
None, # no env variable needed
206
)
207
)
208
+ communicators.append(
209
+ (
210
+ "pynccl-symm",
211
+ lambda t: torch.ops.vllm.all_reduce_symmetric_with_copy(t),
212
+ lambda t: True, # Always available if initialized
213
+ nullcontext(),
214
+ None, # no env variable needed
215
+ )
216
+ )
217
218
if self.symm_mem_comm_multimem is not None:
219
comm = self.symm_mem_comm_multimem
@@ -271,7 +291,9 @@ class CommunicatorBenchmark:
291
# Capture the graph using context manager
292
with context:
293
graph = torch.cuda.CUDAGraph()
274
- with torch.cuda.graph(graph):
294
+ graph_pool = torch.cuda.graph_pool_handle()
295
+ set_graph_pool_id(graph_pool)
296
+ with torch.cuda.graph(graph, pool=graph_pool):
297
for _ in range(CUDA_GRAPH_CAPTURE_CYCLES):
298
allreduce_fn(graph_input)
299
tests/distributed/test_nccl_symm_mem_allreduce.py
new
+94
@@ -0,0 +1,94 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import random
5
+import typing
6
+
7
+import pytest
8
+import torch
9
+import torch.distributed as dist
10
+import torch.multiprocessing as mp
11
+
12
+import vllm.envs as envs
13
+from vllm.distributed import cleanup_dist_env_and_memory
14
+from vllm.distributed.device_communicators.cuda_communicator import (
15
+ CudaCommunicator)
16
+from vllm.distributed.device_communicators.pynccl import (
17
+ register_nccl_symmetric_ops)
18
+from vllm.distributed.device_communicators.pynccl_allocator import (
19
+ get_nccl_mem_pool, is_symmetric_memory_enabled)
20
+from vllm.distributed.parallel_state import (get_tp_group,
21
+ init_distributed_environment,
22
+ initialize_model_parallel)
23
+from vllm.platforms import current_platform
24
+from vllm.utils import update_environment_variables
25
+
26
+torch.manual_seed(42)
27
+random.seed(44)
28
+
29
+test_size_elements = 4 * 1024 * 1024
30
+
31
+
32
+def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
33
+ monkeypatch = pytest.MonkeyPatch()
34
+ with monkeypatch.context() as m:
35
+ m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
36
+ dtype = torch.bfloat16
37
+ device = torch.device(f"cuda:{local_rank}")
38
+ torch.cuda.set_device(device)
39
+ torch.set_default_device(device)
40
+ torch.set_default_dtype(dtype)
41
+ update_environment_variables({
42
+ "RANK": str(local_rank),
43
+ "LOCAL_RANK": str(local_rank),
44
+ "WORLD_SIZE": str(world_size),
45
+ "MASTER_ADDR": "localhost",
46
+ "MASTER_PORT": "12345",
47
+ })
48
+
49
+ init_distributed_environment()
50
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
51
+
52
+ cuda_communicator = typing.cast(CudaCommunicator,
53
+ get_tp_group().device_communicator)
54
+ pynccl_comm = cuda_communicator.pynccl_comm
55
+ if get_nccl_mem_pool() is None:
56
+ pytest.skip("NCCL allocator compilation failed "
57
+ "(probably missing NCCL headers).")
58
+ if not is_symmetric_memory_enabled():
59
+ pytest.skip("NCCL symmetric memory allreduce is disabled.")
60
+
61
+ register_nccl_symmetric_ops(pynccl_comm)
62
+ input = torch.randint(1,
63
+ 23, (test_size_elements, ),
64
+ dtype=dtype,
65
+ device=device)
66
+ input_clone = input.clone()
67
+ output = torch.ops.vllm.all_reduce_symmetric_with_copy(input)
68
+ assert output is not None
69
+
70
+ group = get_tp_group().device_group
71
+ dist.all_reduce(input_clone, group=group)
72
+ torch.testing.assert_close(output, input_clone, atol=2.5, rtol=0.1)
73
+
74
+
75
+@pytest.mark.skipif(
76
+ not current_platform.is_cuda(),
77
+ reason="NCCLSymmMemAllreduce is only available for CUDA platforms.",
78
+)
79
+@pytest.mark.parametrize("world_size", [2])
80
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"],
81
+ reason="Only test on CUDA")
82
+def test_nccl_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, world_size):
83
+ if world_size > torch.cuda.device_count():
84
+ pytest.skip("Not enough GPUs to run the test.")
85
+
86
+ # Enable SymmMemCommunicator
87
+ monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
88
+ monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
89
+ monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
90
+
91
+ mp.spawn(nccl_symm_mem_allreduce_worker,
92
+ args=(world_size, ),
93
+ nprocs=world_size)
94
+ cleanup_dist_env_and_memory()
vllm/compilation/cuda_graph.py
+6
@@ -12,6 +12,8 @@ import vllm.envs as envs
12
from vllm.compilation.counter import compilation_counter
13
from vllm.compilation.monitor import validate_cudagraph_capturing_enabled
14
from vllm.config import CUDAGraphMode, VllmConfig
15
+from vllm.distributed.device_communicators.pynccl_allocator import (
16
+ set_graph_pool_id)
17
from vllm.forward_context import BatchDescriptor, get_forward_context
18
from vllm.logger import init_logger
19
from vllm.platforms import current_platform
@@ -154,6 +156,10 @@ class CUDAGraphWrapper:
156
stack.enter_context(
157
patch("torch.cuda.empty_cache", lambda: None))
158
159
+ if self.graph_pool is not None:
160
+ set_graph_pool_id(self.graph_pool)
161
+ else:
162
+ set_graph_pool_id(current_platform.graph_pool_handle())
163
# mind-exploding: carefully manage the reference and memory.
164
with torch.cuda.graph(cudagraph, pool=self.graph_pool):
165
# `output` is managed by pytorch's cudagraph pool
vllm/distributed/device_communicators/all_reduce_utils.py
+26
-1
@@ -10,8 +10,9 @@ import sys
10
import tempfile
11
from collections.abc import Sequence
12
from itertools import product
13
-from typing import Optional
13
+from typing import Any, Optional
14
15
+import torch
16
import torch.distributed as dist
17
import torch.multiprocessing as mp
18
@@ -56,6 +57,30 @@ SYMM_MEM_ALL_REDUCE_MAX_SIZES = {
57
}
58
}
59
60
+NCCL_SYMM_MEM_ALL_REDUCE_CONFIG: dict[str, Any] = {
61
+ "min_world_size": 4,
62
+ "thresholds": {
63
+ 4: 2 * MiB, # 2 MB
64
+ 8: 1 * MiB, # 1 MB
65
+ },
66
+ "always_use_above_world_size": 8 # Always use symm mem for world_size > 8
67
+}
68
+
69
+
70
+def should_nccl_symm_mem_allreduce(world_size: int,
71
+ input_tensor: torch.Tensor) -> bool:
72
+ from vllm.distributed.device_communicators.pynccl_allocator import (
73
+ is_symmetric_memory_enabled)
74
+ if not is_symmetric_memory_enabled():
75
+ return False
76
+ if world_size < NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["min_world_size"]:
77
+ return False
78
+ threshold = NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["thresholds"].get(world_size)
79
+ if threshold is not None and input_tensor.nbytes >= threshold:
80
+ return True
81
+ return (world_size
82
+ > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"])
83
+
84
85
def producer(batch_src: Sequence[int],
86
producer_queue,
vllm/distributed/device_communicators/cuda_communicator.py
+15
@@ -7,6 +7,12 @@ import torch
7
from torch.distributed import ProcessGroup
8
9
import vllm.envs as envs
10
+from vllm.distributed.device_communicators.all_reduce_utils import (
11
+ should_nccl_symm_mem_allreduce)
12
+from vllm.distributed.device_communicators.pynccl import (
13
+ register_nccl_symmetric_ops)
14
+from vllm.distributed.device_communicators.pynccl_allocator import (
15
+ is_symmetric_memory_enabled)
16
from vllm.logger import init_logger
17
from vllm.platforms import current_platform
18
@@ -53,6 +59,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
59
group=self.cpu_group,
60
device=self.device,
61
)
62
+ if is_symmetric_memory_enabled():
63
+ register_nccl_symmetric_ops(self.pynccl_comm)
64
65
self.ca_comm: Optional[CustomAllreduce] = None
66
self.qr_comm: Optional[QuickAllReduce] = None
@@ -107,6 +115,13 @@ class CudaCommunicator(DeviceCommunicatorBase):
115
raise ValueError(f"Unknown all2all backend: {all2all_backend}")
116
117
def all_reduce(self, input_):
118
+ # since currently we perform copy input -> symm_input -> out-of-place AR
119
+ # return symm_output, we don't need to check if input is symmetric
120
+ if self.pynccl_comm is not None and \
121
+ should_nccl_symm_mem_allreduce(self.pynccl_comm.world_size,input_):
122
+ out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_)
123
+ if out is not None:
124
+ return out
125
# always try quick reduce first, then custom allreduce,
126
# and then pynccl. (quick reduce just for ROCM MI3*)
127
qr_comm = self.qr_comm
vllm/distributed/device_communicators/pynccl.py
+52
-1
@@ -17,6 +17,39 @@ from vllm.utils import current_stream
17
18
logger = init_logger(__name__)
19
20
+_NCCL_SYMM_OPS_REGISTERED = False
21
+
22
+
23
+def register_nccl_symmetric_ops(pynccl_comm):
24
+ from vllm.distributed.device_communicators.pynccl_allocator import (
25
+ nccl_symm_mem_context)
26
+ from vllm.utils import direct_register_custom_op
27
+
28
+ global _NCCL_SYMM_OPS_REGISTERED
29
+ if _NCCL_SYMM_OPS_REGISTERED:
30
+ return
31
+ _NCCL_SYMM_OPS_REGISTERED = True
32
+
33
+ def all_reduce_symmetric_with_copy_impl(
34
+ input_tensor: torch.Tensor) -> torch.Tensor:
35
+ with nccl_symm_mem_context(pynccl_comm):
36
+ symm_input = torch.empty_like(input_tensor)
37
+ symm_output = torch.empty_like(input_tensor)
38
+ symm_input.copy_(input_tensor)
39
+ symm_output = pynccl_comm.all_reduce(symm_input, symm_output)
40
+ return symm_output
41
+
42
+ def all_reduce_symmetric_with_copy_fake(
43
+ input_tensor: torch.Tensor) -> torch.Tensor:
44
+ return torch.empty_like(input_tensor)
45
+
46
+ direct_register_custom_op(
47
+ op_name="all_reduce_symmetric_with_copy",
48
+ op_func=all_reduce_symmetric_with_copy_impl,
49
+ mutates_args=[],
50
+ fake_impl=all_reduce_symmetric_with_copy_fake,
51
+ )
52
+
53
54
class PyNcclCommunicator:
55
@@ -67,6 +100,7 @@ class PyNcclCommunicator:
100
self.available = True
101
self.disabled = False
102
103
+ self.nccl_version = self.nccl.ncclGetRawVersion()
104
logger.info("vLLM is using nccl==%s", self.nccl.ncclGetVersion())
105
106
if self.rank == 0:
@@ -109,6 +143,7 @@ class PyNcclCommunicator:
143
144
def all_reduce(self,
145
in_tensor: torch.Tensor,
146
+ out_tensor: torch.Tensor = None,
147
op: ReduceOp = ReduceOp.SUM,
148
stream=None) -> torch.Tensor:
149
if self.disabled:
@@ -120,7 +155,8 @@ class PyNcclCommunicator:
155
f"this nccl communicator is created to work on {self.device}, "
156
f"but the input tensor is on {in_tensor.device}")
157
123
- out_tensor = torch.empty_like(in_tensor)
158
+ if out_tensor is None:
159
+ out_tensor = torch.empty_like(in_tensor)
160
161
if stream is None:
162
stream = current_stream()
@@ -288,3 +324,18 @@ class PyNcclCommunicator:
324
325
def group_end(self):
326
self.nccl.ncclGroupEnd()
327
+
328
+ def register_comm_window(self, tensor: torch.Tensor):
329
+ return self.nccl.ncclCommWindowRegister(
330
+ self.comm,
331
+ buffer_type(tensor.data_ptr()),
332
+ tensor.numel() * tensor.element_size(),
333
+ 1,
334
+ )
335
+
336
+ def register_comm_window_raw(self, ptr: int, size: int):
337
+ return self.nccl.ncclCommWindowRegister(self.comm, buffer_type(ptr),
338
+ size, 1)
339
+
340
+ def deregister_comm_window(self, window):
341
+ return self.nccl.ncclCommWindowDeregister(self.comm, window)
vllm/distributed/device_communicators/pynccl_allocator.py
new
+186
@@ -0,0 +1,186 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import atexit
4
+import contextlib
5
+import tempfile
6
+from typing import Any, Optional
7
+
8
+import torch
9
+from packaging import version
10
+from torch.cuda.memory import CUDAPluggableAllocator
11
+from torch.utils.cpp_extension import load_inline
12
+
13
+from vllm import envs
14
+from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
15
+from vllm.logger import init_logger
16
+from vllm.platforms import current_platform
17
+from vllm.utils import find_nccl_include_paths
18
+
19
+logger = init_logger(__name__)
20
+
21
+nccl_allocator_source = """
22
+#include <nccl.h>
23
+extern "C" {
24
+
25
+void* nccl_alloc_plug(size_t size, int device, void* stream) {
26
+ void* ptr;
27
+ ncclResult_t err = ncclMemAlloc(&ptr, size);
28
+ return ptr;
29
+
30
+}
31
+
32
+void nccl_free_plug(void* ptr, size_t size, int device, void* stream) {
33
+ ncclResult_t err = ncclMemFree(ptr);
34
+}
35
+
36
+}
37
+"""
38
+
39
+_allocator = None
40
+_allocator_wrapper = None
41
+_mem_pool = None
42
+_registered_base_addrs = set()
43
+_graph_pool_id = None
44
+_nccl_allocator_failed_to_compile = False
45
+_cached_pool_snapshot = None
46
+
47
+
48
+def is_symmetric_memory_enabled():
49
+ global _nccl_allocator_failed_to_compile
50
+ return envs.VLLM_USE_NCCL_SYMM_MEM and not _nccl_allocator_failed_to_compile
51
+
52
+
53
+def is_symmetric_memory_tensor(tensor: torch.Tensor):
54
+ if not is_symmetric_memory_enabled() or _cached_pool_snapshot is None:
55
+ return False
56
+ for segment in _cached_pool_snapshot:
57
+ for block in segment["blocks"]:
58
+ if block["address"] == tensor.untyped_storage().data_ptr():
59
+ return True
60
+ return False
61
+
62
+
63
+def set_graph_pool_id(graph_pool_id):
64
+ global _graph_pool_id
65
+ _graph_pool_id = graph_pool_id
66
+
67
+
68
+def compile_nccl_allocator():
69
+ global _allocator, _allocator_wrapper, _nccl_allocator_failed_to_compile
70
+ if not current_platform.is_cuda():
71
+ _nccl_allocator_failed_to_compile = True
72
+ return
73
+ try:
74
+ out_dir = tempfile.gettempdir()
75
+ nccl_allocator_libname = "nccl_allocator"
76
+ nccl_include_paths = find_nccl_include_paths()
77
+ load_inline(
78
+ name=nccl_allocator_libname,
79
+ cpp_sources=nccl_allocator_source,
80
+ with_cuda=True,
81
+ extra_ldflags=["-lnccl"],
82
+ verbose=envs.VLLM_LOGGING_LEVEL == "DEBUG",
83
+ is_python_module=False,
84
+ build_directory=out_dir,
85
+ extra_include_paths=nccl_include_paths,
86
+ )
87
+ _allocator_wrapper = CUDAPluggableAllocator(
88
+ f"{out_dir}/{nccl_allocator_libname}.so",
89
+ "nccl_alloc_plug",
90
+ "nccl_free_plug",
91
+ )
92
+ _allocator = _allocator_wrapper.allocator()
93
+ except Exception as e:
94
+ _nccl_allocator_failed_to_compile = True
95
+ logger.warning(
96
+ "Failed to compile NCCL memory allocator. "
97
+ "Symmetric memory will be disabled. "
98
+ "This is expected if NCCL headers are not available. "
99
+ "optionally set VLLM_NCCL_INCLUDE_PATH to point to a directory "
100
+ "containing the NCCL header. "
101
+ "Error: %s", str(e))
102
+
103
+
104
+def get_nccl_mem_pool():
105
+ global _mem_pool, _nccl_allocator_failed_to_compile
106
+ if _mem_pool is None and not _nccl_allocator_failed_to_compile:
107
+ compile_nccl_allocator()
108
+ if _allocator is not None:
109
+ _mem_pool = torch.cuda.MemPool(_allocator)
110
+ return _mem_pool
111
+
112
+
113
+def _cleanup_nccl_mem_pool():
114
+ global _mem_pool
115
+ _mem_pool = None
116
+
117
+
118
+def _cleanup_nccl_allocator_wrapper():
119
+ global _allocator_wrapper
120
+ _allocator_wrapper = None
121
+
122
+
123
+atexit.register(_cleanup_nccl_mem_pool)
124
+atexit.register(_cleanup_nccl_allocator_wrapper)
125
+
126
+
127
+class nccl_symm_mem_context:
128
+
129
+ def __init__(
130
+ self,
131
+ pynccl_comm: PyNcclCommunicator,
132
+ disabled: bool = False,
133
+ ):
134
+ self.disabled = (disabled or not is_symmetric_memory_enabled()
135
+ or pynccl_comm.world_size == 1
136
+ or not current_platform.is_cuda()
137
+ or get_nccl_mem_pool() is None or version.parse(
138
+ torch.__version__) < version.parse("2.8.0.a0"))
139
+ if self.disabled:
140
+ self.pynccl_comm: Optional[PyNcclCommunicator] = None
141
+ self._mem_pool_ctx: contextlib.AbstractContextManager[
142
+ Any] = contextlib.nullcontext()
143
+ self.is_graph_capture = None
144
+ self.device = None
145
+ else:
146
+ self.pynccl_comm = pynccl_comm
147
+ self._mem_pool_ctx = torch.cuda.use_mem_pool(get_nccl_mem_pool())
148
+ self.is_graph_capture = torch.cuda.is_current_stream_capturing()
149
+ self.device = torch.cuda.current_device()
150
+
151
+ def __enter__(self):
152
+ if self.disabled:
153
+ return self
154
+ assert (
155
+ self.pynccl_comm
156
+ is not None), "Symmetric memory requires pynccl to be initalized"
157
+ assert (
158
+ self.pynccl_comm.nccl_version >= 22703
159
+ ), "NCCL version 2.27.3 or higher is required for NCCL symmetric memory"
160
+ if self.is_graph_capture:
161
+ assert (
162
+ _graph_pool_id
163
+ is not None), "graph_pool_id is not set under graph capture"
164
+ # Pause graph memory pool to use symmetric memory with cuda graph
165
+ torch._C._cuda_endAllocateToPool(self.device, _graph_pool_id)
166
+ self._mem_pool_ctx.__enter__()
167
+ return self
168
+
169
+ def __exit__(self, exc_type, exc_val, exc_tb):
170
+ if self.disabled:
171
+ return
172
+ global _cached_pool_snapshot
173
+ global _registered_base_addrs
174
+ self._mem_pool_ctx.__exit__(exc_type, exc_val, exc_tb)
175
+ _pool = get_nccl_mem_pool()
176
+ assert _pool is not None
177
+ _cached_pool_snapshot = _pool.snapshot()
178
+ assert self.pynccl_comm is not None
179
+ for segment in _cached_pool_snapshot:
180
+ if segment["address"] not in _registered_base_addrs:
181
+ self.pynccl_comm.register_comm_window_raw(
182
+ segment["address"], segment["total_size"])
183
+ _registered_base_addrs.add(segment["address"])
184
+ if self.is_graph_capture:
185
+ torch._C._cuda_beginAllocateCurrentThreadToPool(
186
+ self.device, _graph_pool_id)
vllm/distributed/device_communicators/pynccl_wrapper.py
+36
-2
@@ -41,6 +41,7 @@ logger = init_logger(__name__)
41
42
ncclResult_t = ctypes.c_int
43
ncclComm_t = ctypes.c_void_p
44
+ncclWindow_t = ctypes.c_void_p
45
46
47
class ncclUniqueId(ctypes.Structure):
@@ -222,6 +223,24 @@ class NCCLLibrary:
223
Function("ncclGroupStart", ncclResult_t, []),
224
# ncclResult_t ncclGroupEnd();
225
Function("ncclGroupEnd", ncclResult_t, []),
226
+ # ncclResult_t ncclCommWindowRegister(
227
+ # ncclComm_t comm, void* buff, size_t size,
228
+ # ncclWindow_t* win, int winFlags);
229
+ Function(
230
+ "ncclCommWindowRegister",
231
+ ncclResult_t,
232
+ [
233
+ ncclComm_t,
234
+ buffer_type,
235
+ ctypes.c_size_t,
236
+ ctypes.POINTER(ncclWindow_t),
237
+ ctypes.c_int,
238
+ ],
239
+ ),
240
+ # ncclResult_t ncclCommWindowDeregister(
241
+ # ncclComm_t comm, ncclWindow_t win);
242
+ Function("ncclCommWindowDeregister", ncclResult_t,
243
+ [ncclComm_t, ncclWindow_t]),
244
]
245
246
# class attribute to store the mapping from the path to the library
@@ -271,10 +290,14 @@ class NCCLLibrary:
290
error_str = self.ncclGetErrorString(result)
291
raise RuntimeError(f"NCCL error: {error_str}")
292
274
- def ncclGetVersion(self) -> str:
293
+ def ncclGetRawVersion(self) -> int:
294
version = ctypes.c_int()
295
self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version)))
277
- version_str = str(version.value)
296
+ # something like 21903
297
+ return version.value
298
+
299
+ def ncclGetVersion(self) -> str:
300
+ version_str = str(self.ncclGetRawVersion())
301
# something like 21903 --> "2.19.3"
302
major = version_str[0].lstrip("0")
303
minor = version_str[1:3].lstrip("0")
@@ -375,6 +398,17 @@ class NCCLLibrary:
398
def ncclGroupEnd(self) -> None:
399
self.NCCL_CHECK(self._funcs["ncclGroupEnd"]())
400
401
+ def ncclCommWindowRegister(self, comm: ncclComm_t, buff: buffer_type,
402
+ size: int, win_flags: int) -> ncclWindow_t:
403
+ window = ncclWindow_t()
404
+ self.NCCL_CHECK(self._funcs["ncclCommWindowRegister"](
405
+ comm, buff, size, ctypes.byref(window), win_flags))
406
+ return window
407
+
408
+ def ncclCommWindowDeregister(self, comm: ncclComm_t,
409
+ window: ncclWindow_t) -> None:
410
+ self.NCCL_CHECK(self._funcs["ncclCommWindowDeregister"](comm, window))
411
+
412
413
__all__ = [
414
"NCCLLibrary", "ncclDataTypeEnum", "ncclRedOpTypeEnum", "ncclUniqueId",
vllm/envs.py
+11
@@ -193,6 +193,8 @@ if TYPE_CHECKING:
193
VLLM_DBO_COMM_SMS: int = 20
194
GPT_OSS_SYSTEM_TOOL_MCP_LABELS: list[str] = []
195
VLLM_PATTERN_MATCH_DEBUG: Optional[str] = None
196
+ VLLM_USE_NCCL_SYMM_MEM: bool = False
197
+ VLLM_NCCL_INCLUDE_PATH: Optional[str] = None
198
199
200
def get_default_cache_root():
@@ -1410,6 +1412,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
1412
["container",
1413
"code_interpreter",
1414
"web_search_preview"]),
1415
+
1416
+ # Flag to enable NCCL symmetric memory allocation and registration
1417
+ "VLLM_USE_NCCL_SYMM_MEM":
1418
+ lambda: bool(int(os.getenv("VLLM_USE_NCCL_SYMM_MEM", "0"))),
1419
+
1420
+ # NCCL header path
1421
+ "VLLM_NCCL_INCLUDE_PATH":
1422
+ lambda: os.environ.get("VLLM_NCCL_INCLUDE_PATH", None),
1423
+
1424
}
1425
1426
# --8<-- [end:env-vars-definition]
vllm/utils/__init__.py
+32
@@ -1383,6 +1383,38 @@ def find_nccl_library() -> str:
1383
return so_file
1384
1385
1386
+def find_nccl_include_paths() -> Optional[list[str]]:
1387
+ """
1388
+ We either use the nccl.h specified by the `VLLM_NCCL_INCLUDE_PATH`
1389
+ environment variable, or we find the library file brought by
1390
+ nvidia-nccl-cuXX. load_inline by default uses
1391
+ torch.utils.cpp_extension.include_paths
1392
+ """
1393
+ paths: list[str] = []
1394
+ inc = envs.VLLM_NCCL_INCLUDE_PATH
1395
+ if inc and os.path.isdir(inc):
1396
+ paths.append(inc)
1397
+
1398
+ try:
1399
+ import importlib.util
1400
+ spec = importlib.util.find_spec("nvidia.nccl")
1401
+ if spec and getattr(spec, "submodule_search_locations", None):
1402
+ for loc in spec.submodule_search_locations:
1403
+ inc_dir = os.path.join(loc, "include")
1404
+ if os.path.exists(os.path.join(inc_dir, "nccl.h")):
1405
+ paths.append(inc_dir)
1406
+ except Exception:
1407
+ pass
1408
+
1409
+ seen = set()
1410
+ out: list[str] = []
1411
+ for p in paths:
1412
+ if p and p not in seen:
1413
+ out.append(p)
1414
+ seen.add(p)
1415
+ return out or None
1416
+
1417
+
1418
prev_set_stream = torch.cuda.set_stream
1419
1420
_current_stream_tls = threading.local()
vllm/v1/worker/gpu_ubatch_wrapper.py
+6
@@ -11,6 +11,8 @@ import vllm.envs as envs
11
from vllm.compilation.cuda_graph import CUDAGraphWrapper
12
from vllm.config import CUDAGraphMode, VllmConfig
13
from vllm.distributed import get_ep_group
14
+from vllm.distributed.device_communicators.pynccl_allocator import (
15
+ set_graph_pool_id)
16
from vllm.forward_context import (create_forward_context, get_forward_context,
17
override_forward_context)
18
from vllm.logger import init_logger
@@ -206,6 +208,10 @@ class UBatchWrapper:
208
cudagraph=torch.cuda.CUDAGraph(),
209
ubatch_metadata=ubatch_metadata,
210
)
211
+ if self.graph_pool is not None:
212
+ set_graph_pool_id(self.graph_pool)
213
+ else:
214
+ set_graph_pool_id(current_platform.graph_pool_handle())
215
with torch.cuda.graph(cudagraph_metadata.cudagraph,
216
stream=compute_stream,
217
pool=self.graph_pool):