Merge remote-tracking branch 'origin/main' into trtllm-attention-test-unify

Author: yzh119

.github/workflows/build-doc.yml

@@ -27,6 +27,11 @@ jobs:
       with:
         submodules: recursive
 
+    - name: Setup Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.12'
+
     - name: Configuring build Environment
       run: |
         sudo apt-get update

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

@@ -7,16 +7,72 @@ jobs:
   build:
     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

.pre-commit-config.yaml

@@ -35,22 +35,25 @@ repos:
       - id: remove-crlf
 
   # Formatters
-  - repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 24.8.0
-    hooks:
-      - id: black
-        exclude: flashinfer/tuning_configs/.*\.py
-
-  - repo: https://github.com/pycqa/isort
-    rev: 5.13.2
-    hooks:
-      - id: isort
-        args: ["--profile=black"] # <-- this one
-
   - repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v19.1.1
     hooks:
       - id: clang-format
         types_or: [c++, c, cuda]
         exclude: |
           (?x)^(3rdparty/.* flashinfer/jit/aot_config.py)$
+
+  -   repo: https://github.com/pre-commit/mirrors-mypy
+      rev: 'v1.17.1'  # Use the sha / tag you want to point at
+      hooks:
+      -   id: mypy
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.12.8
+    hooks:
+      # Run the linter.
+      - id: ruff-check
+      # Run the formatter.
+      - id: ruff-format
+        types_or: [ python, pyi ]

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.