chore: support building torch2.8 wheels (#439)

* support building torch2.8 wheels

* avoid recreating the tag

* rebuild torch2.8 wheels

* support torch2.8 wheels

.github/workflows/nightly-build.yaml

@@ -77,7 +77,7 @@ jobs:
     strategy:
       matrix:
         python: ["3.10", "3.11", "3.12"]
-        torch: ["2.5", "2.6", "2.7"]
+        torch: ["2.5", "2.6", "2.7", "2.8"]
     steps:
       - name: Checkout to the tag
         uses: actions/checkout@v4
@@ -95,6 +95,11 @@ jobs:
             cuda_version="12.4"
           fi
           bash scripts/build_linux_wheel.sh ${{ matrix.python }} ${{ matrix.torch }} $cuda_version
+          if [[ "${{ matrix.torch }}" == "2.8" ]]; then
+            bash scripts/build_linux_wheel_torch2.8.sh ${{ matrix.python }} ${{ matrix.torch }} $cuda_version
+          else
+            bash scripts/build_linux_wheel.sh ${{ matrix.python }} ${{ matrix.torch }} $cuda_version
+          fi
       - name: Upload wheels to GitHub Release
         uses: softprops/action-gh-release@v2
         with:
@@ -115,7 +120,7 @@ jobs:
     strategy:
       matrix:
         python: ["3.10", "3.11", "3.12"]
-        torch: ["2.5", "2.6", "2.7"]
+        torch: ["2.5", "2.6", "2.7", "2.8"]
     steps:
       - name: Checkout to the tag
         uses: actions/checkout@v4
@@ -129,13 +134,18 @@ jobs:
         shell: cmd
         run: |
           SET TORCH_VERSION=${{ matrix.torch }}
+          SET PYTHON_VERSION=${{ matrix.python }}
           IF "%TORCH_VERSION%"=="2.7" (
             SET CUDA_VERSION=12.8
           ) ELSE (
             SET CUDA_VERSION=12.4
           )
           call C:\Users\muyangl\miniconda3\condabin\activate.bat activate
-          call scripts\build_windows_wheel.cmd ${{ matrix.python }} %TORCH_VERSION% %CUDA_VERSION%
+          IF "%TORCH_VERSION%"=="2.8" (
+            call scripts\build_windows_wheel_torch2.8.cmd %PYTHON_VERSION% %TORCH_VERSION% %CUDA_VERSION%
+          ) ELSE (
+            call scripts\build_windows_wheel.cmd %PYTHON_VERSION% %TORCH_VERSION% %CUDA_VERSION%
+          )
       - name: Upload wheels to GitHub Release
         uses: softprops/action-gh-release@v2
         with:

.github/workflows/release-build.yaml

@@ -9,7 +9,7 @@ jobs:
     runs-on: ubuntu-latest
     if: github.repository == 'mit-han-lab/nunchaku'
     outputs:
-      tag_name: ${{ steps.tag.outputs.tag_name }}
+      tag_name: ${{ steps.set-tag.outputs.tag_name }}
     steps:
       - name: Checkout main branch
         uses: actions/checkout@v4
@@ -20,17 +20,30 @@ jobs:
         id: version
         run: |
           version=$(grep '__version__' nunchaku/__version__.py | sed -E 's/.*"([^"]+)".*/\1/')
-          echo "Extracted version: $version"
           echo "version=$version" >> "$GITHUB_OUTPUT"
-      - name: Create and push tag
-        id: tag
+      - name: Check if tag exists
+        id: check-tag
+        run: |
+          tag_name="v${{ steps.version.outputs.version }}"
+          if git rev-parse "$tag_name" >/dev/null 2>&1; then
+            echo "Tag $tag_name already exists."
+            echo "exists=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "Tag $tag_name does not exist."
+            echo "exists=false" >> "$GITHUB_OUTPUT"
+          fi
+      - name: Create and push tag if not exists
+        if: steps.check-tag.outputs.exists == 'false'
         run: |
           tag_name="v${{ steps.version.outputs.version }}"
           git config user.name "github-actions"
           git config user.email "github-actions@users.noreply.github.com"
-          git tag $tag_name
-          git push origin $tag_name
-          echo "tag_name=$tag_name" >> "$GITHUB_OUTPUT"
+          git tag "$tag_name"
+          git push origin "$tag_name"
+      - name: Set tag_name output
+        id: set-tag
+        run: |
+          echo "tag_name=v${{ steps.version.outputs.version }}" >> "$GITHUB_OUTPUT"
   linux-wheels:
     name: Build the linux release wheels
     runs-on: [self-hosted, linux-build]
@@ -56,6 +69,11 @@ jobs:
             cuda_version="12.4"
           fi
           bash scripts/build_linux_wheel.sh ${{ matrix.python }} ${{ matrix.torch }} $cuda_version
+          if [[ "${{ matrix.torch }}" == "2.8" ]]; then
+            bash scripts/build_linux_wheel_torch2.8.sh ${{ matrix.python }} ${{ matrix.torch }} $cuda_version
+          else
+            bash scripts/build_linux_wheel.sh ${{ matrix.python }} ${{ matrix.torch }} $cuda_version
+          fi
       - name: Upload wheels to GitHub Release
         uses: softprops/action-gh-release@v2
         with:
@@ -89,13 +107,18 @@ jobs:
         shell: cmd
         run: |
           SET TORCH_VERSION=${{ matrix.torch }}
+          SET PYTHON_VERSION=${{ matrix.python }}
           IF "%TORCH_VERSION%"=="2.7" (
             SET CUDA_VERSION=12.8
           ) ELSE (
             SET CUDA_VERSION=12.4
           )
           call C:\Users\muyangl\miniconda3\condabin\activate.bat activate
-          call scripts\build_windows_wheel.cmd ${{ matrix.python }} %TORCH_VERSION% %CUDA_VERSION%
+          IF "%TORCH_VERSION%"=="2.8" (
+            call scripts\build_windows_wheel_torch2.8.cmd %PYTHON_VERSION% %TORCH_VERSION% %CUDA_VERSION%
+          ) ELSE (
+            call scripts\build_windows_wheel.cmd %PYTHON_VERSION% %TORCH_VERSION% %CUDA_VERSION%
+          )
       - name: Upload wheels to GitHub Release
         uses: softprops/action-gh-release@v2
         with:

scripts/build_linux_wheel_torch2.8.sh

+#!/bin/bash
+# Modified from https://github.com/sgl-project/sglang/blob/main/sgl-kernel/build.sh
+set -ex
+PYTHON_VERSION=$1
+TORCH_VERSION=$2
+CUDA_VERSION=$3
+MAX_JOBS=${4:-} # optional
+PYTHON_ROOT_PATH=/opt/python/cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}
+
+if [ "$TORCH_VERSION" == "2.8" ]; then
+  TORCHVISION_VERSION="0.23"
+  TORCHAUDIO_VERSION="2.8"
+  echo "TORCH_VERSION is 2.8, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
+else
+  echo "TORCH_VERSION is not 2.8, no changes to versions."
+fi
+
+docker run --rm \
+    -v "$(pwd)":/nunchaku \
+    pytorch/manylinux2_28-builder:cuda${CUDA_VERSION} \
+    bash -c "
+    cd /nunchaku && \
+    rm -rf build && \
+    gcc --version && g++ --version && \
+    ${PYTHON_ROOT_PATH}/bin/pip install --pre --no-cache-dir torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} torchaudio==${TORCHAUDIO_VERSION} --index-url https://download.pytorch.org/whl/nightly/cu${CUDA_VERSION//.} && \
+    ${PYTHON_ROOT_PATH}/bin/pip install build ninja wheel setuptools && \
+    export NUNCHAKU_INSTALL_MODE=ALL && \
+    export NUNCHAKU_BUILD_WHEELS=1 && \
+    export MAX_JOBS=${MAX_JOBS} && \
+    ${PYTHON_ROOT_PATH}/bin/python -m build --wheel --no-isolation
+    "

scripts/build_windows_wheel_torch_nightly.cmd → scripts/build_windows_wheel_torch2.8.cmd

@@ -20,7 +20,7 @@ call conda activate %ENV_NAME%
 :: install dependencies
 call pip install ninja setuptools wheel build
 
-call pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+call pip install --pre torch==2.8 torchvision==0.23 torchaudio==2.8 --index-url https://download.pytorch.org/whl/nightly/cu128
 
 :: set environment variables
 set NUNCHAKU_INSTALL_MODE=ALL
@@ -33,7 +33,6 @@ if exist build rd /s /q build
 :: set up Visual Studio compilation environment
 call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
 set DISTUTILS_USE_SDK=1
-set MAX_JOBS=4
 
 :: build wheels
 python -m build --wheel --no-isolation