dont remove _is_tpu() legacy code

HabanaAI · Sep 26, 2024 · 7453373 · 7453373
1 parent 5bb2b20
commit 7453373
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git 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
@@ -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__)