wehub-resource-sync
219 lines
8.8 KiB
Python
219 lines
8.8 KiB
Python
import ast
|
|
import re
|
|
import os
|
|
import subprocess
|
|
import setuptools
|
|
import importlib
|
|
|
|
from pathlib import Path
|
|
from setuptools.command.build_py import build_py
|
|
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
|
|
|
current_dir = os.path.dirname(os.path.realpath(__file__))
|
|
persistent_env_names = ('EP_JIT_CACHE_DIR', 'EP_JIT_PRINT_COMPILER_COMMAND', 'EP_NUM_TOPK_IDX_BITS', 'EP_NCCL_ROOT_DIR')
|
|
|
|
# Load discover module without triggering `deep_ep.__init__`
|
|
find_pkgs_spec = importlib.util.spec_from_file_location('find_pkgs', os.path.join(current_dir, 'deep_ep', 'utils', 'find_pkgs.py'))
|
|
find_pkgs = importlib.util.module_from_spec(find_pkgs_spec)
|
|
find_pkgs_spec.loader.exec_module(find_pkgs)
|
|
|
|
|
|
# Wheel specific: NVIDIA pip wheels (nvidia-nvshmem-cu12, nvidia-nccl-cu12)
|
|
# only ship the SO name of the host library, e.g. `libnvshmem_host.so.3`,
|
|
# without the unversioned `libnvshmem_host.so` symlink. So `-l:libnvshmem_host.so`
|
|
# (exact-name link) cannot resolve. Resolve the real file name at build time
|
|
# and pass it through to the linker instead.
|
|
def _find_versioned_so(base_dir, prefix):
|
|
"""Return the real filename of the first ``{prefix}.so*`` under ``base_dir/lib``.
|
|
|
|
Prefers an unversioned ``{prefix}.so`` symlink when present so we keep
|
|
behaving identically to the Tarball install. Falls back to the SONAME
|
|
file (``{prefix}.so.X``) shipped by pip wheels.
|
|
"""
|
|
lib_dir = Path(base_dir).joinpath('lib')
|
|
unversioned = lib_dir / f'{prefix}.so'
|
|
if unversioned.exists():
|
|
return unversioned.name
|
|
for file in sorted(lib_dir.rglob(f'{prefix}.so.*')):
|
|
return file.name
|
|
raise ModuleNotFoundError(f'{prefix}.so not found under {lib_dir}')
|
|
|
|
|
|
def get_nvshmem_host_lib_name(base_dir):
|
|
return _find_versioned_so(base_dir, 'libnvshmem_host')
|
|
|
|
|
|
def get_nccl_lib_name(base_dir):
|
|
return _find_versioned_so(base_dir, 'libnccl')
|
|
|
|
|
|
def get_package_version():
|
|
with open(Path(current_dir) / 'deep_ep' / '__init__.py', 'r') as f:
|
|
version_match = re.search(r'^__version__\s*=\s*(.*)$', f.read(), re.MULTILINE)
|
|
public_version = ast.literal_eval(version_match.group(1))
|
|
|
|
# noinspection PyBroadException
|
|
try:
|
|
status_cmd = ['git', 'status', '--porcelain']
|
|
status_output = subprocess.check_output(status_cmd).decode('ascii').strip()
|
|
if status_output:
|
|
print(f'Warning: Git working directory is not clean. Uncommitted changes:\n{status_output}')
|
|
assert False, 'Git working directory is not clean'
|
|
|
|
cmd = ['git', 'rev-parse', '--short', 'HEAD']
|
|
revision = '+' + subprocess.check_output(cmd).decode('ascii').rstrip()
|
|
except:
|
|
revision = '+local'
|
|
return f'{public_version}{revision}'
|
|
|
|
|
|
class CustomBuildPy(build_py):
|
|
def run(self):
|
|
# Make clusters' cache setting default into `envs.py`
|
|
self.generate_default_envs()
|
|
|
|
# Finally, run the regular build
|
|
build_py.run(self)
|
|
|
|
def generate_default_envs(self):
|
|
code = '# Pre-installed environment variables\n'
|
|
code += 'persistent_envs = dict()\n'
|
|
# noinspection PyShadowingNames
|
|
for name in persistent_env_names:
|
|
code += f"persistent_envs['{name}'] = '{os.environ[name]}'\n" if name in os.environ else ''
|
|
|
|
# Create temporary build directory
|
|
build_include_dir = os.path.join(self.build_lib, 'deep_ep')
|
|
os.makedirs(build_include_dir, exist_ok=True)
|
|
with open(os.path.join(self.build_lib, 'deep_ep', 'envs.py'), 'w') as f:
|
|
f.write(code)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# TODO: make NVSHMEM and legacy optional
|
|
nvshmem_root_dir = find_pkgs.find_nvshmem_root()
|
|
nccl_root_dir = find_pkgs.find_nccl_root()
|
|
|
|
# `128,2417` is used to suppress warnings of `fmt`
|
|
cxx_flags = ['-O3', '-Wno-deprecated-declarations', '-Wno-unused-variable', '-Wno-sign-compare', '-Wno-reorder', '-Wno-attributes']
|
|
nvcc_flags = ['-O3', '-Xcompiler', '-O3', '--extended-lambda', '--diag-suppress=128,2417']
|
|
sources = ['csrc/python_api.cpp', 'csrc/kernels/legacy/layout.cu', 'csrc/kernels/legacy/intranode.cu']
|
|
include_dirs = [f'{current_dir}/deep_ep/include',
|
|
f'{current_dir}/third-party/fmt/include',
|
|
'/usr/local/cuda/include/cccl']
|
|
library_dirs = []
|
|
nvcc_dlink = []
|
|
extra_link_args = ['-lcuda']
|
|
|
|
# NVSHMEM flags. Use the real on-disk file name (which may be SONAME-only
|
|
# like ``libnvshmem_host.so.3`` when NVSHMEM came from a pip wheel) so
|
|
# that ``-l:NAME`` can resolve. The static device library always ships
|
|
# under its canonical name, so it stays hard-coded.
|
|
sources.extend(['csrc/kernels/legacy/internode.cu', 'csrc/kernels/legacy/internode_ll.cu', 'csrc/kernels/backend/nvshmem.cu'])
|
|
include_dirs.extend([f'{nvshmem_root_dir}/include'])
|
|
library_dirs.extend([f'{nvshmem_root_dir}/lib'])
|
|
nvcc_dlink.extend(['-dlink', f'-L{nvshmem_root_dir}/lib', '-lnvshmem_device'])
|
|
nvshmem_host_lib = get_nvshmem_host_lib_name(nvshmem_root_dir)
|
|
extra_link_args.extend([f'-l:{nvshmem_host_lib}', '-l:libnvshmem_device.a', f'-Wl,-rpath,{nvshmem_root_dir}/lib'])
|
|
|
|
# NCCL flags. Same story as NVSHMEM above — pip wheels ship
|
|
# ``libnccl.so.2`` only, so resolve the real name dynamically.
|
|
sources.extend(['csrc/kernels/backend/nccl.cu'])
|
|
include_dirs.extend([f'{nccl_root_dir}/include'])
|
|
nccl_lib = get_nccl_lib_name(nccl_root_dir)
|
|
extra_link_args.extend([f'-l:{nccl_lib}', f'-Wl,-rpath,{nccl_root_dir}/lib'])
|
|
|
|
# CUDA driver sources
|
|
sources.extend(['csrc/kernels/backend/cuda_driver.cu'])
|
|
|
|
# TODO: remove these
|
|
if int(os.getenv('DISABLE_SM90_FEATURES', 0)):
|
|
# Prefer A100
|
|
os.environ['TORCH_CUDA_ARCH_LIST'] = os.getenv('TORCH_CUDA_ARCH_LIST', '8.0')
|
|
|
|
# Disable some SM90 features: FP8, launch methods, and TMA
|
|
cxx_flags.append('-DDISABLE_SM90_FEATURES')
|
|
nvcc_flags.append('-DDISABLE_SM90_FEATURES')
|
|
|
|
# Disable internode and low-latency kernels
|
|
assert False, 'Not implemented'
|
|
else:
|
|
# Prefer H800 series
|
|
os.environ['TORCH_CUDA_ARCH_LIST'] = os.getenv('TORCH_CUDA_ARCH_LIST', '9.0')
|
|
|
|
# CUDA 12 flags
|
|
nvcc_flags.extend(['-rdc=true', '--ptxas-options=--register-usage-level=10'])
|
|
|
|
# Disable LD/ST tricks, as some CUDA version does not support `.L1::no_allocate`
|
|
if os.environ['TORCH_CUDA_ARCH_LIST'].strip() != '9.0':
|
|
assert int(os.getenv('DISABLE_AGGRESSIVE_PTX_INSTRS', 1)) == 1
|
|
os.environ['DISABLE_AGGRESSIVE_PTX_INSTRS'] = '1'
|
|
|
|
# Disable aggressive PTX instructions
|
|
if int(os.getenv('DISABLE_AGGRESSIVE_PTX_INSTRS', '1')):
|
|
cxx_flags.append('-DDISABLE_AGGRESSIVE_PTX_INSTRS')
|
|
nvcc_flags.append('-DDISABLE_AGGRESSIVE_PTX_INSTRS')
|
|
|
|
# Legacy environment name
|
|
if 'TOPK_IDX_BITS' in os.environ:
|
|
assert 'EP_NUM_TOPK_IDX_BITS' not in os.environ
|
|
os.environ['EP_NUM_TOPK_IDX_BITS'] = os.environ['TOPK_IDX_BITS']
|
|
|
|
# Bits of `topk_idx.dtype`, choices are 32 and 64
|
|
if 'EP_NUM_TOPK_IDX_BITS' in os.environ:
|
|
num_topk_idx_bits = int(os.environ['EP_NUM_TOPK_IDX_BITS'])
|
|
cxx_flags.append(f'-DEP_NUM_TOPK_IDX_BITS={num_topk_idx_bits}')
|
|
nvcc_flags.append(f'-DEP_NUM_TOPK_IDX_BITS={num_topk_idx_bits}')
|
|
|
|
# Put them together
|
|
extra_compile_args = {
|
|
'cxx': cxx_flags,
|
|
'nvcc': nvcc_flags,
|
|
}
|
|
if len(nvcc_dlink) > 0:
|
|
extra_compile_args['nvcc_dlink'] = nvcc_dlink
|
|
|
|
# Summary
|
|
print('Build summary:')
|
|
print(f' > Sources: {sources}')
|
|
print(f' > Includes: {include_dirs}')
|
|
print(f' > Libraries: {library_dirs}')
|
|
print(f' > Compilation flags: {extra_compile_args}')
|
|
print(f' > Link flags: {extra_link_args}')
|
|
print(f' > Arch list: {os.environ["TORCH_CUDA_ARCH_LIST"]}')
|
|
print(f' > NVSHMEM path: {nvshmem_root_dir}')
|
|
print(f' > NCCL path: {nccl_root_dir}')
|
|
# Print persistent env variables
|
|
persistent_envs = []
|
|
for name in persistent_env_names:
|
|
if name in os.environ:
|
|
persistent_envs.append((name, os.environ[name]))
|
|
if len(persistent_envs) > 0:
|
|
print(f' > Persistent envs:')
|
|
for k, v in persistent_envs:
|
|
print(f' > {k}: {v}')
|
|
print()
|
|
|
|
setuptools.setup(
|
|
name='deep_ep',
|
|
version=get_package_version(),
|
|
packages=setuptools.find_packages(include=['deep_ep', 'deep_ep.*']),
|
|
package_data={
|
|
'deep_ep': [
|
|
'include/deep_ep/**/*',
|
|
]
|
|
},
|
|
ext_modules=[
|
|
CUDAExtension(name='deep_ep._C',
|
|
include_dirs=include_dirs,
|
|
library_dirs=library_dirs,
|
|
sources=sources,
|
|
extra_compile_args=extra_compile_args,
|
|
extra_link_args=extra_link_args)
|
|
],
|
|
cmdclass={
|
|
'build_ext': BuildExtension,
|
|
'build_py': CustomBuildPy
|
|
}
|
|
)
|