Skip to content

GitLab

Unverified Commit 2684efc4 authored by Woosuk Kwon's avatar Woosuk Kwon Committed by GitHub
Browse files

[TPU][Bugfix] Fix tpu type api (#8035)

parent 058344f8
Show whitespace changes
Inline Side-by-side
Showing with 4 additions and 1 deletion
vllm/attention/backends/pallas.py
View file @ 2684efc4
...@@ -124,7 +124,10 @@ class PallasAttentionBackendImpl(AttentionImpl): ...@@ -124,7 +124,10 @@ class PallasAttentionBackendImpl(AttentionImpl):
self.megacore_mode = None self.megacore_mode = None
tpu_env = torch_xla.tpu.get_tpu_env() tpu_env = torch_xla.tpu.get_tpu_env()
tpu_type = tpu_env.get("TYPE") or tpu_env.get("ACCELERATOR_TYPE") tpu_type = (tpu_env.get("ACCELERATOR_TYPE", None)
or tpu_env.get("TYPE", None)
or tpu_env.get("TPU_ACCELERATOR_TYPE", None))
assert tpu_type is not None
tpu_type = tpu_type.lower() tpu_type = tpu_type.lower()
if "lite" not in tpu_type: if "lite" not in tpu_type:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment