Skip to content

Commit

Permalink
[TPU] Update pallas.py to support trillium (vllm-project#8871)
Browse files Browse the repository at this point in the history
  • Loading branch information
bvrockwell authored Sep 27, 2024
1 parent 6d792d2 commit 8df2dc3
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion vllm/attention/backends/pallas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit 8df2dc3

Please sign in to comment.