| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Detect unintended GPU<->CPU syncs in the hot path. |
| 4 | |
| 5 | `torch.cuda.set_sync_debug_mode` is process-global, so we arm it at "warn" |
| 6 | (which never raises) and decide in `_sync_warning_hook` whether a given sync |
| 7 | is a failure. The scoping lives in ContextVars, which are per-thread and |
| 8 | per-asyncio-task, so an allow region opened on one thread is invisible to |
| 9 | every other one. |
| 10 | |
| 11 | torch's mode only instruments explicit synchronizing calls, so it misses |
| 12 | `non_blocking=True` CPU<->CUDA copies that the CUDA driver silently stages |
| 13 | through pageable memory (when the CPU tensor is not pinned or not densely |
| 14 | laid out). Those are caught by the wrappers on `Tensor.to`/`Tensor.cuda`/ |
| 15 | `Tensor.copy_` installed by `_install_copy_checkers()`. |
| 16 | """ |
| 17 | |
| 18 | import functools |
| 19 | import sys |
| 20 | import threading |
| 21 | import warnings |
| 22 | from contextlib import contextmanager, nullcontext |
| 23 | from contextvars import ContextVar |
| 24 | |
| 25 | import torch |
| 26 | |
| 27 | import vllm.envs as envs |
| 28 | from vllm.platforms import current_platform |
| 29 | |
| 30 | SYNC_ERROR_MESSAGE = ( |
| 31 | "GPU<->CPU sync detected - avoid it or wrap with gpu_sync_allowed()" |
| 32 | ) |
| 33 | |
| 34 | # Warning text torch emits in "warn" mode. |
| 35 | _TORCH_SYNC_WARNING = "called a synchronizing CUDA operation" |
| 36 | |
| 37 | _GPU_SYNC_ALLOWED_FIRST_SEEN: set[tuple[str, int]] = set() |
| 38 | |
| 39 | # Read once: `envs.__getattr__` re-evaluates the env lambda on every access, |
| 40 | # which dominates the cost of a disabled `gpu_sync_allowed()`. Tests toggling |
| 41 | # the env var must patch this attribute instead. |
| 42 | _SYNC_CHECK_MODE: str | None = ( |
| 43 | envs.VLLM_GPU_SYNC_CHECK if current_platform.is_cuda_alike() else None |
| 44 | ) |
| 45 | |
| 46 | # Configured mode while a thread is inside a checked call, else None. |
| 47 | _checking: ContextVar[str | None] = ContextVar("vllm_gpu_sync_checking", default=None) |
| 48 | # Depth of nested `gpu_sync_allowed()` regions on this thread. |
| 49 | _allow_depth: ContextVar[int] = ContextVar("vllm_gpu_sync_allow_depth", default=0) |
| 50 | |
| 51 | # Off during engine setup so model load, KV cache init and warmup may sync |
| 52 | # freely; flipped on by `enable_gpu_sync_check()` once warmup completes. |
| 53 | _sync_check_enabled: bool = False |
| 54 | |
| 55 | |
| 56 | def enable_gpu_sync_check() -> None: |
| 57 | """Flip the sync-check gate on, once per worker, after warmup.""" |
| 58 | if _SYNC_CHECK_MODE is None: |
| 59 | return |
| 60 | global _sync_check_enabled |
| 61 | _sync_check_enabled = True |
| 62 | _install_compile_time_sync_suppressors() |
| 63 | _install_copy_checkers() |
| 64 | |
| 65 | |
| 66 | _arm_lock = threading.Lock() |
| 67 | _arm_count: int = 0 |
| 68 | _saved_sync_debug_mode: int = 0 |
| 69 | _prev_showwarning = warnings.showwarning |
| 70 | _sync_filter_head: tuple | None = None |
| 71 | |
| 72 | |
| 73 | def _sync_warning_hook(message, category, filename, lineno, file=None, line=None): |
| 74 | """Turn torch's sync warning into an error, but only on a thread that is |
| 75 | being checked and is outside any allow region.""" |
| 76 | if _TORCH_SYNC_WARNING in str(message): |
| 77 | mode = _checking.get() |
| 78 | if mode is None or _allow_depth.get(): |
| 79 | return None |
| 80 | if mode == "error": |
| 81 | raise RuntimeError(SYNC_ERROR_MESSAGE) |
| 82 | return _prev_showwarning(message, category, filename, lineno, file, line) |
| 83 | |
| 84 | |
| 85 | def _install_warning_hook() -> None: |
| 86 | """(Re)install the hook and a filter that lets torch's warning reach it. |
| 87 | |
| 88 | Done per checked call because pytest runs each test inside |
| 89 | `warnings.catch_warnings()`, which restores both `showwarning` and |
| 90 | `filters`. The hook is left in place afterwards: outside a checked call |
| 91 | the debug mode is disarmed, so torch emits nothing for it to see. |
| 92 | """ |
| 93 | global _prev_showwarning, _sync_filter_head |
| 94 | if warnings.showwarning is not _sync_warning_hook: |
| 95 | _prev_showwarning = warnings.showwarning |
| 96 | warnings.showwarning = _sync_warning_hook |
| 97 | # "always" so the warning survives filtering and isn't deduplicated by |
| 98 | # `__warningregistry__`. `filterwarnings` prepends, so only re-assert it |
| 99 | # once ours is no longer in front. |
| 100 | if warnings.filters[:1] != [_sync_filter_head]: |
| 101 | warnings.filterwarnings( |
| 102 | "always", message=_TORCH_SYNC_WARNING, category=UserWarning |
| 103 | ) |
| 104 | _sync_filter_head = warnings.filters[0] |
| 105 | |
| 106 | |
| 107 | @contextmanager |
| 108 | def _checked_region(mode: str): |
| 109 | """Police syncs on this thread for the duration of the block. |
| 110 | |
| 111 | The debug mode is armed per call so that syncs outside a checked region |
| 112 | emit nothing, and refcounted because `execute_model` and `sample_tokens` |
| 113 | nest. "warn" rather than "error" because torch's error mode raises on |
| 114 | whichever thread synced, with no way to exempt one. |
| 115 | """ |
| 116 | global _arm_count, _saved_sync_debug_mode |
| 117 | _install_warning_hook() |
| 118 | with _arm_lock: |
| 119 | if _arm_count == 0: |
| 120 | _saved_sync_debug_mode = torch.cuda.get_sync_debug_mode() |
| 121 | torch.cuda.set_sync_debug_mode("warn") |
| 122 | _arm_count += 1 |
| 123 | token = _checking.set(mode) |
| 124 | try: |
| 125 | yield |
| 126 | finally: |
| 127 | _checking.reset(token) |
| 128 | with _arm_lock: |
| 129 | _arm_count -= 1 |
| 130 | if _arm_count == 0: |
| 131 | torch.cuda.set_sync_debug_mode(_saved_sync_debug_mode) |
| 132 | |
| 133 | |
| 134 | @contextmanager |
| 135 | def _allow_syncs(): |
| 136 | token = _allow_depth.set(_allow_depth.get() + 1) |
| 137 | try: |
| 138 | yield |
| 139 | finally: |
| 140 | _allow_depth.reset(token) |
| 141 | |
| 142 | |
| 143 | # Shared, since `nullcontext` is stateless: reusing one instance avoids an |
| 144 | # allocation per call, which is most of the cost when the check is disabled. |
| 145 | _NOOP_CM = nullcontext() |
| 146 | |
| 147 | _compile_time_suppressors_installed: bool = False |
| 148 | |
| 149 | |
| 150 | def _suppressing(fn): |
| 151 | """Allow the syncs `fn` performs on its calling thread.""" |
| 152 | |
| 153 | @functools.wraps(fn) |
| 154 | def wrapper(*args, **kwargs): |
| 155 | # Not `gpu_sync_allowed()`, which no-ops while |
| 156 | # `torch.compiler.is_compiling()` -- exactly when these run. |
| 157 | with _allow_syncs(): |
| 158 | return fn(*args, **kwargs) |
| 159 | |
| 160 | return wrapper |
| 161 | |
| 162 | |
| 163 | def _install_compile_time_sync_suppressors() -> None: |
| 164 | """Allow the syncs torch's compile passes perform. |
| 165 | |
| 166 | Warmup-time compiles run before the gate flips, but lazy ones fire inside |
| 167 | `execute_model`. |
| 168 | """ |
| 169 | global _compile_time_suppressors_installed |
| 170 | if _compile_time_suppressors_installed: |
| 171 | return |
| 172 | _compile_time_suppressors_installed = True |
| 173 | |
| 174 | try: |
| 175 | from torch._inductor.fx_passes import joint_graph as _jg |
| 176 | |
| 177 | orig = _jg.joint_graph_passes |
| 178 | wrapped = _suppressing(orig) |
| 179 | # `compile_fx` imports this by value, so patching the defining module |
| 180 | # alone misses that rebind; patch every compile-time module still |
| 181 | # holding the original. |
| 182 | _jg.joint_graph_passes = wrapped |
| 183 | for name, mod in list(sys.modules.items()): |
| 184 | if ( |
| 185 | mod is not None |
| 186 | and name.startswith( |
| 187 | ("torch._inductor", "torch._functorch", "torch._dynamo") |
| 188 | ) |
| 189 | and getattr(mod, "joint_graph_passes", None) is orig |
| 190 | ): |
| 191 | setattr(mod, "joint_graph_passes", wrapped) # noqa: B010 |
| 192 | except Exception: # pragma: no cover |
| 193 | pass |
| 194 | |
| 195 | try: |
| 196 | # Inductor builds its cudagraph tree lazily, so `deferred_cudagraphify` |
| 197 | # and the `capture_begin` sync inside it can fire during |
| 198 | # `execute_model`. It resolves `cudagraphify` as a module global at |
| 199 | # call time, so patching the attribute is enough. |
| 200 | from torch._inductor import cudagraph_trees as _ct |
| 201 | |
| 202 | _ct.cudagraphify = _suppressing(_ct.cudagraphify) |
| 203 | except Exception: # pragma: no cover |
| 204 | pass |
| 205 | |
| 206 | |
| 207 | # --------------------------------------------------------------------------- |
| 208 | # Implicit sync detection in `non_blocking` CPU<->CUDA copies. |
| 209 | # |
| 210 | # A `non_blocking=True` copy between CPU and CUDA still blocks the host when |
| 211 | # the CPU side is not pinned or not densely laid out (gapped strided slices |
| 212 | # stage through a pageable temp): D2H copies into such memory wait for the |
| 213 | # stream, and H2D copies from it do too beyond a driver size threshold. |
| 214 | # None of that goes through an API torch's sync debug mode instruments, so |
| 215 | # wrap the Python-level transfer entry points and flag such copies on threads |
| 216 | # in a checked region. |
| 217 | |
| 218 | IMPLICIT_SYNC_ERROR_MESSAGE = ( |
| 219 | "Implicit GPU<->CPU sync detected - use a pinned, densely laid out CPU " |
| 220 | "tensor, or wrap with gpu_sync_allowed()" |
| 221 | ) |
| 222 | |
| 223 | # Captured at import; `_install_copy_checkers` replaces these attributes. |
| 224 | _torch_to = torch.Tensor.to |
| 225 | _torch_cuda = torch.Tensor.cuda |
| 226 | _torch_copy_ = torch.Tensor.copy_ |
| 227 | |
| 228 | |
| 229 | # Once ``Tensor.to`` is replaced below, Dynamo sees the saved C++ descriptor |
| 230 | # through this module global instead of as a recognized tensor method. Put the |
| 231 | # descriptor behind a weak-referenceable Python callable so ``allow_in_graph`` |
| 232 | # can admit it and the downstream compiler can lower the original operation. |
| 233 | @torch.compiler.allow_in_graph |
| 234 | def _torch_to_in_graph(self, *args, **kwargs): |
| 235 | return _torch_to(self, *args, **kwargs) |
| 236 | |
| 237 | |
| 238 | # torch's own parser for `Tensor.to`'s overloaded signature, returning |
| 239 | # (device, dtype, non_blocking, memory_format) (also used by torch._dynamo). |
| 240 | _parse_to = getattr(torch._C._nn, "_parse_to", None) |
| 241 | |
| 242 | |
| 243 | def _active_check_mode() -> str | None: |
| 244 | """The mode the calling thread is being checked in, if any.""" |
| 245 | if not _sync_check_enabled or torch.compiler.is_compiling() or _allow_depth.get(): |
| 246 | return None |
| 247 | return _checking.get() |
| 248 | |
| 249 | |
| 250 | def _is_dense(t: torch.Tensor) -> bool: |
| 251 | """Whether `t`'s elements fill a contiguous storage range, possibly |
| 252 | permuted (e.g. a transposed matrix). Such layouts can be copied with |
| 253 | pitched cudaMemcpy2D/3DAsync directly from/to pinned memory, whereas |
| 254 | gapped layouts (e.g. strided slices) stage through a pageable temp.""" |
| 255 | if t.is_contiguous(): |
| 256 | # The common case, answered by a memoized TensorImpl flag. |
| 257 | return True |
| 258 | expected_stride = 1 |
| 259 | for stride, size in sorted(zip(t.stride(), t.shape)): |
| 260 | if size <= 1: |
| 261 | continue |
| 262 | if stride != expected_stride: |
| 263 | return False |
| 264 | expected_stride *= size |
| 265 | return True |
| 266 | |
| 267 | |
| 268 | def _cpu_copy_stall_reason( |
| 269 | cpu_tensor: torch.Tensor, device: torch.device | None |
| 270 | ) -> str | None: |
| 271 | """Why a `non_blocking` copy between `cpu_tensor` and a CUDA tensor would |
| 272 | silently block the host, or None if it is genuinely asynchronous. |
| 273 | |
| 274 | `device` describes the other side of the copy; a non-CUDA `device` means |
| 275 | the copy does not involve the GPU at all. |
| 276 | """ |
| 277 | if device is None or device.type != "cuda" or cpu_tensor.device.type != "cpu": |
| 278 | return None |
| 279 | if cpu_tensor.numel() == 0: |
| 280 | return None # Nothing to copy; torch issues no CUDA call at all. |
| 281 | if not _is_dense(cpu_tensor): |
| 282 | return "the CPU tensor is not densely laid out" |
| 283 | if not cpu_tensor.is_pinned(): |
| 284 | return "the CPU tensor is not pinned" |
| 285 | return None |
| 286 | |
| 287 | |
| 288 | def _report_implicit_sync(op: str, direction: str, reason: str, mode: str) -> None: |
| 289 | msg = ( |
| 290 | f"Tensor.{op}: a non_blocking {direction} copy that may be silently " |
| 291 | f"synchronous ({reason}). {IMPLICIT_SYNC_ERROR_MESSAGE}" |
| 292 | ) |
| 293 | if mode == "error": |
| 294 | raise RuntimeError(msg) |
| 295 | warnings.warn(msg, stacklevel=3) |
| 296 | |
| 297 | |
| 298 | def _checked_to(self, *args, **kwargs): |
| 299 | # Only H2D is checked: D2H `to` allocates a pinned destination in |
| 300 | # supported torch versions, so it is genuinely asynchronous. |
| 301 | if (mode := _active_check_mode()) is not None and _parse_to is not None: |
| 302 | try: |
| 303 | device, _, non_blocking, _ = _parse_to(*args, **kwargs) |
| 304 | except Exception: |
| 305 | pass # Invalid args; let the original call raise. |
| 306 | else: |
| 307 | if non_blocking and (reason := _cpu_copy_stall_reason(self, device)): |
| 308 | _report_implicit_sync("to", "H2D", reason, mode) |
| 309 | return _torch_to_in_graph(self, *args, **kwargs) |
| 310 | |
| 311 | |
| 312 | def _checked_cuda(self, *args, **kwargs): |
| 313 | # Signature: cuda(device=None, non_blocking=False, memory_format=None). |
| 314 | if (mode := _active_check_mode()) is not None: |
| 315 | device = args[0] if args else kwargs.get("device") |
| 316 | non_blocking = args[1] if len(args) > 1 else kwargs.get("non_blocking", False) |
| 317 | if non_blocking: |
| 318 | if device is None: |
| 319 | device = torch.device("cuda") |
| 320 | elif isinstance(device, int): |
| 321 | device = torch.device("cuda", device) |
| 322 | elif not isinstance(device, torch.device): |
| 323 | device = torch.device(device) |
| 324 | if reason := _cpu_copy_stall_reason(self, device): |
| 325 | _report_implicit_sync("cuda", "H2D", reason, mode) |
| 326 | return _torch_cuda(self, *args, **kwargs) |
| 327 | |
| 328 | |
| 329 | def _checked_copy_(self, *args, **kwargs): |
| 330 | # Signature: copy_(src, non_blocking=False). |
| 331 | if ( |
| 332 | (mode := _active_check_mode()) is not None |
| 333 | and args |
| 334 | and isinstance(src := args[0], torch.Tensor) |
| 335 | ): |
| 336 | non_blocking = args[1] if len(args) > 1 else kwargs.get("non_blocking", False) |
| 337 | if non_blocking: |
| 338 | if self.device.type == "cuda": |
| 339 | reason = _cpu_copy_stall_reason(src, self.device) |
| 340 | direction = "H2D" |
| 341 | else: |
| 342 | reason = _cpu_copy_stall_reason(self, src.device) |
| 343 | direction = "D2H" |
| 344 | if reason: |
| 345 | _report_implicit_sync("copy_", direction, reason, mode) |
| 346 | return _torch_copy_(self, *args, **kwargs) |
| 347 | |
| 348 | |
| 349 | _copy_checkers_installed: bool = False |
| 350 | |
| 351 | |
| 352 | def _install_copy_checkers() -> None: |
| 353 | """Wrap the Python-level CPU<->CUDA transfer entry points with the checks |
| 354 | above. |
| 355 | |
| 356 | The patch is process-global, but it only reports on threads inside a |
| 357 | checked region, matching the scoping of the torch-level check. Copies |
| 358 | issued from C++ (e.g. inside compiled graphs) bypass it. |
| 359 | """ |
| 360 | global _copy_checkers_installed |
| 361 | if _copy_checkers_installed: |
| 362 | return |
| 363 | _copy_checkers_installed = True |
| 364 | torch.Tensor.to = _checked_to # type: ignore[method-assign] |
| 365 | torch.Tensor.cuda = _checked_cuda # type: ignore[method-assign] |
| 366 | torch.Tensor.copy_ = _checked_copy_ # type: ignore[method-assign] |
| 367 | |
| 368 | |
| 369 | def gpu_sync_allowed(first_only: bool = False): |
| 370 | """Allow GPU<->CPU syncs inside the `with` block, on this thread only. |
| 371 | |
| 372 | With `first_only`, only the first entry from a given call site |
| 373 | (filename, lineno) is allowed, so later syncs there are still reported. |
| 374 | """ |
| 375 | if _SYNC_CHECK_MODE is None or torch.compiler.is_compiling(): |
| 376 | return _NOOP_CM |
| 377 | if first_only: |
| 378 | frame = sys._getframe(1) |
| 379 | key = (frame.f_code.co_filename, frame.f_lineno) |
| 380 | if key in _GPU_SYNC_ALLOWED_FIRST_SEEN: |
| 381 | return _NOOP_CM |
| 382 | _GPU_SYNC_ALLOWED_FIRST_SEEN.add(key) |
| 383 | return _allow_syncs() |
| 384 | |
| 385 | |
| 386 | def with_gpu_sync_check(fn): |
| 387 | """Report GPU<->CPU syncs performed by `fn` on its calling thread. |
| 388 | |
| 389 | Active only once `enable_gpu_sync_check()` has flipped the gate. Other |
| 390 | threads are never policed, so deliberate syncs there (e.g. the EPLB |
| 391 | transfer worker) are unaffected. |
| 392 | """ |
| 393 | if (mode := _SYNC_CHECK_MODE) is None: |
| 394 | return fn |
| 395 | |
| 396 | @functools.wraps(fn) |
| 397 | def wrapper(*args, **kwargs): |
| 398 | if not _sync_check_enabled: |
| 399 | return fn(*args, **kwargs) |
| 400 | with _checked_region(mode): |
| 401 | return fn(*args, **kwargs) |
| 402 | |
| 403 | return wrapper |