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 3, 05:30 UTC
[MISC] Fix symm_mem cap-equal gate; log AR backend selection (#42993)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
May 20, 2026 at 19:50 UTC
9c78c99995b70726f9ea929ff2e535d6303383d6
2 files changed
+68
-1
vllm/distributed/device_communicators/cuda_communicator.py
+67
@@ -7,6 +7,7 @@ from torch.distributed import ProcessGroup
7
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_allreduce,
12
)
13
from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
@@ -114,6 +115,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
115
# currently be an MI300 series.
116
self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device)
117
118
+ if self.world_size > 1:
119
+ self._log_all_reduce_backend_selection()
120
+
121
if self.use_all2all:
122
if self.all2all_backend in ("naive", "allgather_reducescatter"):
123
from .all2all import AgRsAll2AllManager
@@ -171,6 +175,69 @@ class CudaCommunicator(DeviceCommunicatorBase):
175
scope="global",
176
)
177
178
+ def _log_all_reduce_backend_selection(self) -> None:
179
+ """Log the all-reduce backends that are active for this group.
180
+
181
+ The dispatch chain in ``all_reduce`` tries backends in this order and
182
+ falls through to the next one if the current backend rejects the
183
+ input (size/dtype gates) or is disabled. The list of "enabled"
184
+ backends below is the subset of potential backends that may be
185
+ chosen at dispatch time for this group; the actual per-call choice
186
+ depends on the input tensor.
187
+ """
188
+ all_potential_ar_backends = [
189
+ "NCCL_SYMM_MEM",
190
+ "QUICK_REDUCE",
191
+ "FLASHINFER",
192
+ "CUSTOM",
193
+ "SYMM_MEM",
194
+ "PYNCCL",
195
+ ]
196
+ enabled_ar_backends: list[str] = []
197
+ # Mirror the static preconditions of `should_nccl_symm_mem_allreduce`:
198
+ # VLLM_BATCH_INVARIANT off, NCCL symm mem enabled, world_size meets
199
+ # min_world_size, and world_size either has a tuned entry in
200
+ # `custom_ar_preferred_ranges` or is greater than
201
+ # `always_use_above_world_size`. World sizes that fail the latter (e.g.
202
+ # 5/6/7 with the default config) never dispatch NCCL symm mem
203
+ # regardless of input. The per-tensor-size check inside the function
204
+ # stays as a runtime decision.
205
+ nccl_symm_ws_ok = self.world_size >= NCCL_SYMM_MEM_ALL_REDUCE_CONFIG[
206
+ "min_world_size"
207
+ ] and (
208
+ self.world_size
209
+ in NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["custom_ar_preferred_ranges"]
210
+ or self.world_size
211
+ > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"]
212
+ )
213
+ if (
214
+ self.pynccl_comm is not None
215
+ and not self.pynccl_comm.disabled
216
+ and is_symmetric_memory_enabled()
217
+ and not envs.VLLM_BATCH_INVARIANT
218
+ and nccl_symm_ws_ok
219
+ ):
220
+ enabled_ar_backends.append("NCCL_SYMM_MEM")
221
+ if self.qr_comm is not None and not self.qr_comm.disabled:
222
+ enabled_ar_backends.append("QUICK_REDUCE")
223
+ if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled:
224
+ enabled_ar_backends.append("FLASHINFER")
225
+ if self.ca_comm is not None and not self.ca_comm.disabled:
226
+ enabled_ar_backends.append("CUSTOM")
227
+ if self.symm_mem_comm is not None and not self.symm_mem_comm.disabled:
228
+ enabled_ar_backends.append("SYMM_MEM")
229
+ if self.pynccl_comm is not None and not self.pynccl_comm.disabled:
230
+ enabled_ar_backends.append("PYNCCL")
231
+
232
+ logger.info_once(
233
+ "Using %s all-reduce backends (in dispatch order) for group "
234
+ "'%s' out of potential backends: %s.",
235
+ "[" + ", ".join(f"'{b}'" for b in enabled_ar_backends) + "]",
236
+ self.unique_name or "<unnamed>",
237
+ "[" + ", ".join(f"'{b}'" for b in all_potential_ar_backends) + "]",
238
+ scope="global",
239
+ )
240
+
241
def all_reduce(self, input_):
242
# since currently we perform copy input -> symm_input -> out-of-place AR
243
# return symm_output, we don't need to check if input is symmetric
vllm/distributed/device_communicators/symm_mem.py
+1
-1
@@ -121,7 +121,7 @@ class SymmMemCommunicator:
121
inp_size = inp.numel() * inp.element_size()
122
if inp_size % 4 != 0:
123
return False
124
- return inp_size < self.max_size
124
+ return inp_size <= self.max_size
125
126
def all_reduce(
127
self, inp: torch.Tensor, *, out: torch.Tensor | None = None