From a752adf31be21b463bc7a9e7ece8d416652d2b8d Mon Sep 17 00:00:00 2001
From: Ethan Almloff <76850177+NoodlesOfWrath@users.noreply.github.com>
Date: Tue, 30 Sep 2025 22:33:26 -0500
Subject: [PATCH] add cuda support

---
 Cargo.lock | 125 ++++++++++++++++++++++++++++++++++++++++++++++++++---
 Cargo.toml |   3 +-
 2 files changed, 122 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 160283d..6ff3034 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -17,6 +17,17 @@ version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
 
+[[package]]
+name = "bindgen_cuda"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1f8489af5b7d17a81bffe37e0f4d6e1e4de87c87329d05447f22c35d95a1227d"
+dependencies = [
+ "glob",
+ "num_cpus",
+ "rayon",
+]
+
 [[package]]
 name = "bitflags"
 version = "1.3.2"
@@ -101,7 +112,9 @@ version = "0.9.1"
 source = "git+https://github.com/huggingface/candle?rev=41a674c#41a674ce8d1e37a9f8c6a97ff758552eddfd389a"
 dependencies = [
  "byteorder",
- "float8",
+ "candle-kernels",
+ "cudarc 0.17.3",
+ "float8 0.4.2",
  "gemm 0.17.1",
  "half",
  "memmap2",
@@ -112,18 +125,49 @@ dependencies = [
  "rayon",
  "safetensors",
  "thiserror",
- "ug",
+ "ug 0.4.0",
+ "ug-cuda",
  "yoke",
  "zip",
 ]
 
+[[package]]
+name = "candle-core"
+version = "0.9.1"
+source = "git+https://github.com/huggingface/candle?rev=d205fb4#d205fb41ae41b138b2c178b9c5d34591f1ffd07e"
+dependencies = [
+ "byteorder",
+ "float8 0.4.2",
+ "gemm 0.17.1",
+ "half",
+ "memmap2",
+ "num-traits",
+ "num_cpus",
+ "rand",
+ "rand_distr",
+ "rayon",
+ "safetensors",
+ "thiserror",
+ "ug 0.5.0",
+ "yoke",
+ "zip",
+]
+
+[[package]]
+name = "candle-kernels"
+version = "0.9.1"
+source = "git+https://github.com/huggingface/candle?rev=41a674c#41a674ce8d1e37a9f8c6a97ff758552eddfd389a"
+dependencies = [
+ "bindgen_cuda",
+]
+
 [[package]]
 name = "candle-nn"
 version = "0.9.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c1980d53280c8f9e2c6cbe1785855d7ff8010208b46e21252b978badf13ad69d"
 dependencies = [
- "candle-core",
+ "candle-core 0.9.1 (git+https://github.com/huggingface/candle?rev=d205fb4)",
  "half",
  "num-traits",
  "rayon",
@@ -198,6 +242,27 @@ version = "0.2.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
 
+[[package]]
+name = "cudarc"
+version = "0.16.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "17200eb07e7d85a243aa1bf4569a7aa998385ba98d14833973a817a63cc86e92"
+dependencies = [
+ "half",
+ "libloading",
+]
+
+[[package]]
+name = "cudarc"
+version = "0.17.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72ba848ae5c6f3cb36e71eab5f268763e3fabcabe3f7bc683e16f7fa3d46281e"
+dependencies = [
+ "float8 0.3.0",
+ "half",
+ "libloading",
+]
+
 [[package]]
 name = "darling"
 version = "0.21.3"
@@ -335,12 +400,22 @@ version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1ced73b1dacfc750a6db6c0a0c3a3853c8b41997e2e2c563dc90804ae6867959"
 
+[[package]]
+name = "float8"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0f498aec3b227cd892ce18967f4033d9d397d28a80a7ab67e9f6b0176a79654e"
+dependencies = [
+ "half",
+]
+
 [[package]]
 name = "float8"
 version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4203231de188ebbdfb85c11f3c20ca2b063945710de04e7b59268731e728b462"
 dependencies = [
+ "cudarc 0.17.3",
  "half",
  "num-traits",
  "rand",
@@ -691,6 +766,12 @@ dependencies = [
  "wasi",
 ]
 
+[[package]]
+name = "glob"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280"
+
 [[package]]
 name = "half"
 version = "2.6.0"
@@ -878,7 +959,7 @@ version = "0.1.0"
 source = "git+https://github.com/ModuRL/ModuRL#0779192f44246dc6a190ac2c3701e9f61ad92f9e"
 dependencies = [
  "bon",
- "candle-core",
+ "candle-core 0.9.1 (git+https://github.com/huggingface/candle?rev=41a674c)",
  "candle-nn",
  "rand",
 ]
@@ -889,7 +970,7 @@ version = "0.1.0"
 dependencies = [
  "bon",
  "box2d-rs",
- "candle-core",
+ "candle-core 0.9.1 (git+https://github.com/huggingface/candle?rev=41a674c)",
  "log",
  "minifb",
  "modurl",
@@ -1553,6 +1634,40 @@ dependencies = [
  "yoke",
 ]
 
+[[package]]
+name = "ug"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "76b761acf8af3494640d826a8609e2265e19778fb43306c7f15379c78c9b05b0"
+dependencies = [
+ "gemm 0.18.2",
+ "half",
+ "libloading",
+ "memmap2",
+ "num",
+ "num-traits",
+ "num_cpus",
+ "rayon",
+ "safetensors",
+ "serde",
+ "thiserror",
+ "tracing",
+ "yoke",
+]
+
+[[package]]
+name = "ug-cuda"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "14053653d0b7fa7b21015aa9a62edc8af2f60aa6f9c54e66386ecce55f22ed29"
+dependencies = [
+ "cudarc 0.16.6",
+ "half",
+ "serde",
+ "thiserror",
+ "ug 0.4.0",
+]
+
 [[package]]
 name = "unicode-ident"
 version = "1.0.19"
diff --git a/Cargo.toml b/Cargo.toml
index 267cd28..78dfdae 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,10 +18,11 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 
 [patch.crates-io]
-candle-core = { git = "https://github.com/huggingface/candle", rev = "41a674c" }
+candle-core = { git = "https://github.com/huggingface/candle", rev = "d205fb4" }
 
 [features]
 default = []
 tracing = ["dep:tracing"]
 logging = ["dep:log"]
 rendering = ["dep:minifb"]
+cuda = ["candle-core/cuda"]