ci: Improve compatibility with pytorch 2.5 (#711)

- Fix release workflow for pytorch 2.5 wheel by dropping python 3.8.
- Relax pytorch version requirements.
- Right now, install flashinfer will downgrade existing pytorch versions
since it's set the version requirement to exact match.

Author: zifeitong

.github/workflows/release_wheel.yml

@@ -33,6 +33,8 @@ jobs:
             torch: "2.2"
           - cuda: "12.4"
             torch: "2.3"
+          - python: "3.8" # torch 2.5+ drops python 3.8
+            torch: "2.5"
 
     runs-on: [self-hosted]
     steps:

scripts/run-ci-build-wheel.sh

@@ -41,7 +41,7 @@ else
 fi
 
 echo "::group::Install PyTorch"
-pip install torch==$FLASHINFER_CI_TORCH_VERSION --index-url "https://download.pytorch.org/whl/cu${CUDA_MAJOR}${CUDA_MINOR}"
+pip install torch==${FLASHINFER_CI_TORCH_VERSION}.* --index-url "https://download.pytorch.org/whl/cu${CUDA_MAJOR}${CUDA_MINOR}"
 echo "::endgroup::"
 
 echo "::group::Install build system"

setup.py

@@ -135,7 +135,7 @@ def __init__(self, *args, **kwargs) -> None:
     torch_full_version = Version(torch.__version__)
     torch_version = f"{torch_full_version.major}.{torch_full_version.minor}"
     cmdclass["build_ext"] = NinjaBuildExtension
-    install_requires = [f"torch == {torch_version}"]
+    install_requires = [f"torch == {torch_version}.*"]
 
     aot_build_meta = {}
     aot_build_meta["cuda_major"] = cuda_version.major