rwkv5 pre-compiled kernel (for windows)

backend-python/rwkv_pip/model.py

@@ -12,6 +12,56 @@ torch.backends.cudnn.allow_tf32 = True
 torch.backends.cuda.matmul.allow_tf32 = True
 current_path = os.path.dirname(os.path.abspath(__file__))
 
+
+# https://zhuanlan.zhihu.com/p/612879065
+def LoadPreCompileLibrary(file):
+    import importlib
+    import os
+
+    import torch
+
+    # load the custom_op_library and register the custom ops
+    lib_dir = os.path.dirname(__file__)
+    if os.name == "nt":
+        # Register the main torchvision library location on the default DLL path
+        import ctypes
+        import sys
+
+        kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
+        with_load_library_flags = hasattr(kernel32, "AddDllDirectory")
+        prev_error_mode = kernel32.SetErrorMode(0x0001)
+
+        if with_load_library_flags:
+            kernel32.AddDllDirectory.restype = ctypes.c_void_p
+
+        if sys.version_info >= (3, 8):
+            os.add_dll_directory(lib_dir)
+        elif with_load_library_flags:
+            res = kernel32.AddDllDirectory(lib_dir)
+            if res is None:
+                err = ctypes.WinError(ctypes.get_last_error())
+                err.strerror += f' Error adding "{lib_dir}" to the DLL directories.'
+                raise ValueError(err)
+
+        kernel32.SetErrorMode(prev_error_mode)
+
+    loader_details = (
+        importlib.machinery.ExtensionFileLoader,
+        importlib.machinery.EXTENSION_SUFFIXES,
+    )
+
+    extfinder = importlib.machinery.FileFinder(lib_dir, loader_details)
+    ext_specs = extfinder.find_spec(file)
+    if ext_specs is None:
+        return False
+
+    try:
+        torch.ops.load_library(ext_specs.origin)
+    except OSError as exc:
+        return False
+    return True
+
+
 ########################################################################################################
 
 if os.environ.get("RWKV_JIT_ON") != "0":
@@ -29,8 +79,10 @@ else:
     MyStatic = __nop
 
 if os.environ.get("RWKV_CUDA_ON") == "1":
-    from torch.utils.cpp_extension import load
+    DISABLE_CUBLAS_GEMM = False
+    from torch.utils.cpp_extension import load  # L581
 
+    if LoadPreCompileLibrary("wkv_cuda") is False:
         try:
             load(
                 name=f"wkv_cuda",
@@ -40,6 +92,7 @@ if os.environ.get("RWKV_CUDA_ON") == "1":
                     f"{current_path}/cuda/gemm_fp16_cublas.cpp",
                 ],
                 verbose=True,
+                extra_ldflags=["cublas.lib"],
                 extra_cuda_cflags=[
                     "--use_fast_math",
                     "-O3",
@@ -527,8 +580,11 @@ class RWKV(MyModule):
             if self.version == 5.2:
                 assert (
                     os.environ["RWKV_CUDA_ON"] == "1"
-                )  # latest RWKV-5 requires os.environ["RWKV_CUDA_ON"] == '1' (will fix soon)
+                ), "Please Enable Custom CUDA Kernel. Latest RWKV-5 requires os.environ['RWKV_CUDA_ON'] == '1' (will fix soon)"
                 HEAD_SIZE = args.n_att // args.n_head
+                if LoadPreCompileLibrary("rwkv5") is True:
+                    rwkv5 = torch.ops.rwkv5
+                else:
                     rwkv5 = load(
                         name="rwkv5",
                         sources=[
@@ -540,7 +596,6 @@ class RWKV(MyModule):
                             "-res-usage",
                             "--use_fast_math",
                             "-O3",
-                        "-Xptxas -O3",
                             "--extra-device-vectorization",
                             f"-D_N_={HEAD_SIZE}",
                         ],