// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. #[cfg(not(feature = "cuda_kv"))] fn main() {} #[cfg(feature = "cuda_kv")] fn main() { use std::{path::PathBuf, process::Command}; println!("cargo:rerun-if-changed=src/kernels/block_copy.cu"); // first do a which nvcc, if it is in the path // if so, we don't need to set the cuda_lib let nvcc = Command::new("which").arg("nvcc").output().unwrap(); let cuda_lib = if nvcc.status.success() { println!("cargo:info=nvcc found in path"); // Extract the path from nvcc location by removing "bin/nvcc" let nvcc_path = String::from_utf8_lossy(&nvcc.stdout).trim().to_string(); let path = PathBuf::from(nvcc_path); if let Some(parent) = path.parent() { // Remove "nvcc" if let Some(cuda_root) = parent.parent() { // Remove "bin" cuda_root.to_string_lossy().to_string() } else { // Fallback to CUDA_ROOT or default if path extraction fails get_cuda_root_or_default() } } else { // Fallback to CUDA_ROOT or default if path extraction fails get_cuda_root_or_default() } } else { println!("cargo:warning=nvcc not found in path"); get_cuda_root_or_default() }; println!("cargo:info=Using CUDA installation at: {}", cuda_lib); let cuda_lib_path = PathBuf::from(&cuda_lib).join("lib64"); println!("cargo:info=Using CUDA libs: {}", cuda_lib_path.display()); println!("cargo:rustc-link-search=native={}", cuda_lib_path.display()); // Link against multiple CUDA libraries println!("cargo:rustc-link-lib=dylib=cudart"); println!("cargo:rustc-link-lib=dylib=cuda"); println!("cargo:rustc-link-lib=dylib=cudadevrt"); // Make sure the CUDA libraries are found before other system libraries println!( "cargo:rustc-link-arg=-Wl,-rpath,{}", cuda_lib_path.display() ); // Create kernels directory for output if it doesn't exist std::fs::create_dir_all("src/kernels").unwrap_or_else(|_| { println!("Kernels directory already exists"); }); // Compile CUDA code let output = Command::new("nvcc") .arg("src/kernels/block_copy.cu") .arg("-O3") .arg("--compiler-options") .arg("-fPIC") .arg("-o") .arg("src/kernels/libblock_copy.o") .arg("-c") .output() .expect("Failed to compile CUDA code"); if !output.status.success() { panic!( "Failed to compile CUDA kernel: {}", String::from_utf8_lossy(&output.stderr) ); } // Create static library #[cfg(target_os = "windows")] { Command::new("lib") .arg("/OUT:src/kernels/block_copy.lib") .arg("src/kernels/libblock_copy.o") .output() .expect("Failed to create static library"); println!("cargo:rustc-link-search=native=src/kernels"); println!("cargo:rustc-link-lib=static=block_copy"); } #[cfg(not(target_os = "windows"))] { Command::new("ar") .arg("rcs") .arg("src/kernels/libblock_copy.a") .arg("src/kernels/libblock_copy.o") .output() .expect("Failed to create static library"); println!("cargo:rustc-link-search=native=src/kernels"); println!("cargo:rustc-link-lib=static=block_copy"); println!("cargo:rustc-link-lib=dylib=cudart"); println!("cargo:rustc-link-lib=dylib=cuda"); println!("cargo:rustc-link-lib=dylib=cudadevrt"); } } #[cfg(feature = "cuda_kv")] fn get_cuda_root_or_default() -> String { match std::env::var("CUDA_ROOT") { Ok(path) => path, Err(_) => { // Default locations based on OS if cfg!(target_os = "windows") { "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.8".to_string() } else { "/usr/local/cuda".to_string() } } } } #[cfg(feature = "trtllm")] fn main() { extern crate bindgen; use cmake::Config; use std::env; use std::path::PathBuf; let installed_headers = "/usr/local/include/nvidia/nvllm/nvllm_trt.h"; let local_headers = "../bindings/cpp/nvllm-trt/include/nvidia/nvllm/nvllm_trt.h"; let headers_path; if PathBuf::from(installed_headers).exists() { headers_path = installed_headers; println!("cargo:warning=nvllm found. Building with installed version..."); println!("cargo:rustc-link-search=native=/usr/local/lib"); println!("cargo:rustc-link-search=native=/opt/tensorrt_llm/lib"); println!("cargo:rustc-link-lib=dylib=nvllm_trt"); println!("cargo:rustc-link-lib=dylib=tensorrt_llm"); println!("cargo:rustc-link-lib=dylib=tensorrt_llm_nvrtc_wrapper"); println!("cargo:rustc-link-lib=dylib=nvinfer_plugin_tensorrt_llm"); println!("cargo:rustc-link-lib=dylib=decoder_attention"); println!("cargo:rerun-if-changed=/usr/local/lib"); } else if PathBuf::from(local_headers).exists() { headers_path = local_headers; println!("cargo:warning=nvllm not found. Building stub version..."); let dst = Config::new("../bindings/cpp/nvllm-trt") .define("USE_STUBS", "ON") .no_build_target(true) .build(); println!("cargo:warning=building stubs in {}", dst.display()); let dst = dst.canonicalize().unwrap(); println!("cargo:rustc-link-search=native={}/build", dst.display()); println!("cargo:rustc-link-lib=dylib=nvllm_trt"); println!("cargo:rustc-link-lib=dylib=tensorrt_llm"); println!("cargo:rerun-if-changed=../bindings/cpp/nvllm-trt"); } else { panic!("nvllm_trt.h not found"); } // generate bindings for the trtllm c api let bindings = bindgen::Builder::default() .header(headers_path) .generate() .expect("Unable to generate bindings"); // Write the bindings to a file let out_path = PathBuf::from(env::var("OUT_DIR").unwrap()); bindings .write_to_file(out_path.join("bindings.rs")) .expect("Could not write bindings!"); // // Build protobuf // tonic_build::configure() // .build_server(false) // .compile_protos(&["../../proto/trtllm.proto"], &["../../proto"]) // .expect("Failed to compile protos"); }