ci: cross python wheel (#662)

We can build a single python wheel for all supported python versions. - Remove the python versions matrix - Add `--py-limited-api` config in building commands
flashinfer-ai · Dec 14, 2024 · b1b1fb8 · b1b1fb8
1 parent 56f8eca
commit b1b1fb8
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 31 deletions.
diff --git a/.github/workflows/release_wheel.yml b/.github/workflows/release_wheel.yml
@@ -25,7 +25,6 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python: ["3.8", "3.9", "3.10", "3.11", "3.12"]
         cuda: ["11.8", "12.1", "12.4"]
         torch: ["2.2", "2.3", "2.4"]
         exclude: # for cuda 12.4, we only support torch 2.4+
@@ -46,7 +45,6 @@ jobs:
           docker run --rm -t \
               -v "$CI_RUNNER_CACHE_DIR":/ci-cache \
               -v "$GITHUB_WORKSPACE":/app \
-              -e FLASHINFER_CI_PYTHON_VERSION=${{ matrix.python }} \
               -e FLASHINFER_CI_CUDA_VERSION=${{ matrix.cuda }} \
               -e FLASHINFER_CI_TORCH_VERSION=${{ matrix.torch }} \
               -e TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST" \
@@ -59,7 +57,7 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: wheel-cuda${{ matrix.cuda }}-torch${{ matrix.torch }}-python${{ matrix.python }}
+          name: wheel-cuda${{ matrix.cuda }}-torch${{ matrix.torch }}
           path: dist/*
 
   release:
@@ -78,31 +76,7 @@ jobs:
         with:
           tag_name: ${{ inputs.tag_name }}
           files: |
-            dist/flashinfer*cp38*.whl
-
-      - uses: softprops/action-gh-release@v1
-        with:
-          tag_name: ${{ inputs.tag_name }}
-          files: |
-            dist/flashinfer*cp39*.whl
-
-      - uses: softprops/action-gh-release@v1
-        with:
-          tag_name: ${{ inputs.tag_name }}
-          files: |
-            dist/flashinfer*cp310*.whl
-
-      - uses: softprops/action-gh-release@v1
-        with:
-          tag_name: ${{ inputs.tag_name }}
-          files: |
-            dist/flashinfer*cp311*.whl
-
-      - uses: softprops/action-gh-release@v1
-        with:
-          tag_name: ${{ inputs.tag_name }}
-          files: |
-            dist/flashinfer*cp312*.whl
+            dist/flashinfer*.whl
 
       - uses: softprops/action-gh-release@v1
         with:

diff --git a/scripts/run-ci-build-wheel.sh b/scripts/run-ci-build-wheel.sh
@@ -10,7 +10,6 @@ assert_env() {
     fi
 }
 
-assert_env FLASHINFER_CI_PYTHON_VERSION
 assert_env FLASHINFER_CI_TORCH_VERSION
 assert_env FLASHINFER_CI_CUDA_VERSION
 assert_env TORCH_CUDA_ARCH_LIST
@@ -25,8 +24,6 @@ CUDA_MAJOR="${FLASHINFER_CI_CUDA_VERSION%.*}"
 CUDA_MINOR="${FLASHINFER_CI_CUDA_VERSION#*.}"
 TORCH_MAJOR="${FLASHINFER_CI_TORCH_VERSION%.*}"
 TORCH_MINOR="${FLASHINFER_CI_TORCH_VERSION#*.}"
-PYVER="${FLASHINFER_CI_PYTHON_VERSION//./}"
-export PATH="/opt/python/cp${PYVER}-cp${PYVER}/bin:$PATH"
 
 
 echo "::group::Install PyTorch"

diff --git a/setup.py b/setup.py
@@ -198,6 +198,7 @@ def __init__(self, *args, **kwargs) -> None:
                 "cxx": cxx_flags,
                 "nvcc": nvcc_flags,
             },
+            py_limited_api=True,
         ),
         torch_cpp_ext.CUDAExtension(
             name="flashinfer._kernels_sm90",
@@ -207,12 +208,14 @@ def __init__(self, *args, **kwargs) -> None:
                 "cxx": cxx_flags,
                 "nvcc": nvcc_flags + sm90a_flags,
             },
+            py_limited_api=True,
         ),
     ]
 
 setuptools.setup(
     version=get_version(),
     ext_modules=ext_modules,
     cmdclass=cmdclass,
+    options={"bdist_wheel": {"py_limited_api": "cp38"}},
     install_requires=install_requires,
 )