| 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 | |
| 8 | from vllm.config import VllmConfig |
| 9 | from vllm.utils.torch_utils import supports_xpu_graph |
| 10 | from vllm.v1.worker.gpu.model_runner import ( |
| 11 | GPUModelRunner as GPUModelRunnerV2, |
| 12 | ) |
| 13 | from vllm.v1.worker.gpu_model_runner import GPUModelRunner |
| 14 | from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner |
| 15 | |
| 16 | |
| 17 | class XPUModelRunner(GPUModelRunner): |
| 18 | """A model runner for XPU devices.""" |
| 19 | |
| 20 | def __init__( |
| 21 | self, |
| 22 | vllm_config: VllmConfig, |
| 23 | device: torch.device, |
| 24 | ): |
| 25 | with _torch_cuda_wrapper(): |
| 26 | super().__init__(vllm_config, device) |
| 27 | # FIXME: To be verified. |
| 28 | self.cascade_attn_enabled = False |
| 29 | |
| 30 | |
| 31 | class XPUModelRunnerV2(GPUModelRunnerV2): |
| 32 | """A model runner for XPU devices.""" |
| 33 | |
| 34 | def __init__( |
| 35 | self, |
| 36 | vllm_config: VllmConfig, |
| 37 | device: torch.device, |
| 38 | ): |
| 39 | with _torch_cuda_wrapper(): |
| 40 | super().__init__(vllm_config, device) |
| 41 | |
| 42 | |
| 43 | class XPUMMEncoderModelRunner(MMEncoderModelRunner): |
| 44 | """An encoder-only model runner for XPU devices.""" |
| 45 | |
| 46 | def __init__( |
| 47 | self, |
| 48 | vllm_config: VllmConfig, |
| 49 | device: torch.device, |
| 50 | ): |
| 51 | with _torch_cuda_wrapper(): |
| 52 | super().__init__(vllm_config, device) |
| 53 | |
| 54 | |
| 55 | @contextmanager |
| 56 | def _torch_cuda_wrapper(): |
| 57 | # Replace cuda APIs with xpu APIs. Each callable gets its own functools.partial |
| 58 | # so it is not the same object as torch.xpu.* (Torch Dynamo _get_handlers() |
| 59 | # asserts on duplicate registration when cuda aliases xpu directly). |
| 60 | torch.cuda.Stream = torch.xpu.Stream |
| 61 | torch.cuda.default_stream = partial(torch.xpu.current_stream) |
| 62 | torch.cuda.current_stream = partial(torch.xpu.current_stream) |
| 63 | torch.cuda.stream = partial(torch.xpu.stream) |
| 64 | torch.cuda.set_stream = partial(torch.xpu.set_stream) |
| 65 | |
| 66 | # torch.xpu.Event does not accept the ``blocking`` kwarg that |
| 67 | # torch.cuda.Event supports, so drop it here. |
| 68 | def _xpu_event(*args, blocking=None, **kwargs): |
| 69 | return torch.xpu.Event(*args, **kwargs) |
| 70 | |
| 71 | torch.cuda.Event = _xpu_event |
| 72 | if supports_xpu_graph(): |
| 73 | torch.cuda.graph = partial(torch.xpu.graph) |
| 74 | torch.cuda.CUDAGraph = torch.xpu.XPUGraph |
| 75 | torch.cuda.graph_pool_handle = partial(torch.xpu.graph_pool_handle) |
| 76 | torch.cuda.is_current_stream_capturing = torch.xpu.is_current_stream_capturing |
| 77 | yield |