NVIDIA · aybchan · Jun 5, 2026 · Jun 8, 2026 · Jun 8, 2026 · Jun 26, 2026
diff --git a/.github/actions/gke-xpk/action.yml b/.github/actions/gke-xpk/action.yml
@@ -32,6 +32,11 @@ inputs:
     default: 2
     required: false
     type: string
+  NUM_GPUS:
+    description: 'Number of GPUs to use in JobSet'
+    default: 8
+    required: false
+    type: string
   MAIN_CONTAINER_NAME:
     description: 'Name of the main contianer in an XPK JobSet (fixed in xpk)'
     default: gpu-image
@@ -65,7 +70,7 @@ inputs:
   EXIT_COMMAND:
     description: 'Command to set exit code'
     required: false
-    default: 'exit \$EXIT_CODE'
+    default: 'echo EXIT_CODE=\$EXIT_CODE; exit \$EXIT_CODE'
     type: string
   NAME:
     description: 'Workload name prefix for XPK, also used to name uploaded artifact'
@@ -195,34 +200,32 @@ runs:
     run: |
       # install dependencies to enable export artifacts from container to gcs bucket
       PRELUDE="
-          apt install -y ripgrep > /dev/null;
-          curl -LO https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-cli-linux-x86_64.tar.gz;
-          tar xf google-cloud-cli-linux-x86_64.tar.gz;
-          ./google-cloud-sdk/install.sh --quiet > /dev/null;
-          ./google-cloud-sdk/bin/gcloud init;
+          apt update; apt install -y ripgrep > /dev/null;
+          curl -L https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-cli-linux-x86_64.tar.gz -o /tmp/google-cloud-cli-linux-x86_64.tar.gz;
+          tar -C /opt -xf /tmp/google-cloud-cli-linux-x86_64.tar.gz;
+          /opt/google-cloud-sdk/install.sh --quiet > /dev/null;
+          /opt/google-cloud-sdk/bin/gcloud init;
 
           mkdir -p /usr/share/workload;
-          mkdir -p ${{ inputs.CONTAINER_OUTPUT_PATH }};"
+          mkdir -p ${{ inputs.CONTAINER_OUTPUT_PATH }};
 
-      PRELUDE+="
-          test -f \${NCCL_LIB_DIR}/nccl-env-profile.sh && NCCL_LIB_DIR=\${NCCL_LIB_DIR} . \${NCCL_LIB_DIR}/nccl-env-profile.sh;
-          export LD_LIBRARY_PATH=\${NCCL_LIB_DIR}:\${LD_LIBRARY_PATH};
-          export NCCL_TUNER_CONFIG_PATH=\${NCCL_LIB_DIR}/a3plus_tuner_config.textproto;
-          export NCCL_SHIMNET_GUEST_CONFIG_CHECKER_CONFIG_FILE=\${NCCL_LIB_DIR}/a3plus_guest_config.textproto;
+          NCCL_LIB_DIR=\${NCCL_LIB_DIR} . \${NCCL_LIB_DIR}/nccl-env-profile.sh;
       "
 
       # gsutil command to export logs from container's /opt/output to bucket
       POSTLUDE="
-          ./google-cloud-sdk/bin/gsutil cp -r ${{ inputs.CONTAINER_OUTPUT_PATH }}/ ${GCS_ARTIFACT_PATH}/node-0\$NODE_RANK;
+          COMMAND_EXIT_CODE=\$?;
+          if [ -z \${EXIT_CODE:-} ]; then EXIT_CODE=\${COMMAND_EXIT_CODE}; fi;
+          /opt/google-cloud-sdk/bin/gsutil cp -r ${{ inputs.CONTAINER_OUTPUT_PATH }}/ ${GCS_ARTIFACT_PATH}/;
           ${{ inputs.EXIT_COMMAND }}
       "
 
       CMD="${{ inputs.COMMAND }}"
 
       # set container commands in-line
-      PRELUDE=$(echo ${PRELUDE} | sed 's/\n/\ /g')
-      POSTLUDE=$(echo ${POSTLUDE} | sed 's/\n/\ /g')
-      CMD=$(echo ${CMD} | sed 's/\n/\ /g')
+      PRELUDE=$(printf '%s\n' "${PRELUDE}" | paste -sd ' ' - )
+      POSTLUDE=$(printf '%s\n' "${POSTLUDE}" | paste -sd ' ' - )
+      CMD=$(printf '%s\n' "${CMD}" | paste -sd ' ' - )
 
       echo "CMD=${PRELUDE} ${CMD} ${POSTLUDE}" >> ${GITHUB_ENV}
 
