[TPU] Update pallas.py to support trillium (#8871)

2024-09-27 01:16:55 -07:00 · 2024-09-27 01:16:55 -07:00 · 8df2dc3c88
commit 8df2dc3c88
parent 6d792d2f31
1 changed files with 1 additions and 1 deletions
--- a/vllm/attention/backends/pallas.py
+++ b/vllm/attention/backends/pallas.py
@ -130,7 +130,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
        assert tpu_type is not None
        tpu_type = tpu_type.lower()

-        if "lite" not in tpu_type:
+        if (("lite" not in tpu_type) and ("v6" not in tpu_type)):
            if self.num_kv_heads % 2 == 0:
                self.megacore_mode = "kv_head"
            else: