setup.py

# This file was modified for portability to AMDGPU
# Copyright (c) 2022-2026, Advanced Micro Devices, Inc. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

"""Installation script."""

from importlib import metadata
import os
import shutil
import subprocess
import time
from pathlib import Path
from typing import List, Tuple

import setuptools
from setuptools.command.egg_info import egg_info
try:
    from setuptools.command.bdist_wheel import bdist_wheel
except ImportError:
    from wheel.bdist_wheel import bdist_wheel

from build_tools.build_ext import CMakeExtension, get_build_ext
from build_tools.te_version import te_version
from build_tools.utils import (
    rocm_build,
    rocm_version,
    all_files_in_dir,
    cuda_archs,
    cuda_version,
    get_frameworks,
    remove_dups,
    min_python_version_str,
    nccl_ep_enabled,
)

frameworks = get_frameworks()
current_file_path = Path(__file__).parent.resolve()


from setuptools.command.build_ext import build_ext as BuildExtension
from setuptools.command.build_py import build_py as _build_py

os.environ["NVTE_PROJECT_BUILDING"] = "1"

_ROCM_INIT_TEMPLATE = current_file_path / "build_tools" / "templates" / "_rocm_init.py"


class BuildPy(_build_py):
    """Generate _rocm_init.py for ROCm builds only."""

    def run(self):
        # Generated into the source tree so build_py picks it up as an ordinary module of
        # the package, putting the same file in the checkout and in the wheel. Both need
        # it: transformer_engine/__init__.py imports it with `from . import _rocm_init`
        # and tolerates its absence, so an install without it silently skips the rocm-sdk
        # preload and loads the native libraries against an uninitialized ROCm runtime.
        # The generated file is gitignored.
        #
        # The write has to stay ahead of super().run(): build_py globs the package
        # directory as it runs, so a later write would land in the checkout but miss the
        # wheel.
        dest = current_file_path / "transformer_engine" / "_rocm_init.py"
        if rocm_build():
            shutil.copy2(_ROCM_INIT_TEMPLATE, dest)
        else:
            # Drop a copy left behind by an earlier ROCm build in the same tree,
            # so the wheel contents follow this build's config, not build history.
            dest.unlink(missing_ok=True)
        super().run()

if "pytorch" in frameworks:
    from torch.utils.cpp_extension import BuildExtension
elif "jax" in frameworks:
    from pybind11.setup_helpers import build_ext as BuildExtension

class HipifyMeta(egg_info):
    """Custom egg_info command to hipify source files before packaging."""

    def run(self):
        if rocm_build():
            from build_tools.hipify.hipify import do_hipify
            print("Running hipification of installable headers for ROCm build...")
            do_hipify(current_file_path, current_file_path / "transformer_engine/common/include")
        super().run()

CMakeBuildExtension = get_build_ext(BuildExtension)
if not rocm_build():
    archs = cuda_archs()

class TimedBdist(bdist_wheel):
    """Helper class to measure build time"""

    def run(self):
        start_time = time.perf_counter()
        super().run()
        total_time = time.perf_counter() - start_time
        print(f"Total time for bdist_wheel: {total_time:.2f} seconds")


