| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared symmetric-memory infrastructure for context-parallel attention.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import functools |
| 8 | from typing import TYPE_CHECKING, Any |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm.distributed.parallel_state import in_the_same_node_as |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.platforms import current_platform |
| 15 | |
| 16 | if TYPE_CHECKING: |
| 17 | from torch.distributed import ProcessGroup |
| 18 | |
| 19 | from vllm.distributed.parallel_state import GroupCoordinator |
| 20 | |
| 21 | logger = init_logger(__name__) |
| 22 | |
| 23 | try: |
| 24 | import torch.distributed._symmetric_memory as symm_mem |
| 25 | |
| 26 | symm_mem_available = True |
| 27 | except ImportError: |
| 28 | symm_mem = None # type: ignore[assignment] |
| 29 | symm_mem_available = False |
| 30 | |
| 31 | |
| 32 | @functools.cache |
| 33 | def _symm_mem_spans_group(group: GroupCoordinator) -> bool: |
| 34 | """Probe whether the group has NVLS symmetric memory.""" |
| 35 | if not symm_mem_available: |
| 36 | return False |
| 37 | try: |
| 38 | from torch._C._autograd import DeviceType |
| 39 | from torch._C._distributed_c10d import _SymmetricMemory |
| 40 | |
| 41 | device = torch.device("cuda", torch.accelerator.current_device_index()) |
| 42 | if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device.index): |
| 43 | return False |
| 44 | probe = symm_mem.empty(8, dtype=torch.uint8, device=device) |
| 45 | probe.zero_() |
| 46 | torch.accelerator.synchronize() |
| 47 | handle = symm_mem.rendezvous(probe, group.device_group.group_name) |
| 48 | spans = handle is not None and handle.multicast_ptr != 0 |
| 49 | except Exception as error: |
| 50 | logger.debug("Direct CP symmetric-memory probe failed: %s", error) |
| 51 | return False |
| 52 | logger.debug_once( |
| 53 | "Direct CP symmetric memory across %d ranks: %s", |
| 54 | group.world_size, |
| 55 | "available" if spans else "unavailable", |
| 56 | ) |
| 57 | return spans |
| 58 | |
| 59 | |
| 60 | def direct_cp_enabled( |
| 61 | group: GroupCoordinator, |
| 62 | dtype: torch.dtype, |
| 63 | use_direct: bool | None, |
| 64 | supported_dtypes: tuple[torch.dtype, ...] | None = None, |
| 65 | ) -> bool: |
| 66 | if use_direct is not None: |
| 67 | return use_direct |
| 68 | return ( |
| 69 | symm_mem_available |
| 70 | and current_platform.is_cuda() |
| 71 | and (supported_dtypes is None or dtype in supported_dtypes) |
| 72 | and ( |
| 73 | all(in_the_same_node_as(group.cpu_group, source_rank=0)) |
| 74 | or _symm_mem_spans_group(group) |
| 75 | ) |
| 76 | ) |
| 77 | |
| 78 | |
| 79 | def direct_cp_multicast_enabled( |
| 80 | group: GroupCoordinator, |
| 81 | dtype: torch.dtype, |
| 82 | use_direct: bool | None, |
| 83 | supported_dtypes: tuple[torch.dtype, ...] | None = None, |
| 84 | ) -> bool: |
| 85 | return direct_cp_enabled( |
| 86 | group, dtype, use_direct, supported_dtypes |
| 87 | ) and _symm_mem_spans_group(group) |
| 88 | |
| 89 | |
| 90 | class DirectCPWorkspace: |
| 91 | def __init__( |
| 92 | self, |
| 93 | group: ProcessGroup, |
| 94 | device: torch.device, |
| 95 | num_ubatches: int, |
| 96 | ) -> None: |
| 97 | self.group = group |
| 98 | self.world_size = group.size() |
| 99 | self.rank = group.rank() |
| 100 | self.device = torch.device(device) |
| 101 | self.num_ubatches = num_ubatches |
| 102 | self.epoch = torch.zeros(num_ubatches, dtype=torch.int64, device=self.device) |
| 103 | self._allocations: list[tuple[torch.Tensor, Any, list[torch.Tensor]]] = [] |
| 104 | |
| 105 | def _allocate( |
| 106 | self, shape: tuple[int, ...], dtype: torch.dtype |
| 107 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 108 | storage = symm_mem.empty(shape, device=self.device, dtype=dtype) |
| 109 | storage.zero_() |
| 110 | torch.accelerator.synchronize() |
| 111 | handle = symm_mem.rendezvous(storage, self.group.group_name) |
| 112 | assert handle is not None, "CP symmetric memory rendezvous returned None" |
| 113 | handle.barrier() |
| 114 | views = [ |
| 115 | handle.get_buffer(peer, list(shape), dtype, 0) |
| 116 | for peer in range(self.world_size) |
| 117 | ] |
| 118 | self.device = storage.device |
| 119 | peer_ptrs = torch.tensor( |
| 120 | [ |
| 121 | [view[ubatch].data_ptr() for view in views] |
| 122 | for ubatch in range(self.num_ubatches) |
| 123 | ], |
| 124 | dtype=torch.int64, |
| 125 | device=self.device, |
| 126 | ) |
| 127 | self._allocations.append((storage, handle, views)) |
| 128 | return storage, peer_ptrs |
| 129 | |
| 130 | def _multicast_ptrs(self, storage: torch.Tensor) -> list[int]: |
| 131 | disabled = [0] * self.num_ubatches |
| 132 | for allocated, handle, _ in self._allocations: |
| 133 | if allocated is storage: |
| 134 | break |
| 135 | else: |
| 136 | return disabled |
| 137 | try: |
| 138 | from torch._C._autograd import DeviceType |
| 139 | from torch._C._distributed_c10d import _SymmetricMemory |
| 140 | |
| 141 | if not _SymmetricMemory.has_multicast_support( |
| 142 | DeviceType.CUDA, storage.device.index |
| 143 | ): |
| 144 | return disabled |
| 145 | multicast_base = handle.multicast_ptr |
| 146 | except Exception: |
| 147 | return disabled |
| 148 | if not multicast_base: |
| 149 | return disabled |
| 150 | storage_base = storage.data_ptr() |
| 151 | return [ |
| 152 | multicast_base + (storage[ubatch].data_ptr() - storage_base) |
| 153 | for ubatch in range(self.num_ubatches) |
| 154 | ] |