wehub-resource-sync
297 lines
10 KiB
Python
297 lines
10 KiB
Python
import functools
|
|
import os
|
|
import platform
|
|
import subprocess
|
|
import sys
|
|
import urllib.request
|
|
import warnings
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
from packaging.version import Version, parse
|
|
from setuptools import setup
|
|
from setuptools.command.bdist_wheel import bdist_wheel
|
|
from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension
|
|
|
|
BASE_WHEEL_URL = "https://github.com/mit-han-lab/fouroversix/releases/download"
|
|
PACKAGE_NAME = "fouroversix"
|
|
PACKAGE_VERSION = "1.1.0"
|
|
|
|
CUTLASS_DEBUG = os.getenv("CUTLASS_DEBUG", "0") == "1"
|
|
FORCE_BUILD = os.getenv("FORCE_BUILD", "0") == "1"
|
|
FORCE_CXX11_ABI = os.getenv("FORCE_CXX11_ABI", "0") == "1"
|
|
SKIP_CUDA_BUILD = os.getenv("SKIP_CUDA_BUILD", "0") == "1"
|
|
|
|
|
|
@functools.cache
|
|
def get_cuda_archs() -> list[str]:
|
|
return os.getenv("CUDA_ARCHS", "100;103;110;120").split(";")
|
|
|
|
|
|
def get_cuda_bare_metal_version() -> Version | None:
|
|
if CUDA_HOME is None:
|
|
warnings.warn(
|
|
"nvcc was not found. Are you sure your environment has nvcc available? If "
|
|
"you're installing within a container from "
|
|
"https://hub.docker.com/r/pytorch/pytorch, only images with 'devel' in "
|
|
"their name will provide nvcc.",
|
|
stacklevel=1,
|
|
)
|
|
return None
|
|
|
|
raw_output = subprocess.check_output(
|
|
[CUDA_HOME + "/bin/nvcc", "-V"],
|
|
universal_newlines=True,
|
|
)
|
|
output = raw_output.split()
|
|
release_idx = output.index("release") + 1
|
|
return parse(output[release_idx].split(",")[0])
|
|
|
|
|
|
def get_cuda_gencodes() -> list[str]:
|
|
"""
|
|
Add -gencode flags based on nvcc capabilities.
|
|
|
|
Uses the following rules:
|
|
- sm_100/120 on CUDA >= 12.8
|
|
- Use 100f on CUDA >= 12.9 (Blackwell family-specific)
|
|
- Map requested 110 -> 101 if CUDA < 13.0 (Thor rename)
|
|
- Embed PTX for newest arch for forward compatibility
|
|
"""
|
|
|
|
archs = set(get_cuda_archs())
|
|
cuda_version = get_cuda_bare_metal_version()
|
|
cc_flags = []
|
|
|
|
# Blackwell requires >= 12.8
|
|
if cuda_version is not None and cuda_version >= Version("12.8"):
|
|
if "100" in archs:
|
|
cc_flags += ["-gencode", "arch=compute_100a,code=sm_100a"]
|
|
|
|
if "103" in archs:
|
|
cc_flags += ["-gencode", "arch=compute_103a,code=sm_103a"]
|
|
|
|
# Thor rename: 12.9 uses sm_101; 13.0+ uses sm_110
|
|
if "110" in archs:
|
|
if cuda_version >= Version("13.0"):
|
|
cc_flags += ["-gencode", "arch=compute_110f,code=sm_110"]
|
|
elif cuda_version >= Version("12.9"):
|
|
# Provide Thor support for CUDA 12.9 via sm_101
|
|
cc_flags += ["-gencode", "arch=compute_101f,code=sm_101"]
|
|
# else: no Thor support in older toolkits
|
|
|
|
if "120" in archs:
|
|
# sm_120 is supported in CUDA 12.8/12.9+ toolkits
|
|
if cuda_version >= Version("12.9"):
|
|
cc_flags += ["-gencode", "arch=compute_120f,code=sm_120"]
|
|
else:
|
|
cc_flags += ["-gencode", "arch=compute_120a,code=sm_120a"]
|
|
|
|
return cc_flags
|
|
|
|
|
|
def get_platform() -> str:
|
|
if sys.platform.startswith("linux"):
|
|
return f"linux_{platform.uname().machine}"
|
|
if sys.platform == "darwin":
|
|
mac_version = ".".join(platform.mac_ver()[0].split(".")[:2])
|
|
return f"macosx_{mac_version}_x86_64"
|
|
if sys.platform == "win32":
|
|
return "win_amd64"
|
|
|
|
msg = f"Unsupported platform: {sys.platform}"
|
|
raise ValueError(msg)
|
|
|
|
|
|
def get_wheel_url() -> tuple[str, str]:
|
|
torch_version_raw = parse(torch.__version__)
|
|
python_version = f"cp{sys.version_info.major}{sys.version_info.minor}"
|
|
platform_name = get_platform()
|
|
torch_version = f"{torch_version_raw.major}.{torch_version_raw.minor}"
|
|
cxx11_abi = str(torch._C._GLIBCXX_USE_CXX11_ABI).upper() # noqa: SLF001
|
|
|
|
# We only compile for CUDA 12.8 to save CI time. Minor versions should be
|
|
# compatible.
|
|
torch_cuda_version = parse("12.8")
|
|
cuda_version = f"cu{torch_cuda_version.major}"
|
|
|
|
wheel_filename = (
|
|
f"{PACKAGE_NAME}-{PACKAGE_VERSION}+{cuda_version}torch{torch_version}"
|
|
f"cxx11abi{cxx11_abi}-{python_version}-{python_version}-{platform_name}.whl"
|
|
)
|
|
|
|
return f"{BASE_WHEEL_URL}/v{PACKAGE_VERSION}/{wheel_filename}", wheel_filename
|
|
|
|
|
|
class CachedWheelsCommand(bdist_wheel):
|
|
"""
|
|
Custom bdist wheel command that checks for pre-built wheels on GitHub Releases.
|
|
|
|
The CachedWheelsCommand plugs into the default bdist wheel, which is ran by pip
|
|
when it cannot find an existing wheel (which is currently the case for all
|
|
fouroversix installs). We use the environment parameters to detect whether there is
|
|
already a pre-built version of a compatible wheel available and short-circuits the
|
|
standard full build pipeline.
|
|
|
|
Credit: https://github.com/Dao-AILab/flash-attention/blob/main/setup.py
|
|
"""
|
|
|
|
def run(self) -> None:
|
|
"""Run the command."""
|
|
|
|
if FORCE_BUILD:
|
|
return super().run()
|
|
|
|
wheel_url, wheel_filename = get_wheel_url()
|
|
print(f"Guessing wheel URL: {wheel_url}")
|
|
|
|
try:
|
|
urllib.request.urlretrieve(wheel_url, wheel_filename) # noqa: S310
|
|
|
|
# Make the archive
|
|
# Lifted from the root wheel processing command
|
|
# https://github.com/pypa/wheel/blob/cf71108ff9f6ffc36978069acb28824b44ae028e/src/wheel/bdist_wheel.py#LL381C9-L381C85
|
|
if not Path(self.dist_dir).exists():
|
|
Path(self.dist_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
impl_tag, abi_tag, plat_tag = self.get_tag()
|
|
archive_basename = f"{self.wheel_dist_name}-{impl_tag}-{abi_tag}-{plat_tag}"
|
|
|
|
wheel_path = Path(self.dist_dir) / (archive_basename + ".whl")
|
|
print(f"Raw wheel path: {wheel_path}")
|
|
Path(wheel_filename).rename(wheel_path)
|
|
except (urllib.error.HTTPError, urllib.error.URLError):
|
|
print("Precompiled wheel not found. Building from source...")
|
|
# If the wheel could not be downloaded, build from source
|
|
super().run()
|
|
|
|
|
|
class NinjaBuildExtension(BuildExtension):
|
|
"""
|
|
Custom build extension that tells Ninja how many jobs to run.
|
|
|
|
Credit: https://github.com/Dao-AILab/flash-attention/blob/main/setup.py
|
|
"""
|
|
|
|
def __init__(self, *args: list[Any], **kwargs: dict[str, Any]) -> None:
|
|
# do not override env MAX_JOBS if already exists
|
|
if not os.environ.get("MAX_JOBS"):
|
|
try:
|
|
import psutil
|
|
|
|
# calculate the maximum allowed NUM_JOBS based on cores
|
|
max_num_jobs_cores = max(1, os.cpu_count() // 2)
|
|
|
|
# calculate the maximum allowed NUM_JOBS based on free memory
|
|
free_memory_gb = psutil.virtual_memory().available / (
|
|
1024**3
|
|
) # free memory in GB
|
|
max_num_jobs_memory = int(
|
|
free_memory_gb / 9,
|
|
) # each JOB peak memory cost is ~8-9GB when threads = 4
|
|
|
|
# pick lower value of jobs based on cores vs memory metric to minimize
|
|
# oom and swap usage during compilation
|
|
max_jobs = max(1, min(max_num_jobs_cores, max_num_jobs_memory))
|
|
os.environ["MAX_JOBS"] = str(max_jobs)
|
|
except ImportError:
|
|
warnings.warn(
|
|
"psutil not found, install psutil and ninja to get better build "
|
|
"performance",
|
|
stacklevel=1,
|
|
)
|
|
|
|
super().__init__(*args, **kwargs)
|
|
|
|
|
|
if SKIP_CUDA_BUILD:
|
|
warnings.warn(
|
|
"SKIP_CUDA_BUILD is set to 1, installing fouroversix without quantization and "
|
|
"matmul kernels",
|
|
stacklevel=1,
|
|
)
|
|
|
|
ext_modules = None
|
|
else:
|
|
if Path(".git").exists():
|
|
subprocess.run(
|
|
[
|
|
"git",
|
|
"submodule",
|
|
"update",
|
|
"--init",
|
|
"third_party/cutlass",
|
|
],
|
|
check=True,
|
|
)
|
|
elif not Path("third_party/cutlass").exists():
|
|
msg = (
|
|
"third_party/cutlass is missing, please use source distribution or git "
|
|
"clone"
|
|
)
|
|
raise RuntimeError(msg)
|
|
|
|
# The compiler flag -D_GLIBCXX_USE_CXX11_ABI is set to be the same as
|
|
# torch._C._GLIBCXX_USE_CXX11_ABI
|
|
# https://github.com/pytorch/pytorch/blob/8472c24e3b5b60150096486616d98b7bea01500b/torch/utils/cpp_extension.py#L920
|
|
if FORCE_CXX11_ABI:
|
|
torch._C._GLIBCXX_USE_CXX11_ABI = True # noqa: SLF001
|
|
|
|
setup_dir = Path(__file__).parent
|
|
kernels_dir = setup_dir / "src" / "fouroversix" / "csrc"
|
|
sources = [
|
|
path.relative_to(Path(__file__).parent).as_posix()
|
|
for ext in ["**/*.cu", "**/*.cpp"]
|
|
for path in kernels_dir.glob(ext)
|
|
]
|
|
|
|
cxx_compile_args = ["-std=c++17"]
|
|
nvcc_compile_args = [
|
|
"-std=c++17",
|
|
"--expt-relaxed-constexpr",
|
|
"-Xcompiler",
|
|
"-funroll-loops",
|
|
"-Xcompiler",
|
|
"-finline-functions",
|
|
*get_cuda_gencodes(),
|
|
]
|
|
|
|
if CUTLASS_DEBUG:
|
|
nvcc_compile_args.extend(
|
|
[
|
|
"-O0",
|
|
"-DCUTLASS_DEBUG_TRACE_LEVEL=3",
|
|
"-DCUTLASS_DEBUG_ENABLE=1",
|
|
"-g",
|
|
],
|
|
)
|
|
else:
|
|
cxx_compile_args.extend(["-O3"])
|
|
nvcc_compile_args.extend(["-O3", "-DNDEBUG"])
|
|
|
|
ext_modules = [
|
|
CUDAExtension(
|
|
"fouroversix._C",
|
|
sources,
|
|
extra_compile_args={"cxx": cxx_compile_args, "nvcc": nvcc_compile_args},
|
|
include_dirs=[
|
|
setup_dir / "third_party/cutlass/examples/common",
|
|
setup_dir / "third_party/cutlass/include",
|
|
setup_dir / "third_party/cutlass/tools/util/include",
|
|
kernels_dir / "include",
|
|
],
|
|
),
|
|
]
|
|
|
|
setup(
|
|
name=PACKAGE_NAME,
|
|
version=PACKAGE_VERSION,
|
|
ext_modules=ext_modules,
|
|
cmdclass={
|
|
"bdist_wheel": CachedWheelsCommand,
|
|
"build_ext": NinjaBuildExtension,
|
|
},
|
|
include_package_data=True,
|
|
)
|