@@ -318,6 +321,7 @@ runs:
       echo "JobSet ${JOBSET_NAME} has just become active in the cluster"
 
   - name: Set JobSet Pods
+    if: steps.check.outputs.online == 'true'
     shell: bash -u {0}
     run: |
       echo "JOBSET_PODS=($(kubectl get pods -o json | jq -r '.items[].metadata | select(.labels."jobset.sigs.k8s.io/jobset-name" == "'${JOBSET_NAME}'") | .name' | tr '\n' ' '))" >> ${GITHUB_ENV}
@@ -341,7 +345,7 @@ runs:
       mkdir -p ${JOBSET_NAME}
 
       for jobset_pod in ${JOBSET_PODS//[()]/}; do
-          kubectl logs --pod-running-timeout=1m -f --prefix=true --timestamps=true -c gpu-image ${jobset_pod} 2>&1 | tee -a ${JOBSET_NAME}/${jobset_pod}.log &
+          kubectl logs --pod-running-timeout=1m -f --prefix=true --timestamps=true -c ${{ inputs.MAIN_CONTAINER_NAME }} ${jobset_pod} 2>&1 | tee -a ${JOBSET_NAME}/${jobset_pod}.log &
       done
       wait < <(jobs -p)
 
@@ -379,28 +383,28 @@ runs:
       exit 0
 
   - name: Clean up JobSet from cluster
-    if: steps.check.outputs.online == 'true'
+    if: ${{ always() && steps.check.outputs.online == 'true' }}
     shell: bash -x -u {0}
     run: |
       kubectl delete jobset --wait ${JOBSET_NAME} || echo "JobSet ${JOBSET_NAME} does not exist"
 
   - name: Download artifacts from GCS to runner
-    if: steps.check.outputs.online == 'true'
+    if: ${{ !cancelled() && steps.check.outputs.online == 'true' }}
     shell: bash -x -u {0}
     run: |
       mkdir -p ${JOBSET_NAME}/output
-      gsutil cp -r ${GCS_ARTIFACT_PATH} ${JOBSET_NAME}/output
+      gsutil cp -r ${GCS_ARTIFACT_PATH} ${JOBSET_NAME}
       cp ${JOBSET_NAME}/*.log ${JOBSET_NAME}/output
 
   - name: Upload artifacts to GitHub Actions from runner
-    if: steps.check.outputs.online == 'true'
+    if: ${{ !cancelled() && steps.check.outputs.online == 'true' }}
     uses: actions/upload-artifact@v6
     with:
       name: ${{ inputs.NAME }}
       path: ${{ env.JOBSET_NAME }}/output/*
 
   - name: Emit unit test statistics
-    if: ${{ !cancelled() }}
+    if: ${{ !cancelled() && steps.check.outputs.online == 'true' }}
     shell: bash -xe -u {0}
     run: |
       STATISTICS_SCRIPT="${{ inputs.STATISTICS_SCRIPT }}"
@@ -411,7 +415,6 @@ runs:
         BADGE_COLOR=red
         BADGE_MESSAGE=error
         SUMMARY="${{ inputs.NAME }} unit test on ${{ inputs.GPU_ARCH }} did not complete due to errors."
-        continue
       else
         if [[ ${NUM_FAILED} == 0 ]]; then
           BADGE_COLOR=brightgreen
@@ -438,7 +441,7 @@ runs:
       > badge-${{ inputs.NAME }}-${{ inputs.GPU_ARCH }}.json
 
   - name: Upload statistics artifacts
-    if: ${{ !cancelled() }}
+    if: ${{ !cancelled() && steps.check.outputs.online == 'true' }}
     uses: actions/upload-artifact@v6
     with:
       name: ${{ inputs.NAME }}-${{ inputs.GPU_ARCH }}
@@ -450,4 +453,4 @@ runs:
     if: ${{ always() }}
     shell: bash -x -u {0}
     run: |
-      sudo rm -rf ${JOBSET_NAME}
+      sudo rm -rf ${JOBSET_NAME:-} || echo done
diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml
@@ -421,17 +421,24 @@ jobs:
       matrix:
         include:
           - test: L0_jax_unittest
+            name: l0-unit
+            nodes: 1
             gpus: 8
           - test: L0_jax_distributed_unittest
+            name: l0-dist
+            nodes: 1
             gpus: 8
-    runs-on: eks
+    runs-on: gke-a3mega
     steps:
       - uses: actions/checkout@v4
-      - name: Run eks-jobset action
-        uses: ./.github/actions/eks-jobset
+      - name: Run XPK workload on cluster
+        uses: ./.github/actions/gke-xpk
         with:
-          NAME: te-${{ matrix.test }}
+          NAME: te-${{ matrix.name }}
+          NUM_NODES: ${{ matrix.nodes }}
           NUM_GPUS: ${{ matrix.gpus }}
+          ENVS: |
+            CUDNN_FRONTEND_CUDART_LIB_NAME=libcudart.so.13
           IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }}
           COMMAND: |
             set -xu -o pipefail;
@@ -442,8 +449,8 @@ jobs:
             nvidia-cuda-mps-control -d;
             export TE_PATH=\${SRC_PATH_TRANSFORMER_ENGINE};
 
-            [[ "${{ matrix.test }}" == *distributed* ]] && xdist_args=(8 1) || xdist_args=(1 4);
-            pytest-xdist.sh \${xdist_args[@]} /opt/output/pytest-report.jsonl bash \${TE_PATH}/qa/${{ matrix.test }}/test.sh | tee /opt/output/pytest.log
+            [[ ${{ matrix.test }} == *distributed* ]] && xdist_args=(8 1) || xdist_args=(1 4);
+            pytest-xdist.sh \${xdist_args[@]} /opt/output/pytest-report.jsonl bash \${TE_PATH}/qa/${{ matrix.test }}/test.sh | tee /opt/output/pytest.log;
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           NVCR_TOKEN: ${{ secrets.NVCR_TOKEN }}
           STATISTICS_SCRIPT: |
@@ -461,27 +468,30 @@ jobs:
         inputs.MODE == 'full' ||
         inputs.MODE == 'jax'
       )
-    runs-on: eks
+    runs-on: gke-a3mega
     needs: build-jax
     strategy:
       fail-fast: false
       matrix:
         include:
-          - gpus: 8
+          - nodes: 1
+            gpus: 8
     steps:
       - uses: actions/checkout@v4
-      - name: Run eks-jobset action
-        uses: ./.github/actions/eks-jobset
+      - name: Run XPK workload on cluster
+        uses: ./.github/actions/gke-xpk
         with:
           NAME: nsys-jax-unit
+          NUM_NODES: ${{ matrix.nodes }}
           NUM_GPUS: ${{ matrix.gpus }}
           IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }}
           # nsys-jax is already installed, this is just adding the test dependencies
           # abuse knowledge that nsys-jax is installed editable, so the tests exist
           COMMAND: |
             pip install pytest-reportlog nsys-jax[test];
-            test_path=\$(python -c 'import importlib.resources; print(importlib.resources.files(\"nsys_jax\").joinpath(\"..\", \"tests\").resolve())');
-            pytest --basetemp=/opt/output/pytest-tmp --report-log=/opt/output/pytest-report.jsonl \"\${test_path}\" | tee /opt/output/pytest.log
+            test_path=\$(eval \$'python - <<\\\\PY\nimport importlib.resources\nprint(importlib.resources.files(\\\"nsys_jax\\\").joinpath(\\\"..\\\", \\\"tests\\\").resolve())\nPY');
+            pytest --basetemp=/opt/output/pytest-tmp --report-log=/opt/output/pytest-report.jsonl \${test_path} | tee /opt/output/pytest.log;
+            EXIT_CODE=${PIPESTATUS[0]};
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           NVCR_TOKEN: ${{ secrets.NVCR_TOKEN }}
           STATISTICS_SCRIPT: |
@@ -536,7 +546,7 @@ jobs:
         inputs.MODE == 'full' ||
         inputs.MODE == 'jax'
       )
-    runs-on: eks
+    runs-on: gke-a3mega
     strategy:
       fail-fast: false
       matrix:
@@ -546,8 +556,8 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - name: Run eks-jobset action
-        uses: ./.github/actions/eks-jobset
+      - name: Run XPK workload on cluster
+        uses: ./.github/actions/gke-xpk
         with:
           NAME: nsys-jax-nccl
           NUM_NODES: ${{ matrix.nodes }}
@@ -562,7 +572,7 @@ jobs:
               --distributed \
               --gpus-per-process=${{ matrix.gpus }} \
               --process-count=${{ matrix.nodes }} \
-              --process-id=\${NODE_RANK}
+              --process-id=\${NODE_RANK};
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           NVCR_TOKEN: ${{ secrets.NVCR_TOKEN }}
 
@@ -574,7 +584,7 @@ jobs:
         inputs.MODE == 'full' ||
         inputs.MODE == 'jax-cutlass'
       )
-    runs-on: eks
+    runs-on: gke-a3mega
     strategy:
       fail-fast: false
       matrix:
@@ -583,8 +593,8 @@ jobs:
             gpus: 8
     steps:
       - uses: actions/checkout@v4
-      - name: Run eks-jobset action
-        uses: ./.github/actions/eks-jobset
+      - name: Run XPK workload on cluster
+        uses: ./.github/actions/gke-xpk
         with:
           NAME: jax-cutlass
           NUM_NODES: ${{ matrix.nodes }}
@@ -597,17 +607,17 @@ jobs:
 
             any_failed=0;
             for f in /cutlass/examples/python/CuTeDSL/dsl_tutorials/jax/*.py; do
-              echo \"[Executing] \${f}\";
+              echo [Executing] \${f};
               log_output=\$(python \${f} 2>&1);
               exit_code=\${?};
-              if [[ "\${exit_code}" -ne 0 ]]; then any_failed=1; fi;
+              if [[ \${exit_code} -ne 0 ]]; then any_failed=1; fi;
 
-              outcome=\$( [ \${exit_code} -eq 0 ] && echo \"passed\" || echo \"failed\" );
-              echo \"=== \${f} ===\" | tee -a /opt/output/pytest_stdout.log;
-              echo \"\${log_output}\" | tee -a /opt/output/pytest_stdout.log;
-              python3 -c \"import json,sys; print(json.dumps({'outcome': sys.argv[1], 'nodeid': sys.argv[2], 'longrepr': sys.argv[3]}))\" \"\${outcome}\" \"\${f}\" \"\${log_output}\" >> /opt/output/pytest-report.jsonl;
+              outcome=\$( [ \${exit_code} -eq 0 ] && echo passed || echo failed );
+              echo === \${f} === | tee -a /opt/output/pytest_stdout.log;
+              echo \${log_output} | tee -a /opt/output/pytest_stdout.log;
+              eval \$'python3 - \\\"\${outcome}\\\" \\\"\${f}\\\" \\\"\${log_output}\\\" >> /opt/output/pytest-report.jsonl <<\\\\PY\nimport json\nimport sys\nprint(json.dumps({\\\"outcome\\\": sys.argv[1], \\\"nodeid\\\": sys.argv[2], \\\"longrepr\\\": sys.argv[3]}))\nPY';
             done;
-            if [[ \"\${any_failed}\" -ne 0 ]]; then false; else true; fi
+            if [[ \${any_failed} -ne 0 ]]; then false; else true; fi;
 
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           NVCR_TOKEN: ${{ secrets.NVCR_TOKEN }}

diff --git a/.github/workflows/_test_maxtext_gke_xpk.yaml b/.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 \
@@ -56,7 +58,7 @@ jobs:
                                    --model-name=${{ matrix.model }} \
                                    --attn-type=${{ matrix.attention }} \
                                    --remat-policy=${{ matrix.remat-policy }} \
-                                   --steps=10 \
+                                   --steps=20 \
                                    --pipeline-parallel ${{ matrix.pipeline-parallel }} \
                                    --data-parallel ${{ matrix.data-parallel }} \
                                    --fsdp ${{ matrix.fsdp }} \
@@ -68,8 +70,8 @@ jobs:
                                        logits_dot_in_fp32=false
                                        profiler=nsys
                                        upload_all_profiler_results=true
-                                       skip_first_n_steps_for_profiler=3
-                                       profiler_steps=8' |&
+                                       skip_first_n_steps_for_profiler=10
+                                       profiler_steps=10' |&
           tee /opt/output/output.log &> /dev/stdout;
           EXIT_CODE=\${PIPESTATUS[0]};
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}