Add support for candle-flash-attn and candle-flash-attn-v3 in CUDA

.github/workflows/ci.yml

@@ -43,8 +43,8 @@ jobs:
     secrets:
       HF_TOKEN: ${{ github.event_name == 'push' && secrets.HF_TOKEN || '' }}
 
-  build_main:
-    name: Build main
+  build:
+    name: Build
     needs: tests
     if: github.event_name == 'push' && github.ref == 'refs/heads/main'
     permissions:
@@ -82,13 +82,13 @@ jobs:
 
   push:
     name: Push
-    needs: [tests, build_main]
+    needs: [tests, build]
     if: >-
       always() &&
       needs.tests.result == 'success' &&
       github.event_name == 'push' &&
       (github.ref == 'refs/heads/main' || github.ref_type == 'tag') &&
-      (github.ref_type == 'tag' || needs.build_main.result == 'success')
+      (github.ref_type == 'tag' || needs.build.result == 'success')
     permissions:
       contents: read
       packages: write

Cargo.toml

@@ -7,12 +7,14 @@ authors = ["Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>"]
 default-run = "sys1"
 license = "Apache-2.0"
 repository = "https://github.com/alvarobartt/sys1"
-include = ["/src/**", "/Cargo.toml", "/Cargo.lock", "/README.md", "/LICENSE"]
+include = ["/src/**", "/build.rs", "/Cargo.toml", "/Cargo.lock", "/README.md", "/LICENSE"]
 
 [dependencies]
 anyhow = "1"
 axum = "0.8.9"
 candle-core = "0.11.0"
+candle-flash-attn = { version = "0.11.0", optional = true }
+candle-flash-attn-v3 = { version = "0.11.0", optional = true }
 candle-nn = "0.11.0"
 clap = { version = "4.6.7", features = ["derive"] }
 hf-hub = "1.0.0"
@@ -41,6 +43,8 @@ default = ["cpu"]
 cpu = []
 cuda = ["candle-core/cuda", "candle-nn/cuda"]
 cudnn = ["cuda", "candle-core/cudnn", "candle-nn/cudnn"]
+flash-attn-2 = ["cuda", "dep:candle-flash-attn"]
+flash-attn-3 = ["cuda", "dep:candle-flash-attn-v3"]
 metal = ["candle-core/metal", "candle-nn/metal"]
 
 [profile.release]

Dockerfile

@@ -27,6 +27,7 @@ WORKDIR /app
 FROM chef AS planner
 
 COPY Cargo.toml Cargo.lock ./
+COPY build.rs ./
 COPY rust-toolchain.toml ./
 COPY src ./src
 RUN cargo chef prepare --recipe-path recipe.json
@@ -55,6 +56,7 @@ RUN --mount=type=cache,id=sys1-sccache,target=/root/.cache/sccache,sharing=locke
     sccache --show-stats
 
 COPY Cargo.toml Cargo.lock ./
+COPY build.rs ./
 COPY rust-toolchain.toml ./
 COPY src ./src
 

Dockerfile.cuda

@@ -14,7 +14,7 @@ ARG TARGETARCH
 ENV PATH="/root/.cargo/bin:${PATH}"
 
 RUN apt-get update \
