setup.py 2.79 KB
Newer Older
1
2
from pathlib import Path

3
import torch
lukec's avatar
lukec committed
4
from setuptools import find_packages, setup
5
6
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

7
8
9
10
11
12
13
14
15
16
root = Path(__file__).parent.resolve()


def get_version():
    with open(root / "pyproject.toml") as f:
        for line in f:
            if line.startswith("version"):
                return line.split("=")[1].strip().strip('"')


Yineng Zhang's avatar
Yineng Zhang committed
17
18
def update_wheel_platform_tag():
    wheel_dir = Path("dist")
19
20
21
22
23
24
    if wheel_dir.exists() and wheel_dir.is_dir():
        old_wheel = next(wheel_dir.glob("*.whl"))
        new_wheel = wheel_dir / old_wheel.name.replace(
            "linux_x86_64", "manylinux2014_x86_64"
        )
        old_wheel.rename(new_wheel)
Yineng Zhang's avatar
Yineng Zhang committed
25
26


27
cutlass = root / "3rdparty" / "cutlass"
28
flashinfer = root / "3rdparty" / "flashinfer"
29
30
31
include_dirs = [
    cutlass.resolve() / "include",
    cutlass.resolve() / "tools" / "util" / "include",
Ke Bao's avatar
Ke Bao committed
32
    root / "src" / "sgl-kernel" / "csrc",
33
34
    flashinfer.resolve() / "include",
    flashinfer.resolve() / "csrc",
35
]
Ke Bao's avatar
Ke Bao committed
36
nvcc_flags = [
37
    "-DNDEBUG",
Ke Bao's avatar
Ke Bao committed
38
39
40
41
42
43
44
    "-O3",
    "-Xcompiler",
    "-fPIC",
    "-gencode=arch=compute_75,code=sm_75",
    "-gencode=arch=compute_80,code=sm_80",
    "-gencode=arch=compute_89,code=sm_89",
    "-gencode=arch=compute_90,code=sm_90",
45
    "-gencode=arch=compute_90a,code=sm_90a",
46
47
48
49
    "-std=c++17",
    "-use_fast_math",
    "-DFLASHINFER_ENABLE_F16",
    "-DFLASHINFER_ENABLE_BF16",
Ke Bao's avatar
Ke Bao committed
50
]
51
52
53
54
55
56
57
58
59
60
for flag in [
    "-D__CUDA_NO_HALF_OPERATORS__",
    "-D__CUDA_NO_HALF_CONVERSIONS__",
    "-D__CUDA_NO_BFLOAT16_CONVERSIONS__",
    "-D__CUDA_NO_HALF2_OPERATORS__",
]:
    try:
        torch.utils.cpp_extension.COMMON_NVCC_FLAGS.remove(flag)
    except ValueError:
        pass
Ke Bao's avatar
Ke Bao committed
61
cxx_flags = ["-O3"]
62
libraries = ["c10", "torch", "torch_python", "cuda"]
63
extra_link_args = ["-Wl,-rpath,$ORIGIN/../../torch/lib", "-L/usr/lib/x86_64-linux-gnu"]
Ke Bao's avatar
Ke Bao committed
64
65
66
67
68
69
70
ext_modules = [
    CUDAExtension(
        name="sgl_kernel.ops._kernels",
        sources=[
            "src/sgl-kernel/csrc/trt_reduce_internal.cu",
            "src/sgl-kernel/csrc/trt_reduce_kernel.cu",
            "src/sgl-kernel/csrc/moe_align_kernel.cu",
Ke Bao's avatar
Ke Bao committed
71
            "src/sgl-kernel/csrc/int8_gemm_kernel.cu",
72
            "src/sgl-kernel/csrc/sampling_scaling_penalties.cu",
Ke Bao's avatar
Ke Bao committed
73
            "src/sgl-kernel/csrc/sgl_kernel_ops.cu",
74
            "src/sgl-kernel/csrc/rotary_embedding.cu",
75
            "src/sgl-kernel/csrc/norm.cu",
Ke Bao's avatar
Ke Bao committed
76
        ],
77
        include_dirs=include_dirs,
Ke Bao's avatar
Ke Bao committed
78
79
80
81
82
83
84
85
86
        extra_compile_args={
            "nvcc": nvcc_flags,
            "cxx": cxx_flags,
        },
        libraries=libraries,
        extra_link_args=extra_link_args,
    ),
]

87
88
setup(
    name="sgl-kernel",
89
    version=get_version(),
lukec's avatar
lukec committed
90
    packages=find_packages(),
91
    package_dir={"": "src"},
Ke Bao's avatar
Ke Bao committed
92
    ext_modules=ext_modules,
93
94
95
    cmdclass={"build_ext": BuildExtension},
    install_requires=["torch"],
)
Yineng Zhang's avatar
Yineng Zhang committed
96
97

update_wheel_platform_tag()