Update cuDNN FE in TE

Author: aybchan

.github/container/Dockerfile.jax

@@ -47,7 +47,7 @@ RUN --mount=type=ssh \
     git-clone.sh ${URLREF_XLA} ${SRC_PATH_XLA}
 EOF
 
-ADD build-jax.sh build-te.sh local_cuda_arch cudnn_frontend.patch /usr/local/bin/
+ADD build-jax.sh build-te.sh local_cuda_arch /usr/local/bin/ 
 # Install bazelisk
 RUN ARCH="$(dpkg --print-architecture)" && \
     wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-${ARCH} && \
@@ -70,7 +70,7 @@ RUN mkdir -p /builder/extra-targets/{bin,python} && \
 RUN --mount=type=ssh \
     --mount=type=secret,id=SSH_KNOWN_HOSTS,target=/root/.ssh/known_hosts \
     git-clone.sh ${URLREF_TRANSFORMER_ENGINE} ${SRC_PATH_TRANSFORMER_ENGINE}
-RUN pushd ${SRC_PATH_TRANSFORMER_ENGINE}/3rdparty/cudnn-frontend && git apply /usr/local/bin/cudnn_frontend.patch && popd
+RUN pushd ${SRC_PATH_TRANSFORMER_ENGINE}/3rdparty/cudnn-frontend && git apply fetch && git checkout f06a0766b664b73860c48f7aea2b6b09499cea8d && popd
 # Populate ${SRC_PATH_TRANSFORMER_ENGINE}/dist with [a] .whl file(s); --no-install
 # because (a) this is the builder stage, and (b) pip-finalize.sh does the install
 RUN build-te.sh \

.github/workflows/_test_maxtext_gke_xpk.yaml

@@ -45,6 +45,8 @@ jobs:
       with:
         IMAGE: ${{ inputs.MAXTEXT_IMAGE }}
         NAME: ${{ matrix.test }}
+        ENVS: |
+          CUDNN_FRONTEND_CUDART_LIB_NAME=libcudart.so.13;
         COMMAND: |
           source \${NCCL_LIB_DIR}/nccl-env-profile.sh;
           nsys-jax --capture-range=cudaProfilerApi \