Switches to cudnn-frontend pypi package (Lightning-AI#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

@@ -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"

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

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`_.