Merge branch 'main' into fern-migration

Author: zhenghax

.github/container/test-jax.sh

@@ -177,7 +177,7 @@ case "${BATTERY}" in
         # TODO: unclear if JOBS_PER_GPU>1 will cause deadlocks for any tests in the suite
         set_default VISIBLE_GPUS all
         FLAGS+=(
-            "--test_tag_filters=multiaccelerator"
+            "--test_tag_filters=multiaccelerator,multiaccelerator-only"
             "//tests:gpu_tests"
             "//tests/pallas:gpu_tests"
             "//tests/mosaic:gpu_tests"
@@ -187,7 +187,7 @@ case "${BATTERY}" in
     single-gpu)
         set_default VISIBLE_GPUS 1
         FLAGS+=(
-            "--test_tag_filters=-multiaccelerator"
+            "--test_tag_filters=-multiaccelerator,-multiaccelerator-only"
             "//tests:gpu_tests"
             "//tests/pallas:gpu_tests"
             "//tests/mosaic:gpu_tests"

.github/workflows/ngc-release-testing.yaml

@@ -104,6 +104,18 @@ jobs:
             tensor-parallel: 4
             attention: dot_product
             remat-policy: minimal_with_context
+          - test: maxtext-cudnn-flash-te-gke
+            nodes: 2
+            gpus: 8
+            model: llama2-7b
+            pipeline-parallel: 1
+            data-parallel: 1
+            fsdp: 2
+            tensor-parallel: 4
+            attention: cudnn_flash_te
+            remat-policy: minimal_with_context
+            envs: |-
+              CUDNN_FRONTEND_CUDART_LIB_NAME=libcudart.so.13
     steps:
     - uses: actions/checkout@v6
 
@@ -112,6 +124,8 @@ jobs:
       with:
         IMAGE: ${{ inputs.MAXTEXT_IMAGE }}
         NAME: ${{ matrix.test }}
+        ENVS: |
+          ${{ matrix.envs }}
         COMMAND: |
           source \${NCCL_LIB_DIR}/nccl-env-profile.sh;
           nsys-jax --capture-range=cudaProfilerApi \