From be0f9e01d9d45ff5a78a692d325e4035f4540716 Mon Sep 17 00:00:00 2001
From: Casper <casperbh.96@gmail.com>
Date: Thu, 15 Feb 2024 21:48:51 +0100
Subject: [PATCH] Fix workflow (#345)

---
 .github/workflows/build.yaml |  2 +-
 scripts/download_wheels.sh   |  2 +-
 setup.py                     | 20 +++++++++++++++++---
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index 6b35f802..eae16a25 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -113,7 +113,7 @@ jobs:
           }
 
           python setup.py sdist bdist_wheel
-      
+
       - name: Upload Assets
         uses: shogo82148/actions-upload-release-asset@v1
         with:
diff --git a/scripts/download_wheels.sh b/scripts/download_wheels.sh
index 432a0b2b..9df78268 100644
--- a/scripts/download_wheels.sh
+++ b/scripts/download_wheels.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 
 # Set variables
-AWQ_VERSION="0.1.6"
+AWQ_VERSION="0.2.0"
 RELEASE_URL="https://api.github.com/repos/casper-hansen/AutoAWQ/releases/tags/v${AWQ_VERSION}"
 
 # Create a directory to download the wheels
diff --git a/setup.py b/setup.py
index 0e702a01..7307ecc2 100644
--- a/setup.py
+++ b/setup.py
@@ -2,9 +2,9 @@
 import torch
 import platform
 import requests
-import importlib_metadata
 from pathlib import Path
 from setuptools import setup, find_packages
+from torch.utils.cpp_extension import CUDAExtension
 
 
 def get_latest_kernels_version(repo):
@@ -88,15 +88,20 @@ def get_kernels_whl_url(
     "torch>=2.0.1",
     "transformers>=4.35.0",
     "tokenizers>=0.12.1",
+    "typing_extensions>=4.8.0"
     "accelerate",
     "datasets",
     "zstandard",
 ]
 
 try:
-    importlib_metadata.version("autoawq-kernels")
+    if ROCM_VERSION:
+        import exlv2_ext
+    else:
+        import awq_ext
+
     KERNELS_INSTALLED = True
-except importlib_metadata.PackageNotFoundError:
+except ImportError:
     KERNELS_INSTALLED = False
 
 # kernels can be downloaded from pypi for cuda+121 only
@@ -133,5 +138,14 @@ def get_kernels_whl_url(
         "eval": ["lm_eval>=0.4.0", "tabulate", "protobuf", "evaluate", "scipy"],
         "dev": ["black", "mkdocstrings-python", "mkdocs-material", "griffe-typingdoc"]
     },
+    # NOTE: We create an empty CUDAExtension because torch helps us with
+    # creating the right boilerplate to enable correct targeting of
+    # the autoawq-kernels package
+    ext_modules=[
+        CUDAExtension(
+            name="__build_artifact_for_awq_kernel_targeting",
+            sources=[],
+        )
+    ],
     **common_setup_kwargs,
 )