[XPU] Fix CUDA API shims breaking Torch Dynamo during AOT compile (#43092)
Signed-off-by: Łukasz Ślusarczyk <lukasz.slusarczyk@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Łukasz Ślusarczyk committed
Jul 6, 2026 at 02:29 UTC
78a04c208d1f6fad27a949aa1a7f53ae638e86e3
2 files changed
+56
-7
tests/v1/worker/test_xpu_model_runner.py
new
+46
@@ -0,0 +1,46 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Unit tests for ``vllm.v1.worker.xpu_model_runner`` (XPU worker / CUDA shims)."""
5
+
6
+import pytest
7
+import torch
8
+from torch._dynamo.variables.torch import TorchInGraphFunctionVariable
9
+
10
+from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper
11
+
12
+# XPU-only: needs distinct torch.cuda vs torch.xpu current_stream symbols.
13
+pytestmark = pytest.mark.skipif(
14
+ not hasattr(torch, "xpu") or not hasattr(torch.xpu, "current_stream"),
15
+ reason="torch.xpu.current_stream is required",
16
+)
17
+
18
+
19
+# Child process: patched torch.cuda must not leak to other tests in the session.
20
+@pytest.mark.forked
21
+def test_torch_cuda_wrapper_allows_dynamo_handler_registration() -> None:
22
+ """Guard against XPU CUDA shim breaking Torch Dynamo during AOT compile.
23
+
24
+ Before the fix, ``_torch_cuda_wrapper`` assigned
25
+ ``torch.cuda.current_stream = torch.xpu.current_stream`` (same function object).
26
+ On the first AOT/profile run, Dynamo builds its in-graph handler table and
27
+ registers ``torch.cuda.current_stream`` and ``torch.xpu.current_stream``
28
+ separately; duplicate identity triggers::
29
+
30
+ AssertionError: Handler already registered for <function current_stream ...>
31
+
32
+ That surfaced as EngineCore failing in ``profile_run`` / ``_get_handlers()``.
33
+ The fix uses distinct shim callables so both can be registered.
34
+
35
+ This test replays the post-init state (wrapper applied, patches left on
36
+ ``torch.cuda``) and checks that Dynamo's real ``_get_handlers()`` succeeds.
37
+ """
38
+ # Same entry point as XPUModelRunner.__init__ (patches persist after exit).
39
+ with _torch_cuda_wrapper():
40
+ pass
41
+
42
+ # Fresh handler table build, as on first torch.compile / AOT in the worker.
43
+ # Registers torch.cuda.current_stream and torch.xpu.current_stream separately;
44
+ # if they are the same object (pre-fix alias), raises Handler already registered.
45
+ TorchInGraphFunctionVariable._get_handlers.cache_clear()
46
+ TorchInGraphFunctionVariable._get_handlers()
vllm/v1/worker/xpu_model_runner.py
+10
-7
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from contextlib import contextmanager
4
+from functools import partial
5
6
import torch
7
@@ -40,14 +41,16 @@ class XPUModelRunnerV2(GPUModelRunnerV2):
41
42
@contextmanager
43
def _torch_cuda_wrapper():
43
- # replace cuda APIs with xpu APIs, this should work by default
44
+ # Replace cuda APIs with xpu APIs. Each callable gets its own functools.partial
45
+ # so it is not the same object as torch.xpu.* (Torch Dynamo _get_handlers()
46
+ # asserts on duplicate registration when cuda aliases xpu directly).
47
torch.cuda.Stream = torch.xpu.Stream
45
- torch.cuda.default_stream = torch.xpu.current_stream
46
- torch.cuda.current_stream = torch.xpu.current_stream
47
- torch.cuda.stream = torch.xpu.stream
48
- torch.cuda.set_stream = torch.xpu.set_stream
48
+ torch.cuda.default_stream = partial(torch.xpu.current_stream)
49
+ torch.cuda.current_stream = partial(torch.xpu.current_stream)
50
+ torch.cuda.stream = partial(torch.xpu.stream)
51
+ torch.cuda.set_stream = partial(torch.xpu.set_stream)
52
if supports_xpu_graph():
50
- torch.cuda.graph = torch.xpu.graph
53
+ torch.cuda.graph = partial(torch.xpu.graph)
54
torch.cuda.CUDAGraph = torch.xpu.XPUGraph
52
- torch.cuda.graph_pool_handle = torch.xpu.graph_pool_handle
55
+ torch.cuda.graph_pool_handle = partial(torch.xpu.graph_pool_handle)
56
yield