265 lines
8.2 KiB
Rust
265 lines
8.2 KiB
Rust
//! Build script for Flash Attention CUDA kernels
|
|
//! Optimized for RTX 5090 (sm_89/sm_90) architecture
|
|
|
|
use std::env;
|
|
use std::path::{Path, PathBuf};
|
|
use std::process::Command;
|
|
|
|
fn main() {
|
|
println!("cargo:rerun-if-changed=cuda/");
|
|
println!("cargo:rerun-if-changed=metal/");
|
|
println!("cargo:rerun-if-changed=build.rs");
|
|
|
|
// Only compile CUDA if the cuda feature is enabled
|
|
let cuda_feature = env::var("CARGO_FEATURE_CUDA").is_ok();
|
|
let metal_feature = env::var("CARGO_FEATURE_METAL").is_ok();
|
|
|
|
if cuda_feature {
|
|
let cuda_available = detect_cuda();
|
|
if cuda_available {
|
|
println!("cargo:rustc-cfg=cuda_available");
|
|
compile_cuda_kernels();
|
|
link_cuda_libraries();
|
|
} else {
|
|
println!("cargo:warning=CUDA feature enabled but CUDA not found on system");
|
|
}
|
|
}
|
|
|
|
if metal_feature {
|
|
println!("cargo:rustc-cfg=metal_available");
|
|
compile_metal_shaders();
|
|
println!("cargo:warning=Metal GPU support enabled for macOS");
|
|
}
|
|
|
|
if !cuda_feature && !metal_feature {
|
|
println!(
|
|
"cargo:warning=No GPU backend enabled. Enable 'cuda' or 'metal' feature for GPU acceleration"
|
|
);
|
|
}
|
|
}
|
|
|
|
fn detect_cuda() -> bool {
|
|
// Try to find nvcc compiler
|
|
if let Ok(output) = Command::new("nvcc").arg("--version").output()
|
|
&& output.status.success()
|
|
{
|
|
let version_output = String::from_utf8_lossy(&output.stdout);
|
|
println!(
|
|
"cargo:warning=Found CUDA: {}",
|
|
version_output.lines().next().unwrap_or("unknown")
|
|
);
|
|
return true;
|
|
}
|
|
|
|
// Check for CUDA_PATH environment variable
|
|
if let Ok(cuda_path) = env::var("CUDA_PATH") {
|
|
let nvcc_path = Path::new(&cuda_path).join("bin").join("nvcc");
|
|
if nvcc_path.exists() {
|
|
println!("cargo:warning=Found CUDA at: {cuda_path}");
|
|
return true;
|
|
}
|
|
}
|
|
|
|
false
|
|
}
|
|
|
|
fn compile_cuda_kernels() {
|
|
let out_dir = env::var("OUT_DIR").unwrap();
|
|
let cuda_dir = PathBuf::from("cuda");
|
|
|
|
// RTX 5090 optimized compilation flags
|
|
let cuda_flags = vec![
|
|
"--gpu-architecture=compute_90", // RTX 5090 architecture (Ada Lovelace)
|
|
"--gpu-code=sm_90", // RTX 5090 compute capability
|
|
"--use_fast_math",
|
|
"--maxrregcount=255",
|
|
"-O3",
|
|
"--extra-device-vectorization",
|
|
"--fmad=true",
|
|
"--prec-div=false",
|
|
"--prec-sqrt=false",
|
|
"--ftz=true", // Flush denormal to zero for performance
|
|
"-DRTX_5090_OPTIMIZATION=1",
|
|
"-DUSE_TENSOR_CORE_4TH_GEN=1",
|
|
"-DMAX_SHARED_MEMORY=163840", // 160KB for RTX 5090
|
|
];
|
|
|
|
let kernels = vec![
|
|
"flash_attention_forward.cu",
|
|
"flash_attention_backward.cu",
|
|
"online_softmax.cu",
|
|
"utils.cu",
|
|
];
|
|
|
|
for kernel in kernels {
|
|
let input_path = cuda_dir.join(kernel);
|
|
let output_name = kernel.replace(".cu", ".ptx");
|
|
let output_path = Path::new(&out_dir).join(&output_name);
|
|
|
|
println!(
|
|
"cargo:warning=Compiling CUDA kernel: {} -> {}",
|
|
input_path.display(),
|
|
output_path.display()
|
|
);
|
|
|
|
let mut cmd = Command::new("nvcc");
|
|
cmd.arg("--ptx")
|
|
.arg(&input_path)
|
|
.arg("-o")
|
|
.arg(&output_path);
|
|
|
|
// Add optimization flags
|
|
for flag in &cuda_flags {
|
|
cmd.arg(flag);
|
|
}
|
|
|
|
// Add include directories
|
|
if let Ok(cuda_path) = env::var("CUDA_PATH") {
|
|
cmd.arg(format!("-I{cuda_path}/include"));
|
|
}
|
|
|
|
let output = cmd.output().expect("Failed to execute nvcc");
|
|
|
|
if output.status.success() {
|
|
println!("cargo:warning=Successfully compiled {kernel}");
|
|
} else {
|
|
let stderr = String::from_utf8_lossy(&output.stderr);
|
|
panic!("Failed to compile CUDA kernel {kernel}: {stderr}");
|
|
}
|
|
}
|
|
|
|
// Create a combined PTX file with all kernels
|
|
combine_ptx_files(&out_dir);
|
|
}
|
|
|
|
fn combine_ptx_files(out_dir: &str) {
|
|
use std::fs::{File, OpenOptions};
|
|
use std::io::{BufWriter, Read, Write};
|
|
|
|
let combined_path = Path::new(out_dir).join("flash_attention_kernels.ptx");
|
|
let mut combined_file = BufWriter::new(
|
|
OpenOptions::new()
|
|
.create(true)
|
|
.write(true)
|
|
.truncate(true)
|
|
.open(&combined_path)
|
|
.expect("Failed to create combined PTX file"),
|
|
);
|
|
|
|
let ptx_files = vec![
|
|
"flash_attention_forward.ptx",
|
|
"flash_attention_backward.ptx",
|
|
"online_softmax.ptx",
|
|
"utils.ptx",
|
|
];
|
|
|
|
// Write PTX header
|
|
writeln!(combined_file, ".version 7.8").unwrap();
|
|
writeln!(combined_file, ".target sm_90").unwrap();
|
|
writeln!(combined_file, ".address_size 64").unwrap();
|
|
writeln!(combined_file).unwrap();
|
|
|
|
for ptx_file in ptx_files {
|
|
let ptx_path = Path::new(out_dir).join(ptx_file);
|
|
if ptx_path.exists() {
|
|
let mut contents = String::new();
|
|
File::open(&ptx_path)
|
|
.expect("Failed to open PTX file")
|
|
.read_to_string(&mut contents)
|
|
.expect("Failed to read PTX file");
|
|
|
|
// Skip header lines and write content
|
|
let lines: Vec<&str> = contents.lines().collect();
|
|
let content_start = lines.iter().position(|&line| line.is_empty()).unwrap_or(0) + 1;
|
|
|
|
for line in &lines[content_start..] {
|
|
writeln!(combined_file, "{line}").unwrap();
|
|
}
|
|
|
|
writeln!(combined_file).unwrap();
|
|
}
|
|
}
|
|
|
|
combined_file.flush().unwrap();
|
|
drop(combined_file);
|
|
|
|
println!(
|
|
"cargo:warning=Created combined PTX file: {}",
|
|
combined_path.display()
|
|
);
|
|
|
|
// Tell Rust to include the PTX file
|
|
println!(
|
|
"cargo:rustc-env=FLASH_ATTENTION_PTX_PATH={}",
|
|
combined_path.display()
|
|
);
|
|
}
|
|
|
|
fn link_cuda_libraries() {
|
|
// Add CUDA library paths
|
|
if let Ok(cuda_path) = env::var("CUDA_PATH") {
|
|
println!("cargo:rustc-link-search=native={cuda_path}/lib64");
|
|
} else {
|
|
// Try common CUDA installation paths
|
|
let cuda_paths = vec![
|
|
"/usr/local/cuda/lib64",
|
|
"/opt/cuda/lib64",
|
|
"/usr/lib/x86_64-linux-gnu",
|
|
];
|
|
|
|
for path in cuda_paths {
|
|
if Path::new(path).exists() {
|
|
println!("cargo:rustc-link-search=native={path}");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Link CUDA libraries
|
|
println!("cargo:rustc-link-lib=cudart");
|
|
println!("cargo:rustc-link-lib=cublas");
|
|
println!("cargo:rustc-link-lib=curand");
|
|
println!("cargo:rustc-link-lib=cusparse");
|
|
println!("cargo:rustc-link-lib=cufft");
|
|
|
|
// RTX 5090 specific optimizations
|
|
println!("cargo:rustc-cfg=rtx_5090_optimized");
|
|
println!("cargo:rustc-cfg=tensor_core_4th_gen");
|
|
}
|
|
|
|
fn compile_metal_shaders() {
|
|
// Metal shaders can be compiled using metallib
|
|
// For now, we'll just mark them as available
|
|
// In a full implementation, you would compile .metal files to .metallib
|
|
|
|
let _out_dir = env::var("OUT_DIR").unwrap();
|
|
let metal_dir = PathBuf::from("metal");
|
|
|
|
// Check if Metal shader source exists
|
|
let flash_attention_metal = metal_dir.join("flash_attention.metal");
|
|
if flash_attention_metal.exists() {
|
|
println!(
|
|
"cargo:warning=Found Metal shader source: {}",
|
|
flash_attention_metal.display()
|
|
);
|
|
|
|
// In production, compile with:
|
|
// xcrun -sdk macosx metal -c flash_attention.metal -o flash_attention.air
|
|
// xcrun -sdk macosx metallib flash_attention.air -o flash_attention.metallib
|
|
|
|
// For now, just copy the source path
|
|
let metal_source_path = std::fs::canonicalize(&flash_attention_metal)
|
|
.expect("Failed to get absolute path to Metal shader");
|
|
println!(
|
|
"cargo:rustc-env=METAL_SHADER_PATH={}",
|
|
metal_source_path.display()
|
|
);
|
|
println!(
|
|
"cargo:warning=Metal shader source available at: {}",
|
|
metal_source_path.display()
|
|
);
|
|
} else {
|
|
println!("cargo:warning=Metal shader source not found, using stub implementation");
|
|
}
|
|
}
|