Skip to content

Commit 8df2dc3

Browse files
authored
[TPU] Update pallas.py to support trillium (#8871)
1 parent 6d792d2 commit 8df2dc3

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

vllm/attention/backends/pallas.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ def __init__(
130130
assert tpu_type is not None
131131
tpu_type = tpu_type.lower()
132132

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

0 commit comments

Comments
 (0)