fix linux kernel (partial revert 68228a45)

2023-10-05 00:08:18 +08:00 · 2023-10-05 00:08:18 +08:00 · 0005816c1d
commit 0005816c1d
parent f70672e5a0
1 changed files with 2 additions and 1 deletions
--- a/backend-python/rwkv_pip/model.py
+++ b/backend-python/rwkv_pip/model.py
@ -92,7 +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_ldflags=["cublas.lib" if os.name == "nt" else ""],
                extra_cuda_cflags=[
                    "--use_fast_math",
                    "-O3",
@ -596,6 +596,7 @@ class RWKV(MyModule):
                            "-res-usage",
                            "--use_fast_math",
                            "-O3",
+                            "-Xptxas -O3" if os.name != "nt" else "",
                            "--extra-device-vectorization",
                            f"-D_N_={HEAD_SIZE}",
                        ],