[Core] Add `VLLM_GPU_SYNC_CHECK` env var (#44800)
Nick Hill committed
Jun 26, 2026 at 08:24 UTC
4e07ca2c9284ad0661a44d33e1e8a1c597c48686
5 files changed
+279
tests/utils_/test_gpu_sync_debug.py
new
+61
@@ -0,0 +1,61 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+import vllm.utils.gpu_sync_debug as gsd
7
+from vllm.utils.gpu_sync_debug import (
8
+ SYNC_ERROR_MESSAGE,
9
+ gpu_sync_allowed,
10
+ with_gpu_sync_check,
11
+)
12
+
13
+from ..utils import create_new_process_for_each_test
14
+
15
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
16
+
17
+
18
+def _no_sync():
19
+ # Pure on-GPU compute, no implicit CPU sync...
20
+ x = torch.ones(4, device="cuda") + 1
21
+ # ...plus a sync that we explicitly allow.
22
+ with gpu_sync_allowed():
23
+ return x.cpu()
24
+
25
+
26
+def _causes_sync():
27
+ x = torch.ones(4, device="cuda")
28
+ # An allowed sync (suppressed)...
29
+ with gpu_sync_allowed():
30
+ x.cpu()
31
+ # ...then an un-allowed sync that should trip the check.
32
+ return x.cpu()
33
+
34
+
35
+@pytest.mark.parametrize("mode", ["warn", "error"])
36
+@create_new_process_for_each_test()
37
+def test_with_env_set(monkeypatch, mode):
38
+ # Env set + gate flipped on: the unguarded sync is detected.
39
+ monkeypatch.setenv("VLLM_GPU_SYNC_CHECK", mode)
40
+ monkeypatch.setattr(gsd, "_sync_check_enabled", True)
41
+
42
+ # Guarded syncs always pass.
43
+ with_gpu_sync_check(_no_sync)()
44
+
45
+ if mode == "error":
46
+ # "error" mode turns the stray sync into a RuntimeError.
47
+ with pytest.raises(RuntimeError, match=SYNC_ERROR_MESSAGE):
48
+ with_gpu_sync_check(_causes_sync)()
49
+ else:
50
+ # "warn" mode only warns, so the call still succeeds.
51
+ with_gpu_sync_check(_causes_sync)()
52
+
53
+
54
+@create_new_process_for_each_test()
55
+def test_without_env_set(monkeypatch):
56
+ # Env unset: the decorator is a pass-through, no sync is detected.
57
+ monkeypatch.delenv("VLLM_GPU_SYNC_CHECK", raising=False)
58
+ monkeypatch.setattr(gsd, "_sync_check_enabled", True)
59
+
60
+ with_gpu_sync_check(_no_sync)()
61
+ with_gpu_sync_check(_causes_sync)()
vllm/compilation/compiler_interface.py
+28
@@ -765,6 +765,34 @@ def set_functorch_config() -> None:
765
setattr(torch._functorch.config, k, v)
766
767
768
+def trigger_inductor_lazy_init(device: torch.device | None = None) -> None:
769
+ """Eagerly trigger inductor's once-per-process lazy inits (SFDP pattern
770
+ matcher, pad_mm, misc patterns).
771
+
772
+ These normally fire on the first torch.compile invocation and include
773
+ CUDA syncs. If warmup hits the on-disk compile cache, no compile actually
774
+ runs so these never fire during warmup, and they'd blow up on the first
775
+ real-request cache miss once the sync-check gate is on.
776
+
777
+ Private torch API; best-effort. Newer torch versions take an
778
+ `input_device` argument and cache per-device, so pass the current CUDA
779
+ device to ensure the cache key matches later compile calls.
780
+ """
781
+ try:
782
+ import inspect
783
+
784
+ from torch._inductor.fx_passes.joint_graph import (
785
+ lazy_init as _inductor_lazy_init,
786
+ )
787
+
788
+ if inspect.signature(_inductor_lazy_init).parameters:
789
+ _inductor_lazy_init(device)
790
+ else:
791
+ _inductor_lazy_init()
792
+ except Exception as e: # noqa: BLE001
793
+ logger.info("Skipping inductor lazy_init pre-trigger: %s", e)
794
+
795
+
796
class EagerAdaptor(CompilerInterface):
797
name = "eager"
798
vllm/envs.py
+8
@@ -87,6 +87,7 @@ if TYPE_CHECKING:
87
VLLM_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest"
88
VLLM_BATCH_INVARIANT: bool = False
89
VLLM_TRITON_ATTN_USE_TD: bool | None = None
90
+ VLLM_GPU_SYNC_CHECK: Literal["warn", "error"] | None = None
91
MAX_JOBS: str | None = None
92
NVCC_THREADS: str | None = None
93
VLLM_USE_PRECOMPILED: bool = False
@@ -584,6 +585,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
585
"VLLM_TRITON_ATTN_USE_TD": lambda: {"1": True, "0": False}.get(
586
os.getenv("VLLM_TRITON_ATTN_USE_TD", "").strip()
587
),
588
+ # If set, enable PyTorch's GPU<->CPU synchronization debug mode around
589
+ # the worker's `execute_model` and `sample_tokens` calls. Valid values
590
+ # are "warn" (print a warning on each sync) or "error" (raise on sync).
591
+ # Unset disables the check. See `torch.cuda.set_sync_debug_mode`.
592
+ "VLLM_GPU_SYNC_CHECK": env_with_choices(
593
+ "VLLM_GPU_SYNC_CHECK", None, ["warn", "error"]
594
+ ),
595
# Maximum number of compilation jobs to run in parallel.
596
# By default this is the number of CPUs
597
"MAX_JOBS": lambda: os.getenv("MAX_JOBS", None),
vllm/utils/gpu_sync_debug.py
new
+165
@@ -0,0 +1,165 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import functools
4
+import sys
5
+from contextlib import contextmanager
6
+
7
+import torch
8
+
9
+import vllm.envs as envs
10
+from vllm.platforms import current_platform
11
+
12
+SYNC_ERROR_MESSAGE = (
13
+ "GPU<->CPU sync detected - avoid it or wrap with gpu_sync_allowed()"
14
+)
15
+
16
+_GPU_SYNC_ALLOWED_FIRST_SEEN: set[tuple[str, int]] = set()
17
+
18
+# Global sync-check gate. Off during engine setup (model load, KV cache
19
+# init, warmup/compile) so first-compile and lazy-init syncs pass through;
20
+# flipped on by `enable_gpu_sync_check()` at the end of
21
+# `GPUWorker.compile_or_warm_up_model`, after which `with_gpu_sync_check`-
22
+# decorated functions activate the configured debug mode.
23
+_sync_check_enabled: bool = False
24
+
25
+
26
+def enable_gpu_sync_check() -> None:
27
+ """Flip the sync-check gate on. Call once per worker, after warmup /
28
+ first-compile is complete. No-op unless `VLLM_GPU_SYNC_CHECK` is set."""
29
+ if envs.VLLM_GPU_SYNC_CHECK is None:
30
+ return
31
+ global _sync_check_enabled
32
+ _sync_check_enabled = True
33
+ _install_compile_time_sync_suppressors()
34
+
35
+
36
+_compile_time_suppressors_installed: bool = False
37
+
38
+
39
+def _install_compile_time_sync_suppressors() -> None:
40
+ """Wrap torch inductor/aot_autograd compile entry points so the
41
+ synchronizing ops those passes perform don't trip the
42
+ sync-check mode we set around `execute_model` / `sample_tokens`.
43
+
44
+ Warmup-time compiles already run under the gate (before
45
+ `enable_gpu_sync_check`), but post-warmup compiles fire inside
46
+ `execute_model` and we want to avoid this tripping the sync check.
47
+ """
48
+ global _compile_time_suppressors_installed
49
+ if _compile_time_suppressors_installed:
50
+ return
51
+ _compile_time_suppressors_installed = True
52
+
53
+ try: # noqa: BLE001
54
+ from torch._inductor.fx_passes import joint_graph as _jg
55
+
56
+ _orig_joint = _jg.joint_graph_passes
57
+
58
+ @functools.wraps(_orig_joint)
59
+ def _wrapped_joint(*args, **kwargs):
60
+ prev_mode = torch.cuda.get_sync_debug_mode()
61
+ if not prev_mode:
62
+ return _orig_joint(*args, **kwargs)
63
+ torch.cuda.set_sync_debug_mode(0)
64
+ try:
65
+ return _orig_joint(*args, **kwargs)
66
+ finally:
67
+ torch.cuda.set_sync_debug_mode(prev_mode)
68
+
69
+ # `compile_fx` does `from .fx_passes.joint_graph import
70
+ # joint_graph_passes`, which binds the *function object* at import
71
+ # time. Patching just the module attribute won't update that rebind,
72
+ # so patch every already-imported reference we can find. Restrict
73
+ # the scan to torch's compile-time modules.
74
+ import sys as _sys
75
+
76
+ setattr(_jg, "joint_graph_passes", _wrapped_joint) # noqa: B010
77
+ for _name, _mod in list(_sys.modules.items()):
78
+ if _mod is None:
79
+ continue
80
+ if not (
81
+ _name.startswith("torch._inductor")
82
+ or _name.startswith("torch._functorch")
83
+ or _name.startswith("torch._dynamo")
84
+ ):
85
+ continue
86
+ if getattr(_mod, "joint_graph_passes", None) is _orig_joint:
87
+ setattr(_mod, "joint_graph_passes", _wrapped_joint) # noqa: B010
88
+ except Exception: # pragma: no cover
89
+ pass
90
+
91
+
92
+@contextmanager
93
+def _suppress_gpu_sync_check(prev_mode: int):
94
+ torch.cuda.set_sync_debug_mode(0)
95
+ try:
96
+ yield
97
+ finally:
98
+ torch.cuda.set_sync_debug_mode(prev_mode)
99
+
100
+
101
+@contextmanager
102
+def _noop_cm():
103
+ yield
104
+
105
+
106
+if current_platform.is_cuda_alike():
107
+
108
+ def gpu_sync_allowed(first_only: bool = False):
109
+ """Context manager that suppresses `torch.cuda.set_sync_debug_mode` for the
110
+ duration of the `with` block.
111
+
112
+ If `first_only` is True, only the first entry from this call site
113
+ suppresses the sync check; subsequent entries from the same site are
114
+ no-ops so any further GPU syncs will be reported. The "site" is the
115
+ caller's (filename, lineno), so different
116
+ `with gpu_sync_allowed(first_only=True):` lines track independently.
117
+ """
118
+ if envs.VLLM_GPU_SYNC_CHECK is None or torch.compiler.is_compiling():
119
+ return _noop_cm()
120
+ prev_mode = torch.cuda.get_sync_debug_mode()
121
+ if not prev_mode:
122
+ return _noop_cm()
123
+ if first_only:
124
+ frame = sys._getframe(1)
125
+ key = (frame.f_code.co_filename, frame.f_lineno)
126
+ if key in _GPU_SYNC_ALLOWED_FIRST_SEEN:
127
+ return _noop_cm()
128
+ _GPU_SYNC_ALLOWED_FIRST_SEEN.add(key)
129
+ return _suppress_gpu_sync_check(prev_mode)
130
+
131
+ def with_gpu_sync_check(fn):
132
+ """Decorator that enables `torch.cuda.set_sync_debug_mode` around `fn`
133
+ when `VLLM_GPU_SYNC_CHECK` is set *and* the gate has been flipped by
134
+ `enable_gpu_sync_check()`. Before the gate flips (i.e. during
135
+ engine setup / warmup) the decorated function runs as-is.
136
+ """
137
+ mode = envs.VLLM_GPU_SYNC_CHECK
138
+ if mode is None:
139
+ return fn
140
+
141
+ @functools.wraps(fn)
142
+ def wrapper(*args, **kwargs):
143
+ if not _sync_check_enabled:
144
+ return fn(*args, **kwargs)
145
+ prev_mode = torch.cuda.get_sync_debug_mode()
146
+ torch.cuda.set_sync_debug_mode(mode)
147
+ try:
148
+ return fn(*args, **kwargs)
149
+ except RuntimeError as re:
150
+ if str(re) == "called a synchronizing CUDA operation":
151
+ raise RuntimeError(SYNC_ERROR_MESSAGE) from re
152
+ raise re
153
+ finally:
154
+ torch.cuda.set_sync_debug_mode(prev_mode)
155
+
156
+ return wrapper
157
+
158
+else:
159
+ # No-op the methods in non-CUDA cases.
160
+
161
+ def gpu_sync_allowed(first_only: bool = False):
162
+ return _noop_cm()
163
+
164
+ def with_gpu_sync_check(fn):
165
+ return fn
vllm/v1/worker/gpu_worker.py
+17
@@ -56,6 +56,7 @@ from vllm.sequence import IntermediateTensors
56
from vllm.tasks import SupportedTask
57
from vllm.tracing import instrument
58
from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback
59
+from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check
60
from vllm.utils.mem_constants import GiB_bytes
61
from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling
62
from vllm.utils.torch_utils import set_random_seed
@@ -759,6 +760,16 @@ class Worker(WorkerBase):
760
# the model initialization and profiling.
761
set_random_seed(self.model_config.seed)
762
763
+ # Eagerly trigger inductor's once-per-process lazy inits during
764
+ # warmup (rather than on a later compile cache-miss at runtime).
765
+ c_config = self.compilation_config
766
+ if c_config.mode != CompilationMode.NONE and c_config.backend == "inductor":
767
+ from vllm.compilation.compiler_interface import (
768
+ trigger_inductor_lazy_init,
769
+ )
770
+
771
+ trigger_inductor_lazy_init(self.device)
772
+
773
# All warmup is done — start monitoring for unexpected JIT
774
# compilations that would cause latency spikes during inference.
775
from vllm.utils.jit_monitor import activate as activate_jit_monitor
@@ -773,6 +784,10 @@ class Worker(WorkerBase):
784
freeze_gc_heap()
785
maybe_attach_gc_debug_callback()
786
787
+ # Warmup / first-compile is done — activate the `VLLM_GPU_SYNC_CHECK`
788
+ # gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
789
+ enable_gpu_sync_check()
790
+
791
return CompilationTimes(
792
language_model=self.compilation_config.compilation_time,
793
encoder=self.compilation_config.encoder_compilation_time,
@@ -826,12 +841,14 @@ class Worker(WorkerBase):
841
return self.profiler.annotate_context_manager(annotation)
842
843
@torch.inference_mode()
844
+ @with_gpu_sync_check
845
def sample_tokens(
846
self, grammar_output: "GrammarOutput | None"
847
) -> ModelRunnerOutput | AsyncModelRunnerOutput:
848
return self.model_runner.sample_tokens(grammar_output)
849
850
@torch.inference_mode()
851
+ @with_gpu_sync_check
852
def execute_model(
853
self, scheduler_output: "SchedulerOutput"
854
) -> ModelRunnerOutput | AsyncModelRunnerOutput | None: