setup.py 7.67 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Copyright 2025 SGLang Team. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

16
import os
17
import shutil
18
import sys
19
20
from pathlib import Path

21
import torch
lukec's avatar
lukec committed
22
from setuptools import find_packages, setup
23
from setuptools.command.build_py import build_py
24
25
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

26
27
28
root = Path(__file__).parent.resolve()


29
30
if "bdist_wheel" in sys.argv and "--plat-name" not in sys.argv:
    sys.argv.extend(["--plat-name", "manylinux2014_x86_64"])
Yineng Zhang's avatar
Yineng Zhang committed
31
32


33
def _get_cuda_version():
34
35
36
37
38
    if torch.version.cuda:
        return tuple(map(int, torch.version.cuda.split(".")))
    return (0, 0)


39
def _get_device_sm():
40
41
42
43
44
45
    if torch.cuda.is_available():
        major, minor = torch.cuda.get_device_capability()
        return major * 10 + minor
    return 0


46
47
48
49
50
51
def _get_version():
    with open(root / "pyproject.toml") as f:
        for line in f:
            if line.startswith("version"):
                return line.split("=")[1].strip().strip('"')

52

53
operator_namespace = "sgl_kernel"
54
55
cutlass_default = root / "3rdparty" / "cutlass"
cutlass = Path(os.environ.get("CUSTOM_CUTLASS_SRC_DIR", default=cutlass_default))
56
flashinfer = root / "3rdparty" / "flashinfer"
57
deepgemm = root / "3rdparty" / "deepgemm"
58
include_dirs = [
59
60
    root / "include",
    root / "csrc",
61
62
    cutlass.resolve() / "include",
    cutlass.resolve() / "tools" / "util" / "include",
63
    flashinfer.resolve() / "include",
64
    flashinfer.resolve() / "include" / "gemm",
65
    flashinfer.resolve() / "csrc",
66
    "cublas",
67
]
68

69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

class CustomBuildPy(build_py):
    def run(self):
        self.copy_deepgemm_to_build_lib()
        self.make_jit_include_symlinks()
        build_py.run(self)

    def make_jit_include_symlinks(self):
        # Make symbolic links of third-party include directories
        build_include_dir = os.path.join(self.build_lib, "deep_gemm/include")
        os.makedirs(build_include_dir, exist_ok=True)

        third_party_include_dirs = [
            cutlass.resolve() / "include" / "cute",
            cutlass.resolve() / "include" / "cutlass",
        ]

        for d in third_party_include_dirs:
            dirname = str(d).split("/")[-1]
            src_dir = d
            dst_dir = f"{build_include_dir}/{dirname}"
            assert os.path.exists(src_dir)
            if os.path.exists(dst_dir):
                assert os.path.islink(dst_dir)
                os.unlink(dst_dir)
            os.symlink(src_dir, dst_dir, target_is_directory=True)

    def copy_deepgemm_to_build_lib(self):
        """
        This function copies DeepGemm to python's site-packages
        """
        dst_dir = os.path.join(self.build_lib, "deep_gemm")
        os.makedirs(dst_dir, exist_ok=True)

        # Copy deepgemm/deep_gemm to the build directory
        src_dir = os.path.join(str(deepgemm.resolve()), "deep_gemm")

        # Remove existing directory if it exists
        if os.path.exists(dst_dir):
            shutil.rmtree(dst_dir)

        # Copy the directory
        shutil.copytree(src_dir, dst_dir)


Ke Bao's avatar
Ke Bao committed
114
nvcc_flags = [
115
    "-DNDEBUG",
116
    f"-DOPERATOR_NAMESPACE={operator_namespace}",
Ke Bao's avatar
Ke Bao committed
117
118
119
120
121
122
123
    "-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",
124
125
    "-std=c++17",
    "-DFLASHINFER_ENABLE_F16",
126
    "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1",
127
128
129
130
131
132
    "-DCUTLASS_VERSIONS_GENERATED",
    "-DCUTE_USE_PACKED_TUPLE=1",
    "-DCUTLASS_TEST_LEVEL=0",
    "-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1",
    "-DCUTLASS_DEBUG_TRACE_LEVEL=0",
    "--ptxas-options=-v",
133
    "--expt-relaxed-constexpr",
134
135
    "-Xcompiler=-Wconversion",
    "-Xcompiler=-fno-strict-aliasing",
Ke Bao's avatar
Ke Bao committed
136
]
137
138
139
140
141
nvcc_flags_fp8 = [
    "-DFLASHINFER_ENABLE_FP8",
    "-DFLASHINFER_ENABLE_FP8_E4M3",
    "-DFLASHINFER_ENABLE_FP8_E5M2",
]
142

