xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-24 05:51:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add is_tpu

Browse Source
This commit is contained in:
Woosuk Kwon 2024-04-01 03:18:36 +00:00
parent d148c2ef00
commit 3b8f43024f
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/utils.py
View File

@ -1,6 +1,7 @@
import asyncio import asyncio
import enum import enum
import gc import gc
import importlib
import os import os
import socket import socket
import subprocess import subprocess
@ -126,6 +127,11 @@ def is_neuron() -> bool:
return transformers_neuronx is not None return transformers_neuronx is not None
@lru_cache(maxsize=None)
def is_tpu() -> bool:
return importlib.util.find_spec("libtpu") is not None
@lru_cache(maxsize=None) @lru_cache(maxsize=None)
def get_max_shared_memory_bytes(gpu: int = 0) -> int: def get_max_shared_memory_bytes(gpu: int = 0) -> int:
"""Returns the maximum shared memory per thread block in bytes.""" """Returns the maximum shared memory per thread block in bytes."""