chore: import upstream snapshot with attribution

vllm/utils/gpu_sync_debug.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import functools
+import sys
+from contextlib import contextmanager
+
+import torch
+
+import vllm.envs as envs
+from vllm.platforms import current_platform
+
+SYNC_ERROR_MESSAGE = (
+    "GPU<->CPU sync detected - avoid it or wrap with gpu_sync_allowed()"
+)
+
+_GPU_SYNC_ALLOWED_FIRST_SEEN: set[tuple[str, int]] = set()
+
+# Global sync-check gate. Off during engine setup (model load, KV cache
+# init, warmup/compile) so first-compile and lazy-init syncs pass through;
+# flipped on by `enable_gpu_sync_check()` at the end of
+# `GPUWorker.compile_or_warm_up_model`, after which `with_gpu_sync_check`-
+# decorated functions activate the configured debug mode.
+_sync_check_enabled: bool = False
+
+
+def enable_gpu_sync_check() -> None:
+    """Flip the sync-check gate on. Call once per worker, after warmup /
+    first-compile is complete. No-op unless `VLLM_GPU_SYNC_CHECK` is set."""
+    if envs.VLLM_GPU_SYNC_CHECK is None:
+        return
+    global _sync_check_enabled
+    _sync_check_enabled = True
+    _install_compile_time_sync_suppressors()
+
+
+_compile_time_suppressors_installed: bool = False
+
+
+def _install_compile_time_sync_suppressors() -> None:
+    """Wrap torch inductor/aot_autograd compile entry points so the
+    synchronizing ops those passes perform don't trip the
+    sync-check mode we set around `execute_model` / `sample_tokens`.
+
+    Warmup-time compiles already run under the gate (before
+    `enable_gpu_sync_check`), but post-warmup compiles fire inside
+    `execute_model` and we want to avoid this tripping the sync check.
+    """
+    global _compile_time_suppressors_installed
+    if _compile_time_suppressors_installed:
+        return
+    _compile_time_suppressors_installed = True
+
+    try:  # noqa: BLE001
+        from torch._inductor.fx_passes import joint_graph as _jg
+
+        _orig_joint = _jg.joint_graph_passes
+
+        @functools.wraps(_orig_joint)
+        def _wrapped_joint(*args, **kwargs):
+            prev_mode = torch.cuda.get_sync_debug_mode()
+            if not prev_mode:
+                return _orig_joint(*args, **kwargs)
+            torch.cuda.set_sync_debug_mode(0)
+            try:
+                return _orig_joint(*args, **kwargs)
+            finally:
+                torch.cuda.set_sync_debug_mode(prev_mode)
+
+        # `compile_fx` does `from .fx_passes.joint_graph import
+        # joint_graph_passes`, which binds the *function object* at import
+        # time. Patching just the module attribute won't update that rebind,
+        # so patch every already-imported reference we can find. Restrict
+        # the scan to torch's compile-time modules.
+        import sys as _sys
+
+        setattr(_jg, "joint_graph_passes", _wrapped_joint)  # noqa: B010
+        for _name, _mod in list(_sys.modules.items()):
+            if _mod is None:
+                continue
+            if not (
+                _name.startswith("torch._inductor")
+                or _name.startswith("torch._functorch")
+                or _name.startswith("torch._dynamo")
+            ):
+                continue
+            if getattr(_mod, "joint_graph_passes", None) is _orig_joint:
+                setattr(_mod, "joint_graph_passes", _wrapped_joint)  # noqa: B010
+    except Exception:  # pragma: no cover
+        pass
+
+
+@contextmanager
+def _suppress_gpu_sync_check(prev_mode: int):
+    torch.cuda.set_sync_debug_mode(0)
+    try:
+        yield
+    finally:
+        torch.cuda.set_sync_debug_mode(prev_mode)
+
+
+@contextmanager
+def _noop_cm():
+    yield
+
+
+if current_platform.is_cuda_alike():
+
+    def gpu_sync_allowed(first_only: bool = False):
+        """Context manager that suppresses `torch.cuda.set_sync_debug_mode` for the
+        duration of the `with` block.
+
+        If `first_only` is True, only the first entry from this call site
+        suppresses the sync check; subsequent entries from the same site are
+        no-ops so any further GPU syncs will be reported. The "site" is the
+        caller's (filename, lineno), so different
+        `with gpu_sync_allowed(first_only=True):` lines track independently.
+        """
+        if envs.VLLM_GPU_SYNC_CHECK is None or torch.compiler.is_compiling():
+            return _noop_cm()
+        prev_mode = torch.cuda.get_sync_debug_mode()
+        if not prev_mode:
+            return _noop_cm()
+        if first_only:
+            frame = sys._getframe(1)
+            key = (frame.f_code.co_filename, frame.f_lineno)
+            if key in _GPU_SYNC_ALLOWED_FIRST_SEEN:
+                return _noop_cm()
+            _GPU_SYNC_ALLOWED_FIRST_SEEN.add(key)
+        return _suppress_gpu_sync_check(prev_mode)
+
+    def with_gpu_sync_check(fn):
+        """Decorator that enables `torch.cuda.set_sync_debug_mode` around `fn`
+        when `VLLM_GPU_SYNC_CHECK` is set *and* the gate has been flipped by
+        `enable_gpu_sync_check()`. Before the gate flips (i.e. during
+        engine setup / warmup) the decorated function runs as-is.
+        """
+        mode = envs.VLLM_GPU_SYNC_CHECK
+        if mode is None:
+            return fn
+
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            if not _sync_check_enabled:
+                return fn(*args, **kwargs)
+            prev_mode = torch.cuda.get_sync_debug_mode()
+            torch.cuda.set_sync_debug_mode(mode)
+            try:
+                return fn(*args, **kwargs)
+            except RuntimeError as re:
+                if str(re) == "called a synchronizing CUDA operation":
+                    raise RuntimeError(SYNC_ERROR_MESSAGE) from re
+                raise re
+            finally:
+                torch.cuda.set_sync_debug_mode(prev_mode)
+
+        return wrapper
+
+else:
+    # No-op the methods in non-CUDA cases.
+
+    def gpu_sync_allowed(first_only: bool = False):
+        return _noop_cm()
+
+    def with_gpu_sync_check(fn):
+        return fn