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

[core][misc] update libcudart finding (#7620)

Browse Source 
Co-authored-by: cjackal <44624812+cjackal@users.noreply.github.com>
This commit is contained in:
youkaichao 2024-08-16 23:01:35 -07:00 committed by GitHub
parent 5bf45db7df
commit d95cc0a55c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
vllm/distributed/device_communicators/cuda_wrapper.py
View File

@ -49,8 +49,13 @@ def find_loaded_library(lib_name) -> Optional[str]:
if not found:
# the library is not loaded in the current process
return None
# if lib_name is libcudart, we need to match a line with:
# address /path/to/libcudart-hash.so.11.0
start = line.index("/")
path = line[start:].strip()
filename = path.split("/")[-1]
assert filename.rpartition(".so")[0].startswith(lib_name), \
f"Unexpected filename: {filename} for library {lib_name}"
return path
@ -98,9 +103,9 @@ class CudaRTLibrary:
def __init__(self, so_file: Optional[str] = None):
if so_file is None:
so_file = find_loaded_library("libcudart.so")
so_file = find_loaded_library("libcudart")
assert so_file is not None, \
"libcudart.so is not loaded in the current process"
"libcudart is not loaded in the current process"
if so_file not in CudaRTLibrary.path_to_library_cache:
lib = ctypes.CDLL(so_file)
CudaRTLibrary.path_to_library_cache[so_file] = lib