Update build process for flexibility and torch version compatibility (#29)

* add the support for Jetson Orin devices whose compute capability is 87

* don't force CC and CXX in env and be more flexible with torch version required

* update torch version requirement to be even more flexible

* update setup.py to support multiple compute capabilities

* impove compute_capabilities set definition

* Add support for Python 3.12 in classifiers tags

* (first draft) update readme with new building procedure

* Add install notes

---------

Co-authored-by: Shi Hui <[email protected]>
Co-authored-by: Casper <[email protected]>

Author: 3 people

README.md

@@ -23,12 +23,25 @@ pip install autoawq-kernels
 ```
 
 ### Build from source
-You can also build from source:
+
+To build the kernels from source, you first need to setup an environment containing the necessary dependencies.
+
+#### Build Requirements
+
+- Python>=3.8.0
+- Numpy
+- Wheel
+- PyTorch
+- ROCm: You need to install the following packages `rocsparse-dev hipsparse-dev rocthrust-dev rocblas-dev hipblas-dev`.
+
+#### Building process
 
 ```
-git clone https://github.com/casper-hansen/AutoAWQ_kernels
-cd AutoAWQ_kernels
-pip install -e .
+pip install git+https://github.com/casper-hansen/AutoAWQ_kernels.git
 ```
 
-To build for ROCm, you need to first install the following packages `rocsparse-dev hipsparse-dev rocthrust-dev rocblas-dev hipblas-dev`.
+Notes on environment variables:
+- `TORCH_VERSION`: By default, we build using the current version of torch by `torch.__version__`. You can override it with `TORCH_VERSION`.
+    - `CUDA_VERSION` or `ROCM_VERSION` can also be used to build for a specific version of CUDA or ROCm.
+- `CC` and `CXX`: You can specify which build system to use for the C code, e.g. `CC=g++-13 CXX=g++-13 pip install -e .`
+- `COMPUTE_CAPABILITIES`: You can specify specific compute capabilities to compile for: `COMPUTE_CAPABILITIES="75,80,86,87,89,90" pip install -e .`

setup.py

@@ -5,10 +5,14 @@
 from distutils.sysconfig import get_python_lib
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
-os.environ["CC"] = "g++"
-os.environ["CXX"] = "g++"
+if "CC" not in os.environ:
+    os.environ["CC"] = "g++"
+if "CXX" not in os.environ:
+    os.environ["CXX"] = "g++"
 AUTOAWQ_KERNELS_VERSION = "0.0.8"
 PYPI_BUILD = os.getenv("PYPI_BUILD", "0") == "1"
+COMPUTE_CAPABILITIES = os.getenv("COMPUTE_CAPABILITIES", "75,80,86,87,89,90")
+TORCH_VERSION = str(os.getenv("TORCH_VERSION", None) or torch.__version__).split('+', maxsplit=1)[0]
 CUDA_VERSION = os.getenv("CUDA_VERSION", None) or torch.version.cuda
 ROCM_VERSION = os.environ.get("ROCM_VERSION", None) or torch.version.hip
 
@@ -57,7 +61,7 @@
 }
 
 requirements = [
-    "torch==2.4.1",
+    f"torch>={TORCH_VERSION}",
 ]
 
 
@@ -91,7 +95,7 @@ def get_generator_flag():
 
 
 def get_compute_capabilities(
-    compute_capabilities={75, 80, 86, 89, 90}
+    compute_capabilities=set(map(int, COMPUTE_CAPABILITIES.split(",")))
 ):
     capability_flags = []