[XPU][Bugfix] Make Omni AR async output path device-agnostic (#5569)

Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

vllm_omni/platforms/xpu/utils.py

@@ -1,16 +1,27 @@
 from contextlib import contextmanager
+from functools import partial
 
 import torch
+from vllm.utils.torch_utils import supports_xpu_graph
 
 
 @contextmanager
 def torch_cuda_wrapper():
-    try:
-        # replace cuda APIs with xpu APIs, this should work by default
-        torch.cuda.Stream = torch.xpu.Stream
-        torch.cuda.default_stream = torch.xpu.current_stream
-        torch.cuda.current_stream = torch.xpu.current_stream
-        torch.cuda.stream = torch.xpu.stream
-        yield
-    finally:
-        pass
+    # Use partial() so callables differ from torch.xpu.* (Dynamo registration).
+    torch.cuda.Stream = torch.xpu.Stream
+    torch.cuda.default_stream = partial(torch.xpu.current_stream)
+    torch.cuda.current_stream = partial(torch.xpu.current_stream)
+    torch.cuda.stream = partial(torch.xpu.stream)
+    torch.cuda.set_stream = partial(torch.xpu.set_stream)
+    torch.cuda.set_device = partial(torch.xpu.set_device)
+
+    def _xpu_event(*args, blocking=None, **kwargs):
+        return torch.xpu.Event(*args, **kwargs)
+
+    torch.cuda.Event = _xpu_event
+    if supports_xpu_graph():
+        torch.cuda.graph = partial(torch.xpu.graph)
+        torch.cuda.CUDAGraph = torch.xpu.XPUGraph
+        torch.cuda.graph_pool_handle = partial(torch.xpu.graph_pool_handle)
+        torch.cuda.is_current_stream_capturing = torch.xpu.is_current_stream_capturing
+    yield