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

Signed-off-by: Sumit Dubey <sumit.dubey2@ibm.com>
sumitd2 · Nov 14, 2024 · 00e1241 · 00e1241
1 parent b3dffde
commit 00e1241
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm/attention/backends/pallas.py b/vllm/attention/backends/pallas.py
@@ -130,7 +130,7 @@ def __init__(
         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: