From 745337328e9aa8f604678f00180ccb5c44ecf26f Mon Sep 17 00:00:00 2001 From: Harish Subramony Date: Thu, 26 Sep 2024 17:06:39 +0000 Subject: [PATCH] dont remove _is_tpu() legacy code --- vllm/platforms/__init__.py | 5 ++++- vllm/utils.py | 1 - 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index d69e98abc3d89..d597e00faea1a 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -1,7 +1,7 @@ from typing import Optional import torch - +from vllm.utils import is_tpu from .interface import Platform, PlatformEnum, UnspecifiedPlatform current_platform: Optional[Platform] @@ -21,6 +21,9 @@ elif torch.version.hip is not None: from .rocm import RocmPlatform current_platform = RocmPlatform() +elif is_tpu(): + from .tpu import TpuPlatform + current_platform = TpuPlatform() elif is_hpu: from .hpu import HpuPlatform current_platform = HpuPlatform() diff --git a/vllm/utils.py b/vllm/utils.py index fa6e132dd3522..11b3478b6db7b 100644 --- a/vllm/utils.py +++ b/vllm/utils.py @@ -27,7 +27,6 @@ from typing_extensions import ParamSpec import vllm.envs as envs -from vllm import _custom_ops as ops from vllm.logger import enable_trace_function_call, init_logger logger = init_logger(__name__)