Update to compile with torch 2.4.1

.github/workflows/build.yaml

@@ -40,8 +40,8 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-20.04, windows-latest]
-        pyver: ["3.8", "3.9", "3.10", "3.11"]
-        cuda: ["11.8.0", "12.1.1"]
+        pyver: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        cuda: ["12.4.1"]
     defaults:
       run:
         shell: pwsh
@@ -91,7 +91,7 @@ jobs:
           # Install torch
           $cudaVersion = $env:CUDA_VERSION.Replace('.', '')
           $cudaVersionPytorch = $cudaVersion.Substring(0, $cudaVersion.Length - 1)
-          $pytorchVersion = "torch==2.3.1"
+          $pytorchVersion = "torch==2.4.1"
           python -m pip install --upgrade --no-cache-dir $pytorchVersion+cu$cudaVersionPytorch --index-url https://download.pytorch.org/whl/cu$cudaVersionPytorch
           python -m pip install build setuptools wheel ninja
 
@@ -119,114 +119,3 @@ jobs:
         with:
           upload_url: ${{ needs.release.outputs.upload_url }}
           asset_path: ./dist/*.whl
-
-  build_rocm_wheels:
-    name: Build AWQ with ROCm
-    runs-on: ${{ matrix.os }}
-    needs: release
-
-    strategy:
-      matrix:
-        os: [ubuntu-20.04]
-        python: ["3.8", "3.9", "3.10", "3.11"]
-        rocm: ["5.7.1"]
-    defaults:
-      run:
-        shell: bash
-    env:
-      ROCM_VERSION: ${{ matrix.rocm }}
-
-    steps:
-      - uses: actions/checkout@v3
-
-      - name: Free Disk Space
-        run: |
-          df -h
-          echo "Removing large packages"
-          sudo apt-get remove -y '^dotnet-.*'
-          sudo apt-get remove -y 'php.*'
-          sudo apt-get remove -y azure-cli google-chrome-stable firefox powershell mono-devel
-          df -h
-          sudo apt-get autoremove -y >/dev/null 2>&1
-          sudo apt-get clean
-          sudo apt-get autoremove -y >/dev/null 2>&1
-          sudo apt-get autoclean -y >/dev/null 2>&1
-          df -h
-          echo "https://github.com/actions/virtual-environments/issues/709"
-          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
-          df -h
-          echo "remove big /usr/local"
-          sudo rm -rf "/usr/local/share/boost"
-          sudo rm -rf /usr/local/lib/android >/dev/null 2>&1
-          df -h
-          sudo rm -rf /usr/share/dotnet/sdk > /dev/null 2>&1
-          sudo rm -rf /usr/share/dotnet/shared > /dev/null 2>&1
-          sudo rm -rf /usr/share/swift > /dev/null 2>&1
-          df -h
-
-      - uses: actions/setup-python@v3
-        with:
-          python-version: ${{ matrix.python }}
-
-      - name: Setup Mamba
-        uses: conda-incubator/setup-miniconda@v2.2.0
-        with:
-          activate-environment: "build"
-          python-version: ${{ matrix.python }}
-          mamba-version: "*"
-          use-mamba: false
-          channels: conda-forge,defaults
-          channel-priority: true
-          add-pip-as-python-dependency: true
-          auto-activate-base: false
-
-      - name: Set up ROCm
-        run: |
-          echo "Using python:"
-          python --version
-          which python
-
-          if [[ "${{ matrix.rocm }}" == "5.4.2" ]]; then
-            export ROCM_DL_FILE=amdgpu-install_5.4.50402-1_all.deb
-          elif [[ "${{ matrix.rocm }}" == "5.6.1" ]]; then
-            export ROCM_DL_FILE=amdgpu-install_5.6.50601-1_all.deb
-          elif [[ "${{ matrix.rocm }}" == "5.7.1" ]]; then
-            export ROCM_DL_FILE=amdgpu-install_5.7.50701-1_all.deb
-          else
-            echo Unknown rocm version
-            exit 1
-          fi
-
-          curl -O https://repo.radeon.com/amdgpu-install/${{ matrix.rocm }}/ubuntu/focal/$ROCM_DL_FILE
-          sudo dpkg -i $ROCM_DL_FILE
-          sudo DEBIAN_FRONTEND=noninteractive amdgpu-install --usecase=rocm --no-dkms --no-32 -y
-
-      - name: Install Dependencies
-        run: |
-          sudo apt-get update
-          sudo apt-get install -y --no-install-recommends rocsparse-dev rocthrust-dev rocblas-dev hipblas-dev hipsparse-dev
-
-          python -m pip install --upgrade build setuptools wheel
-
-          if [[ "${{ matrix.rocm }}" == "5.7.1" ]]; then
-            python -m pip install torch==2.3.1 --index-url https://download.pytorch.org/whl/rocm5.7
-          elif [[ "${{ matrix.rocm }}" == "5.6.1" ]]; then
-            python -m pip install torch==2.3.1 --index-url https://download.pytorch.org/whl/rocm5.6
-          else
-            echo Unknown rocm version for python install
-            exit 1
-          fi
-
-      - name: Build Wheel
-        run: |
-          echo "Using python for build:"
-          python --version
-          which python
-
-          ROCM_VERSION=${{ matrix.rocm }} python setup.py sdist bdist_wheel
-
-      - name: Upload Assets
-        uses: shogo82148/actions-upload-release-asset@v1
-        with:
-          upload_url: ${{ needs.release.outputs.upload_url }}
-          asset_path: ./dist/*.whl

README.md

@@ -10,26 +10,18 @@ AutoAWQ Kernels is a new package that is split up from the [main repository](htt
   - GPU: Must be compute capability 7.5 or higher.
   - CUDA Toolkit: Must be 11.8 or higher.
 - AMD:
-  - ROCm: Must be 5.6 or higher.
+  - ROCm: Must be 5.6 or higher. [Build from source](#build-from-source)
 
 ## Install
 
 ### Install from PyPi
 
-The package is available on PyPi with CUDA 12.1.1 wheels:
+The package is available on PyPi with CUDA 12.4.1 wheels:
 
 ```
 pip install autoawq-kernels
 ```
 
-### Install release wheels
-
-For ROCm and other CUDA versions, you can use the wheels published at each [release](https://github.com/casper-hansen/AutoAWQ_kernels/releases/):
-
-```
-pip install https://github.com/casper-hansen/AutoAWQ_kernels/releases/download/v0.0.2/autoawq_kernels-0.0.2+rocm561-cp310-cp310-linux_x86_64.whl
-```
-
 ### Build from source
 You can also build from source:
 

scripts/download_wheels.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
 # Set variables
-AWQ_KERNELS_VERSION="0.0.7"
+AWQ_KERNELS_VERSION="0.0.8"
 RELEASE_URL="https://api.github.com/repos/casper-hansen/AutoAWQ_kernels/releases/tags/v${AWQ_KERNELS_VERSION}"
 
 # Create a directory to download the wheels

setup.py

@@ -7,7 +7,7 @@
 
 os.environ["CC"] = "g++"
 os.environ["CXX"] = "g++"
-AUTOAWQ_KERNELS_VERSION = "0.0.7"
+AUTOAWQ_KERNELS_VERSION = "0.0.8"
 PYPI_BUILD = os.getenv("PYPI_BUILD", "0") == "1"
 CUDA_VERSION = os.getenv("CUDA_VERSION", None) or torch.version.cuda
 ROCM_VERSION = os.environ.get("ROCM_VERSION", None) or torch.version.hip
@@ -51,12 +51,13 @@
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
         "Programming Language :: C++",
     ],
 }
 
 requirements = [
-    "torch==2.3.1",
+    "torch==2.4.1",
 ]