Build/add torch 2.8 2.9 cuda 128

.github/workflows/wheels.yml

@@ -4,22 +4,19 @@ on:
   workflow_dispatch:
 
 jobs:
-  build:
+  build-local:
     strategy:
       fail-fast: false
       matrix:
         os:
           - ubuntu-22.04
           # - windows-2019
-        python:
-          - "3.9"
-          - "3.10"
-          - "3.11"
-          - "3.12"
-        torch_version:
-          - "2.7.0"
-        cuda_short_version:
-          - "126"
+        python: ['3.10', '3.11', '3.12']
+        torch_version: ['2.7.0', '2.8.0', '2.9.0']
+        cuda_short_version: ['126']
+        exclude:
+          - torch_version: '2.9.0'
+            python: '3.9'
 
     uses: ./.github/workflows/wheels_build.yml
     with:
@@ -28,44 +25,82 @@ jobs:
       torch_version: ${{ matrix.torch_version }}
       cuda_short_version: ${{ matrix.cuda_short_version }}
 
+  build-pypi:
+    # Single canonical build intended for PyPI: no local CUDA/torch suffix
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ['ubuntu-22.04']
+        python: ['3.10', '3.11', '3.12']
+
+    uses: ./.github/workflows/wheels_build.yml
+    with:
+      os: ${{ matrix.os }}
+      python: ${{ matrix.python }}
+      torch_version: '2.9.0'
+      cuda_short_version: '128'
+      append_local_version: '0'  # 0 to disable local version suffix
+
   # publish to GitHub Release
-  gh_release:
-    name: gh_release
-    needs: build
-    runs-on: ubuntu-20.04
-
-    timeout-minutes: 360
-    defaults:
-      run:
-        shell: bash
+  # gh_release:
+  #   name: gh_release
+  #   needs: build
+  #   runs-on: ubuntu-20.04
+
+  #   timeout-minutes: 360
+  #   defaults:
+  #     run:
+  #       shell: bash
+  #   steps:
+  #     - uses: actions/download-artifact@v4
+  #       with:
+  #         path: dist
+
+  #     - run: ls -R dist/
+
+  #     # create night release if it's a push to main
+  #     - if: github.event_name == 'push' && github.ref == 'refs/heads/main'
+  #       name: Nightly Release
+  #       uses: andelf/nightly-release@v1
+  #       env:
+  #         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  #       with:
+  #         tag_name: nightly
+  #         name: 'stable-fast Nightly Release $$'
+  #         prerelease: true
+  #         body: 'TODO: Add nightly release notes'
+  #         files: |
+  #           dist/*/*.whl
+
+  #     # create release if it's a tag like vx.y.z
+  #     - if: github.ref_type == 'tag' && startsWith(github.ref, 'refs/tags/v')
+  #       name: Release
+  #       uses: softprops/action-gh-release@v1
+  #       with :
+  #         files: |
+  #           dist/*/*.whl
+
+
+  consolidate-wheels:
+    needs: [build-local, build-pypi]
+    runs-on: ubuntu-latest
     steps:
-      - uses: actions/download-artifact@v4
+      - name: Download all wheel artifacts
+        uses: actions/download-artifact@v4
         with:
           path: dist
 
-      - run: ls -R dist/
+      - name: Consolidate wheels into a single folder
+        run: |
+          mkdir -p consolidated_wheels
+          find dist -name '*.whl' -exec cp {} consolidated_wheels/ \;
+          ls -l consolidated_wheels
 
-      # create night release if it's a push to main
-      - if: github.event_name == 'push' && github.ref == 'refs/heads/main'
-        name: Nightly Release
-        uses: andelf/nightly-release@v1
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      - name: Upload consolidated wheels
+        uses: actions/upload-artifact@v4
         with:
-          tag_name: nightly
-          name: 'stable-fast Nightly Release $$'
-          prerelease: true
-          body: 'TODO: Add nightly release notes'
-          files: |
-            dist/*/*.whl
-
-      # create release if it's a tag like vx.y.z
-      - if: github.ref_type == 'tag' && startsWith(github.ref, 'refs/tags/v')
-        name: Release
-        uses: softprops/action-gh-release@v1
-        with :
-          files: |
-            dist/*/*.whl
+          name: built-wheels
+          path: consolidated_wheels
 
   # upload_pip:
   #   needs: build

.github/workflows/wheels_build.yml

@@ -17,6 +17,11 @@ on:
         required: true
         type: string
         description: "Example: 117 for 11.7"
+      append_local_version:
+        required: false
+        type: string
+        default: '1'
+        description: "Set to '0' to disable SFAST local version suffix"
       cudnn_version_major:
         required: false
         default: '8'
@@ -38,6 +43,11 @@ on:
         required: true
         type: string
         description: "Example: 117 for 11.7"
+      append_local_version:
+        required: false
+        type: string
+        default: '1'
+        description: "Set to '0' to disable SFAST local version suffix"
       cudnn_version_major:
         required: false
         default: '8'
