[rust] Use fused linear

deepjavalibrary · Jun 24, 2024 · d072013 · d072013
1 parent 0b768e3
commit d072013
Show file tree

Hide file tree

Showing 16 changed files with 1,175 additions and 29 deletions.
diff --git a/extensions/tokenizers/rust/Cargo.toml b/extensions/tokenizers/rust/Cargo.toml
@@ -6,10 +6,11 @@ edition = "2021"
 
 [dependencies]
 jni = "0.21.1"
-candle-core = "0.5.1"
+candle = { version = "0.5.1", package = "candle-core" }
 candle-nn = "0.5.1"
 candle-transformers = "0.5.1"
 candle-flash-attn = { version = "0.5.1", optional = true }
+cudarc = { version = "0.11.6", default-features = false, features = [ "cublaslt", "f16" ]}
 tokenizers = { path = "../tokenizers/tokenizers", version = "*", features = ["http"] }
 half = "2.4.0"
 tracing = "0.1.40"
@@ -25,5 +26,5 @@ openssl = { version = "0.10", features = ["vendored"] }
 crate_type = ["cdylib"]
 
 [features]
-cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
+cuda = ["candle/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
 flash-attn = ["cuda", "candle-transformers/flash-attn", "dep:candle-flash-attn"]
diff --git a/extensions/tokenizers/rust/src/compute_cap.rs b/extensions/tokenizers/rust/src/compute_cap.rs
@@ -1,7 +1,7 @@
-use candle_core::cuda_backend::cudarc::driver::sys::CUdevice_attribute::{
+use candle::cuda_backend::cudarc::driver::sys::CUdevice_attribute::{
     CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
 };
-use candle_core::cuda_backend::cudarc::driver::CudaDevice;
+use candle::cuda_backend::cudarc::driver::CudaDevice;
 use std::sync::Once;
 
 static INIT: Once = Once::new();