xpu.py 3.29 KB
Newer Older
1
from typing import TYPE_CHECKING, Optional
2

3
4
import torch

5
6
7
8
from vllm.logger import init_logger

from .interface import DeviceCapability, Platform, PlatformEnum, _Backend

9
10
11
12
13
if TYPE_CHECKING:
    from vllm.config import VllmConfig
else:
    VllmConfig = None

14
logger = init_logger(__name__)
15
16
17
18


class XPUPlatform(Platform):
    _enum = PlatformEnum.XPU
19
    device_name: str = "xpu"
20
    device_type: str = "xpu"
21
    dispatch_key: str = "XPU"
22

23
    @classmethod
24
25
26
    def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
                             dtype: torch.dtype, kv_cache_dtype: Optional[str],
                             block_size: int, use_v1: bool) -> str:
27
28
        if selected_backend != _Backend.IPEX:
            logger.info("Cannot use %s backend on XPU.", selected_backend)
29
30
        logger.info("Using IPEX attention backend.")
        return "vllm.attention.backends.ipex_attn.IpexAttnBackend"
31

32
33
    @staticmethod
    def get_device_capability(device_id: int = 0) -> DeviceCapability:
34
35
36
        major, minor, *_ = torch.xpu.get_device_capability(
            device_id)['version'].split('.')
        return DeviceCapability(major=int(major), minor=int(minor))
37
38
39
40

    @staticmethod
    def get_device_name(device_id: int = 0) -> str:
        return torch.xpu.get_device_name(device_id)
41
42
43
44
45

    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        device_props = torch.xpu.get_device_properties(device_id)
        return device_props.total_memory
46

47
48
49
50
    @classmethod
    def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
        return True

51
52
53
    @staticmethod
    def inference_mode():
        return torch.no_grad()
54
55
56

    @classmethod
    def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
57
58
59
60
        cache_config = vllm_config.cache_config
        if cache_config and cache_config.block_size is None:
            cache_config.block_size = 16

61
62
63
64
65
66
67
68
69
70
71
        # check and update model config
        model_config = vllm_config.model_config
        if model_config.dtype == torch.bfloat16:
            logger.warning(
                "bfloat16 is not fully supported on XPU, casting to float16.")
            model_config.dtype = torch.float16
        if not model_config.enforce_eager:
            logger.warning(
                "CUDA graph is not supported on XPU, fallback to the eager "
                "mode.")
            model_config.enforce_eager = True
72

73
74
75
76
        if vllm_config.speculative_config is not None:
            raise NotImplementedError(
                "XPU does not support speculative decoding")

77
78
79
80
81
82
83
84
85
        # check and update parallel config
        parallel_config = vllm_config.parallel_config
        if (parallel_config.distributed_executor_backend is not None
                and parallel_config.distributed_executor_backend != "ray"):
            logger.warning(
                "%s is not supported on XPU, fallback to ray distributed"
                " executor backend.",
                parallel_config.distributed_executor_backend)
            parallel_config.distributed_executor_backend = "ray"
86
87
        if parallel_config.worker_cls == "auto":
            parallel_config.worker_cls = "vllm.worker.xpu_worker.XPUWorker"
88
89
90
91
92

    @classmethod
    def is_pin_memory_available(cls):
        logger.warning("Pin memory is not supported on XPU.")
        return False