[CI/CD] Add torch 2.2.0 wheels (#110)

Fix issue #105

Author: yzh119

.github/workflows/release_wheel.yml

@@ -19,7 +19,6 @@ on:
 
 env:
   TORCH_CUDA_ARCH_LIST: "8.0 8.6 8.9 9.0+PTX"
-  FLASHINFER_CI_TORCH_VERSION: "2.1.0"
 
 jobs:
   build:
@@ -28,6 +27,7 @@ jobs:
       matrix:
         python: ["3.10", "3.11"]
         cuda: ["11.8", "12.1"]
+        torch: ["2.1.0", "2.2.0"]
     runs-on: [self-hosted]
     steps:
       - uses: actions/checkout@v4
@@ -43,7 +43,7 @@ jobs:
               -v "$GITHUB_WORKSPACE":/app \
               -e FLASHINFER_CI_PYTHON_VERSION=${{ matrix.python }} \
               -e FLASHINFER_CI_CUDA_VERSION=${{ matrix.cuda }} \
-              -e FLASHINFER_CI_TORCH_VERSION=$FLASHINFER_CI_TORCH_VERSION \
+              -e FLASHINFER_CI_TORCH_VERSION=${{ matrix.torch }} \
               -e FLASHINFER_BUILD_VERSION=$version \
               -e TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST" \
               --user $CI_UID:$CI_GID \
@@ -54,7 +54,7 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: wheel-cuda${{ matrix.cuda }}-python${{ matrix.python }}
+          name: wheel-cuda${{ matrix.cuda }}-torch${{ matrix.torch }}-python${{ matrix.python }}
           path: python/dist/*
 
   release:
@@ -90,8 +90,8 @@ jobs:
           for path in sorted(pathlib.Path("python/dist").glob("*.whl")):
             with open(path, "rb") as f:
               sha256 = hashlib.sha256(f.read()).hexdigest()
-            ver, cu = re.findall(r"flashinfer-([0-9.]+)\+cu(\d+)-", path.name)[0]
-            with open(f"flashinfer-whl/cu{cu}/flashinfer/index.html", "a") as f:
+            ver, cu, torch = re.findall(r"flashinfer-([0-9.]+)\+cu(\d+)\+torch([0-9.]+)-", path.name)[0]
+            with open(f"flashinfer-whl/cu{cu}/torch{torch}/flashinfer/index.html", "a") as f:
               f.write(f'<a href="https://github.com/flashinfer-ai/flashinfer/releases/download/v{ver}/{path.name}#sha256={sha256}">{path.name}</a><br>\n')
 
       - name: Push wheel index

scripts/run-ci-build-wheel.sh

@@ -40,7 +40,7 @@ echo "::endgroup::"
 
 echo "::group::Build wheel for FlashInfer"
 cd "$PROJECT_ROOT/python"
-MAX_JOBS=32 FLASHINFER_BUILD_VERSION="${FLASHINFER_BUILD_VERSION}+cu${CUDA_MAJOR}${CUDA_MINOR}" python -m build --no-isolation
+MAX_JOBS=32 FLASHINFER_BUILD_VERSION="${FLASHINFER_BUILD_VERSION}+cu${CUDA_MAJOR}${CUDA_MINOR}+torch${FLASHINFER_CI_TORCH_VERSION}" python -m build --no-isolation
 rm -f dist/*.tar.gz
 python -m build --no-isolation --sdist
 echo "::endgroup::"

version.txt

@@ -1 +1 @@
-0.0.1
+0.0.2