@tduncan / Nvidia-Optimization-Stack-VLLM / commits / fbff187d59

[Core] Make the GPU sync check thread-local and fix its suppressors (#51455)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Aug 8, 2026 at 17:13 UTC fbff187d593c322c4a221864ce5756bd849e174e
2 files changed +288 -120
tests/utils_/test_gpu_sync_debug.py
+101 -2
@@ -1,5 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import threading
4 +import warnings
5 +
6 import pytest
7 import torch
8
@@ -36,7 +39,9 @@ def _causes_sync():
39 @create_new_process_for_each_test()
40 def test_with_env_set(monkeypatch, mode):
41 # Env set + gate flipped on: the unguarded sync is detected.
39 - monkeypatch.setenv("VLLM_GPU_SYNC_CHECK", mode)
42 + # `_SYNC_CHECK_MODE` is read from the env once at import, so patch the
43 + # module attribute rather than the environment.
44 + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", mode)
45 monkeypatch.setattr(gsd, "_sync_check_enabled", True)
46
47 # Guarded syncs always pass.
@@ -51,10 +56,104 @@ def test_with_env_set(monkeypatch, mode):
56 with_gpu_sync_check(_causes_sync)()
57
58
59 +@create_new_process_for_each_test()
60 +def test_other_threads_are_not_policed(monkeypatch):
61 + """A background thread that syncs deliberately must not be broken by the
62 + check being armed on the thread running the decorated function."""
63 + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
64 + monkeypatch.setattr(gsd, "_sync_check_enabled", True)
65 +
66 + def sync_on_worker():
67 + failure: list[BaseException] = []
68 +
69 + def worker():
70 + try:
71 + torch.ones(4, device="cuda").cpu()
72 + except BaseException as exc: # pragma: no cover - failure path
73 + failure.append(exc)
74 +
75 + thread = threading.Thread(target=worker)
76 + thread.start()
77 + thread.join()
78 + assert not failure, f"background thread raised: {failure[0]!r}"
79 +
80 + with_gpu_sync_check(sync_on_worker)()
81 +
82 +
83 +@create_new_process_for_each_test()
84 +def test_allow_on_other_thread_does_not_disarm(monkeypatch):
85 + """`gpu_sync_allowed()` on one thread must not suppress the check on
86 + another. It is scoped by ContextVar rather than torch's process-global
87 + sync debug mode, which a previous implementation mutated."""
88 + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
89 + monkeypatch.setattr(gsd, "_sync_check_enabled", True)
90 +
91 + def main_syncs_while_worker_allows():
92 + stop = threading.Event()
93 +
94 + def worker():
95 + with gpu_sync_allowed():
96 + while not stop.is_set():
97 + torch.ones(4, device="cuda").cpu()
98 +
99 + thread = threading.Thread(target=worker)
100 + thread.start()
101 + try:
102 + # Must still be reported despite the worker's open allow region.
103 + torch.ones(4, device="cuda").cpu()
104 + finally:
105 + stop.set()
106 + thread.join()
107 +
108 + with pytest.raises(RuntimeError, match=SYNC_ERROR_MESSAGE):
109 + with_gpu_sync_check(main_syncs_while_worker_allows)()
110 +
111 +
112 +@create_new_process_for_each_test()
113 +def test_suppressing_works_while_compiling(monkeypatch):
114 + """`_suppressing` wraps torch compile entry points, which run with
115 + `torch.compiler.is_compiling()` true. `gpu_sync_allowed()` deliberately
116 + no-ops in that state, so `_suppressing` must not route through it."""
117 + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
118 + monkeypatch.setattr(gsd, "_sync_check_enabled", True)
119 + # Emulate being inside a torch compile, as inductor passes are.
120 + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
121 +
122 + suppressed = gsd._suppressing(lambda: torch.ones(4, device="cuda").cpu())
123 + with_gpu_sync_check(suppressed)()
124 +
125 +
126 +@create_new_process_for_each_test()
127 +def test_sync_debug_mode_restored_after_checked_call(monkeypatch):
128 + """The mode is armed only for the duration of a checked call. Leaving it
129 + on process-wide made every sync outside a checked region emit a
130 + `UserWarning` whenever our handler was not the installed one."""
131 + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
132 + monkeypatch.setattr(gsd, "_sync_check_enabled", True)
133 +
134 + before = torch.cuda.get_sync_debug_mode()
135 +
136 + def nested():
137 + # `execute_model` -> `sample_tokens` both carry the decorator.
138 + assert torch.cuda.get_sync_debug_mode() != 0, "armed inside"
139 + with_gpu_sync_check(lambda: None)()
140 + assert torch.cuda.get_sync_debug_mode() != 0, "still armed after inner"
141 +
142 + with_gpu_sync_check(nested)()
143 + assert torch.cuda.get_sync_debug_mode() == before
144 +
145 + # With the mode back to its original value, torch emits nothing for a
146 + # sync outside a checked region.
147 + with warnings.catch_warnings(record=True) as caught:
148 + warnings.simplefilter("always")
149 + torch.ones(4, device="cuda").cpu()
150 + assert not [r for r in caught if gsd._TORCH_SYNC_WARNING in str(r.message)]
151 +
152 +
153 @create_new_process_for_each_test()
154 def test_without_env_set(monkeypatch):
155 # Env unset: the decorator is a pass-through, no sync is detected.
57 - monkeypatch.delenv("VLLM_GPU_SYNC_CHECK", raising=False)
156 + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", None)
157 monkeypatch.setattr(gsd, "_sync_check_enabled", True)
158
159 with_gpu_sync_check(_no_sync)()
vllm/utils/gpu_sync_debug.py
+187 -118
@@ -1,8 +1,20 @@
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 +
12 import functools
13 import sys
5 -from contextlib import contextmanager
14 +import threading
15 +import warnings
16 +from contextlib import contextmanager, nullcontext
17 +from contextvars import ContextVar
18
19 import torch
20
@@ -13,153 +25,210 @@ SYNC_ERROR_MESSAGE = (
25 "GPU<->CPU sync detected - avoid it or wrap with gpu_sync_allowed()"
26 )
27
28 +# Warning text torch emits in "warn" mode.
29 +_TORCH_SYNC_WARNING = "called a synchronizing CUDA operation"
30 +
31 _GPU_SYNC_ALLOWED_FIRST_SEEN: set[tuple[str, int]] = set()
32
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.
33 +# Read once: `envs.__getattr__` re-evaluates the env lambda on every access,
34 +# which dominates the cost of a disabled `gpu_sync_allowed()`. Tests toggling
35 +# the env var must patch this attribute instead.
36 +_SYNC_CHECK_MODE: str | None = (
37 + envs.VLLM_GPU_SYNC_CHECK if current_platform.is_cuda_alike() else None
38 +)
39 +
40 +# Configured mode while a thread is inside a checked call, else None.
41 +_checking: ContextVar[str | None] = ContextVar("vllm_gpu_sync_checking", default=None)
42 +# Depth of nested `gpu_sync_allowed()` regions on this thread.
43 +_allow_depth: ContextVar[int] = ContextVar("vllm_gpu_sync_allow_depth", default=0)
44 +
45 +# Off during engine setup so model load, KV cache init and warmup may sync
46 +# freely; flipped on by `enable_gpu_sync_check()` once warmup completes.
47 _sync_check_enabled: bool = False
48
49
50 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:
51 + """Flip the sync-check gate on, once per worker, after warmup."""
52 + if _SYNC_CHECK_MODE is None:
53 return
54 global _sync_check_enabled
55 _sync_check_enabled = True
56 _install_compile_time_sync_suppressors()
57
58
59 +_arm_lock = threading.Lock()
60 +_arm_count: int = 0
61 +_saved_sync_debug_mode: int = 0
62 +_prev_showwarning = warnings.showwarning
63 +_sync_filter_head: tuple | None = None
64 +
65 +
66 +def _sync_warning_hook(message, category, filename, lineno, file=None, line=None):
67 + """Turn torch's sync warning into an error, but only on a thread that is
68 + being checked and is outside any allow region."""
69 + if _TORCH_SYNC_WARNING in str(message):
70 + mode = _checking.get()
71 + if mode is None or _allow_depth.get():
72 + return None
73 + if mode == "error":
74 + raise RuntimeError(SYNC_ERROR_MESSAGE)
75 + return _prev_showwarning(message, category, filename, lineno, file, line)
76 +
77 +
78 +def _install_warning_hook() -> None:
79 + """(Re)install the hook and a filter that lets torch's warning reach it.
80 +
81 + Done per checked call because pytest runs each test inside
82 + `warnings.catch_warnings()`, which restores both `showwarning` and
83 + `filters`. The hook is left in place afterwards: outside a checked call
84 + the debug mode is disarmed, so torch emits nothing for it to see.
85 + """
86 + global _prev_showwarning, _sync_filter_head
87 + if warnings.showwarning is not _sync_warning_hook:
88 + _prev_showwarning = warnings.showwarning
89 + warnings.showwarning = _sync_warning_hook
90 + # "always" so the warning survives filtering and isn't deduplicated by
91 + # `__warningregistry__`. `filterwarnings` prepends, so only re-assert it
92 + # once ours is no longer in front.
93 + if warnings.filters[:1] != [_sync_filter_head]:
94 + warnings.filterwarnings(
95 + "always", message=_TORCH_SYNC_WARNING, category=UserWarning
96 + )
97 + _sync_filter_head = warnings.filters[0]
98 +
99 +
100 +@contextmanager
101 +def _checked_region(mode: str):
102 + """Police syncs on this thread for the duration of the block.
103 +
104 + The debug mode is armed per call so that syncs outside a checked region
105 + emit nothing, and refcounted because `execute_model` and `sample_tokens`
106 + nest. "warn" rather than "error" because torch's error mode raises on
107 + whichever thread synced, with no way to exempt one.
108 + """
109 + global _arm_count, _saved_sync_debug_mode
110 + _install_warning_hook()
111 + with _arm_lock:
112 + if _arm_count == 0:
113 + _saved_sync_debug_mode = torch.cuda.get_sync_debug_mode()
114 + torch.cuda.set_sync_debug_mode("warn")
115 + _arm_count += 1
116 + token = _checking.set(mode)
117 + try:
118 + yield
119 + finally:
120 + _checking.reset(token)
121 + with _arm_lock:
122 + _arm_count -= 1
123 + if _arm_count == 0:
124 + torch.cuda.set_sync_debug_mode(_saved_sync_debug_mode)
125 +
126 +
127 +@contextmanager
128 +def _allow_syncs():
129 + token = _allow_depth.set(_allow_depth.get() + 1)
130 + try:
131 + yield
132 + finally:
133 + _allow_depth.reset(token)
134 +
135 +
136 +# Shared, since `nullcontext` is stateless: reusing one instance avoids an
137 +# allocation per call, which is most of the cost when the check is disabled.
138 +_NOOP_CM = nullcontext()
139 +
140 _compile_time_suppressors_installed: bool = False
141
142
143 +def _suppressing(fn):
144 + """Allow the syncs `fn` performs on its calling thread."""
145 +
146 + @functools.wraps(fn)
147 + def wrapper(*args, **kwargs):
148 + # Not `gpu_sync_allowed()`, which no-ops while
149 + # `torch.compiler.is_compiling()` -- exactly when these run.
150 + with _allow_syncs():
151 + return fn(*args, **kwargs)
152 +
153 + return wrapper
154 +
155 +
156 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`.
157 + """Allow the syncs torch's compile passes perform.
158
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.
159 + Warmup-time compiles run before the gate flips, but lazy ones fire inside
160 + `execute_model`.
161 """
162 global _compile_time_suppressors_installed
163 if _compile_time_suppressors_installed:
164 return
165 _compile_time_suppressors_installed = True
166
53 - try: # noqa: BLE001
167 + try:
168 from torch._inductor.fx_passes import joint_graph as _jg
169
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")
170 + orig = _jg.joint_graph_passes
171 + wrapped = _suppressing(orig)
172 + # `compile_fx` imports this by value, so patching the defining module
173 + # alone misses that rebind; patch every compile-time module still
174 + # holding the original.
175 + _jg.joint_graph_passes = wrapped
176 + for name, mod in list(sys.modules.items()):
177 + if (
178 + mod is not None
179 + and name.startswith(
180 + ("torch._inductor", "torch._functorch", "torch._dynamo")
181 + )
182 + and getattr(mod, "joint_graph_passes", None) is orig
183 ):
85 - continue
86 - if getattr(_mod, "joint_graph_passes", None) is _orig_joint:
87 - setattr(_mod, "joint_graph_passes", _wrapped_joint) # noqa: B010
184 + setattr(mod, "joint_graph_passes", wrapped) # noqa: B010
185 except Exception: # pragma: no cover
186 pass
187
91 -
92 -@contextmanager
93 -def _suppress_gpu_sync_check(prev_mode: int):
94 - torch.cuda.set_sync_debug_mode(0)
188 try:
96 - yield
97 - finally:
98 - torch.cuda.set_sync_debug_mode(prev_mode)
189 + # Inductor builds its cudagraph tree lazily, so `deferred_cudagraphify`
190 + # and the `capture_begin` sync inside it can fire during
191 + # `execute_model`. It resolves `cudagraphify` as a module global at
192 + # call time, so patching the attribute is enough.
193 + from torch._inductor import cudagraph_trees as _ct
194 +
195 + _ct.cudagraphify = _suppressing(_ct.cudagraphify)
196 + except Exception: # pragma: no cover
197 + pass
198
199
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):
200 +def gpu_sync_allowed(first_only: bool = False):
201 + """Allow GPU<->CPU syncs inside the `with` block, on this thread only.
202 +
203 + With `first_only`, only the first entry from a given call site
204 + (filename, lineno) is allowed, so later syncs there are still reported.
205 + """
206 + if _SYNC_CHECK_MODE is None or torch.compiler.is_compiling():
207 + return _NOOP_CM
208 + if first_only:
209 + frame = sys._getframe(1)
210 + key = (frame.f_code.co_filename, frame.f_lineno)
211 + if key in _GPU_SYNC_ALLOWED_FIRST_SEEN:
212 + return _NOOP_CM
213 + _GPU_SYNC_ALLOWED_FIRST_SEEN.add(key)
214 + return _allow_syncs()
215 +
216 +
217 +def with_gpu_sync_check(fn):
218 + """Report GPU<->CPU syncs performed by `fn` on its calling thread.
219 +
220 + Active only once `enable_gpu_sync_check()` has flipped the gate. Other
221 + threads are never policed, so deliberate syncs there (e.g. the EPLB
222 + transfer worker) are unaffected.
223 + """
224 + if (mode := _SYNC_CHECK_MODE) is None:
225 return fn
226 +
227 + @functools.wraps(fn)
228 + def wrapper(*args, **kwargs):
229 + if not _sync_check_enabled:
230 + return fn(*args, **kwargs)
231 + with _checked_region(mode):
232 + return fn(*args, **kwargs)
233 +
234 + return wrapper