[Bugfix] Make GPU sync checks safe under torch.compile (#56904)
Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Sherlock <sherlock@raft.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Kevin H. Luu committed
Sep 15, 2026 at 21:37 UTC
af1c01499b289be555c475669ba50a88e96d846e
3 files changed
+32
-3
tests/models/multimodal/generation/test_voxtral_realtime.py
+2
-1
@@ -21,7 +21,7 @@ from vllm.utils.math_utils import cdiv
21
from vllm.v1.engine.async_llm import AsyncLLM
22
from vllm.v1.kv_cache_interface import SlidingWindowSpec
23
24
-from ....utils import ROCM_ENGINE_KWARGS
24
+from ....utils import ROCM_ENGINE_KWARGS, create_new_process_for_each_test
25
26
MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
27
AUDIO_LAYER_NAME = "whisper_encoder.whisper_encoder.layers.0.layers.self_attn.attn"
@@ -192,6 +192,7 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypa
192
"cudagraph_mode",
193
[CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE],
194
)
195
+@create_new_process_for_each_test(method="spawn")
196
def test_voxtral_realtime_cudagraph(
197
audio_assets, tokenizer, vllm_runner, monkeypatch, cudagraph_mode
198
):
tests/utils_/test_gpu_sync_debug.py
+18
@@ -123,6 +123,24 @@ def test_suppressing_works_while_compiling(monkeypatch):
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
vllm/utils/gpu_sync_debug.py
+12
-2
@@ -225,6 +225,16 @@ _torch_to = torch.Tensor.to
225
_torch_cuda = torch.Tensor.cuda
226
_torch_copy_ = torch.Tensor.copy_
227
228
+
229
+# Once ``Tensor.to`` is replaced below, Dynamo sees the saved C++ descriptor
230
+# through this module global instead of as a recognized tensor method. Put the
231
+# descriptor behind a weak-referenceable Python callable so ``allow_in_graph``
232
+# can admit it and the downstream compiler can lower the original operation.
233
+@torch.compiler.allow_in_graph
234
+def _torch_to_in_graph(self, *args, **kwargs):
235
+ return _torch_to(self, *args, **kwargs)
236
+
237
+
238
# torch's own parser for `Tensor.to`'s overloaded signature, returning
239
# (device, dtype, non_blocking, memory_format) (also used by torch._dynamo).
240
_parse_to = getattr(torch._C._nn, "_parse_to", None)
@@ -232,7 +242,7 @@ _parse_to = getattr(torch._C._nn, "_parse_to", None)
242
243
def _active_check_mode() -> str | None:
244
"""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():
245
+ if not _sync_check_enabled or torch.compiler.is_compiling() or _allow_depth.get():
246
return None
247
return _checking.get()
248
@@ -296,7 +306,7 @@ def _checked_to(self, *args, **kwargs):
306
else:
307
if non_blocking and (reason := _cpu_copy_stall_reason(self, device)):
308
_report_implicit_sync("to", "H2D", reason, mode)
299
- return _torch_to(self, *args, **kwargs)
309
+ return _torch_to_in_graph(self, *args, **kwargs)
310
311
312
def _checked_cuda(self, *args, **kwargs):