From 040526310ed1b502647510648464d2673de8ad63 Mon Sep 17 00:00:00 2001
From: Matthew Douglas <38992547+matthewdouglas@users.noreply.github.com>
Date: Tue, 26 Mar 2024 17:42:25 -0400
Subject: [PATCH] Add CUDA 12.4 to docs/install helper (#1136)

* Add CUDA 12.4 download to utility script, docs

* (ci) Add CUDA 12.4.0 build to workflow

* Apply ruff format to install_cuda.py
---
 docs/source/installation.mdx | 2 +-
 install_cuda.py              | 9 +++++++--
 install_cuda.sh              | 7 +++++--
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/docs/source/installation.mdx b/docs/source/installation.mdx
index 49d8b4ebd..d0dd7ba76 100644
--- a/docs/source/installation.mdx
+++ b/docs/source/installation.mdx
@@ -84,7 +84,7 @@ Then locally install the CUDA version you need with this script from bitsandbyte
 ```bash
 wget https://raw.githubusercontent.com/TimDettmers/bitsandbytes/main/install_cuda.sh
 # Syntax cuda_install CUDA_VERSION INSTALL_PREFIX EXPORT_TO_BASH
-#   CUDA_VERSION in {110, 111, 112, 113, 114, 115, 116, 117, 118, 120, 121, 122, 123}
+#   CUDA_VERSION in {110, 111, 112, 113, 114, 115, 116, 117, 118, 120, 121, 122, 123, 124}
 #   EXPORT_TO_BASH in {0, 1} with 0=False and 1=True
 
 # For example, the following installs CUDA 11.7 to ~/local/cuda-11.7 and exports the path to your .bashrc
diff --git a/install_cuda.py b/install_cuda.py
index 9e426cbd7..a5d09356d 100644
--- a/install_cuda.py
+++ b/install_cuda.py
@@ -17,6 +17,7 @@
     "121": "https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run",
     "122": "https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run",
     "123": "https://developer.download.nvidia.com/compute/cuda/12.3.2/local_installers/cuda_12.3.2_545.23.08_linux.run",
+    "124": "https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run",
 }
 
 
@@ -76,7 +77,9 @@ def main():
     download_path = "/tmp"  # default download path
 
     if len(sys.argv) < 2:
-        print("Usage: python install_cuda.py <version/all> [user/system] [download_path]")
+        print(
+            "Usage: python install_cuda.py <version/all> [user/system] [download_path]"
+        )
         sys.exit(1)
 
     version = sys.argv[1]
@@ -97,7 +100,9 @@ def main():
     elif version in cuda_versions:
         install_cuda(version, base_path, download_path)
     else:
-        print(f"Invalid CUDA version: {version}. Available versions are: {', '.join(cuda_versions.keys())}")
+        print(
+            f"Invalid CUDA version: {version}. Available versions are: {', '.join(cuda_versions.keys())}"
+        )
         sys.exit(1)
 
 
diff --git a/install_cuda.sh b/install_cuda.sh
index 8ffbc8478..2e7fe8ed2 100644
--- a/install_cuda.sh
+++ b/install_cuda.sh
@@ -11,7 +11,7 @@ URL120=https://developer.download.nvidia.com/compute/cuda/12.0.1/local_installer
 URL121=https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run
 URL122=https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run
 URL123=https://developer.download.nvidia.com/compute/cuda/12.3.2/local_installers/cuda_12.3.2_545.23.08_linux.run
-
+URL124=https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run
 
 CUDA_VERSION=$1
 BASE_PATH=$2
@@ -57,8 +57,11 @@ if [[ -n "$CUDA_VERSION" ]]; then
   elif [[ "$CUDA_VERSION" -eq "123" ]]; then
     URL=$URL123
     FOLDER=cuda-12.3
+  elif [[ "$CUDA_VERSION" -eq "124" ]]; then
+    URL=$URL124
+    FOLDER=cuda-12.4
   else
-    echo "argument error: No cuda version passed as input. Choose among versions 92 to 123"
+    echo "argument error: No cuda version passed as input. Choose among versions 110 to 124"
   fi
 else
     echo "argument error: No cuda version passed as input. Choose among versions 92 to 123"