vllm/vllm/triton_utils/importing.py
Yan Ma 056a68c7db
[XPU] avoid triton import for xpu (#9440)
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
2024-10-24 05:14:00 +00:00

17 lines
480 B
Python

from importlib.util import find_spec
from vllm.logger import init_logger
from vllm.platforms import current_platform
logger = init_logger(__name__)
HAS_TRITON = (
find_spec("triton") is not None
and not current_platform.is_xpu() # Not compatible
and not current_platform.is_neuron() # neuron has too old torch
)
if not HAS_TRITON:
logger.info("Triton not installed or not compatible; certain GPU-related"
" functions will not be available.")