143
sources = [
144
145
146
    "csrc/allreduce/trt_reduce_internal.cu",
    "csrc/allreduce/trt_reduce_kernel.cu",
    "csrc/attention/lightning_attention_decode_kernel.cu",
147
    "csrc/elementwise/activation.cu",
148
    "csrc/elementwise/fused_add_rms_norm_kernel.cu",
149
150
    "csrc/elementwise/rope.cu",
    "csrc/gemm/bmm_fp8.cu",
151
    "csrc/gemm/cublas_grouped_gemm.cu",
152
    "csrc/gemm/awq_kernel.cu",
153
154
155
156
157
158
159
    "csrc/gemm/fp8_gemm_kernel.cu",
    "csrc/gemm/fp8_blockwise_gemm_kernel.cu",
    "csrc/gemm/int8_gemm_kernel.cu",
    "csrc/gemm/per_token_group_quant_fp8.cu",
    "csrc/gemm/per_token_quant_fp8.cu",
    "csrc/gemm/per_tensor_quant_fp8.cu",
    "csrc/moe/moe_align_kernel.cu",
160
    "csrc/moe/moe_topk_softmax_kernels.cu",
161
162
    "csrc/speculative/eagle_utils.cu",
    "csrc/speculative/speculative_sampling.cu",
163
    "csrc/speculative/packbit.cu",
164
    "csrc/torch_extension.cc",
165
166
    "3rdparty/flashinfer/csrc/norm.cu",
    "3rdparty/flashinfer/csrc/renorm.cu",
167
    "3rdparty/flashinfer/csrc/sampling.cu",
168
169
]

170
171
172
enable_bf16 = os.getenv("SGL_KERNEL_ENABLE_BF16", "0") == "1"
enable_fp8 = os.getenv("SGL_KERNEL_ENABLE_FP8", "0") == "1"
enable_sm90a = os.getenv("SGL_KERNEL_ENABLE_SM90A", "0") == "1"
173
enable_sm100a = os.getenv("SGL_KERNEL_ENABLE_SM100A", "0") == "1"
174
175
176
177
178
179
cuda_version = _get_cuda_version()
sm_version = _get_device_sm()

if torch.cuda.is_available():
    if cuda_version >= (12, 0) and sm_version >= 90:
        nvcc_flags.append("-gencode=arch=compute_90a,code=sm_90a")
180
181
182
    if cuda_version >= (12, 8) and sm_version >= 100:
        nvcc_flags.append("-gencode=arch=compute_100,code=sm_100")
        nvcc_flags.append("-gencode=arch=compute_100a,code=sm_100a")
Yineng Zhang's avatar
Yineng Zhang committed
183
184
    else:
        nvcc_flags.append("-use_fast_math")
185
    if sm_version >= 90:
186
        nvcc_flags.extend(nvcc_flags_fp8)
187
188
189
190
191
192
    if sm_version >= 80:
        nvcc_flags.append("-DFLASHINFER_ENABLE_BF16")
else:
    # compilation environment without GPU
    if enable_sm90a:
        nvcc_flags.append("-gencode=arch=compute_90a,code=sm_90a")
193
194
    if enable_sm100a:
        nvcc_flags.append("-gencode=arch=compute_100a,code=sm_100a")
Yineng Zhang's avatar
Yineng Zhang committed
195
196
    else:
        nvcc_flags.append("-use_fast_math")
197
    if enable_fp8:
198
        nvcc_flags.extend(nvcc_flags_fp8)
199
200
    if enable_bf16:
        nvcc_flags.append("-DFLASHINFER_ENABLE_BF16")
201

202
203
204
205
206
207
208
209
210
211
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
212

Ke Bao's avatar
Ke Bao committed
213
cxx_flags = ["-O3"]
Yineng Zhang's avatar
Yineng Zhang committed
214
libraries = ["c10", "torch", "torch_python", "cuda", "cublas"]
215
extra_link_args = ["-Wl,-rpath,$ORIGIN/../../torch/lib", "-L/usr/lib/x86_64-linux-gnu"]
216

Ke Bao's avatar
Ke Bao committed
217
218
ext_modules = [
    CUDAExtension(
219
        name="sgl_kernel.common_ops",
220
        sources=sources,
221
        include_dirs=include_dirs,
Ke Bao's avatar
Ke Bao committed
222
223
224
225
226
227
        extra_compile_args={
            "nvcc": nvcc_flags,
            "cxx": cxx_flags,
        },
        libraries=libraries,
        extra_link_args=extra_link_args,
228
        py_limited_api=True,
Ke Bao's avatar
Ke Bao committed
229
230
231
    ),
]

232
233
setup(
    name="sgl-kernel",
234
    version=_get_version(),
235
236
    packages=find_packages(where="python"),
    package_dir={"": "python"},
Ke Bao's avatar
Ke Bao committed
237
    ext_modules=ext_modules,
238
239
240
241
    cmdclass={
        "build_ext": BuildExtension.with_options(use_ninja=True),
        "build_py": CustomBuildPy,
    },
242
    options={"bdist_wheel": {"py_limited_api": "cp39"}},
243
)