Skip to content

Commit f8f9ff5

Browse files
authored
[Bugfix][TPU] Fix megacore setting for v5e-litepod (#6397)
1 parent 6bc9710 commit f8f9ff5

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
@@ -116,7 +116,7 @@ def __init__(
116116

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

0 commit comments

Comments
 (0)