| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import atexit |
| 4 | import contextlib |
| 5 | import logging |
| 6 | import tempfile |
| 7 | from typing import Any |
| 8 | |
| 9 | import torch |
| 10 | from packaging import version |
| 11 | from torch.cuda.memory import CUDAPluggableAllocator |
| 12 | from torch.utils.cpp_extension import load_inline |
| 13 | |
| 14 | from vllm import envs |
| 15 | from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.platforms import current_platform |
| 18 | from vllm.utils.nccl import find_nccl_include_paths, find_nccl_library_paths |
| 19 | |
| 20 | logger = init_logger(__name__) |
| 21 | |
| 22 | nccl_allocator_source = """ |
| 23 | #include <nccl.h> |
| 24 | extern "C" { |
| 25 | |
| 26 | void* nccl_alloc_plug(size_t size, int device, void* stream) { |
| 27 | void* ptr; |
| 28 | ncclResult_t err = ncclMemAlloc(&ptr, size); |
| 29 | return ptr; |
| 30 | |
| 31 | } |
| 32 | |
| 33 | void nccl_free_plug(void* ptr, size_t size, int device, void* stream) { |
| 34 | ncclResult_t err = ncclMemFree(ptr); |
| 35 | } |
| 36 | |
| 37 | } |
| 38 | """ |
| 39 | |
| 40 | _allocator = None |
| 41 | _allocator_wrapper = None |
| 42 | _mem_pool = None |
| 43 | _registered_base_addrs: dict[bytes, set] = {} |
| 44 | _graph_pool_id = None |
| 45 | _nccl_allocator_failed_to_compile = False |
| 46 | _cached_pool_snapshot = None |
| 47 | |
| 48 | |
| 49 | def is_symmetric_memory_enabled(): |
| 50 | global _nccl_allocator_failed_to_compile |
| 51 | return envs.VLLM_USE_NCCL_SYMM_MEM and not _nccl_allocator_failed_to_compile |
| 52 | |
| 53 | |
| 54 | def is_symmetric_memory_tensor(tensor: torch.Tensor): |
| 55 | if not is_symmetric_memory_enabled() or _cached_pool_snapshot is None: |
| 56 | return False |
| 57 | for segment in _cached_pool_snapshot: |
| 58 | for block in segment["blocks"]: |
| 59 | if block["address"] == tensor.untyped_storage().data_ptr(): |
| 60 | return True |
| 61 | return False |
| 62 | |
| 63 | |
| 64 | def set_graph_pool_id(graph_pool_id: Any) -> None: |
| 65 | global _graph_pool_id |
| 66 | _graph_pool_id = graph_pool_id |
| 67 | |
| 68 | |
| 69 | def compile_nccl_allocator(): |
| 70 | global _allocator, _allocator_wrapper, _nccl_allocator_failed_to_compile |
| 71 | if not current_platform.is_cuda(): |
| 72 | _nccl_allocator_failed_to_compile = True |
| 73 | return |
| 74 | try: |
| 75 | out_dir = tempfile.gettempdir() |
| 76 | nccl_allocator_libname = "nccl_allocator" |
| 77 | nccl_include_paths = find_nccl_include_paths() |
| 78 | ldflags = ["-l:libnccl.so.2"] |
| 79 | nccl_lib_paths = find_nccl_library_paths() |
| 80 | if nccl_lib_paths: |
| 81 | ldflags = [f"-L{p}" for p in nccl_lib_paths] + ldflags |
| 82 | load_inline( |
| 83 | name=nccl_allocator_libname, |
| 84 | cpp_sources=nccl_allocator_source, |
| 85 | with_cuda=True, |
| 86 | extra_ldflags=ldflags, |
| 87 | verbose=logger.isEnabledFor(logging.DEBUG), |
| 88 | is_python_module=False, |
| 89 | build_directory=out_dir, |
| 90 | extra_include_paths=nccl_include_paths, |
| 91 | ) |
| 92 | _allocator_wrapper = CUDAPluggableAllocator( |
| 93 | f"{out_dir}/{nccl_allocator_libname}.so", |
| 94 | "nccl_alloc_plug", |
| 95 | "nccl_free_plug", |
| 96 | ) |
| 97 | _allocator = _allocator_wrapper.allocator() |
| 98 | except Exception as e: |
| 99 | _nccl_allocator_failed_to_compile = True |
| 100 | logger.warning( |
| 101 | "Failed to compile NCCL memory allocator. " |
| 102 | "Symmetric memory will be disabled. " |
| 103 | "This is expected if NCCL headers are not available. " |
| 104 | "optionally set VLLM_NCCL_INCLUDE_PATH to point to a directory " |
| 105 | "containing the NCCL header. " |
| 106 | "Error: %s", |
| 107 | str(e), |
| 108 | ) |
| 109 | |
| 110 | |
| 111 | def get_nccl_mem_pool(): |
| 112 | global _mem_pool, _nccl_allocator_failed_to_compile |
| 113 | if _mem_pool is None and not _nccl_allocator_failed_to_compile: |
| 114 | compile_nccl_allocator() |
| 115 | if _allocator is not None: |
| 116 | _mem_pool = torch.cuda.MemPool(_allocator) |
| 117 | return _mem_pool |
| 118 | |
| 119 | |
| 120 | def _cleanup_nccl_mem_pool(): |
| 121 | global _mem_pool |
| 122 | _mem_pool = None |
| 123 | |
| 124 | |
| 125 | def _cleanup_nccl_allocator_wrapper(): |
| 126 | global _allocator_wrapper |
| 127 | _allocator_wrapper = None |
| 128 | |
| 129 | |
| 130 | atexit.register(_cleanup_nccl_mem_pool) |
| 131 | atexit.register(_cleanup_nccl_allocator_wrapper) |
| 132 | |
| 133 | |
| 134 | class nccl_symm_mem_context: |
| 135 | def __init__( |
| 136 | self, |
| 137 | pynccl_comm: PyNcclCommunicator, |
| 138 | disabled: bool = False, |
| 139 | ): |
| 140 | self.disabled = ( |
| 141 | disabled |
| 142 | or not is_symmetric_memory_enabled() |
| 143 | or pynccl_comm.world_size == 1 |
| 144 | or not current_platform.is_cuda() |
| 145 | or get_nccl_mem_pool() is None |
| 146 | or version.parse(torch.__version__) < version.parse("2.8.0.a0") |
| 147 | ) |
| 148 | if self.disabled: |
| 149 | self.pynccl_comm: PyNcclCommunicator | None = None |
| 150 | self._mem_pool_ctx: contextlib.AbstractContextManager[Any] = ( |
| 151 | contextlib.nullcontext() |
| 152 | ) |
| 153 | self.is_graph_capture = None |
| 154 | self.device = None |
| 155 | else: |
| 156 | self.pynccl_comm = pynccl_comm |
| 157 | self._mem_pool_ctx = torch.cuda.use_mem_pool(get_nccl_mem_pool()) |
| 158 | self.is_graph_capture = torch.cuda.is_current_stream_capturing() |
| 159 | self.device = torch.accelerator.current_device_index() |
| 160 | |
| 161 | def __enter__(self): |
| 162 | if self.disabled: |
| 163 | return self |
| 164 | assert self.pynccl_comm is not None, ( |
| 165 | "Symmetric memory requires pynccl to be initialized" |
| 166 | ) |
| 167 | assert self.pynccl_comm.nccl_version >= 22703, ( |
| 168 | "NCCL version 2.27.3 or higher is required for NCCL symmetric memory" |
| 169 | ) |
| 170 | if self.is_graph_capture: |
| 171 | assert _graph_pool_id is not None, ( |
| 172 | "graph_pool_id is not set under graph capture" |
| 173 | ) |
| 174 | # Pause graph memory pool to use symmetric memory with cuda graph |
| 175 | torch._C._cuda_endAllocateToPool(self.device, _graph_pool_id) |
| 176 | self._mem_pool_ctx.__enter__() |
| 177 | return self |
| 178 | |
| 179 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 180 | if self.disabled: |
| 181 | return |
| 182 | global _cached_pool_snapshot |
| 183 | global _registered_base_addrs |
| 184 | self._mem_pool_ctx.__exit__(exc_type, exc_val, exc_tb) |
| 185 | _pool = get_nccl_mem_pool() |
| 186 | assert _pool is not None |
| 187 | _cached_pool_snapshot = _pool.snapshot() |
| 188 | assert self.pynccl_comm is not None |
| 189 | comm_key = bytes(self.pynccl_comm.unique_id.internal) |
| 190 | if comm_key not in _registered_base_addrs: |
| 191 | _registered_base_addrs[comm_key] = set() |
| 192 | for segment in _cached_pool_snapshot: |
| 193 | if segment["address"] not in _registered_base_addrs[comm_key]: |
| 194 | self.pynccl_comm.register_comm_window_raw( |
| 195 | segment["address"], segment["total_size"] |
| 196 | ) |
| 197 | _registered_base_addrs[comm_key].add(segment["address"]) |
| 198 | if self.is_graph_capture: |
| 199 | torch._C._cuda_beginAllocateCurrentThreadToPool(self.device, _graph_pool_id) |