From dcd09cbf8b3d7dfe72820dab1b7f615f6bd40afd Mon Sep 17 00:00:00 2001
From: Guoqing Bao <topon@outlook.com>
Date: Thu, 9 Jan 2025 03:43:13 +0000
Subject: [PATCH] Fix kernel rebuild bugs

---
 .gitignore       | 4 +++-
 kernels/build.rs | 5 +++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 8856ec9..21c3d8b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,6 @@
 __pycache__
 *.ptx
 launch.json
-libpagedattention.a
\ No newline at end of file
+libpagedattention.a
+*.gz
+kernels/src/lib.rs
\ No newline at end of file
diff --git a/kernels/build.rs b/kernels/build.rs
index 2f37096..193cef6 100644
--- a/kernels/build.rs
+++ b/kernels/build.rs
@@ -21,10 +21,10 @@ fn main() -> Result<()> {
     println!("cargo:rerun-if-changed=src/reshape_and_cache_kernel.cu");
     println!("cargo:rerun-if-changed=src/marlin_cuda_kernel.cu");
     println!("cargo:rerun-if-changed=src/gptq_cuda_kernel.cu");
-
+    let build_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap_or("".to_string()));
     let builder = bindgen_cuda::Builder::default().arg("--expt-relaxed-constexpr");
     println!("cargo:info={builder:?}");
-    builder.build_lib("libpagedattention.a");
+    builder.build_lib(build_dir.join("libpagedattention.a"));
 
     let bindings = builder.build_ptx().unwrap();
     bindings.write("src/lib.rs").unwrap();
@@ -36,6 +36,7 @@ fn main() -> Result<()> {
         "cargo:rustc-link-search=native={}",
         absolute_kernel_dir.display()
     );
+    println!("cargo:rustc-link-search={}", build_dir.display());
     println!("cargo:rustc-link-lib=pagedattention");
     println!("cargo:rustc-link-lib=dylib=cudart");