def setup_common_extension() -> CMakeExtension:
    """Setup CMake extension for common library"""
    cmake_flags = []
    if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))):
            assert (
                os.getenv("MPI_HOME") is not None
            ), "MPI_HOME must be set when compiling with NVTE_UB_WITH_MPI=1"
            cmake_flags.append("-DNVTE_UB_WITH_MPI=ON")

    if rocm_build():
        cmake_flags.append("-DUSE_ROCM=ON")
        cmake_flags.append(
            f"-DCK_FUSED_ATTN_FLOAT_TO_BFLOAT16_DEFAULT={os.getenv('NVTE_CK_FUSED_ATTN_FLOAT_TO_BFLOAT16_DEFAULT', '3')}"
        )

        if int(os.getenv("NVTE_FUSED_ATTN_AOTRITON", "1"))==0 or int(os.getenv("NVTE_FUSED_ATTN", "1"))==0:
            cmake_flags.append("-DUSE_FUSED_ATTN_AOTRITON=OFF")
        elif os.getenv("NVTE_FUSED_ATTN_AOTRITON") or os.getenv("NVTE_FUSED_ATTN"):
            cmake_flags.append("-DUSE_FUSED_ATTN_AOTRITON=ON")

        if int(os.getenv("NVTE_FUSED_ATTN_CK", "1"))==0 or int(os.getenv("NVTE_FUSED_ATTN", "1"))==0:
            cmake_flags.append("-DUSE_FUSED_ATTN_CK=OFF")
        elif os.getenv("NVTE_FUSED_ATTN_CK") or os.getenv("NVTE_FUSED_ATTN"):
            cmake_flags.append("-DUSE_FUSED_ATTN_CK=ON")

        if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", "0"))) and os.getenv("NVTE_ENABLE_ROCSHMEM") is None:
            os.environ["NVTE_ENABLE_ROCSHMEM"] = '1'
            os.environ["NVTE_ENABLE_NVSHMEM"] = '0'
            print("Turning NVTE_ENABLE_ROCSHMEM on, disabling NVTE_ENABLE_NVSHMEM")
        if bool(int(os.getenv("NVTE_ENABLE_ROCSHMEM", "0"))):
            cmake_flags.append("-DNVTE_ENABLE_ROCSHMEM=ON")

    else:
        cmake_flags.extend(("-DUSE_ROCM=OFF", "-DCMAKE_CUDA_ARCHITECTURES={}".format(archs)))

        if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", "0"))):
            assert (
                os.getenv("NVSHMEM_HOME") is not None
            ), "NVSHMEM_HOME must be set when compiling with NVTE_ENABLE_NVSHMEM=1"
            cmake_flags.append("-DNVTE_ENABLE_NVSHMEM=ON")

        if bool(int(os.getenv("NVTE_BUILD_ACTIVATION_WITH_FAST_MATH", "0"))):
            cmake_flags.append("-DNVTE_BUILD_ACTIVATION_WITH_FAST_MATH=ON")

    if bool(int(os.getenv("NVTE_WITH_CUBLASMP", "0"))):
        cmake_flags.append("-DNVTE_WITH_CUBLASMP=ON")
        cublasmp_dir = os.getenv("CUBLASMP_HOME") or metadata.distribution(
            f"nvidia-cublasmp-cu{cuda_version()[0]}"
        ).locate_file(f"nvidia/cublasmp/cu{cuda_version()[0]}")
        cmake_flags.append(f"-DCUBLASMP_DIR={cublasmp_dir}")

    if bool(int(os.getenv("NVTE_WITH_CUSOLVERMP", "0"))):
        cmake_flags.append("-DNVTE_WITH_CUSOLVERMP=ON")
        cusolvermp_dir = os.getenv("CUSOLVERMP_HOME", "/usr")
        cmake_flags.append(f"-DCUSOLVERMP_DIR={cusolvermp_dir}")

    # NCCL EP (Hopper+): on by default; auto-skipped when no arch >= 90 is
    # targeted. Set NVTE_WITH_NCCL_EP=0 to force off.
    # Disabled on ROCm
    if rocm_build():
        cmake_flags.append("-DNVTE_WITH_NCCL_EP=OFF")
    elif nccl_ep_enabled(archs):
        nccl_home = build_nccl_ep_submodule()
        cmake_flags.append(f"-DNCCL_INCLUDE_DIR={nccl_home}/include")
    else:
        cmake_flags.append("-DNVTE_WITH_NCCL_EP=OFF")

    # Add custom CMake arguments from environment variable
    nvte_cmake_extra_args = os.getenv("NVTE_CMAKE_EXTRA_ARGS")
    if nvte_cmake_extra_args:
        cmake_flags.extend(nvte_cmake_extra_args.split())

    # Project directory root
    root_path = Path(__file__).resolve().parent

    return CMakeExtension(
        name="transformer_engine",
        cmake_path=root_path / Path("transformer_engine/common"),
        cmake_flags=cmake_flags,
    )


def setup_requirements() -> Tuple[List[str], List[str]]:
    """Setup Python dependencies

    Returns dependencies for runtime and testing.
    """

    # Common requirements
    install_reqs: List[str] = [
        "pydantic",
        "importlib-metadata>=1.0",
        "packaging",
    ]
    test_reqs: List[str] = ["pytest>=8.2.1"]

    # Framework-specific requirements
    if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))):
        if "pytorch" in frameworks:
            from build_tools.pytorch import install_requirements, test_requirements

            install_reqs.extend(install_requirements())
            test_reqs.extend(test_requirements())
        if "jax" in frameworks:
            from build_tools.jax import install_requirements, test_requirements

            install_reqs.extend(install_requirements())
            test_reqs.extend(test_requirements())

    return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]]


def _discover_nccl_home() -> str:
    """Resolve NCCL_HOME, preferring the NCCL the dynamic loader resolves at runtime.

    Probes in order: NCCL_HOME env var, ldconfig cache, well-known prefixes, then a
    pip-installed nvidia-nccl-cu* wheel. To test a non-default NCCL (e.g. a wheel), set
    NCCL_HOME and ensure the runtime loader resolves the same lib (e.g. LD_LIBRARY_PATH).
    """
    env_home = os.environ.get("NCCL_HOME")
    if env_home:
        if (Path(env_home) / "include" / "nccl.h").exists():
            return env_home
        print(
            f"[NCCL EP] WARNING: NCCL_HOME='{env_home}' is set but "
            f"'{env_home}/include/nccl.h' was not found; falling back to system probes."
        )

    lib_names = ("libnccl.so", "libnccl.so.2")
    # Include Debian/Ubuntu multiarch subdirs (e.g. lib/aarch64-linux-gnu).
    lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu")

    # Prefer the NCCL the dynamic loader will actually resolve at runtime so the
    # EP build links against the same libnccl that gets loaded. libtransformer_engine
    # carries no NCCL RUNPATH, so the loader uses ldconfig/system paths; building
    # against a different NCCL (e.g. a pip wheel) causes ABI mismatches. ldconfig is
    # the ground truth for runtime resolution, so consult it before well-known prefixes.
    try:
        out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode()
        for line in out.splitlines():
            if "libnccl.so" in line and "=>" in line:
                lib_path = Path(line.split("=>")[-1].strip())
                # Walk upward so multiarch layouts (.../lib/<triplet>/libnccl.so)
                # resolve to the prefix that contains include/nccl.h.
                for root in (lib_path.parent.parent, lib_path.parent.parent.parent):
                    if (root / "include" / "nccl.h").exists():
                        return str(root)
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass

    for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"):
        p = Path(cand)
        if (p / "include" / "nccl.h").exists() and any(
            (p / sub / name).exists() for sub in lib_subdirs for name in lib_names
        ):
            return str(p)

    # Fall back to a pip-installed NCCL (nvidia-nccl-cu* wheel) under nvidia/nccl
    # in site-packages, used only when no system NCCL is present.
    try:
        import importlib.util

        spec = importlib.util.find_spec("nvidia.nccl")
        if spec is not None and spec.submodule_search_locations:
            pip_root = Path(next(iter(spec.submodule_search_locations)))
            if (pip_root / "include" / "nccl.h").exists() and any(
                (pip_root / sub / name).exists() for sub in lib_subdirs for name in lib_names
            ):
                return str(pip_root)
    except (ImportError, ValueError):
        pass

    raise RuntimeError(
        "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix."
    )


def build_nccl_ep_submodule() -> str:
    """Build libnccl_ep.a from the 3rdparty/nccl submodule and return NCCL_HOME."""
    nccl_root = current_file_path / "3rdparty" / "nccl"
    if not (nccl_root / "Makefile").exists():
        raise RuntimeError(
            f"NCCL submodule not found at {nccl_root}. "
            "Run `git submodule update --init --recursive`."
        )

    build_dir = nccl_root / "build"
    nccl_ep_lib = build_dir / "lib" / "libnccl_ep.a"
    gencode_stamp = build_dir / "lib" / "libnccl_ep.gencode"

    # Caller gates on arch >= 90 or "native"; expand "native" to the host's
    # actual sm_XX so the build stamp distinguishes machines.
    arch_tokens = [a.strip() for a in str(cuda_archs() or "").split(";") if a.strip()]
    arch_list: list[str] = []
    for t in arch_tokens:
        if t.lower() == "native":
            try:
                out = subprocess.check_output(
                    ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
                    stderr=subprocess.DEVNULL,
                ).decode()
            except (subprocess.CalledProcessError, FileNotFoundError) as e:
                raise RuntimeError(
                    "NVTE_CUDA_ARCHS=native requires nvidia-smi to resolve the host arch."
                ) from e
            for line in out.splitlines():
                cap = line.strip().replace(".", "")
                if cap.isdigit() and int(cap) >= 90 and cap not in arch_list:
                    arch_list.append(cap)
        else:
            bare = t.rstrip("af")
            if bare.isdigit() and int(bare) >= 90 and bare not in arch_list:
                arch_list.append(bare)
    if not arch_list:
        raise RuntimeError(
            "NCCL EP requires Hopper or newer (SM >= 90); none found in"
            f" NVTE_CUDA_ARCHS={cuda_archs()!r}. Re-run with NVTE_WITH_NCCL_EP=0 to skip the NCCL"
            " EP build (the rest of TE still builds)."
        )
    gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list)

    nproc = os.cpu_count() or 8
    env = os.environ.copy()
    env["NVCC_GENCODE"] = gencode
    # NCCL EP needs the core NCCL headers + libnccl.so; write NCCL EP build
    # outputs to the submodule's local build/ tree.
    nccl_home = _discover_nccl_home()
    env["NCCL_HOME"] = nccl_home
    env["NCCL_EP_BUILDDIR"] = str(build_dir)

    prev_gencode = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None
    if not nccl_ep_lib.exists() or prev_gencode != gencode:
        if nccl_ep_lib.exists() and prev_gencode != gencode:
            print(
                f"[NCCL EP] gencode changed ('{prev_gencode}' -> '{gencode}'); "
                "rebuilding libnccl_ep.a"
            )
            subprocess.check_call(
                ["make", "-C", "contrib/nccl_ep", "clean"],
                cwd=str(nccl_root),
                env=env,
            )
        print(f"[NCCL EP] Building libnccl_ep.a (gencode='{gencode}')")
        subprocess.check_call(
            ["make", "-j", str(nproc), "-C", "contrib/nccl_ep", "lib"],
            cwd=str(nccl_root),
            env=env,
        )
        gencode_stamp.parent.mkdir(parents=True, exist_ok=True)
        gencode_stamp.write_text(gencode)

    return nccl_home


def git_check_submodules() -> None:
    """
    Attempt to checkout git submodules automatically during setup.

    This runs successfully only if the submodules are
    either in the correct or uninitialized state.

    Note to devs: With this, any updates to the submodules itself, e.g. moving to a newer
    commit, must be commited before build. This also ensures that stale submodules aren't
    being silently used by developers.
    """

    # Provide an option to skip these checks for development.
    if bool(int(os.getenv("NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD", "0"))):
        return

    # Require git executable.
    if shutil.which("git") is None:
        return

    # Require a .gitmodules file.
    if not (current_file_path / ".gitmodules").exists():
        return

    try:
        submodules = subprocess.check_output(
            ["git", "submodule", "status", "--recursive"],
            cwd=str(current_file_path),
            text=True,
        ).splitlines()

        for submodule in submodules:
            # '-' start is for an uninitialized submodule.
            # ' ' start is for a submodule on the correct commit.
            assert submodule[0] in (
                " ",
                "-",
            ), (
                "Submodules are initialized incorrectly. If this is intended, set the "
                "environment variable `NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD` to a "
                "non-zero value to skip these checks during development. Otherwise, "
                "run `git submodule update --init --recursive` to checkout the correct"
                " submodule commits."
            )

        subprocess.check_call(
            ["git", "submodule", "update", "--init", "--recursive"],
            cwd=str(current_file_path),
        )
    except subprocess.CalledProcessError:
        return


