better install/import dependencies

lmdeploy/__init__.py

@@ -10,6 +10,16 @@
 ]
 
 
+def find_win_cuda_path(cuda_ver, start, stop, step):
+    import os
+    for i in range(start, stop, step):
+        env_name = 'CUDA_PATH_V' + cuda_ver.replace('.x', '_' + str(i))
+        cuda_path = os.getenv(env_name)
+        if cuda_path is not None:
+            return cuda_path
+    return None
+
+
 def bootstrap():
     import os
     import sys
@@ -20,8 +30,14 @@ def bootstrap():
         has_turbomind = True
     if os.name == 'nt' and has_turbomind:
         if sys.version_info[:2] >= (3, 8):
-            CUDA_PATH = os.getenv('CUDA_PATH')
-            os.add_dll_directory(os.path.join(CUDA_PATH, 'bin'))
+            from .version import cuda_ver
+            cuda_path = None
+            if cuda_ver == '11.x':
+                cuda_path = find_win_cuda_path(cuda_ver, 8, 2, -1)
+            elif cuda_ver == '12.x':
+                cuda_path = find_win_cuda_path(cuda_ver, 1, 9, 1)
+            print(f'Find cuda path: {cuda_path}')
+            os.add_dll_directory(os.path.join(cuda_path, 'bin'))
 
 
 bootstrap()

setup.py

@@ -42,7 +42,7 @@ def parse_requirements(fname='requirements.txt', with_version=True):
     """
     require_fpath = fname
 
-    def get_nccl_pkg():
+    def get_cuda_pkgs():
         arg_name = '--cuda='
         arg_value = None
         for arg in sys.argv[1:]:
@@ -51,11 +51,23 @@ def get_nccl_pkg():
                 sys.argv.remove(arg)
                 break
 
+        cuda_ver = ''
+        cuda_pkgs = []
         if arg_value == '11':
-            return 'nvidia-nccl-cu11'
+            cuda_pkgs = [
+                'nvidia-nccl-cu11', 'nvidia-cuda-runtime-cu11',
+                'nvidia-cublas-cu11'
+            ]
+            cuda_ver = '11.x'
         elif arg_value == '12':
-            return 'nvidia-nccl-cu12'
-        return None
+            cuda_pkgs = [
+                'nvidia-nccl-cu12', 'nvidia-cuda-runtime-cu12',
+                'nvidia-cublas-cu12'
+            ]
+            cuda_ver = '12.x'
+        with open('lmdeploy/version.py', 'a') as f:
+            f.write(cuda_ver + '\n')
+        return cuda_pkgs
 
     def parse_line(line):
         """Parse information from a line in a requirements text file."""
@@ -113,9 +125,8 @@ def gen_packages_items():
                 yield item
 
     packages = list(gen_packages_items())
-    nccl_pkg = get_nccl_pkg()
-    if nccl_pkg is not None:
-        packages += [nccl_pkg]
+    cuda_pkgs = get_cuda_pkgs()
+    packages += cuda_pkgs
     return packages
 
 

src/turbomind/python/CMakeLists.txt

@@ -17,6 +17,13 @@ target_link_libraries(${PROJECT_NAME} PRIVATE TransformerTritonBackend
     LlamaTritonBackend)
 target_compile_features(${PROJECT_NAME} PRIVATE cxx_std_14)
 
+set(_INSTALL_CUDA_RPATH
+    "\$ORIGIN"
+    "\$ORIGIN/../../nvidia/nccl/lib/"
+    "\$ORIGIN/../../nvidia/cuda_runtime/lib/"
+    "\$ORIGIN/../../nvidia/cublas/lib/"
+)
 set_target_properties(${PROJECT_NAME} PROPERTIES
-        BUILD_RPATH "\$ORIGIN"
-	INSTALL_RPATH "\$ORIGIN;\$ORIGIN/../../nvidia/nccl/lib/")
+    BUILD_RPATH "\$ORIGIN"
+	INSTALL_RPATH ${_INSTALL_CUDA_RPATH}
+)