EffectiveKernels/setup.py at main · Kwai-Keye/EffectiveKernels · 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
81
82
83
84
85
86
87
88
89
90
91
92
"""
Setup script for effective_kernels package.

Uses PyTorch's native cpp_extension for building CUDA extensions.
"""
from setuptools import setup, find_packages
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os

with open("effective_kernels/__init__.py", "r") as f:
    for line in f:
        if line.startswith("__version__"):
            version = line.split('"')[1]
            break

cuda_arch = os.environ.get("TORCH_CUDA_ARCH_LIST", "9.0a")

extra_cuda_cflags = [
    "-O3",
    "--use_fast_math",
    "--expt-relaxed-constexpr",
    "--expt-extended-lambda",
]

sources = [
    "csrc/bindings.cpp",
    "csrc/indexing/instantiations/topk_128_block_64_unique.cu",
    "csrc/indexing/instantiations/topk_2048_block_16_unique.cu",
    "csrc/indexing/instantiations/topk_128_index_sort.cu",
    "csrc/indexing/instantiations/topk_2048_index_sort.cu",
]

ext_modules = [
    CUDAExtension(
        name="effective_kernels._C",
        sources=sources,
        extra_compile_args={
            "cxx": ["-O3", "-Wall", "-fPIC"],
            "nvcc": extra_cuda_cflags,
        },
        include_dirs=["csrc"],
        define_macros=[],
    ),
]

aot_enabled = os.environ.get("EFFECTIVE_KERNELS_AOT") == "1"

if aot_enabled:
    print("Building AOT kernels...")
    import sys
    import importlib.util
    from pathlib import Path
    _setup_dir = Path(__file__).parent.resolve()
    sys.path.insert(0, str(_setup_dir))
    from effective_kernels.ops.aot import export_all_aot_kernels
    so_paths = export_all_aot_kernels()
    for p in so_paths:
        print(f"  Built: {p}")

def _get_packages():
    all_packages = find_packages()
    all_packages.append("effective_kernels.aot_kernels")
    if aot_enabled:
        all_packages = [p for p in all_packages if not p.startswith("effective_kernels.kernel")]
    return all_packages

setup(
    name="effective_kernels",
    version=version,
    description="Effective CUDA kernels for machine learning",
    author="Effective Kernels Team",
    packages=_get_packages(),
    package_data={"effective_kernels.aot_kernels": ["*.so"]} if aot_enabled else {},
    ext_modules=ext_modules,
    cmdclass={"build_ext": BuildExtension},
    python_requires=">=3.8",
    classifiers=[
        "Development Status :: 3 - Alpha",
        "Intended Audience :: Developers",
        "Intended Audience :: Science/Research",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "Programming Language :: C++",
        "Programming Language :: CUDA",
        "Topic :: Scientific/Engineering",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
    zip_safe=False,
)