[Bugfix][TPU] Fix megacore setting for v5e-litepod (#6397)

2024-07-12 15:59:47 -07:00 · 2024-07-12 15:59:47 -07:00 · f8f9ff57ee
commit f8f9ff57ee
parent 6bc9710f6e
1 changed files with 1 additions and 1 deletions
--- a/vllm/attention/backends/pallas.py
+++ b/vllm/attention/backends/pallas.py
@ -116,7 +116,7 @@ class PallasAttentionBackendImpl(AttentionImpl):

        self.megacore_mode = None
        tpu_type = torch_xla.tpu.get_tpu_env()["TYPE"].lower()
-        if not tpu_type.endswith("lite"):
+        if "lite" not in tpu_type:
            if self.num_kv_heads % 2 == 0:
                self.megacore_mode = "kv_head"
            else: