Skip to content

Commit

Permalink
Fix kernel rebuild bugs
Browse files Browse the repository at this point in the history
  • Loading branch information
guoqingbao committed Jan 9, 2025
1 parent 605c785 commit dcd09cb
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 3 deletions.
4 changes: 3 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -3,4 +3,6 @@
__pycache__
*.ptx
launch.json
libpagedattention.a
libpagedattention.a
*.gz
kernels/src/lib.rs
5 changes: 3 additions & 2 deletions kernels/build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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");

Expand Down

0 comments on commit dcd09cb

Please sign in to comment.