Added sparse_matmul e2e tests for DS-v2.

Author: gpolovets1

.github/workflows/run_maxtext_jetstream_tests.yaml

@@ -107,19 +107,19 @@ jobs:
           bash -c "
             bash JetStream/.github/workflows/test_moe_benchmarks.sh
           "
-    - name: Test llama 70b Benchmarks
-      run: |
-        # sync with the image uploaded from build_stable_stack stage
-        # Report should generated in OUTPUT_DIR depend on ENV
-        DOCKER_OUTPUT_DIR=/output
-        docker run \
-          -v ${OUTPUT_DIR}:${DOCKER_OUTPUT_DIR} \
-          --env OUTPUT_DIR=${DOCKER_OUTPUT_DIR} \
-          --privileged --net=host --rm -i \
-          gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }} \
-          bash -c "
-            bash JetStream/.github/workflows/test_llama_benchmarks.sh
-          "
+    # - name: Test llama 70b Benchmarks
+    #   run: |
+    #     # sync with the image uploaded from build_stable_stack stage
+    #     # Report should generated in OUTPUT_DIR depend on ENV
+    #     DOCKER_OUTPUT_DIR=/output
+    #     docker run \
+    #       -v ${OUTPUT_DIR}:${DOCKER_OUTPUT_DIR} \
+    #       --env OUTPUT_DIR=${DOCKER_OUTPUT_DIR} \
+    #       --privileged --net=host --rm -i \
+    #       gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }} \
+    #       bash -c "
+    #         bash JetStream/.github/workflows/test_llama_benchmarks.sh
+    #       "
     - name: Upload build artifact
       uses: actions/upload-artifact@v4
       with:
@@ -205,5 +205,6 @@ jobs:
         to: singhvijaya@google.com, yuyanpeng@google.com, vipannalla@google.com
         from: JetStream Runs
         secure: true
-        attachments: ${{ env.BUILD_MANIFEST_DIR }}/${{ needs.build_stable_stack.outputs.manifest_name }},${{ env.BENCHMARK_REPORT_DIR }}/moe_8x7b.txt,${{ env.BENCHMARK_REPORT_DIR }}/moe_8x22b.txt,${{ env.BENCHMARK_REPORT_DIR }}/moe_8x22b_long_context_8k_prefill.txt,${{ env.BENCHMARK_REPORT_DIR }}/moe_8x7b_jetstream.txt,${{ env.BENCHMARK_REPORT_DIR }}/llama_70b_jetstream.txt,  ${OUTPUT_DIR}/golden-numbers.txt,  ${OUTPUT_DIR}/result_comparison.txt
+        # attachments: ${{ env.BUILD_MANIFEST_DIR }}/${{ needs.build_stable_stack.outputs.manifest_name }},${{ env.BENCHMARK_REPORT_DIR }}/moe_8x7b.txt,${{ env.BENCHMARK_REPORT_DIR }}/moe_8x22b.txt,${{ env.BENCHMARK_REPORT_DIR }}/moe_8x22b_long_context_8k_prefill.txt,${{ env.BENCHMARK_REPORT_DIR }}/moe_8x7b_jetstream.txt,${{ env.BENCHMARK_REPORT_DIR }}/llama_70b_jetstream.txt,  ${OUTPUT_DIR}/golden-numbers.txt,  ${OUTPUT_DIR}/result_comparison.txt
+        attachments: ${{ env.BUILD_MANIFEST_DIR }}/${{ needs.build_stable_stack.outputs.manifest_name }},${{ env.BENCHMARK_REPORT_DIR }}/deepseekv2-16b_ragged_dot.txt,${{ env.BENCHMARK_REPORT_DIR }}/deepseekv2-16b_megablox.txt, ${OUTPUT_DIR}/golden-numbers.txt,  ${OUTPUT_DIR}/result_comparison.txt
         body: workflow for ${{github.repository}} completed successfully!