| 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 | |
| 9 | import vllm.utils.gpu_sync_debug as gsd |
| 10 | from vllm.utils.gpu_sync_debug import ( |
| 11 | SYNC_ERROR_MESSAGE, |
| 12 | gpu_sync_allowed, |
| 13 | with_gpu_sync_check, |
| 14 | ) |
| 15 | |
| 16 | from ..utils import create_new_process_for_each_test |
| 17 | |
| 18 | pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 19 | |
| 20 | |
| 21 | def _no_sync(): |
| 22 | # Pure on-GPU compute, no implicit CPU sync... |
| 23 | x = torch.ones(4, device="cuda") + 1 |
| 24 | # ...plus a sync that we explicitly allow. |
| 25 | with gpu_sync_allowed(): |
| 26 | return x.cpu() |
| 27 | |
| 28 | |
| 29 | def _causes_sync(): |
| 30 | x = torch.ones(4, device="cuda") |
| 31 | # An allowed sync (suppressed)... |
| 32 | with gpu_sync_allowed(): |
| 33 | x.cpu() |
| 34 | # ...then an un-allowed sync that should trip the check. |
| 35 | return x.cpu() |
| 36 | |
| 37 | |
| 38 | @pytest.mark.parametrize("mode", ["warn", "error"]) |
| 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. |
| 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. |
| 48 | with_gpu_sync_check(_no_sync)() |
| 49 | |
| 50 | if mode == "error": |
| 51 | # "error" mode turns the stray sync into a RuntimeError. |
| 52 | with pytest.raises(RuntimeError, match=SYNC_ERROR_MESSAGE): |
| 53 | with_gpu_sync_check(_causes_sync)() |
| 54 | else: |
| 55 | # "warn" mode only warns, so the call still succeeds. |
| 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_copy_checker_is_traceable_while_compiling(monkeypatch): |
| 128 | """The copy wrapper must bypass ContextVars while Dynamo is tracing.""" |
| 129 | monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") |
| 130 | monkeypatch.setattr(gsd, "_sync_check_enabled", True) |
| 131 | gsd._install_copy_checkers() |
| 132 | |
| 133 | def convert_like(x, like): |
| 134 | return x.to(device=like.device, dtype=like.dtype) |
| 135 | |
| 136 | compiled = torch.compile(convert_like, fullgraph=True) |
| 137 | result = compiled( |
| 138 | torch.ones(4, device="cuda", dtype=torch.float64), |
| 139 | torch.ones(4, device="cuda", dtype=torch.float32), |
| 140 | ) |
| 141 | assert result.dtype == torch.float32 |
| 142 | |
| 143 | |
| 144 | @create_new_process_for_each_test() |
| 145 | def test_sync_debug_mode_restored_after_checked_call(monkeypatch): |
| 146 | """The mode is armed only for the duration of a checked call. Leaving it |
| 147 | on process-wide made every sync outside a checked region emit a |
| 148 | `UserWarning` whenever our handler was not the installed one.""" |
| 149 | monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") |
| 150 | monkeypatch.setattr(gsd, "_sync_check_enabled", True) |
| 151 | |
| 152 | before = torch.cuda.get_sync_debug_mode() |
| 153 | |
| 154 | def nested(): |
| 155 | # `execute_model` -> `sample_tokens` both carry the decorator. |
| 156 | assert torch.cuda.get_sync_debug_mode() != 0, "armed inside" |
| 157 | with_gpu_sync_check(lambda: None)() |
| 158 | assert torch.cuda.get_sync_debug_mode() != 0, "still armed after inner" |
| 159 | |
| 160 | with_gpu_sync_check(nested)() |
| 161 | assert torch.cuda.get_sync_debug_mode() == before |
| 162 | |
| 163 | # With the mode back to its original value, torch emits nothing for a |
| 164 | # sync outside a checked region. |
| 165 | with warnings.catch_warnings(record=True) as caught: |
| 166 | warnings.simplefilter("always") |
| 167 | torch.ones(4, device="cuda").cpu() |
| 168 | assert not [r for r in caught if gsd._TORCH_SYNC_WARNING in str(r.message)] |
| 169 | |
| 170 | |
| 171 | @create_new_process_for_each_test() |
| 172 | def test_without_env_set(monkeypatch): |
| 173 | # Env unset: the decorator is a pass-through, no sync is detected. |
| 174 | monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", None) |
| 175 | monkeypatch.setattr(gsd, "_sync_check_enabled", True) |
| 176 | |
| 177 | with_gpu_sync_check(_no_sync)() |
| 178 | with_gpu_sync_check(_causes_sync)() |
| 179 | |
| 180 | |
| 181 | def _pageable_h2d_nonblocking(): |
| 182 | # Pageable source: the "async" copy is staged through pageable memory. |
| 183 | torch.zeros(4).to("cuda", non_blocking=True) |
| 184 | |
| 185 | |
| 186 | def _pinned_h2d_nonblocking(): |
| 187 | # Pinned, contiguous, same dtype: genuinely asynchronous. |
| 188 | torch.zeros(4).pin_memory().to("cuda", non_blocking=True) |
| 189 | |
| 190 | |
| 191 | def _noncontiguous_h2d_nonblocking(): |
| 192 | # Pinned but gapped layout (strided slice): staged through a pageable temp. |
| 193 | torch.zeros(8).pin_memory()[::2].to("cuda", non_blocking=True) |
| 194 | |
| 195 | |
| 196 | def _transposed_h2d_nonblocking(): |
| 197 | # Dense permuted layout: copied with pitched cudaMemcpy2DAsync, stays async. |
| 198 | torch.zeros(4, 4).pin_memory().t().to("cuda", non_blocking=True) |
| 199 | |
| 200 | |
| 201 | def _dtype_converting_h2d_nonblocking(): |
| 202 | # Pinned and contiguous: the conversion runs GPU-side, staying async. |
| 203 | torch.zeros(4, dtype=torch.float64).pin_memory().to( |
| 204 | "cuda", dtype=torch.float32, non_blocking=True |
| 205 | ) |
| 206 | |
| 207 | |
| 208 | def _pageable_h2d_via_cuda(): |
| 209 | torch.zeros(4).cuda(non_blocking=True) |
| 210 | |
| 211 | |
| 212 | def _pageable_h2d_via_copy_(): |
| 213 | torch.zeros(4, device="cuda").copy_(torch.zeros(4), non_blocking=True) |
| 214 | |
| 215 | |
| 216 | def _d2h_via_to(): |
| 217 | # Genuinely asynchronous: torch allocates the destination pinned. |
| 218 | torch.zeros(4, device="cuda").to("cpu", non_blocking=True) |
| 219 | |
| 220 | |
| 221 | def _pageable_d2h_via_copy_(): |
| 222 | torch.zeros(4).copy_(torch.zeros(4, device="cuda"), non_blocking=True) |
| 223 | |
| 224 | |
| 225 | def _pinned_d2h_via_copy_(): |
| 226 | # Pinned, contiguous, same dtype: genuinely asynchronous. |
| 227 | torch.zeros(4).pin_memory().copy_(torch.zeros(4, device="cuda"), non_blocking=True) |
| 228 | |
| 229 | |
| 230 | def _noncontiguous_d2h_via_copy_(): |
| 231 | # Pinned but gapped (strided slice) destination. |
| 232 | torch.zeros(8).pin_memory()[::2].copy_( |
| 233 | torch.zeros(4, device="cuda"), non_blocking=True |
| 234 | ) |
| 235 | |
| 236 | |
| 237 | def _transposed_d2h_via_copy_(): |
| 238 | # Dense permuted destination: pitched cudaMemcpy2DAsync, stays async. |
| 239 | torch.zeros(4, 4).pin_memory().t().copy_( |
| 240 | torch.zeros(4, 4, device="cuda"), non_blocking=True |
| 241 | ) |
| 242 | |
| 243 | |
| 244 | def _empty_h2d_nonblocking(): |
| 245 | # Empty (e.g. first-step penalties): no CUDA call is issued at all. |
| 246 | torch.zeros(0).to("cuda", non_blocking=True) |
| 247 | |
| 248 | |
| 249 | def _dtype_converting_d2h_via_copy_(): |
| 250 | # Pinned and contiguous: the conversion runs GPU-side, staying async. |
| 251 | torch.zeros(4).pin_memory().copy_( |
| 252 | torch.zeros(4, dtype=torch.float64, device="cuda"), non_blocking=True |
| 253 | ) |
| 254 | |
| 255 | |
| 256 | @pytest.mark.parametrize("mode", ["warn", "error"]) |
| 257 | @pytest.mark.parametrize( |
| 258 | "fn", |
| 259 | [ |
| 260 | _pageable_h2d_nonblocking, |
| 261 | _noncontiguous_h2d_nonblocking, |
| 262 | _pageable_h2d_via_cuda, |
| 263 | _pageable_h2d_via_copy_, |
| 264 | _pageable_d2h_via_copy_, |
| 265 | _noncontiguous_d2h_via_copy_, |
| 266 | ], |
| 267 | ) |
| 268 | @create_new_process_for_each_test() |
| 269 | def test_implicit_copy_sync_detected(monkeypatch, mode, fn): |
| 270 | """`non_blocking=True` CPU<->CUDA copies with a pageable or |
| 271 | non-densely-laid-out CPU tensor may block the host without tripping |
| 272 | torch's sync debug mode; the `Tensor.to`/`cuda`/`copy_` wrappers must |
| 273 | flag them. |
| 274 | """ |
| 275 | monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", mode) |
| 276 | monkeypatch.setattr(gsd, "_sync_check_enabled", True) |
| 277 | gsd._install_copy_checkers() |
| 278 | |
| 279 | if mode == "error": |
| 280 | with pytest.raises(RuntimeError, match="Implicit GPU<->CPU sync"): |
| 281 | with_gpu_sync_check(fn)() |
| 282 | else: |
| 283 | with pytest.warns(UserWarning, match="Implicit GPU<->CPU sync"): |
| 284 | with_gpu_sync_check(fn)() |
| 285 | |
| 286 | |
| 287 | @create_new_process_for_each_test() |
| 288 | def test_genuinely_async_transfers_pass(monkeypatch): |
| 289 | """Pinned, densely laid out CPU tensors make `non_blocking=True` truly |
| 290 | asynchronous in both directions -- even with a dtype conversion, which |
| 291 | runs GPU-side, and even permuted (e.g. transposed), which uses pitched |
| 292 | cudaMemcpy2D/3DAsync; D2H `Tensor.to` allocates a pinned destination. |
| 293 | None of these may be flagged.""" |
| 294 | monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") |
| 295 | monkeypatch.setattr(gsd, "_sync_check_enabled", True) |
| 296 | gsd._install_copy_checkers() |
| 297 | |
| 298 | for fn in ( |
| 299 | _pinned_h2d_nonblocking, |
| 300 | _pinned_d2h_via_copy_, |
| 301 | _d2h_via_to, |
| 302 | _dtype_converting_h2d_nonblocking, |
| 303 | _dtype_converting_d2h_via_copy_, |
| 304 | _transposed_h2d_nonblocking, |
| 305 | _transposed_d2h_via_copy_, |
| 306 | _empty_h2d_nonblocking, |
| 307 | ): |
| 308 | fn() # Warm up outside the checked region. |
| 309 | with_gpu_sync_check(fn)() |
| 310 | |
| 311 | |
| 312 | @create_new_process_for_each_test() |
| 313 | def test_implicit_copy_sync_can_be_allowed(monkeypatch): |
| 314 | """`gpu_sync_allowed()` must exempt implicit copy syncs too.""" |
| 315 | monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") |
| 316 | monkeypatch.setattr(gsd, "_sync_check_enabled", True) |
| 317 | gsd._install_copy_checkers() |
| 318 | |
| 319 | def allowed_pageable_copies(): |
| 320 | with gpu_sync_allowed(): |
| 321 | _pageable_h2d_nonblocking() |
| 322 | _pageable_d2h_via_copy_() |
| 323 | |
| 324 | with_gpu_sync_check(allowed_pageable_copies)() |