From 2684efc4678eb46d1dc7fe4311365a99215e2dc6 Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Fri, 30 Aug 2024 09:01:26 -0700
Subject: [PATCH] [TPU][Bugfix] Fix tpu type api (#8035)

---
 vllm/attention/backends/pallas.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/vllm/attention/backends/pallas.py b/vllm/attention/backends/pallas.py
index c324d62d..83fdef16 100644
--- a/vllm/attention/backends/pallas.py
+++ b/vllm/attention/backends/pallas.py
@@ -124,7 +124,10 @@ class PallasAttentionBackendImpl(AttentionImpl):
 
         self.megacore_mode = None
         tpu_env = torch_xla.tpu.get_tpu_env()
-        tpu_type = tpu_env.get("TYPE") or tpu_env.get("ACCELERATOR_TYPE")
+        tpu_type = (tpu_env.get("ACCELERATOR_TYPE", None)
+                    or tpu_env.get("TYPE", None)
+                    or tpu_env.get("TPU_ACCELERATOR_TYPE", None))
+        assert tpu_type is not None
         tpu_type = tpu_type.lower()
 
         if "lite" not in tpu_type: