Add flash attention (#241)

* Add some flash-attn kernel, import the code for flash-attn v2 from Dao-AILab.

* More flash attn.

* Set up the flash attn parameters.

* Get things to compile locally.

* Move the flash attention files in a different directory.

* Build the static C library with nvcc.

* Add more flash attention.

* Update the build part.

* Better caching.

* Exclude flash attention from the default workspace.

* Put flash-attn behind a feature gate.

* Get the flash attn kernel to run.

* Move the flags to a more appropriate place.

* Enable flash attention in llama.

* Use flash attention in llama.

candle-examples/Cargo.toml

@@ -14,6 +14,7 @@ readme = "README.md"
 candle = { path = "../candle-core" }
 candle-nn = { path = "../candle-nn" }
 candle-transformers = { path = "../candle-transformers" }
+candle-flash-attn = { path = "../candle-flash-attn", optional = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 num-traits = { workspace = true }
@@ -37,4 +38,5 @@ anyhow = { workspace = true }
 [features]
 default = []
 cuda = ["candle/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
+flash-attn = ["cuda", "dep:candle-flash-attn"]
 mkl = ["dep:intel-mkl-src", "candle/mkl", "candle-nn/mkl", "candle-transformers/mkl"]

candle-examples/build.rs

@@ -6,11 +6,13 @@ use std::path::PathBuf;
 struct KernelDirectories {
     kernel_dir: &'static str,
     rust_target: &'static str,
+    include_dirs: &'static [&'static str],
 }
 
 const DIRS: [KernelDirectories; 1] = [KernelDirectories {
     kernel_dir: "examples/custom-ops/kernels/",
     rust_target: "examples/custom-ops/cuda_kernels.rs",
+    include_dirs: &[],
 }];
 
 impl KernelDirectories {
@@ -32,12 +34,15 @@ impl KernelDirectories {
             {
                 let mut command = std::process::Command::new("nvcc");
                 let out_dir = ptx_file.parent().context("no parent for ptx file")?;
+                let include_dirs: Vec<String> =
+                    self.include_dirs.iter().map(|c| format!("-I{c}")).collect();
                 command
                     .arg(format!("--gpu-architecture=sm_{compute_cap}"))
                     .arg("--ptx")
                     .args(["--default-stream", "per-thread"])
                     .args(["--output-directory", out_dir.to_str().unwrap()])
                     .arg(format!("-I/{}", self.kernel_dir))
+                    .args(include_dirs)
                     .arg(cu_file);
                 let output = command
                     .spawn()
@@ -221,6 +226,7 @@ fn compute_cap() -> Result<usize> {
     }
 
     println!("cargo:rerun-if-env-changed=CUDA_COMPUTE_CAP");
+
     if let Ok(compute_cap_str) = std::env::var("CUDA_COMPUTE_CAP") {
         compute_cap = compute_cap_str
             .parse::<usize>()

candle-examples/examples/llama/main.rs

@@ -116,6 +116,9 @@ struct Args {
 
     #[arg(long)]
     v2: bool,
+
+    #[arg(long)]
+    use_flash_attn: bool,
 }
 
 fn main() -> Result<()> {
@@ -124,7 +127,7 @@ fn main() -> Result<()> {
     let args = Args::parse();
 
     let device = candle_examples::device(args.cpu)?;
-    let config = Config::config_7b();
+    let config = Config::config_7b(args.use_flash_attn);
     let cache = model::Cache::new(!args.no_kv_cache, &config, &device)?;
     let dtype = if args.use_f32 { DType::F32 } else { DType::F16 };
     let (llama, tokenizer_filename) = match args.npy {

candle-examples/examples/llama/model.rs

@@ -13,10 +13,11 @@ pub struct Config {
     pub n_head: usize,
     pub n_embd: usize,
     pub n_key_value_head: usize,
+    pub use_flash_attn: bool,
 }
 
 impl Config {
-    pub fn config_7b() -> Self {
+    pub fn config_7b(use_flash_attn: bool) -> Self {
         Self {
             hidden_size: 4096,
             intermediate_size: 11008,
@@ -25,6 +26,7 @@ impl Config {
             n_head: 32,
             n_embd: 4096,
             n_key_value_head: 32,
+            use_flash_attn,
         }
     }
 }
@@ -140,6 +142,17 @@ struct CausalSelfAttention {
     n_key_value_head: usize,
     head_dim: usize,
     cache: Cache,
+    use_flash_attn: bool,
+}
+
+#[cfg(feature = "flash-attn")]
+fn flash_attn(q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
+    q.custom_op3(k, v, candle_flash_attn::FlashHdim32Sm80)
+}
+
+#[cfg(not(feature = "flash-attn"))]
+fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor) -> Result<Tensor> {
+    unimplemented!("compile with '--features flash-attn'")
 }
 
 impl CausalSelfAttention {
@@ -202,12 +215,17 @@ impl CausalSelfAttention {
 
         let k = self.repeat_kv(k)?;
         let v = self.repeat_kv(v)?;
-        let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?;
-        let mask = self.cache.mask(seq_len)?.broadcast_as(att.shape())?;
-        let att = masked_fill(&att, &mask, f32::NEG_INFINITY)?;
-        let att = att.softmax(D::Minus1)?;
-        // Convert to contiguous as matmul doesn't support strided vs for now.
-        let y = att.matmul(&v.contiguous()?)?;
+
+        let y = if self.use_flash_attn {
+            flash_attn(&q, &k, &v)?
+        } else {
+            let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?;
+            let mask = self.cache.mask(seq_len)?.broadcast_as(att.shape())?;
+            let att = masked_fill(&att, &mask, f32::NEG_INFINITY)?;
+            let att = att.softmax(D::Minus1)?;
+            // Convert to contiguous as matmul doesn't support strided vs for now.
+            att.matmul(&v.contiguous()?)?
+        };
         let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, n_embd])?;
         let y = y.to_dtype(x_dtype)?;
         let y = self.o_proj.forward(&y)?;
@@ -245,6 +263,7 @@ impl CausalSelfAttention {
             n_key_value_head: cfg.n_key_value_head,
             head_dim: cfg.hidden_size / cfg.n_head,
             cache: cache.clone(),
+            use_flash_attn: cfg.use_flash_attn,
         })
     }
 }