Add cuDNN FE patch for TE GKE build

.github/container/Dockerfile.jax

@@ -39,6 +39,7 @@ ARG EXTRA_BAZEL_TARGETS
 ARG EXTRA_BUILD_JAX_ARGS
 ARG GIT_USER_NAME
 ARG GIT_USER_EMAIL
+ARG NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1
 
 RUN --mount=type=ssh \
     --mount=type=secret,id=SSH_KNOWN_HOSTS,target=/root/.ssh/known_hosts \
@@ -47,7 +48,7 @@ RUN --mount=type=ssh \
     git-clone.sh ${URLREF_XLA} ${SRC_PATH_XLA}
 EOF
 
-ADD build-jax.sh build-te.sh local_cuda_arch /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,6 +71,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 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/container/cudnn_frontend.patch

@@ -0,0 +1,23 @@
+diff --git a/include/cudnn_frontend_shim.h b/include/cudnn_frontend_shim.h
+index ea838d3..c8f9de7 100644
+--- a/include/cudnn_frontend_shim.h
++++ b/include/cudnn_frontend_shim.h
+@@ -96,12 +96,12 @@ load_cudart_so() {
+         const char *error = reinterpret_cast<const char *>(dlerror());
+
+         if (handle && !error) {
+-            if (lib_handle) {
+-                // Already loaded one -> multiple found
+-                dlclose(handle);
+-                throw std::runtime_error("Multiple libcudart libraries found: " + std::string(libs[loaded_index]) +
+-                                         " and " + std::string(libs[i]));
+-            }
++            // if (lib_handle) {
++            //     // Already loaded one -> multiple found
++            //     dlclose(handle);
++            //     throw std::runtime_error("Multiple libcudart libraries found: " + std::string(libs[loaded_index]) +
++            //                              " and " + std::string(libs[i]));
++            // }
+             lib_handle   = handle;
+             loaded_index = static_cast<int>(i);
+         }

.github/workflows/_test_maxtext_gke_xpk.yaml

@@ -45,8 +45,10 @@ 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;
+          NCCL_LIB_DIR=\${NCCL_LIB_DIR} . \${NCCL_LIB_DIR}/nccl-env-profile.sh;
           nsys-jax --capture-range=cudaProfilerApi \
                    --capture-range-end=stop \
                    -o /opt/output/profile.zip \