fq-diffusion/setup.py at main · ag2718/fq-diffusion · GitHub

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
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
from setuptools import setup
import torch.utils.cpp_extension as torch_cpp_ext
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os
import pathlib

setup_dir = os.path.dirname(os.path.realpath(__file__))
HERE = pathlib.Path(__file__).absolute().parent


def remove_unwanted_pytorch_nvcc_flags():
    REMOVE_NVCC_FLAGS = [
        "-D__CUDA_NO_HALF_OPERATORS__",
        "-D__CUDA_NO_HALF_CONVERSIONS__",
        "-D__CUDA_NO_BFLOAT16_CONVERSIONS__",
        "-D__CUDA_NO_HALF2_OPERATORS__",
    ]
    for flag in REMOVE_NVCC_FLAGS:
        try:
            torch_cpp_ext.COMMON_NVCC_FLAGS.remove(flag)
        except ValueError:
            pass


def get_cuda_arch_flags():
    return [
        "-gencode",
        "arch=compute_75,code=sm_75",  # Turing
        "-gencode",
        "arch=compute_80,code=sm_80",  # Ampere
        "-gencode",
        "arch=compute_86,code=sm_86",  # Ampere
    ]


def third_party_cmake():
    import subprocess, sys, shutil

    cmake = shutil.which("cmake")
    if cmake is None:
        raise RuntimeError("Cannot find CMake executable.")

    retcode = subprocess.call([cmake, HERE])
    if retcode != 0:
        sys.stderr.write("Error: CMake configuration failed.\n")
        sys.exit(1)

    # install fast hadamard transform
    hadamard_dir = os.path.join(HERE, "third-party/fast-hadamard-transform")
    pip = shutil.which("pip")
    retcode = subprocess.call([pip, "install", "-e", hadamard_dir])


if __name__ == "__main__":
    third_party_cmake()
    remove_unwanted_pytorch_nvcc_flags()
    setup(
        name="deploy",
        ext_modules=[
            CUDAExtension(
                name="deploy._CUDA",
                sources=[
                    "deploy/kernels/bindings.cpp",
                    "deploy/kernels/gemm.cu",
                    "deploy/kernels/quant.cu",
                    "deploy/kernels/flashinfer.cu",
                ],
                include_dirs=[
                    os.path.join(setup_dir, "deploy/kernels/include"),
                    os.path.join(setup_dir, "third-party/cutlass/include"),
                    os.path.join(setup_dir, "third-party/cutlass/tools/util/include"),
                ],
                extra_compile_args={
                    "cxx": [],
                    "nvcc": get_cuda_arch_flags(),
                },
            )
        ],
        cmdclass={"build_ext": BuildExtension},
    )