| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | from torch.distributed import ProcessGroup |
| 7 | |
| 8 | import vllm.envs as envs |
| 9 | from vllm._aiter_ops import rocm_aiter_ops |
| 10 | from vllm.distributed.device_communicators.all_reduce_utils import ( |
| 11 | NCCL_SYMM_MEM_ALL_REDUCE_CONFIG, |
| 12 | should_nccl_symm_mem_ag_rs, |
| 13 | should_nccl_symm_mem_allreduce, |
| 14 | ) |
| 15 | from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops |
| 16 | from vllm.distributed.device_communicators.pynccl_allocator import ( |
| 17 | is_symmetric_memory_enabled, |
| 18 | ) |
| 19 | from vllm.logger import init_logger |
| 20 | from vllm.platforms import current_platform |
| 21 | |
| 22 | from ..utils import StatelessProcessGroup |
| 23 | from .aiter_custom_all_reduce import AiterCustomAllreduce |
| 24 | from .base_device_communicator import DeviceCommunicatorBase |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | |
| 28 | |
| 29 | class CudaCommunicator(DeviceCommunicatorBase): |
| 30 | def __init__( |
| 31 | self, |
| 32 | cpu_group: ProcessGroup, |
| 33 | device: torch.device | None = None, |
| 34 | device_group: ProcessGroup | None = None, |
| 35 | unique_name: str = "", |
| 36 | global_ranks: list[int] | None = None, |
| 37 | global_world_size: int | None = None, |
| 38 | tcp_store_group: StatelessProcessGroup | None = None, |
| 39 | use_all2all: bool = False, |
| 40 | ): |
| 41 | super().__init__( |
| 42 | cpu_group, |
| 43 | device, |
| 44 | device_group, |
| 45 | unique_name, |
| 46 | global_ranks, |
| 47 | global_world_size, |
| 48 | use_all2all=use_all2all, |
| 49 | ) |
| 50 | # Match the group name exactly so ETP does not enable TP-only backends. |
| 51 | if unique_name.split(":")[0] != "tp": |
| 52 | # custom allreduce or torch symm mem can be used only by tp |
| 53 | use_custom_allreduce = False |
| 54 | use_torch_symm_mem = False |
| 55 | use_flashinfer_allreduce = False |
| 56 | use_flashinfer_pcie_ipc_allreduce = False |
| 57 | use_aiter_allreduce = False |
| 58 | else: |
| 59 | from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE |
| 60 | |
| 61 | use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE |
| 62 | use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM |
| 63 | # FlashInfer all-reduce does not provide a fixed reduction order. |
| 64 | use_flashinfer_allreduce = ( |
| 65 | envs.VLLM_ALLREDUCE_USE_FLASHINFER and not envs.VLLM_BATCH_INVARIANT |
| 66 | ) |
| 67 | use_flashinfer_pcie_ipc_allreduce = ( |
| 68 | envs.VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC |
| 69 | and not envs.VLLM_BATCH_INVARIANT |
| 70 | ) |
| 71 | # Neither AITER nor QuickReduce all-reduce has a fixed reduction order. |
| 72 | use_aiter_allreduce = ( |
| 73 | use_custom_allreduce |
| 74 | and not envs.VLLM_BATCH_INVARIANT |
| 75 | and bool(rocm_aiter_ops.is_custom_all_reduce_enabled()) |
| 76 | ) |
| 77 | |
| 78 | self.use_custom_allreduce = use_custom_allreduce |
| 79 | self.use_torch_symm_mem = use_torch_symm_mem |
| 80 | self.use_flashinfer_allreduce = use_flashinfer_allreduce |
| 81 | self.use_flashinfer_pcie_ipc_allreduce = use_flashinfer_pcie_ipc_allreduce |
| 82 | self.use_aiter_allreduce = use_aiter_allreduce |
| 83 | |
| 84 | # lazy import to avoid documentation build error |
| 85 | from vllm.distributed.device_communicators.custom_all_reduce import ( |
| 86 | CustomAllreduce, |
| 87 | ) |
| 88 | from vllm.distributed.device_communicators.flashinfer_all_reduce import ( |
| 89 | FlashInferAllReduce, |
| 90 | ) |
| 91 | from vllm.distributed.device_communicators.flashinfer_pcie_ipc_all_reduce import ( # noqa: E501 |
| 92 | FlashInferPcieIpcAllReduce, |
| 93 | ) |
| 94 | from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator |
| 95 | from vllm.distributed.device_communicators.quick_all_reduce import ( |
| 96 | QuickAllReduce, |
| 97 | ) |
| 98 | from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator |
| 99 | |
| 100 | self.pynccl_comm: PyNcclCommunicator | None = None |
| 101 | if self.world_size > 1: |
| 102 | self.pynccl_comm = PyNcclCommunicator( |
| 103 | group=self.cpu_group if tcp_store_group is None else tcp_store_group, |
| 104 | device=self.device, |
| 105 | ) |
| 106 | if is_symmetric_memory_enabled(): |
| 107 | register_nccl_symmetric_ops(self.pynccl_comm) |
| 108 | |
| 109 | self.ca_comm: CustomAllreduce | None = None |
| 110 | self.qr_comm: QuickAllReduce | None = None |
| 111 | self.symm_mem_comm: SymmMemCommunicator | None = None |
| 112 | self.fi_ar_comm: FlashInferAllReduce | None = None |
| 113 | self.fi_pcie_ipc_ar_comm: FlashInferPcieIpcAllReduce | None = None |
| 114 | self.aiter_ar_comm: AiterCustomAllreduce | None = None |
| 115 | self.use_aiter_ag_rs: bool = False |
| 116 | |
| 117 | if use_torch_symm_mem and current_platform.is_cuda(): |
| 118 | self.symm_mem_comm = SymmMemCommunicator( |
| 119 | group=self.cpu_group, |
| 120 | device=self.device, |
| 121 | ) |
| 122 | |
| 123 | if self.use_flashinfer_allreduce and self.world_size > 1: |
| 124 | self.fi_ar_comm = FlashInferAllReduce( |
| 125 | group=self.cpu_group, |
| 126 | device=self.device, |
| 127 | ) |
| 128 | |
| 129 | if ( |
| 130 | self.use_flashinfer_pcie_ipc_allreduce |
| 131 | and self.world_size > 1 |
| 132 | and self.device_group is not None |
| 133 | ): |
| 134 | self.fi_pcie_ipc_ar_comm = FlashInferPcieIpcAllReduce( |
| 135 | group=self.device_group, |
| 136 | tune_group=self.cpu_group, |
| 137 | device=self.device, |
| 138 | ) |
| 139 | |
| 140 | if self.use_aiter_allreduce and self.world_size > 1: |
| 141 | self.aiter_ar_comm = AiterCustomAllreduce( |
| 142 | group=self.cpu_group, |
| 143 | device=self.device, |
| 144 | ) |
| 145 | |
| 146 | if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1: |
| 147 | # Initialize a custom fast all-reduce implementation. |
| 148 | self.ca_comm = CustomAllreduce( |
| 149 | group=self.cpu_group, |
| 150 | device=self.device, |
| 151 | symm_mem_enabled=( |
| 152 | self.symm_mem_comm is not None and not self.symm_mem_comm.disabled |
| 153 | ), |
| 154 | ) |
| 155 | |
| 156 | # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine |
| 157 | if ( |
| 158 | "dp" in unique_name |
| 159 | and self.world_size in (2, 4, 8) |
| 160 | and current_platform.is_rocm() |
| 161 | and rocm_aiter_ops.is_custom_all_reduce_enabled() |
| 162 | ): |
| 163 | self.aiter_ar_comm = AiterCustomAllreduce( |
| 164 | group=self.cpu_group, |
| 165 | device=self.device, |
| 166 | ) |
| 167 | if self.aiter_ar_comm.disabled: |
| 168 | self.aiter_ar_comm = None |
| 169 | else: |
| 170 | self.use_aiter_ag_rs = True |
| 171 | |
| 172 | if ( |
| 173 | use_custom_allreduce |
| 174 | and self.world_size > 1 |
| 175 | and current_platform.is_rocm() |
| 176 | and not envs.VLLM_BATCH_INVARIANT |
| 177 | ): |
| 178 | # Initialize a custom quick all-reduce implementation for AMD. |
| 179 | # Quick reduce is designed as a complement to custom allreduce |
| 180 | # (vLLM's or AITER's), so it is initialized for either backend. |
| 181 | # Based on quickreduce (https://github.com/mk1-project/quickreduce). |
| 182 | # On ROCm, 'use_custom_allreduce==True' means it must currently be |
| 183 | # an MI300 series. |
| 184 | self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device) |
| 185 | |
| 186 | if self.world_size > 1: |
| 187 | self._log_all_reduce_backend_selection() |
| 188 | |
| 189 | if self.use_all2all: |
| 190 | if self.all2all_backend in ("naive", "allgather_reducescatter"): |
| 191 | from .all2all import AgRsAll2AllManager |
| 192 | |
| 193 | self.all2all_manager = AgRsAll2AllManager( |
| 194 | self.cpu_group, tcp_store_group |
| 195 | ) |
| 196 | elif self.all2all_backend == "deepep_high_throughput": |
| 197 | from .all2all import DeepEPHTAll2AllManager |
| 198 | |
| 199 | self.all2all_manager = DeepEPHTAll2AllManager( |
| 200 | self.cpu_group, tcp_store_group |
| 201 | ) |
| 202 | elif self.all2all_backend == "deepep_low_latency": |
| 203 | from .all2all import DeepEPLLAll2AllManager |
| 204 | |
| 205 | self.all2all_manager = DeepEPLLAll2AllManager( |
| 206 | self.cpu_group, tcp_store_group |
| 207 | ) |
| 208 | elif self.all2all_backend in ( |
| 209 | "mori_high_throughput", |
| 210 | "mori_low_latency", |
| 211 | ): |
| 212 | from .all2all import MoriAll2AllManager |
| 213 | |
| 214 | self.all2all_manager = MoriAll2AllManager( |
| 215 | self.cpu_group, self.all2all_backend |
| 216 | ) |
| 217 | elif self.all2all_backend == "deepep_v2": |
| 218 | from .all2all import DeepEPV2All2AllManager |
| 219 | |
| 220 | self.all2all_manager = DeepEPV2All2AllManager( |
| 221 | self.cpu_group, |
| 222 | tcp_store_group, |
| 223 | device_group=self.device_group, |
| 224 | ) |
| 225 | elif self.all2all_backend == "moonep": |
| 226 | from .all2all import MoonEPAll2AllManager |
| 227 | |
| 228 | self.all2all_manager = MoonEPAll2AllManager( |
| 229 | self.cpu_group, |
| 230 | tcp_store_group, |
| 231 | device_group=self.device_group, |
| 232 | ) |
| 233 | elif self.all2all_backend == "nixl_ep": |
| 234 | from .all2all import NixlEPAll2AllManager |
| 235 | |
| 236 | self.all2all_manager = NixlEPAll2AllManager( |
| 237 | self.cpu_group, tcp_store_group |
| 238 | ) |
| 239 | elif ( |
| 240 | self.all2all_backend == "flashinfer_all2allv" |
| 241 | or self.all2all_backend == "flashinfer_nvlink_two_sided" |
| 242 | ): |
| 243 | if self.all2all_backend == "flashinfer_all2allv": |
| 244 | logger.warning_once( |
| 245 | "'flashinfer_all2allv' is deprecated and has been renamed to" |
| 246 | "'flashinfer_nvlink_two_sided'. It will be removed in a future" |
| 247 | "release." |
| 248 | ) |
| 249 | from .all2all import FlashInferNVLinkTwoSidedManager |
| 250 | |
| 251 | self.all2all_manager = FlashInferNVLinkTwoSidedManager( |
| 252 | self.cpu_group, tcp_store_group |
| 253 | ) |
| 254 | elif self.all2all_backend == "flashinfer_nvlink_one_sided": |
| 255 | from .all2all import FlashInferNVLinkOneSidedManager |
| 256 | |
| 257 | self.all2all_manager = FlashInferNVLinkOneSidedManager(self.cpu_group) |
| 258 | else: |
| 259 | raise ValueError(f"Unknown all2all backend: {self.all2all_backend}") |
| 260 | |
| 261 | logger.info_once( |
| 262 | "Using %s all2all manager.", |
| 263 | self.all2all_manager.__class__.__name__, |
| 264 | scope="global", |
| 265 | ) |
| 266 | |
| 267 | def _log_all_reduce_backend_selection(self) -> None: |
| 268 | """Log the all-reduce backends that are active for this group. |
| 269 | |
| 270 | The dispatch chain in ``all_reduce`` tries backends in this order and |
| 271 | falls through to the next one if the current backend rejects the |
| 272 | input (size/dtype gates) or is disabled. The list of "enabled" |
| 273 | backends below is the subset of potential backends that may be |
| 274 | chosen at dispatch time for this group; the actual per-call choice |
| 275 | depends on the input tensor. |
| 276 | """ |
| 277 | all_potential_ar_backends = [ |
| 278 | "FLASHINFER_PCIE_IPC", |
| 279 | "FLASHINFER", |
| 280 | "NCCL_SYMM_MEM", |
| 281 | "QUICK_REDUCE", |
| 282 | "AITER_CUSTOM", |
| 283 | "CUSTOM", |
| 284 | "SYMM_MEM", |
| 285 | "PYNCCL", |
| 286 | ] |
| 287 | enabled_ar_backends: list[str] = [] |
| 288 | if ( |
| 289 | self.fi_pcie_ipc_ar_comm is not None |
| 290 | and not self.fi_pcie_ipc_ar_comm.disabled |
| 291 | ): |
| 292 | enabled_ar_backends.append("FLASHINFER_PCIE_IPC") |
| 293 | if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled: |
| 294 | enabled_ar_backends.append("FLASHINFER") |
| 295 | # Mirror the static preconditions of `should_nccl_symm_mem_allreduce`: |
| 296 | # VLLM_BATCH_INVARIANT off, NCCL symm mem enabled, world_size meets |
| 297 | # min_world_size, and world_size either has a tuned entry in |
| 298 | # `custom_ar_preferred_ranges` or is greater than |
| 299 | # `always_use_above_world_size`. World sizes that fail the latter (e.g. |
| 300 | # 5/6/7 with the default config) never dispatch NCCL symm mem |
| 301 | # regardless of input. The per-tensor-size check inside the function |
| 302 | # stays as a runtime decision. |
| 303 | nccl_symm_ws_ok = self.world_size >= NCCL_SYMM_MEM_ALL_REDUCE_CONFIG[ |
| 304 | "min_world_size" |
| 305 | ] and ( |
| 306 | self.world_size |
| 307 | in NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["custom_ar_preferred_ranges"] |
| 308 | or self.world_size |
| 309 | > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"] |
| 310 | ) |
| 311 | if ( |
| 312 | self.pynccl_comm is not None |
| 313 | and not self.pynccl_comm.disabled |
| 314 | and is_symmetric_memory_enabled() |
| 315 | and not envs.VLLM_BATCH_INVARIANT |
| 316 | and nccl_symm_ws_ok |
| 317 | ): |
| 318 | enabled_ar_backends.append("NCCL_SYMM_MEM") |
| 319 | if self.qr_comm is not None and not self.qr_comm.disabled: |
| 320 | enabled_ar_backends.append("QUICK_REDUCE") |
| 321 | if ( |
| 322 | self.use_aiter_allreduce |
| 323 | and self.aiter_ar_comm is not None |
| 324 | and not self.aiter_ar_comm.disabled |
| 325 | ): |
| 326 | enabled_ar_backends.append("AITER_CUSTOM") |
| 327 | if self.ca_comm is not None and not self.ca_comm.disabled: |
| 328 | enabled_ar_backends.append("CUSTOM") |
| 329 | if self.symm_mem_comm is not None and not self.symm_mem_comm.disabled: |
| 330 | enabled_ar_backends.append("SYMM_MEM") |
| 331 | if self.pynccl_comm is not None and not self.pynccl_comm.disabled: |
| 332 | enabled_ar_backends.append("PYNCCL") |
| 333 | |
| 334 | logger.info_once( |
| 335 | "Using %s all-reduce backends (in dispatch order) for group " |
| 336 | "'%s' out of potential backends: %s.", |
| 337 | "[" + ", ".join(f"'{b}'" for b in enabled_ar_backends) + "]", |
| 338 | self.unique_name or "<unnamed>", |
| 339 | "[" + ", ".join(f"'{b}'" for b in all_potential_ar_backends) + "]", |
| 340 | scope="global", |
| 341 | ) |
| 342 | |
| 343 | def all_reduce(self, input_): |
| 344 | fi_ar_comm = self.fi_ar_comm |
| 345 | use_fi_ar = ( |
| 346 | fi_ar_comm is not None |
| 347 | and not fi_ar_comm.disabled |
| 348 | and fi_ar_comm.should_use_fi_ar(input_) |
| 349 | ) |
| 350 | |
| 351 | # since currently we perform copy input -> symm_input -> out-of-place AR |
| 352 | # return symm_output, we don't need to check if input is symmetric |
| 353 | if ( |
| 354 | self.pynccl_comm is not None |
| 355 | and not use_fi_ar |
| 356 | and should_nccl_symm_mem_allreduce(self.pynccl_comm.world_size, input_) |
| 357 | ): |
| 358 | out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_) |
| 359 | if out is not None: |
| 360 | return out |
| 361 | qr_comm = self.qr_comm |
| 362 | if ( |
| 363 | qr_comm is not None |
| 364 | and not qr_comm.disabled |
| 365 | and qr_comm.should_quick_allreduce(input_) |
| 366 | ): |
| 367 | out = qr_comm.quick_all_reduce(input_) |
| 368 | assert out is not None |
| 369 | return out |
| 370 | fi_pcie_ipc_ar_comm = self.fi_pcie_ipc_ar_comm |
| 371 | if fi_pcie_ipc_ar_comm is not None and fi_pcie_ipc_ar_comm.should_use(input_): |
| 372 | return fi_pcie_ipc_ar_comm.all_reduce(input_) |
| 373 | if use_fi_ar: |
| 374 | assert fi_ar_comm is not None |
| 375 | out = fi_ar_comm.all_reduce(input_) |
| 376 | assert out is not None |
| 377 | return out |
| 378 | aiter_ar_comm = self.aiter_ar_comm |
| 379 | if ( |
| 380 | self.use_aiter_allreduce |
| 381 | and aiter_ar_comm is not None |
| 382 | and not aiter_ar_comm.disabled |
| 383 | and aiter_ar_comm.should_custom_ar(input_) |
| 384 | ): |
| 385 | out = aiter_ar_comm.custom_all_reduce(input_) |
| 386 | assert out is not None |
| 387 | return out |
| 388 | ca_comm = self.ca_comm |
| 389 | if ( |
| 390 | ca_comm is not None |
| 391 | and not ca_comm.disabled |
| 392 | and ca_comm.should_custom_ar(input_) |
| 393 | ): |
| 394 | out = ca_comm.custom_all_reduce(input_) |
| 395 | assert out is not None |
| 396 | return out |
| 397 | symm_mem_comm = self.symm_mem_comm |
| 398 | if symm_mem_comm is not None and symm_mem_comm.should_use_symm_mem(input_): |
| 399 | out = symm_mem_comm.all_reduce(input_) |
| 400 | assert out is not None |
| 401 | return out |
| 402 | pynccl_comm = self.pynccl_comm |
| 403 | if pynccl_comm is None or pynccl_comm.disabled: |
| 404 | out = input_.clone() |
| 405 | torch.distributed.all_reduce(out, group=self.device_group) |
| 406 | return out |
| 407 | assert pynccl_comm is not None |
| 408 | out = pynccl_comm.all_reduce(input_) |
| 409 | if out is None: |
| 410 | # fall back to the default all-reduce using PyTorch. |
| 411 | # this usually happens during testing. |
| 412 | # when we run the model, allreduce only happens for the TP |
| 413 | # group, where we always have either custom allreduce or pynccl. |
| 414 | out = input_.clone() |
| 415 | torch.distributed.all_reduce(out, group=self.device_group) |
| 416 | return out |
| 417 | |
| 418 | def custom_all_gather(self, input_: torch.Tensor) -> torch.Tensor | None: |
| 419 | ca_comm = self.ca_comm |
| 420 | if ca_comm is None: |
| 421 | return None |
| 422 | return ca_comm.custom_all_gather(input_.contiguous()) |
| 423 | |
| 424 | def custom_reduce_scatter(self, input_: torch.Tensor) -> torch.Tensor | None: |
| 425 | ca_comm = self.ca_comm |
| 426 | if ca_comm is None: |
| 427 | return None |
| 428 | return ca_comm.custom_reduce_scatter(input_.contiguous()) |
| 429 | |
| 430 | def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: |
| 431 | # Route uniform dim-0 all-gathers through NVLS symmetric memory when |
| 432 | # enabled (mirrors reduce_scatter); otherwise fall back to the |
| 433 | # PyNccl/base-class all-gather. Sequence parallelism's |
| 434 | # gather-before-GEMM uses dim=0 with tp-aligned (uniform) shards. |
| 435 | if dim < 0: |
| 436 | dim += input_.dim() |
| 437 | if dim == 0 and should_nccl_symm_mem_ag_rs(): |
| 438 | return self._all_gather_symm_mem(input_.contiguous()) |
| 439 | |
| 440 | pynccl_comm = self.pynccl_comm |
| 441 | if pynccl_comm is None or pynccl_comm.disabled: |
| 442 | return super().all_gather(input_, dim) |
| 443 | |
| 444 | # On ROCm, the base-class all_gather (all_gather_into_tensor) is faster |
| 445 | # than the manual pynccl + torch.empty + movedim + reshape path below, |
| 446 | # which adds a per-call output allocation and (for dim != 0) an extra |
| 447 | # copy on every step. This is on the hot path for TP forward passes, so |
| 448 | # keep ROCm on the base-class collective to avoid a decode regression. |
| 449 | if current_platform.is_rocm(): |
| 450 | return super().all_gather(input_, dim) |
| 451 | |
| 452 | input_size = input_.size() |
| 453 | output_size = (input_size[0] * self.world_size,) + input_size[1:] |
| 454 | output_tensor = torch.empty( |
| 455 | output_size, dtype=input_.dtype, device=input_.device |
| 456 | ) |
| 457 | pynccl_comm.all_gather(output_tensor, input_.contiguous()) |
| 458 | output_tensor = output_tensor.reshape((self.world_size,) + input_size) |
| 459 | output_tensor = output_tensor.movedim(0, dim) |
| 460 | return output_tensor.reshape( |
| 461 | input_size[:dim] |
| 462 | + (self.world_size * input_size[dim],) |
| 463 | + input_size[dim + 1 :] |
| 464 | ) |
| 465 | |
| 466 | def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): |
| 467 | world_size = self.world_size |
| 468 | pynccl_comm = self.pynccl_comm |
| 469 | assert pynccl_comm is not None |
| 470 | if dim < 0: |
| 471 | # Convert negative dim to positive. |
| 472 | dim += input_.dim() |
| 473 | |
| 474 | # Note: This will produce an incorrect answer if we don't make |
| 475 | # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? |
| 476 | input_tensor = input_.movedim(0, dim).contiguous() |
| 477 | |
| 478 | assert input_tensor.shape[0] % world_size == 0 |
| 479 | chunk_size = input_tensor.shape[0] // world_size |
| 480 | output_shape = (chunk_size,) + input_tensor.shape[1:] |
| 481 | |
| 482 | if should_nccl_symm_mem_ag_rs(): |
| 483 | output = self._reduce_scatter_symm_mem(input_tensor) |
| 484 | else: |
| 485 | output = torch.empty( |
| 486 | output_shape, dtype=input_tensor.dtype, device=input_tensor.device |
| 487 | ) |
| 488 | pynccl_comm.reduce_scatter(output, input_tensor) |
| 489 | |
| 490 | # Reshape before returning |
| 491 | return output.movedim(0, dim).contiguous() |
| 492 | |
| 493 | def reduce_scatterv( |
| 494 | self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None |
| 495 | ): |
| 496 | world_size = self.world_size |
| 497 | pynccl_comm = self.pynccl_comm |
| 498 | assert pynccl_comm is not None |
| 499 | if dim < 0: |
| 500 | # Convert negative dim to positive. |
| 501 | dim += input_.dim() |
| 502 | |
| 503 | # 'sizes' is not needed if all inputs in the same group have the same |
| 504 | # shape |
| 505 | if sizes is not None and all(s == sizes[0] for s in sizes): |
| 506 | sizes = None |
| 507 | |
| 508 | # Note: This will produce an incorrect answer if we don't make |
| 509 | # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? |
| 510 | input_tensor = input_.movedim(0, dim).contiguous() |
| 511 | |
| 512 | if sizes is not None: |
| 513 | assert len(sizes) == world_size, f"{len(sizes)} == {world_size}" |
| 514 | assert input_tensor.shape[0] == sum(sizes) |
| 515 | chunk_size = sizes[self.rank_in_group] |
| 516 | else: |
| 517 | assert input_tensor.shape[0] % world_size == 0 |
| 518 | chunk_size = input_tensor.shape[0] // world_size |
| 519 | output_shape = (chunk_size,) + input_tensor.shape[1:] |
| 520 | |
| 521 | if self._can_use_aiter_ag_rs(sizes): |
| 522 | aiter_comm = self.aiter_ar_comm |
| 523 | assert aiter_comm is not None |
| 524 | if aiter_comm.should_custom_rs(input_tensor, dim=0): |
| 525 | output = torch.empty( |
| 526 | output_shape, dtype=input_tensor.dtype, device=input_tensor.device |
| 527 | ) |
| 528 | aiter_comm.custom_reduce_scatter(input_tensor, output, dim=0) |
| 529 | return output.movedim(0, dim).contiguous() |
| 530 | |
| 531 | # Symmetric memory is only used when all ranks have uniform sizes. |
| 532 | # ncclCommWindowRegister is collective: asymmetric pool allocations |
| 533 | # from variable per-rank sizes cause deadlocks. |
| 534 | use_symm_mem = sizes is None and should_nccl_symm_mem_ag_rs() |
| 535 | if use_symm_mem: |
| 536 | output = self._reduce_scatter_symm_mem(input_tensor) |
| 537 | else: |
| 538 | output = torch.empty( |
| 539 | output_shape, dtype=input_tensor.dtype, device=input_tensor.device |
| 540 | ) |
| 541 | use_deterministic_rs = envs.VLLM_BATCH_INVARIANT and world_size > 2 |
| 542 | if use_deterministic_rs: |
| 543 | # Reduce to a fixed root (0) for determinism |
| 544 | reduced = torch.empty_like(input_tensor) |
| 545 | sizes = sizes if sizes else [chunk_size] * world_size |
| 546 | pynccl_comm.reduce(reduced, input_tensor, root=0) |
| 547 | pynccl_comm.scatter(output, reduced, sizes, root=0) |
| 548 | elif sizes is not None and sizes.count(sizes[0]) != len(sizes): |
| 549 | pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes) |
| 550 | else: |
| 551 | pynccl_comm.reduce_scatter(output, input_tensor) |
| 552 | |
| 553 | # Reshape before returning |
| 554 | return output.movedim(0, dim).contiguous() |
| 555 | |
| 556 | def _get_symm_scratch( |
| 557 | self, |
| 558 | role: str, |
| 559 | shape: tuple[int, ...], |
| 560 | dtype: torch.dtype, |
| 561 | device: torch.device, |
| 562 | ) -> torch.Tensor: |
| 563 | """Persistent, pre-registered NCCL symmetric-memory scratch buffer. |
| 564 | |
| 565 | Allocating a fresh symm tensor per collective pays the |
| 566 | ``nccl_symm_mem_context`` snapshot + window-registration scan on every |
| 567 | call (~0.5 ms/RS+AG pair, dwarfing the NVLS transfer itself). Instead we |
| 568 | allocate once per ``(role, shape, dtype)``, register once, and reuse. |
| 569 | |
| 570 | Safe for serial (eager) sequence parallelism: each collective's result |
| 571 | is consumed on the same stream before the next same-role collective |
| 572 | reuses the buffer. Distinct roles (e.g. ``rs_in`` vs ``ag_out``, both |
| 573 | full-size) get distinct buffers so a reduce-scatter input copy never |
| 574 | clobbers a still-live all-gather output. |
| 575 | """ |
| 576 | from vllm.distributed.device_communicators.pynccl_allocator import ( |
| 577 | nccl_symm_mem_context, |
| 578 | ) |
| 579 | |
| 580 | pynccl_comm = self.pynccl_comm |
| 581 | assert pynccl_comm is not None |
| 582 | cache = self.__dict__.setdefault("_symm_scratch_bufs", {}) |
| 583 | key = (role, tuple(shape), dtype) |
| 584 | buf = cache.get(key) |
| 585 | if buf is None: |
| 586 | with nccl_symm_mem_context(pynccl_comm): |
| 587 | buf = torch.empty(shape, dtype=dtype, device=device) |
| 588 | cache[key] = buf |
| 589 | return buf |
| 590 | |
| 591 | def _reduce_scatter_symm_mem( |
| 592 | self, |
| 593 | input_tensor: torch.Tensor, |
| 594 | ) -> torch.Tensor: |
| 595 | """ReduceScatter using NCCL symmetric memory (NVLS). |
| 596 | |
| 597 | Only called for uniform-size reduce_scatter (variable sizes are |
| 598 | guarded out by the caller to avoid asymmetric ncclCommWindowRegister). |
| 599 | Uses persistent pre-registered scratch (see _get_symm_scratch). |
| 600 | """ |
| 601 | from vllm.distributed.device_communicators.pynccl_allocator import ( |
| 602 | is_symmetric_memory_tensor, |
| 603 | ) |
| 604 | |
| 605 | pynccl_comm = self.pynccl_comm |
| 606 | assert pynccl_comm is not None |
| 607 | |
| 608 | chunk = input_tensor.shape[0] // self.world_size |
| 609 | output_shape = (chunk,) + tuple(input_tensor.shape[1:]) |
| 610 | |
| 611 | symm_output = self._get_symm_scratch( |
| 612 | "rs_out", output_shape, input_tensor.dtype, input_tensor.device |
| 613 | ) |
| 614 | # NVLS reduce-scatter (LDMC) requires the input in symmetric memory. |
| 615 | if is_symmetric_memory_tensor(input_tensor): |
| 616 | symm_input = input_tensor |
| 617 | else: |
| 618 | symm_input = self._get_symm_scratch( |
| 619 | "rs_in", |
| 620 | tuple(input_tensor.shape), |
| 621 | input_tensor.dtype, |
| 622 | input_tensor.device, |
| 623 | ) |
| 624 | symm_input.copy_(input_tensor) |
| 625 | |
| 626 | pynccl_comm.reduce_scatter(symm_output, symm_input) |
| 627 | return symm_output |
| 628 | |
| 629 | def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: |
| 630 | """Sends a tensor to the destination rank in a blocking way.""" |
| 631 | """NOTE: `dst` is the local rank of the destination rank.""" |
| 632 | if dst is None: |
| 633 | dst = (self.rank_in_group + 1) % self.world_size |
| 634 | |
| 635 | pynccl_comm = self.pynccl_comm |
| 636 | if pynccl_comm is not None and not pynccl_comm.disabled: |
| 637 | pynccl_comm.send(tensor, dst) |
| 638 | else: |
| 639 | torch.distributed.send(tensor, self.ranks[dst], self.device_group) |
| 640 | |
| 641 | def recv( |
| 642 | self, size: torch.Size, dtype: torch.dtype, src: int | None = None |
| 643 | ) -> torch.Tensor: |
| 644 | """Receives a tensor from the source rank.""" |
| 645 | """NOTE: `src` is the local rank of the source rank.""" |
| 646 | if src is None: |
| 647 | src = (self.rank_in_group - 1) % self.world_size |
| 648 | |
| 649 | tensor = torch.empty(size, dtype=dtype, device=self.device) |
| 650 | pynccl_comm = self.pynccl_comm |
| 651 | if pynccl_comm is not None and not pynccl_comm.disabled: |
| 652 | pynccl_comm.recv(tensor, src) |
| 653 | else: |
| 654 | torch.distributed.recv(tensor, self.ranks[src], self.device_group) |
| 655 | return tensor |
| 656 | |
| 657 | def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor: |
| 658 | """Broadcast a tensor from source rank to all ranks.""" |
| 659 | if self.world_size == 1: |
| 660 | return tensor |
| 661 | |
| 662 | pynccl_comm = self.pynccl_comm |
| 663 | if pynccl_comm is not None and not pynccl_comm.disabled: |
| 664 | pynccl_comm.broadcast(tensor, src) |
| 665 | return tensor |
| 666 | else: |
| 667 | raise ValueError("No PyNCCL communicator found") |
| 668 | |
| 669 | def destroy(self): |
| 670 | if self.pynccl_comm is not None: |
| 671 | self.pynccl_comm.destroy() |
| 672 | self.pynccl_comm = None |
| 673 | if self.ca_comm is not None: |
| 674 | self.ca_comm = None |
| 675 | if self.aiter_ar_comm is not None: |
| 676 | self.aiter_ar_comm.close() |
| 677 | self.aiter_ar_comm = None |
| 678 | if self.fi_ar_comm is not None: |
| 679 | self.fi_ar_comm.destroy() |
| 680 | self.fi_ar_comm = None |
| 681 | if self.fi_pcie_ipc_ar_comm is not None: |
| 682 | self.fi_pcie_ipc_ar_comm.destroy() |
| 683 | self.fi_pcie_ipc_ar_comm = None |
| 684 | if self.all2all_manager is not None: |
| 685 | self.all2all_manager.destroy() |
| 686 | self.all2all_manager = None # type: ignore[assignment] |
| 687 | |
| 688 | def _can_use_aiter_ag_rs(self, sizes: list[int] | None) -> bool: |
| 689 | """Whether the AITER custom AG/RS fast path may run for this collective. |
| 690 | |
| 691 | Requires: |
| 692 | - uniform batches |
| 693 | - FULL CUDAgraphs |
| 694 | """ |
| 695 | if ( |
| 696 | not self.use_aiter_ag_rs |
| 697 | or self.aiter_ar_comm is None |
| 698 | or self.aiter_ar_comm.disabled |
| 699 | ): |
| 700 | return False |
| 701 | if sizes is not None: |
| 702 | return False |
| 703 | |
| 704 | from vllm.config.compilation import CUDAGraphMode |
| 705 | from vllm.forward_context import get_forward_context |
| 706 | |
| 707 | try: |
| 708 | ctx = get_forward_context() |
| 709 | except AssertionError: |
| 710 | return False |
| 711 | if ctx.cudagraph_runtime_mode != CUDAGraphMode.FULL: |
| 712 | return False |
| 713 | bd = ctx.batch_descriptor |
| 714 | return bd is not None and bd.uniform |
| 715 | |
| 716 | def suspend(self) -> None: |
| 717 | if self.pynccl_comm is not None: |
| 718 | self.pynccl_comm.suspend() |
| 719 | |
| 720 | def resume(self) -> None: |
| 721 | if self.pynccl_comm is not None: |
| 722 | self.pynccl_comm.resume() |
| 723 | |
| 724 | def checkpoint_prepare(self) -> None: |
| 725 | # Only FlashInfer all-reduce and FlashInfer all2all are supported for now. |
| 726 | from .flashinfer_all_reduce import checkpoint_prepare_fi_ar_workspaces |
| 727 | |
| 728 | checkpoint_prepare_fi_ar_workspaces(self.cpu_group) |
| 729 | if self.all2all_manager is not None: |
| 730 | self.all2all_manager.checkpoint_prepare() |
| 731 | |
| 732 | def checkpoint_restore(self) -> None: |
| 733 | # Only FlashInfer all-reduce and FlashInfer all2all are supported for now. |
| 734 | from .flashinfer_all_reduce import checkpoint_restore_fi_ar_workspaces |
| 735 | |
| 736 | checkpoint_restore_fi_ar_workspaces(self.cpu_group) |
| 737 | if self.all2all_manager is not None: |
| 738 | self.all2all_manager.checkpoint_restore() |
| 739 | |
| 740 | def all_gatherv( |
| 741 | self, |
| 742 | input_: torch.Tensor | list[torch.Tensor], |
| 743 | dim: int = 0, |
| 744 | sizes: list[int] | None = None, |
| 745 | ): |
| 746 | if dim != 0: |
| 747 | raise NotImplementedError("only dim 0 all-gatherv is supported") |
| 748 | world_size = self.world_size |
| 749 | |
| 750 | # 'sizes' is not needed if all inputs in the same group have the same |
| 751 | # shape |
| 752 | if sizes is not None and all(s == sizes[0] for s in sizes): |
| 753 | sizes = None |
| 754 | |
| 755 | if self._can_use_aiter_ag_rs(sizes): |
| 756 | aiter_comm = self.aiter_ar_comm |
| 757 | assert aiter_comm is not None |
| 758 | if isinstance(input_, torch.Tensor): |
| 759 | if aiter_comm.should_custom_ag(input_): |
| 760 | out = aiter_comm.custom_all_gather(input_, dim=0) |
| 761 | if out is not None: |
| 762 | return out |
| 763 | elif all(aiter_comm.should_custom_ag(inp) for inp in input_): |
| 764 | outs = [aiter_comm.custom_all_gather(inp, dim=0) for inp in input_] |
| 765 | if all(o is not None for o in outs): |
| 766 | return outs |
| 767 | |
| 768 | pynccl_comm = self.pynccl_comm |
| 769 | assert pynccl_comm is not None and not pynccl_comm.disabled |
| 770 | |
| 771 | # Symmetric memory is only used when all ranks have uniform sizes. |
| 772 | # ncclCommWindowRegister is collective: asymmetric pool allocations |
| 773 | # from variable per-rank sizes cause deadlocks. |
| 774 | if sizes is None and should_nccl_symm_mem_ag_rs(): |
| 775 | if isinstance(input_, torch.Tensor): |
| 776 | return self._all_gather_symm_mem(input_) |
| 777 | return self._all_gather_batched_symm_mem(input_) |
| 778 | |
| 779 | def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None): |
| 780 | input_size = input_.size() |
| 781 | if sizes is not None: |
| 782 | assert len(sizes) == world_size |
| 783 | assert input_.shape[dim] == sizes[self.rank_in_group], ( |
| 784 | f"{input_.shape[dim]} != {sizes[self.rank_in_group]}" |
| 785 | ) |
| 786 | output_size = (sum(sizes),) + input_size[1:] |
| 787 | else: |
| 788 | output_size = (input_size[0] * world_size,) + input_size[1:] |
| 789 | # Allocate output tensor. |
| 790 | output_tensor = torch.empty( |
| 791 | output_size, dtype=input_.dtype, device=input_.device |
| 792 | ) |
| 793 | if sizes is not None: |
| 794 | pynccl_comm.all_gatherv(output_tensor, input_, sizes=sizes) |
| 795 | else: |
| 796 | pynccl_comm.all_gather(output_tensor, input_) |
| 797 | return output_tensor |
| 798 | |
| 799 | if isinstance(input_, torch.Tensor): |
| 800 | return _all_gather_single(input_, sizes) |
| 801 | |
| 802 | output_list = [] |
| 803 | pynccl_comm.group_start() |
| 804 | for inp in input_: |
| 805 | output_list.append(_all_gather_single(inp, sizes=sizes)) |
| 806 | pynccl_comm.group_end() |
| 807 | |
| 808 | return output_list |
| 809 | |
| 810 | def _all_gather_symm_mem(self, input_: torch.Tensor) -> torch.Tensor: |
| 811 | """AllGather a single tensor using NCCL symmetric memory (NVLS). |
| 812 | |
| 813 | Only the output needs to be in symmetric memory; NCCL does not |
| 814 | require the AG input to be symmetrically allocated. |
| 815 | """ |
| 816 | pynccl_comm = self.pynccl_comm |
| 817 | assert pynccl_comm is not None |
| 818 | |
| 819 | out_size = (input_.size(0) * self.world_size,) + tuple(input_.size()[1:]) |
| 820 | # Persistent pre-registered scratch avoids the per-call symm-mem context |
| 821 | # snapshot/registration overhead (see _get_symm_scratch). |
| 822 | symm_output = self._get_symm_scratch( |
| 823 | "ag_out", out_size, input_.dtype, input_.device |
| 824 | ) |
| 825 | pynccl_comm.all_gather(symm_output, input_) |
| 826 | return symm_output |
| 827 | |
| 828 | def _all_gather_batched_symm_mem( |
| 829 | self, inputs: list[torch.Tensor] |
| 830 | ) -> list[torch.Tensor]: |
| 831 | """AllGather a list of tensors using NCCL symmetric memory (NVLS). |
| 832 | |
| 833 | Uses group_start/group_end to batch the collectives. |
| 834 | Only the output needs to be in symmetric memory (see |
| 835 | _all_gather_symm_mem). |
| 836 | """ |
| 837 | from vllm.distributed.device_communicators.pynccl_allocator import ( |
| 838 | nccl_symm_mem_context, |
| 839 | ) |
| 840 | |
| 841 | pynccl_comm = self.pynccl_comm |
| 842 | assert pynccl_comm is not None |
| 843 | world_size = self.world_size |
| 844 | |
| 845 | symm_outputs = [] |
| 846 | with nccl_symm_mem_context(pynccl_comm): |
| 847 | for inp in inputs: |
| 848 | out_size = (inp.size(0) * world_size,) + inp.size()[1:] |
| 849 | symm_outputs.append( |
| 850 | torch.empty(out_size, dtype=inp.dtype, device=inp.device) |
| 851 | ) |
| 852 | |
| 853 | pynccl_comm.group_start() |
| 854 | for symm_out, inp in zip(symm_outputs, inputs): |
| 855 | pynccl_comm.all_gather(symm_out, inp) |
| 856 | pynccl_comm.group_end() |
| 857 | |
| 858 | return symm_outputs |
| 859 | |
| 860 | def dispatch_router_logits( |
| 861 | self, |
| 862 | hidden_states: torch.Tensor, |
| 863 | router_logits: torch.Tensor, |
| 864 | is_sequence_parallel: bool = False, |
| 865 | extra_tensors: list[torch.Tensor] | None = None, |
| 866 | ) -> ( |
| 867 | tuple[torch.Tensor, torch.Tensor] |
| 868 | | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 869 | ): |
| 870 | """Dispatch the hidden states and router logits to the appropriate device. |
| 871 | This is a no-op in the base class. |
| 872 | """ |
| 873 | assert self.all2all_manager is not None |
| 874 | return self.all2all_manager.dispatch_router_logits( |
| 875 | hidden_states, |
| 876 | router_logits, |
| 877 | is_sequence_parallel, |
| 878 | extra_tensors, |
| 879 | ) |
| 880 | |
| 881 | def dispatch( |
| 882 | self, |
| 883 | hidden_states: torch.Tensor, |
| 884 | topk_weights: torch.Tensor, |
| 885 | topk_ids: torch.Tensor, |
| 886 | is_sequence_parallel: bool = False, |
| 887 | extra_tensors: list[torch.Tensor] | None = None, |
| 888 | ) -> ( |
| 889 | tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 890 | | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 891 | ): |
| 892 | """Dispatch the hidden states and topk weights/ids to the appropriate device. |
| 893 | This is a no-op in the base class. |
| 894 | """ |
| 895 | assert self.all2all_manager is not None |
| 896 | return self.all2all_manager.dispatch( |
| 897 | hidden_states, |
| 898 | topk_weights, |
| 899 | topk_ids, |
| 900 | is_sequence_parallel, |
| 901 | extra_tensors=extra_tensors, |
| 902 | ) |
| 903 | |
| 904 | def combine( |
| 905 | self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False |
| 906 | ) -> torch.Tensor: |
| 907 | """Combine the hidden states and router logits from the appropriate device. |
| 908 | This is a no-op in the base class. |
| 909 | """ |
| 910 | assert self.all2all_manager is not None |
| 911 | return self.all2all_manager.combine( |
| 912 | hidden_states, |
| 913 | is_sequence_parallel, |
| 914 | ) |
| 915 | |
| 916 | def batch_isend_irecv(self, p2p_ops: list): |
| 917 | pynccl_comm = self.pynccl_comm |
| 918 | if pynccl_comm is not None and not pynccl_comm.disabled: |
| 919 | pynccl_comm.batch_isend_irecv(p2p_ops) |
| 920 | else: |
| 921 | raise ValueError("No PyNCCL communicator found") |