diff --git a/vllm/v1/worker/xpu_model_runner.py b/vllm/v1/worker/xpu_model_runner.py index fb892211f..7becdd392 100644 --- a/vllm/v1/worker/xpu_model_runner.py +++ b/vllm/v1/worker/xpu_model_runner.py @@ -45,8 +45,12 @@ def _torch_cuda_wrapper(): self.synchronize = lambda: None try: - # replace cuda Event with xpu Event, this should work by default + # replace cuda APIs with xpu APIs, this should work by default torch.cuda.Event = torch.xpu.Event + 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: # if anything goes wrong, just patch it with a placeholder