From 58a09447ae161d752dcc00b261b31ea09feaa87c Mon Sep 17 00:00:00 2001
From: Zihao Ye <expye@outlook.com>
Date: Wed, 18 Dec 2024 11:18:56 -0500
Subject: [PATCH] bugfix: Revert "ci: cross python wheel (#662)" (#681)

This reverts PR #662 as we found there being some issue with cross
python wheel #669 , we should fix the issue later but we have to revert
it first to unblock v0.2 release job.

cc @xslingcn
---
 .github/workflows/release_wheel.yml | 30 +++++++++++++++++++++++++++--
 scripts/run-ci-build-wheel.sh       |  3 +++
 setup.py                            |  3 ---
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/release_wheel.yml b/.github/workflows/release_wheel.yml
index 187e97c8..9fbbf8bf 100644
--- a/.github/workflows/release_wheel.yml
+++ b/.github/workflows/release_wheel.yml
@@ -25,6 +25,7 @@ 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,6 +47,7 @@ jobs:
               -v "$CI_RUNNER_CACHE_DIR":/ci-cache \
               -v "$GITHUB_WORKSPACE":/app \
               -e FLASHINFER_CI_CACHE=/ci-cache \
+              -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" \
@@ -58,7 +60,7 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: wheel-cuda${{ matrix.cuda }}-torch${{ matrix.torch }}
+          name: wheel-cuda${{ matrix.cuda }}-torch${{ matrix.torch }}-python${{ matrix.python }}
           path: dist/*
 
   release:
@@ -77,7 +79,31 @@ jobs:
         with:
           tag_name: ${{ inputs.tag_name }}
           files: |
-            dist/flashinfer*.whl
+            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
 
       - uses: softprops/action-gh-release@v1
         with:
diff --git a/scripts/run-ci-build-wheel.sh b/scripts/run-ci-build-wheel.sh
index bb64f8e6..a1a24c17 100644
--- a/scripts/run-ci-build-wheel.sh
+++ b/scripts/run-ci-build-wheel.sh
@@ -11,6 +11,7 @@ assert_env() {
 }
 
 assert_env FLASHINFER_CI_CACHE
+assert_env FLASHINFER_CI_PYTHON_VERSION
 assert_env FLASHINFER_CI_CUDA_VERSION
 assert_env FLASHINFER_CI_TORCH_VERSION
 assert_env TORCH_CUDA_ARCH_LIST
@@ -25,6 +26,8 @@ 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"
 
 FLASHINFER_LOCAL_VERSION="cu${CUDA_MAJOR}${CUDA_MINOR}torch${FLASHINFER_CI_TORCH_VERSION}"
 if [ -n "${FLASHINFER_GIT_SHA}" ]; then
diff --git a/setup.py b/setup.py
index 4c1b8c58..02c18e17 100644
--- a/setup.py
+++ b/setup.py
@@ -213,7 +213,6 @@ def __init__(self, *args, **kwargs) -> None:
                 "cxx": cxx_flags,
                 "nvcc": nvcc_flags,
             },
-            py_limited_api=True,
         ),
         torch_cpp_ext.CUDAExtension(
             name="flashinfer._kernels_sm90",
@@ -223,7 +222,6 @@ def __init__(self, *args, **kwargs) -> None:
                 "cxx": cxx_flags,
                 "nvcc": nvcc_flags + sm90a_flags,
             },
-            py_limited_api=True,
         ),
     ]
 
@@ -231,6 +229,5 @@ def __init__(self, *args, **kwargs) -> None:
     version=get_version(),
     ext_modules=ext_modules,
     cmdclass=cmdclass,
-    options={"bdist_wheel": {"py_limited_api": "cp38"}},
     install_requires=install_requires,
 )