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
)
# Neither AITER nor QuickReduce all-reduce has a fixed reduction order.
use_aiter_allreduce = (
use_custom_allreduce
and not envs.VLLM_BATCH_INVARIANT
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)
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
# cuMem graph buffers cannot be IPC-registered; capture copies them instead.
config = get_current_vllm_config_or_none()
register = config is None or not config.use_cumem_cudagraph_pool
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
),
register_graph_buffers=register,
)
# 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()
and not envs.VLLM_BATCH_INVARIANT
):
# 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)
elif self.all2all_backend == "passthrough":
from .all2all import PassThroughAll2AllManager
self.all2all_manager = PassThroughAll2AllManager(
self.cpu_group, tcp_store_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 "<unnamed>",
"[" + ", ".join(f"'{b}'" for b in all_potential_ar_backends) + "]",
scope="global",
)
def all_reduce(self, 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:
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:
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
):
return self._reduce_scatterv(input_, dim, sizes)
def reduce_scatterv_into_output(
self,
input_: torch.Tensor,
output: torch.Tensor,
dim: int = -1,
sizes: list[int] | None = None,
) -> torch.Tensor:
return self._reduce_scatterv(input_, dim, sizes, output)
def _reduce_scatterv(
self,
input_: torch.Tensor,
dim: int,
sizes: list[int] | None,
output: torch.Tensor | None = None,
) -> torch.Tensor:
world_size = self.world_size
pynccl_comm = self.pynccl_comm
assert pynccl_comm is not None
assert not pynccl_comm.disabled, "reduce_scatterv requires PyNccl"
if dim < 0:
# Convert negative dim to positive.
dim += input_.dim()
sizes_are_explicit = sizes is not None
# '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 output is not None:
assert dim == 0, "preallocated reduce-scatter output requires dim=0"
assert output.shape == output_shape
assert output.dtype == input_tensor.dtype
assert output.device == input_tensor.device
assert output.is_contiguous()
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):
if output is None:
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.
uniform_sizes = sizes is None or all(size == sizes[0] for size in sizes)
direct_output_supported = (
output is None
or is_symmetric_memory_tensor(output)
or pynccl_comm.nccl_version >= NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION
)
use_symm_mem = (
uniform_sizes
and direct_output_supported
and (
output is None
or not sizes_are_explicit
or is_symmetric_memory_tensor(input_tensor)
)
and should_nccl_symm_mem_ag_rs()
)
if use_symm_mem:
output = self._reduce_scatter_symm_mem(input_tensor, output)
else:
if output is None:
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)
scatter_sizes = sizes
if scatter_sizes is None:
scatter_sizes = [chunk_size] * world_size
pynccl_comm.reduce(reduced, input_tensor, root=0)
pynccl_comm.scatter(output, reduced, scatter_sizes, root=0)
elif not uniform_sizes:
assert sizes is not None
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,
each ``(role, shape[1:], dtype, device)`` keeps a geometric high-water
allocation and returns a leading slice. Superseded allocations remain
referenced because a captured CUDA graph may still use their pointers;
geometric growth bounds their total capacity to less than twice the
current allocation.
Safe across serial MoE layers and eager sequence parallelism: the
producer and collective are ordered on the same stream before the next
same-role operation reuses the buffer. DBO microbatches use distinct
cache entries. Any future cross-layer communication overlap must also
use distinct roles.
"""
from vllm.distributed.device_communicators.pynccl_allocator import (
nccl_symm_mem_context,
)
from vllm.v1.worker.ubatching import dbo_current_ubatch_id
pynccl_comm = self.pynccl_comm
assert pynccl_comm is not None
assert shape, "symmetric scratch buffers require at least one dimension"
cache = self.__dict__.setdefault("_symm_scratch_bufs", {})
key = (role, dbo_current_ubatch_id(), tuple(shape[1:]), dtype, device)
buf = cache.get(key)
requested_rows = shape[0]
if buf is None or buf.shape[0] < requested_rows:
capacity = (
requested_rows if buf is None else max(requested_rows, 2 * buf.shape[0])
)
with nccl_symm_mem_context(pynccl_comm):
new_buf = torch.empty(
(capacity, *shape[1:]), dtype=dtype, device=device
)
if buf is not None:
retired = self.__dict__.setdefault("_retired_symm_scratch_bufs", {})
retired.setdefault(key, []).append(buf)
buf = new_buf
cache[key] = buf
return buf[:requested_rows]
def _reduce_scatter_symm_mem(
self,
input_tensor: torch.Tensor,
output: torch.Tensor | None = None,
) -> torch.Tensor:
"""ReduceScatter using NCCL symmetric memory (NVLS).
The MoE reduce_scatterv path passes explicit uniform sizes and calls
this only with an already-registered input, so that path never stages.
Uniform calls without a caller-owned output retain their existing
opt-in behavior and stage ordinary input into registered scratch.
The output may be ordinary memory.
"""
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:])
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)
if output is None:
output = self._get_symm_scratch(
"rs_out", output_shape, input_tensor.dtype, input_tensor.device
)
pynccl_comm.reduce_scatter(output, symm_input)
return output
def get_symmetric_memory_buffer(
self,
role: str,
shape: tuple[int, ...],
dtype: torch.dtype,
device: torch.device,
) -> torch.Tensor | None:
pynccl_comm = self.pynccl_comm
if (
pynccl_comm is None
or pynccl_comm.disabled
or pynccl_comm.world_size == 1
or pynccl_comm.nccl_version < NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION
or not should_nccl_symm_mem_ag_rs()
):
return None
output = self._get_symm_scratch(role, shape, dtype, device)
return output if is_symmetric_memory_tensor(output) else None
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 allocate_combine_input(
self,
shape: tuple[int, ...],
dtype: torch.dtype,
device: torch.device,
is_sequence_parallel: bool = False,
) -> torch.Tensor | None:
assert self.all2all_manager is not None
return self.all2all_manager.allocate_combine_input(
shape, dtype, device, is_sequence_parallel
)
def combine_into_output(
self,
hidden_states: torch.Tensor,
output: torch.Tensor,
is_sequence_parallel: bool = False,
) -> torch.Tensor:
assert self.all2all_manager is not None
return self.all2all_manager.combine_into_output(
hidden_states, output, 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")