obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-16 02:38:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Nick Hill <[email protected]>
This commit is contained in:
Nick Hill
2026-08-08 17:13:40 -07:00
committed by GitHub
parent 9b0afeb4f6
commit fbff187d59
2 changed files with 288 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/utils_/test_gpu_sync_debug.py
+101 -2
View File
@@ -1,5 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import threading
import warnings
import pytest
import torch
@@ -36,7 +39,9 @@ def _causes_sync():
@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)
# `_SYNC_CHECK_MODE` is read from the env once at import, so patch the
# module attribute rather than the environment.
monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", mode)
monkeypatch.setattr(gsd, "_sync_check_enabled", True)
# Guarded syncs always pass.
@@ -51,10 +56,104 @@ def test_with_env_set(monkeypatch, mode):
with_gpu_sync_check(_causes_sync)()
@create_new_process_for_each_test()
def test_other_threads_are_not_policed(monkeypatch):
"""A background thread that syncs deliberately must not be broken by the
check being armed on the thread running the decorated function."""
monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
monkeypatch.setattr(gsd, "_sync_check_enabled", True)
def sync_on_worker():
failure: list[BaseException] = []
def worker():
try:
torch.ones(4, device="cuda").cpu()
except BaseException as exc: # pragma: no cover - failure path
failure.append(exc)
thread = threading.Thread(target=worker)
thread.start()
thread.join()
assert not failure, f"background thread raised: {failure[0]!r}"
with_gpu_sync_check(sync_on_worker)()
@create_new_process_for_each_test()
def test_allow_on_other_thread_does_not_disarm(monkeypatch):
"""`gpu_sync_allowed()` on one thread must not suppress the check on
another. It is scoped by ContextVar rather than torch's process-global
sync debug mode, which a previous implementation mutated."""
monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
monkeypatch.setattr(gsd, "_sync_check_enabled", True)
def main_syncs_while_worker_allows():
stop = threading.Event()
def worker():
with gpu_sync_allowed():
while not stop.is_set():
torch.ones(4, device="cuda").cpu()
thread = threading.Thread(target=worker)
thread.start()
try:
# Must still be reported despite the worker's open allow region.
torch.ones(4, device="cuda").cpu()
finally:
stop.set()
thread.join()
with pytest.raises(RuntimeError, match=SYNC_ERROR_MESSAGE):
with_gpu_sync_check(main_syncs_while_worker_allows)()
@create_new_process_for_each_test()
def test_suppressing_works_while_compiling(monkeypatch):
"""`_suppressing` wraps torch compile entry points, which run with
`torch.compiler.is_compiling()` true. `gpu_sync_allowed()` deliberately
no-ops in that state, so `_suppressing` must not route through it."""
monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
monkeypatch.setattr(gsd, "_sync_check_enabled", True)
# Emulate being inside a torch compile, as inductor passes are.
monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
suppressed = gsd._suppressing(lambda: torch.ones(4, device="cuda").cpu())
with_gpu_sync_check(suppressed)()
@create_new_process_for_each_test()
def test_sync_debug_mode_restored_after_checked_call(monkeypatch):
"""The mode is armed only for the duration of a checked call. Leaving it
on process-wide made every sync outside a checked region emit a
`UserWarning` whenever our handler was not the installed one."""
monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
monkeypatch.setattr(gsd, "_sync_check_enabled", True)
before = torch.cuda.get_sync_debug_mode()
def nested():
# `execute_model` -> `sample_tokens` both carry the decorator.
assert torch.cuda.get_sync_debug_mode() != 0, "armed inside"
with_gpu_sync_check(lambda: None)()
assert torch.cuda.get_sync_debug_mode() != 0, "still armed after inner"
with_gpu_sync_check(nested)()
assert torch.cuda.get_sync_debug_mode() == before
# With the mode back to its original value, torch emits nothing for a
# sync outside a checked region.
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
torch.ones(4, device="cuda").cpu()
assert not [r for r in caught if gsd._TORCH_SYNC_WARNING in str(r.message)]
@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_MODE", None)
monkeypatch.setattr(gsd, "_sync_check_enabled", True)
with_gpu_sync_check(_no_sync)()
vllm/utils/gpu_sync_debug.py
+187 -118
View File
@@ -1,8 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Detect unintended GPU<->CPU syncs in the hot path.
`torch.cuda.set_sync_debug_mode` is process-global, so we arm it at "warn"
(which never raises) and decide in `_sync_warning_hook` whether a given sync
is a failure. The scoping lives in ContextVars, which are per-thread and
per-asyncio-task, so an allow region opened on one thread is invisible to
every other one.
"""
import functools
import sys
from contextlib import contextmanager
import threading
import warnings
from contextlib import contextmanager, nullcontext
from contextvars import ContextVar
import torch
@@ -13,153 +25,210 @@ SYNC_ERROR_MESSAGE = (
"GPU<->CPU sync detected - avoid it or wrap with gpu_sync_allowed()"
)
# Warning text torch emits in "warn" mode.
_TORCH_SYNC_WARNING = "called a synchronizing CUDA operation"
_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.
# Read once: `envs.__getattr__` re-evaluates the env lambda on every access,
# which dominates the cost of a disabled `gpu_sync_allowed()`. Tests toggling
# the env var must patch this attribute instead.
_SYNC_CHECK_MODE: str | None = (
envs.VLLM_GPU_SYNC_CHECK if current_platform.is_cuda_alike() else None
)
# Configured mode while a thread is inside a checked call, else None.
_checking: ContextVar[str | None] = ContextVar("vllm_gpu_sync_checking", default=None)
# Depth of nested `gpu_sync_allowed()` regions on this thread.
_allow_depth: ContextVar[int] = ContextVar("vllm_gpu_sync_allow_depth", default=0)
# Off during engine setup so model load, KV cache init and warmup may sync
# freely; flipped on by `enable_gpu_sync_check()` once warmup completes.
_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:
"""Flip the sync-check gate on, once per worker, after warmup."""
if _SYNC_CHECK_MODE is None:
return
global _sync_check_enabled
_sync_check_enabled = True
_install_compile_time_sync_suppressors()
_arm_lock = threading.Lock()
_arm_count: int = 0
_saved_sync_debug_mode: int = 0
_prev_showwarning = warnings.showwarning
_sync_filter_head: tuple | None = None
def _sync_warning_hook(message, category, filename, lineno, file=None, line=None):
"""Turn torch's sync warning into an error, but only on a thread that is
being checked and is outside any allow region."""
if _TORCH_SYNC_WARNING in str(message):
mode = _checking.get()
if mode is None or _allow_depth.get():
return None
if mode == "error":
raise RuntimeError(SYNC_ERROR_MESSAGE)
return _prev_showwarning(message, category, filename, lineno, file, line)
def _install_warning_hook() -> None:
"""(Re)install the hook and a filter that lets torch's warning reach it.
Done per checked call because pytest runs each test inside
`warnings.catch_warnings()`, which restores both `showwarning` and
`filters`. The hook is left in place afterwards: outside a checked call
the debug mode is disarmed, so torch emits nothing for it to see.
"""
global _prev_showwarning, _sync_filter_head
if warnings.showwarning is not _sync_warning_hook:
_prev_showwarning = warnings.showwarning
warnings.showwarning = _sync_warning_hook
# "always" so the warning survives filtering and isn't deduplicated by
# `__warningregistry__`. `filterwarnings` prepends, so only re-assert it
# once ours is no longer in front.
if warnings.filters[:1] != [_sync_filter_head]:
warnings.filterwarnings(
"always", message=_TORCH_SYNC_WARNING, category=UserWarning
)
_sync_filter_head = warnings.filters[0]
@contextmanager
def _checked_region(mode: str):
"""Police syncs on this thread for the duration of the block.
The debug mode is armed per call so that syncs outside a checked region
emit nothing, and refcounted because `execute_model` and `sample_tokens`
nest. "warn" rather than "error" because torch's error mode raises on
whichever thread synced, with no way to exempt one.
"""
global _arm_count, _saved_sync_debug_mode
_install_warning_hook()
with _arm_lock:
if _arm_count == 0:
_saved_sync_debug_mode = torch.cuda.get_sync_debug_mode()
torch.cuda.set_sync_debug_mode("warn")
_arm_count += 1
token = _checking.set(mode)
try:
yield
finally:
_checking.reset(token)
with _arm_lock:
_arm_count -= 1
if _arm_count == 0:
torch.cuda.set_sync_debug_mode(_saved_sync_debug_mode)
@contextmanager
def _allow_syncs():
token = _allow_depth.set(_allow_depth.get() + 1)
try:
yield
finally:
_allow_depth.reset(token)
# Shared, since `nullcontext` is stateless: reusing one instance avoids an
# allocation per call, which is most of the cost when the check is disabled.
_NOOP_CM = nullcontext()
_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`.
def _suppressing(fn):
"""Allow the syncs `fn` performs on its calling thread."""
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.
@functools.wraps(fn)
def wrapper(*args, **kwargs):
# Not `gpu_sync_allowed()`, which no-ops while
# `torch.compiler.is_compiling()` -- exactly when these run.
with _allow_syncs():
return fn(*args, **kwargs)
return wrapper
def _install_compile_time_sync_suppressors() -> None:
"""Allow the syncs torch's compile passes perform.
Warmup-time compiles run before the gate flips, but lazy ones fire inside
`execute_model`.
"""
global _compile_time_suppressors_installed
if _compile_time_suppressors_installed:
return
_compile_time_suppressors_installed = True
try: # noqa: BLE001
try:
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")
orig = _jg.joint_graph_passes
wrapped = _suppressing(orig)
# `compile_fx` imports this by value, so patching the defining module
# alone misses that rebind; patch every compile-time module still
# holding the original.
_jg.joint_graph_passes = wrapped
for name, mod in list(sys.modules.items()):
if (
mod is not None
and name.startswith(
("torch._inductor", "torch._functorch", "torch._dynamo")
)
and getattr(mod, "joint_graph_passes", None) is orig
):
continue
if getattr(_mod, "joint_graph_passes", None) is _orig_joint:
setattr(_mod, "joint_graph_passes", _wrapped_joint) # noqa: B010
setattr(mod, "joint_graph_passes", wrapped) # noqa: B010
except Exception: # pragma: no cover
pass
try:
# Inductor builds its cudagraph tree lazily, so `deferred_cudagraphify`
# and the `capture_begin` sync inside it can fire during
# `execute_model`. It resolves `cudagraphify` as a module global at
# call time, so patching the attribute is enough.
from torch._inductor import cudagraph_trees as _ct
_ct.cudagraphify = _suppressing(_ct.cudagraphify)
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)
def gpu_sync_allowed(first_only: bool = False):
"""Allow GPU<->CPU syncs inside the `with` block, on this thread only.
With `first_only`, only the first entry from a given call site
(filename, lineno) is allowed, so later syncs there are still reported.
"""
if _SYNC_CHECK_MODE is None or torch.compiler.is_compiling():
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 _allow_syncs()
@contextmanager
def _noop_cm():
yield
def with_gpu_sync_check(fn):
"""Report GPU<->CPU syncs performed by `fn` on its calling thread.
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):
Active only once `enable_gpu_sync_check()` has flipped the gate. Other
threads are never policed, so deliberate syncs there (e.g. the EPLB
transfer worker) are unaffected.
"""
if (mode := _SYNC_CHECK_MODE) is None:
return fn
@functools.wraps(fn)
def wrapper(*args, **kwargs):
if not _sync_check_enabled:
return fn(*args, **kwargs)
with _checked_region(mode):
return fn(*args, **kwargs)
return wrapper