[Core] Add VLLM_GPU_SYNC_CHECK env var (#44800)

tests/utils_/test_gpu_sync_debug.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+
+import vllm.utils.gpu_sync_debug as gsd
+from vllm.utils.gpu_sync_debug import (
+    SYNC_ERROR_MESSAGE,
+    gpu_sync_allowed,
+    with_gpu_sync_check,
+)
+
+from ..utils import create_new_process_for_each_test
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+
+
+def _no_sync():
+    # Pure on-GPU compute, no implicit CPU sync...
+    x = torch.ones(4, device="cuda") + 1
+    # ...plus a sync that we explicitly allow.
+    with gpu_sync_allowed():
+        return x.cpu()
+
+
+def _causes_sync():
+    x = torch.ones(4, device="cuda")
+    # An allowed sync (suppressed)...
+    with gpu_sync_allowed():
+        x.cpu()
+    # ...then an un-allowed sync that should trip the check.
+    return x.cpu()
+
+
+@pytest.mark.parametrize("mode", ["warn", "error"])
+@create_new_process_for_each_test()
+def test_with_env_set(monkeypatch, mode):
+    # Env set + gate flipped on: the unguarded sync is detected.
+    monkeypatch.setenv("VLLM_GPU_SYNC_CHECK", mode)
+    monkeypatch.setattr(gsd, "_sync_check_enabled", True)
+
+    # Guarded syncs always pass.
+    with_gpu_sync_check(_no_sync)()
+
+    if mode == "error":
+        # "error" mode turns the stray sync into a RuntimeError.
+        with pytest.raises(RuntimeError, match=SYNC_ERROR_MESSAGE):
+            with_gpu_sync_check(_causes_sync)()
+    else:
+        # "warn" mode only warns, so the call still succeeds.
+        with_gpu_sync_check(_causes_sync)()
+
+
+@create_new_process_for_each_test()
+def test_without_env_set(monkeypatch):
+    # Env unset: the decorator is a pass-through, no sync is detected.
+    monkeypatch.delenv("VLLM_GPU_SYNC_CHECK", raising=False)
+    monkeypatch.setattr(gsd, "_sync_check_enabled", True)
+
+    with_gpu_sync_check(_no_sync)()
+    with_gpu_sync_check(_causes_sync)()

vllm/compilation/compiler_interface.py

@@ -765,6 +765,34 @@ def set_functorch_config() -> None:
         setattr(torch._functorch.config, k, v)
 
 
+def trigger_inductor_lazy_init(device: torch.device | None = None) -> None:
+    """Eagerly trigger inductor's once-per-process lazy inits (SFDP pattern
+    matcher, pad_mm, misc patterns).
+
+    These normally fire on the first torch.compile invocation and include
+    CUDA syncs. If warmup hits the on-disk compile cache, no compile actually
+    runs so these never fire during warmup, and they'd blow up on the first
+    real-request cache miss once the sync-check gate is on.
+
+    Private torch API; best-effort. Newer torch versions take an
+    `input_device` argument and cache per-device, so pass the current CUDA
+    device to ensure the cache key matches later compile calls.
+    """
+    try:
+        import inspect
+
+        from torch._inductor.fx_passes.joint_graph import (
+            lazy_init as _inductor_lazy_init,
+        )
+
+        if inspect.signature(_inductor_lazy_init).parameters:
+            _inductor_lazy_init(device)
+        else:
+            _inductor_lazy_init()
+    except Exception as e:  # noqa: BLE001
+        logger.info("Skipping inductor lazy_init pre-trigger: %s", e)
+
+
 class EagerAdaptor(CompilerInterface):
     name = "eager"
 

vllm/envs.py

@@ -87,6 +87,7 @@ if TYPE_CHECKING:
     VLLM_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest"
     VLLM_BATCH_INVARIANT: bool = False
     VLLM_TRITON_ATTN_USE_TD: bool | None = None
+    VLLM_GPU_SYNC_CHECK: Literal["warn", "error"] | None = None
     MAX_JOBS: str | None = None
     NVCC_THREADS: str | None = None
     VLLM_USE_PRECOMPILED: bool = False
@@ -584,6 +585,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_TRITON_ATTN_USE_TD": lambda: {"1": True, "0": False}.get(
         os.getenv("VLLM_TRITON_ATTN_USE_TD", "").strip()
     ),
+    # If set, enable PyTorch's GPU<->CPU synchronization debug mode around
+    # the worker's `execute_model` and `sample_tokens` calls. Valid values
+    # are "warn" (print a warning on each sync) or "error" (raise on sync).
+    # Unset disables the check. See `torch.cuda.set_sync_debug_mode`.
+    "VLLM_GPU_SYNC_CHECK": env_with_choices(
+        "VLLM_GPU_SYNC_CHECK", None, ["warn", "error"]
+    ),
     # Maximum number of compilation jobs to run in parallel.
     # By default this is the number of CPUs
     "MAX_JOBS": lambda: os.getenv("MAX_JOBS", None),

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

vllm/v1/worker/gpu_worker.py

@@ -56,6 +56,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
 from vllm.tracing import instrument
 from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback
+from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling
 from vllm.utils.torch_utils import set_random_seed
@@ -759,6 +760,16 @@ class Worker(WorkerBase):
         # the model initialization and profiling.
         set_random_seed(self.model_config.seed)
 
+        # Eagerly trigger inductor's once-per-process lazy inits during
+        # warmup (rather than on a later compile cache-miss at runtime).
+        c_config = self.compilation_config
+        if c_config.mode != CompilationMode.NONE and c_config.backend == "inductor":
+            from vllm.compilation.compiler_interface import (
+                trigger_inductor_lazy_init,
+            )
+
+            trigger_inductor_lazy_init(self.device)
+
         # All warmup is done — start monitoring for unexpected JIT
         # compilations that would cause latency spikes during inference.
         from vllm.utils.jit_monitor import activate as activate_jit_monitor
@@ -773,6 +784,10 @@ class Worker(WorkerBase):
         freeze_gc_heap()
         maybe_attach_gc_debug_callback()
 
+        # Warmup / first-compile is done — activate the `VLLM_GPU_SYNC_CHECK`
+        # gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
+        enable_gpu_sync_check()
+
         return CompilationTimes(
             language_model=self.compilation_config.compilation_time,
             encoder=self.compilation_config.encoder_compilation_time,
@@ -826,12 +841,14 @@ class Worker(WorkerBase):
         return self.profiler.annotate_context_manager(annotation)
 
     @torch.inference_mode()
+    @with_gpu_sync_check
     def sample_tokens(
         self, grammar_output: "GrammarOutput | None"
     ) -> ModelRunnerOutput | AsyncModelRunnerOutput:
         return self.model_runner.sample_tokens(grammar_output)
 
     @torch.inference_mode()
+    @with_gpu_sync_check
     def execute_model(
         self, scheduler_output: "SchedulerOutput"
     ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None: