upd

Author: yzh119

.github/workflows/nightly-release.yml

@@ -52,7 +52,22 @@ jobs:
           echo "VERSION=$version" >> $GITHUB_ENV
           echo "version=$version" >> $GITHUB_OUTPUT
         id: get_version
+
+      - name: Set torch and cuda version value
+        id: set_torch_and_cuda_version
+        run: |
+          IFS='.' read -r major minor <<< "${{ matrix.torch }}"
+          version_value=$((major * 100 + minor))
+          echo "TORCH_VERSION=$version_value" >> $GITHUB_OUTPUT
+          IFS='.' read -r major minor <<< "${{ matrix.cuda }}"
+          version_value=$((major * 100 + minor))
+          echo "CUDA_VERSION=$version_value" >> $GITHUB_OUTPUT
+
       - name: Build wheel
+        env:
+          TORCH_VERSION: ${{ steps.set_torch_and_cuda_version.outputs.TORCH_VERSION }}
+          CUDA_VERSION: ${{ steps.set_torch_and_cuda_version.outputs.CUDA_VERSION }}
+          BUILDER_IMAGE: ${{ ((steps.set_torch_and_cuda_version.outputs.TORCH_VERSION == 206 && steps.set_torch_and_cuda_version.outputs.CUDA_VERSION >= 1206) || steps.set_torch_and_cuda_version.outputs.TORCH_VERSION > 206) && 'pytorch/manylinux2_28-builder:cuda' || 'pytorch/manylinux-builder:cuda' }}${{ matrix.cuda }}
         run: |
           pushd flashinfer
           docker run --rm -t \
@@ -65,7 +80,7 @@ jobs:
               -e TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST" \
               -e MAX_JOBS=192 \
               --user $(id -u):$(id -g) \
-              pytorch/manylinux-builder:cuda${{ matrix.cuda }} \
+              $BUILDER_IMAGE \
               bash /app/scripts/run-ci-build-wheel.sh
           popd
         timeout-minutes: 30