Chendi.Xueandgcanlin
482f9b89a1
Signed-off-by: gcanlin <canlinguosdu@gmail.com> Signed-off-by: Chendi Xue <chendi.xue@intel.com> Co-authored-by: gcanlin <canlinguosdu@gmail.com>
78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
from vllm.platforms.xpu import XPUPlatform
|
|
|
|
from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
|
|
from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class XPUOmniPlatform(OmniPlatform, XPUPlatform):
|
|
"""XPU/Intel GPU implementation of OmniPlatform.
|
|
|
|
Inherits all XPU-specific implementations from vLLM's XPUPlatform,
|
|
and adds Omni-specific interfaces from OmniPlatform.
|
|
"""
|
|
|
|
_omni_enum = OmniPlatformEnum.XPU
|
|
|
|
@classmethod
|
|
def get_omni_ar_worker_cls(cls) -> str:
|
|
return "vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker"
|
|
|
|
@classmethod
|
|
def get_omni_generation_worker_cls(cls) -> str:
|
|
return "vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker"
|
|
|
|
@classmethod
|
|
def get_diffusion_attn_backend_cls(
|
|
cls,
|
|
selected_backend: str | None,
|
|
head_size: int,
|
|
) -> str:
|
|
if selected_backend is not None:
|
|
backend_upper = selected_backend.upper()
|
|
backend = DiffusionAttentionBackendEnum[backend_upper]
|
|
logger.info("Using diffusion attention backend '%s'", backend_upper)
|
|
return backend.get_path()
|
|
|
|
logger.info("Defaulting to diffusion attention backend SDPA")
|
|
return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
|
|
|
|
@classmethod
|
|
def supports_torch_inductor(cls) -> bool:
|
|
# TODO: Enable this when torch compile bugs are resolved
|
|
return False
|
|
|
|
@classmethod
|
|
def get_default_stage_config_path(cls) -> str:
|
|
return "vllm_omni/platforms/xpu/stage_configs"
|
|
|
|
@classmethod
|
|
def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
|
|
if local_rank is None:
|
|
return torch.device("xpu")
|
|
return torch.device("xpu", local_rank)
|
|
|
|
@classmethod
|
|
def get_device_count(cls) -> int:
|
|
return torch.xpu.device_count()
|
|
|
|
@classmethod
|
|
def get_device_version(cls) -> str | None:
|
|
# XPU does not have a version string like CUDA
|
|
return None
|
|
|
|
@classmethod
|
|
def synchronize(cls) -> None:
|
|
torch.xpu.synchronize()
|
|
|
|
@classmethod
|
|
def get_free_memory(cls, device: torch.device | None = None) -> int:
|
|
free, _ = torch.xpu.mem_get_info(device)
|
|
return free
|