| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from contextlib import contextmanager, nullcontext |
| 5 | from contextvars import ContextVar |
| 6 | |
| 7 | import torch |
| 8 | import torch.distributed as dist |
| 9 | |
| 10 | from vllm.config import ParallelConfig |
| 11 | from vllm.distributed.parallel_state import get_dp_group |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 14 | from vllm.utils.torch_utils import PIN_MEMORY |
| 15 | from vllm.v1.worker.ubatch_utils import ( |
| 16 | check_ubatch_thresholds, |
| 17 | is_last_ubatch_empty, |
| 18 | ) |
| 19 | |
| 20 | logger = init_logger(__name__) |
| 21 | |
| 22 | _skip_dp_coordination = ContextVar("skip_dp_coordination", default=False) |
| 23 | |
| 24 | |
| 25 | @contextmanager |
| 26 | def skip_dp_coordination(): |
| 27 | """Run without coordinating DP metadata with other DP ranks.""" |
| 28 | token = _skip_dp_coordination.set(True) |
| 29 | try: |
| 30 | yield |
| 31 | finally: |
| 32 | _skip_dp_coordination.reset(token) |
| 33 | |
| 34 | |
| 35 | def should_skip_dp_coordination() -> bool: |
| 36 | return _skip_dp_coordination.get() |
| 37 | |
| 38 | |
| 39 | def _get_device_and_group(parallel_config: ParallelConfig): |
| 40 | # Use the actual device assigned to the DP group, not just the device type |
| 41 | device = get_dp_group().device |
| 42 | group = get_dp_group().device_group |
| 43 | |
| 44 | # Transferring this tensor from GPU to CPU will introduce a GPU sync |
| 45 | # point that could adversely affect performance of vllm with asynch |
| 46 | # scheduling. This environment variable exists to quickly disable |
| 47 | # this optimization if we run into this case. |
| 48 | if parallel_config.disable_nccl_for_dp_synchronization: |
| 49 | logger.info_once( |
| 50 | "Using CPU all reduce to synchronize DP padding between ranks.", |
| 51 | ) |
| 52 | device = "cpu" |
| 53 | group = get_dp_group().cpu_group |
| 54 | return device, group |
| 55 | |
| 56 | |
| 57 | def _run_ar( |
| 58 | should_ubatch: bool, |
| 59 | orig_num_tokens_per_ubatch: int, |
| 60 | padded_num_tokens_per_ubatch: int, |
| 61 | cudagraph_mode: int, |
| 62 | parallel_config: ParallelConfig, |
| 63 | ) -> torch.Tensor: |
| 64 | dp_size = parallel_config.data_parallel_size |
| 65 | dp_rank = parallel_config.data_parallel_rank |
| 66 | device, group = _get_device_and_group(parallel_config) |
| 67 | # Populate this rank's contribution on CPU to reduce GPU syncs. |
| 68 | pin_memory = PIN_MEMORY and torch.device(device).type == "cuda" |
| 69 | tensor_cpu = torch.zeros(4, dp_size, dtype=torch.int32, pin_memory=pin_memory) |
| 70 | tensor_cpu[0][dp_rank] = orig_num_tokens_per_ubatch |
| 71 | tensor_cpu[1][dp_rank] = padded_num_tokens_per_ubatch |
| 72 | tensor_cpu[2][dp_rank] = 1 if should_ubatch else 0 |
| 73 | tensor_cpu[3][dp_rank] = cudagraph_mode |
| 74 | tensor = tensor_cpu.to(device, non_blocking=True) |
| 75 | dist.all_reduce(tensor, group=group) |
| 76 | return tensor |
| 77 | |
| 78 | |
| 79 | def _post_process_ubatch(tensor: torch.Tensor, num_ubatches: int) -> bool: |
| 80 | orig_num_tokens_tensor = tensor[0, :] |
| 81 | padded_num_tokens_tensor = tensor[1, :] |
| 82 | |
| 83 | # First determine if we are going to be ubatching. |
| 84 | should_ubatch: bool = bool(torch.all(tensor[2] == 1).item()) |
| 85 | if not should_ubatch: |
| 86 | return False |
| 87 | # If the DP ranks are planning to ubatch, make sure that |
| 88 | # there are no "empty" second ubatches |
| 89 | orig_min_num_tokens = int(orig_num_tokens_tensor.min().item()) |
| 90 | padded_max_num_tokens = int(padded_num_tokens_tensor.max().item()) |
| 91 | if is_last_ubatch_empty(orig_min_num_tokens, padded_max_num_tokens, num_ubatches): |
| 92 | logger.debug( |
| 93 | "Aborting ubatching %s %s", orig_min_num_tokens, padded_max_num_tokens |
| 94 | ) |
| 95 | should_ubatch = False |
| 96 | return should_ubatch |
| 97 | |
| 98 | |
| 99 | def _post_process_dp_padding(tensor: torch.Tensor, should_dp_pad: bool) -> torch.Tensor: |
| 100 | num_tokens_across_dp = tensor[1, :] |
| 101 | if should_dp_pad: |
| 102 | # If DP padding is enabled, ensure that each rank is processing the same number |
| 103 | # of tokens |
| 104 | max_num_tokens = int(num_tokens_across_dp.max().item()) |
| 105 | return torch.tensor( |
| 106 | [max_num_tokens] * len(num_tokens_across_dp), |
| 107 | device="cpu", |
| 108 | dtype=torch.int32, |
| 109 | ) |
| 110 | else: |
| 111 | return num_tokens_across_dp.cpu() |
| 112 | |
| 113 | |
| 114 | def _post_process_cudagraph_mode(tensor: torch.Tensor) -> int: |
| 115 | """Synchronize cudagraph_mode across DP ranks by taking the minimum. |
| 116 | If any rank has NONE (0), all ranks use NONE. |
| 117 | This ensures all ranks send consistent values (all padded or all unpadded). |
| 118 | """ |
| 119 | return int(tensor[3, :].min().item()) |
| 120 | |
| 121 | |
| 122 | def _synchronize_dp_ranks( |
| 123 | num_tokens_unpadded: int, |
| 124 | num_tokens_padded: int, |
| 125 | should_attempt_ubatching: bool, |
| 126 | cudagraph_mode: int, |
| 127 | parallel_config: ParallelConfig, |
| 128 | ) -> tuple[bool, torch.Tensor | None, int]: |
| 129 | """1. Decides if each DP rank is going to microbatch. Either all ranks |
| 130 | run with microbatching or none of them do. |
| 131 | |
| 132 | 2. Determines the total number of tokens that each rank will run. |
| 133 | When running microbatched or if cudagraph is enabled (synced across ranks), |
| 134 | all ranks will be padded out so that they run with the same number of tokens. |
| 135 | |
| 136 | 3. Synchronizes cudagraph_mode across ranks by taking the minimum. |
| 137 | |
| 138 | Returns: tuple[ |
| 139 | should_ubatch: Are all DP ranks going to microbatch |
| 140 | num_tokens_after_padding: A tensor containing the total number of |
| 141 | tokens per-microbatch for each DP rank including any DP padding. |
| 142 | synced_cudagraph_mode: The synchronized cudagraph mode (min across ranks) |
| 143 | ] |
| 144 | |
| 145 | """ |
| 146 | assert num_tokens_padded >= num_tokens_unpadded |
| 147 | |
| 148 | # Coordinate between the DP ranks via an All Reduce |
| 149 | # to determine the total number of tokens that each rank |
| 150 | # will run and if we are using ubatching or not. |
| 151 | tensor = _run_ar( |
| 152 | should_ubatch=should_attempt_ubatching, |
| 153 | orig_num_tokens_per_ubatch=num_tokens_unpadded, |
| 154 | padded_num_tokens_per_ubatch=num_tokens_padded, |
| 155 | cudagraph_mode=cudagraph_mode, |
| 156 | parallel_config=parallel_config, |
| 157 | ) |
| 158 | |
| 159 | # Only the NCCL path leaves `tensor` on device. With Gloo -- the default |
| 160 | # under async scheduling -- the all-reduce runs on CPU, so the reads below |
| 161 | # are host-side and the check should stay armed. |
| 162 | with ( |
| 163 | nullcontext() |
| 164 | if parallel_config.disable_nccl_for_dp_synchronization |
| 165 | else gpu_sync_allowed() |
| 166 | ): |
| 167 | # Synchronize cudagraph_mode across ranks first (take min). |
| 168 | # This is needed before DP padding decision since we use the synced |
| 169 | # cudagraph mode to determine whether DP padding is needed. |
| 170 | synced_cudagraph_mode = _post_process_cudagraph_mode(tensor) |
| 171 | |
| 172 | # Check conditions for microbatching |
| 173 | should_ubatch = _post_process_ubatch(tensor, parallel_config.num_ubatches) |
| 174 | |
| 175 | # DP padding is needed when cudagraph is enabled (synced across ranks) |
| 176 | # or when ubatching/DBO is active (ubatching requires uniform batch |
| 177 | # sizes across DP ranks currently). |
| 178 | # Use the synced runtime cudagraph mode rather than the compilation |
| 179 | # config so we can avoid padding when cudagraph is not enabled for |
| 180 | # this step. |
| 181 | should_dp_pad = synced_cudagraph_mode != 0 or should_ubatch |
| 182 | |
| 183 | # Pad all DP ranks up to the maximum token count across ranks if |
| 184 | # should_dp_pad is True |
| 185 | num_tokens_after_padding = _post_process_dp_padding(tensor, should_dp_pad) |
| 186 | |
| 187 | return should_ubatch, num_tokens_after_padding, synced_cudagraph_mode |
| 188 | |
| 189 | |
| 190 | def coordinate_batch_across_dp( |
| 191 | num_tokens_unpadded: int, |
| 192 | allow_microbatching: bool, |
| 193 | parallel_config: ParallelConfig, |
| 194 | num_tokens_padded: int | None = None, |
| 195 | uniform_decode: bool | None = None, |
| 196 | cudagraph_mode: int = 0, |
| 197 | ) -> tuple[bool, torch.Tensor | None, int]: |
| 198 | """Coordinates amongst all DP ranks to determine if and how the full batch |
| 199 | should be split into microbatches. |
| 200 | |
| 201 | Args: |
| 202 | num_tokens_unpadded: Number of tokens without accounting for padding |
| 203 | allow_microbatching: If microbatching should be attempted |
| 204 | parallel_config: The parallel config |
| 205 | num_tokens_padded: Number of tokens including any non-DP padding (CUDA graphs, |
| 206 | TP, etc) |
| 207 | uniform_decode: Only used if allow_microbatching is True. True if the batch |
| 208 | only contains single token decodes |
| 209 | cudagraph_mode: The cudagraph mode for this rank (0=NONE, 1=PIECEWISE, 2=FULL). |
| 210 | DP padding is enabled when synced cudagraph mode across ranks is not NONE. |
| 211 | |
| 212 | Returns: tuple[ |
| 213 | ubatch_slices: if this is set then all DP ranks have agreed to |
| 214 | microbatch |
| 215 | num_tokens_after_padding: A tensor containing the total number of |
| 216 | tokens per-microbatch for each DP rank including padding. Will be |
| 217 | padded up to the max value across all DP ranks when cudagraph is enabled. |
| 218 | synced_cudagraph_mode: The synchronized cudagraph mode (min across ranks) |
| 219 | ] |
| 220 | |
| 221 | """ |
| 222 | if parallel_config.data_parallel_size == 1: |
| 223 | # Early exit. |
| 224 | return False, None, cudagraph_mode |
| 225 | |
| 226 | # If the caller has explicitly enabled microbatching. |
| 227 | should_attempt_ubatching = False |
| 228 | if allow_microbatching: |
| 229 | # Check preconditions for microbatching |
| 230 | assert uniform_decode is not None |
| 231 | should_attempt_ubatching = check_ubatch_thresholds( |
| 232 | parallel_config, |
| 233 | num_tokens_unpadded, |
| 234 | uniform_decode=uniform_decode, |
| 235 | ) |
| 236 | |
| 237 | if num_tokens_padded is None: |
| 238 | num_tokens_padded = num_tokens_unpadded |
| 239 | |
| 240 | if should_skip_dp_coordination(): |
| 241 | should_ubatch = should_attempt_ubatching and not is_last_ubatch_empty( |
| 242 | num_tokens_unpadded, |
| 243 | num_tokens_padded, |
| 244 | parallel_config.num_ubatches, |
| 245 | ) |
| 246 | num_tokens_after_padding = torch.full( |
| 247 | (parallel_config.data_parallel_size,), |
| 248 | num_tokens_padded, |
| 249 | dtype=torch.int32, |
| 250 | ) |
| 251 | return should_ubatch, num_tokens_after_padding, cudagraph_mode |
| 252 | |
| 253 | (should_ubatch, num_tokens_after_padding, synced_cudagraph_mode) = ( |
| 254 | _synchronize_dp_ranks( |
| 255 | num_tokens_unpadded, |
| 256 | num_tokens_padded, |
| 257 | should_attempt_ubatching, |
| 258 | cudagraph_mode, |
| 259 | parallel_config, |
| 260 | ) |
| 261 | ) |
| 262 | |
| 263 | return (should_ubatch, num_tokens_after_padding, synced_cudagraph_mode) |