| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import atexit |
| 6 | import os |
| 7 | import random |
| 8 | import threading |
| 9 | from typing import Any |
| 10 | |
| 11 | import torch |
| 12 | import torch.distributed as dist |
| 13 | from torch.distributed import ProcessGroup |
| 14 | |
| 15 | import vllm.envs as envs |
| 16 | from vllm.config.compilation import PassConfig |
| 17 | from vllm.distributed.device_communicators.all_reduce_utils import ( |
| 18 | FI_MNNVL_ALLREDUCE_MAX_SIZE_MB, |
| 19 | ) |
| 20 | from vllm.distributed.parallel_state import _node_count, get_node_count |
| 21 | from vllm.logger import init_logger |
| 22 | from vllm.platforms import current_platform |
| 23 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 24 | |
| 25 | logger = init_logger(__name__) |
| 26 | |
| 27 | # The empirical value for small batch |
| 28 | PDL_ADVANCE_LAUNCH_TOKENS = 16 |
| 29 | |
| 30 | MiB = 1024 * 1024 |
| 31 | |
| 32 | fi_ar_available = False |
| 33 | try: |
| 34 | import flashinfer.comm as flashinfer_comm # type: ignore[no-redef] |
| 35 | from flashinfer.comm.mnnvl import ( |
| 36 | TorchDistBackend, # type: ignore[import-not-found, no-redef] |
| 37 | ) |
| 38 | |
| 39 | fi_ar_available = hasattr(flashinfer_comm, "allreduce_fusion") |
| 40 | except ImportError: |
| 41 | pass |
| 42 | |
| 43 | # Workspace for standalone allreduce and non-quant ar+rms fusion |
| 44 | _fi_ar_workspace = None |
| 45 | # Extra workspace for quant fusion patterns. This may use either the primary |
| 46 | # allreduce backend or a fallback backend when the primary workspace is not |
| 47 | # available on the current topology. |
| 48 | _fi_ar_quant_workspace = None |
| 49 | # Creation is collective, so a failed attempt must not be retried on every |
| 50 | # all-reduce call. |
| 51 | _fi_ar_workspace_failed = False |
| 52 | _fi_ar_quant_workspace_failed = False |
| 53 | _fi_ar_workspace_groups: dict[int, ProcessGroup] = {} |
| 54 | |
| 55 | |
| 56 | def _get_tuned_standalone_max_size( |
| 57 | world_size: int, |
| 58 | backend: str, |
| 59 | group: ProcessGroup, |
| 60 | ) -> int | None: |
| 61 | if backend != "mnnvl": |
| 62 | return None |
| 63 | capability = current_platform.get_device_capability() |
| 64 | if capability is None: |
| 65 | return None |
| 66 | max_size_mb = FI_MNNVL_ALLREDUCE_MAX_SIZE_MB.get( |
| 67 | (capability.to_int(), world_size, _node_count(group)) |
| 68 | ) |
| 69 | # Tuned cutoffs are exclusive; store the largest accepted size. |
| 70 | return None if max_size_mb is None else int(max_size_mb * MiB) - 1 |
| 71 | |
| 72 | |
| 73 | def _create_workspace( |
| 74 | backend: str, |
| 75 | world_size: int, |
| 76 | rank: int, |
| 77 | max_token_num: int, |
| 78 | hidden_dim: int, |
| 79 | dtype: torch.dtype, |
| 80 | group: ProcessGroup, |
| 81 | ): |
| 82 | """Create a flashinfer allreduce workspace, returning None on failure.""" |
| 83 | comm_backend = TorchDistBackend(group=group) |
| 84 | rng_state = random.getstate() |
| 85 | try: |
| 86 | random.seed(int.from_bytes(os.urandom(16), byteorder="big")) |
| 87 | # Creation may run lazily inside the first sync-checked forward (e.g. |
| 88 | # with enforce_eager, which skips warmup). |
| 89 | with gpu_sync_allowed(first_only=True): |
| 90 | workspace = flashinfer_comm.create_allreduce_fusion_workspace( |
| 91 | backend=backend, |
| 92 | world_size=world_size, |
| 93 | rank=rank, |
| 94 | max_token_num=max_token_num, |
| 95 | hidden_dim=hidden_dim, |
| 96 | dtype=dtype, |
| 97 | comm_backend=comm_backend, |
| 98 | group=group, |
| 99 | ) |
| 100 | if backend == "mnnvl" and not getattr(workspace, "mc_ptr", 0): |
| 101 | workspace.destroy() |
| 102 | logger.warning_once( |
| 103 | "FlashInfer MNNVL multicast is unavailable on the current topology." |
| 104 | ) |
| 105 | return None |
| 106 | except Exception as e: |
| 107 | if "multicast" in str(e).lower(): |
| 108 | logger.warning_once( |
| 109 | "Failed to initialize FlashInfer All Reduce workspace: %s. " |
| 110 | "This is expected on GPUs without NVSwitch (e.g., NVLink " |
| 111 | "bridge-only or PCIe topologies).", |
| 112 | e, |
| 113 | ) |
| 114 | else: |
| 115 | logger.warning_once( |
| 116 | "Failed to initialize FlashInfer All Reduce workspace: %s.", |
| 117 | e, |
| 118 | ) |
| 119 | return None |
| 120 | finally: |
| 121 | random.setstate(rng_state) |
| 122 | workspace_id = id(workspace) |
| 123 | workspace_group = _fi_ar_workspace_groups.get(workspace_id) |
| 124 | if workspace_group is not None and workspace_group is not group: |
| 125 | raise RuntimeError( |
| 126 | "FlashInfer returned an all-reduce workspace already associated " |
| 127 | "with a different process group" |
| 128 | ) |
| 129 | _fi_ar_workspace_groups[workspace_id] = group |
| 130 | logger.debug( |
| 131 | "Initialized FlashInfer All Reduce workspace: backend=%s, " |
| 132 | "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s", |
| 133 | backend, |
| 134 | world_size, |
| 135 | rank, |
| 136 | max_token_num, |
| 137 | hidden_dim, |
| 138 | dtype, |
| 139 | ) |
| 140 | return workspace |
| 141 | |
| 142 | |
| 143 | def _resolve_fi_ar_backend() -> tuple[str, bool]: |
| 144 | """Resolve the flashinfer allreduce backend for the current setup. |
| 145 | |
| 146 | Returns: |
| 147 | A ``(backend, allow_trtllm_fallback)`` tuple. ``allow_trtllm_fallback`` |
| 148 | is True only when ``auto`` selects mnnvl for a single node, so that |
| 149 | workspace creation can fall back to trtllm on single-node topologies |
| 150 | without NVSwitch multicast support (where mnnvl is unavailable). |
| 151 | |
| 152 | """ |
| 153 | backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND |
| 154 | if backend != "auto": |
| 155 | logger.debug_once("Using flashinfer allreduce backend: %s", backend) |
| 156 | return backend, False |
| 157 | |
| 158 | # Default to mnnvl for both single- and multi-node setups. The mnnvl |
| 159 | # cudagraph hang that previously forced single-node to trtllm |
| 160 | # (https://github.com/vllm-project/vllm/issues/35772) was fixed upstream in |
| 161 | # FlashInfer (>= 0.6.12, vLLM pins 0.6.15), so mnnvl is safe here. trtllm |
| 162 | # does not support multi-node allreduce, so mnnvl is required there anyway. |
| 163 | # mnnvl needs NVSwitch multicast; on single-node topologies without it, |
| 164 | # fall back to trtllm so fused allreduce stays enabled. |
| 165 | backend = "mnnvl" |
| 166 | allow_trtllm_fallback = get_node_count() == 1 |
| 167 | |
| 168 | logger.debug_once("Auto-selected flashinfer allreduce backend: %s", backend) |
| 169 | return backend, allow_trtllm_fallback |
| 170 | |
| 171 | |
| 172 | def get_fi_ar_workspace( |
| 173 | world_size: int, |
| 174 | rank: int, |
| 175 | max_token_num: int, |
| 176 | hidden_dim: int, |
| 177 | dtype: torch.dtype, |
| 178 | group: ProcessGroup, |
| 179 | ): |
| 180 | """Return the allreduce workspace for non-quant patterns, initializing if needed. |
| 181 | |
| 182 | Used by AllReduceFusionPass (non-quant patterns) and FlashInferAllReduce |
| 183 | for standalone allreduce. Backend is controlled by |
| 184 | VLLM_FLASHINFER_ALLREDUCE_BACKEND env var. |
| 185 | """ |
| 186 | global _fi_ar_workspace, _fi_ar_workspace_failed |
| 187 | if _fi_ar_workspace is not None or _fi_ar_workspace_failed: |
| 188 | return _fi_ar_workspace |
| 189 | |
| 190 | backend, allow_trtllm_fallback = _resolve_fi_ar_backend() |
| 191 | |
| 192 | if get_node_count() > 1 and backend == "trtllm": |
| 193 | raise ValueError( |
| 194 | "Flashinfer allreduce is not supported for multi-node allreduce with " |
| 195 | "'trtllm' backend. Please use 'mnnvl' backend instead." |
| 196 | ) |
| 197 | |
| 198 | if ( |
| 199 | envs.VLLM_ALLREDUCE_USE_FLASHINFER |
| 200 | and (max_size := _get_tuned_standalone_max_size(world_size, backend, group)) |
| 201 | is not None |
| 202 | ): |
| 203 | element_size = torch.empty((), dtype=dtype, device="cpu").element_size() |
| 204 | max_token_num = max(max_token_num, max_size // (hidden_dim * element_size)) |
| 205 | |
| 206 | def _get_or_create(be: str): |
| 207 | # Reuse the quant workspace if it was already created with the same backend |
| 208 | if _fi_ar_quant_workspace is not None and _fi_ar_quant_workspace.backend == be: |
| 209 | return _fi_ar_quant_workspace |
| 210 | return _create_workspace( |
| 211 | be, world_size, rank, max_token_num, hidden_dim, dtype, group |
| 212 | ) |
| 213 | |
| 214 | _fi_ar_workspace = _get_or_create(backend) |
| 215 | if _fi_ar_workspace is None and allow_trtllm_fallback and backend != "trtllm": |
| 216 | logger.warning_once( |
| 217 | "FlashInfer mnnvl allreduce workspace unavailable (likely no NVSwitch " |
| 218 | "multicast support); falling back to trtllm backend for single node." |
| 219 | ) |
| 220 | backend = "trtllm" |
| 221 | _fi_ar_workspace = _get_or_create(backend) |
| 222 | |
| 223 | if _fi_ar_workspace is not None: |
| 224 | logger.info_once( |
| 225 | "Initialized FlashInfer Allreduce norm fusion workspace " |
| 226 | f"with backend={backend}" |
| 227 | ) |
| 228 | else: |
| 229 | _fi_ar_workspace_failed = True |
| 230 | logger.error_once( |
| 231 | "Failed to initialize FlashInfer Allreduce norm fusion workspace " |
| 232 | f"with backend={backend}" |
| 233 | ) |
| 234 | |
| 235 | return _fi_ar_workspace |
| 236 | |
| 237 | |
| 238 | def get_fi_ar_quant_workspace( |
| 239 | world_size: int, |
| 240 | rank: int, |
| 241 | max_token_num: int, |
| 242 | hidden_dim: int, |
| 243 | dtype: torch.dtype, |
| 244 | group: ProcessGroup, |
| 245 | ): |
| 246 | """Return the allreduce workspace for quant patterns, initializing if needed. |
| 247 | |
| 248 | Backend is controlled by VLLM_FLASHINFER_ALLREDUCE_BACKEND env var, matching |
| 249 | non-quant fusion. With ``auto`` this prefers mnnvl and falls back to trtllm |
| 250 | only on single-node topologies where mnnvl multicast is unavailable. |
| 251 | """ |
| 252 | global _fi_ar_quant_workspace, _fi_ar_quant_workspace_failed |
| 253 | if _fi_ar_quant_workspace is not None or _fi_ar_quant_workspace_failed: |
| 254 | return _fi_ar_quant_workspace |
| 255 | |
| 256 | backend, allow_trtllm_fallback = _resolve_fi_ar_backend() |
| 257 | |
| 258 | if get_node_count() > 1 and backend == "trtllm": |
| 259 | raise ValueError( |
| 260 | "Flashinfer allreduce quantization fusion is not supported for " |
| 261 | "multi-node allreduce with 'trtllm' backend. Please use 'mnnvl' " |
| 262 | "backend instead." |
| 263 | ) |
| 264 | |
| 265 | # Reuse the non-quant workspace if it was already created with the same |
| 266 | # backend. |
| 267 | if _fi_ar_workspace is not None and _fi_ar_workspace.backend == backend: |
| 268 | _fi_ar_quant_workspace = _fi_ar_workspace |
| 269 | return _fi_ar_quant_workspace |
| 270 | |
| 271 | if ( |
| 272 | _fi_ar_workspace is not None |
| 273 | and _fi_ar_workspace.backend == "trtllm" |
| 274 | and allow_trtllm_fallback |
| 275 | and backend != "trtllm" |
| 276 | ): |
| 277 | _fi_ar_quant_workspace = _fi_ar_workspace |
| 278 | return _fi_ar_quant_workspace |
| 279 | |
| 280 | _fi_ar_quant_workspace = _create_workspace( |
| 281 | backend, world_size, rank, max_token_num, hidden_dim, dtype, group |
| 282 | ) |
| 283 | if _fi_ar_quant_workspace is None and allow_trtllm_fallback and backend != "trtllm": |
| 284 | logger.warning_once( |
| 285 | "FlashInfer mnnvl allreduce quantization fusion workspace unavailable " |
| 286 | "(likely no NVSwitch multicast support); falling back to trtllm " |
| 287 | "backend for single node." |
| 288 | ) |
| 289 | backend = "trtllm" |
| 290 | if _fi_ar_workspace is not None and _fi_ar_workspace.backend == backend: |
| 291 | _fi_ar_quant_workspace = _fi_ar_workspace |
| 292 | else: |
| 293 | _fi_ar_quant_workspace = _create_workspace( |
| 294 | backend, world_size, rank, max_token_num, hidden_dim, dtype, group |
| 295 | ) |
| 296 | |
| 297 | if _fi_ar_quant_workspace is not None: |
| 298 | logger.info_once( |
| 299 | "Initialized FlashInfer Allreduce norm quantization " |
| 300 | f"fusion workspace with backend={backend}" |
| 301 | ) |
| 302 | else: |
| 303 | _fi_ar_quant_workspace_failed = True |
| 304 | logger.error_once( |
| 305 | "Failed to initialize FlashInfer Allreduce norm quantization " |
| 306 | f"fusion workspace with backend={backend}" |
| 307 | ) |
| 308 | |
| 309 | return _fi_ar_quant_workspace |
| 310 | |
| 311 | |
| 312 | _fi_ar_workspace_lock = threading.Lock() |
| 313 | |
| 314 | |
| 315 | def destroy_fi_ar_workspace(): |
| 316 | global _fi_ar_workspace, _fi_ar_quant_workspace |
| 317 | global _fi_ar_workspace_failed, _fi_ar_quant_workspace_failed |
| 318 | with _fi_ar_workspace_lock: |
| 319 | is_alias = _fi_ar_workspace is _fi_ar_quant_workspace |
| 320 | |
| 321 | if _fi_ar_workspace is not None: |
| 322 | _fi_ar_workspace.destroy() |
| 323 | if _fi_ar_quant_workspace is not None and not is_alias: |
| 324 | _fi_ar_quant_workspace.destroy() |
| 325 | |
| 326 | _fi_ar_workspace = _fi_ar_quant_workspace = None |
| 327 | _fi_ar_workspace_failed = _fi_ar_quant_workspace_failed = False |
| 328 | _fi_ar_workspace_groups.clear() |
| 329 | |
| 330 | |
| 331 | def _fi_ar_workspaces_for_group(group: ProcessGroup) -> list[Any]: |
| 332 | workspaces = [_fi_ar_workspace] |
| 333 | if _fi_ar_quant_workspace is not _fi_ar_workspace: |
| 334 | workspaces.append(_fi_ar_quant_workspace) |
| 335 | |
| 336 | group_workspaces = [] |
| 337 | for workspace in workspaces: |
| 338 | if workspace is None: |
| 339 | continue |
| 340 | workspace_group = _fi_ar_workspace_groups.get(id(workspace)) |
| 341 | if workspace_group is None: |
| 342 | raise RuntimeError( |
| 343 | "FlashInfer all-reduce workspace process group was not retained" |
| 344 | ) |
| 345 | if workspace_group is group: |
| 346 | group_workspaces.append(workspace) |
| 347 | return group_workspaces |
| 348 | |
| 349 | |
| 350 | def checkpoint_prepare_fi_ar_workspaces(group: ProcessGroup) -> None: |
| 351 | for workspace in _fi_ar_workspaces_for_group(group): |
| 352 | workspace.checkpoint_prepare() |
| 353 | |
| 354 | |
| 355 | def checkpoint_restore_fi_ar_workspaces(group: ProcessGroup) -> None: |
| 356 | for workspace in _fi_ar_workspaces_for_group(group): |
| 357 | workspace.checkpoint_restore(TorchDistBackend(group=group)) |
| 358 | |
| 359 | |
| 360 | atexit.register(destroy_fi_ar_workspace) |
| 361 | |
| 362 | |
| 363 | class FlashInferAllReduce: |
| 364 | def __init__( |
| 365 | self, |
| 366 | group: ProcessGroup, |
| 367 | device: int | str | torch.device, |
| 368 | ): |
| 369 | self.disabled = True |
| 370 | |
| 371 | if not fi_ar_available: |
| 372 | logger.info( |
| 373 | "FlashInfer All Reduce is disabled because flashinfer is not available" |
| 374 | ) |
| 375 | return |
| 376 | |
| 377 | if not current_platform.is_cuda(): |
| 378 | logger.info( |
| 379 | "FlashInfer All Reduce is disabled because it requires CUDA platform" |
| 380 | ) |
| 381 | return |
| 382 | |
| 383 | self.group = group |
| 384 | self.world_size = dist.get_world_size(self.group) |
| 385 | self.rank = dist.get_rank(self.group) |
| 386 | self.device = device |
| 387 | if self.world_size == 1: |
| 388 | return |
| 389 | |
| 390 | default_max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get( |
| 391 | self.world_size |
| 392 | ) |
| 393 | if not default_max_size_mb: |
| 394 | logger.warning( |
| 395 | "FlashInfer All Reduce is disabled because it " |
| 396 | "is not supported for world_size=%d.", |
| 397 | self.world_size, |
| 398 | ) |
| 399 | return |
| 400 | backend, _ = _resolve_fi_ar_backend() |
| 401 | tuned_max_size = _get_tuned_standalone_max_size( |
| 402 | self.world_size, |
| 403 | backend, |
| 404 | self.group, |
| 405 | ) |
| 406 | self.max_workspace_size = ( |
| 407 | tuned_max_size |
| 408 | if tuned_max_size is not None |
| 409 | else int(default_max_size_mb * MiB) |
| 410 | ) |
| 411 | self.max_num_tokens = 0 |
| 412 | self.disabled = False |
| 413 | |
| 414 | def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool: |
| 415 | """Ensure the all reduce workspace is initialized.""" |
| 416 | if self.max_num_tokens == 0: |
| 417 | element_size = torch.tensor([], dtype=dtype, device="cpu").element_size() |
| 418 | self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size) |
| 419 | workspace = get_fi_ar_workspace( |
| 420 | world_size=self.world_size, |
| 421 | rank=self.rank, |
| 422 | max_token_num=self.max_num_tokens, |
| 423 | hidden_dim=hidden_dim, |
| 424 | dtype=dtype, |
| 425 | group=self.group, |
| 426 | ) |
| 427 | if workspace is None: |
| 428 | self.disabled = True |
| 429 | return False |
| 430 | return True |
| 431 | |
| 432 | def should_use_fi_ar(self, input_tensor: torch.Tensor) -> bool: |
| 433 | if self.disabled: |
| 434 | return False |
| 435 | |
| 436 | if not input_tensor.is_cuda: |
| 437 | return False |
| 438 | |
| 439 | if not input_tensor.is_contiguous(): |
| 440 | return False |
| 441 | |
| 442 | if len(input_tensor.shape) != 2: |
| 443 | return False |
| 444 | |
| 445 | if input_tensor.dtype not in ( |
| 446 | torch.float16, |
| 447 | torch.bfloat16, |
| 448 | torch.float32, |
| 449 | ): |
| 450 | return False |
| 451 | |
| 452 | if input_tensor.nbytes > self.max_workspace_size: |
| 453 | return False |
| 454 | |
| 455 | num_tokens, hidden_dim = input_tensor.shape |
| 456 | if not self.max_num_tokens: |
| 457 | element_size = torch.tensor([], dtype=input_tensor.dtype).element_size() |
| 458 | self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size) |
| 459 | |
| 460 | if num_tokens > self.max_num_tokens: |
| 461 | return False |
| 462 | |
| 463 | if not self._ensure_workspace(hidden_dim, input_tensor.dtype): |
| 464 | return False |
| 465 | |
| 466 | workspace = get_fi_ar_workspace( |
| 467 | world_size=self.world_size, |
| 468 | rank=self.rank, |
| 469 | max_token_num=self.max_num_tokens, |
| 470 | hidden_dim=hidden_dim, |
| 471 | dtype=input_tensor.dtype, |
| 472 | group=self.group, |
| 473 | ) |
| 474 | assert workspace is not None |
| 475 | # The token bound above uses the full allocation budget, but mnnvl's |
| 476 | # Lamport buffers rotate through three slots, so only ~1/3 is usable per |
| 477 | # call. Asking the workspace directly to reject sizes the kernel can't fit. |
| 478 | return workspace.is_buffer_size_sufficient( |
| 479 | tp_size=self.world_size, |
| 480 | num_tokens=num_tokens, |
| 481 | hidden_dim=hidden_dim, |
| 482 | dtype=input_tensor.dtype, |
| 483 | ) |
| 484 | |
| 485 | def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor: |
| 486 | num_tokens, hidden_dim = input_tensor.shape |
| 487 | workspace = get_fi_ar_workspace( |
| 488 | world_size=self.world_size, |
| 489 | rank=self.rank, |
| 490 | max_token_num=self.max_num_tokens, |
| 491 | hidden_dim=hidden_dim, |
| 492 | dtype=input_tensor.dtype, |
| 493 | group=self.group, |
| 494 | ) |
| 495 | return flashinfer_comm.allreduce_fusion( |
| 496 | input=input_tensor, |
| 497 | workspace=workspace, |
| 498 | pattern=flashinfer_comm.AllReduceFusionPattern.kAllReduce, |
| 499 | launch_with_pdl=True, |
| 500 | trigger_completion_at_end=num_tokens > PDL_ADVANCE_LAUNCH_TOKENS, |
| 501 | ) |
| 502 | |
| 503 | def destroy(self): |
| 504 | if not self.disabled: |
| 505 | destroy_fi_ar_workspace() |