Use is_flashinfer_available to replace is_hip for flashinfer check (#1596)

Co-authored-by: Zhang Liangang <liangang.zhang@intel.com>
This commit is contained in:
Lianmin Zheng
2024-10-06 22:54:05 -07:00
committed by GitHub
parent 565b05f02f
commit 6a5b352aaf
9 changed files with 29 additions and 28 deletions

View File

@@ -50,11 +50,19 @@ show_time_cost = False
time_infos = {}
# torch flag AMD GPU
def is_hip() -> bool:
"""Return whether it is HIP on the AMD ROCm platform."""
return torch.version.hip is not None
def is_flashinfer_available():
"""
Check whether flashinfer is available.
As of Oct. 6, 2024, it is only available on NVIDIA GPUs.
"""
return torch.cuda.is_available() and not is_hip()
def is_ipv6(address):
try:
ipaddress.IPv6Address(address)