Skip to content

Commit

Permalink
[TPU][Bugfix] Fix tpu type api (vllm-project#8035)
Browse files Browse the repository at this point in the history
  • Loading branch information
WoosukKwon authored Aug 30, 2024
1 parent 058344f commit 2684efc
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion vllm/attention/backends/pallas.py
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,10 @@ def __init__(

self.megacore_mode = None
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()

if "lite" not in tpu_type:
Expand Down

0 comments on commit 2684efc

Please sign in to comment.