@@ -54,7 +64,7 @@ env:
   TORCH_CUDA_ARCH_LIST: "6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX"
   MAX_JOBS: 2
   DISTUTILS_USE_SDK: 1 # otherwise distutils will complain on windows about multiple versions of msvc
-  SFAST_APPEND_VERSION: 1
+  SFAST_APPEND_VERSION: ${{ inputs.append_local_version }}
   TWINE_USERNAME: __token__
 
 jobs:
@@ -81,12 +91,14 @@ jobs:
           cushort = "${{ inputs.cuda_short_version }}"
           # https://github.com/Jimver/cuda-toolkit/blob/master/src/links/linux-links.ts
           full_version, install_script = {
+            "128": ("12.8.0", "https://developer.download.nvidia.com/compute/cuda/12.8.0/local_installers/cuda_12.8.0_570.86.10_linux.run"),
             "126": ("12.6.0", "https://developer.download.nvidia.com/compute/cuda/12.6.0/local_installers/cuda_12.6.0_560.28.03_linux.run"),
             "118": ("11.8.0", "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"),
             "117": ("11.7.1", "https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda_11.7.1_515.65.01_linux.run"),
             "116": ("11.6.2", "https://developer.download.nvidia.com/compute/cuda/11.6.2/local_installers/cuda_11.6.2_510.47.03_linux.run"),
           }[cushort]
           cudnn_pypi_package = {
+            "128": "nvidia-cudnn-cu12",
             "126": "nvidia-cudnn-cu12",
             "121": "nvidia-cudnn-cu12",
             "118": "nvidia-cudnn-cu11",
@@ -184,17 +196,9 @@ jobs:
           set -Eeuo pipefail
           git config --global --add safe.directory "*"
           version=`cat version.txt`
-          torch_version_suffix=torch$(echo ${{ inputs.torch_version }} | sed 's/\.//g')
-          cuda_version_suffix=${{ steps.cuda_info.outputs.CUDA_VERSION_SUFFIX }}
-          nightly_tag=$([[ ${VERSION_SOURCE} == 'tag' ]] && echo '' || echo '.dev'`date +%Y%m%d`)
-
-          if [[ "${{ inputs.cuda_short_version }}" == "126" ]]; then
-            echo "BUILD_VERSION=${version}" >> ${GITHUB_ENV}
-            echo "BUILD_VERSION=${version}" >> ${GITHUB_OUTPUT}
-          else
-            echo "BUILD_VERSION=${version}+${torch_version_suffix}${cuda_version_suffix}" >> ${GITHUB_ENV}
-            echo "BUILD_VERSION=${version}+${torch_version_suffix}${cuda_version_suffix}" >> ${GITHUB_OUTPUT}
-          fi
+          echo "BUILD_VERSION=${version}" >> ${GITHUB_ENV}
+          echo "BUILD_VERSION=${version}" >> ${GITHUB_OUTPUT}
+
       - run: echo "sfast-${BUILD_VERSION}"
       - run: echo "release version"
         if: ${{ !contains(steps.sfast_version.outputs.BUILD_VERSION, '.dev') }}
@@ -218,7 +222,7 @@ jobs:
           CUDNN_PYPI_PACKAGE: ${{ steps.cuda_info.outputs.CUDNN_PYPI_PACKAGE }}
         run: |
           cudnn_next_version_major=$((${CUDNN_VERSION_MAJOR} + 1))
-          cudnn_package_name="${CUDNN_PYPI_PACKAGE}>=${CUDNN_VERSION_MAJOR}.0.0.0,<9.6.0.0"
+          cudnn_package_name="${CUDNN_PYPI_PACKAGE}>=${CUDNN_VERSION_MAJOR}.0.0.0,<=9.10.2.21"
           $PY -m pip install --upgrade pip
           $PY -m pip install wheel setuptools ninja twine "torch==${{ inputs.torch_version }}" "${cudnn_package_name}" -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu${{ inputs.cuda_short_version }} --no-cache-dir
 
@@ -229,8 +233,18 @@ jobs:
           PLAT_ARG: ${{ contains(inputs.os, 'ubuntu') && '--plat-name manylinux2014_x86_64' || '' }}
 
       - run: du -h dist/*
-      - uses: actions/upload-artifact@v4
+
+      - name: Upload artifact (local build)
+        if: ${{ inputs.append_local_version != '0' }}
+        uses: actions/upload-artifact@v4
         with:
           name: ${{ inputs.os }}-py${{ inputs.python }}-torch${{ inputs.torch_version }}+cu${{ inputs.cuda_short_version }}
           path: dist/*.whl
+
+      - name: Upload artifact (pypi build)
+        if: ${{ inputs.append_local_version == '0' }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ inputs.os }}-py${{ inputs.python }}
+          path: dist/*.whl
 # Note: it might be helpful to have additional steps that test if the built wheels actually work