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

[Hardware][AMD]: Replace HIPCC version with more precise ROCm version (#11515)

Browse Source 
Signed-off-by: hjwei <hjwei_xd@163.com>
This commit is contained in:
hj-wei 2024-12-28 19:17:35 +08:00 committed by GitHub
parent b7dcc003dc
commit 59d6bb4c86
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 29 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
setup.py
View File

@ -1,3 +1,4 @@
import ctypes
import importlib.util import importlib.util
import logging import logging
import os import os
@ -13,7 +14,7 @@ from packaging.version import Version, parse
from setuptools import Extension, find_packages, setup from setuptools import Extension, find_packages, setup
from setuptools.command.build_ext import build_ext from setuptools.command.build_ext import build_ext
from setuptools_scm import get_version from setuptools_scm import get_version
from torch.utils.cpp_extension import CUDA_HOME from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
def load_module_from_path(module_name, path): def load_module_from_path(module_name, path):
@ -379,25 +380,31 @@ def _build_custom_ops() -> bool:
return _is_cuda() or _is_hip() or _is_cpu() return _is_cuda() or _is_hip() or _is_cpu()
def get_hipcc_rocm_version(): def get_rocm_version():
# Run the hipcc --version command # Get the Rocm version from the ROCM_HOME/bin/librocm-core.so
result = subprocess.run(['hipcc', '--version'], # see https://github.com/ROCm/rocm-core/blob/d11f5c20d500f729c393680a01fa902ebf92094b/rocm_version.cpp#L21
stdout=subprocess.PIPE, try:
stderr=subprocess.STDOUT, librocm_core_file = Path(ROCM_HOME) / "lib" / "librocm-core.so"
text=True) if not librocm_core_file.is_file():
# Check if the command was executed successfully
if result.returncode != 0:
print("Error running 'hipcc --version'")
return None return None
librocm_core = ctypes.CDLL(librocm_core_file)
VerErrors = ctypes.c_uint32
get_rocm_core_version = librocm_core.getROCmVersion
get_rocm_core_version.restype = VerErrors
get_rocm_core_version.argtypes = [
ctypes.POINTER(ctypes.c_uint32),
ctypes.POINTER(ctypes.c_uint32),
ctypes.POINTER(ctypes.c_uint32),
]
major = ctypes.c_uint32()
minor = ctypes.c_uint32()
patch = ctypes.c_uint32()
# Extract the version using a regular expression if (get_rocm_core_version(ctypes.byref(major), ctypes.byref(minor),
match = re.search(r'HIP version: (\S+)', result.stdout) ctypes.byref(patch)) == 0):
if match: return "%d.%d.%d" % (major.value, minor.value, patch.value)
# Return the version string return None
return match.group(1) except Exception:
else:
print("Could not find HIP version in the output")
return None return None
@ -479,11 +486,10 @@ def get_vllm_version() -> str:
if "sdist" not in sys.argv: if "sdist" not in sys.argv:
version += f"{sep}cu{cuda_version_str}" version += f"{sep}cu{cuda_version_str}"
elif _is_hip(): elif _is_hip():
# Get the HIP version # Get the Rocm Version
hipcc_version = get_hipcc_rocm_version() rocm_version = get_rocm_version() or torch.version.hip
if hipcc_version != MAIN_CUDA_VERSION: if rocm_version and rocm_version != MAIN_CUDA_VERSION:
rocm_version_str = hipcc_version.replace(".", "")[:3] version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}"
version += f"{sep}rocm{rocm_version_str}"
elif _is_neuron(): elif _is_neuron():
# Get the Neuron version # Get the Neuron version
neuron_version = str(get_neuronxcc_version()) neuron_version = str(get_neuronxcc_version())