| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import threading |
| 4 | from collections.abc import Iterator |
| 5 | from contextlib import contextmanager |
| 6 | from weakref import WeakValueDictionary |
| 7 | |
| 8 | import torch |
| 9 | import torch.distributed as dist |
| 10 | from torch.distributed import ProcessGroup |
| 11 | |
| 12 | from vllm.logger import init_logger |
| 13 | |
| 14 | logger = init_logger(__name__) |
| 15 | |
| 16 | |
| 17 | class Cache: |
| 18 | def __init__(self): |
| 19 | self._cache: WeakValueDictionary = WeakValueDictionary() |
| 20 | self._lock = threading.RLock() # Reentrant lock for thread safety |
| 21 | |
| 22 | def get_or_create(self, kwargs, func): |
| 23 | # Create a hashable key from the kwargs |
| 24 | key = tuple(sorted((k, v) for k, v in kwargs.items())) |
| 25 | |
| 26 | with self._lock: |
| 27 | instance = self._cache.get(key) |
| 28 | if instance is None: |
| 29 | instance = func(**kwargs) |
| 30 | self._cache[key] = instance |
| 31 | return instance |
| 32 | |
| 33 | |
| 34 | class All2AllManagerBase: |
| 35 | rank: int |
| 36 | world_size: int |
| 37 | |
| 38 | def __init__(self, cpu_group, tcp_store_group=None): |
| 39 | self.cpu_group = cpu_group |
| 40 | self.tcp_store_group = tcp_store_group |
| 41 | |
| 42 | # compute some common properties |
| 43 | from vllm.distributed.parallel_state import ( |
| 44 | get_dp_group, |
| 45 | get_tp_group, |
| 46 | in_the_same_node_as, |
| 47 | ) |
| 48 | |
| 49 | # all2all lives in ep group, which is merged from dp and tp group |
| 50 | self.dp_group = get_dp_group() |
| 51 | self.tp_group = get_tp_group() |
| 52 | |
| 53 | # no self.ep_group since self.ep_group is still in construction |
| 54 | # when we create this object |
| 55 | self.dp_rank = self.dp_group.rank_in_group |
| 56 | self.dp_world_size = self.dp_group.world_size |
| 57 | self.rank = cpu_group.rank() |
| 58 | self.world_size = cpu_group.size() |
| 59 | |
| 60 | # all2all communication often has separate implementations for |
| 61 | # intra-node and inter-node communication |
| 62 | if tcp_store_group is None: |
| 63 | self.internode = not all(in_the_same_node_as(cpu_group, source_rank=0)) |
| 64 | else: |
| 65 | self.internode = not all( |
| 66 | in_the_same_node_as(tcp_store_group, source_rank=0) |
| 67 | ) |
| 68 | |
| 69 | self.support_fault_tolerance = False |
| 70 | |
| 71 | def get_handle(self, kwargs): |
| 72 | # get a handle for the all2all communication, |
| 73 | # based on the kwargs. |
| 74 | # different layers can have different configs, |
| 75 | # e.g. one layer has hidden size 1024, another has 2048. |
| 76 | # usually the underlying implementation caches the handle |
| 77 | # and reuse it for the same config. |
| 78 | raise NotImplementedError |
| 79 | |
| 80 | def dispatch_router_logits( |
| 81 | self, |
| 82 | hidden_states: torch.Tensor, |
| 83 | router_logits: torch.Tensor, |
| 84 | is_sequence_parallel: bool = False, |
| 85 | extra_tensors: list[torch.Tensor] | None = None, |
| 86 | ) -> ( |
| 87 | tuple[torch.Tensor, torch.Tensor] |
| 88 | | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 89 | ): |
| 90 | # Subclasses should either: |
| 91 | # - implement handling for extra_tensors, or |
| 92 | # - raise a clear error if extra_tensors is not supported. |
| 93 | raise NotImplementedError |
| 94 | |
| 95 | def dispatch( |
| 96 | self, |
| 97 | hidden_states: torch.Tensor, |
| 98 | topk_weights: torch.Tensor, |
| 99 | topk_ids: torch.Tensor, |
| 100 | is_sequence_parallel: bool = False, |
| 101 | extra_tensors: list[torch.Tensor] | None = None, |
| 102 | ) -> ( |
| 103 | tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 104 | | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 105 | ): |
| 106 | # Subclasses should either: |
| 107 | # - implement handling for extra_tensors, or |
| 108 | # - raise a clear error if extra_tensors is not supported. |
| 109 | raise NotImplementedError |
| 110 | |
| 111 | def query_active_mask(self) -> torch.Tensor: |
| 112 | """Return the all2all liveness mask for the EP ranks. |
| 113 | |
| 114 | Returns: |
| 115 | An int32 device tensor where 0 marks a live rank and 1 marks a |
| 116 | masked (dead/unreachable) rank. |
| 117 | |
| 118 | """ |
| 119 | raise NotImplementedError |
| 120 | |
| 121 | def query_fault(self) -> torch.Tensor: |
| 122 | """Return a scalar bool tensor, True if a new fault appeared. |
| 123 | |
| 124 | Compares the current mask against the baseline recorded at the last |
| 125 | recovery point. |
| 126 | """ |
| 127 | raise NotImplementedError |
| 128 | |
| 129 | def clean_buffers(self) -> None: |
| 130 | """Reset this rank's RDMA buffers and all2all mask state (rank-local). |
| 131 | |
| 132 | Post-fault cleanup: a dispatch/combine that hit a dead peer or timed |
| 133 | out can leave partially-written or stale tokens in the RDMA receive |
| 134 | buffer, so it is zeroed to stop the next forward from reading that |
| 135 | contaminated data. |
| 136 | """ |
| 137 | raise NotImplementedError |
| 138 | |
| 139 | def set_num_sms(self, num_sms: int): |
| 140 | pass |
| 141 | |
| 142 | def max_sms_used(self) -> int | None: |
| 143 | return None # None means it could use the whole GPU |
| 144 | |
| 145 | def checkpoint_prepare(self) -> None: |
| 146 | logger.warning_once( |
| 147 | "%s.checkpoint_prepare is not implemented; skipping.", |
| 148 | type(self).__name__, |
| 149 | ) |
| 150 | |
| 151 | def checkpoint_restore(self) -> None: |
| 152 | logger.warning_once( |
| 153 | "%s.checkpoint_restore is not implemented; skipping.", |
| 154 | type(self).__name__, |
| 155 | ) |
| 156 | |
| 157 | def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False): |
| 158 | raise NotImplementedError |
| 159 | |
| 160 | def destroy(self): |
| 161 | pass |
| 162 | |
| 163 | def stage_ep_size(self) -> None: |
| 164 | pass |
| 165 | |
| 166 | def commit_ep_size(self) -> None: |
| 167 | pass |
| 168 | |
| 169 | @contextmanager |
| 170 | def mask_remote_ranks(self) -> Iterator[None]: |
| 171 | yield |
| 172 | |
| 173 | |
| 174 | class DeviceCommunicatorBase: |
| 175 | """Base class for device-specific communicator. |
| 176 | It can use the `cpu_group` to initialize the communicator. |
| 177 | If the device has PyTorch integration (PyTorch can recognize its |
| 178 | communication backend), the `device_group` will also be given. |
| 179 | """ |
| 180 | |
| 181 | def __init__( |
| 182 | self, |
| 183 | cpu_group: ProcessGroup, |
| 184 | device: torch.device | None = None, |
| 185 | device_group: ProcessGroup | None = None, |
| 186 | unique_name: str = "", |
| 187 | global_ranks: list[int] | None = None, |
| 188 | global_world_size: int | None = None, |
| 189 | use_all2all: bool = False, |
| 190 | ): |
| 191 | self.device = device or torch.device("cpu") |
| 192 | self.cpu_group = cpu_group |
| 193 | self.device_group = device_group |
| 194 | self.unique_name = unique_name |
| 195 | |
| 196 | # Check if this is a stateless process group |
| 197 | from torch.distributed.distributed_c10d import _world |
| 198 | |
| 199 | is_stateless = _world.pg_map.get(cpu_group, None) is None |
| 200 | |
| 201 | if is_stateless: |
| 202 | # For stateless groups, we can't use torch.distributed methods |
| 203 | self.rank = cpu_group.rank() |
| 204 | self.world_size = cpu_group.size() |
| 205 | assert global_ranks is not None |
| 206 | assert global_world_size is not None |
| 207 | self.ranks = global_ranks |
| 208 | self.global_rank = self.ranks[self.rank] |
| 209 | self.global_world_size = global_world_size |
| 210 | self.rank_in_group = self.rank |
| 211 | else: |
| 212 | self.rank = dist.get_rank(cpu_group) |
| 213 | self.world_size = dist.get_world_size(cpu_group) |
| 214 | self.ranks = dist.get_process_group_ranks(cpu_group) |
| 215 | self.global_rank = dist.get_rank() |
| 216 | self.global_world_size = dist.get_world_size() |
| 217 | self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank) |
| 218 | |
| 219 | all2all_backend = None |
| 220 | from vllm.config import get_current_vllm_config_or_none |
| 221 | |
| 222 | config = get_current_vllm_config_or_none() |
| 223 | if config is not None: |
| 224 | all2all_backend = config.parallel_config.all2all_backend |
| 225 | |
| 226 | self.is_ep_communicator = unique_name.split(":")[0] == "ep" |
| 227 | self.use_all2all = self.is_ep_communicator and use_all2all |
| 228 | self.all2all_backend = all2all_backend |
| 229 | self.all2all_manager: All2AllManagerBase | None = None |
| 230 | |
| 231 | def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: |
| 232 | dist.all_reduce(input_, group=self.device_group) |
| 233 | return input_ |
| 234 | |
| 235 | def checkpoint_prepare(self) -> None: |
| 236 | """Prepare reclaimable communicator state for checkpoint (default: no-op).""" |
| 237 | |
| 238 | def checkpoint_restore(self) -> None: |
| 239 | """Restore communicator state after checkpoint (default: no-op).""" |
| 240 | |
| 241 | def suspend(self) -> None: |
| 242 | """Release reclaimable communicator memory (default: no-op).""" |
| 243 | |
| 244 | def resume(self) -> None: |
| 245 | """Restore memory released by ``suspend`` (default: no-op).""" |
| 246 | |
| 247 | def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: |
| 248 | if dim < 0: |
| 249 | # Convert negative dim to positive. |
| 250 | dim += input_.dim() |
| 251 | input_size = input_.size() |
| 252 | # NOTE: we have to use concat-style all-gather here, |
| 253 | # stack-style all-gather has compatibility issues with |
| 254 | # torch.compile . see https://github.com/pytorch/pytorch/issues/138795 |
| 255 | output_size = (input_size[0] * self.world_size,) + input_size[1:] |
| 256 | # Allocate output tensor. |
| 257 | output_tensor = torch.empty( |
| 258 | output_size, dtype=input_.dtype, device=input_.device |
| 259 | ) |
| 260 | # All-gather. |
| 261 | dist.all_gather_into_tensor(output_tensor, input_, group=self.device_group) |
| 262 | # Reshape |
| 263 | output_tensor = output_tensor.reshape((self.world_size,) + input_size) |
| 264 | output_tensor = output_tensor.movedim(0, dim) |
| 265 | output_tensor = output_tensor.reshape( |
| 266 | input_size[:dim] |
| 267 | + (self.world_size * input_size[dim],) |
| 268 | + input_size[dim + 1 :] |
| 269 | ) |
| 270 | return output_tensor |
| 271 | |
| 272 | def all_gatherv( |
| 273 | self, |
| 274 | input_: torch.Tensor | list[torch.Tensor], |
| 275 | dim: int = 0, |
| 276 | sizes: list[int] | None = None, |
| 277 | ) -> torch.Tensor | list[torch.Tensor]: |
| 278 | raise NotImplementedError |
| 279 | |
| 280 | def reduce_scatter(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: |
| 281 | world_size = self.world_size |
| 282 | # Bypass the function if we are using only 1 GPU. |
| 283 | if world_size == 1: |
| 284 | return input_ |
| 285 | assert -input_.dim() <= dim < input_.dim(), ( |
| 286 | f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" |
| 287 | ) |
| 288 | |
| 289 | if dim < 0: |
| 290 | # Convert negative dim to positive. |
| 291 | dim += input_.dim() |
| 292 | |
| 293 | # Note: This will produce an incorrect answer if we don't make |
| 294 | # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? |
| 295 | input_tensor = input_.movedim(0, dim).contiguous() |
| 296 | |
| 297 | assert input_tensor.shape[0] % world_size == 0 |
| 298 | chunk_size = input_tensor.shape[0] // world_size |
| 299 | output_shape = (chunk_size,) + input_tensor.shape[1:] |
| 300 | |
| 301 | output_tensor = torch.empty( |
| 302 | output_shape, dtype=input_tensor.dtype, device=input_tensor.device |
| 303 | ) |
| 304 | |
| 305 | # Perform reduce-scatter operation |
| 306 | torch.distributed.reduce_scatter_tensor( |
| 307 | output_tensor, input_tensor, group=self.device_group |
| 308 | ) |
| 309 | |
| 310 | # Reshape before returning |
| 311 | return output_tensor.movedim(0, dim).contiguous() |
| 312 | |
| 313 | def reduce_scatterv( |
| 314 | self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None |
| 315 | ) -> torch.Tensor: |
| 316 | raise NotImplementedError |
| 317 | |
| 318 | def gather( |
| 319 | self, input_: torch.Tensor, dst: int = 0, dim: int = -1 |
| 320 | ) -> torch.Tensor | None: |
| 321 | """NOTE: We assume that the input tensor is on the same device across |
| 322 | all the ranks. |
| 323 | NOTE: `dst` is the local rank of the destination rank. |
| 324 | """ |
| 325 | world_size = self.world_size |
| 326 | assert -input_.dim() <= dim < input_.dim(), ( |
| 327 | f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" |
| 328 | ) |
| 329 | if dim < 0: |
| 330 | # Convert negative dim to positive. |
| 331 | dim += input_.dim() |
| 332 | |
| 333 | # Allocate output tensor. |
| 334 | if self.rank_in_group == dst: |
| 335 | gather_list = [torch.empty_like(input_) for _ in range(world_size)] |
| 336 | else: |
| 337 | gather_list = None |
| 338 | # Gather. |
| 339 | torch.distributed.gather( |
| 340 | input_, gather_list, dst=self.ranks[dst], group=self.device_group |
| 341 | ) |
| 342 | if self.rank_in_group == dst: |
| 343 | output_tensor = torch.cat(gather_list, dim=dim) |
| 344 | else: |
| 345 | output_tensor = None |
| 346 | return output_tensor |
| 347 | |
| 348 | def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: |
| 349 | """Sends a tensor to the destination rank in a blocking way.""" |
| 350 | """NOTE: `dst` is the local rank of the destination rank.""" |
| 351 | if dst is None: |
| 352 | dst = (self.rank_in_group + 1) % self.world_size |
| 353 | torch.distributed.send(tensor, self.ranks[dst], self.device_group) |
| 354 | |
| 355 | def recv( |
| 356 | self, size: torch.Size, dtype: torch.dtype, src: int | None = None |
| 357 | ) -> torch.Tensor: |
| 358 | """Receives a tensor from the source rank.""" |
| 359 | """NOTE: `src` is the local rank of the source rank.""" |
| 360 | if src is None: |
| 361 | src = (self.rank_in_group - 1) % self.world_size |
| 362 | |
| 363 | tensor = torch.empty(size, dtype=dtype, device=self.device) |
| 364 | torch.distributed.recv(tensor, self.ranks[src], self.device_group) |
| 365 | return tensor |
| 366 | |
| 367 | def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor: |
| 368 | """Broadcast a tensor from source rank to all ranks.""" |
| 369 | if self.world_size == 1: |
| 370 | return tensor |
| 371 | torch.distributed.broadcast(tensor, self.ranks[src], self.device_group) |
| 372 | return tensor |
| 373 | |
| 374 | def destroy(self): |
| 375 | pass |
| 376 | |
| 377 | def dispatch_router_logits( |
| 378 | self, |
| 379 | hidden_states: torch.Tensor, |
| 380 | router_logits: torch.Tensor, |
| 381 | is_sequence_parallel: bool = False, |
| 382 | extra_tensors: list[torch.Tensor] | None = None, |
| 383 | ) -> ( |
| 384 | tuple[torch.Tensor, torch.Tensor] |
| 385 | | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 386 | ): |
| 387 | """Dispatch the hidden states and router logits to the appropriate device. |
| 388 | This is a no-op in the base class. |
| 389 | """ |
| 390 | if extra_tensors is not None: |
| 391 | return hidden_states, router_logits, extra_tensors |
| 392 | return hidden_states, router_logits |
| 393 | |
| 394 | def dispatch( |
| 395 | self, |
| 396 | hidden_states: torch.Tensor, |
| 397 | topk_weights: torch.Tensor, |
| 398 | topk_ids: torch.Tensor, |
| 399 | is_sequence_parallel: bool = False, |
| 400 | extra_tensors: list[torch.Tensor] | None = None, |
| 401 | ) -> ( |
| 402 | tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 403 | | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 404 | ): |
| 405 | """Dispatch the hidden states and topk weights/ids to the appropriate device. |
| 406 | This is a no-op in the base class. |
| 407 | """ |
| 408 | if extra_tensors is not None: |
| 409 | return hidden_states, topk_weights, topk_ids, extra_tensors |
| 410 | return hidden_states, topk_weights, topk_ids |
| 411 | |
| 412 | def combine( |
| 413 | self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False |
| 414 | ) -> torch.Tensor: |
| 415 | """Combine the hidden states and router logits from the appropriate device. |
| 416 | This is a no-op in the base class. |
| 417 | """ |
| 418 | return hidden_states |
| 419 | |
| 420 | def batch_isend_irecv(self, p2p_ops: list): |
| 421 | raise NotImplementedError |