-    && apt-get install -y --no-install-recommends build-essential ca-certificates curl pkg-config \
+    && apt-get install -y --no-install-recommends build-essential ca-certificates curl git pkg-config \
     && rm -rf /var/lib/apt/lists/*
 
 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain "${RUST_VERSION}" --profile minimal
@@ -41,6 +41,7 @@ WORKDIR /app
 FROM chef AS planner
 
 COPY Cargo.toml Cargo.lock ./
+COPY build.rs ./
 COPY rust-toolchain.toml ./
 COPY src ./src
 
@@ -74,17 +75,22 @@ RUN --mount=type=cache,id=sys1-sccache,target=/root/.cache/sccache,sharing=locke
     fi; \
     compute_caps="$(printf '%s' "${CUDA_COMPUTE_CAPS}" | tr ';' ' ')"; \
     for compute_cap in ${compute_caps}; do \
+        case "${compute_cap}" in \
+            80|86|89|90) features=cuda,flash-attn-2 ;; \
+            *) features=cuda ;; \
+        esac; \
         CUDA_COMPUTE_CAP="${compute_cap}" \
         CARGO_TARGET_DIR="target/${compute_cap}" \
         cargo chef cook \
             --release \
             --no-default-features \
-            --features cuda \
+            --features "${features}" \
             --recipe-path recipe.json; \
     done; \
     sccache --show-stats
 
 COPY Cargo.toml Cargo.lock ./
+COPY build.rs ./
 COPY rust-toolchain.toml ./
 COPY src ./src
 
@@ -106,13 +112,17 @@ RUN --mount=type=cache,id=sys1-sccache,target=/root/.cache/sccache,sharing=locke
     set -- ${compute_caps}; \
     install -d /app/bin; \
     for compute_cap do \
+        case "${compute_cap}" in \
+            80|86|89|90) features=cuda,flash-attn-2 ;; \
+            *) features=cuda ;; \
+        esac; \
         CUDA_COMPUTE_CAP="${compute_cap}" \
         CARGO_TARGET_DIR="target/${compute_cap}" \
         cargo build \
             --release \
             --locked \
             --no-default-features \
-            --features cuda \
+            --features "${features}" \
             --bin sys1; \
         if [ "$#" -eq 1 ]; then \
             install "target/${compute_cap}/release/sys1" /app/bin/sys1; \

README.md

@@ -1,14 +1,12 @@
 <div align="center">
   <img
     src="https://github.com/user-attachments/assets/8b2774a7-6439-422e-bbd3-3fdd85117c5d"
-    alt="System One"
+    alt="sys1"
     width="1200"
   />
   <br/>
   <em>
-    System One compatible API for open decision models, e.g.
-    <a href="https://huggingface.co/convaiinnovations/laya">Laya</a>,
-    written in Rust.
+    System One compatible API for open decision models, written in Rust.
   </em>
 </div>
 
@@ -18,9 +16,8 @@
 - System One compatible API Spec
 - `candle` with [`tokenizers` release candidate](https://huggingface.co/blog/tokenizers-v1)!
 - Dynamic, token-based batching
-- Support for ModernBert with Laya custom decision heads
-- CPU, CUDA and Metal (MPS) supported
-- ~14ms per query on NVIDIA RTX Pro 6000
+- SDPA on CPU, Metal, and CUDA
+- Flash Attention on Ampere, Ada Lovelace, and Hopper
 
 ## Get started
 
@@ -30,6 +27,8 @@ Install it with support for CPU, Metal or CUDA.
 cargo install sys1 --features cpu
 # cargo install sys1 --no-default-features --features metal
 # cargo install sys1 --no-default-features --features cuda
+# cargo install sys1 --no-default-features --features cuda,flash-attn-2 # Ampere, Ada Lovelace, or Hopper
+# cargo install sys1 --no-default-features --features cuda,flash-attn-3 # Hopper
 ```
 
 Then run it with any of the supported models (more coming soon!).

build.rs

@@ -0,0 +1,45 @@
+use std::env;
+
+fn main() {
+    println!("cargo:rerun-if-env-changed=CUDA_COMPUTE_CAP");
+
+    let flash_attention_2 = env::var_os("CARGO_FEATURE_FLASH_ATTN_2").is_some();
+    let flash_attention_3 = env::var_os("CARGO_FEATURE_FLASH_ATTN_3").is_some();
+    if !flash_attention_2 && !flash_attention_3 {
+        return;
+    }
+
+    let Ok(value) = env::var("CUDA_COMPUTE_CAP") else {
+        println!(
+            "cargo:warning=CUDA_COMPUTE_CAP is not set; the Flash Attention architecture will be checked at startup"
+        );
+        return;
+    };
+    let capability = parse_compute_capability(&value);
+    match (flash_attention_2, capability) {
+        (true, 80..=99) | (false, 90) => {}
+        (true, _) => panic!(
+            "`flash-attn-2` supports compute capability 8.x or 9.x; got {:?}",
+            value
+        ),
+        (false, _) => panic!(
+            "`flash-attn-3` requires Hopper compute capability 9.0; got {:?}",
+            value
+        ),
+    }
+}
+
+fn parse_compute_capability(value: &str) -> u32 {
+    let normalized = value.trim().to_ascii_lowercase();
+    let normalized = normalized
+        .strip_prefix("sm_")
+        .unwrap_or(&normalized)
+        .trim_end_matches(['a', 'f'])
+        .replace('.', "");
+    normalized.parse::<u32>().unwrap_or_else(|_| {
+        panic!(
+            "invalid CUDA_COMPUTE_CAP {:?}; expected values such as 80, 89, or 90",
+            value
+        )
+    })
+}

src/device.rs

@@ -1,5 +1,19 @@
 use candle_core::Device;
 
+pub fn validate_available() -> anyhow::Result<()> {
+    #[cfg(feature = "cuda")]
+    anyhow::ensure!(
+        candle_core::utils::cuda_is_available(),
+        "CUDA backend selected, but no CUDA device is available"
+    );
+    #[cfg(feature = "metal")]
+    anyhow::ensure!(
+        candle_core::utils::metal_is_available(),
+        "Metal backend selected, but no Metal device is available"
+    );
+    Ok(())
+}
+
 #[cfg(all(feature = "cpu", not(any(feature = "cuda", feature = "metal"))))]
 pub fn load() -> anyhow::Result<Device> {
     Ok(Device::Cpu)

src/lib.rs

@@ -6,6 +6,14 @@ compile_error!("enable one backend feature: cpu, cuda, or metal");
     all(feature = "cuda", feature = "metal")
 ))]
 compile_error!("backend features are mutually exclusive: choose cpu, cuda, or metal");
+#[cfg(all(feature = "flash-attn-2", feature = "flash-attn-3"))]
+compile_error!(
+    "Flash Attention features are mutually exclusive: choose flash-attn-2 or flash-attn-3"
+);
+#[cfg(all(feature = "metal", not(target_os = "macos")))]
+compile_error!("the `metal` feature is only supported when targeting macOS");
+#[cfg(all(feature = "cuda", target_os = "macos"))]
+compile_error!("the `cuda` feature is not supported when targeting macOS");
 
 pub mod api;
 pub mod batching;
@@ -14,3 +22,7 @@ pub mod hub;
 pub mod models;
 pub mod schema;
 pub mod tokenizer;
+
+pub fn validate_backend() -> anyhow::Result<()> {
+    device::validate_available()
+}

src/main.rs

@@ -52,6 +52,8 @@ struct Args {
     max_model_len: Option<usize>,
     #[arg(long, value_enum, default_value_t = Precision::Auto)]
     dtype: Precision,
+    #[arg(long, value_enum, default_value_t = models::AttentionImplementation::Eager)]
+    attention: models::AttentionImplementation,
 }
 
 impl Args {
@@ -78,6 +80,7 @@ impl Args {
             self.max_model_len != Some(0),
             "--max-model-len must be positive"
         );
+        self.attention.validate(self.dtype.resolve())?;
         Ok(())
     }
 }
@@ -111,6 +114,7 @@ async fn main() -> anyhow::Result<()> {
         .init();
     let args = Args::parse();
     args.validate()?;
+    sys1::validate_backend()?;
     let dtype = args.dtype.resolve();
     info!(
         version = env!("CARGO_PKG_VERSION"),
@@ -152,8 +156,14 @@ async fn main() -> anyhow::Result<()> {
         path = %model_path.display(),
         "loading model"
     );
-    let model = models::load(&model_path, architecture, dtype, args.max_model_len)
-        .with_context(|| format!("failed to load model from {}", model_path.display()))?;
+    let model = models::load(
+        &model_path,
+        architecture,
+        dtype,
+        args.max_model_len,
+        args.attention,
+    )
+    .with_context(|| format!("failed to load model from {}", model_path.display()))?;
     info!(
         model = %served_model_name,
         elapsed_ms = started.elapsed().as_millis(),
@@ -244,6 +254,7 @@ mod tests {
         assert_eq!(args.max_request_bytes, 1_048_576);
         assert_eq!(args.max_model_len, None);
         assert_eq!(args.dtype, Precision::Auto);
+        assert_eq!(args.attention, models::AttentionImplementation::Eager);
     }
 
     #[test]
@@ -261,6 +272,12 @@ mod tests {
 
         let args = Args::try_parse_from(["sys1", "--dtype", "bf16"]).unwrap();
         assert_eq!(args.dtype, Precision::Bf16);
+
+        let args = Args::try_parse_from(["sys1", "--attention", "flash-attn-2"]).unwrap();
+        assert_eq!(
+            args.attention,
+            models::AttentionImplementation::FlashAttention2
+        );
     }
 
     #[test]
@@ -302,4 +319,11 @@ mod tests {
         let args = Args::try_parse_from(["sys1", "--max-queue-size", "0"]).unwrap();
         assert!(args.validate().is_err());
     }
+
+    #[test]
+    fn rejects_flash_attention_with_f32() {
+        let args = Args::try_parse_from(["sys1", "--attention", "flash-attn-3", "--dtype", "f32"])
+            .unwrap();
+        assert!(args.validate().is_err());
+    }
 }