[Bugfix] Fix typo in Pallas backend (vllm-project#5558)

xjpang · Jul 24, 2024 · 0e20159 · 0e20159
1 parent 5d2a1ca
commit 0e20159
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm/attention/backends/pallas.py b/vllm/attention/backends/pallas.py
@@ -110,7 +110,7 @@ def __init__(
             raise NotImplementedError("TPU version must be 4 or higher.")
 
         self.megacore_mode = None
-        tpu_type = torch_xla.tpu.get_tp_groupu_env()["TYPE"].lower()
+        tpu_type = torch_xla.tpu.get_tpu_env()["TYPE"].lower()
         if not tpu_type.endswith("lite"):
             if self.num_kv_heads % 2 == 0:
                 self.megacore_mode = "kv_head"