| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import ctypes |
| 5 | import json |
| 6 | import os |
| 7 | import pickle |
| 8 | import subprocess |
| 9 | import sys |
| 10 | import tempfile |
| 11 | from collections.abc import Sequence |
| 12 | from itertools import product |
| 13 | from typing import Any |
| 14 | |
| 15 | import torch |
| 16 | import torch.distributed as dist |
| 17 | import torch.multiprocessing as mp |
| 18 | |
| 19 | import vllm.envs as envs |
| 20 | from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary |
| 21 | from vllm.logger import init_logger |
| 22 | from vllm.platforms import current_platform |
| 23 | from vllm.utils.system_utils import update_environment_variables |
| 24 | |
| 25 | logger = init_logger(__name__) |
| 26 | |
| 27 | KiB = 1024 |
| 28 | MiB = 1024 * 1024 |
| 29 | # Max size for each world size in case symmetric memory is available |
| 30 | # For different SM architectures |
| 31 | CUSTOM_ALL_REDUCE_MAX_SIZES = { |
| 32 | "9.0": { |
| 33 | 2: 64 * MiB, # 64 MB |
| 34 | 4: 32 * MiB, # 32 MB |
| 35 | 6: MiB // 2, # 512 KB |
| 36 | 8: MiB // 4, # 256 KB |
| 37 | }, |
| 38 | "10.0": { |
| 39 | 2: 2 * MiB, # 2 MB |
| 40 | 4: 2 * MiB, # 2 MB |
| 41 | 6: 1 * MiB, # 1 MB |
| 42 | 8: 1 * MiB, # 1 MB |
| 43 | }, |
| 44 | "10.3": { |
| 45 | 2: 4 * MiB, # 4 MB |
| 46 | 4: 4 * MiB, # 4 MB |
| 47 | 6: 8 * MiB, # 8 MB |
| 48 | 8: 4 * MiB, # 4 MB |
| 49 | }, |
| 50 | "10.7": { # sm_107 (Rubin): reuse 10.3 all-reduce thresholds |
| 51 | 2: 4 * MiB, # 4 MB |
| 52 | 4: 4 * MiB, # 4 MB |
| 53 | 6: 8 * MiB, # 8 MB |
| 54 | 8: 4 * MiB, # 4 MB |
| 55 | }, |
| 56 | } |
| 57 | |
| 58 | SYMM_MEM_ALL_REDUCE_MAX_SIZES = { |
| 59 | "9.0": { |
| 60 | 2: 64 * MiB, # 64 MB |
| 61 | 4: 32 * MiB, # 32 MB |
| 62 | 6: 64 * MiB, # 64 MB |
| 63 | 8: 64 * MiB, # 64 MB |
| 64 | }, |
| 65 | "10.0": { |
| 66 | 2: 8 * MiB, # 8 MB |
| 67 | 4: 32 * MiB, # 32 MB |
| 68 | 6: 128 * MiB, # 128 MB |
| 69 | 8: 128 * MiB, # 128 MB |
| 70 | }, |
| 71 | "10.3": { |
| 72 | 2: 4 * MiB, # 4 MB |
| 73 | 4: 32 * MiB, # 32 MB |
| 74 | 6: 32 * MiB, # 32 MB |
| 75 | 8: 64 * MiB, # 64 MB |
| 76 | }, |
| 77 | "10.7": { # sm_107 (Rubin): reuse 10.3 all-reduce thresholds |
| 78 | 2: 4 * MiB, # 4 MB |
| 79 | 4: 32 * MiB, # 32 MB |
| 80 | 6: 32 * MiB, # 32 MB |
| 81 | 8: 64 * MiB, # 64 MB |
| 82 | }, |
| 83 | } |
| 84 | |
| 85 | # Per-rank input limit for standalone FlashInfer MNNVL all-reduce. |
| 86 | # The key is (compute capability, world size, node count). |
| 87 | FI_MNNVL_ALLREDUCE_MAX_SIZE_MB: dict[tuple[int, int, int], float] = { |
| 88 | (103, 4, 1): 80, |
| 89 | (103, 8, 1): 64, |
| 90 | (103, 8, 2): 64, |
| 91 | (103, 16, 4): 8, |
| 92 | } |
| 93 | |
| 94 | # NCCL symmetric memory allreduce configuration based on H100 and GB200 benchmarks. |
| 95 | # PyNCCL-symm outperforms custom_AR for small and large tensor sizes, |
| 96 | # while custom_AR wins for mid-range sizes. |
| 97 | # |
| 98 | # Benchmark results (8 GPUs): |
| 99 | # 2K - 16K: PyNCCL-symm wins (1.35x - 1.48x faster) |
| 100 | # 32K - 64K: custom_AR wins |
| 101 | # 128K - 1G: PyNCCL-symm wins (1.12x - 6.14x faster) |
| 102 | # |
| 103 | # Benchmark results (4 GPUs): |
| 104 | # 2K - 16K: PyNCCL-symm wins (1.21x - 1.30x faster) |
| 105 | # 32K - 256K: custom_AR wins (1.07x - 1.35x faster) |
| 106 | # 512K - 1G: PyNCCL-symm wins (1.10x - 2.32x faster) |
| 107 | # |
| 108 | # The config defines ranges where custom_AR is preferred (symm_mem disabled). |
| 109 | NCCL_SYMM_MEM_ALL_REDUCE_CONFIG: dict[str, Any] = { |
| 110 | "min_world_size": 4, |
| 111 | # Ranges where custom_AR outperforms NCCL symm_mem: (lower_bound, upper_bound) |
| 112 | # NCCL symm_mem will NOT be used for sizes in range: lower < size < upper |
| 113 | "custom_ar_preferred_ranges": { |
| 114 | 4: (16 * KiB, 512 * KiB), # custom_AR wins for 32K-256K |
| 115 | 8: (16 * KiB, 128 * KiB), # custom_AR wins for 32K-64K |
| 116 | }, |
| 117 | "always_use_above_world_size": 8, # Always use symm mem for world_size > 8 |
| 118 | } |
| 119 | |
| 120 | |
| 121 | def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor) -> bool: |
| 122 | """Determine if NCCL symmetric memory allreduce should be used. |
| 123 | |
| 124 | Based on H100 and GB200 benchmarks, NCCL symm_mem is preferred for: |
| 125 | - Small tensors (≤16K): Lower latency than custom_AR |
| 126 | - Large tensors (≥128K for 8 GPUs, ≥512K for 4 GPUs): Better bandwidth |
| 127 | |
| 128 | Custom_AR is preferred for mid-range sizes where its P2P approach |
| 129 | has lower overhead than the symm_mem copy-in/copy-out pattern. |
| 130 | """ |
| 131 | from vllm.distributed.device_communicators.pynccl_allocator import ( |
| 132 | is_symmetric_memory_enabled, |
| 133 | ) |
| 134 | |
| 135 | if envs.VLLM_BATCH_INVARIANT: |
| 136 | return False |
| 137 | |
| 138 | if not is_symmetric_memory_enabled(): |
| 139 | return False |
| 140 | |
| 141 | if world_size < NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["min_world_size"]: |
| 142 | return False |
| 143 | |
| 144 | tensor_size = input_tensor.nbytes |
| 145 | custom_ar_range = NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["custom_ar_preferred_ranges"].get( |
| 146 | world_size |
| 147 | ) |
| 148 | |
| 149 | if custom_ar_range is not None: |
| 150 | lower_bound, upper_bound = custom_ar_range |
| 151 | # Use symm_mem for small sizes (≤ lower_bound) and large sizes (≥ upper_bound) |
| 152 | # Use custom_AR (not symm_mem) for mid-range sizes |
| 153 | return tensor_size <= lower_bound or tensor_size >= upper_bound |
| 154 | return world_size > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"] |
| 155 | |
| 156 | |
| 157 | def should_nccl_symm_mem_ag_rs() -> bool: |
| 158 | """Check whether NCCL symmetric memory should be used for |
| 159 | AllGather / ReduceScatter collectives.""" |
| 160 | from vllm.distributed.device_communicators.pynccl_allocator import ( |
| 161 | is_symmetric_memory_enabled, |
| 162 | ) |
| 163 | |
| 164 | if envs.VLLM_BATCH_INVARIANT: |
| 165 | return False |
| 166 | return is_symmetric_memory_enabled() |
| 167 | |
| 168 | |
| 169 | def producer( |
| 170 | batch_src: Sequence[int], |
| 171 | producer_queue, |
| 172 | consumer_queue, |
| 173 | result_queue, |
| 174 | cuda_visible_devices: str | None = None, |
| 175 | ): |
| 176 | if cuda_visible_devices is not None: |
| 177 | update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices}) |
| 178 | |
| 179 | lib = CudaRTLibrary() |
| 180 | for i in batch_src: |
| 181 | lib.cudaSetDevice(i) |
| 182 | pointer = lib.cudaMalloc(1024) |
| 183 | lib.cudaMemset(pointer, 1, 1024) |
| 184 | lib.cudaDeviceSynchronize() |
| 185 | handle = lib.cudaIpcGetMemHandle(pointer) |
| 186 | producer_queue.put(handle) |
| 187 | open_success = consumer_queue.get() |
| 188 | if open_success: |
| 189 | # use two queues to simulate barrier |
| 190 | producer_queue.put(0) |
| 191 | consumer_queue.get() |
| 192 | # check if the memory is modified |
| 193 | host_data = (ctypes.c_char * 1024)() |
| 194 | lib.cudaMemcpy(host_data, pointer, 1024) # type: ignore |
| 195 | for i in range(1024): |
| 196 | if ord(host_data[i]) != 2: |
| 197 | open_success = False |
| 198 | break |
| 199 | result_queue.put(open_success) |
| 200 | lib.cudaDeviceReset() |
| 201 | |
| 202 | |
| 203 | def consumer( |
| 204 | batch_tgt: Sequence[int], |
| 205 | producer_queue, |
| 206 | consumer_queue, |
| 207 | result_queue, |
| 208 | cuda_visible_devices: str | None = None, |
| 209 | ): |
| 210 | if cuda_visible_devices is not None: |
| 211 | update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices}) |
| 212 | |
| 213 | lib = CudaRTLibrary() |
| 214 | for j in batch_tgt: |
| 215 | lib.cudaSetDevice(j) |
| 216 | handle = producer_queue.get() |
| 217 | open_success = False |
| 218 | try: |
| 219 | pointer = lib.cudaIpcOpenMemHandle(handle) # type: ignore |
| 220 | open_success = True |
| 221 | except RuntimeError: |
| 222 | # cannot error out here, because the producer process |
| 223 | # is still waiting for the response. |
| 224 | pass |
| 225 | consumer_queue.put(open_success) |
| 226 | if open_success: |
| 227 | # modify the memory |
| 228 | lib.cudaMemset(pointer, 2, 1024) |
| 229 | lib.cudaDeviceSynchronize() |
| 230 | # use two queues to simulate barrier |
| 231 | producer_queue.get() |
| 232 | consumer_queue.put(0) |
| 233 | # check if the memory is modified |
| 234 | host_data = (ctypes.c_char * 1024)() |
| 235 | lib.cudaMemcpy(host_data, pointer, 1024) # type: ignore |
| 236 | for i in range(1024): |
| 237 | if ord(host_data[i]) != 2: |
| 238 | open_success = False |
| 239 | break |
| 240 | result_queue.put(open_success) |
| 241 | lib.cudaDeviceReset() |
| 242 | |
| 243 | |
| 244 | def can_actually_p2p( |
| 245 | batch_src: Sequence[int], |
| 246 | batch_tgt: Sequence[int], |
| 247 | ) -> Sequence[bool]: |
| 248 | """ |
| 249 | Usually, checking if P2P access is enabled can be done by |
| 250 | `torch.cuda.can_device_access_peer(src, tgt)`. However, sometimes |
| 251 | the driver might be broken, and `torch.cuda.can_device_access_peer(src, tgt)` |
| 252 | returns `True` even if P2P access is not actually possible. |
| 253 | See https://github.com/vllm-project/vllm/issues/2728 and |
| 254 | https://forums.developer.nvidia.com/t/direct-gpu-gpu-communication-does-not-seem-to-work-properly/283264/10 |
| 255 | Therefore, we have to perform a real P2P access to check if it is actually |
| 256 | possible. |
| 257 | |
| 258 | Note on p2p and cuda IPC: |
| 259 | Usually, one process uses one GPU: |
| 260 | GPU src --> cuda context src --> tensor src --> process src |
| 261 | |
| 262 | We need to combine p2p and cuda IPC, so that: |
| 263 | GPU src --> cuda context src --> tensor src --> process src |
| 264 | |shared| |
| 265 | GPU tgt --> cuda context tgt --> tensor tgt --> process tgt |
| 266 | That is to say, process src creates a tensor in GPU src, passes IPC handle to |
| 267 | process tgt, and process tgt accesses the tensor in GPU tgt. Any operation on the |
| 268 | tensor in process tgt will be reflected in the tensor in process src, because |
| 269 | they are the same memory segment. |
| 270 | It is important to note that process tgt accesses the tensor in GPU tgt, not |
| 271 | GPU src. That's why we need p2p access. |
| 272 | |
| 273 | The most time-consuming part is the process creation. To avoid creating |
| 274 | processes for every pair of GPUs, we use batched testing. We create two |
| 275 | processes for testing all pairs of GPUs in batch. The trick is to reset |
| 276 | the device after each test (which is not available in PyTorch). |
| 277 | """ # noqa |
| 278 | cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES |
| 279 | # pass the CUDA_VISIBLE_DEVICES to the child process |
| 280 | # to make sure they see the same set of GPUs |
| 281 | |
| 282 | # make sure the processes are spawned |
| 283 | smp = mp.get_context("spawn") |
| 284 | producer_queue = smp.Queue() |
| 285 | consumer_queue = smp.Queue() |
| 286 | result_queue = smp.Queue() |
| 287 | p_src = smp.Process( |
| 288 | target=producer, |
| 289 | args=( |
| 290 | batch_src, |
| 291 | producer_queue, |
| 292 | consumer_queue, |
| 293 | result_queue, |
| 294 | cuda_visible_devices, |
| 295 | ), |
| 296 | ) |
| 297 | p_tgt = smp.Process( |
| 298 | target=consumer, |
| 299 | args=( |
| 300 | batch_tgt, |
| 301 | producer_queue, |
| 302 | consumer_queue, |
| 303 | result_queue, |
| 304 | cuda_visible_devices, |
| 305 | ), |
| 306 | ) |
| 307 | p_src.start() |
| 308 | p_tgt.start() |
| 309 | p_src.join() |
| 310 | p_tgt.join() |
| 311 | assert p_src.exitcode == 0 and p_tgt.exitcode == 0 |
| 312 | result: list[bool] = [] |
| 313 | for src, tgt in zip(batch_src, batch_tgt): |
| 314 | a = result_queue.get() |
| 315 | b = result_queue.get() |
| 316 | if a != b: |
| 317 | logger.warning( |
| 318 | "Two processes do not agree on the P2P access" |
| 319 | " status on %d -> %d, treat as disabled.", |
| 320 | src, |
| 321 | tgt, |
| 322 | ) |
| 323 | result.append(False) |
| 324 | else: |
| 325 | result.append(a) |
| 326 | return result |
| 327 | |
| 328 | |
| 329 | # why do we need this cache? |
| 330 | # we are testing peer-to-peer (p2p) access between GPUs,across processes. |
| 331 | # if we test it every time, it will be very slow, because we need to create |
| 332 | # N * N * 2 processes, where N is the world size. This is very slow. |
| 333 | # to reduce the time, we use a cache file to store the p2p access status. |
| 334 | # the cache file is generated by the master process if it does not exist. |
| 335 | # then all the processes can read the cache file to check the p2p access status. |
| 336 | # Note that the cache file is suffixed by the CUDA_VISIBLE_DEVICES, so that we |
| 337 | # can have different cache files for different CUDA_VISIBLE_DEVICES settings, |
| 338 | # e.g. used by different vllm engines. The device id in the cache file is a |
| 339 | # **local** device id, i.e. from 0 to num_dev-1, where num_dev is the number |
| 340 | # of visible devices in the vllm engine. |
| 341 | _gpu_p2p_access_cache: dict[str, bool] | None = None |
| 342 | |
| 343 | |
| 344 | def gpu_p2p_access_check(src: int, tgt: int) -> bool: |
| 345 | """Check if GPU src can access GPU tgt.""" |
| 346 | # if the cache variable is already calculated, |
| 347 | # read from the cache instead of checking it again |
| 348 | global _gpu_p2p_access_cache |
| 349 | if _gpu_p2p_access_cache is not None: |
| 350 | return _gpu_p2p_access_cache[f"{src}->{tgt}"] |
| 351 | |
| 352 | is_distributed = dist.is_initialized() |
| 353 | |
| 354 | from vllm.platforms.interface import get_assigned_physical_gpu_ids |
| 355 | |
| 356 | assigned_physical_gpu_ids = get_assigned_physical_gpu_ids() |
| 357 | if assigned_physical_gpu_ids is not None: |
| 358 | # Key by the ordered list: the cache stores directed local-index |
| 359 | # pairs, so permutations of the same set are distinct mappings. |
| 360 | cache_key = ",".join(str(i) for i in assigned_physical_gpu_ids) |
| 361 | num_dev = len(assigned_physical_gpu_ids) |
| 362 | else: |
| 363 | num_dev = current_platform.device_count() |
| 364 | cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES |
| 365 | cache_key = cuda_visible_devices or ",".join(str(i) for i in range(num_dev)) |
| 366 | |
| 367 | path = os.path.join( |
| 368 | envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cache_key}.json" |
| 369 | ) |
| 370 | os.makedirs(os.path.dirname(path), exist_ok=True) |
| 371 | from vllm.distributed.parallel_state import get_world_group |
| 372 | |
| 373 | if (not is_distributed or get_world_group().local_rank == 0) and ( |
| 374 | not os.path.exists(path) |
| 375 | ): |
| 376 | # only the local master process (with local_rank == 0) can |
| 377 | # enter this block to calculate the cache |
| 378 | logger.info("generating GPU P2P access cache in %s", path) |
| 379 | cache: dict[str, bool] = {} |
| 380 | # The probe subprocesses inherit this process's device-control env |
| 381 | # var, so they must be given visible ordinals, not physical IDs. |
| 382 | if assigned_physical_gpu_ids is not None: |
| 383 | ids = [ |
| 384 | current_platform.logical_device_id_to_visible_device_id(local) |
| 385 | for local in range(num_dev) |
| 386 | ] |
| 387 | else: |
| 388 | ids = list(range(num_dev)) |
| 389 | # batch of all pairs of GPUs |
| 390 | batch_src, batch_tgt = zip(*list(product(ids, ids))) |
| 391 | # NOTE: we use `subprocess` rather than `multiprocessing` here |
| 392 | # because the caller might not have `if __name__ == "__main__":`, |
| 393 | # in that case we cannot use spawn method in multiprocessing. |
| 394 | # However, `can_actually_p2p` requires spawn method. |
| 395 | # The fix is, we use `subprocess` to call the function, |
| 396 | # where we have `if __name__ == "__main__":` in this file. |
| 397 | |
| 398 | # use a temporary file to store the result |
| 399 | # we don't use the output of the subprocess directly, |
| 400 | # because the subprocess might produce logging output |
| 401 | with tempfile.NamedTemporaryFile() as output_file: |
| 402 | input_bytes = pickle.dumps((batch_src, batch_tgt, output_file.name)) |
| 403 | returned = subprocess.run( |
| 404 | [sys.executable, __file__], input=input_bytes, capture_output=True |
| 405 | ) |
| 406 | # check if the subprocess is successful |
| 407 | try: |
| 408 | returned.check_returncode() |
| 409 | except Exception as e: |
| 410 | # wrap raised exception to provide more information |
| 411 | raise RuntimeError( |
| 412 | f"Error happened when batch testing " |
| 413 | f"peer-to-peer access from {batch_src} to {batch_tgt}:\n" |
| 414 | f"{returned.stderr.decode()}" |
| 415 | ) from e |
| 416 | with open(output_file.name, "rb") as f: |
| 417 | result = pickle.load(f) |
| 418 | # Cache entries must be keyed by local indices (0..N-1) because |
| 419 | # gpu_p2p_access_check() is called with local ranks. |
| 420 | id_to_local = {device_id: local for local, device_id in enumerate(ids)} |
| 421 | for _i, _j, r in zip(batch_src, batch_tgt, result): |
| 422 | cache[f"{id_to_local[_i]}->{id_to_local[_j]}"] = r |
| 423 | with open(path, "w") as f: |
| 424 | json.dump(cache, f, indent=4) |
| 425 | if is_distributed: |
| 426 | get_world_group().barrier() |
| 427 | logger.info("reading GPU P2P access cache from %s", path) |
| 428 | with open(path) as f: |
| 429 | cache = json.load(f) |
| 430 | _gpu_p2p_access_cache = cache |
| 431 | return _gpu_p2p_access_cache[f"{src}->{tgt}"] |
| 432 | |
| 433 | |
| 434 | __all__ = ["gpu_p2p_access_check"] |
| 435 | |
| 436 | if __name__ == "__main__": |
| 437 | batch_src, batch_tgt, output_file = pickle.loads(sys.stdin.buffer.read()) |
| 438 | result = can_actually_p2p(batch_src, batch_tgt) |
| 439 | with open(output_file, "wb") as f: |
| 440 | f.write(pickle.dumps(result)) |