diff --git a/python/setup.py b/python/setup.py index 736aeab78fb7..6901fbf8648c 100644 --- a/python/setup.py +++ b/python/setup.py @@ -55,6 +55,11 @@ def get_lib_path(): libs.append(name) break + # Add byoc shared libraries, if present + for name in lib_path: + if "3rdparty" in name: + libs.append(name) + # Add standalone_crt, if present for name in lib_path: candidate_path = os.path.join(os.path.dirname(name), "standalone_crt") diff --git a/python/tvm/_ffi/libinfo.py b/python/tvm/_ffi/libinfo.py index 4cf3d3c22a76..5fe501c8bd55 100644 --- a/python/tvm/_ffi/libinfo.py +++ b/python/tvm/_ffi/libinfo.py @@ -112,25 +112,39 @@ def find_lib_path(name=None, search_path=None, optional=False): else: lib_dll_path = [os.path.join(p, name) for p in dll_path] runtime_dll_path = [] + ext_lib_dll_path = [] else: if sys.platform.startswith("win32"): lib_dll_names = ["libtvm.dll", "tvm.dll"] runtime_dll_names = ["libtvm_runtime.dll", "tvm_runtime.dll"] + ext_lib_dll_names = [ + "3rdparty/cutlass_fpA_intB_gemm/cutlass_kernels/libfpA_intB_gemm.dll", + "3rdparty/libflash_attn/src/libflash_attn.dll", + ] elif sys.platform.startswith("darwin"): lib_dll_names = ["libtvm.dylib"] runtime_dll_names = ["libtvm_runtime.dylib"] + ext_lib_dll_names = [ + "3rdparty/cutlass_fpA_intB_gemm/cutlass_kernels/libfpA_intB_gemm.dylib", + "3rdparty/libflash_attn/src/libflash_attn.dylib", + ] else: lib_dll_names = ["libtvm.so"] runtime_dll_names = ["libtvm_runtime.so"] + ext_lib_dll_names = [ + "3rdparty/cutlass_fpA_intB_gemm/cutlass_kernels/libfpA_intB_gemm.so", + "3rdparty/libflash_attn/src/libflash_attn.so", + ] - name = lib_dll_names + runtime_dll_names + name = lib_dll_names + runtime_dll_names + ext_lib_dll_names lib_dll_path = [os.path.join(p, name) for name in lib_dll_names for p in dll_path] runtime_dll_path = [os.path.join(p, name) for name in runtime_dll_names for p in dll_path] - + ext_lib_dll_path = [os.path.join(p, name) for name in ext_lib_dll_names for p in dll_path] if not use_runtime: # try to find lib_dll_path lib_found = [p for p in lib_dll_path if os.path.exists(p) and os.path.isfile(p)] lib_found += [p for p in runtime_dll_path if os.path.exists(p) and os.path.isfile(p)] + lib_found += [p for p in ext_lib_dll_path if os.path.exists(p) and os.path.isfile(p)] else: # try to find runtime_dll_path use_runtime = True