[TPU][Bugfix] Fix tpu type api (#8035)

2684efc4 · Woosuk Kwon · GitHub · 058344f8 · 2684efc4
Unverified Commit 2684efc4 authored Aug 30, 2024 by Woosuk Kwon Committed by GitHub Aug 30, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

vllm/attention/backends/pallas.py vllm/attention/backends/pallas.py +4 -1

No files found.
--- 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: