| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """FlashInfer PCIe CUDA-IPC all-reduce integration. |
| 4 | |
| 5 | This backend targets small tensor-parallel all-reduces on a single PCIe-only |
| 6 | node. Its workspace has a stricter lifetime and stream contract than the |
| 7 | existing FlashInfer MNNVL/TRT-LLM all-reduce implementation, so it intentionally |
| 8 | lives behind a separate wrapper and an opt-in environment variable. |
| 9 | """ |
| 10 | |
| 11 | from collections.abc import Sequence |
| 12 | from contextlib import contextmanager |
| 13 | from pathlib import Path |
| 14 | from typing import TYPE_CHECKING, Any |
| 15 | |
| 16 | import torch |
| 17 | import torch.distributed as dist |
| 18 | from torch.distributed import ProcessGroup |
| 19 | |
| 20 | from vllm.distributed.parallel_state import in_the_same_node_as |
| 21 | from vllm.logger import init_logger |
| 22 | from vllm.platforms import current_platform |
| 23 | |
| 24 | if TYPE_CHECKING: |
| 25 | from vllm.v1.worker.gpu_worker import Worker |
| 26 | |
| 27 | logger = init_logger(__name__) |
| 28 | |
| 29 | _SUPPORTED_WORLD_SIZES = (2, 4, 8) |
| 30 | _SUPPORTED_DTYPES = (torch.bfloat16, torch.float16) |
| 31 | |
| 32 | try: |
| 33 | import flashinfer.comm as flashinfer_comm |
| 34 | |
| 35 | _pcie_ipc_available = hasattr(flashinfer_comm, "PcieIpcAllReduceWorkspace") |
| 36 | except ImportError: |
| 37 | flashinfer_comm = None # type: ignore[assignment] |
| 38 | _pcie_ipc_available = False |
| 39 | |
| 40 | |
| 41 | class FlashInferPcieIpcAllReduce: |
| 42 | """vLLM lifecycle wrapper for FlashInfer's PCIe IPC all-reduce.""" |
| 43 | |
| 44 | def __init__( |
| 45 | self, |
| 46 | group: ProcessGroup, |
| 47 | tune_group: ProcessGroup, |
| 48 | device: int | str | torch.device, |
| 49 | ) -> None: |
| 50 | self.disabled = True |
| 51 | self.group = group |
| 52 | self.tune_group = tune_group |
| 53 | self.device = torch.device(device) |
| 54 | self.world_size = dist.get_world_size(group) |
| 55 | self.rank = dist.get_rank(group) |
| 56 | self.workspace: Any | None = None |
| 57 | self.hidden_dim = 0 |
| 58 | self.dtype: torch.dtype | None = None |
| 59 | |
| 60 | if not _pcie_ipc_available: |
| 61 | logger.warning_once( |
| 62 | "FlashInfer PCIe IPC all-reduce was requested but this " |
| 63 | "FlashInfer build does not provide PcieIpcAllReduceWorkspace; " |
| 64 | "falling back to another all-reduce backend." |
| 65 | ) |
| 66 | return |
| 67 | if not current_platform.is_cuda(): |
| 68 | logger.warning_once( |
| 69 | "FlashInfer PCIe IPC all-reduce requires the CUDA platform." |
| 70 | ) |
| 71 | return |
| 72 | if self.world_size not in _SUPPORTED_WORLD_SIZES: |
| 73 | logger.warning_once( |
| 74 | "FlashInfer PCIe IPC all-reduce does not support world_size=%d; " |
| 75 | "supported sizes are %s.", |
| 76 | self.world_size, |
| 77 | _SUPPORTED_WORLD_SIZES, |
| 78 | ) |
| 79 | return |
| 80 | if not all(in_the_same_node_as(tune_group, source_rank=0)): |
| 81 | logger.warning_once( |
| 82 | "FlashInfer PCIe IPC all-reduce requires every rank in the TP " |
| 83 | "group to be on one node." |
| 84 | ) |
| 85 | return |
| 86 | |
| 87 | # Setup is deferred until kernel_warmup, where the model hidden size and |
| 88 | # exact CUDA Graph buckets are known. Until then dispatch falls through. |
| 89 | self.disabled = False |
| 90 | |
| 91 | @property |
| 92 | def initialized(self) -> bool: |
| 93 | return not self.disabled and self.workspace is not None |
| 94 | |
| 95 | def setup( |
| 96 | self, |
| 97 | *, |
| 98 | hidden_dim: int, |
| 99 | dtype: torch.dtype, |
| 100 | capture_sizes: Sequence[int], |
| 101 | tune_cache: Path, |
| 102 | ) -> None: |
| 103 | """Allocate, tune, and prepare the exact graph-capture shapes.""" |
| 104 | if self.disabled or self.workspace is not None: |
| 105 | return |
| 106 | if dtype not in _SUPPORTED_DTYPES: |
| 107 | logger.warning_once( |
| 108 | "FlashInfer PCIe IPC all-reduce does not support dtype=%s; " |
| 109 | "falling back to another all-reduce backend.", |
| 110 | dtype, |
| 111 | ) |
| 112 | self.disabled = True |
| 113 | return |
| 114 | |
| 115 | batches = tuple(sorted({int(size) for size in capture_sizes if size > 0})) |
| 116 | if not batches: |
| 117 | logger.warning_once( |
| 118 | "FlashInfer PCIe IPC all-reduce has no CUDA Graph capture sizes " |
| 119 | "to prepare; falling back to another all-reduce backend." |
| 120 | ) |
| 121 | self.disabled = True |
| 122 | return |
| 123 | |
| 124 | self.hidden_dim = int(hidden_dim) |
| 125 | self.dtype = dtype |
| 126 | max_numel = batches[-1] * self.hidden_dim |
| 127 | workspace = flashinfer_comm.PcieIpcAllReduceWorkspace( |
| 128 | group=self.group, |
| 129 | max_numel=max_numel, |
| 130 | dtype=dtype, |
| 131 | tune_batches=batches, |
| 132 | tune_cache=str(tune_cache), |
| 133 | ) |
| 134 | self.workspace = workspace |
| 135 | |
| 136 | # tune() reuses a complete persisted cache and profiles only cache |
| 137 | # misses. It is deliberately mandatory: FlashInfer's seed policy was |
| 138 | # substantially slower than NCCL for the target TP4 decode workload. |
| 139 | torch.accelerator.synchronize(self.device) |
| 140 | workspace.rebind_stream() |
| 141 | workspace.tune( |
| 142 | [self.hidden_dim], |
| 143 | dtype=dtype, |
| 144 | tune_group=self.tune_group, |
| 145 | ) |
| 146 | workspace.prepare([(batch, self.hidden_dim) for batch in batches], dtype=dtype) |
| 147 | torch.accelerator.synchronize(self.device) |
| 148 | workspace.rebind_stream() |
| 149 | logger.info_once( |
| 150 | "Initialized FlashInfer PCIe IPC all-reduce for TP%d, hidden_dim=%d, " |
| 151 | "dtype=%s, max_tokens=%d.", |
| 152 | self.world_size, |
| 153 | self.hidden_dim, |
| 154 | dtype, |
| 155 | batches[-1], |
| 156 | ) |
| 157 | |
| 158 | def should_use(self, inp: torch.Tensor) -> bool: |
| 159 | workspace = self.workspace |
| 160 | return bool( |
| 161 | not self.disabled |
| 162 | and workspace is not None |
| 163 | and inp.is_cuda |
| 164 | and inp.is_contiguous() |
| 165 | and inp.dim() == 2 |
| 166 | and inp.shape[1] == self.hidden_dim |
| 167 | and inp.dtype == self.dtype |
| 168 | and workspace.supports(inp) |
| 169 | ) |
| 170 | |
| 171 | def all_reduce(self, inp: torch.Tensor) -> torch.Tensor: |
| 172 | assert self.workspace is not None |
| 173 | return self.workspace.all_reduce(inp) |
| 174 | |
| 175 | @contextmanager |
| 176 | def capture(self): |
| 177 | """Move the single-stream workspace to and from the capture stream.""" |
| 178 | workspace = self.workspace |
| 179 | if workspace is None: |
| 180 | yield |
| 181 | return |
| 182 | |
| 183 | # The caller orders graph capture against normal execution. A device |
| 184 | # sync makes that assertion explicit before relaxing FlashInfer's |
| 185 | # single-stream guard. Captured launches themselves are exempt from the |
| 186 | # guard, while capture warmups execute and bind to the capture stream. |
| 187 | torch.accelerator.synchronize(self.device) |
| 188 | workspace.rebind_stream() |
| 189 | try: |
| 190 | yield |
| 191 | finally: |
| 192 | torch.accelerator.synchronize(self.device) |
| 193 | workspace.rebind_stream() |
| 194 | |
| 195 | def destroy(self) -> None: |
| 196 | workspace = self.workspace |
| 197 | if workspace is not None: |
| 198 | workspace.destroy() |
| 199 | self.workspace = None |
| 200 | |
| 201 | |
| 202 | def warmup_flashinfer_pcie_ipc_allreduce(worker: "Worker") -> None: |
| 203 | """Initialize the TP PCIe IPC backend immediately before graph capture.""" |
| 204 | from vllm.distributed.device_communicators.cuda_communicator import ( |
| 205 | CudaCommunicator, |
| 206 | ) |
| 207 | from vllm.distributed.parallel_state import get_tp_group |
| 208 | from vllm.model_executor.warmup.flashinfer_autotune_cache import ( |
| 209 | resolve_flashinfer_autotune_file, |
| 210 | ) |
| 211 | |
| 212 | tp_group = get_tp_group() |
| 213 | communicator = tp_group.device_communicator |
| 214 | if not isinstance(communicator, CudaCommunicator): |
| 215 | return |
| 216 | pcie_comm = communicator.fi_pcie_ipc_ar_comm |
| 217 | if pcie_comm is None or pcie_comm.disabled: |
| 218 | return |
| 219 | if worker.vllm_config.parallel_config.use_ubatching: |
| 220 | logger.warning_once( |
| 221 | "FlashInfer PCIe IPC all-reduce does not yet support DBO or " |
| 222 | "multi-ubatch execution; falling back to another backend." |
| 223 | ) |
| 224 | pcie_comm.disabled = True |
| 225 | return |
| 226 | |
| 227 | capture_sizes = worker.vllm_config.compilation_config.cudagraph_capture_sizes |
| 228 | if not capture_sizes: |
| 229 | return |
| 230 | |
| 231 | base_cache = resolve_flashinfer_autotune_file(worker.model_runner) |
| 232 | ranks = "-".join(str(rank) for rank in tp_group.ranks) |
| 233 | tune_cache = base_cache.with_name(f"pcie_ipc_allreduce_tp_{ranks}.json") |
| 234 | pcie_comm.setup( |
| 235 | hidden_dim=worker.model_config.get_hidden_size(), |
| 236 | dtype=worker.model_config.dtype, |
| 237 | capture_sizes=capture_sizes, |
| 238 | tune_cache=tune_cache, |
| 239 | ) |