diff --git a/tests/utils_/test_gpu_sync_debug.py b/tests/utils_/test_gpu_sync_debug.py new file mode 100644 index 0000000000..ea9b76f34b --- /dev/null +++ b/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)() diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index 2348ff3191..742ec55e6e 100644 --- a/vllm/compilation/compiler_interface.py +++ b/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" diff --git a/vllm/envs.py b/vllm/envs.py index 08314a8c88..ab6184d22d 100755 --- a/vllm/envs.py +++ b/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), diff --git a/vllm/utils/gpu_sync_debug.py b/vllm/utils/gpu_sync_debug.py new file mode 100644 index 0000000000..1e2114f3b5 --- /dev/null +++ b/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 diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 4aa8ca5ca3..589a16576e 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/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: