setup.py

import importlib.util

if importlib.util.find_spec("setuptools_scm") is None:
  raise ImportError("setuptools-scm is not installed. Install it by `pip3 install setuptools-scm`")

import os
import re
import subprocess
from os import path
from pathlib import Path

from setuptools import find_packages, setup
from setuptools_scm.version import get_local_dirty_tag

try:
  from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
except ImportError:
  _bdist_wheel = None

ROOT_DIR = Path(__file__).resolve().parent
RELEASE_FLAVOR_ENV = "CACHE_DIT_RELEASE_FLAVOR"
VERSION_WRITE_TO_ENV = "CACHE_DIT_VERSION_WRITE_TO"
SVDQUANT_BUILD_FLAG = "CACHE_DIT_BUILD_SVDQUANT"
WHEEL_PLATFORM_TAG_ENV = "CACHE_DIT_WHEEL_PLATFORM_TAG"
# Distribution name is intentionally not managed via a PACKAGE_NAME env var in
# setup.py. Under PEP 621, setuptools reads the project name from the active
# workspace's pyproject.toml. Default builds therefore use the repository root
# pyproject.toml, while release builds obtain cache-dit-cu13 by preparing a
# temporary workspace whose pyproject.toml is rewritten by
# tools/release_workspace.py.
SPDLOG_SUBMODULE_PATH = ROOT_DIR / "csrc" / "third_party" / "spdlog"
SPDLOG_HEADER_PATH = SPDLOG_SUBMODULE_PATH / "include" / "spdlog" / "spdlog.h"
# sm100 (Blackwell B100/B200, CC 10.0) is not included in the default
# target list because the NVFP4 block-scaled MMA instruction used by the
# SVDQuant FP4 kernel requires sm_120a or higher (PTX ISA, warp-level mma
# Target ISA Notes: ".kind::mxf4nvf4 and .kind::mxf4 are supported on
# sm_120a and sm_121a"). INT4 kernels work correctly on sm100. Users who
# only need INT4 kernels can build for sm100 explicitly via
# CACHE_DIT_CUDA_ARCH_LIST=100 or TORCH_CUDA_ARCH_LIST=100.
CUDA_ARCH_ALIASES = {
  "maxwell": "50",
  "pascal": "60",
  "volta": "70",
  "turing": "75",
  "ampere": "80",
  "ada": "89",
  "hopper": "90",
  "blackwell": "100",
}


def _normalize_svdquant_sm_target(sm: str, *, support_sm120: bool, support_sm121: bool) -> str:
  if sm == "120" and support_sm120:
    return "120a"
  if sm == "121" and support_sm121:
    return "121a"
  return sm


def _is_supported_svdquant_sm_target(sm: str) -> bool:
  normalized = sm.removesuffix("a")
  if not normalized.isdigit():
    return False

  capability = int(normalized)
  return capability >= 80 and capability != 90


def _env_flag(name: str) -> bool:
  return os.getenv(name, "0").strip().lower() in {"1", "true", "yes", "on"}


def _get_release_flavor() -> str:
  flavor = os.getenv(RELEASE_FLAVOR_ENV, "base").strip().lower()
  if flavor not in {"base", "cu13"}:
    raise RuntimeError(f"Unsupported {RELEASE_FLAVOR_ENV}={flavor!r}. Expected 'base' or 'cu13'.")
  return flavor


RELEASE_FLAVOR = _get_release_flavor()
IS_CU13_RELEASE = RELEASE_FLAVOR == "cu13"


def _get_release_platform_tag() -> str:
  if not IS_CU13_RELEASE:
    return ""

  platform_tag = os.getenv(WHEEL_PLATFORM_TAG_ENV, "").strip()
  return platform_tag or "linux_x86_64"


def _should_build_svdquant() -> bool:
  return _env_flag(SVDQUANT_BUILD_FLAG)


if _should_build_svdquant():
  try:
    import torch
    from packaging import version as packaging_version
    from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension
  except ImportError as exc:
    raise RuntimeError(
      "Building the optional SVDQuant extension requires `torch` and `packaging` "
      "to be installed in the active environment. Run `conda activate cdit` and "
      "initialize submodules with `git submodule update --init --recursive --force`, then install with "
      "`CACHE_DIT_BUILD_SVDQUANT=1 python3 -m pip install -e . --no-build-isolation`.") from exc
else:
  torch = None
  packaging_version = None
  CUDA_HOME = None
  BuildExtension = None
  CUDAExtension = None


def _ensure_spdlog_submodule() -> None:
  if SPDLOG_HEADER_PATH.exists():
    return

  try:
    subprocess.check_call(
      ["git", "submodule", "update", "--init", "--recursive", "--force"],
      cwd=ROOT_DIR,
    )
  except (OSError, subprocess.CalledProcessError) as exc:
    raise RuntimeError(
      "The SVDQuant build requires the git submodule `csrc/third_party/spdlog`, "
      "and automatic submodule initialization failed. Run "
      "`git submodule update --init --recursive --force` manually, then rerun the install command."
    ) from exc

  if SPDLOG_HEADER_PATH.exists():
    return

  raise RuntimeError(
    "The SVDQuant build requires the git submodule `csrc/third_party/spdlog`, "
    "but the expected header is still missing after automatic submodule initialization. "
    "Run `git submodule update --init --recursive --force` manually and verify the submodule checkout, "
    "then rerun the install command.")


def _parse_arch_list(raw_arch_list: str) -> list[str]:
  arch_list: list[str] = []
  for arch in re.split(r"[;, ]+", raw_arch_list):
    normalized = arch.strip().lower()
    if not normalized:
      continue
    normalized = normalized.removesuffix("+ptx")
    normalized = normalized.removeprefix("sm_").removeprefix("compute_")
    normalized = normalized.replace(".", "")
    normalized = CUDA_ARCH_ALIASES.get(normalized, normalized)
    arch_list.append(normalized)
  return arch_list


def _svdquant_v2_launch_sources() -> list[str]:
  return [
    f"csrc/kernels/svdq/gemm_w4a4_v2_launch_{dtype}_int4_stage{stage}.cu"
    for dtype in ("fp16", "bf16") for stage in (1, 2, 3)
  ]


def _required_svdquant_sources() -> list[str]:
  return [
    "csrc/kernels/svdq/pybind.cpp",
    "csrc/kernels/svdq/torch.cpp",
    "csrc/kernels/svdq/gemm_w4a4_v2.cu",
    *_svdquant_v2_launch_sources(),
    # Build both INT4 and NVFP4 kernels
    "csrc/kernels/svdq/gemm_w4a4.cu",
    "csrc/kernels/svdq/gemm_w4a4_launch_fp16_int4.cu",
    "csrc/kernels/svdq/gemm_w4a4_launch_fp16_int4_fasteri2f.cu",
    "csrc/kernels/svdq/gemm_w4a4_launch_fp16_fp4.cu",
    "csrc/kernels/svdq/gemm_w4a4_launch_bf16_int4.cu",
    "csrc/kernels/svdq/gemm_w4a4_launch_bf16_fp4.cu",
  ]


def _missing_svdquant_sources() -> list[str]:
  return [
    relative_path for relative_path in _required_svdquant_sources()
    if not (ROOT_DIR / relative_path).exists()
  ]


def _get_nvcc_version(cuda_home: str | None) -> str:
  nvcc_path = path.join(cuda_home, "bin", "nvcc") if cuda_home else "nvcc"
  try:
    nvcc_output = subprocess.check_output([nvcc_path, "--version"]).decode()
  except (OSError, subprocess.CalledProcessError) as exc:
    raise RuntimeError(
      "`CACHE_DIT_BUILD_SVDQUANT=1` was set, but `nvcc` was not found. "
      "Activate the CUDA toolchain inside `conda activate cdit` before building.") from exc

  match = re.search(r"release (\d+\.\d+), V(\d+\.\d+\.\d+)", nvcc_output)
  if match is None:
    raise RuntimeError(f"Unable to parse nvcc version from: {nvcc_output!r}")

  return match.group(2)


def _get_svdquant_nvcc_threads() -> int:
  raw_threads = os.getenv("CACHE_DIT_SVDQ_NVCC_THREADS", "8").strip()
  try:
    threads = int(raw_threads)
  except ValueError as exc:
    raise RuntimeError("CACHE_DIT_SVDQ_NVCC_THREADS must be a positive integer.") from exc

  if threads < 1:
    raise RuntimeError("CACHE_DIT_SVDQ_NVCC_THREADS must be a positive integer.")

  return threads


def _get_version_write_to() -> str | None:
  if VERSION_WRITE_TO_ENV not in os.environ:
    return path.join("src", "cache_dit", "_version.py")

  raw_value = os.getenv(VERSION_WRITE_TO_ENV, "").strip()
  return raw_value or None


def _get_cmdclass() -> dict[str, object]:
  cmdclass = dict(CMDCLASS)
  if _bdist_wheel is None:
    return cmdclass

  class CacheDiTBdistWheel(_bdist_wheel):

    def finalize_options(self):
      super().finalize_options()
      if IS_CU13_RELEASE:
        self.root_is_pure = False

    def get_tag(self):
      python_tag, abi_tag, platform_tag = super().get_tag()
      if IS_CU13_RELEASE:
        return python_tag, abi_tag, _get_release_platform_tag()
      return python_tag, abi_tag, platform_tag

  cmdclass["bdist_wheel"] = CacheDiTBdistWheel
  return cmdclass


def _validate_release_configuration() -> None:
  if not IS_CU13_RELEASE:
    return

  if not _should_build_svdquant():
    raise RuntimeError("The cache-dit-cu13 release flavor requires CACHE_DIT_BUILD_SVDQUANT=1.")

  build_mode = os.getenv("CACHE_DIT_SVDQUANT_BUILD_MODE", "").strip().upper()
  if build_mode != "ALL":
    raise RuntimeError(
      "The cache-dit-cu13 release flavor requires CACHE_DIT_SVDQUANT_BUILD_MODE=ALL.")


def _get_sm_targets() -> list[str]:
  assert packaging_version is not None
  assert CUDA_HOME is not None
  assert torch is not None

  nvcc_version = _get_nvcc_version(CUDA_HOME)
  support_sm120 = packaging_version.parse(nvcc_version) >= packaging_version.parse("12.8")
  support_sm121 = packaging_version.parse(nvcc_version) >= packaging_version.parse("13.0")

  explicit_arch_list = os.getenv("CACHE_DIT_CUDA_ARCH_LIST") or os.getenv("TORCH_CUDA_ARCH_LIST")
  if explicit_arch_list:
    sm_targets = [
      _normalize_svdquant_sm_target(
        sm,
        support_sm120=support_sm120,
        support_sm121=support_sm121,
      ) for sm in _parse_arch_list(explicit_arch_list)
    ]
  else:
    sm_targets = []

  build_mode = os.getenv("CACHE_DIT_SVDQUANT_BUILD_MODE", "FAST").upper()
  if build_mode not in {"FAST", "ALL"}:
    raise RuntimeError(
      f"Unsupported CACHE_DIT_SVDQUANT_BUILD_MODE={build_mode!r}. Expected 'FAST' or 'ALL'.")

  if not sm_targets and build_mode == "FAST" and torch.cuda.is_available(
  ) and torch.cuda.device_count() > 0:
    for device_index in range(torch.cuda.device_count()):
      capability = torch.cuda.get_device_capability(device_index)
      sm = _normalize_svdquant_sm_target(
        f"{capability[0]}{capability[1]}",
        support_sm120=support_sm120,
        support_sm121=support_sm121,
      )
      if sm not in sm_targets:
        sm_targets.append(sm)
  elif not sm_targets:
    sm_targets = ["80", "86", "89"]
    if support_sm120:
      sm_targets.append("120a")
    if support_sm121:
      sm_targets.append("121a")

  unsupported_targets = [sm for sm in sm_targets if not _is_supported_svdquant_sm_target(sm)]
  if unsupported_targets:
    raise RuntimeError("SVDQuant INT4 kernels require >=sm80 and exclude Hopper (sm90). "
                       "Unsupported CUDA SM targets: "
                       f"{', '.join(unsupported_targets)}")

  if not sm_targets:
    raise RuntimeError("No CUDA SM targets were resolved for the SVDQuant extension build.")

  return sm_targets


def _require_min_cuda_for_cu13_release() -> None:
  if not IS_CU13_RELEASE:
    return

  assert packaging_version is not None
  nvcc_version = _get_nvcc_version(CUDA_HOME)
  if packaging_version.parse(nvcc_version) < packaging_version.parse("13.0"):
    raise RuntimeError("The cache-dit-cu13 release flavor requires CUDA 13.0+ for compilation, "
                       f"but nvcc reports {nvcc_version}.")


def _get_svdquant_extension():
  if not _should_build_svdquant():
    return [], {}

  _ensure_spdlog_submodule()
  _require_min_cuda_for_cu13_release()

  missing_sources = _missing_svdquant_sources()
  if missing_sources:
    preview = ", ".join(missing_sources[:4])
    if len(missing_sources) > 4:
      preview += ", ..."
    raise RuntimeError(
      "`CACHE_DIT_BUILD_SVDQUANT=1` was set, but the migrated native source tree is incomplete. "
      f"Missing files under {ROOT_DIR}: {preview}")

  assert BuildExtension is not None
  assert CUDAExtension is not None

  class CacheDiTBuildExtension(BuildExtension):

    def build_extensions(self):
      for ext in self.extensions:
        ext.extra_compile_args.setdefault("cxx", [])
        ext.extra_compile_args.setdefault("nvcc", [])
        if self.compiler.compiler_type == "msvc":
          ext.extra_compile_args["cxx"] += ext.extra_compile_args.pop("msvc", [])
          ext.extra_compile_args["nvcc"] += ext.extra_compile_args.pop("nvcc_msvc", [])
        else:
          ext.extra_compile_args["cxx"] += ext.extra_compile_args.pop("gcc", [])
      super().build_extensions()

  sm_targets = _get_sm_targets()
  gcc_flags = [
    "-DENABLE_BF16=1",
    "-DBUILD_CACHE_DIT_SVDQUANT=1",
    "-fvisibility=hidden",
    "-O3",
    "-std=c++20",
  ]
  msvc_flags = [
    "/DENABLE_BF16=1",
    "/DBUILD_CACHE_DIT_SVDQUANT=1",
    "/O2",
    "/std:c++20",
    "/Zc:__cplusplus",
    "/FS",
  ]
  nvcc_flags = [
    "-DENABLE_BF16=1",
    "-DBUILD_CACHE_DIT_SVDQUANT=1",
    "-O3",
    "-std=c++20",
    "-Xcudafe",
    "--diag_suppress=20208",
    "-U__CUDA_NO_HALF_OPERATORS__",
    "-U__CUDA_NO_HALF_CONVERSIONS__",
    "-U__CUDA_NO_HALF2_OPERATORS__",
    "-U__CUDA_NO_HALF2_CONVERSIONS__",
    "-U__CUDA_NO_BFLOAT16_OPERATORS__",
    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
    "-U__CUDA_NO_BFLOAT162_OPERATORS__",
    "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
    f"--threads={_get_svdquant_nvcc_threads()}",
    "--use_fast_math",
    "--expt-relaxed-constexpr",
    "--expt-extended-lambda",
    "--ptxas-options=--allow-expensive-optimizations=true",
  ]
  if os.getenv("CACHE_DIT_BUILD_WHEELS", "0") == "0":
    nvcc_flags.append("--generate-line-info")
  for sm_target in sm_targets:
    nvcc_flags += ["-gencode", f"arch=compute_{sm_target},code=sm_{sm_target}"]

  extension = CUDAExtension(
    name="cache_dit._C_svdquant",
    sources=_required_svdquant_sources(),
    include_dirs=[
      str(ROOT_DIR / "csrc" / "kernels"),
      str(ROOT_DIR / "csrc" / "kernels" / "svdq"),
      str(ROOT_DIR / "csrc" / "third_party" / "spdlog" / "include"),
    ],
    extra_compile_args={
      "gcc": gcc_flags,
      "msvc": msvc_flags,
      "nvcc": nvcc_flags,
      "nvcc_msvc": [
        "-Xcompiler",
        "/Zc:__cplusplus",
        "-Xcompiler",
        "/FS",
        "-Xcompiler",
        "/bigobj",
      ],
    },
  )

  return [extension], {"build_ext": CacheDiTBuildExtension}


def is_git_directory(path="."):
  return (subprocess.call(
    ["git", "-C", path, "status"],
    stderr=subprocess.STDOUT,
    stdout=open(os.devnull, "w"),
  ) == 0)


def my_local_scheme(version):
  # The following is used to build release packages.
  # Users should never use it.
  local_version = os.getenv("CACHE_DIT_BUILD_LOCAL_VERSION")
  if local_version is None:
    return get_local_dirty_tag(version)
  return f"+{local_version}"


_validate_release_configuration()
EXT_MODULES, CMDCLASS = _get_svdquant_extension()
USE_SCM_VERSION = {
  "version_scheme": "python-simplified-semver",
  "local_scheme": my_local_scheme,
}
VERSION_WRITE_TO = _get_version_write_to()
if VERSION_WRITE_TO is not None:
  USE_SCM_VERSION["write_to"] = VERSION_WRITE_TO

setup(
  description=
  "Cache-DiT: Cache-DiT: A PyTorch-native Inference Engine with Cache, Parallelism, Quantization and CPU Offload for DiTs.",
  author="DefTruth, vipshop.com",
  use_scm_version=USE_SCM_VERSION,
  package_dir={"": "src"},
  packages=find_packages(where="src"),
  include_package_data=True,
  python_requires=">=3.10",
  ext_modules=EXT_MODULES,
  cmdclass=_get_cmdclass(),
)