if __name__ == "__main__":
    __version__ = te_version()

    git_check_submodules()

    with open("README.rst", encoding="utf-8") as f:
        long_description = f.read()

    # Settings for building top level empty package for dependency management.
    if bool(int(os.getenv("NVTE_BUILD_METAPACKAGE", "0"))):
        assert bool(
            int(os.getenv("NVTE_RELEASE_BUILD", "0"))
        ), "NVTE_RELEASE_BUILD env must be set for metapackage build."
        ext_modules = []
        package_data = {}
        include_package_data = False
        install_requires = []
        extras_require = {
            "core": [f"transformer_engine_cu12=={__version__}"],
            "core_cu12": [f"transformer_engine_cu12=={__version__}"],
            "core_cu13": [f"transformer_engine_cu13=={__version__}"],
            "pytorch": [f"transformer_engine_torch=={__version__}"],
            "jax": [f"transformer_engine_jax=={__version__}"],
        } if not rocm_build() else {
            "rocm": [f"transformer_engine_rocm7=={__version__}"],
            "rocm7": [f"transformer_engine_rocm7=={__version__}"],
            "rocm_pytorch": [f"transformer_engine_rocm7[pytorch]=={__version__}"],
            "rocm_jax": [f"transformer_engine_rocm7[jax]=={__version__}"],
        }
    else:
        install_requires, test_requires = setup_requirements()
        ext_modules = [setup_common_extension()]
        package_data = {
            "": ["VERSION.txt"],
            "transformer_engine.pytorch.triton_kernels.gmm": ["configs/*.json"],
        }
        include_package_data = True
        extras_require = {"test": test_requires}

        if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))):
            if "pytorch" in frameworks:
                from build_tools.pytorch import setup_pytorch_extension

                ext_modules.append(
                    setup_pytorch_extension(
                        "transformer_engine/pytorch/csrc",
                        current_file_path / "transformer_engine" / "pytorch" / "csrc",
                        current_file_path / "transformer_engine",
                    )
                )
            if "jax" in frameworks:
                from build_tools.jax import setup_jax_extension

                ext_modules.append(
                    setup_jax_extension(
                        "transformer_engine/jax/csrc",
                        current_file_path / "transformer_engine" / "jax" / "csrc",
                        current_file_path / "transformer_engine",
                    )
                )

    PACKAGE_NAME="transformer_engine"
    if (rocm_build() and bool(int(os.getenv("NVTE_RELEASE_BUILD", "0")))
        and not bool(int(os.getenv("NVTE_BUILD_METAPACKAGE", "0"))) ):
        PACKAGE_NAME=f"transformer_engine_rocm{rocm_version()[0]}"
        #On ROCm add extras to core package so it can be installed w/o metapackage
        extras_require.update({
            "pytorch": [f"transformer_engine_rocm_torch=={__version__}"],
            "jax": [f"transformer_engine_rocm_jax=={__version__}"],
        })

    # Configure package
    setuptools.setup(
        name=PACKAGE_NAME,
        version=__version__,
        packages=setuptools.find_packages(
            include=[
                "transformer_engine",
                "transformer_engine.*",
                "transformer_engine/build_tools",
            ],
        ),
        extras_require=extras_require,
        description="Transformer acceleration library",
        long_description=long_description,
        long_description_content_type="text/x-rst",
        ext_modules=ext_modules,
        cmdclass={
            "egg_info": HipifyMeta,
            "build_py": BuildPy,
            "build_ext": CMakeBuildExtension,
            "bdist_wheel": TimedBdist,
        },
        python_requires=f">={min_python_version_str()}",
        classifiers=["Programming Language :: Python :: 3"],
        install_requires=install_requires,
        license_files=("LICENSE",),
        include_package_data=include_package_data,
        package_data=package_data,
    )