[hardware] unify usage of is_tpu to current_platform.is_tpu() (#7102)

This commit is contained in:
youkaichao
2024-08-13 00:16:42 -07:00
committed by GitHub
parent 7025b11d94
commit 4d2dc5072b
8 changed files with 29 additions and 33 deletions

View File

@@ -10,8 +10,7 @@ import vllm.envs as envs
from vllm.attention.backends.abstract import AttentionBackend
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils import (STR_BACKEND_ENV_VAR, is_cpu, is_hip, is_openvino,
is_tpu, is_xpu)
from vllm.utils import STR_BACKEND_ENV_VAR, is_cpu, is_hip, is_openvino, is_xpu
logger = init_logger(__name__)
@@ -194,7 +193,7 @@ def which_attn_to_use(
logger.info("Cannot use %s backend on XPU.", selected_backend)
return _Backend.IPEX
if is_tpu():
if current_platform.is_tpu():
if selected_backend != _Backend.PALLAS:
logger.info("Cannot use %s backend on TPU.", selected_backend)
return _Backend.PALLAS