Merge remote-tracking branch 'origin/main' into update/BatchPrefillPagedWrapper

Author: yzh119

.github/workflows/release-ci-docker.yml

@@ -5,25 +5,74 @@ on:
 
 jobs:
   build:
-    strategy:
-      matrix:
-        include:
-          - os: ubuntu-latest
-            arch: amd64
-          - os: [self-hosted, Linux, ARM64]
-            arch: arm64
-    runs-on: ${{ matrix.os }}
+    runs-on: ubuntu-latest
     steps:
+      - name: Free Disk Space
+        uses: jlumbroso/free-disk-space@main
+        with:
+          tool-cache: false
+          android: true
+          dotnet: true
+          haskell: true
+          large-packages: true
+          docker-images: true
+          swap-storage: true
+
       - uses: actions/checkout@v4
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
       - name: Login to Docker Hub
         uses: docker/login-action@v3
         with:
           username: flashinfer
           password: ${{ secrets.DOCKERHUB_TOKEN }}
 
-      - uses: docker/build-push-action@v4
+      - name: Build and push AMD64 image
+        uses: docker/build-push-action@v5
         with:
           context: docker
           file: docker/Dockerfile.ci_gpu
+          platforms: linux/amd64
           push: true
-          tags: flashinfer/flashinfer-ci:latest
+          tags: |
+            flashinfer/flashinfer-ci:amd64
+          cache-from: type=registry,ref=flashinfer/flashinfer-ci:buildcache-amd64
+          cache-to: type=registry,ref=flashinfer/flashinfer-ci:buildcache-amd64,mode=max
+          provenance: false
+          sbom: false
+
+      - name: Clean up after AMD64 build
+        run: |
+          df -h /
+          docker buildx prune -af
+          docker system prune -af --volumes
+          sudo rm -rf /tmp/* /var/tmp/* ~/.docker/buildx
+          df -h /
+
+      - name: Re-setup Docker Buildx for ARM64
+        uses: docker/setup-buildx-action@v3
+
+      - name: Build and push ARM64 image
+        uses: docker/build-push-action@v5
+        with:
+          context: docker
+          file: docker/Dockerfile.ci_gpu
+          platforms: linux/arm64
+          push: true
+          tags: |
+            flashinfer/flashinfer-ci:arm64
+          cache-from: type=registry,ref=flashinfer/flashinfer-ci:buildcache-arm64
+          cache-to: type=registry,ref=flashinfer/flashinfer-ci:buildcache-arm64,mode=max
+          provenance: false
+          sbom: false
+
+      - name: Create and push multi-arch manifest
+        run: |
+          docker buildx imagetools create -t flashinfer/flashinfer-ci:latest \
+            flashinfer/flashinfer-ci:amd64 \
+            flashinfer/flashinfer-ci:arm64

3rdparty/composable_kernels

@@ -133,7 +133,7 @@ stage('Unittest') {
   cancel_previous_build()
   parallel(
     failFast: true,
-    'AOT-Build-Import': {
+    'AOT-Build-Import-x86-64': {
       try {
         run_unittest_CPU_AOT_COMPILE('CPU-LARGE-SPOT')
       } catch (Throwable ex) {
@@ -151,6 +151,24 @@ stage('Unittest') {
         }
       }
     },
+    'AOT-Build-Import-aarch64': {
+      try {
+        run_unittest_CPU_AOT_COMPILE('ARM-LARGE-SPOT')
+      } catch (Throwable ex) {
+        echo 'Exception during SPOT run ' + ex.toString()
+        if (is_last_build()) {
+          // retry if we are currently at last build
+          // mark the current stage as success
+          // and try again via on demand node
+          echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand'
+          currentBuild.result = 'SUCCESS'
+          run_unittest_CPU_AOT_COMPILE('ARM-LARGE')
+        } else {
+          echo 'Exit since it is not last build'
+          throw ex
+        }
+      }
+    },
     'JIT-Unittest-1': {
       try {
         shard_run_unittest_GPU('GPU-G5-SPOT', 1)

Jenkinsfile

@@ -0,0 +1,8 @@
+FlashInfer
+Copyright 2025 NVIDIA
+Copyright 2023-2025 FlashInfer community (https://flashinfer.ai/)
+
+-------------------------------------------------------------------------------------------------
+Some of the code in this project are adapted from other open-source projects with different
+licenses. This product also bundles some third-party components under other open source licenses.
+See licenses/ for text of these licenses.