Upgrade candle2 (#543)

Author: Narsil

.github/workflows/build.yaml

@@ -58,6 +58,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
+        with:
+          submodules: true
 
       - name: Initialize Docker Buildx
         uses: docker/setup-buildx-action@v3

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "candle-extensions/candle-flash-attn-v1/cutlass"]
+	path = candle-extensions/candle-flash-attn-v1/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

Cargo.lock

@@ -7,6 +7,9 @@ members = [
     "backends/python",
     "backends/grpc-client",
     "candle-extensions/candle-cublaslt",
+    "candle-extensions/candle-flash-attn-v1",
+    "candle-extensions/candle-layer-norm",
+    "candle-extensions/candle-rotary",
     "core",
     "router",
 ]
@@ -49,13 +52,7 @@ candle = { git = "https://github.com/OlivierDehaene/candle", rev = "7e02ad856104
 candle-nn = { git = "https://github.com/OlivierDehaene/candle", rev = "7e02ad856104799b73a946ac1e153f0de77feaaf", package = "candle-nn" }
 candle-transformers = { git = "https://github.com/OlivierDehaene/candle", rev = "7e02ad856104799b73a946ac1e153f0de77feaaf", package = "candle-transformers" }
 candle-flash-attn = { git = "https://github.com/OlivierDehaene/candle", rev = "7e02ad856104799b73a946ac1e153f0de77feaaf", package = "candle-flash-attn" }
-
-
-[patch.crates-io]
-candle = { git = "https://github.com/OlivierDehaene/candle", rev = "7e02ad856104799b73a946ac1e153f0de77feaaf", package = "candle-core" }
-candle-nn = { git = "https://github.com/OlivierDehaene/candle", rev = "7e02ad856104799b73a946ac1e153f0de77feaaf", package = "candle-nn" }
-candle-transformers = { git = "https://github.com/OlivierDehaene/candle", rev = "7e02ad856104799b73a946ac1e153f0de77feaaf", package = "candle-transformers" }
-candle-flash-attn = { git = "https://github.com/OlivierDehaene/candle", rev = "7e02ad856104799b73a946ac1e153f0de77feaaf", package = "candle-flash-attn" }
+half = { version = "2.3.1", features = ["num-traits"] }
 
 [profile.release]
 debug = 0

Cargo.toml

@@ -9,15 +9,14 @@ homepage.workspace = true
 anyhow = { workspace = true }
 accelerate-src = { version = "0.3.2", optional = true }
 intel-mkl-src = { version = "0.8.1", optional = true  }
-candle = { version = "*", package = "candle-core", default-features = false }
-candle-nn = { version = "*" }
-candle-transformers = { version = "*" }
-candle-flash-attn = { version = "*", optional = true }
-candle-flash-attn-v1 = { git = "https://github.com/huggingface/candle-flash-attn-v1", rev = "3f1870b0d708579904c76e41745c659c3f9fa038", optional = true }
-# candle-cublaslt = { git = "https://github.com/huggingface/candle-cublaslt", rev = "cf789b7dd6d4abb19b03b9556442f94f0588b4a0", optional = true }
+candle = { workspace = true }
+candle-nn = { workspace = true }
+candle-transformers = { workspace = true }
+candle-flash-attn = { workspace = true, optional = true}
+candle-flash-attn-v1 = { path = "../../candle-extensions/candle-flash-attn-v1", optional = true }
 candle-cublaslt = { path = "../../candle-extensions/candle-cublaslt", optional = true }
-candle-layer-norm = { git = "https://github.com/huggingface/candle-layer-norm", rev = "94c2add7d94c2d63aebde77f7534614e04dbaea1", optional = true }
-candle-rotary = { git = "https://github.com/huggingface/candle-rotary", rev = "0a718a0856569a92f3112e64f10d07e4447822e8", optional = true }
+candle-layer-norm = { path = "../../candle-extensions/candle-layer-norm", optional = true }
+candle-rotary = { path = "../../candle-extensions/candle-rotary", optional = true }
 nohash-hasher = { workspace = true }
 text-embeddings-backend-core = { path = "../core" }
 tracing = { workspace = true }

backends/candle/Cargo.toml

@@ -6,10 +6,9 @@ edition = "2021"
 description = "CUBLASLt gemm for the candle ML framework."
 
 [dependencies]
-# candle = { version = "0.8", package = "candle-core", features = ["cuda"]}
 candle = { workspace=true, features = ["cuda"]}
 cudarc = { workspace = true, features = [ "cublaslt", "f16" ]}
-half = { version = "2.3.1", features = ["num-traits"] }
+half = { workspace = true}
 
 [features]
 static-linking = ["cudarc/static-linking"]

candle-extensions/candle-cublaslt/Cargo.toml

@@ -0,0 +1,3 @@
+.idea
+target
+Cargo.lock

candle-extensions/candle-flash-attn-v1/.gitignore

@@ -0,0 +1,3 @@
+[submodule "cutlass"]
+	path = cutlass
+	url = https://github.com/NVIDIA/cutlass.git

candle-extensions/candle-flash-attn-v1/.gitmodules

@@ -0,0 +1,23 @@
+[package]
+name = "candle-flash-attn-v1"
+version = "0.0.1"
+edition = "2021"
+
+description = "Flash attention V1 layer for the candle ML framework."
+keywords = ["blas", "tensor", "machine-learning"]
+categories = ["science"]
+license = "MIT OR Apache-2.0"
+readme = "README.md"
+
+[dependencies]
+candle = { workspace = true }
+half = { workspace = true }
+
+[build-dependencies]
+anyhow = { version = "1", features = ["backtrace"] }
+num_cpus = "1.15.0"
+rayon = "1.7.0"
+
+[dev-dependencies]
+anyhow = { version = "1", features = ["backtrace"] }
+candle-nn = { version = "0.3.0", features = ["cuda"] }