[PERF] Extend NCCL symmetric memory to AllGather and ReduceScatter (#46703)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: snordmann <snordmann@nvidia.com>
Woosuk Kwon committed
Jun 30, 2026 at 11:38 UTC
953bba488dee45b1136379a823ef8a93d6d4e19e
6 files changed
+405
-112
.buildkite/test_areas/distributed.yaml
+1
-1
@@ -233,7 +233,7 @@ steps:
233
num_devices: 2
234
commands:
235
- pytest -v -s tests/distributed/test_context_parallel.py
236
- - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
236
+ - pytest -v -s tests/distributed/test_nccl_symm_mem.py
237
- pytest -v -s tests/v1/distributed/test_dbo.py
238
- pytest -v -s tests/distributed/test_mnnvl_alltoall.py
239
tests/distributed/test_nccl_symm_mem.py
new
+223
@@ -0,0 +1,223 @@
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 tests.utils import ensure_current_vllm_config
14
+from vllm.distributed import cleanup_dist_env_and_memory
15
+from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
16
+from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
17
+from vllm.distributed.device_communicators.pynccl_allocator import (
18
+ get_nccl_mem_pool,
19
+ is_symmetric_memory_enabled,
20
+)
21
+from vllm.distributed.parallel_state import (
22
+ get_tp_group,
23
+ init_distributed_environment,
24
+ initialize_model_parallel,
25
+)
26
+from vllm.platforms import current_platform
27
+from vllm.utils.system_utils import update_environment_variables
28
+
29
+torch.manual_seed(42)
30
+random.seed(44)
31
+
32
+test_size_elements = 4 * 1024 * 1024
33
+
34
+
35
+def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
36
+ monkeypatch = pytest.MonkeyPatch()
37
+ with monkeypatch.context() as m:
38
+ m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
39
+ dtype = torch.bfloat16
40
+ device = torch.device(f"cuda:{local_rank}")
41
+ torch.accelerator.set_device_index(device)
42
+ torch.set_default_device(device)
43
+ torch.set_default_dtype(dtype)
44
+ update_environment_variables(
45
+ {
46
+ "RANK": str(local_rank),
47
+ "LOCAL_RANK": str(local_rank),
48
+ "WORLD_SIZE": str(world_size),
49
+ "MASTER_ADDR": "localhost",
50
+ "MASTER_PORT": "12345",
51
+ }
52
+ )
53
+
54
+ init_distributed_environment()
55
+ with ensure_current_vllm_config():
56
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
57
+
58
+ cuda_communicator = typing.cast(
59
+ CudaCommunicator, get_tp_group().device_communicator
60
+ )
61
+ pynccl_comm = cuda_communicator.pynccl_comm
62
+ if get_nccl_mem_pool() is None:
63
+ pytest.skip(
64
+ "NCCL allocator compilation failed (probably missing NCCL headers)."
65
+ )
66
+ if not is_symmetric_memory_enabled():
67
+ pytest.skip("NCCL symmetric memory allreduce is disabled.")
68
+
69
+ register_nccl_symmetric_ops(pynccl_comm)
70
+ input = torch.randint(1, 23, (test_size_elements,), dtype=dtype, device=device)
71
+ input_clone = input.clone()
72
+ output = torch.ops.vllm.all_reduce_symmetric_with_copy(input)
73
+ assert output is not None
74
+
75
+ group = get_tp_group().device_group
76
+ dist.all_reduce(input_clone, group=group)
77
+ torch.testing.assert_close(output, input_clone, atol=2.5, rtol=0.1)
78
+
79
+
80
+@pytest.mark.skipif(
81
+ not current_platform.is_cuda(),
82
+ reason="NCCLSymmMemAllreduce is only available for CUDA platforms.",
83
+)
84
+@pytest.mark.parametrize("world_size", [2])
85
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
86
+def test_nccl_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, world_size):
87
+ if world_size > torch.accelerator.device_count():
88
+ pytest.skip("Not enough GPUs to run the test.")
89
+
90
+ # Enable SymmMemCommunicator
91
+ monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
92
+ monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
93
+ monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
94
+
95
+ mp.spawn(nccl_symm_mem_allreduce_worker, args=(world_size,), nprocs=world_size)
96
+ cleanup_dist_env_and_memory()
97
+
98
+
99
+def nccl_symm_mem_allgather_worker(local_rank: int, world_size: int):
100
+ monkeypatch = pytest.MonkeyPatch()
101
+ with monkeypatch.context() as m:
102
+ m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
103
+ dtype = torch.bfloat16
104
+ device = torch.device(f"cuda:{local_rank}")
105
+ torch.accelerator.set_device_index(device)
106
+ torch.set_default_device(device)
107
+ torch.set_default_dtype(dtype)
108
+ update_environment_variables(
109
+ {
110
+ "RANK": str(local_rank),
111
+ "LOCAL_RANK": str(local_rank),
112
+ "WORLD_SIZE": str(world_size),
113
+ "MASTER_ADDR": "localhost",
114
+ "MASTER_PORT": "12346",
115
+ }
116
+ )
117
+
118
+ init_distributed_environment()
119
+ with ensure_current_vllm_config():
120
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
121
+
122
+ cuda_communicator = typing.cast(
123
+ CudaCommunicator, get_tp_group().device_communicator
124
+ )
125
+ if get_nccl_mem_pool() is None:
126
+ pytest.skip(
127
+ "NCCL allocator compilation failed (probably missing NCCL headers)."
128
+ )
129
+ if not is_symmetric_memory_enabled():
130
+ pytest.skip("NCCL symmetric memory is disabled.")
131
+
132
+ per_rank_size = test_size_elements // world_size
133
+ input_tensor = torch.randint(
134
+ 1, 23, (per_rank_size,), dtype=dtype, device=device
135
+ )
136
+ output = cuda_communicator.all_gatherv(input_tensor, dim=0)
137
+
138
+ group = get_tp_group().device_group
139
+ expected = torch.empty(test_size_elements, dtype=dtype, device=device)
140
+ dist.all_gather_into_tensor(expected, input_tensor, group=group)
141
+ torch.testing.assert_close(output, expected, atol=0.0, rtol=0.0)
142
+
143
+
144
+@pytest.mark.skipif(
145
+ not current_platform.is_cuda(),
146
+ reason="NCCL symmetric memory is only available for CUDA platforms.",
147
+)
148
+@pytest.mark.parametrize("world_size", [2])
149
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
150
+def test_nccl_symm_mem_allgather(monkeypatch: pytest.MonkeyPatch, world_size):
151
+ if world_size > torch.accelerator.device_count():
152
+ pytest.skip("Not enough GPUs to run the test.")
153
+
154
+ monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
155
+ monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
156
+ monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
157
+
158
+ mp.spawn(nccl_symm_mem_allgather_worker, args=(world_size,), nprocs=world_size)
159
+ cleanup_dist_env_and_memory()
160
+
161
+
162
+def nccl_symm_mem_reduce_scatter_worker(local_rank: int, world_size: int):
163
+ monkeypatch = pytest.MonkeyPatch()
164
+ with monkeypatch.context() as m:
165
+ m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
166
+ dtype = torch.bfloat16
167
+ device = torch.device(f"cuda:{local_rank}")
168
+ torch.accelerator.set_device_index(device)
169
+ torch.set_default_device(device)
170
+ torch.set_default_dtype(dtype)
171
+ update_environment_variables(
172
+ {
173
+ "RANK": str(local_rank),
174
+ "LOCAL_RANK": str(local_rank),
175
+ "WORLD_SIZE": str(world_size),
176
+ "MASTER_ADDR": "localhost",
177
+ "MASTER_PORT": "12347",
178
+ }
179
+ )
180
+
181
+ init_distributed_environment()
182
+ with ensure_current_vllm_config():
183
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
184
+
185
+ cuda_communicator = typing.cast(
186
+ CudaCommunicator, get_tp_group().device_communicator
187
+ )
188
+ if get_nccl_mem_pool() is None:
189
+ pytest.skip(
190
+ "NCCL allocator compilation failed (probably missing NCCL headers)."
191
+ )
192
+ if not is_symmetric_memory_enabled():
193
+ pytest.skip("NCCL symmetric memory is disabled.")
194
+
195
+ per_rank_size = test_size_elements // world_size
196
+ input_tensor = torch.randint(
197
+ 1, 23, (test_size_elements,), dtype=dtype, device=device
198
+ )
199
+ input_clone = input_tensor.clone()
200
+ output = cuda_communicator.reduce_scatter(input_tensor, dim=0)
201
+
202
+ group = get_tp_group().device_group
203
+ expected = torch.empty(per_rank_size, dtype=dtype, device=device)
204
+ dist.reduce_scatter_tensor(expected, input_clone, group=group)
205
+ torch.testing.assert_close(output, expected, atol=2.5, rtol=0.1)
206
+
207
+
208
+@pytest.mark.skipif(
209
+ not current_platform.is_cuda(),
210
+ reason="NCCL symmetric memory is only available for CUDA platforms.",
211
+)
212
+@pytest.mark.parametrize("world_size", [2])
213
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
214
+def test_nccl_symm_mem_reduce_scatter(monkeypatch: pytest.MonkeyPatch, world_size):
215
+ if world_size > torch.accelerator.device_count():
216
+ pytest.skip("Not enough GPUs to run the test.")
217
+
218
+ monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
219
+ monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
220
+ monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
221
+
222
+ mp.spawn(nccl_symm_mem_reduce_scatter_worker, args=(world_size,), nprocs=world_size)
223
+ cleanup_dist_env_and_memory()
tests/distributed/test_nccl_symm_mem_allreduce.py
deleted
-96
@@ -1,96 +0,0 @@
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 tests.utils import ensure_current_vllm_config
14
-from vllm.distributed import cleanup_dist_env_and_memory
15
-from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
16
-from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
17
-from vllm.distributed.device_communicators.pynccl_allocator import (
18
- get_nccl_mem_pool,
19
- is_symmetric_memory_enabled,
20
-)
21
-from vllm.distributed.parallel_state import (
22
- get_tp_group,
23
- init_distributed_environment,
24
- initialize_model_parallel,
25
-)
26
-from vllm.platforms import current_platform
27
-from vllm.utils.system_utils import update_environment_variables
28
-
29
-torch.manual_seed(42)
30
-random.seed(44)
31
-
32
-test_size_elements = 4 * 1024 * 1024
33
-
34
-
35
-def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
36
- monkeypatch = pytest.MonkeyPatch()
37
- with monkeypatch.context() as m:
38
- m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
39
- dtype = torch.bfloat16
40
- device = torch.device(f"cuda:{local_rank}")
41
- torch.accelerator.set_device_index(device)
42
- torch.set_default_device(device)
43
- torch.set_default_dtype(dtype)
44
- update_environment_variables(
45
- {
46
- "RANK": str(local_rank),
47
- "LOCAL_RANK": str(local_rank),
48
- "WORLD_SIZE": str(world_size),
49
- "MASTER_ADDR": "localhost",
50
- "MASTER_PORT": "12345",
51
- }
52
- )
53
-
54
- init_distributed_environment()
55
- with ensure_current_vllm_config():
56
- initialize_model_parallel(tensor_model_parallel_size=world_size)
57
-
58
- cuda_communicator = typing.cast(
59
- CudaCommunicator, get_tp_group().device_communicator
60
- )
61
- pynccl_comm = cuda_communicator.pynccl_comm
62
- if get_nccl_mem_pool() is None:
63
- pytest.skip(
64
- "NCCL allocator compilation failed (probably missing NCCL headers)."
65
- )
66
- if not is_symmetric_memory_enabled():
67
- pytest.skip("NCCL symmetric memory allreduce is disabled.")
68
-
69
- register_nccl_symmetric_ops(pynccl_comm)
70
- input = torch.randint(1, 23, (test_size_elements,), dtype=dtype, device=device)
71
- input_clone = input.clone()
72
- output = torch.ops.vllm.all_reduce_symmetric_with_copy(input)
73
- assert output is not None
74
-
75
- group = get_tp_group().device_group
76
- dist.all_reduce(input_clone, group=group)
77
- torch.testing.assert_close(output, input_clone, atol=2.5, rtol=0.1)
78
-
79
-
80
-@pytest.mark.skipif(
81
- not current_platform.is_cuda(),
82
- reason="NCCLSymmMemAllreduce is only available for CUDA platforms.",
83
-)
84
-@pytest.mark.parametrize("world_size", [2])
85
-@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
86
-def test_nccl_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, world_size):
87
- if world_size > torch.accelerator.device_count():
88
- pytest.skip("Not enough GPUs to run the test.")
89
-
90
- # Enable SymmMemCommunicator
91
- monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
92
- monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
93
- monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
94
-
95
- mp.spawn(nccl_symm_mem_allreduce_worker, args=(world_size,), nprocs=world_size)
96
- cleanup_dist_env_and_memory()
vllm/distributed/device_communicators/all_reduce_utils.py
+12
@@ -134,6 +134,18 @@ def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor)
134
return world_size > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"]
135
136
137
+def should_nccl_symm_mem_ag_rs() -> bool:
138
+ """Check whether NCCL symmetric memory should be used for
139
+ AllGather / ReduceScatter collectives."""
140
+ from vllm.distributed.device_communicators.pynccl_allocator import (
141
+ is_symmetric_memory_enabled,
142
+ )
143
+
144
+ if envs.VLLM_BATCH_INVARIANT:
145
+ return False
146
+ return is_symmetric_memory_enabled()
147
+
148
+
149
def producer(
150
batch_src: Sequence[int],
151
producer_queue,
vllm/distributed/device_communicators/cuda_communicator.py
+163
-12
@@ -8,6 +8,7 @@ from torch.distributed import ProcessGroup
8
import vllm.envs as envs
9
from vllm.distributed.device_communicators.all_reduce_utils import (
10
NCCL_SYMM_MEM_ALL_REDUCE_CONFIG,
11
+ should_nccl_symm_mem_ag_rs,
12
should_nccl_symm_mem_allreduce,
13
)
14
from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
@@ -310,6 +311,17 @@ class CudaCommunicator(DeviceCommunicatorBase):
311
torch.distributed.all_reduce(out, group=self.device_group)
312
return out
313
314
+ def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
315
+ # Route uniform dim-0 all-gathers through NVLS symmetric memory when
316
+ # enabled (mirrors reduce_scatter); otherwise fall back to the
317
+ # base-class ring all-gather. Sequence parallelism's gather-before-GEMM
318
+ # uses dim=0 with tp-aligned (uniform) shards.
319
+ if dim < 0:
320
+ dim += input_.dim()
321
+ if dim == 0 and should_nccl_symm_mem_ag_rs():
322
+ return self._all_gather_symm_mem(input_.contiguous())
323
+ return super().all_gather(input_, dim)
324
+
325
def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
326
world_size = self.world_size
327
pynccl_comm = self.pynccl_comm
@@ -326,11 +338,13 @@ class CudaCommunicator(DeviceCommunicatorBase):
338
chunk_size = input_tensor.shape[0] // world_size
339
output_shape = (chunk_size,) + input_tensor.shape[1:]
340
329
- output = torch.empty(
330
- output_shape, dtype=input_tensor.dtype, device=input_tensor.device
331
- )
332
-
333
- pynccl_comm.reduce_scatter(output, input_tensor)
341
+ if should_nccl_symm_mem_ag_rs():
342
+ output = self._reduce_scatter_symm_mem(input_tensor)
343
+ else:
344
+ output = torch.empty(
345
+ output_shape, dtype=input_tensor.dtype, device=input_tensor.device
346
+ )
347
+ pynccl_comm.reduce_scatter(output, input_tensor)
348
349
# Reshape before returning
350
return output.movedim(0, dim).contiguous()
@@ -358,18 +372,97 @@ class CudaCommunicator(DeviceCommunicatorBase):
372
chunk_size = input_tensor.shape[0] // world_size
373
output_shape = (chunk_size,) + input_tensor.shape[1:]
374
361
- output = torch.empty(
362
- output_shape, dtype=input_tensor.dtype, device=input_tensor.device
363
- )
364
-
365
- if sizes is not None and sizes.count(sizes[0]) != len(sizes):
366
- pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes)
375
+ # Symmetric memory is only used when all ranks have uniform sizes.
376
+ # ncclCommWindowRegister is collective: asymmetric pool allocations
377
+ # from variable per-rank sizes cause deadlocks.
378
+ use_symm_mem = sizes is None and should_nccl_symm_mem_ag_rs()
379
+ if use_symm_mem:
380
+ output = self._reduce_scatter_symm_mem(input_tensor)
381
else:
368
- pynccl_comm.reduce_scatter(output, input_tensor)
382
+ output = torch.empty(
383
+ output_shape, dtype=input_tensor.dtype, device=input_tensor.device
384
+ )
385
+ if sizes is not None and sizes.count(sizes[0]) != len(sizes):
386
+ pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes)
387
+ else:
388
+ pynccl_comm.reduce_scatter(output, input_tensor)
389
390
# Reshape before returning
391
return output.movedim(0, dim).contiguous()
392
393
+ def _get_symm_scratch(
394
+ self,
395
+ role: str,
396
+ shape: tuple[int, ...],
397
+ dtype: torch.dtype,
398
+ device: torch.device,
399
+ ) -> torch.Tensor:
400
+ """Persistent, pre-registered NCCL symmetric-memory scratch buffer.
401
+
402
+ Allocating a fresh symm tensor per collective pays the
403
+ ``nccl_symm_mem_context`` snapshot + window-registration scan on every
404
+ call (~0.5 ms/RS+AG pair, dwarfing the NVLS transfer itself). Instead we
405
+ allocate once per ``(role, shape, dtype)``, register once, and reuse.
406
+
407
+ Safe for serial (eager) sequence parallelism: each collective's result
408
+ is consumed on the same stream before the next same-role collective
409
+ reuses the buffer. Distinct roles (e.g. ``rs_in`` vs ``ag_out``, both
410
+ full-size) get distinct buffers so a reduce-scatter input copy never
411
+ clobbers a still-live all-gather output.
412
+ """
413
+ from vllm.distributed.device_communicators.pynccl_allocator import (
414
+ nccl_symm_mem_context,
415
+ )
416
+
417
+ pynccl_comm = self.pynccl_comm
418
+ assert pynccl_comm is not None
419
+ cache = self.__dict__.setdefault("_symm_scratch_bufs", {})
420
+ key = (role, tuple(shape), dtype)
421
+ buf = cache.get(key)
422
+ if buf is None:
423
+ with nccl_symm_mem_context(pynccl_comm):
424
+ buf = torch.empty(shape, dtype=dtype, device=device)
425
+ cache[key] = buf
426
+ return buf
427
+
428
+ def _reduce_scatter_symm_mem(
429
+ self,
430
+ input_tensor: torch.Tensor,
431
+ ) -> torch.Tensor:
432
+ """ReduceScatter using NCCL symmetric memory (NVLS).
433
+
434
+ Only called for uniform-size reduce_scatter (variable sizes are
435
+ guarded out by the caller to avoid asymmetric ncclCommWindowRegister).
436
+ Uses persistent pre-registered scratch (see _get_symm_scratch).
437
+ """
438
+ from vllm.distributed.device_communicators.pynccl_allocator import (
439
+ is_symmetric_memory_tensor,
440
+ )
441
+
442
+ pynccl_comm = self.pynccl_comm
443
+ assert pynccl_comm is not None
444
+
445
+ chunk = input_tensor.shape[0] // self.world_size
446
+ output_shape = (chunk,) + tuple(input_tensor.shape[1:])
447
+
448
+ symm_output = self._get_symm_scratch(
449
+ "rs_out", output_shape, input_tensor.dtype, input_tensor.device
450
+ )
451
+ # NVLS reduce-scatter (LDMC) requires the input in symmetric memory.
452
+ if is_symmetric_memory_tensor(input_tensor):
453
+ symm_input = input_tensor
454
+ else:
455
+ symm_input = self._get_symm_scratch(
456
+ "rs_in",
457
+ tuple(input_tensor.shape),
458
+ input_tensor.dtype,
459
+ input_tensor.device,
460
+ )
461
+ symm_input.copy_(input_tensor)
462
+
463
+ pynccl_comm.reduce_scatter(symm_output, symm_input)
464
+ return symm_output
465
+
466
def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
467
"""Sends a tensor to the destination rank in a blocking way"""
468
"""NOTE: `dst` is the local rank of the destination rank."""
@@ -440,6 +533,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
533
if sizes is not None and all(s == sizes[0] for s in sizes):
534
sizes = None
535
536
+ # Symmetric memory is only used when all ranks have uniform sizes.
537
+ # ncclCommWindowRegister is collective: asymmetric pool allocations
538
+ # from variable per-rank sizes cause deadlocks.
539
+ if sizes is None and should_nccl_symm_mem_ag_rs():
540
+ if isinstance(input_, torch.Tensor):
541
+ return self._all_gather_symm_mem(input_)
542
+ return self._all_gather_batched_symm_mem(input_)
543
+
544
def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None):
545
input_size = input_.size()
546
if sizes is not None:
@@ -471,6 +572,56 @@ class CudaCommunicator(DeviceCommunicatorBase):
572
573
return output_list
574
575
+ def _all_gather_symm_mem(self, input_: torch.Tensor) -> torch.Tensor:
576
+ """AllGather a single tensor using NCCL symmetric memory (NVLS).
577
+
578
+ Only the output needs to be in symmetric memory; NCCL does not
579
+ require the AG input to be symmetrically allocated.
580
+ """
581
+ pynccl_comm = self.pynccl_comm
582
+ assert pynccl_comm is not None
583
+
584
+ out_size = (input_.size(0) * self.world_size,) + tuple(input_.size()[1:])
585
+ # Persistent pre-registered scratch avoids the per-call symm-mem context
586
+ # snapshot/registration overhead (see _get_symm_scratch).
587
+ symm_output = self._get_symm_scratch(
588
+ "ag_out", out_size, input_.dtype, input_.device
589
+ )
590
+ pynccl_comm.all_gather(symm_output, input_)
591
+ return symm_output
592
+
593
+ def _all_gather_batched_symm_mem(
594
+ self, inputs: list[torch.Tensor]
595
+ ) -> list[torch.Tensor]:
596
+ """AllGather a list of tensors using NCCL symmetric memory (NVLS).
597
+
598
+ Uses group_start/group_end to batch the collectives.
599
+ Only the output needs to be in symmetric memory (see
600
+ _all_gather_symm_mem).
601
+ """
602
+ from vllm.distributed.device_communicators.pynccl_allocator import (
603
+ nccl_symm_mem_context,
604
+ )
605
+
606
+ pynccl_comm = self.pynccl_comm
607
+ assert pynccl_comm is not None
608
+ world_size = self.world_size
609
+
610
+ symm_outputs = []
611
+ with nccl_symm_mem_context(pynccl_comm):
612
+ for inp in inputs:
613
+ out_size = (inp.size(0) * world_size,) + inp.size()[1:]
614
+ symm_outputs.append(
615
+ torch.empty(out_size, dtype=inp.dtype, device=inp.device)
616
+ )
617
+
618
+ pynccl_comm.group_start()
619
+ for symm_out, inp in zip(symm_outputs, inputs):
620
+ pynccl_comm.all_gather(symm_out, inp)
621
+ pynccl_comm.group_end()
622
+
623
+ return symm_outputs
624
+
625
def dispatch_router_logits(
626
self,
627
hidden_states: torch.Tensor,
vllm/distributed/device_communicators/pynccl_allocator.py
+6
-3
@@ -39,7 +39,7 @@ void nccl_free_plug(void* ptr, size_t size, int device, void* stream) {
39
_allocator = None
40
_allocator_wrapper = None
41
_mem_pool = None
42
-_registered_base_addrs = set()
42
+_registered_base_addrs: dict[bytes, set] = {}
43
_graph_pool_id = None
44
_nccl_allocator_failed_to_compile = False
45
_cached_pool_snapshot = None
@@ -181,11 +181,14 @@ class nccl_symm_mem_context:
181
assert _pool is not None
182
_cached_pool_snapshot = _pool.snapshot()
183
assert self.pynccl_comm is not None
184
+ comm_key = bytes(self.pynccl_comm.unique_id.internal)
185
+ if comm_key not in _registered_base_addrs:
186
+ _registered_base_addrs[comm_key] = set()
187
for segment in _cached_pool_snapshot:
185
- if segment["address"] not in _registered_base_addrs:
188
+ if segment["address"] not in _registered_base_addrs[comm_key]:
189
self.pynccl_comm.register_comm_window_raw(
190
segment["address"], segment["total_size"]
191
)
189
- _registered_base_addrs.add(segment["address"])
192
+ _registered_base_addrs[comm_key].add(segment["address"])
193
if self.is_graph_capture:
194
torch._C._cuda_beginAllocateCurrentThreadToPool(self.device, _graph_pool_id)