@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4aadb0f14d

[Core] Enhance cpu<->gpu sync checking to include paged async copies (#53491)

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

Nick Hill committed Sep 8, 2026 at 17:26 UTC 4aadb0f14d4c0d6397cfba271f11906e980b18c1
5 files changed +332 -14
tests/utils_/test_gpu_sync_debug.py
+146
@@ -158,3 +158,149 @@ def test_without_env_set(monkeypatch):
158
159 with_gpu_sync_check(_no_sync)()
160 with_gpu_sync_check(_causes_sync)()
161 +
162 +
163 +def _pageable_h2d_nonblocking():
164 + # Pageable source: the "async" copy is staged through pageable memory.
165 + torch.zeros(4).to("cuda", non_blocking=True)
166 +
167 +
168 +def _pinned_h2d_nonblocking():
169 + # Pinned, contiguous, same dtype: genuinely asynchronous.
170 + torch.zeros(4).pin_memory().to("cuda", non_blocking=True)
171 +
172 +
173 +def _noncontiguous_h2d_nonblocking():
174 + # Pinned but gapped layout (strided slice): staged through a pageable temp.
175 + torch.zeros(8).pin_memory()[::2].to("cuda", non_blocking=True)
176 +
177 +
178 +def _transposed_h2d_nonblocking():
179 + # Dense permuted layout: copied with pitched cudaMemcpy2DAsync, stays async.
180 + torch.zeros(4, 4).pin_memory().t().to("cuda", non_blocking=True)
181 +
182 +
183 +def _dtype_converting_h2d_nonblocking():
184 + # Pinned and contiguous: the conversion runs GPU-side, staying async.
185 + torch.zeros(4, dtype=torch.float64).pin_memory().to(
186 + "cuda", dtype=torch.float32, non_blocking=True
187 + )
188 +
189 +
190 +def _pageable_h2d_via_cuda():
191 + torch.zeros(4).cuda(non_blocking=True)
192 +
193 +
194 +def _pageable_h2d_via_copy_():
195 + torch.zeros(4, device="cuda").copy_(torch.zeros(4), non_blocking=True)
196 +
197 +
198 +def _d2h_via_to():
199 + # Genuinely asynchronous: torch allocates the destination pinned.
200 + torch.zeros(4, device="cuda").to("cpu", non_blocking=True)
201 +
202 +
203 +def _pageable_d2h_via_copy_():
204 + torch.zeros(4).copy_(torch.zeros(4, device="cuda"), non_blocking=True)
205 +
206 +
207 +def _pinned_d2h_via_copy_():
208 + # Pinned, contiguous, same dtype: genuinely asynchronous.
209 + torch.zeros(4).pin_memory().copy_(torch.zeros(4, device="cuda"), non_blocking=True)
210 +
211 +
212 +def _noncontiguous_d2h_via_copy_():
213 + # Pinned but gapped (strided slice) destination.
214 + torch.zeros(8).pin_memory()[::2].copy_(
215 + torch.zeros(4, device="cuda"), non_blocking=True
216 + )
217 +
218 +
219 +def _transposed_d2h_via_copy_():
220 + # Dense permuted destination: pitched cudaMemcpy2DAsync, stays async.
221 + torch.zeros(4, 4).pin_memory().t().copy_(
222 + torch.zeros(4, 4, device="cuda"), non_blocking=True
223 + )
224 +
225 +
226 +def _empty_h2d_nonblocking():
227 + # Empty (e.g. first-step penalties): no CUDA call is issued at all.
228 + torch.zeros(0).to("cuda", non_blocking=True)
229 +
230 +
231 +def _dtype_converting_d2h_via_copy_():
232 + # Pinned and contiguous: the conversion runs GPU-side, staying async.
233 + torch.zeros(4).pin_memory().copy_(
234 + torch.zeros(4, dtype=torch.float64, device="cuda"), non_blocking=True
235 + )
236 +
237 +
238 +@pytest.mark.parametrize("mode", ["warn", "error"])
239 +@pytest.mark.parametrize(
240 + "fn",
241 + [
242 + _pageable_h2d_nonblocking,
243 + _noncontiguous_h2d_nonblocking,
244 + _pageable_h2d_via_cuda,
245 + _pageable_h2d_via_copy_,
246 + _pageable_d2h_via_copy_,
247 + _noncontiguous_d2h_via_copy_,
248 + ],
249 +)
250 +@create_new_process_for_each_test()
251 +def test_implicit_copy_sync_detected(monkeypatch, mode, fn):
252 + """`non_blocking=True` CPU<->CUDA copies with a pageable or
253 + non-densely-laid-out CPU tensor may block the host without tripping
254 + torch's sync debug mode; the `Tensor.to`/`cuda`/`copy_` wrappers must
255 + flag them.
256 + """
257 + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", mode)
258 + monkeypatch.setattr(gsd, "_sync_check_enabled", True)
259 + gsd._install_copy_checkers()
260 +
261 + if mode == "error":
262 + with pytest.raises(RuntimeError, match="Implicit GPU<->CPU sync"):
263 + with_gpu_sync_check(fn)()
264 + else:
265 + with pytest.warns(UserWarning, match="Implicit GPU<->CPU sync"):
266 + with_gpu_sync_check(fn)()
267 +
268 +
269 +@create_new_process_for_each_test()
270 +def test_genuinely_async_transfers_pass(monkeypatch):
271 + """Pinned, densely laid out CPU tensors make `non_blocking=True` truly
272 + asynchronous in both directions -- even with a dtype conversion, which
273 + runs GPU-side, and even permuted (e.g. transposed), which uses pitched
274 + cudaMemcpy2D/3DAsync; D2H `Tensor.to` allocates a pinned destination.
275 + None of these may be flagged."""
276 + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
277 + monkeypatch.setattr(gsd, "_sync_check_enabled", True)
278 + gsd._install_copy_checkers()
279 +
280 + for fn in (
281 + _pinned_h2d_nonblocking,
282 + _pinned_d2h_via_copy_,
283 + _d2h_via_to,
284 + _dtype_converting_h2d_nonblocking,
285 + _dtype_converting_d2h_via_copy_,
286 + _transposed_h2d_nonblocking,
287 + _transposed_d2h_via_copy_,
288 + _empty_h2d_nonblocking,
289 + ):
290 + fn() # Warm up outside the checked region.
291 + with_gpu_sync_check(fn)()
292 +
293 +
294 +@create_new_process_for_each_test()
295 +def test_implicit_copy_sync_can_be_allowed(monkeypatch):
296 + """`gpu_sync_allowed()` must exempt implicit copy syncs too."""
297 + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
298 + monkeypatch.setattr(gsd, "_sync_check_enabled", True)
299 + gsd._install_copy_checkers()
300 +
301 + def allowed_pageable_copies():
302 + with gpu_sync_allowed():
303 + _pageable_h2d_nonblocking()
304 + _pageable_d2h_via_copy_()
305 +
306 + with_gpu_sync_check(allowed_pageable_copies)()
vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
+12 -4
@@ -135,7 +135,11 @@ class ExampleConnector(KVConnectorBase_V1):
135 slot_mapping (torch.Tensor): the slot mapping. In shape
136 [num_tokens].
137 """
138 - slot_mapping = slot_mapping.to(dst_kv_cache_layer.device, non_blocking=True)
138 + # Deliberate bulk transfers in this reference connector.
139 + with gpu_sync_allowed():
140 + slot_mapping = slot_mapping.to(
141 + dst_kv_cache_layer.device, non_blocking=True
142 + )
143 if isinstance(attn_metadata, MLACommonMetadata):
144 # [B, 1, N, C] -> [B * N, C]; slot_mapping indexes B * N slots.
145 dst_kv_cache_layer = dst_kv_cache_layer.reshape(
@@ -177,8 +181,10 @@ class ExampleConnector(KVConnectorBase_V1):
181 filename = self._generate_filename_debug(
182 layer_name, request.token_ids, request.mm_hashes
183 )
180 - kv_cache_cpu = safetensors.torch.load_file(filename)["kv_cache"]
181 - kv_cache = kv_cache_cpu.to(kv_cache_layer.device, non_blocking=True)
184 + # Deliberate bulk load in this reference connector.
185 + with gpu_sync_allowed():
186 + kv_cache_cpu = safetensors.torch.load_file(filename)["kv_cache"]
187 + kv_cache = kv_cache_cpu.to(kv_cache_layer.device, non_blocking=True)
188 if isinstance(attn_metadata, dict):
189 inject_kv_into_layer(
190 kv_cache_layer,
@@ -224,7 +230,9 @@ class ExampleConnector(KVConnectorBase_V1):
230
231 The layer is a standardized [B, H, N, C] per-layer view (H == 1 for MLA).
232 """
227 - slot_mapping = slot_mapping.to(layer.device, non_blocking=True)
233 + # Deliberate bulk transfer in this reference connector.
234 + with gpu_sync_allowed():
235 + slot_mapping = slot_mapping.to(layer.device, non_blocking=True)
236 if isinstance(attn_metadata, MLACommonMetadata):
237 # [B, 1, N, C] -> [B * N, C]; slot_mapping indexes B * N slots.
238 return layer.reshape(-1, layer.shape[-1])[slot_mapping, ...]
vllm/envs.py
+5 -3
@@ -642,9 +642,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
642 "VLLM_TRITON_USE_TD": lambda: {"1": True, "0": False}.get(
643 os.getenv("VLLM_TRITON_USE_TD", "").strip()
644 ),
645 - # If set, enable PyTorch's GPU<->CPU synchronization debug mode around
646 - # the worker's `execute_model` and `sample_tokens` calls. Valid values
647 - # are "warn" (print a warning on each sync) or "error" (raise on sync).
645 + # If set, enable GPU<->CPU synchronization checking around the worker's
646 + # `execute_model` and `sample_tokens` calls, via PyTorch's sync debug mode
647 + # plus wrappers flagging `non_blocking` CPU<->CUDA copies that silently
648 + # block the host (CPU tensors that are not pinned or not densely laid out).
649 + # Valid values are "warn" (warn on each sync) or "error" (raise on sync).
650 # Unset disables the check. See `torch.cuda.set_sync_debug_mode`.
651 "VLLM_GPU_SYNC_CHECK": env_with_choices(
652 "VLLM_GPU_SYNC_CHECK", None, ["warn", "error"]
vllm/model_executor/offloader/uva.py
+10 -7
@@ -3,6 +3,7 @@
3 """UVA-based CPU offloading using Unified Virtual Addressing."""
4
5 from collections.abc import Generator
6 +from contextlib import nullcontext
7
8 import torch
9 import torch.nn as nn
@@ -11,6 +12,7 @@ from torch.func import functional_call
12 import vllm.envs as envs
13 from vllm.logger import init_logger
14 from vllm.model_executor.offloader.base import BaseOffloader, should_pin_memory
15 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed
16 from vllm.utils.mem_utils import format_gib
17 from vllm.utils.platform_utils import is_uva_available
18 from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
@@ -126,13 +128,14 @@ class UVAOffloader(BaseOffloader):
128
129 def forward(*args, **kwargs):
130 module.forward = original_forward
129 - device_state = {
130 - # here we blindly call `to(device)`
131 - # if the parameter is already on the device,
132 - # it will be a no-op
133 - k: v.to(device, non_blocking=True)
134 - for k, v in module.state_dict().items()
135 - }
131 + with nullcontext() if self.pin_memory else gpu_sync_allowed():
132 + device_state = {
133 + # here we blindly call `to(device)`
134 + # if the parameter is already on the device,
135 + # it will be a no-op
136 + k: v.to(device, non_blocking=True)
137 + for k, v in module.state_dict().items()
138 + }
139
140 # set `tie_weights=False` as tied weights in original model
141 # become untied when calling .to(device) individually
vllm/utils/gpu_sync_debug.py
+159
@@ -7,6 +7,12 @@
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 +torch's mode only instruments explicit synchronizing calls, so it misses
12 +`non_blocking=True` CPU<->CUDA copies that the CUDA driver silently stages
13 +through pageable memory (when the CPU tensor is not pinned or not densely
14 +laid out). Those are caught by the wrappers on `Tensor.to`/`Tensor.cuda`/
15 +`Tensor.copy_` installed by `_install_copy_checkers()`.
16 """
17
18 import functools
@@ -54,6 +60,7 @@ def enable_gpu_sync_check() -> None:
60 global _sync_check_enabled
61 _sync_check_enabled = True
62 _install_compile_time_sync_suppressors()
63 + _install_copy_checkers()
64
65
66 _arm_lock = threading.Lock()
@@ -197,6 +204,158 @@ def _install_compile_time_sync_suppressors() -> None:
204 pass
205
206
207 +# ---------------------------------------------------------------------------
208 +# Implicit sync detection in `non_blocking` CPU<->CUDA copies.
209 +#
210 +# A `non_blocking=True` copy between CPU and CUDA still blocks the host when
211 +# the CPU side is not pinned or not densely laid out (gapped strided slices
212 +# stage through a pageable temp): D2H copies into such memory wait for the
213 +# stream, and H2D copies from it do too beyond a driver size threshold.
214 +# None of that goes through an API torch's sync debug mode instruments, so
215 +# wrap the Python-level transfer entry points and flag such copies on threads
216 +# in a checked region.
217 +
218 +IMPLICIT_SYNC_ERROR_MESSAGE = (
219 + "Implicit GPU<->CPU sync detected - use a pinned, densely laid out CPU "
220 + "tensor, or wrap with gpu_sync_allowed()"
221 +)
222 +
223 +# Captured at import; `_install_copy_checkers` replaces these attributes.
224 +_torch_to = torch.Tensor.to
225 +_torch_cuda = torch.Tensor.cuda
226 +_torch_copy_ = torch.Tensor.copy_
227 +
228 +# torch's own parser for `Tensor.to`'s overloaded signature, returning
229 +# (device, dtype, non_blocking, memory_format) (also used by torch._dynamo).
230 +_parse_to = getattr(torch._C._nn, "_parse_to", None)
231 +
232 +
233 +def _active_check_mode() -> str | None:
234 + """The mode the calling thread is being checked in, if any."""
235 + if not _sync_check_enabled or _allow_depth.get() or torch.compiler.is_compiling():
236 + return None
237 + return _checking.get()
238 +
239 +
240 +def _is_dense(t: torch.Tensor) -> bool:
241 + """Whether `t`'s elements fill a contiguous storage range, possibly
242 + permuted (e.g. a transposed matrix). Such layouts can be copied with
243 + pitched cudaMemcpy2D/3DAsync directly from/to pinned memory, whereas
244 + gapped layouts (e.g. strided slices) stage through a pageable temp."""
245 + if t.is_contiguous():
246 + # The common case, answered by a memoized TensorImpl flag.
247 + return True
248 + expected_stride = 1
249 + for stride, size in sorted(zip(t.stride(), t.shape)):
250 + if size <= 1:
251 + continue
252 + if stride != expected_stride:
253 + return False
254 + expected_stride *= size
255 + return True
256 +
257 +
258 +def _cpu_copy_stall_reason(
259 + cpu_tensor: torch.Tensor, device: torch.device | None
260 +) -> str | None:
261 + """Why a `non_blocking` copy between `cpu_tensor` and a CUDA tensor would
262 + silently block the host, or None if it is genuinely asynchronous.
263 +
264 + `device` describes the other side of the copy; a non-CUDA `device` means
265 + the copy does not involve the GPU at all.
266 + """
267 + if device is None or device.type != "cuda" or cpu_tensor.device.type != "cpu":
268 + return None
269 + if cpu_tensor.numel() == 0:
270 + return None # Nothing to copy; torch issues no CUDA call at all.
271 + if not _is_dense(cpu_tensor):
272 + return "the CPU tensor is not densely laid out"
273 + if not cpu_tensor.is_pinned():
274 + return "the CPU tensor is not pinned"
275 + return None
276 +
277 +
278 +def _report_implicit_sync(op: str, direction: str, reason: str, mode: str) -> None:
279 + msg = (
280 + f"Tensor.{op}: a non_blocking {direction} copy that may be silently "
281 + f"synchronous ({reason}). {IMPLICIT_SYNC_ERROR_MESSAGE}"
282 + )
283 + if mode == "error":
284 + raise RuntimeError(msg)
285 + warnings.warn(msg, stacklevel=3)
286 +
287 +
288 +def _checked_to(self, *args, **kwargs):
289 + # Only H2D is checked: D2H `to` allocates a pinned destination in
290 + # supported torch versions, so it is genuinely asynchronous.
291 + if (mode := _active_check_mode()) is not None and _parse_to is not None:
292 + try:
293 + device, _, non_blocking, _ = _parse_to(*args, **kwargs)
294 + except Exception:
295 + pass # Invalid args; let the original call raise.
296 + else:
297 + if non_blocking and (reason := _cpu_copy_stall_reason(self, device)):
298 + _report_implicit_sync("to", "H2D", reason, mode)
299 + return _torch_to(self, *args, **kwargs)
300 +
301 +
302 +def _checked_cuda(self, *args, **kwargs):
303 + # Signature: cuda(device=None, non_blocking=False, memory_format=None).
304 + if (mode := _active_check_mode()) is not None:
305 + device = args[0] if args else kwargs.get("device")
306 + non_blocking = args[1] if len(args) > 1 else kwargs.get("non_blocking", False)
307 + if non_blocking:
308 + if device is None:
309 + device = torch.device("cuda")
310 + elif isinstance(device, int):
311 + device = torch.device("cuda", device)
312 + elif not isinstance(device, torch.device):
313 + device = torch.device(device)
314 + if reason := _cpu_copy_stall_reason(self, device):
315 + _report_implicit_sync("cuda", "H2D", reason, mode)
316 + return _torch_cuda(self, *args, **kwargs)
317 +
318 +
319 +def _checked_copy_(self, *args, **kwargs):
320 + # Signature: copy_(src, non_blocking=False).
321 + if (
322 + (mode := _active_check_mode()) is not None
323 + and args
324 + and isinstance(src := args[0], torch.Tensor)
325 + ):
326 + non_blocking = args[1] if len(args) > 1 else kwargs.get("non_blocking", False)
327 + if non_blocking:
328 + if self.device.type == "cuda":
329 + reason = _cpu_copy_stall_reason(src, self.device)
330 + direction = "H2D"
331 + else:
332 + reason = _cpu_copy_stall_reason(self, src.device)
333 + direction = "D2H"
334 + if reason:
335 + _report_implicit_sync("copy_", direction, reason, mode)
336 + return _torch_copy_(self, *args, **kwargs)
337 +
338 +
339 +_copy_checkers_installed: bool = False
340 +
341 +
342 +def _install_copy_checkers() -> None:
343 + """Wrap the Python-level CPU<->CUDA transfer entry points with the checks
344 + above.
345 +
346 + The patch is process-global, but it only reports on threads inside a
347 + checked region, matching the scoping of the torch-level check. Copies
348 + issued from C++ (e.g. inside compiled graphs) bypass it.
349 + """
350 + global _copy_checkers_installed
351 + if _copy_checkers_installed:
352 + return
353 + _copy_checkers_installed = True
354 + torch.Tensor.to = _checked_to # type: ignore[method-assign]
355 + torch.Tensor.cuda = _checked_cuda # type: ignore[method-assign]
356 + torch.Tensor.copy_ = _checked_copy_ # type: ignore[method-assign]
357 +
358 +
359 def gpu_sync_allowed(first_only: bool = False):
360 """Allow GPU<->CPU syncs inside the `with` block, on this thread only.
361