From eeee1c3b1ae30a9714dffe7a58bdbed10b1e2e38 Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Wed, 21 Aug 2024 21:31:49 -0700
Subject: [PATCH] [TPU] Avoid initializing TPU runtime in is_tpu (#7763)

---
 vllm/platforms/__init__.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py
index 958f6c51..aedf3c3a 100644
--- a/vllm/platforms/__init__.py
+++ b/vllm/platforms/__init__.py
@@ -8,8 +8,10 @@ current_platform: Platform
 
 is_tpu = False
 try:
-    import torch_xla.core.xla_model as xm
-    xm.xla_device(devkind="TPU")
+    # While it's technically possible to install libtpu on a non-TPU machine,
+    # this is a very uncommon scenario. Therefore, we assume that libtpu is
+    # installed if and only if the machine has TPUs.
+    import libtpu  # noqa: F401
     is_tpu = True
 except Exception:
     pass