vLLM-Kunlun/setup.py at main · fore1337/vLLM-Kunlun · 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
#
# setup.py for vllm_kunlun
#

import os
import shutil

from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

ROOT_DIR = os.path.dirname(__file__)

ext_modules = [
    CppExtension(
        name="vllm_kunlun._kunlun",
        sources=["vllm_kunlun/csrc/utils.cpp"],
        include_dirs=[
            "vllm_kunlun/csrc",
            "/usr/local/cuda/include",
        ],
        library_dirs=["/usr/local/cuda/lib64"],
        extra_compile_args=["-O3"],
    )
]


class CustomBuildExt(BuildExtension):
    def run(self):
        super().run()
        for ext in self.extensions:
            ext_path = self.get_ext_fullpath(ext.name)
            file_name = os.path.basename(ext_path)
            target_path = os.path.join("vllm_kunlun", file_name)

            if os.path.exists(target_path):
                os.remove(target_path)
            shutil.copyfile(ext_path, target_path)
            print(f"[BuildExt] Copied {ext_path} -> {target_path}")


if __name__ == "__main__":

    setup(
        name="vllm_kunlun",
        version="v1.0",
        author="vLLM-Kunlun team",
        license="Apache 2.0",
        description="vLLM Kunlun3 backend plugin",
        packages=find_packages(exclude=("docs", "examples", "tests*")),
        package_data={"vllm_kunlun": ["_kunlun.so", "so/*.so", "include/*.h"]},
        python_requires=">=3.10",
        ext_modules=ext_modules,
        cmdclass={
            "build_ext": CustomBuildExt,
        },
        entry_points={
            "vllm.platform_plugins": ["kunlun = vllm_kunlun:register"],
            "vllm.general_plugins": [
                "kunlun_model = vllm_kunlun:register_model",
                "kunlun_quant = vllm_kunlun:register_quant_method",
                "kunlun_tool_parser = vllm_kunlun:register_tool_parser",
                "kunlun_reasoning_parser = vllm_kunlun:register_reasoning_parser",
            ],
            "console_scripts": ["vllm_kunlun = vllm_kunlun.entrypoints.main:main"],
        },
    )