61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
import os
|
|
import torch
|
|
from pathlib import Path
|
|
from setuptools import setup, find_packages
|
|
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
|
|
|
cwd = Path(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
nvcc_flags = [
|
|
"-std=c++17", # NOTE: CUTLASS requires c++17
|
|
"-DENABLE_BF16", # Enable BF16 for cuda_version >= 11
|
|
]
|
|
|
|
env_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST", None)
|
|
|
|
if env_arch_list:
|
|
# Let PyTorch builder to choose device to target for.
|
|
device_capability = ""
|
|
else:
|
|
device_capability = torch.cuda.get_device_capability()
|
|
device_capability = f"{device_capability[0]}{device_capability[1]}"
|
|
|
|
if device_capability:
|
|
nvcc_flags.extend(
|
|
[
|
|
f"--generate-code=arch=compute_{device_capability},code=sm_{device_capability}",
|
|
f"-DGROUPED_GEMM_DEVICE_CAPABILITY={device_capability}",
|
|
]
|
|
)
|
|
|
|
ext_modules = [
|
|
CUDAExtension(
|
|
"wallx_csrc",
|
|
[
|
|
"csrc/ops.cu",
|
|
"csrc/dual_asym_grouped_gemm.cu",
|
|
"csrc/permute.cu",
|
|
"csrc/rope.cu",
|
|
],
|
|
include_dirs=[f"{cwd}/3rdparty/cutlass/include/", f"{cwd}/csrc"],
|
|
extra_compile_args={
|
|
"cxx": ["-fopenmp", "-fPIC", "-Wno-strict-aliasing"],
|
|
"nvcc": nvcc_flags,
|
|
},
|
|
)
|
|
]
|
|
|
|
setup(
|
|
name="wall_x",
|
|
version="1.0.0",
|
|
author="X2Robot Team",
|
|
classifiers=[
|
|
"Programming Language :: Python :: 3",
|
|
"License :: OSI Approved :: BSD License",
|
|
"Operating System :: Unix",
|
|
],
|
|
packages=find_packages(),
|
|
ext_modules=ext_modules,
|
|
cmdclass={"build_ext": BuildExtension},
|
|
)
|