From 73ca5dc79179773c542d71b3f81d02ae0a220329 Mon Sep 17 00:00:00 2001
From: Vedaanta Agarwalla <142048820+vedaanta@users.noreply.github.com>
Date: Mon, 13 May 2024 14:45:01 -0700
Subject: [PATCH] Switches to cudnn-frontend pypi package (#400)

Co-authored-by: Vedaanta Agarwalla <vagarwalla@ipp2-1949.nvidia.com>
Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
---
 .azure/docker-build.yml                   |  8 ++++----
 dockers/ubuntu-cuda/Dockerfile            | 11 +++++------
 docs/source/fundamentals/installation.rst |  9 ++-------
 3 files changed, 11 insertions(+), 17 deletions(-)

diff --git a/.azure/docker-build.yml b/.azure/docker-build.yml
index 169dae0b25..afd186fefa 100644
--- a/.azure/docker-build.yml
+++ b/.azure/docker-build.yml
@@ -40,10 +40,10 @@ jobs:
       #maxParallel: "3"
       matrix:
         # CUDA 12.1
-        "cuda 12.1 | torch 2.3 | cudnn FE v1.3":
-          { CUDA_VERSION: "12.1.1", TORCH_VERSION: "2.3.0", TRITON_VERSION: "2.3.0", CUDNN_FRONTEND: "1.4.0" }
-        "cuda 12.1 | torch 2.4 /nightly | cudnn FE v1.3":
-          { CUDA_VERSION: "12.1.1", TORCH_VERSION: "main", TORCH_INSTALL: "source", CUDNN_FRONTEND: "1.4.0" }
+        "cuda 12.1 | torch 2.3 | cudnn FE v1.4":
+          { CUDA_VERSION: "12.1.1", TORCH_VERSION: "2.3.0", TRITON_VERSION: "2.3.0", CUDNN_FRONTEND_VERSION: "1.4.0" }
+        "cuda 12.1 | torch 2.4 /nightly | cudnn FE v1.4":
+          { CUDA_VERSION: "12.1.1", TORCH_VERSION: "main", TORCH_INSTALL: "source", CUDNN_FRONTEND_VERSION: "1.4.0" }
         #'cuda 12.1': # this version - '8.9.5.29-1+cuda12.1' for 'libcudnn8' was not found
     # how much time to give 'run always even if cancelled tasks' before stopping them
     cancelTimeoutInMinutes: "2"
diff --git a/dockers/ubuntu-cuda/Dockerfile b/dockers/ubuntu-cuda/Dockerfile
index d2baffb268..6da7b4418f 100644
--- a/dockers/ubuntu-cuda/Dockerfile
+++ b/dockers/ubuntu-cuda/Dockerfile
@@ -20,7 +20,7 @@ ARG IMAGE_TYPE="devel"
 FROM nvidia/cuda:${CUDA_VERSION}-${IMAGE_TYPE}-ubuntu${UBUNTU_VERSION}
 
 ARG CUDNN_VERSION="9.1.0.70"
-ARG CUDNN_FRONTEND_CHECKOUT="v1.4.0"
+ARG CUDNN_FRONTEND_VERSION="1.4.0"
 ARG PYTHON_VERSION="3.10"
 ARG TORCH_VERSION="2.2.1"
 ARG TRITON_VERSION="2.2.0"
@@ -74,7 +74,9 @@ RUN \
     rm -rf /root/.cache && \
     rm -rf /var/lib/apt/lists/*;
 
-# cuda TK is only needed for building from source, otherwise PyTorch wheels include it
+# cuDNN package is used by both torch and cuDNN frontend python package.
+# When using torch wheels, cuDNN pure-lib python package is fetched as a dependency.
+# But when building torch from source, cuDNN needs to be manually installed first.
 RUN if [ "${TORCH_INSTALL}" == "source" ]; then \
         echo "CUDA_VERSION=$CUDA_VERSION ; CUDNN_VERSION=$CUDNN_VERSION" && \
         CUDNN_BASE_VER=${CUDNN_VERSION%%.*} && \
@@ -159,10 +161,7 @@ RUN \
     fi
 
 RUN \
-    # building cudnn-frontend from source
-    git clone --branch ${CUDNN_FRONTEND_CHECKOUT} https://github.com/NVIDIA/cudnn-frontend.git && \
-    CMAKE_BUILD_PARALLEL_LEVEL=16 pip install cudnn-frontend/ -v && \
-    rm -rf cudnn-frontend
+    pip install nvidia-cudnn-frontend==${CUDNN_FRONTEND_VERSION}
 
 RUN \
     # Show what we have
diff --git a/docs/source/fundamentals/installation.rst b/docs/source/fundamentals/installation.rst
index afc531c1b8..3c531ccca5 100644
--- a/docs/source/fundamentals/installation.rst
+++ b/docs/source/fundamentals/installation.rst
@@ -33,15 +33,10 @@ Thunder can use NVIDIA's Apex to accelerate some PyTorch operations. To install
 Install cuDNN
 -------------
 
-Thunder can use NVIDIA's cuDNN Python frontend bindings to accelerate some PyTorch operations. cuDNN's Python frontend currently requires being built from source. See the Git repository for instructions. Below is a template for CUDA 12.x that requires setting the ``CUDAToolkit_ROOT`` environment variable::
+Thunder can use NVIDIA's cuDNN Python frontend bindings to accelerate some PyTorch operations. cuDNN backend is a pure-lib python package which needs to downloaded separately.
 
   pip install nvidia-cudnn-cu12
-  export CUDNN_PATH=/usr/local/lib/python3.10/dist-packages/nvidia/cudnn/
-  for file in $CUDNN_PATH/lib/*.so.[0-9]; do filename_without_version="${file%??}"; ln -s $file $filename_without_version; done
-
-  git clone -b v1.4.0 https://github.com/NVIDIA/cudnn-frontend.git
-  export CUDAToolkit_ROOT=/path/to/cuda
-  CMAKE_BUILD_PARALLEL_LEVEL=16 pip install cudnn_frontend/ -v
+  pip install nvidia-cudnn-frontend
 
 You're all set, now follow `Install Thunder`_.