diff --git a/vllm_omni/platforms/xpu/utils.py b/vllm_omni/platforms/xpu/utils.py index 18437cae3..6eb314745 100644 --- a/vllm_omni/platforms/xpu/utils.py +++ b/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