From 51e3d1f588bda7af0453f385a1994ff4649bcc5d Mon Sep 17 00:00:00 2001
From: nsaintot <neoternon@gmail.com>
Date: Tue, 10 Nov 2026 13:55:23 +0100
Subject: [PATCH] Leave out the cuDNN libraries MLX never loads

Drop libcudnn adv, ops, cnn and ext from the Linux packages and from the
Windows --install-cuda unpack, about 296 MB unpacked. On a 3090 Ti,
htdemucs_ft and bs_polarformer load only libcudnn, graph, heuristic and
the three engine libraries, and run at the same speed without the four.
stage.sh stops if another library links one of them.
---
 crates/stemd-server/src/cuda.rs | 30 +++++++++++++++++++++++++++++-
 docs/building.md                |  2 +-
 packaging/linux/stage.sh        |  8 ++++++++
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/crates/stemd-server/src/cuda.rs b/crates/stemd-server/src/cuda.rs
index 5e0152a..32936ca 100644
--- a/crates/stemd-server/src/cuda.rs
+++ b/crates/stemd-server/src/cuda.rs
@@ -175,7 +175,7 @@ fn unpack_libraries(archive: &Path, into: &Path) -> Result<usize> {
         let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
             continue;
         };
-        if !name.to_ascii_lowercase().ends_with(".dll") {
+        if !wanted(name) {
             continue;
         }
 
@@ -194,6 +194,21 @@ fn unpack_libraries(archive: &Path, into: &Path) -> Result<usize> {
     Ok(placed)
 }
 
+/// cuDNN libraries MLX never loads: it uses only the graph API, and no other
+/// cuDNN library links these.
+const UNUSED: &[&str] = &[
+    "cudnn_adv64_9.dll",
+    "cudnn_ops64_9.dll",
+    "cudnn_cnn64_9.dll",
+    "cudnn_ext64_9.dll",
+];
+
+/// Whether an archive entry is a DLL to install.
+fn wanted(name: &str) -> bool {
+    let name = name.to_ascii_lowercase();
+    name.ends_with(".dll") && !UNUSED.contains(&name.as_str())
+}
+
 /// Fail on a read-only install directory before spending a gigabyte finding out.
 ///
 /// Under `Program Files` this is the ordinary case rather than a strange one,
@@ -271,6 +286,19 @@ mod tests {
         }
     }
 
+    #[test]
+    fn the_unused_cudnn_libraries_are_skipped() {
+        assert!(wanted("cudnn64_9.dll"));
+        assert!(wanted("cudnn_graph64_9.dll"));
+        assert!(wanted("cublasLt64_13.dll"));
+        assert!(!wanted("cudnn_ops64_9.dll"));
+        assert!(!wanted("CUDNN_ADV64_9.DLL"));
+        assert!(!wanted("cudnn_cnn64_9.dll"));
+        assert!(!wanted("cudnn_ext64_9.dll"));
+        assert!(wanted("cudnn_engines_tensor_ir64_9.dll"));
+        assert!(!wanted("cudnn.h"));
+    }
+
     /// Distinct names, since they share one staging directory.
     #[test]
     fn no_two_components_stage_under_the_same_name() {
diff --git a/docs/building.md b/docs/building.md
index f206258..2603513 100644
--- a/docs/building.md
+++ b/docs/building.md
@@ -161,7 +161,7 @@ at run time.
 | `jit-headers.sh` | writes `jit-headers.txt`, the toolkit headers MLX's run-time compiler reads; rerun after an MLX or CUDA update |
 
 Every package carries CUDA, cuDNN, BLAS and avahi-compat in `/opt/stemd/lib`,
-about 2 GB unpacked. The driver (`libcuda.so.1`), the display stack and the C++
+about 1.7 GB unpacked. The driver (`libcuda.so.1`), the display stack and the C++
 runtime come from the machine. NVIDIA's libraries are copied unmodified, as
 their licences require; the executables find them through `DT_RPATH`
 (`$ORIGIN/../lib`), which also covers what cuDNN opens at run time.
diff --git a/packaging/linux/stage.sh b/packaging/linux/stage.sh
index a6ab755..b802af9 100755
--- a/packaging/linux/stage.sh
+++ b/packaging/linux/stage.sh
@@ -82,9 +82,16 @@ soname() { readelf -d "$1" | sed -n 's/.*(SONAME).*\[\(.*\)\]/\1/p'; }
 # cuDNN opens its engine libraries itself and NVRTC opens its builtins, so no
 # NEEDED entry names them. They go into the walk so their own dependencies come
 # too. Copied under their soname, which is the name they are opened by.
+# cuDNN's adv, ops, cnn and ext libraries are left out: MLX uses only the
+# graph API, which never loads them, and no other cuDNN library links them.
+unused_cudnn() {
+    case "$1" in libcudnn_adv.so*|libcudnn_ops.so*|libcudnn_cnn.so*|libcudnn_ext.so*) return 0 ;; esac
+    return 1
+}
 queue=("$exe")
 for f in "$cudnn"/lib/libcudnn*.so.9 "$cuda"/lib64/libnvrtc-builtins.so; do
     [[ -f "$f" ]] || continue
+    unused_cudnn "$(basename "$f")" && continue
     name="$(soname "$f")"
     [[ -f "$prefix/lib/$name" ]] && continue
     cp -L "$f" "$prefix/lib/$name"
@@ -102,6 +109,7 @@ while [[ ${#queue[@]} -gt 0 ]]; do
     resolved="$(ldd "$target" | awk '$2 == "=>" {print $1, ($3 == "not" ? "" : $3)}')"
     for name in $(readelf -d "$target" | sed -n 's/.*(NEEDED).*\[\(.*\)\]/\1/p'); do
         host_owned "$name" && continue
+        unused_cudnn "$name" && die "$(basename "$target") links $name, which is left out as unused"
         [[ "$seen" == *" $name "* ]] && continue
         seen="$seen$name "
         path="$(awk -v n="$name" '$1 == n {print $2}' <<<"$resolved")"