# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os import torch from torch.distributed import ProcessGroup import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.device_communicators.all_reduce_utils import ( NCCL_SYMM_MEM_ALL_REDUCE_CONFIG, should_nccl_symm_mem_ag_rs, should_nccl_symm_mem_allreduce, ) from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops from vllm.distributed.device_communicators.pynccl_allocator import ( is_symmetric_memory_enabled, ) from vllm.logger import init_logger from vllm.platforms import current_platform from ..utils import StatelessProcessGroup from .aiter_custom_all_reduce import AiterCustomAllreduce from .base_device_communicator import DeviceCommunicatorBase logger = init_logger(__name__) class CudaCommunicator(DeviceCommunicatorBase): def __init__( self, cpu_group: ProcessGroup, device: torch.device | None = None, device_group: ProcessGroup | None = None, unique_name: str = "", global_ranks: list[int] | None = None, global_world_size: int | None = None, tcp_store_group: StatelessProcessGroup | None = None, use_all2all: bool = False, ): super().__init__( cpu_group, device, device_group, unique_name, global_ranks, global_world_size, use_all2all=use_all2all, ) # Match the group name exactly so ETP does not enable TP-only backends. if unique_name.split(":")[0] != "tp": # custom allreduce or torch symm mem can be used only by tp use_custom_allreduce = False use_torch_symm_mem = False use_flashinfer_allreduce = False use_flashinfer_pcie_ipc_allreduce = False use_aiter_allreduce = False else: from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM # FlashInfer all-reduce does not provide a fixed reduction order. use_flashinfer_allreduce = ( envs.VLLM_ALLREDUCE_USE_FLASHINFER and not envs.VLLM_BATCH_INVARIANT ) use_flashinfer_pcie_ipc_allreduce = ( envs.VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC and not envs.VLLM_BATCH_INVARIANT ) use_aiter_allreduce = use_custom_allreduce and bool( rocm_aiter_ops.is_custom_all_reduce_enabled() ) self.use_custom_allreduce = use_custom_allreduce self.use_torch_symm_mem = use_torch_symm_mem self.use_flashinfer_allreduce = use_flashinfer_allreduce self.use_flashinfer_pcie_ipc_allreduce = use_flashinfer_pcie_ipc_allreduce self.use_aiter_allreduce = use_aiter_allreduce # lazy import to avoid documentation build error from vllm.distributed.device_communicators.custom_all_reduce import ( CustomAllreduce, ) from vllm.distributed.device_communicators.flashinfer_all_reduce import ( FlashInferAllReduce, ) from vllm.distributed.device_communicators.flashinfer_pcie_ipc_all_reduce import ( # noqa: E501 FlashInferPcieIpcAllReduce, ) from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.device_communicators.quick_all_reduce import ( QuickAllReduce, ) from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator self.pynccl_comm: PyNcclCommunicator | None = None if self.world_size > 1: self.pynccl_comm = PyNcclCommunicator( group=self.cpu_group if tcp_store_group is None else tcp_store_group, device=self.device, ) if is_symmetric_memory_enabled(): register_nccl_symmetric_ops(self.pynccl_comm) # GB10: RDMA all-reduce for decode-sized tensors (arx.py). self.arx_comm = None if ( unique_name.split(":")[0] == "tp" and self.world_size > 1 and os.environ.get("VLLM_ARX_ALLREDUCE") == "1" ): from vllm.distributed.device_communicators.arx import ArxCommunicator self.arx_comm = ArxCommunicator(self.cpu_group, self.device) self.arxbig = None if self.arx_comm is not None and os.environ.get("VLLM_ARXBIG") == "1": from vllm.distributed.device_communicators.arx import ArxBig self.arxbig = ArxBig(self.cpu_group, self.device) self.ca_comm: CustomAllreduce | None = None self.qr_comm: QuickAllReduce | None = None self.symm_mem_comm: SymmMemCommunicator | None = None self.fi_ar_comm: FlashInferAllReduce | None = None self.fi_pcie_ipc_ar_comm: FlashInferPcieIpcAllReduce | None = None self.aiter_ar_comm: AiterCustomAllreduce | None = None self.use_aiter_ag_rs: bool = False if use_torch_symm_mem and current_platform.is_cuda(): self.symm_mem_comm = SymmMemCommunicator( group=self.cpu_group, device=self.device, ) if self.use_flashinfer_allreduce and self.world_size > 1: self.fi_ar_comm = FlashInferAllReduce( group=self.cpu_group, device=self.device, ) if ( self.use_flashinfer_pcie_ipc_allreduce and self.world_size > 1 and self.device_group is not None ): self.fi_pcie_ipc_ar_comm = FlashInferPcieIpcAllReduce( group=self.device_group, tune_group=self.cpu_group, device=self.device, ) if self.use_aiter_allreduce and self.world_size > 1: self.aiter_ar_comm = AiterCustomAllreduce( group=self.cpu_group, device=self.device, ) if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1: # Initialize a custom fast all-reduce implementation. self.ca_comm = CustomAllreduce( group=self.cpu_group, device=self.device, symm_mem_enabled=( self.symm_mem_comm is not None and not self.symm_mem_comm.disabled ), ) # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine if ( "dp" in unique_name and self.world_size in (2, 4, 8) and current_platform.is_rocm() and rocm_aiter_ops.is_custom_all_reduce_enabled() ): self.aiter_ar_comm = AiterCustomAllreduce( group=self.cpu_group, device=self.device, ) if self.aiter_ar_comm.disabled: self.aiter_ar_comm = None else: self.use_aiter_ag_rs = True if use_custom_allreduce and self.world_size > 1 and current_platform.is_rocm(): # Initialize a custom quick all-reduce implementation for AMD. # Quick reduce is designed as a complement to custom allreduce # (vLLM's or AITER's), so it is initialized for either backend. # Based on quickreduce (https://github.com/mk1-project/quickreduce). # On ROCm, 'use_custom_allreduce==True' means it must currently be # an MI300 series. self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device) if self.world_size > 1: self._log_all_reduce_backend_selection() if self.use_all2all: if self.all2all_backend in ("naive", "allgather_reducescatter"): from .all2all import AgRsAll2AllManager self.all2all_manager = AgRsAll2AllManager( self.cpu_group, tcp_store_group ) elif self.all2all_backend == "deepep_high_throughput": from .all2all import DeepEPHTAll2AllManager self.all2all_manager = DeepEPHTAll2AllManager( self.cpu_group, tcp_store_group ) elif self.all2all_backend == "deepep_low_latency": from .all2all import DeepEPLLAll2AllManager self.all2all_manager = DeepEPLLAll2AllManager( self.cpu_group, tcp_store_group ) elif self.all2all_backend in ( "mori_high_throughput", "mori_low_latency", ): from .all2all import MoriAll2AllManager self.all2all_manager = MoriAll2AllManager( self.cpu_group, self.all2all_backend ) elif self.all2all_backend == "deepep_v2": from .all2all import DeepEPV2All2AllManager self.all2all_manager = DeepEPV2All2AllManager( self.cpu_group, tcp_store_group, device_group=self.device_group, ) elif self.all2all_backend == "moonep": from .all2all import MoonEPAll2AllManager self.all2all_manager = MoonEPAll2AllManager( self.cpu_group, tcp_store_group, device_group=self.device_group, ) elif self.all2all_backend == "nixl_ep": from .all2all import NixlEPAll2AllManager self.all2all_manager = NixlEPAll2AllManager( self.cpu_group, tcp_store_group ) elif ( self.all2all_backend == "flashinfer_all2allv" or self.all2all_backend == "flashinfer_nvlink_two_sided" ): if self.all2all_backend == "flashinfer_all2allv": logger.warning_once( "'flashinfer_all2allv' is deprecated and has been renamed to" "'flashinfer_nvlink_two_sided'. It will be removed in a future" "release." ) from .all2all import FlashInferNVLinkTwoSidedManager self.all2all_manager = FlashInferNVLinkTwoSidedManager( self.cpu_group, tcp_store_group ) elif self.all2all_backend == "flashinfer_nvlink_one_sided": from .all2all import FlashInferNVLinkOneSidedManager self.all2all_manager = FlashInferNVLinkOneSidedManager(self.cpu_group) else: raise ValueError(f"Unknown all2all backend: {self.all2all_backend}") logger.info_once( "Using %s all2all manager.", self.all2all_manager.__class__.__name__, scope="global", ) def _log_all_reduce_backend_selection(self) -> None: """Log the all-reduce backends that are active for this group. The dispatch chain in ``all_reduce`` tries backends in this order and falls through to the next one if the current backend rejects the input (size/dtype gates) or is disabled. The list of "enabled" backends below is the subset of potential backends that may be chosen at dispatch time for this group; the actual per-call choice depends on the input tensor. """ all_potential_ar_backends = [ "FLASHINFER_PCIE_IPC", "FLASHINFER", "NCCL_SYMM_MEM", "QUICK_REDUCE", "AITER_CUSTOM", "CUSTOM", "SYMM_MEM", "PYNCCL", ] enabled_ar_backends: list[str] = [] if ( self.fi_pcie_ipc_ar_comm is not None and not self.fi_pcie_ipc_ar_comm.disabled ): enabled_ar_backends.append("FLASHINFER_PCIE_IPC") if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled: enabled_ar_backends.append("FLASHINFER") # Mirror the static preconditions of `should_nccl_symm_mem_allreduce`: # VLLM_BATCH_INVARIANT off, NCCL symm mem enabled, world_size meets # min_world_size, and world_size either has a tuned entry in # `custom_ar_preferred_ranges` or is greater than # `always_use_above_world_size`. World sizes that fail the latter (e.g. # 5/6/7 with the default config) never dispatch NCCL symm mem # regardless of input. The per-tensor-size check inside the function # stays as a runtime decision. nccl_symm_ws_ok = self.world_size >= NCCL_SYMM_MEM_ALL_REDUCE_CONFIG[ "min_world_size" ] and ( self.world_size in NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["custom_ar_preferred_ranges"] or self.world_size > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"] ) if ( self.pynccl_comm is not None and not self.pynccl_comm.disabled and is_symmetric_memory_enabled() and not envs.VLLM_BATCH_INVARIANT and nccl_symm_ws_ok ): enabled_ar_backends.append("NCCL_SYMM_MEM") if self.qr_comm is not None and not self.qr_comm.disabled: enabled_ar_backends.append("QUICK_REDUCE") if ( self.use_aiter_allreduce and self.aiter_ar_comm is not None and not self.aiter_ar_comm.disabled ): enabled_ar_backends.append("AITER_CUSTOM") if self.ca_comm is not None and not self.ca_comm.disabled: enabled_ar_backends.append("CUSTOM") if self.symm_mem_comm is not None and not self.symm_mem_comm.disabled: enabled_ar_backends.append("SYMM_MEM") if self.pynccl_comm is not None and not self.pynccl_comm.disabled: enabled_ar_backends.append("PYNCCL") logger.info_once( "Using %s all-reduce backends (in dispatch order) for group " "'%s' out of potential backends: %s.", "[" + ", ".join(f"'{b}'" for b in enabled_ar_backends) + "]", self.unique_name or "", "[" + ", ".join(f"'{b}'" for b in all_potential_ar_backends) + "]", scope="global", ) # GB10: sequence-parallel collectives over RDMA, tried first by # custom_all_gather / custom_reduce_scatter below: arx all-reduces for # decode sizes (~13-27 us against ~80 for NCCL), arxbig for prefill-sized # gathers. Returning None hands the call on. def _arx_reduce_scatter(self, input_): arx_comm = self.arx_comm if arx_comm is None or arx_comm.disabled or not arx_comm.should_use(input_): return None chunk = input_.shape[0] // self.world_size return arx_comm.all_reduce(input_)[self.rank_in_group * chunk:(self.rank_in_group + 1) * chunk] def _arx_all_gather(self, input_): # arxbig's gather lands in pinned memory, and the GEMMs that read it # there ran 16% slower overall than NCCL's gather into device memory. big = self.arxbig if big is not None and not big.disabled and big.gather and big.should_gather(input_): return big.all_gather(input_) # Each rank's rows in their own slot of a zero tensor: the sum is the # concatenation, exactly (x + 0 is x). arx_comm = self.arx_comm if (arx_comm is None or arx_comm.disabled or input_.dtype != torch.bfloat16 or input_.numel() * 2 * self.world_size > arx_comm.max_bytes): return None full = input_.new_zeros(input_.shape[0] * self.world_size, *input_.shape[1:]) if not arx_comm.should_use(full): return None chunk = input_.shape[0] full[self.rank_in_group * chunk:(self.rank_in_group + 1) * chunk] = input_ return arx_comm.all_reduce(full) def all_reduce(self, input_): arx_comm = self.arx_comm if arx_comm is not None and not arx_comm.disabled and arx_comm.should_use(input_, two_shot=True): return arx_comm.all_reduce(input_) fi_ar_comm = self.fi_ar_comm use_fi_ar = ( fi_ar_comm is not None and not fi_ar_comm.disabled and fi_ar_comm.should_use_fi_ar(input_) ) # since currently we perform copy input -> symm_input -> out-of-place AR # return symm_output, we don't need to check if input is symmetric if ( self.pynccl_comm is not None and not use_fi_ar and should_nccl_symm_mem_allreduce(self.pynccl_comm.world_size, input_) ): out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_) if out is not None: return out qr_comm = self.qr_comm if ( qr_comm is not None and not qr_comm.disabled and qr_comm.should_quick_allreduce(input_) ): out = qr_comm.quick_all_reduce(input_) assert out is not None return out fi_pcie_ipc_ar_comm = self.fi_pcie_ipc_ar_comm if fi_pcie_ipc_ar_comm is not None and fi_pcie_ipc_ar_comm.should_use(input_): return fi_pcie_ipc_ar_comm.all_reduce(input_) if use_fi_ar: assert fi_ar_comm is not None out = fi_ar_comm.all_reduce(input_) assert out is not None return out aiter_ar_comm = self.aiter_ar_comm if ( self.use_aiter_allreduce and aiter_ar_comm is not None and not aiter_ar_comm.disabled and aiter_ar_comm.should_custom_ar(input_) ): out = aiter_ar_comm.custom_all_reduce(input_) assert out is not None return out ca_comm = self.ca_comm if ( ca_comm is not None and not ca_comm.disabled and ca_comm.should_custom_ar(input_) ): out = ca_comm.custom_all_reduce(input_) assert out is not None return out symm_mem_comm = self.symm_mem_comm if symm_mem_comm is not None and symm_mem_comm.should_use_symm_mem(input_): out = symm_mem_comm.all_reduce(input_) assert out is not None return out pynccl_comm = self.pynccl_comm if pynccl_comm is None or pynccl_comm.disabled: out = input_.clone() torch.distributed.all_reduce(out, group=self.device_group) return out assert pynccl_comm is not None out = pynccl_comm.all_reduce(input_) if out is None: # fall back to the default all-reduce using PyTorch. # this usually happens during testing. # when we run the model, allreduce only happens for the TP # group, where we always have either custom allreduce or pynccl. out = input_.clone() torch.distributed.all_reduce(out, group=self.device_group) return out def custom_all_gather(self, input_: torch.Tensor) -> torch.Tensor | None: out = self._arx_all_gather(input_) if out is not None: return out ca_comm = self.ca_comm if ca_comm is None: return None return ca_comm.custom_all_gather(input_.contiguous()) def custom_reduce_scatter(self, input_: torch.Tensor) -> torch.Tensor | None: out = self._arx_reduce_scatter(input_) if out is not None: return out ca_comm = self.ca_comm if ca_comm is None: return None return ca_comm.custom_reduce_scatter(input_.contiguous()) def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: # Route uniform dim-0 all-gathers through NVLS symmetric memory when # enabled (mirrors reduce_scatter); otherwise fall back to the # PyNccl/base-class all-gather. Sequence parallelism's # gather-before-GEMM uses dim=0 with tp-aligned (uniform) shards. if dim < 0: dim += input_.dim() if dim == 0 and should_nccl_symm_mem_ag_rs(): return self._all_gather_symm_mem(input_.contiguous()) pynccl_comm = self.pynccl_comm if pynccl_comm is None or pynccl_comm.disabled: return super().all_gather(input_, dim) # On ROCm, the base-class all_gather (all_gather_into_tensor) is faster # than the manual pynccl + torch.empty + movedim + reshape path below, # which adds a per-call output allocation and (for dim != 0) an extra # copy on every step. This is on the hot path for TP forward passes, so # keep ROCm on the base-class collective to avoid a decode regression. if current_platform.is_rocm(): return super().all_gather(input_, dim) input_size = input_.size() output_size = (input_size[0] * self.world_size,) + input_size[1:] output_tensor = torch.empty( output_size, dtype=input_.dtype, device=input_.device ) pynccl_comm.all_gather(output_tensor, input_.contiguous()) output_tensor = output_tensor.reshape((self.world_size,) + input_size) output_tensor = output_tensor.movedim(0, dim) return output_tensor.reshape( input_size[:dim] + (self.world_size * input_size[dim],) + input_size[dim + 1 :] ) def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): world_size = self.world_size pynccl_comm = self.pynccl_comm assert pynccl_comm is not None if dim < 0: # Convert negative dim to positive. dim += input_.dim() # Note: This will produce an incorrect answer if we don't make # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? input_tensor = input_.movedim(0, dim).contiguous() assert input_tensor.shape[0] % world_size == 0 chunk_size = input_tensor.shape[0] // world_size output_shape = (chunk_size,) + input_tensor.shape[1:] if should_nccl_symm_mem_ag_rs(): output = self._reduce_scatter_symm_mem(input_tensor) else: output = torch.empty( output_shape, dtype=input_tensor.dtype, device=input_tensor.device ) pynccl_comm.reduce_scatter(output, input_tensor) # Reshape before returning return output.movedim(0, dim).contiguous() def reduce_scatterv( self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None ): world_size = self.world_size pynccl_comm = self.pynccl_comm assert pynccl_comm is not None if dim < 0: # Convert negative dim to positive. dim += input_.dim() # 'sizes' is not needed if all inputs in the same group have the same # shape if sizes is not None and all(s == sizes[0] for s in sizes): sizes = None # Note: This will produce an incorrect answer if we don't make # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? input_tensor = input_.movedim(0, dim).contiguous() if sizes is not None: assert len(sizes) == world_size, f"{len(sizes)} == {world_size}" assert input_tensor.shape[0] == sum(sizes) chunk_size = sizes[self.rank_in_group] else: assert input_tensor.shape[0] % world_size == 0 chunk_size = input_tensor.shape[0] // world_size output_shape = (chunk_size,) + input_tensor.shape[1:] if self._can_use_aiter_ag_rs(sizes): aiter_comm = self.aiter_ar_comm assert aiter_comm is not None if aiter_comm.should_custom_rs(input_tensor, dim=0): output = torch.empty( output_shape, dtype=input_tensor.dtype, device=input_tensor.device ) aiter_comm.custom_reduce_scatter(input_tensor, output, dim=0) return output.movedim(0, dim).contiguous() # Symmetric memory is only used when all ranks have uniform sizes. # ncclCommWindowRegister is collective: asymmetric pool allocations # from variable per-rank sizes cause deadlocks. use_symm_mem = sizes is None and should_nccl_symm_mem_ag_rs() if use_symm_mem: output = self._reduce_scatter_symm_mem(input_tensor) else: output = torch.empty( output_shape, dtype=input_tensor.dtype, device=input_tensor.device ) use_deterministic_rs = envs.VLLM_BATCH_INVARIANT and world_size > 2 if use_deterministic_rs: # Reduce to a fixed root (0) for determinism reduced = torch.empty_like(input_tensor) sizes = sizes if sizes else [chunk_size] * world_size pynccl_comm.reduce(reduced, input_tensor, root=0) pynccl_comm.scatter(output, reduced, sizes, root=0) elif sizes is not None and sizes.count(sizes[0]) != len(sizes): pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes) else: pynccl_comm.reduce_scatter(output, input_tensor) # Reshape before returning return output.movedim(0, dim).contiguous() def _get_symm_scratch( self, role: str, shape: tuple[int, ...], dtype: torch.dtype, device: torch.device, ) -> torch.Tensor: """Persistent, pre-registered NCCL symmetric-memory scratch buffer. Allocating a fresh symm tensor per collective pays the ``nccl_symm_mem_context`` snapshot + window-registration scan on every call (~0.5 ms/RS+AG pair, dwarfing the NVLS transfer itself). Instead we allocate once per ``(role, shape, dtype)``, register once, and reuse. Safe for serial (eager) sequence parallelism: each collective's result is consumed on the same stream before the next same-role collective reuses the buffer. Distinct roles (e.g. ``rs_in`` vs ``ag_out``, both full-size) get distinct buffers so a reduce-scatter input copy never clobbers a still-live all-gather output. """ from vllm.distributed.device_communicators.pynccl_allocator import ( nccl_symm_mem_context, ) pynccl_comm = self.pynccl_comm assert pynccl_comm is not None cache = self.__dict__.setdefault("_symm_scratch_bufs", {}) key = (role, tuple(shape), dtype) buf = cache.get(key) if buf is None: with nccl_symm_mem_context(pynccl_comm): buf = torch.empty(shape, dtype=dtype, device=device) cache[key] = buf return buf def _reduce_scatter_symm_mem( self, input_tensor: torch.Tensor, ) -> torch.Tensor: """ReduceScatter using NCCL symmetric memory (NVLS). Only called for uniform-size reduce_scatter (variable sizes are guarded out by the caller to avoid asymmetric ncclCommWindowRegister). Uses persistent pre-registered scratch (see _get_symm_scratch). """ from vllm.distributed.device_communicators.pynccl_allocator import ( is_symmetric_memory_tensor, ) pynccl_comm = self.pynccl_comm assert pynccl_comm is not None chunk = input_tensor.shape[0] // self.world_size output_shape = (chunk,) + tuple(input_tensor.shape[1:]) symm_output = self._get_symm_scratch( "rs_out", output_shape, input_tensor.dtype, input_tensor.device ) # NVLS reduce-scatter (LDMC) requires the input in symmetric memory. if is_symmetric_memory_tensor(input_tensor): symm_input = input_tensor else: symm_input = self._get_symm_scratch( "rs_in", tuple(input_tensor.shape), input_tensor.dtype, input_tensor.device, ) symm_input.copy_(input_tensor) pynccl_comm.reduce_scatter(symm_output, symm_input) return symm_output def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: """Sends a tensor to the destination rank in a blocking way.""" """NOTE: `dst` is the local rank of the destination rank.""" if dst is None: dst = (self.rank_in_group + 1) % self.world_size pynccl_comm = self.pynccl_comm if pynccl_comm is not None and not pynccl_comm.disabled: pynccl_comm.send(tensor, dst) else: torch.distributed.send(tensor, self.ranks[dst], self.device_group) def recv( self, size: torch.Size, dtype: torch.dtype, src: int | None = None ) -> torch.Tensor: """Receives a tensor from the source rank.""" """NOTE: `src` is the local rank of the source rank.""" if src is None: src = (self.rank_in_group - 1) % self.world_size tensor = torch.empty(size, dtype=dtype, device=self.device) pynccl_comm = self.pynccl_comm if pynccl_comm is not None and not pynccl_comm.disabled: pynccl_comm.recv(tensor, src) else: torch.distributed.recv(tensor, self.ranks[src], self.device_group) return tensor def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor: """Broadcast a tensor from source rank to all ranks.""" if self.world_size == 1: return tensor pynccl_comm = self.pynccl_comm if pynccl_comm is not None and not pynccl_comm.disabled: pynccl_comm.broadcast(tensor, src) return tensor else: raise ValueError("No PyNCCL communicator found") def destroy(self): if self.pynccl_comm is not None: self.pynccl_comm.destroy() self.pynccl_comm = None if self.ca_comm is not None: self.ca_comm = None if self.aiter_ar_comm is not None: self.aiter_ar_comm.close() self.aiter_ar_comm = None if self.fi_ar_comm is not None: self.fi_ar_comm.destroy() self.fi_ar_comm = None if self.fi_pcie_ipc_ar_comm is not None: self.fi_pcie_ipc_ar_comm.destroy() self.fi_pcie_ipc_ar_comm = None if self.all2all_manager is not None: self.all2all_manager.destroy() self.all2all_manager = None # type: ignore[assignment] def _can_use_aiter_ag_rs(self, sizes: list[int] | None) -> bool: """Whether the AITER custom AG/RS fast path may run for this collective. Requires: - uniform batches - FULL CUDAgraphs """ if ( not self.use_aiter_ag_rs or self.aiter_ar_comm is None or self.aiter_ar_comm.disabled ): return False if sizes is not None: return False from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import get_forward_context try: ctx = get_forward_context() except AssertionError: return False if ctx.cudagraph_runtime_mode != CUDAGraphMode.FULL: return False bd = ctx.batch_descriptor return bd is not None and bd.uniform def suspend(self) -> None: if self.pynccl_comm is not None: self.pynccl_comm.suspend() def resume(self) -> None: if self.pynccl_comm is not None: self.pynccl_comm.resume() def checkpoint_prepare(self) -> None: # Only FlashInfer all-reduce and FlashInfer all2all are supported for now. from .flashinfer_all_reduce import checkpoint_prepare_fi_ar_workspaces checkpoint_prepare_fi_ar_workspaces(self.cpu_group) if self.all2all_manager is not None: self.all2all_manager.checkpoint_prepare() def checkpoint_restore(self) -> None: # Only FlashInfer all-reduce and FlashInfer all2all are supported for now. from .flashinfer_all_reduce import checkpoint_restore_fi_ar_workspaces checkpoint_restore_fi_ar_workspaces(self.cpu_group) if self.all2all_manager is not None: self.all2all_manager.checkpoint_restore() def all_gatherv( self, input_: torch.Tensor | list[torch.Tensor], dim: int = 0, sizes: list[int] | None = None, ): if dim != 0: raise NotImplementedError("only dim 0 all-gatherv is supported") world_size = self.world_size # 'sizes' is not needed if all inputs in the same group have the same # shape if sizes is not None and all(s == sizes[0] for s in sizes): sizes = None if self._can_use_aiter_ag_rs(sizes): aiter_comm = self.aiter_ar_comm assert aiter_comm is not None if isinstance(input_, torch.Tensor): if aiter_comm.should_custom_ag(input_): out = aiter_comm.custom_all_gather(input_, dim=0) if out is not None: return out elif all(aiter_comm.should_custom_ag(inp) for inp in input_): outs = [aiter_comm.custom_all_gather(inp, dim=0) for inp in input_] if all(o is not None for o in outs): return outs pynccl_comm = self.pynccl_comm assert pynccl_comm is not None and not pynccl_comm.disabled # Symmetric memory is only used when all ranks have uniform sizes. # ncclCommWindowRegister is collective: asymmetric pool allocations # from variable per-rank sizes cause deadlocks. if sizes is None and should_nccl_symm_mem_ag_rs(): if isinstance(input_, torch.Tensor): return self._all_gather_symm_mem(input_) return self._all_gather_batched_symm_mem(input_) def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None): input_size = input_.size() if sizes is not None: assert len(sizes) == world_size assert input_.shape[dim] == sizes[self.rank_in_group], ( f"{input_.shape[dim]} != {sizes[self.rank_in_group]}" ) output_size = (sum(sizes),) + input_size[1:] else: output_size = (input_size[0] * world_size,) + input_size[1:] # Allocate output tensor. output_tensor = torch.empty( output_size, dtype=input_.dtype, device=input_.device ) if sizes is not None: pynccl_comm.all_gatherv(output_tensor, input_, sizes=sizes) else: pynccl_comm.all_gather(output_tensor, input_) return output_tensor if isinstance(input_, torch.Tensor): return _all_gather_single(input_, sizes) output_list = [] pynccl_comm.group_start() for inp in input_: output_list.append(_all_gather_single(inp, sizes=sizes)) pynccl_comm.group_end() return output_list def _all_gather_symm_mem(self, input_: torch.Tensor) -> torch.Tensor: """AllGather a single tensor using NCCL symmetric memory (NVLS). Only the output needs to be in symmetric memory; NCCL does not require the AG input to be symmetrically allocated. """ pynccl_comm = self.pynccl_comm assert pynccl_comm is not None out_size = (input_.size(0) * self.world_size,) + tuple(input_.size()[1:]) # Persistent pre-registered scratch avoids the per-call symm-mem context # snapshot/registration overhead (see _get_symm_scratch). symm_output = self._get_symm_scratch( "ag_out", out_size, input_.dtype, input_.device ) pynccl_comm.all_gather(symm_output, input_) return symm_output def _all_gather_batched_symm_mem( self, inputs: list[torch.Tensor] ) -> list[torch.Tensor]: """AllGather a list of tensors using NCCL symmetric memory (NVLS). Uses group_start/group_end to batch the collectives. Only the output needs to be in symmetric memory (see _all_gather_symm_mem). """ from vllm.distributed.device_communicators.pynccl_allocator import ( nccl_symm_mem_context, ) pynccl_comm = self.pynccl_comm assert pynccl_comm is not None world_size = self.world_size symm_outputs = [] with nccl_symm_mem_context(pynccl_comm): for inp in inputs: out_size = (inp.size(0) * world_size,) + inp.size()[1:] symm_outputs.append( torch.empty(out_size, dtype=inp.dtype, device=inp.device) ) pynccl_comm.group_start() for symm_out, inp in zip(symm_outputs, inputs): pynccl_comm.all_gather(symm_out, inp) pynccl_comm.group_end() return symm_outputs def dispatch_router_logits( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, is_sequence_parallel: bool = False, extra_tensors: list[torch.Tensor] | None = None, ) -> ( tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] ): """Dispatch the hidden states and router logits to the appropriate device. This is a no-op in the base class. """ assert self.all2all_manager is not None return self.all2all_manager.dispatch_router_logits( hidden_states, router_logits, is_sequence_parallel, extra_tensors, ) def dispatch( self, hidden_states: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, is_sequence_parallel: bool = False, extra_tensors: list[torch.Tensor] | None = None, ) -> ( tuple[torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] ): """Dispatch the hidden states and topk weights/ids to the appropriate device. This is a no-op in the base class. """ assert self.all2all_manager is not None return self.all2all_manager.dispatch( hidden_states, topk_weights, topk_ids, is_sequence_parallel, extra_tensors=extra_tensors, ) def combine( self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False ) -> torch.Tensor: """Combine the hidden states and router logits from the appropriate device. This is a no-op in the base class. """ assert self.all2all_manager is not None return self.all2all_manager.combine( hidden_states, is_sequence_parallel, ) def batch_isend_irecv(self, p2p_ops: list): pynccl_comm = self.pynccl_comm if pynccl_comm is not None and not pynccl_comm.disabled: pynccl_comm.batch_isend_irecv(p2p_ops) else: raise ValueError("No PyNCCL communicator found")