squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Woosuk Kwon 2024-08-30 09:01:26 -07:00 committed by GitHub
parent 058344f89a
commit 2684efc467
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/attention/backends/pallas.py
View File

@ -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: