setup_rocm.py 2.72 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 platform
17
18
19
import sys
from pathlib import Path

20
import torch
21
22
23
24
from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

root = Path(__file__).parent.resolve()
25
arch = platform.machine().lower()
26
27
28
29
30
31
32
33
34


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


35
operator_namespace = "sgl_kernel"
36
include_dirs = [
37
38
    root / "include",
    root / "csrc",
39
40
41
]

sources = [
42
43
    "csrc/allreduce/custom_all_reduce.hip",
    "csrc/moe/moe_align_kernel.cu",
44
    "csrc/moe/moe_topk_softmax_kernels.cu",
45
    "csrc/torch_extension_rocm.cc",
46
    "csrc/speculative/eagle_utils.cu",
47
48
49
50
]

cxx_flags = ["-O3"]
libraries = ["hiprtc", "amdhip64", "c10", "torch", "torch_python"]
51
extra_link_args = ["-Wl,-rpath,$ORIGIN/../../torch/lib", f"-L/usr/lib/{arch}-linux-gnu"]
52

53
54
55
56
57
58
59
amdgpu_target = torch.cuda.get_device_properties("cuda").gcnArchName.split(":")[0]
if amdgpu_target not in ["gfx942", "gfx950"]:
    print(
        f"Warning: Unsupported GPU architecture detected '{amdgpu_target}'. Expected 'gfx942' or 'gfx950'."
    )
    sys.exit(1)

60
61
62
63
64
65
66
67
hipcc_flags = [
    "-DNDEBUG",
    f"-DOPERATOR_NAMESPACE={operator_namespace}",
    "-O3",
    "-Xcompiler",
    "-fPIC",
    "-std=c++17",
    "-D__HIP_PLATFORM_AMD__=1",
68
    f"--amdgpu-target={amdgpu_target}",
69
70
71
72
    "-DENABLE_BF16",
    "-DENABLE_FP8",
]

73
74
75
76
77
78
79
80
81
82
83
ext_modules = [
    CUDAExtension(
        name="sgl_kernel.common_ops",
        sources=sources,
        include_dirs=include_dirs,
        extra_compile_args={
            "nvcc": hipcc_flags,
            "cxx": cxx_flags,
        },
        libraries=libraries,
        extra_link_args=extra_link_args,
84
        py_limited_api=False,
85
86
87
    ),
]

88
89
90
setup(
    name="sgl-kernel",
    version=_get_version(),
91
    packages=find_packages(where="python"),
92
93
    package_dir={"": "python"},
    ext_modules=ext_modules,
94
    cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},
95
96
    options={"bdist_wheel": {"py_limited_api": "cp39"}},
)