-
Notifications
You must be signed in to change notification settings - Fork 7
/
Copy pathsetup.py
27 lines (26 loc) · 1.08 KB
/
setup.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
from setuptools import setup, find_packages
from torch.utils import cpp_extension
setup(
name='autosmoothquant',
ext_modules=[
cpp_extension.CUDAExtension(
name='autosmoothquant._CUDA',
sources=[
'csrc/int8gemm/cuda_utils.cc',
'csrc/int8gemm/cublasINT8MMWrapper.cc',
'csrc/int8gemm/cublasAlgoMap.cc',
'csrc/int8gemm/bindings.cpp',
],
extra_link_args=['-lcublas_static', '-lcublasLt_static',
'-lculibos', '-lcudart', '-lcudart_static',
'-lrt', '-lpthread', '-ldl', '-L/usr/lib/x86_64-linux-gnu/'],
extra_compile_args={'cxx': ['-std=c++17', '-O3'],
'nvcc': ['-O3', '-std=c++17', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__']},
),
],
cmdclass={
'build_ext': cpp_extension.BuildExtension.with_options(use_ninja=False)
},
packages=find_packages(
exclude=['notebook', 'scripts', 'tests']),
)