| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from contextlib import contextmanager |
| 5 | from typing import Literal, cast |
| 6 | |
| 7 | import torch |
| 8 | import torch.distributed as dist |
| 9 | from torch.distributed import ProcessGroup |
| 10 | |
| 11 | import vllm.envs as envs |
| 12 | from vllm import _custom_ops as ops |
| 13 | from vllm.distributed.device_communicators.all_reduce_utils import ( |
| 14 | CUSTOM_ALL_REDUCE_MAX_SIZES, |
| 15 | gpu_p2p_access_check, |
| 16 | ) |
| 17 | from vllm.distributed.parallel_state import in_the_same_node_as |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.platforms import current_platform |
| 20 | |
| 21 | try: |
| 22 | ops.meta_size() |
| 23 | custom_ar = True |
| 24 | except Exception: |
| 25 | # For CPUs |
| 26 | custom_ar = False |
| 27 | |
| 28 | try: |
| 29 | import torch.distributed._symmetric_memory as torch_symm_mem |
| 30 | except ImportError: |
| 31 | torch_symm_mem = None |
| 32 | |
| 33 | logger = init_logger(__name__) |
| 34 | |
| 35 | _ReduceScatterBackend = Literal["legacy", "mnnvl_lamport", "mnnvl_multimem"] |
| 36 | _MNNVL_MULTIMEM_REDUCE_SCATTER_WORLD_SIZES = (2, 4, 8) |
| 37 | |
| 38 | |
| 39 | def _has_local_multicast_support(device: torch.device) -> bool: |
| 40 | """Return whether this CUDA device can allocate multicast memory.""" |
| 41 | if torch_symm_mem is None or not current_platform.is_cuda(): |
| 42 | return False |
| 43 | if device.index is None: |
| 44 | return False |
| 45 | |
| 46 | try: |
| 47 | from torch._C._autograd import DeviceType |
| 48 | from torch._C._distributed_c10d import _SymmetricMemory |
| 49 | |
| 50 | return _SymmetricMemory.has_multicast_support( |
| 51 | DeviceType.CUDA, |
| 52 | device.index, |
| 53 | ) |
| 54 | except Exception as error: |
| 55 | logger.debug("MNNVL capability probe failed: %s", error) |
| 56 | return False |
| 57 | |
| 58 | |
| 59 | def _all_ranks_true(group: ProcessGroup, local_value: bool) -> bool: |
| 60 | group_value = torch.tensor(int(local_value), dtype=torch.int32, device="cpu") |
| 61 | dist.all_reduce(group_value, op=dist.ReduceOp.MIN, group=group) |
| 62 | return bool(group_value.item()) |
| 63 | |
| 64 | |
| 65 | def _group_can_attempt_mnnvl( |
| 66 | group: ProcessGroup, |
| 67 | device: torch.device, |
| 68 | ) -> bool: |
| 69 | """Return whether every rank can enter the cross-node MNNVL path. |
| 70 | |
| 71 | MNNVL is available only on Blackwell-class GPUs. Local multicast support |
| 72 | is necessary but does not establish that the process group spans an MNNVL |
| 73 | domain; the symmetric-memory rendezvous below performs that group-level |
| 74 | check. The CPU all-reduce keeps every rank on the same control-flow path |
| 75 | when a heterogeneous or partially configured group is encountered. |
| 76 | """ |
| 77 | device_index = device.index |
| 78 | local_support = ( |
| 79 | device_index is not None |
| 80 | and current_platform.has_device_capability(100, device_index) |
| 81 | and _has_local_multicast_support(device) |
| 82 | ) |
| 83 | return _all_ranks_true(group, local_support) |
| 84 | |
| 85 | |
| 86 | def _can_p2p(rank: int, world_size: int) -> bool: |
| 87 | for i in range(world_size): |
| 88 | if i == rank: |
| 89 | continue |
| 90 | if envs.VLLM_SKIP_P2P_CHECK: |
| 91 | logger.debug("Skipping P2P check and trusting the driver's P2P report.") |
| 92 | # can_device_access_peer takes visible device ordinals, while |
| 93 | # rank and i are logical local IDs. |
| 94 | return torch.cuda.can_device_access_peer( |
| 95 | current_platform.logical_device_id_to_visible_device_id(rank), |
| 96 | current_platform.logical_device_id_to_visible_device_id(i), |
| 97 | ) |
| 98 | if not gpu_p2p_access_check(rank, i): |
| 99 | return False |
| 100 | return True |
| 101 | |
| 102 | |
| 103 | def _supports_mnnvl_multimem_reduce_scatter( |
| 104 | device: torch.device, world_size: int |
| 105 | ) -> bool: |
| 106 | return ( |
| 107 | world_size in _MNNVL_MULTIMEM_REDUCE_SCATTER_WORLD_SIZES |
| 108 | and device.index is not None |
| 109 | and ( |
| 110 | current_platform.is_device_capability((10, 0), device.index) |
| 111 | or current_platform.is_device_capability((10, 3), device.index) |
| 112 | ) |
| 113 | and _has_local_multicast_support(device) |
| 114 | ) |
| 115 | |
| 116 | |
| 117 | from vllm.distributed.utils import is_weak_contiguous # noqa: E402 |
| 118 | |
| 119 | |
| 120 | class CustomAllreduce: |
| 121 | _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8, 16] |
| 122 | batch_invariant = False |
| 123 | _DEFAULT_ALL_GATHER_MAX_SIZE = 2 * 1024 * 1024 |
| 124 | _DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES = { |
| 125 | 2: 8 * 1024 * 1024, |
| 126 | 4: 4 * 1024 * 1024, |
| 127 | 6: 2 * 1024 * 1024, |
| 128 | 8: 2 * 1024 * 1024, |
| 129 | 16: 2 * 1024 * 1024, |
| 130 | } |
| 131 | _DEFAULT_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024 |
| 132 | _DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024 |
| 133 | _DEFAULT_MNNVL_MULTIMEM_REDUCE_SCATTER_MAX_SIZE = 64 * 1024 * 1024 |
| 134 | _MNNVL_MULTIMEM_REDUCE_SCATTER_BLOCKS = 8 |
| 135 | |
| 136 | # max_size: max supported allreduce size |
| 137 | def __init__( |
| 138 | self, |
| 139 | group: ProcessGroup, |
| 140 | device: int | str | torch.device, |
| 141 | max_size=8192 * 1024, |
| 142 | max_all_gather_size=_DEFAULT_ALL_GATHER_MAX_SIZE, |
| 143 | max_mnnvl_all_gather_size=None, |
| 144 | max_reduce_scatter_size=_DEFAULT_REDUCE_SCATTER_MAX_SIZE, |
| 145 | max_mnnvl_reduce_scatter_size=_DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE, |
| 146 | max_mnnvl_multimem_reduce_scatter_size=( |
| 147 | _DEFAULT_MNNVL_MULTIMEM_REDUCE_SCATTER_MAX_SIZE |
| 148 | ), |
| 149 | symm_mem_enabled=False, |
| 150 | ) -> None: |
| 151 | """Args: |
| 152 | group: the process group to work on. If None, it will use the |
| 153 | default process group. |
| 154 | device: the device to bind the CustomAllreduce to. If None, |
| 155 | it will be bound to f"cuda:{local_rank}". |
| 156 | It is the caller's responsibility to make sure each communicator |
| 157 | is bind to a unique device, and all communicators in this group |
| 158 | are in the same node. |
| 159 | |
| 160 | """ |
| 161 | self._IS_CAPTURING = False |
| 162 | self._ptr = 0 |
| 163 | self.disabled = True |
| 164 | self.mnnvl_buffer = None |
| 165 | self.mnnvl_handle = None |
| 166 | self.mnnvl_peer_buffers: list[torch.Tensor] | None = None |
| 167 | self.mnnvl_multicast_ptr = 0 |
| 168 | self.mnnvl_buffer_size = 0 |
| 169 | self.mnnvl_lamport_ag_local_ptr = 0 |
| 170 | self.mnnvl_lamport_ag_multicast_ptr = 0 |
| 171 | self.mnnvl_lamport_rs_local_ptr = 0 |
| 172 | self.mnnvl_lamport_epochs = None |
| 173 | self.mnnvl_lamport_ag_epoch_ptr = 0 |
| 174 | self.mnnvl_lamport_rs_epoch_ptr = 0 |
| 175 | self.mnnvl_multimem_rs_supported = False |
| 176 | self.mnnvl_multimem_rs_initialized = False |
| 177 | self.mnnvl_multimem_rs_buffer = None |
| 178 | self.mnnvl_multimem_rs_handle = None |
| 179 | self.mnnvl_multimem_rs_buffer_size = 0 |
| 180 | self.mnnvl_multimem_rs_local_ptr = 0 |
| 181 | self.mnnvl_multimem_rs_multicast_ptr = 0 |
| 182 | self.mnnvl_only = False |
| 183 | self.batch_invariant = envs.VLLM_BATCH_INVARIANT |
| 184 | |
| 185 | if not custom_ar: |
| 186 | # disable because of missing custom allreduce library |
| 187 | # e.g. in a non-GPU environment |
| 188 | logger.info_once( |
| 189 | "Custom allreduce is disabled because " |
| 190 | "of missing custom allreduce library" |
| 191 | ) |
| 192 | return |
| 193 | |
| 194 | self.group = group |
| 195 | |
| 196 | assert dist.get_backend(group) != dist.Backend.NCCL, ( |
| 197 | "CustomAllreduce should be attached to a non-NCCL group." |
| 198 | ) |
| 199 | |
| 200 | same_node = all(in_the_same_node_as(group, source_rank=0)) |
| 201 | self.mnnvl_only = not same_node |
| 202 | |
| 203 | rank = dist.get_rank(group=self.group) |
| 204 | self.rank = rank |
| 205 | world_size = dist.get_world_size(group=self.group) |
| 206 | if world_size == 1: |
| 207 | # No need to initialize custom allreduce for single GPU case. |
| 208 | return |
| 209 | |
| 210 | if world_size not in CustomAllreduce._SUPPORTED_WORLD_SIZES: |
| 211 | logger.warning_once( |
| 212 | "Custom allreduce is disabled due to an unsupported world" |
| 213 | " size: %d. Supported world sizes: %s. To silence this " |
| 214 | "warning, specify disable_custom_all_reduce=True explicitly.", |
| 215 | world_size, |
| 216 | str(CustomAllreduce._SUPPORTED_WORLD_SIZES), |
| 217 | ) |
| 218 | return |
| 219 | |
| 220 | if isinstance(device, int): |
| 221 | device = torch.device(f"cuda:{device}") |
| 222 | elif isinstance(device, str): |
| 223 | device = torch.device(device) |
| 224 | # now `device` is a `torch.device` object |
| 225 | assert isinstance(device, torch.device) |
| 226 | self.device = device |
| 227 | mnnvl_multimem_rs_supported = _supports_mnnvl_multimem_reduce_scatter( |
| 228 | device, world_size |
| 229 | ) |
| 230 | if not same_node and not _group_can_attempt_mnnvl(group, device): |
| 231 | logger.warning( |
| 232 | "Custom collectives are disabled because this multi-node " |
| 233 | "group does not have MNNVL-capable GPUs on every rank." |
| 234 | ) |
| 235 | return |
| 236 | device_capability = current_platform.get_device_capability() |
| 237 | if ( |
| 238 | current_platform.is_cuda() |
| 239 | and symm_mem_enabled |
| 240 | and device_capability is not None |
| 241 | ): |
| 242 | device_capability_str = device_capability.as_version_str() |
| 243 | if ( |
| 244 | device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES |
| 245 | and world_size in CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str] |
| 246 | ): |
| 247 | max_size = min( |
| 248 | CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size], |
| 249 | max_size, |
| 250 | ) |
| 251 | # device.index is a visible ordinal, not a logical local ID. |
| 252 | fully_connected = False |
| 253 | if same_node: |
| 254 | physical_device_id = ( |
| 255 | current_platform.visible_device_id_to_physical_device_id(device.index) |
| 256 | ) |
| 257 | tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") |
| 258 | gather_list = [ |
| 259 | torch.tensor([0], dtype=torch.int, device="cpu") |
| 260 | for _ in range(world_size) |
| 261 | ] |
| 262 | dist.all_gather(gather_list, tensor, group=self.group) |
| 263 | physical_device_ids = [t.item() for t in gather_list] |
| 264 | assert current_platform.is_cuda_alike() |
| 265 | fully_connected = current_platform.is_fully_connected(physical_device_ids) |
| 266 | if same_node and world_size > 2 and not fully_connected: |
| 267 | logger.warning( |
| 268 | "Custom allreduce is disabled because it's not supported on" |
| 269 | " more than two PCIe-only GPUs. To silence this warning, " |
| 270 | "specify disable_custom_all_reduce=True explicitly." |
| 271 | ) |
| 272 | return |
| 273 | # test P2P capability, this checks software/cudaruntime support |
| 274 | # this is expensive to compute at the first time |
| 275 | # then we cache the result |
| 276 | # On AMD GPU, p2p is always enabled between XGMI connected GPUs |
| 277 | if ( |
| 278 | same_node |
| 279 | and not current_platform.is_rocm() |
| 280 | and not _can_p2p(rank, world_size) |
| 281 | ): |
| 282 | logger.warning( |
| 283 | "Custom allreduce is disabled because your platform lacks " |
| 284 | "GPU P2P capability or P2P test failed. To silence this " |
| 285 | "warning, specify disable_custom_all_reduce=True explicitly." |
| 286 | ) |
| 287 | return |
| 288 | |
| 289 | self.disabled = False |
| 290 | # Buffers memory are owned by this Python class and passed to C++. |
| 291 | # Metadata composes of two parts: metadata for synchronization and a |
| 292 | # temporary buffer for storing intermediate allreduce results. |
| 293 | if same_node: |
| 294 | self.meta_ptrs = self.create_shared_buffer( |
| 295 | ops.meta_size() + max_size, group=group, uncached=True |
| 296 | ) |
| 297 | else: |
| 298 | meta_ptr, _ = ops.allocate_shared_buffer_and_handle(ops.meta_size()) |
| 299 | self.meta_ptrs = [meta_ptr] * world_size |
| 300 | # This is a pre-registered IPC buffer. In eager mode, input tensors |
| 301 | # are first copied into this buffer before the operation is performed |
| 302 | legacy_buffer_size = max(max_size, max_all_gather_size, max_reduce_scatter_size) |
| 303 | if same_node: |
| 304 | self.buffer_ptrs = self.create_shared_buffer( |
| 305 | legacy_buffer_size, |
| 306 | group=group, |
| 307 | ) |
| 308 | else: |
| 309 | buffer_ptr, _ = ops.allocate_shared_buffer_and_handle(legacy_buffer_size) |
| 310 | self.buffer_ptrs = [buffer_ptr] * world_size |
| 311 | # This stores tuples of pointers to IPC buffers from all ranks. |
| 312 | # Each registered tuple contains at most 16 addresses. |
| 313 | # Allocating 8MB is enough for 65536 such tuples. The largest model uses |
| 314 | # fewer than 10000 registered tuples. |
| 315 | self.rank_data = torch.empty( |
| 316 | 8 * 1024 * 1024, dtype=torch.uint8, device=self.device |
| 317 | ) |
| 318 | self.max_size = max_size |
| 319 | self.max_all_gather_size = max_all_gather_size |
| 320 | if max_mnnvl_all_gather_size is None: |
| 321 | max_mnnvl_all_gather_size = self._DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES[ |
| 322 | world_size |
| 323 | ] |
| 324 | self.max_mnnvl_all_gather_size = max_mnnvl_all_gather_size |
| 325 | self.max_reduce_scatter_size = max_reduce_scatter_size |
| 326 | self.max_mnnvl_reduce_scatter_size = max_mnnvl_reduce_scatter_size |
| 327 | self.max_mnnvl_multimem_reduce_scatter_size = ( |
| 328 | max_mnnvl_multimem_reduce_scatter_size |
| 329 | ) |
| 330 | self.rank = rank |
| 331 | self.world_size = world_size |
| 332 | self.fully_connected = fully_connected |
| 333 | self._ptr = ops.init_custom_ar( |
| 334 | self.meta_ptrs, self.rank_data, rank, self.fully_connected |
| 335 | ) |
| 336 | ops.register_buffer(self._ptr, self.buffer_ptrs) |
| 337 | self._init_mnnvl_buffer( |
| 338 | max( |
| 339 | max_mnnvl_all_gather_size * world_size, |
| 340 | max_mnnvl_reduce_scatter_size, |
| 341 | ) |
| 342 | ) |
| 343 | if world_size in _MNNVL_MULTIMEM_REDUCE_SCATTER_WORLD_SIZES: |
| 344 | self.mnnvl_multimem_rs_supported = _all_ranks_true( |
| 345 | self.group, |
| 346 | mnnvl_multimem_rs_supported and bool(self.mnnvl_multicast_ptr), |
| 347 | ) |
| 348 | if not same_node and not self.mnnvl_multicast_ptr: |
| 349 | logger.warning( |
| 350 | "Custom collectives are disabled because this multi-node " |
| 351 | "group does not support MNNVL multicast." |
| 352 | ) |
| 353 | self.close() |
| 354 | self.disabled = True |
| 355 | |
| 356 | def _init_mnnvl_buffer(self, stage_size: int) -> None: |
| 357 | if torch_symm_mem is None or not current_platform.is_cuda(): |
| 358 | return |
| 359 | try: |
| 360 | buffer_size = stage_size * 6 |
| 361 | buffer = torch_symm_mem.empty( |
| 362 | buffer_size, dtype=torch.uint8, device=self.device |
| 363 | ) |
| 364 | handle = torch_symm_mem.rendezvous(buffer, self.group.group_name) |
| 365 | if handle.multicast_ptr == 0: |
| 366 | return |
| 367 | peer_buffers = [ |
| 368 | handle.get_buffer( |
| 369 | peer, |
| 370 | (buffer_size,), |
| 371 | torch.uint8, |
| 372 | storage_offset=0, |
| 373 | ) |
| 374 | for peer in range(self.world_size) |
| 375 | ] |
| 376 | ptrs = [peer_buffer.data_ptr() for peer_buffer in peer_buffers] |
| 377 | lamport_ag_offset = 0 |
| 378 | lamport_rs_offset = stage_size * 3 |
| 379 | lamport_ag_ptrs = [ptr + lamport_ag_offset for ptr in ptrs] |
| 380 | lamport_rs_ptrs = [ptr + lamport_rs_offset for ptr in ptrs] |
| 381 | ops.register_buffer(self._ptr, lamport_ag_ptrs) |
| 382 | ops.register_buffer(self._ptr, lamport_rs_ptrs) |
| 383 | |
| 384 | buffer.view(torch.int32).fill_(-2147483648) |
| 385 | epochs = torch.zeros( |
| 386 | (2, 32), |
| 387 | dtype=torch.int32, |
| 388 | device=self.device, |
| 389 | ) |
| 390 | torch.accelerator.synchronize() |
| 391 | dist.barrier(group=self.group) |
| 392 | |
| 393 | self.mnnvl_buffer = buffer |
| 394 | self.mnnvl_handle = handle |
| 395 | self.mnnvl_peer_buffers = peer_buffers |
| 396 | self.mnnvl_multicast_ptr = handle.multicast_ptr |
| 397 | self.mnnvl_buffer_size = stage_size |
| 398 | self.mnnvl_lamport_ag_local_ptr = lamport_ag_ptrs[self.rank] |
| 399 | self.mnnvl_lamport_ag_multicast_ptr = ( |
| 400 | handle.multicast_ptr + lamport_ag_offset |
| 401 | ) |
| 402 | self.mnnvl_lamport_rs_local_ptr = lamport_rs_ptrs[self.rank] |
| 403 | self.mnnvl_lamport_epochs = epochs |
| 404 | self.mnnvl_lamport_ag_epoch_ptr = epochs[0].data_ptr() |
| 405 | self.mnnvl_lamport_rs_epoch_ptr = epochs[1].data_ptr() |
| 406 | except RuntimeError as error: |
| 407 | logger.debug("MNNVL AG/RS initialization failed: %s", error) |
| 408 | |
| 409 | def _init_mnnvl_multimem_reduce_scatter_buffer(self) -> None: |
| 410 | if self.mnnvl_multimem_rs_initialized: |
| 411 | return |
| 412 | self.mnnvl_multimem_rs_initialized = True |
| 413 | if not self.mnnvl_multimem_rs_supported: |
| 414 | return |
| 415 | |
| 416 | assert torch_symm_mem is not None |
| 417 | buffer = None |
| 418 | signal_size = ops.meta_size() |
| 419 | try: |
| 420 | buffer = torch_symm_mem.empty( |
| 421 | signal_size + self.max_mnnvl_multimem_reduce_scatter_size, |
| 422 | dtype=torch.uint8, |
| 423 | device=self.device, |
| 424 | ) |
| 425 | except RuntimeError as error: |
| 426 | logger.debug("MNNVL multimem RS allocation failed: %s", error) |
| 427 | if not _all_ranks_true(self.group, buffer is not None): |
| 428 | logger.warning_once( |
| 429 | "MNNVL multimem reduce-scatter symmetric-memory allocation " |
| 430 | "failed on at least one rank; falling back to NCCL.", |
| 431 | scope="global", |
| 432 | ) |
| 433 | return |
| 434 | assert buffer is not None |
| 435 | |
| 436 | handle = None |
| 437 | try: |
| 438 | handle = torch_symm_mem.rendezvous(buffer, self.group.group_name) |
| 439 | if handle is not None: |
| 440 | buffer[:signal_size].zero_() |
| 441 | torch.accelerator.synchronize() |
| 442 | except RuntimeError as error: |
| 443 | logger.debug("MNNVL multimem RS rendezvous failed: %s", error) |
| 444 | if not _all_ranks_true( |
| 445 | self.group, |
| 446 | handle is not None and bool(handle.multicast_ptr), |
| 447 | ): |
| 448 | logger.warning_once( |
| 449 | "MNNVL multimem reduce-scatter symmetric-memory rendezvous " |
| 450 | "failed on at least one rank; falling back to NCCL.", |
| 451 | scope="global", |
| 452 | ) |
| 453 | return |
| 454 | assert handle is not None |
| 455 | |
| 456 | self.mnnvl_multimem_rs_buffer = buffer |
| 457 | self.mnnvl_multimem_rs_handle = handle |
| 458 | self.mnnvl_multimem_rs_buffer_size = self.max_mnnvl_multimem_reduce_scatter_size |
| 459 | self.mnnvl_multimem_rs_local_ptr = buffer.data_ptr() + signal_size |
| 460 | self.mnnvl_multimem_rs_multicast_ptr = handle.multicast_ptr + signal_size |
| 461 | |
| 462 | @contextmanager |
| 463 | def capture(self): |
| 464 | """The main responsibility of this context manager is the |
| 465 | `register_graph_buffers` call at the end of the context. |
| 466 | It records all the buffer addresses used in the CUDA graph. |
| 467 | """ |
| 468 | try: |
| 469 | self._IS_CAPTURING = True |
| 470 | yield |
| 471 | finally: |
| 472 | self._IS_CAPTURING = False |
| 473 | if not self.disabled: |
| 474 | self.register_graph_buffers() |
| 475 | |
| 476 | def register_graph_buffers(self): |
| 477 | handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr) |
| 478 | logger.debug("Registering %d cuda graph addresses", len(offset)) |
| 479 | # We cannot directly use `dist.all_gather_object` here |
| 480 | # because it is incompatible with `gloo` backend under inference mode. |
| 481 | # see https://github.com/pytorch/pytorch/issues/126032 for details. |
| 482 | all_data: list[list[list[int] | None]] |
| 483 | all_data = [[None, None] for _ in range(dist.get_world_size(group=self.group))] |
| 484 | all_data[self.rank] = [handle, offset] |
| 485 | ranks = sorted(dist.get_process_group_ranks(group=self.group)) |
| 486 | for i, rank in enumerate(ranks): |
| 487 | dist.broadcast_object_list( |
| 488 | all_data[i], src=rank, group=self.group, device="cpu" |
| 489 | ) |
| 490 | # Unpack list of tuples to tuple of lists. |
| 491 | handles = cast(list[list[int]], [d[0] for d in all_data]) |
| 492 | offsets = cast(list[list[int]], [d[1] for d in all_data]) |
| 493 | ops.register_graph_buffers(self._ptr, handles, offsets) |
| 494 | |
| 495 | def should_custom_ar(self, inp: torch.Tensor): |
| 496 | if self.disabled or self.world_size > 8: |
| 497 | return False |
| 498 | if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16): |
| 499 | return False |
| 500 | inp_size = inp.numel() * inp.element_size() |
| 501 | # custom allreduce requires input byte size to be multiples of 16 |
| 502 | if inp_size % 16 != 0: |
| 503 | return False |
| 504 | if not is_weak_contiguous(inp): |
| 505 | return False |
| 506 | # for 4 or more non NVLink-capable GPUs, custom allreduce provides |
| 507 | # little performance improvement over NCCL. |
| 508 | if self.world_size == 2 or self.fully_connected: |
| 509 | return self.batch_invariant or inp_size < self.max_size |
| 510 | return False |
| 511 | |
| 512 | def all_reduce( |
| 513 | self, inp: torch.Tensor, *, out: torch.Tensor = None, registered: bool = False |
| 514 | ): |
| 515 | """Performs an out-of-place all reduce. |
| 516 | |
| 517 | If registered is True, this assumes inp's pointer is already |
| 518 | IPC-registered. Otherwise, inp is first copied into a pre-registered |
| 519 | buffer. |
| 520 | """ |
| 521 | if out is None: |
| 522 | out = torch.empty_like(inp) |
| 523 | chunk_numel = self.max_size // inp.element_size() |
| 524 | if inp.numel() <= chunk_numel: |
| 525 | self._all_reduce_chunk(inp, out, registered) |
| 526 | return out |
| 527 | # Storage-order views: the gate only admits one contiguous block. |
| 528 | flat_inp = inp.as_strided((inp.numel(),), (1,), inp.storage_offset()) |
| 529 | flat_out = out.as_strided((out.numel(),), (1,), out.storage_offset()) |
| 530 | for start in range(0, flat_inp.numel(), chunk_numel): |
| 531 | end = start + chunk_numel |
| 532 | self._all_reduce_chunk(flat_inp[start:end], flat_out[start:end], registered) |
| 533 | return out |
| 534 | |
| 535 | def _all_reduce_chunk( |
| 536 | self, inp: torch.Tensor, out: torch.Tensor, registered: bool |
| 537 | ) -> None: |
| 538 | if registered: |
| 539 | ops.all_reduce(self._ptr, inp, out, 0, 0) |
| 540 | else: |
| 541 | ops.all_reduce( |
| 542 | self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size |
| 543 | ) |
| 544 | |
| 545 | def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor | None: |
| 546 | """The main allreduce API that provides support for cuda graph.""" |
| 547 | # When custom allreduce is disabled, this will be None. |
| 548 | if self.disabled or not self.should_custom_ar(input): |
| 549 | return None |
| 550 | if self._IS_CAPTURING: |
| 551 | if torch.cuda.is_current_stream_capturing(): |
| 552 | return self.all_reduce(input, registered=True) |
| 553 | else: |
| 554 | # If warm up, mimic the allocation pattern since custom |
| 555 | # allreduce is out-of-place. |
| 556 | return torch.empty_like(input) |
| 557 | else: |
| 558 | # Note: outside of cuda graph context, custom allreduce incurs a |
| 559 | # cost of cudaMemcpy, which should be small (<=1% of overall |
| 560 | # latency) compared to the performance gain of using custom kernels |
| 561 | return self.all_reduce(input, registered=False) |
| 562 | |
| 563 | def should_custom_all_gather(self, inp: torch.Tensor) -> bool: |
| 564 | if self.disabled or not current_platform.is_cuda(): |
| 565 | return False |
| 566 | if self.world_size == 16 and not self.mnnvl_only: |
| 567 | return False |
| 568 | inp_size = inp.nbytes |
| 569 | if inp.dtype not in ( |
| 570 | torch.float32, |
| 571 | torch.float16, |
| 572 | torch.bfloat16, |
| 573 | ): |
| 574 | return False |
| 575 | max_size = ( |
| 576 | self.max_mnnvl_all_gather_size |
| 577 | if self.mnnvl_multicast_ptr |
| 578 | else self.max_all_gather_size |
| 579 | ) |
| 580 | return ( |
| 581 | 0 < inp_size <= max_size |
| 582 | and inp_size % 16 == 0 |
| 583 | and is_weak_contiguous(inp) |
| 584 | and (self.fully_connected or bool(self.mnnvl_multicast_ptr)) |
| 585 | ) |
| 586 | |
| 587 | def custom_all_gather(self, inp: torch.Tensor) -> torch.Tensor | None: |
| 588 | if not self.should_custom_all_gather(inp): |
| 589 | return None |
| 590 | out_shape = (inp.shape[0] * self.world_size,) + inp.shape[1:] |
| 591 | if self.mnnvl_multicast_ptr: |
| 592 | logger.info_once( |
| 593 | "Using the MNNVL Lamport all-gather kernel.", |
| 594 | scope="global", |
| 595 | ) |
| 596 | out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device) |
| 597 | ops.mnnvl_lamport_all_gather( |
| 598 | self._ptr, |
| 599 | inp, |
| 600 | out, |
| 601 | self.mnnvl_lamport_ag_local_ptr, |
| 602 | self.mnnvl_lamport_ag_multicast_ptr, |
| 603 | self.mnnvl_lamport_ag_epoch_ptr, |
| 604 | self.mnnvl_buffer_size, |
| 605 | ) |
| 606 | else: |
| 607 | out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device) |
| 608 | ops.custom_all_gather( |
| 609 | self._ptr, |
| 610 | inp, |
| 611 | out, |
| 612 | self.buffer_ptrs[self.rank], |
| 613 | self.max_all_gather_size, |
| 614 | ) |
| 615 | return out |
| 616 | |
| 617 | def _select_reduce_scatter_backend( |
| 618 | self, inp: torch.Tensor |
| 619 | ) -> _ReduceScatterBackend | None: |
| 620 | if self.disabled or not current_platform.is_cuda(): |
| 621 | return None |
| 622 | # The size gates below would switch backends per batch. |
| 623 | if self.batch_invariant: |
| 624 | return None |
| 625 | if self.world_size == 16 and not self.mnnvl_only: |
| 626 | return None |
| 627 | inp_size = inp.nbytes |
| 628 | if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16): |
| 629 | return None |
| 630 | if inp.shape[0] % self.world_size != 0: |
| 631 | return None |
| 632 | output_size = inp_size // self.world_size |
| 633 | if inp_size <= 0 or output_size % 16 != 0 or not is_weak_contiguous(inp): |
| 634 | return None |
| 635 | |
| 636 | if self.mnnvl_multicast_ptr: |
| 637 | if inp_size <= self.max_mnnvl_reduce_scatter_size: |
| 638 | return "mnnvl_lamport" |
| 639 | if ( |
| 640 | ( |
| 641 | self.mnnvl_multimem_rs_multicast_ptr |
| 642 | or ( |
| 643 | self.mnnvl_multimem_rs_supported |
| 644 | and not self.mnnvl_multimem_rs_initialized |
| 645 | ) |
| 646 | ) |
| 647 | and not envs.VLLM_BATCH_INVARIANT |
| 648 | and inp_size <= self.max_mnnvl_multimem_reduce_scatter_size |
| 649 | ): |
| 650 | return "mnnvl_multimem" |
| 651 | return None |
| 652 | |
| 653 | if self.fully_connected and inp_size <= self.max_reduce_scatter_size: |
| 654 | return "legacy" |
| 655 | return None |
| 656 | |
| 657 | def should_custom_reduce_scatter(self, inp: torch.Tensor) -> bool: |
| 658 | return self._select_reduce_scatter_backend(inp) is not None |
| 659 | |
| 660 | def should_mnnvl_multimem_reduce_scatter(self, inp: torch.Tensor) -> bool: |
| 661 | return self._select_reduce_scatter_backend(inp) == "mnnvl_multimem" |
| 662 | |
| 663 | def custom_reduce_scatter(self, inp: torch.Tensor) -> torch.Tensor | None: |
| 664 | backend = self._select_reduce_scatter_backend(inp) |
| 665 | if backend is None: |
| 666 | return None |
| 667 | if backend == "mnnvl_multimem" and not self.mnnvl_multimem_rs_multicast_ptr: |
| 668 | if self._IS_CAPTURING: |
| 669 | return None |
| 670 | self._init_mnnvl_multimem_reduce_scatter_buffer() |
| 671 | if not self.mnnvl_multimem_rs_multicast_ptr: |
| 672 | return None |
| 673 | |
| 674 | out_shape = (inp.shape[0] // self.world_size,) + inp.shape[1:] |
| 675 | out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device) |
| 676 | if backend == "mnnvl_multimem": |
| 677 | logger.info_once( |
| 678 | "Using the low-SM MNNVL multimem reduce-scatter kernel.", |
| 679 | scope="global", |
| 680 | ) |
| 681 | ops.mnnvl_multimem_reduce_scatter( |
| 682 | self._ptr, |
| 683 | inp, |
| 684 | out, |
| 685 | self.mnnvl_multimem_rs_local_ptr, |
| 686 | self.mnnvl_multimem_rs_multicast_ptr, |
| 687 | self.mnnvl_multimem_rs_buffer_size, |
| 688 | self._MNNVL_MULTIMEM_REDUCE_SCATTER_BLOCKS, |
| 689 | ) |
| 690 | elif backend == "mnnvl_lamport": |
| 691 | logger.info_once( |
| 692 | "Using the MNNVL Lamport reduce-scatter kernel.", |
| 693 | scope="global", |
| 694 | ) |
| 695 | ops.mnnvl_lamport_reduce_scatter( |
| 696 | self._ptr, |
| 697 | inp, |
| 698 | out, |
| 699 | self.mnnvl_lamport_rs_local_ptr, |
| 700 | self.mnnvl_lamport_rs_epoch_ptr, |
| 701 | self.mnnvl_buffer_size, |
| 702 | ) |
| 703 | else: |
| 704 | ops.custom_reduce_scatter( |
| 705 | self._ptr, |
| 706 | inp, |
| 707 | out, |
| 708 | self.buffer_ptrs[self.rank], |
| 709 | self.max_reduce_scatter_size, |
| 710 | ) |
| 711 | return out |
| 712 | |
| 713 | def close(self): |
| 714 | if not self.disabled and self._ptr: |
| 715 | if ops is not None: |
| 716 | ops.dispose(self._ptr) |
| 717 | self._ptr = 0 |
| 718 | self.free_shared_buffer(self.meta_ptrs, rank=self.rank) |
| 719 | self.free_shared_buffer(self.buffer_ptrs, rank=self.rank) |
| 720 | self.mnnvl_peer_buffers = None |
| 721 | self.mnnvl_handle = None |
| 722 | self.mnnvl_buffer = None |
| 723 | self.mnnvl_lamport_epochs = None |
| 724 | self.mnnvl_multimem_rs_handle = None |
| 725 | self.mnnvl_multimem_rs_buffer = None |
| 726 | |
| 727 | def __del__(self): |
| 728 | self.close() |
| 729 | |
| 730 | @staticmethod |
| 731 | def create_shared_buffer( |
| 732 | size_in_bytes: int, |
| 733 | group: ProcessGroup | None = None, |
| 734 | uncached: bool | None = False, |
| 735 | ) -> list[int]: |
| 736 | pointer, handle = ops.allocate_shared_buffer_and_handle(size_in_bytes) |
| 737 | |
| 738 | world_size = dist.get_world_size(group=group) |
| 739 | rank = dist.get_rank(group=group) |
| 740 | handles = [None] * world_size |
| 741 | dist.all_gather_object(handles, handle, group=group) |
| 742 | |
| 743 | pointers: list[int] = [] |
| 744 | for i, h in enumerate(handles): |
| 745 | if i == rank: |
| 746 | pointers.append(pointer) # type: ignore |
| 747 | else: |
| 748 | pointers.append(ops.open_mem_handle(h)) |
| 749 | return pointers |
| 750 | |
| 751 | @staticmethod |
| 752 | def free_shared_buffer( |
| 753 | pointers: list[int], |
| 754 | group: ProcessGroup | None = None, |
| 755 | rank: int | None = None, |
| 756 | ) -> None: |
| 757 | if rank is None: |
| 758 | rank = dist.get_rank(group=group) |
| 759 | if ops is not None: |
| 760 | ops.free_shared_buffer(pointers[rank]) |