From 1a6efa123132a9bedd45c859a64bc44ff38eb8e3 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 24 Jul 2026 12:26:41 -0500 Subject: [PATCH 01/24] Add CUDA wheels as build-time requirements for build isolation Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 261 +++++++++++++++++++++++++++++++++++ docs/envvars.rst | 6 + pyproject.toml | 7 +- 3 files changed, 271 insertions(+), 3 deletions(-) create mode 100644 build_tools/build_backend.py diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py new file mode 100644 index 00000000000..0b892949799 --- /dev/null +++ b/build_tools/build_backend.py @@ -0,0 +1,261 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PEP 517 backend with CUDA-version-specific build requirements.""" + +import os +import re +import shutil +import subprocess +from pathlib import Path +from typing import List, Mapping, Optional, Union + +from setuptools import build_meta + + +_SETUPTOOLS_BACKEND = build_meta.__legacy__ +_CUDA_BUILD_PACKAGES = { + "12": [ + # Also contains nvvm and crt + "nvidia-cuda-nvcc-cu12", + "nvidia-cuda-runtime-cu12", + "nvidia-cuda-cccl-cu12", + "nvidia-cuda-profiler-api-cu12", + "nvidia-nvml-dev-cu12", + ], + "13": [ + "nvidia-cuda-nvcc", + "nvidia-cuda-runtime", + "nvidia-cuda-crt", + "nvidia-cuda-cccl", + "nvidia-cuda-profiler-api", + "nvidia-nvml-dev", + "nvidia-nvvm", + ], +} +_DEFAULT_CUDA_VERSION = "13.3" + +ConfigValue = Union[str, List[str]] +ConfigSettings = Optional[Mapping[str, ConfigValue]] + + +def _setuptools_config_settings(config_settings: ConfigSettings) -> ConfigSettings: + """Remove settings that are private to this backend.""" + if config_settings is None: + return None + + settings = dict(config_settings) + settings.pop("cuda-version", None) + return settings or None + + +def _normalize_cuda_version(value: str) -> str: + """Validate and normalize a CUDA major/minor version.""" + value = value.lower().removeprefix("cu") + if re.fullmatch(r"\d+\.\d+", value) is None: + raise ValueError( + f"Invalid CUDA version {value!r}; expected ., for example 12.8" + ) + + major = value.split(".", maxsplit=1)[0] + if major not in _CUDA_BUILD_PACKAGES: + supported = ", ".join(sorted(_CUDA_BUILD_PACKAGES)) + raise ValueError(f"Unsupported CUDA major version {major!r}; expected one of: {supported}") + + return value + + +def _config_cuda_version(config_settings: ConfigSettings) -> Optional[str]: + """Get the CUDA version from PEP 517 config settings.""" + settings: Mapping[str, ConfigValue] = config_settings or {} + value = settings.get("cuda-version") + if value is None: + return None + + if isinstance(value, list): + if not value: + raise ValueError("CUDA package version cannot be empty") + value = value[-1] + + return _normalize_cuda_version(str(value)) + + +def _environment_cuda_version() -> Optional[str]: + """Get the CUDA version from the environment.""" + value = os.getenv("NVTE_CUDA_VERSION") + return _normalize_cuda_version(value) if value is not None else None + + +def _system_cuda_version() -> Optional[str]: + """Get the version reported by a system-installed NVCC.""" + candidates = [] + if cuda_compiler := os.getenv("CUDACXX"): + candidates.append(Path(cuda_compiler)) + if cuda_home := os.getenv("CUDA_HOME"): + candidates.append(Path(cuda_home) / "bin" / "nvcc") + if nvcc := shutil.which("nvcc"): + candidates.append(Path(nvcc)) + candidates.append(Path("/usr/local/cuda/bin/nvcc")) + + checked = set() + for candidate in candidates: + if candidate in checked or not candidate.is_file(): + continue + + checked.add(candidate) + try: + result = subprocess.run( + [candidate, "--version"], + capture_output=True, + check=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + continue + + match = re.search(r"release\s+(\d+\.\d+)", result.stdout) + if match is not None: + return _normalize_cuda_version(match.group(1)) + + return None + + +def _torch_cuda_version() -> Optional[str]: + """Get the CUDA version required by PyTorch, if available.""" + try: + import torch + except (ImportError, OSError): + return None + value = getattr(getattr(torch, "version", None), "cuda", None) + return _normalize_cuda_version(value) if value else None + + +def _framework_cuda_version() -> Optional[str]: + """Get a CUDA version required by the selected framework.""" + frameworks = { + framework.strip().lower() + for framework in os.getenv("NVTE_FRAMEWORK", "").split(",") + if framework.strip() + } + + if not frameworks or frameworks.intersection({"all", "pytorch"}): + if cuda_version := _torch_cuda_version(): + return cuda_version + + # JAX CUDA plugin names identify the CUDA major version, but not the + # major/minor toolkit release needed to select component packages. + return None + + +def _cuda_version(config_settings: ConfigSettings) -> str: + """Resolve the CUDA version in descending order of precedence.""" + return ( + _config_cuda_version(config_settings) + or _environment_cuda_version() + or _system_cuda_version() + or _framework_cuda_version() + or _DEFAULT_CUDA_VERSION + ) + + +def _cuda_build_requirements(config_settings: ConfigSettings) -> List[str]: + """Get build requirements for the requested CUDA package version.""" + cuda_version = _cuda_version(config_settings) + cuda_major = cuda_version.split(".", maxsplit=1)[0] + component_constraint = f"=={cuda_version}.*" + + packages = _CUDA_BUILD_PACKAGES[cuda_major] + requirements = [f"{package}{component_constraint}" for package in packages] + requirements.append(f"nvidia-nccl-cu{cuda_major}>=2") + return requirements + +################################################################################################### +# PEP 517 and PEP 660 defined functions +################################################################################################### + +# Defined by PEP 517 +def get_requires_for_build_wheel(config_settings: ConfigSettings = None) -> List[str]: + """Get CUDA requirements for building a wheel.""" + return _cuda_build_requirements(config_settings) + + +# Defined by PEP 517 +def get_requires_for_build_sdist(config_settings: ConfigSettings = None) -> List[str]: + """Get CUDA requirements needed while evaluating setup.py for an sdist.""" + # CUDA build requirements are needed even for sdist because setup.py currently always + # depends on these packages. It could be refactored to only depend on them when building + # a wheel, but that would require a more invasive refactor. + return _cuda_build_requirements(config_settings) + + +# Defined by PEP 517 +def build_wheel( + wheel_directory: str, + config_settings: ConfigSettings = None, + metadata_directory: Optional[str] = None, +) -> str: + """Build a wheel with setuptools.""" + return _SETUPTOOLS_BACKEND.build_wheel( + wheel_directory, + _setuptools_config_settings(config_settings), + metadata_directory, + ) + + +# Defined by PEP 517 +def prepare_metadata_for_build_wheel( + metadata_directory: str, + config_settings: ConfigSettings = None, +) -> str: + """Prepare wheel metadata with setuptools.""" + return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_wheel( + metadata_directory, + _setuptools_config_settings(config_settings), + ) + + +# Defined by PEP 517 +def build_sdist( + sdist_directory: str, + config_settings: ConfigSettings = None, +) -> str: + """Build an sdist with setuptools.""" + return _SETUPTOOLS_BACKEND.build_sdist( + sdist_directory, + _setuptools_config_settings(config_settings), + ) + + +if hasattr(_SETUPTOOLS_BACKEND, "build_editable"): + + # Defined by PEP 660 + def get_requires_for_build_editable(config_settings: ConfigSettings = None) -> List[str]: + """Get CUDA requirements for building an editable wheel.""" + return _cuda_build_requirements(config_settings) + + + # Defined by PEP 660 + def build_editable( + wheel_directory: str, + config_settings: ConfigSettings = None, + metadata_directory: Optional[str] = None, + ) -> str: + """Build an editable wheel with setuptools.""" + return _SETUPTOOLS_BACKEND.build_editable( + wheel_directory, + _setuptools_config_settings(config_settings), + metadata_directory, + ) + + + # Defined by PEP 660 + def prepare_metadata_for_build_editable( + metadata_directory: str, + config_settings: ConfigSettings = None, + ) -> str: + """Prepare editable-wheel metadata with setuptools.""" + return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_editable( + metadata_directory, + _setuptools_config_settings(config_settings), + ) diff --git a/docs/envvars.rst b/docs/envvars.rst index 5ff430eabea..32fc3a2960f 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -78,6 +78,12 @@ Build Configuration :Default: ``10`` :Description: Number of Philox4x32 rounds used by stochastic rounding kernels. Must be a positive integer. +.. envvar:: NVTE_CUDA_VERSION + + :Type: ``str`` + :Default: Auto-detected based on system CUDA version, falling back to ``"13.3"`` + :Description: CUDA version to build for. Valid values: ``"12.*"``, ``"13.*"``. If not set, automatically determined based on the installed CUDA Toolkit version. Not intended for external use. + Optional Dependencies ^^^^^^^^^^^^^^^^^^^^^ diff --git a/pyproject.toml b/pyproject.toml index 2f5f8979703..4f7278a7fc4 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,8 @@ # See LICENSE for license information. [build-system] -requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.29.0"] +# CUDA build-time requirements are defined in `build_tools/build_backend.py`. -# Use legacy backend to import local packages in setup.py -build-backend = "setuptools.build_meta:__legacy__" +# Select build dependencies with --config-setting cuda-version=. +build-backend = "build_tools.build_backend" +backend-path = ["."] From 573e5ebf17d6bf5d9f737b1ef897e99c486c085b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 24 Jul 2026 18:00:35 +0000 Subject: [PATCH 02/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- build_tools/build_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index 0b892949799..f340b47ba98 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -170,10 +170,12 @@ def _cuda_build_requirements(config_settings: ConfigSettings) -> List[str]: requirements.append(f"nvidia-nccl-cu{cuda_major}>=2") return requirements + ################################################################################################### # PEP 517 and PEP 660 defined functions ################################################################################################### + # Defined by PEP 517 def get_requires_for_build_wheel(config_settings: ConfigSettings = None) -> List[str]: """Get CUDA requirements for building a wheel.""" @@ -234,7 +236,6 @@ def get_requires_for_build_editable(config_settings: ConfigSettings = None) -> L """Get CUDA requirements for building an editable wheel.""" return _cuda_build_requirements(config_settings) - # Defined by PEP 660 def build_editable( wheel_directory: str, @@ -248,7 +249,6 @@ def build_editable( metadata_directory, ) - # Defined by PEP 660 def prepare_metadata_for_build_editable( metadata_directory: str, From e8e4d2a625d875313f2631e400585f4f527bf932 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 4 Sep 2026 12:48:49 -0500 Subject: [PATCH 03/24] Support framework-specific isolated builds Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 185 +++++++++++++++++++++++++++-------- docs/envvars.rst | 4 +- pyproject.toml | 13 ++- 3 files changed, 156 insertions(+), 46 deletions(-) diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index f340b47ba98..a943acad023 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -2,18 +2,18 @@ # # See LICENSE for license information. -"""PEP 517 backend with CUDA-version-specific build requirements.""" +"""PEP 517 backend with CUDA-version- and framework-specific build requirements.""" +from contextlib import contextmanager import os import re import shutil import subprocess from pathlib import Path -from typing import List, Mapping, Optional, Union +from typing import Iterator, List, Mapping, Optional, Union from setuptools import build_meta - _SETUPTOOLS_BACKEND = build_meta.__legacy__ _CUDA_BUILD_PACKAGES = { "12": [ @@ -34,6 +34,11 @@ "nvidia-nvvm", ], } +_FRAMEWORK_BUILD_PACKAGES = { + "pytorch": ["torch>=2.1"], + "jax": ["jax>=0.5.0", "flax>=0.7.1"], +} +_SUPPORTED_FRAMEWORKS = tuple(_FRAMEWORK_BUILD_PACKAGES) _DEFAULT_CUDA_VERSION = "13.3" ConfigValue = Union[str, List[str]] @@ -47,6 +52,7 @@ def _setuptools_config_settings(config_settings: ConfigSettings) -> ConfigSettin settings = dict(config_settings) settings.pop("cuda-version", None) + settings.pop("framework", None) return settings or None @@ -87,6 +93,78 @@ def _environment_cuda_version() -> Optional[str]: return _normalize_cuda_version(value) if value is not None else None +def _normalize_frameworks(value: str) -> List[str]: + """Validate and normalize a framework selection.""" + frameworks = list( + dict.fromkeys( + framework.strip().lower() for framework in value.split(",") if framework.strip() + ) + ) + if not frameworks: + raise ValueError("Framework selection cannot be empty") + + special_frameworks = {"all", "none"}.intersection(frameworks) + if special_frameworks: + if len(frameworks) != 1: + raise ValueError("'all' and 'none' cannot be combined with other frameworks") + return list(_SUPPORTED_FRAMEWORKS) if frameworks[0] == "all" else [] + + unsupported = [framework for framework in frameworks if framework not in _SUPPORTED_FRAMEWORKS] + if unsupported: + supported = ", ".join((*_SUPPORTED_FRAMEWORKS, "all", "none")) + raise ValueError( + f"Unsupported framework {unsupported[0]!r}; expected one or more of: {supported}" + ) + + return frameworks + + +def _config_frameworks(config_settings: ConfigSettings) -> Optional[List[str]]: + """Get the target frameworks from PEP 517 config settings.""" + settings: Mapping[str, ConfigValue] = config_settings or {} + value = settings.get("framework") + if value is None: + return None + + if isinstance(value, list): + if not value: + raise ValueError("Framework selection cannot be empty") + value = value[-1] + + return _normalize_frameworks(str(value)) + + +def _environment_frameworks() -> Optional[List[str]]: + """Get the target frameworks from the environment.""" + value = os.getenv("NVTE_FRAMEWORK") + return _normalize_frameworks(value) if value is not None else None + + +def _requested_frameworks(config_settings: ConfigSettings) -> Optional[List[str]]: + """Resolve an explicit framework selection in descending order of precedence.""" + configured = _config_frameworks(config_settings) + return configured if configured is not None else _environment_frameworks() + + +@contextmanager +def _framework_environment(config_settings: ConfigSettings) -> Iterator[None]: + """Apply a config-setting framework selection while running setuptools.""" + frameworks = _config_frameworks(config_settings) + if frameworks is None: + yield + return + + previous = os.environ.get("NVTE_FRAMEWORK") + os.environ["NVTE_FRAMEWORK"] = ",".join(frameworks) if frameworks else "none" + try: + yield + finally: + if previous is None: + os.environ.pop("NVTE_FRAMEWORK", None) + else: + os.environ["NVTE_FRAMEWORK"] = previous + + def _system_cuda_version() -> Optional[str]: """Get the version reported by a system-installed NVCC.""" candidates = [] @@ -131,15 +209,11 @@ def _torch_cuda_version() -> Optional[str]: return _normalize_cuda_version(value) if value else None -def _framework_cuda_version() -> Optional[str]: +def _framework_cuda_version(config_settings: ConfigSettings) -> Optional[str]: """Get a CUDA version required by the selected framework.""" - frameworks = { - framework.strip().lower() - for framework in os.getenv("NVTE_FRAMEWORK", "").split(",") - if framework.strip() - } + frameworks = _requested_frameworks(config_settings) - if not frameworks or frameworks.intersection({"all", "pytorch"}): + if frameworks is None or "pytorch" in frameworks: if cuda_version := _torch_cuda_version(): return cuda_version @@ -154,7 +228,7 @@ def _cuda_version(config_settings: ConfigSettings) -> str: _config_cuda_version(config_settings) or _environment_cuda_version() or _system_cuda_version() - or _framework_cuda_version() + or _framework_cuda_version(config_settings) or _DEFAULT_CUDA_VERSION ) @@ -171,6 +245,28 @@ def _cuda_build_requirements(config_settings: ConfigSettings) -> List[str]: return requirements +def _framework_build_requirements(config_settings: ConfigSettings) -> List[str]: + """Get build requirements for the requested frameworks.""" + frameworks = _requested_frameworks(config_settings) + if frameworks is None: + raise ValueError( + "Framework must be selected for an isolated build; set NVTE_FRAMEWORK or pass " + "--config-settings framework=pytorch (or jax, all, none)" + ) + + return [ + requirement + for framework in frameworks + for requirement in _FRAMEWORK_BUILD_PACKAGES[framework] + ] + + +def _build_requirements(config_settings: ConfigSettings) -> List[str]: + """Get CUDA and framework-specific build requirements.""" + framework_requirements = _framework_build_requirements(config_settings) + return _cuda_build_requirements(config_settings) + framework_requirements + + ################################################################################################### # PEP 517 and PEP 660 defined functions ################################################################################################### @@ -178,17 +274,17 @@ def _cuda_build_requirements(config_settings: ConfigSettings) -> List[str]: # Defined by PEP 517 def get_requires_for_build_wheel(config_settings: ConfigSettings = None) -> List[str]: - """Get CUDA requirements for building a wheel.""" - return _cuda_build_requirements(config_settings) + """Get requirements for building a wheel.""" + return _build_requirements(config_settings) # Defined by PEP 517 def get_requires_for_build_sdist(config_settings: ConfigSettings = None) -> List[str]: - """Get CUDA requirements needed while evaluating setup.py for an sdist.""" - # CUDA build requirements are needed even for sdist because setup.py currently always - # depends on these packages. It could be refactored to only depend on them when building - # a wheel, but that would require a more invasive refactor. - return _cuda_build_requirements(config_settings) + """Get requirements needed while evaluating setup.py for an sdist.""" + # Build requirements are needed even for sdist because setup.py currently always depends + # on them. It could be refactored to only depend on them when building a wheel, but that + # would require a more invasive refactor. + return _build_requirements(config_settings) # Defined by PEP 517 @@ -198,11 +294,12 @@ def build_wheel( metadata_directory: Optional[str] = None, ) -> str: """Build a wheel with setuptools.""" - return _SETUPTOOLS_BACKEND.build_wheel( - wheel_directory, - _setuptools_config_settings(config_settings), - metadata_directory, - ) + with _framework_environment(config_settings): + return _SETUPTOOLS_BACKEND.build_wheel( + wheel_directory, + _setuptools_config_settings(config_settings), + metadata_directory, + ) # Defined by PEP 517 @@ -211,10 +308,11 @@ def prepare_metadata_for_build_wheel( config_settings: ConfigSettings = None, ) -> str: """Prepare wheel metadata with setuptools.""" - return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_wheel( - metadata_directory, - _setuptools_config_settings(config_settings), - ) + with _framework_environment(config_settings): + return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_wheel( + metadata_directory, + _setuptools_config_settings(config_settings), + ) # Defined by PEP 517 @@ -223,18 +321,19 @@ def build_sdist( config_settings: ConfigSettings = None, ) -> str: """Build an sdist with setuptools.""" - return _SETUPTOOLS_BACKEND.build_sdist( - sdist_directory, - _setuptools_config_settings(config_settings), - ) + with _framework_environment(config_settings): + return _SETUPTOOLS_BACKEND.build_sdist( + sdist_directory, + _setuptools_config_settings(config_settings), + ) if hasattr(_SETUPTOOLS_BACKEND, "build_editable"): # Defined by PEP 660 def get_requires_for_build_editable(config_settings: ConfigSettings = None) -> List[str]: - """Get CUDA requirements for building an editable wheel.""" - return _cuda_build_requirements(config_settings) + """Get requirements for building an editable wheel.""" + return _build_requirements(config_settings) # Defined by PEP 660 def build_editable( @@ -243,11 +342,12 @@ def build_editable( metadata_directory: Optional[str] = None, ) -> str: """Build an editable wheel with setuptools.""" - return _SETUPTOOLS_BACKEND.build_editable( - wheel_directory, - _setuptools_config_settings(config_settings), - metadata_directory, - ) + with _framework_environment(config_settings): + return _SETUPTOOLS_BACKEND.build_editable( + wheel_directory, + _setuptools_config_settings(config_settings), + metadata_directory, + ) # Defined by PEP 660 def prepare_metadata_for_build_editable( @@ -255,7 +355,8 @@ def prepare_metadata_for_build_editable( config_settings: ConfigSettings = None, ) -> str: """Prepare editable-wheel metadata with setuptools.""" - return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_editable( - metadata_directory, - _setuptools_config_settings(config_settings), - ) + with _framework_environment(config_settings): + return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_editable( + metadata_directory, + _setuptools_config_settings(config_settings), + ) diff --git a/docs/envvars.rst b/docs/envvars.rst index 32fc3a2960f..7bb44beecb0 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -39,8 +39,8 @@ Build Configuration .. envvar:: NVTE_FRAMEWORK :Type: ``str`` - :Default: Auto-detected - :Description: Comma-separated list of frameworks to build support for (``pytorch``, ``jax``, ``all``, or ``none``). If not specified, automatically detects installed frameworks. + :Default: Auto-detected without build isolation; required with build isolation + :Description: Comma-separated list of frameworks to build support for (``pytorch``, ``jax``, ``all``, or ``none``). Isolated PEP 517 builds cannot inspect frameworks in the caller's environment, so an explicit selection is required through this variable or pip's ``--config-settings framework=`` option. The config setting takes precedence over this variable. .. envvar:: NVTE_USE_CCACHE diff --git a/pyproject.toml b/pyproject.toml index 4f7278a7fc4..3c1c890f78f 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,8 +3,17 @@ # See LICENSE for license information. [build-system] -# CUDA build-time requirements are defined in `build_tools/build_backend.py`. +requires = [ + "setuptools>=61.0", + "cmake>=3.21", + "wheel", + "pybind11[global]", + "ninja", + "nvidia-cudnn-frontend>=1.28.0", +] -# Select build dependencies with --config-setting cuda-version=. +# CUDA and framework-specific requirements are defined in `build_tools/build_backend.py`. +# Isolated builds must select a framework with --config-settings framework= +# or NVTE_FRAMEWORK. Select CUDA with --config-settings cuda-version=. build-backend = "build_tools.build_backend" backend-path = ["."] From ddb31786988f8ef2dedc4f5a9ce5004efd33e382 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 4 Sep 2026 14:26:05 -0500 Subject: [PATCH 04/24] Exclude cmake 3.25.0 due to but with finding cuda toolkit Signed-off-by: Fred Heinecke --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3c1c890f78f..c374023fe81 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,8 @@ [build-system] requires = [ "setuptools>=61.0", - "cmake>=3.21", + # CMake 3.25.0 breaks FindCUDAToolkit when nvptxcompiler_static is absent from CUDA wheels. + "cmake>=3.21,!=3.25.0", "wheel", "pybind11[global]", "ninja", From ecfe945a5e22357896a964c3a46b2507d1b035e0 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 4 Sep 2026 14:28:39 -0500 Subject: [PATCH 05/24] fixed cuda library resolution when using wheels as shared lib source Signed-off-by: Fred Heinecke --- build_tools/build_ext.py | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index 4b7ec153be1..acb5469c28d 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -70,10 +70,19 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: if discovered_cuda_home is not None: configure_command.append(f"-DCUDAToolkit_ROOT={discovered_cuda_home}") - # CUDA wheels use `lib`, while toolkit installations typically use - # `lib64`. Only override CMake's library discovery for a wheel-style - # layout with the expected libraries present. cuda_lib_dir = discovered_cuda_home / "lib" + if cuda_lib_dir.is_dir(): + # NVCC searches `lib64` for its static runtime libraries, while + # CUDA wheels install them in `lib`. + cuda_static_libraries = ( + cuda_lib_dir / "libcudadevrt.a", + cuda_lib_dir / "libcudart_static.a", + ) + if all(library.is_file() for library in cuda_static_libraries): + configure_command.append(f"-DCMAKE_CUDA_FLAGS=-L{cuda_lib_dir}") + + # Override CMake's shared-library discovery for each library that is + # present in the wheel-style layout. if cuda_lib_dir.is_dir() and (cuda_full_version := cuda_version()): cuda_major_version = cuda_full_version[0] cuda_libraries = { @@ -82,11 +91,11 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: "CUDA_cublas_LIBRARY": cuda_lib_dir / f"libcublas.so.{cuda_major_version}", "CUDA_cublasLt_LIBRARY": cuda_lib_dir / f"libcublasLt.so.{cuda_major_version}", } - if all(library.is_file() for library in cuda_libraries.values()): - configure_command.append(f"-DCMAKE_CUDA_FLAGS=-L{cuda_lib_dir}") - configure_command.extend( - f"-D{variable}={library}" for variable, library in cuda_libraries.items() - ) + configure_command.extend( + f"-D{variable}={library}" + for variable, library in cuda_libraries.items() + if library.is_file() + ) discovered_nvcc_path = nvcc_path() if discovered_nvcc_path is not None: From 1501892a3a505e1755033b237ac944130ad8492b Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 4 Sep 2026 14:47:38 -0500 Subject: [PATCH 06/24] Standardize CUDA lookup order based on feedback from #3251 Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 51 +++++++++++++----------------------- 1 file changed, 18 insertions(+), 33 deletions(-) diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index a943acad023..d8af39823dc 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -7,13 +7,13 @@ from contextlib import contextmanager import os import re -import shutil import subprocess -from pathlib import Path from typing import Iterator, List, Mapping, Optional, Union from setuptools import build_meta +from build_tools.utils import nvcc_path + _SETUPTOOLS_BACKEND = build_meta.__legacy__ _CUDA_BUILD_PACKAGES = { "12": [ @@ -165,38 +165,23 @@ def _framework_environment(config_settings: ConfigSettings) -> Iterator[None]: os.environ["NVTE_FRAMEWORK"] = previous -def _system_cuda_version() -> Optional[str]: - """Get the version reported by a system-installed NVCC.""" - candidates = [] - if cuda_compiler := os.getenv("CUDACXX"): - candidates.append(Path(cuda_compiler)) - if cuda_home := os.getenv("CUDA_HOME"): - candidates.append(Path(cuda_home) / "bin" / "nvcc") - if nvcc := shutil.which("nvcc"): - candidates.append(Path(nvcc)) - candidates.append(Path("/usr/local/cuda/bin/nvcc")) - - checked = set() - for candidate in candidates: - if candidate in checked or not candidate.is_file(): - continue - - checked.add(candidate) - try: - result = subprocess.run( - [candidate, "--version"], - capture_output=True, - check=True, - text=True, - ) - except (OSError, subprocess.CalledProcessError): - continue +def _nvcc_cuda_version() -> Optional[str]: + """Get the version reported by the selected NVCC.""" + if (cuda_compiler := nvcc_path()) is None: + return None - match = re.search(r"release\s+(\d+\.\d+)", result.stdout) - if match is not None: - return _normalize_cuda_version(match.group(1)) + try: + result = subprocess.run( + [cuda_compiler, "--version"], + capture_output=True, + check=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return None - return None + match = re.search(r"release\s+(\d+\.\d+)", result.stdout) + return _normalize_cuda_version(match.group(1)) if match is not None else None def _torch_cuda_version() -> Optional[str]: @@ -227,7 +212,7 @@ def _cuda_version(config_settings: ConfigSettings) -> str: return ( _config_cuda_version(config_settings) or _environment_cuda_version() - or _system_cuda_version() + or _nvcc_cuda_version() or _framework_cuda_version(config_settings) or _DEFAULT_CUDA_VERSION ) From aca8e2066091560d6f1cd7bcd05521f1513f3497 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 4 Sep 2026 14:47:59 -0500 Subject: [PATCH 07/24] Fix CUDA include dir resolution Signed-off-by: Fred Heinecke --- build_tools/utils.py | 39 +++++++++++++++++++++++---------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/build_tools/utils.py b/build_tools/utils.py index 6841a595802..c61f5274484 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -377,30 +377,37 @@ def cuda_toolkit_include_path() -> Optional[Path]: @functools.lru_cache(maxsize=None) -def get_cuda_include_dirs() -> Tuple[str, str]: - """Returns the CUDA header directory.""" +def get_cuda_include_dirs() -> List[Path]: + """Returns the CUDA header directories.""" force_wheels = bool(int(os.getenv("NVTE_BUILD_USE_NVIDIA_WHEELS", "0"))) - # If cuda is installed via toolkit, all necessary headers - # are bundled inside the top level cuda directory. - if not force_wheels and cuda_toolkit_include_path() is not None: - return [cuda_toolkit_include_path()] - - # Use pip wheels to include all headers. try: import nvidia - except ModuleNotFoundError as e: - raise RuntimeError("CUDA not found.") - - if nvidia.__file__ is not None: - cuda_root = Path(nvidia.__file__).parent + except ModuleNotFoundError: + nvidia_roots = [] else: - cuda_root = Path(nvidia.__path__[0]) # namespace - return [ + nvidia_roots = [Path(path).resolve() for path in nvidia.__path__] + + cuda_home = cuda_home_path() + cuda_home_from_wheel = cuda_home is not None and any( + cuda_home.is_relative_to(root) for root in nvidia_roots + ) + + # A system toolkit contains all necessary headers under one include + # directory. NVIDIA wheels split them across namespace packages. + if not force_wheels and not cuda_home_from_wheel: + if (include_path := cuda_toolkit_include_path()) is not None: + return [include_path] + + include_dirs = [ subdir / "include" - for subdir in cuda_root.iterdir() + for root in nvidia_roots + for subdir in root.iterdir() if subdir.is_dir() and (subdir / "include").is_dir() ] + if not include_dirs: + raise RuntimeError("CUDA not found.") + return include_dirs @functools.lru_cache(maxsize=None) From da9f4af0397bf1e833987451b90964a5c7035f75 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 6 Nov 2026 15:42:05 -0600 Subject: [PATCH 08/24] Require cuDNN frontend 1.29 for isolated builds Match the nvidia-cudnn-frontend version that main requires. Signed-off-by: Fred Heinecke --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c374023fe81..6a7bdae625e 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ requires = [ "wheel", "pybind11[global]", "ninja", - "nvidia-cudnn-frontend>=1.28.0", + "nvidia-cudnn-frontend>=1.29.0", ] # CUDA and framework-specific requirements are defined in `build_tools/build_backend.py`. From 28047c0cad196d8e56aa122f69d2cc175aec90fc Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 6 Nov 2026 15:42:39 -0600 Subject: [PATCH 09/24] Add build requirements and config settings for optional features Enable cuBLASMp, cuSOLVERMp, NVSHMEM, and CuTeDSL with either their NVTE_* variables or matching pip config settings (with-cublasmp, with-cusolvermp, enable-nvshmem, with-cutedsl). In isolated builds, install the corresponding Python packages as build requirements unless the feature's *_HOME variable selects an existing installation. Without build isolation, fall back to installed Python packages when *_HOME is not set. Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 121 ++++++++++++++++++++++++++++------- build_tools/pytorch.py | 10 +-- build_tools/utils.py | 51 +++++++++++++++ docs/envvars.rst | 23 ++++++- pyproject.toml | 7 +- setup.py | 24 ++++--- 6 files changed, 195 insertions(+), 41 deletions(-) diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index d8af39823dc..c8f636a06d9 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -2,13 +2,13 @@ # # See LICENSE for license information. -"""PEP 517 backend with CUDA-version- and framework-specific build requirements.""" +"""PEP 517 backend with CUDA-version-, framework-, and feature-specific build requirements.""" from contextlib import contextmanager import os import re import subprocess -from typing import Iterator, List, Mapping, Optional, Union +from typing import Iterator, List, Mapping, NamedTuple, Optional, Tuple, Union from setuptools import build_meta @@ -41,6 +41,35 @@ _SUPPORTED_FRAMEWORKS = tuple(_FRAMEWORK_BUILD_PACKAGES) _DEFAULT_CUDA_VERSION = "13.3" + +class _Feature(NamedTuple): + """Optional feature that can be enabled with an environment variable or config setting.""" + + environment_variable: str + # Selects an existing installation instead of installing the requirements. + home_variable: Optional[str] + # Formatted with the CUDA major version. + requirements: Tuple[str, ...] + + +# Keyed by config setting, e.g. --config-settings with-cublasmp=1. +_FEATURES = { + "with-cublasmp": _Feature( + "NVTE_WITH_CUBLASMP", + "CUBLASMP_HOME", + # cuBLASMp's overlap algorithms require NCCL 2.30. + ("nvidia-cublasmp-cu{cuda_major}>=0.8.1", "nvidia-nccl-cu{cuda_major}>=2.30"), + ), + "with-cusolvermp": _Feature( + "NVTE_WITH_CUSOLVERMP", "CUSOLVERMP_HOME", ("nvidia-cusolvermp-cu{cuda_major}",) + ), + "enable-nvshmem": _Feature( + "NVTE_ENABLE_NVSHMEM", "NVSHMEM_HOME", ("nvidia-nvshmem-cu{cuda_major}",) + ), + "with-cutedsl": _Feature("NVTE_WITH_CUTEDSL", None, ("apache-tvm-ffi>=0.1.12",)), +} +_PRIVATE_SETTINGS = ("cuda-version", "framework", *_FEATURES) + ConfigValue = Union[str, List[str]] ConfigSettings = Optional[Mapping[str, ConfigValue]] @@ -51,8 +80,8 @@ def _setuptools_config_settings(config_settings: ConfigSettings) -> ConfigSettin return None settings = dict(config_settings) - settings.pop("cuda-version", None) - settings.pop("framework", None) + for setting in _PRIVATE_SETTINGS: + settings.pop(setting, None) return settings or None @@ -146,23 +175,53 @@ def _requested_frameworks(config_settings: ConfigSettings) -> Optional[List[str] return configured if configured is not None else _environment_frameworks() +def _config_feature(config_settings: ConfigSettings, setting: str) -> Optional[bool]: + """Get whether PEP 517 config settings enable or disable an optional feature.""" + settings: Mapping[str, ConfigValue] = config_settings or {} + value = settings.get(setting) + if value is None: + return None + + if isinstance(value, list): + if not value: + raise ValueError(f"{setting} cannot be empty") + value = value[-1] + + if value not in ("0", "1"): + raise ValueError(f"Invalid value {value!r} for {setting}; expected 0 or 1") + return value == "1" + + +def _feature_enabled(config_settings: ConfigSettings, setting: str) -> bool: + """Resolve whether an optional feature is enabled, preferring config settings.""" + configured = _config_feature(config_settings, setting) + if configured is not None: + return configured + return bool(int(os.getenv(_FEATURES[setting].environment_variable, "0"))) + + @contextmanager -def _framework_environment(config_settings: ConfigSettings) -> Iterator[None]: - """Apply a config-setting framework selection while running setuptools.""" +def _build_environment(config_settings: ConfigSettings) -> Iterator[None]: + """Apply config-setting framework and feature selections while running setuptools.""" + overrides = {} frameworks = _config_frameworks(config_settings) - if frameworks is None: - yield - return - - previous = os.environ.get("NVTE_FRAMEWORK") - os.environ["NVTE_FRAMEWORK"] = ",".join(frameworks) if frameworks else "none" + if frameworks is not None: + overrides["NVTE_FRAMEWORK"] = ",".join(frameworks) if frameworks else "none" + for setting, feature in _FEATURES.items(): + enabled = _config_feature(config_settings, setting) + if enabled is not None: + overrides[feature.environment_variable] = "1" if enabled else "0" + + previous = {name: os.environ.get(name) for name in overrides} + os.environ.update(overrides) try: yield finally: - if previous is None: - os.environ.pop("NVTE_FRAMEWORK", None) - else: - os.environ["NVTE_FRAMEWORK"] = previous + for name, value in previous.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value def _nvcc_cuda_version() -> Optional[str]: @@ -246,10 +305,26 @@ def _framework_build_requirements(config_settings: ConfigSettings) -> List[str]: ] +def _feature_build_requirements(config_settings: ConfigSettings) -> List[str]: + """Get build requirements for the enabled optional features.""" + cuda_major = _cuda_version(config_settings).split(".", maxsplit=1)[0] + return [ + requirement.format(cuda_major=cuda_major) + for setting, feature in _FEATURES.items() + if _feature_enabled(config_settings, setting) + and not (feature.home_variable and os.getenv(feature.home_variable)) + for requirement in feature.requirements + ] + + def _build_requirements(config_settings: ConfigSettings) -> List[str]: - """Get CUDA and framework-specific build requirements.""" + """Get CUDA, framework, and optional feature build requirements.""" framework_requirements = _framework_build_requirements(config_settings) - return _cuda_build_requirements(config_settings) + framework_requirements + return ( + _cuda_build_requirements(config_settings) + + framework_requirements + + _feature_build_requirements(config_settings) + ) ################################################################################################### @@ -279,7 +354,7 @@ def build_wheel( metadata_directory: Optional[str] = None, ) -> str: """Build a wheel with setuptools.""" - with _framework_environment(config_settings): + with _build_environment(config_settings): return _SETUPTOOLS_BACKEND.build_wheel( wheel_directory, _setuptools_config_settings(config_settings), @@ -293,7 +368,7 @@ def prepare_metadata_for_build_wheel( config_settings: ConfigSettings = None, ) -> str: """Prepare wheel metadata with setuptools.""" - with _framework_environment(config_settings): + with _build_environment(config_settings): return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_wheel( metadata_directory, _setuptools_config_settings(config_settings), @@ -306,7 +381,7 @@ def build_sdist( config_settings: ConfigSettings = None, ) -> str: """Build an sdist with setuptools.""" - with _framework_environment(config_settings): + with _build_environment(config_settings): return _SETUPTOOLS_BACKEND.build_sdist( sdist_directory, _setuptools_config_settings(config_settings), @@ -327,7 +402,7 @@ def build_editable( metadata_directory: Optional[str] = None, ) -> str: """Build an editable wheel with setuptools.""" - with _framework_environment(config_settings): + with _build_environment(config_settings): return _SETUPTOOLS_BACKEND.build_editable( wheel_directory, _setuptools_config_settings(config_settings), @@ -340,7 +415,7 @@ def prepare_metadata_for_build_editable( config_settings: ConfigSettings = None, ) -> str: """Prepare editable-wheel metadata with setuptools.""" - with _framework_environment(config_settings): + with _build_environment(config_settings): return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_editable( metadata_directory, _setuptools_config_settings(config_settings), diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 173f9cee708..b5588194ba2 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -17,6 +17,7 @@ debug_build_enabled, get_bolt_build_flags, nccl_ep_enabled, + nvshmem_home_path, setup_mpi_flags, ) from typing import List @@ -109,10 +110,11 @@ def setup_pytorch_extension( library_dirs = [] libraries = [] 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" - nvshmem_home = Path(os.getenv("NVSHMEM_HOME")) + nvshmem_home = nvshmem_home_path() + assert nvshmem_home is not None, ( + "NVSHMEM_HOME must be set, or the nvidia-nvshmem-cu* Python package installed, when" + " compiling with NVTE_ENABLE_NVSHMEM=1" + ) include_dirs.append(nvshmem_home / "include") library_dirs.append(nvshmem_home / "lib") libraries.append("nvshmem_host") diff --git a/build_tools/utils.py b/build_tools/utils.py index c61f5274484..d53d26b1841 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -366,6 +366,57 @@ def nccl_lib_path() -> Optional[Path]: return None +def _python_package_path(package_name: str, relative_path: str) -> Optional[Path]: + """Return a path inside an installed Python package, or `None` if it does not exist.""" + try: + package_distribution = distribution(package_name) + except PackageNotFoundError: + return None + + path = Path(package_distribution.locate_file(relative_path)) + return path if path.exists() else None + + +def cublasmp_home_path() -> Optional[Path]: + """Return the cuBLASMp installation root. + + Prefers CUBLASMP_HOME, then the nvidia-cublasmp-cu* Python package. + Returns `None` if cuBLASMp is not found.""" + if cublasmp_home := os.getenv("CUBLASMP_HOME"): + return Path(cublasmp_home) + + cuda_major = cuda_version()[0] + return _python_package_path( + f"nvidia-cublasmp-cu{cuda_major}", f"nvidia/cublasmp/cu{cuda_major}" + ) + + +def cusolvermp_home_path() -> Path: + """Return the cuSOLVERMp installation root. + + Prefers CUSOLVERMP_HOME, then the nvidia-cusolvermp-cu* Python package, + then `/usr`.""" + if cusolvermp_home := os.getenv("CUSOLVERMP_HOME"): + return Path(cusolvermp_home) + + cuda_major = cuda_version()[0] + package_path = _python_package_path( + f"nvidia-cusolvermp-cu{cuda_major}", f"nvidia/cu{cuda_major}" + ) + return package_path if package_path is not None else Path("/usr") + + +def nvshmem_home_path() -> Optional[Path]: + """Return the NVSHMEM installation root. + + Prefers NVSHMEM_HOME, then the nvidia-nvshmem-cu* Python package. + Returns `None` if NVSHMEM is not found.""" + if nvshmem_home := os.getenv("NVSHMEM_HOME"): + return Path(nvshmem_home) + + return _python_package_path(f"nvidia-nvshmem-cu{cuda_version()[0]}", "nvidia/nvshmem") + + @functools.lru_cache(maxsize=None) def cuda_toolkit_include_path() -> Optional[Path]: """Return the CUDA Toolkit include directory.""" diff --git a/docs/envvars.rst b/docs/envvars.rst index 7bb44beecb0..94b80a7942b 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -87,6 +87,13 @@ Build Configuration Optional Dependencies ^^^^^^^^^^^^^^^^^^^^^ +Each of the following options can also be set for an isolated PEP 517 build with pip's +``--config-settings =<0 or 1>`` option, which takes precedence over the variable. +In an isolated build, the Python package listed for an option is installed as a build +requirement unless the option's installation directory variable is set. Without build +isolation, the installation directory variable is used if set, and the package is used if +installed. Not intended for external use. + .. envvar:: NVTE_UB_WITH_MPI :Type: ``int`` (0 or 1) @@ -97,14 +104,26 @@ Optional Dependencies :Type: ``int`` (0 or 1) :Default: ``0`` - :Description: Enable NVSHMEM support. When set to ``1``, requires ``NVSHMEM_HOME`` to be set to the NVSHMEM installation directory. + :Description: Enable NVSHMEM support. Requires the NVSHMEM installation directory in ``NVSHMEM_HOME`` or the ``nvidia-nvshmem-cu`` Python package. Config setting: ``enable-nvshmem``. + +.. envvar:: NVTE_WITH_CUBLASMP + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Use cuBLASMp for tensor-parallel GEMMs. Requires cuBLASMp 0.8.1 or newer, from the installation directory in ``CUBLASMP_HOME`` or the ``nvidia-cublasmp-cu`` Python package, and NCCL 2.30 or newer. Config setting: ``with-cublasmp``. + +.. envvar:: NVTE_WITH_CUSOLVERMP + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Use cuSOLVERMp for distributed Newton-Schulz. Requires cuSOLVERMp from the installation directory in ``CUSOLVERMP_HOME``, the ``nvidia-cusolvermp-cu`` Python package, or ``/usr``, in that order. Config setting: ``with-cusolvermp``. .. envvar:: NVTE_WITH_CUTEDSL :Type: ``int`` (0 or 1) :Default: ``0`` :Description: Enable CuTeDSL support when building ``libtransformer_engine.so``. When set to ``1``, requires - CPython Stable ABIs from Python 3.10 and above, and ``apache-tvm-ffi`` installed. + CPython Stable ABIs from Python 3.10 and above, and ``apache-tvm-ffi`` installed. Config setting: ``with-cutedsl``. Note: when linking ``libtransformer_engine.so`` built with ``NVTE_WITH_CUTEDSL=1`` to a standalone C++ executable, you must provide the definition of these CPython symbols. See ``tests/cpp/CMakeLists.txt`` for an example. diff --git a/pyproject.toml b/pyproject.toml index 6a7bdae625e..0687f3d8524 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,8 +13,9 @@ requires = [ "nvidia-cudnn-frontend>=1.29.0", ] -# CUDA and framework-specific requirements are defined in `build_tools/build_backend.py`. -# Isolated builds must select a framework with --config-settings framework= -# or NVTE_FRAMEWORK. Select CUDA with --config-settings cuda-version=. +# CUDA, framework, and optional feature requirements are defined in +# `build_tools/build_backend.py`. Isolated builds must select a framework with +# --config-settings framework= or NVTE_FRAMEWORK. Select CUDA with +# --config-settings cuda-version=. build-backend = "build_tools.build_backend" backend-path = ["."] diff --git a/setup.py b/setup.py index 2dabf37b322..ede76359c9b 100644 --- a/setup.py +++ b/setup.py @@ -18,10 +18,11 @@ from build_tools.build_ext import CMakeExtension, get_build_ext from build_tools.te_version import te_version from build_tools.utils import ( + cublasmp_home_path, cuda_archs, cuda_home_path, - cuda_version, cudnn_frontend_include_path, + cusolvermp_home_path, get_bolt_build_flags, get_frameworks, remove_dups, @@ -29,6 +30,7 @@ nccl_ep_enabled, get_max_jobs_for_parallel_build, nvcc_path, + nvshmem_home_path, ) frameworks = get_frameworks() @@ -81,25 +83,29 @@ def setup_common_extension() -> CMakeExtension: cmake_flags.append("-DNVTE_UB_WITH_MPI=ON") 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" + nvshmem_home = nvshmem_home_path() + assert nvshmem_home is not None, ( + "NVSHMEM_HOME must be set, or the nvidia-nvshmem-cu* Python package installed, when" + " compiling with NVTE_ENABLE_NVSHMEM=1" + ) cmake_flags.append("-DNVTE_ENABLE_NVSHMEM=ON") + cmake_flags.append(f"-DNVSHMEM_HOME={nvshmem_home}") 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"))): + cublasmp_dir = cublasmp_home_path() + assert cublasmp_dir is not None, ( + "CUBLASMP_HOME must be set, or the nvidia-cublasmp-cu* Python package installed, when" + " compiling with NVTE_WITH_CUBLASMP=1" + ) 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}") + cmake_flags.append(f"-DCUSOLVERMP_DIR={cusolvermp_home_path()}") # NCCL EP (Hopper+): on by default; auto-skipped when no arch >= 90 is # targeted. Set NVTE_WITH_NCCL_EP=0 to force off. From 412650776e96caa5a51c031044bcf3cbd7596424 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 6 Nov 2026 15:42:46 -0600 Subject: [PATCH 10/24] Support the openmpi Python package for MPI bootstrapping Add the ub-with-mpi config setting for NVTE_UB_WITH_MPI. Isolated builds install the openmpi package unless MPI_HOME is set, and builds without isolation use an installed openmpi package when MPI_HOME is unset. The package installs into the environment prefix, so locate it through its installed mpi.h. Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 1 + build_tools/utils.py | 29 +++++++++++++++++++++++++---- docs/envvars.rst | 2 +- setup.py | 10 +++++++--- 4 files changed, 34 insertions(+), 8 deletions(-) diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index c8f636a06d9..5b884c91aea 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -66,6 +66,7 @@ class _Feature(NamedTuple): "enable-nvshmem": _Feature( "NVTE_ENABLE_NVSHMEM", "NVSHMEM_HOME", ("nvidia-nvshmem-cu{cuda_major}",) ), + "ub-with-mpi": _Feature("NVTE_UB_WITH_MPI", "MPI_HOME", ("openmpi",)), "with-cutedsl": _Feature("NVTE_WITH_CUTEDSL", None, ("apache-tvm-ffi>=0.1.12",)), } _PRIVATE_SETTINGS = ("cuda-version", "framework", *_FEATURES) diff --git a/build_tools/utils.py b/build_tools/utils.py index d53d26b1841..98c9e64cb45 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -417,6 +417,26 @@ def nvshmem_home_path() -> Optional[Path]: return _python_package_path(f"nvidia-nvshmem-cu{cuda_version()[0]}", "nvidia/nvshmem") +def mpi_home_path() -> Optional[Path]: + """Return the MPI installation root. + + Prefers MPI_HOME, then the openmpi Python package, which installs into the + environment prefix rather than site-packages. Returns `None` if MPI is not found.""" + if mpi_home := os.getenv("MPI_HOME"): + return Path(mpi_home) + + try: + openmpi_distribution = distribution("openmpi") + except PackageNotFoundError: + return None + + for package_path in openmpi_distribution.files or []: + if package_path.name == "mpi.h" and package_path.parent.name == "include": + return Path(openmpi_distribution.locate_file(package_path)).resolve().parent.parent + + return None + + @functools.lru_cache(maxsize=None) def cuda_toolkit_include_path() -> Optional[Path]: """Return the CUDA Toolkit include directory.""" @@ -609,10 +629,11 @@ def get_frameworks() -> List[str]: def setup_mpi_flags(include_dirs: List, cxx_flags: List) -> None: """Add MPI include path and compile definition if NVTE_UB_WITH_MPI is enabled.""" if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))): - assert ( - os.getenv("MPI_HOME") is not None - ), "MPI_HOME=/path/to/mpi must be set when compiling with NVTE_UB_WITH_MPI=1!" - mpi_path = Path(os.getenv("MPI_HOME")) + mpi_path = mpi_home_path() + assert mpi_path is not None, ( + "MPI_HOME=/path/to/mpi must be set, or the openmpi Python package installed, " + "when compiling with NVTE_UB_WITH_MPI=1!" + ) include_dirs.append(mpi_path / "include") cxx_flags.append("-DNVTE_UB_WITH_MPI") diff --git a/docs/envvars.rst b/docs/envvars.rst index 94b80a7942b..b606bf86552 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -98,7 +98,7 @@ installed. Not intended for external use. :Type: ``int`` (0 or 1) :Default: ``0`` - :Description: Enable MPI support for userbuffers. When set to ``1``, requires ``MPI_HOME`` to be set to the MPI installation directory. + :Description: Enable MPI support for userbuffers. Requires the MPI installation directory in ``MPI_HOME`` or the ``openmpi`` Python package. Config setting: ``ub-with-mpi``. .. envvar:: NVTE_ENABLE_NVSHMEM diff --git a/setup.py b/setup.py index ede76359c9b..5dc4808c7b9 100644 --- a/setup.py +++ b/setup.py @@ -27,6 +27,7 @@ get_frameworks, remove_dups, min_python_version_str, + mpi_home_path, nccl_ep_enabled, get_max_jobs_for_parallel_build, nvcc_path, @@ -77,10 +78,13 @@ def setup_common_extension() -> CMakeExtension: cmake_flags.append("-DNVTE_WITH_CUTEDSL=OFF") 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" + mpi_home = mpi_home_path() + assert mpi_home is not None, ( + "MPI_HOME must be set, or the openmpi Python package installed, when compiling with" + " NVTE_UB_WITH_MPI=1" + ) cmake_flags.append("-DNVTE_UB_WITH_MPI=ON") + cmake_flags.append(f"-DMPI_HOME={mpi_home}") if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", "0"))): nvshmem_home = nvshmem_home_path() From 3360be9a413f48c4aee256de9dacf0d333271075 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 6 Nov 2026 15:42:52 -0600 Subject: [PATCH 11/24] Find versioned shared libraries from Python packages The cuSOLVERMp, NVSHMEM, and NCCL Python packages only provide versioned shared libraries (libcusolverMp.so.0, libnvshmem_host.so.3, libnccl.so.2), which CMake's find_library and -l linker flags do not match. Look for the versioned names as well, and find NCCL's library next to its headers. Signed-off-by: Fred Heinecke --- build_tools/pytorch.py | 6 +++++- transformer_engine/common/CMakeLists.txt | 18 ++++++++++++------ .../common/nvshmem_api/CMakeLists.txt | 9 ++++++++- 3 files changed, 25 insertions(+), 8 deletions(-) diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index b5588194ba2..5819c6937cf 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -117,7 +117,11 @@ def setup_pytorch_extension( ) include_dirs.append(nvshmem_home / "include") library_dirs.append(nvshmem_home / "lib") - libraries.append("nvshmem_host") + if (nvshmem_home / "lib" / "libnvshmem_host.so").is_file(): + libraries.append("nvshmem_host") + else: + # NVSHMEM Python packages only provide the versioned shared library. + libraries.append(":libnvshmem_host.so.3") cxx_flags.append("-DNVTE_ENABLE_NVSHMEM") if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))): diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index f90f74449c5..55d208d703c 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -534,9 +534,11 @@ if (NVTE_WITH_CUBLASMP) PATHS ${CUBLASMP_DIR} PATH_SUFFIXES lib lib64 lib/aarch64-linux-gnu lib/sbsa-linux-gnu lib/x86_64-linux-gnu REQUIRED) + # NCCL Python packages only provide the versioned shared library. find_library(NCCL_LIB - NAMES nccl libnccl - PATH_SUFFIXES lib + NAMES nccl libnccl.so.2 + HINTS "${NCCL_INCLUDE_DIR}/.." + PATH_SUFFIXES lib lib64 REQUIRED) # cuBLASMp 0.8.1 adds MXFP8 support to the CUDA-graph-safe overlap algos, # and NCCL 2.30 is the first release with graph-safe one-sided RMA @@ -561,14 +563,16 @@ option(NVTE_WITH_CUSOLVERMP "Use cuSolverMp for distributed Newton-Schulz" OFF) if (NVTE_WITH_CUSOLVERMP) target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUSOLVERMP) target_include_directories(transformer_engine PRIVATE ${CUSOLVERMP_DIR}/include) + # cuSOLVERMp and NCCL Python packages only provide the versioned shared libraries. find_library(CUSOLVERMP_LIB - NAMES cusolverMp libcusolverMp + NAMES cusolverMp libcusolverMp.so.0 PATHS ${CUSOLVERMP_DIR} PATH_SUFFIXES lib REQUIRED) find_library(NCCL_LIB - NAMES nccl libnccl - PATH_SUFFIXES lib + NAMES nccl libnccl.so.2 + HINTS "${NCCL_INCLUDE_DIR}/.." + PATH_SUFFIXES lib lib64 REQUIRED) target_link_libraries(transformer_engine PRIVATE ${NCCL_LIB} ${CUSOLVERMP_LIB}) message(STATUS "Using cuSolverMp at: ${CUSOLVERMP_DIR}") @@ -608,8 +612,10 @@ file(REAL_PATH "${NCCL_EP_SHARED_LINK}" NCCL_EP_SHARED_LIB) # -- NCCL core library ------------------------------------------------------- if(NOT NCCL_LIB) + # NCCL Python packages only provide the versioned shared library. find_library(NCCL_LIB - NAMES nccl libnccl + NAMES nccl libnccl.so.2 + HINTS "${NCCL_INCLUDE_DIR}/.." PATH_SUFFIXES lib lib64 REQUIRED) endif() diff --git a/transformer_engine/common/nvshmem_api/CMakeLists.txt b/transformer_engine/common/nvshmem_api/CMakeLists.txt index 3d9b6b5ec44..54d22732485 100644 --- a/transformer_engine/common/nvshmem_api/CMakeLists.txt +++ b/transformer_engine/common/nvshmem_api/CMakeLists.txt @@ -13,8 +13,15 @@ set(NVSHMEM_HOME "$ENV{NVSHMEM_HOME}" CACHE STRING "Location of NVSHMEM installa add_library(nvshmemapi STATIC nvshmem_waitkernel.cu) set(NVSHMEMAPI_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}" PARENT_SCOPE) +# NVSHMEM Python packages only provide the versioned host library. +find_library(NVSHMEM_HOST_LIB + NAMES nvshmem_host libnvshmem_host.so.3 + PATHS ${NVSHMEM_HOME} + PATH_SUFFIXES lib lib64 + NO_DEFAULT_PATH + REQUIRED) target_link_directories(nvshmemapi PUBLIC ${NVSHMEM_HOME}/lib) -target_link_libraries(nvshmemapi PUBLIC -static-libstdc++ nvshmem_device nvshmem_host CUDA::nvml CUDA::cublas CUDA::cuda_driver) +target_link_libraries(nvshmemapi PUBLIC -static-libstdc++ nvshmem_device ${NVSHMEM_HOST_LIB} CUDA::nvml CUDA::cublas CUDA::cuda_driver) target_include_directories(nvshmemapi PRIVATE ${NVSHMEM_HOME}/include/ ${CMAKE_CURRENT_SOURCE_DIR}/../include) From db6548ee790838c782a34697d341797a0b3c85e2 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 6 Nov 2026 15:42:58 -0600 Subject: [PATCH 12/24] Set RPATH for dependencies installed from Python packages Add RUNPATH entries to libtransformer_engine.so for the cuBLASMp, cuSOLVERMp, NCCL, NVSHMEM, and openmpi Python packages when the corresponding features are enabled, relative to the installed library's location in site-packages. NCCL is included whenever the library links it, including the default NCCL EP build, so an NCCL Python package installed next to Transformer Engine takes precedence over a system NCCL. Signed-off-by: Fred Heinecke --- build_tools/build_ext.py | 32 +++++++++++++++++------- docs/envvars.rst | 3 ++- setup.py | 9 ++++--- transformer_engine/common/CMakeLists.txt | 28 +++++++++++++++++++++ 4 files changed, 58 insertions(+), 14 deletions(-) diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index acb5469c28d..b6559b80d54 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -44,7 +44,12 @@ def __init__( self.cmake_path: Path = cmake_path self.cmake_flags: List[str] = [] if cmake_flags is None else cmake_flags - def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: + def _build_cmake( + self, + build_dir: Path, + install_dir: Path, + site_packages_rpath: Optional[str] = None, + ) -> None: # Make sure paths are str _cmake_bin = str(cmake_bin()) cmake_path = str(self.cmake_path) @@ -108,6 +113,8 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: f"-DCMAKE_CUDA_COMPILER_LAUNCHER={ccache_bin}", ] configure_command += self.cmake_flags + if site_packages_rpath is not None: + configure_command.append(f"-DNVTE_SITE_PACKAGES_RPATH={site_packages_rpath}") import pybind11 @@ -147,6 +154,19 @@ class _CMakeBuildExtension(extension_cls): """Setuptools command with support for CMake extension modules""" def run(self) -> None: + # Ensure that shared objects files for source and PyPI installations live + # in separate directories to avoid conflicts during install and runtime. + lib_dir = ( + "wheel_lib" + if bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))) or framework_extension_only + else "" + ) + + # Non-inplace builds move shared objects into site-packages/transformer_engine/. + site_packages_rpath = None + if not self.inplace: + site_packages_rpath = "$ORIGIN/../.." if lib_dir else "$ORIGIN/.." + # Build CMake extensions for ext in self.extensions: package_path = Path(self.get_ext_fullpath(ext.name)) @@ -167,6 +187,7 @@ def run(self) -> None: ext._build_cmake( build_dir=build_dir, install_dir=install_dir, + site_packages_rpath=site_packages_rpath, ) continue @@ -181,6 +202,7 @@ def run(self) -> None: ext._build_cmake( build_dir=Path(build_dir), install_dir=install_dir, + site_packages_rpath=site_packages_rpath, ) # Build non-CMake extensions as usual @@ -191,14 +213,6 @@ def run(self) -> None: super().run() self.extensions = all_extensions - # Ensure that shared objects files for source and PyPI installations live - # in separate directories to avoid conflicts during install and runtime. - lib_dir = ( - "wheel_lib" - if bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))) or framework_extension_only - else "" - ) - # Ensure that binaries are not in global package space. # For editable/inplace builds this is not a concern as # the SOs will be in a local directory anyway. diff --git a/docs/envvars.rst b/docs/envvars.rst index b606bf86552..cc95b3d4f90 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -92,7 +92,8 @@ Each of the following options can also be set for an isolated PEP 517 build with In an isolated build, the Python package listed for an option is installed as a build requirement unless the option's installation directory variable is set. Without build isolation, the installation directory variable is used if set, and the package is used if -installed. Not intended for external use. +installed. ``libtransformer_engine.so`` finds these packages at runtime when they are installed +in the same environment. Not intended for external use. .. envvar:: NVTE_UB_WITH_MPI diff --git a/setup.py b/setup.py index 5dc4808c7b9..8031f8031bb 100644 --- a/setup.py +++ b/setup.py @@ -185,10 +185,11 @@ def _discover_nccl_home() -> str: 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. + # EP build links against the same libnccl that gets loaded. Unless an nvidia-nccl + # wheel is installed next to it, which libtransformer_engine's RUNPATH prefers, + # the loader uses ldconfig/system paths; building against a different NCCL 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(): diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 55d208d703c..b1e3608e84f 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -844,6 +844,34 @@ endif() set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --threads ${BUILD_THREADS_PER_JOB}") message(STATUS "Threads per parallel build job: ${BUILD_THREADS_PER_JOB}") +# Find optional dependencies that are installed from Python packages into the +# same environment as the installed library. +set(NVTE_SITE_PACKAGES_RPATH "" CACHE STRING + "Python site-packages directory relative to the installed library, e.g. \$ORIGIN/..") +if (NVTE_SITE_PACKAGES_RPATH) + set(nvte_python_package_rpath) + if (NVTE_WITH_CUBLASMP) + list(APPEND nvte_python_package_rpath + "${NVTE_SITE_PACKAGES_RPATH}/nvidia/cublasmp/cu${CUDAToolkit_VERSION_MAJOR}/lib") + endif() + if (NVTE_WITH_CUSOLVERMP) + list(APPEND nvte_python_package_rpath + "${NVTE_SITE_PACKAGES_RPATH}/nvidia/cu${CUDAToolkit_VERSION_MAJOR}/lib") + endif() + if (NVTE_WITH_NCCL_EP OR NVTE_WITH_CUBLASMP OR NVTE_WITH_CUSOLVERMP) + list(APPEND nvte_python_package_rpath "${NVTE_SITE_PACKAGES_RPATH}/nvidia/nccl/lib") + endif() + if (NVTE_ENABLE_NVSHMEM) + list(APPEND nvte_python_package_rpath "${NVTE_SITE_PACKAGES_RPATH}/nvidia/nvshmem/lib") + endif() + if (NVTE_UB_WITH_MPI) + # The openmpi package installs into /lib, two levels above site-packages. + list(APPEND nvte_python_package_rpath "${NVTE_SITE_PACKAGES_RPATH}/../..") + endif() + set_target_properties(transformer_engine PROPERTIES + INSTALL_RPATH "${nvte_python_package_rpath}") +endif() + # Install library install(TARGETS transformer_engine DESTINATION .) if(NVTE_WITH_NCCL_EP) From bb92a24ef1cc8f7dcaf1c30533225e24a924d857 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 6 Nov 2026 15:43:05 -0600 Subject: [PATCH 13/24] Build release core wheels with the isolated build backend Build the core library wheel with pip wheel so build_tools/build_backend.py installs its CUDA and optional feature build requirements, instead of installing cuDNN frontend and TVM-FFI into the build interpreter. Signed-off-by: Fred Heinecke --- build_tools/wheel_utils/build_wheels.sh | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh index 761f12390e1..72467caf200 100644 --- a/build_tools/wheel_utils/build_wheels.sh +++ b/build_tools/wheel_utils/build_wheels.sh @@ -23,8 +23,9 @@ cd /TransformerEngine git checkout $TARGET_BRANCH git submodule update --init --recursive -# Install deps -/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.25.0' 'apache-tvm-ffi>=0.1.12' +# Install deps. The core library wheel is built in an isolated environment, but the +# metapackage, wheel repacking, and framework sdists still use this interpreter. +/opt/python/cp310-cp310/bin/pip install pybind11[global] setuptools wheel if $BUILD_METAPACKAGE ; then cd /TransformerEngine @@ -44,8 +45,12 @@ if $BUILD_COMMON ; then WHEEL_ABI_TAG=none fi - # Create the wheel. - /opt/python/cp310-cp310/bin/python setup.py bdist_wheel --verbose --python-tag=$WHEEL_PYTHON_TAG --plat-name=$PLATFORM 2>&1 | tee /wheelhouse/logs/common.txt + # Create the wheel in an isolated build environment. build_tools/build_backend.py + # installs the CUDA and optional feature (NVTE_WITH_*) build requirements. + /opt/python/cp310-cp310/bin/python -m pip wheel . --no-deps --wheel-dir dist --verbose \ + --config-settings framework=none \ + --config-settings=--build-option=--python-tag=$WHEEL_PYTHON_TAG \ + --config-settings=--build-option=--plat-name=$PLATFORM 2>&1 | tee /wheelhouse/logs/common.txt # Repack the wheel for specific cuda version. /opt/python/cp310-cp310/bin/wheel unpack dist/* From 152b4fb0d233be7e8961b5181eeee2b8840b1b9e Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 13:53:00 -0600 Subject: [PATCH 14/24] Build only the core library when no framework is selected in isolated builds Isolated builds previously failed unless NVTE_FRAMEWORK or the framework config setting was set. Instead, add no framework build requirements, so setup.py's framework detection finds none in the isolated build environment and only the core library is built. Builds without isolation still detect installed frameworks. Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 7 +++---- docs/envvars.rst | 4 ++-- pyproject.toml | 4 ++-- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index 5b884c91aea..61c3f627b9f 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -294,10 +294,9 @@ def _framework_build_requirements(config_settings: ConfigSettings) -> List[str]: """Get build requirements for the requested frameworks.""" frameworks = _requested_frameworks(config_settings) if frameworks is None: - raise ValueError( - "Framework must be selected for an isolated build; set NVTE_FRAMEWORK or pass " - "--config-settings framework=pytorch (or jax, all, none)" - ) + # Without a selection, setup.py detects frameworks in the isolated build + # environment, which has none, so only the core library is built. + return [] return [ requirement diff --git a/docs/envvars.rst b/docs/envvars.rst index cc95b3d4f90..4f8e66cc2a8 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -39,8 +39,8 @@ Build Configuration .. envvar:: NVTE_FRAMEWORK :Type: ``str`` - :Default: Auto-detected without build isolation; required with build isolation - :Description: Comma-separated list of frameworks to build support for (``pytorch``, ``jax``, ``all``, or ``none``). Isolated PEP 517 builds cannot inspect frameworks in the caller's environment, so an explicit selection is required through this variable or pip's ``--config-settings framework=`` option. The config setting takes precedence over this variable. + :Default: Auto-detected without build isolation; ``none`` with build isolation + :Description: Comma-separated list of frameworks to build support for (``pytorch``, ``jax``, ``all``, or ``none``). Isolated PEP 517 builds cannot inspect frameworks in the caller's environment, so they only build the core library unless frameworks are selected through this variable or pip's ``--config-settings framework=`` option. The config setting takes precedence over this variable. .. envvar:: NVTE_USE_CCACHE diff --git a/pyproject.toml b/pyproject.toml index 0687f3d8524..7be88f7988b 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,8 +14,8 @@ requires = [ ] # CUDA, framework, and optional feature requirements are defined in -# `build_tools/build_backend.py`. Isolated builds must select a framework with -# --config-settings framework= or NVTE_FRAMEWORK. Select CUDA with +# `build_tools/build_backend.py`. Isolated builds only build framework support +# selected with --config-settings framework= or NVTE_FRAMEWORK. Select CUDA with # --config-settings cuda-version=. build-backend = "build_tools.build_backend" backend-path = ["."] From c23b2150ce71e773b588fd94649f71dbaef094d4 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 16:44:00 -0600 Subject: [PATCH 15/24] Require an explicit CUDA version when none can be detected in isolated builds Remove the hardcoded CUDA 13.3 fallback. A guessed CUDA version can build successfully but fail at runtime against an older driver or a framework built for a different CUDA major version. The version is still detected from the cuda-version config setting, NVTE_CUDA_VERSION, or the CUDA Toolkit's nvcc; otherwise the build fails with instructions. Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 11 ++++++++--- docs/envvars.rst | 4 ++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index 61c3f627b9f..8bf227d0f1f 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -39,7 +39,6 @@ "jax": ["jax>=0.5.0", "flax>=0.7.1"], } _SUPPORTED_FRAMEWORKS = tuple(_FRAMEWORK_BUILD_PACKAGES) -_DEFAULT_CUDA_VERSION = "13.3" class _Feature(NamedTuple): @@ -269,13 +268,19 @@ def _framework_cuda_version(config_settings: ConfigSettings) -> Optional[str]: def _cuda_version(config_settings: ConfigSettings) -> str: """Resolve the CUDA version in descending order of precedence.""" - return ( + cuda_version = ( _config_cuda_version(config_settings) or _environment_cuda_version() or _nvcc_cuda_version() or _framework_cuda_version(config_settings) - or _DEFAULT_CUDA_VERSION ) + if cuda_version is None: + raise ValueError( + "Could not determine the CUDA version for an isolated build; install a CUDA " + "toolkit, set CUDA_HOME, or pass --config-settings cuda-version=. " + "(or set NVTE_CUDA_VERSION)" + ) + return cuda_version def _cuda_build_requirements(config_settings: ConfigSettings) -> List[str]: diff --git a/docs/envvars.rst b/docs/envvars.rst index 4f8e66cc2a8..5b2b5fa25fe 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -81,8 +81,8 @@ Build Configuration .. envvar:: NVTE_CUDA_VERSION :Type: ``str`` - :Default: Auto-detected based on system CUDA version, falling back to ``"13.3"`` - :Description: CUDA version to build for. Valid values: ``"12.*"``, ``"13.*"``. If not set, automatically determined based on the installed CUDA Toolkit version. Not intended for external use. + :Default: Auto-detected based on system CUDA version + :Description: CUDA version to build for. Valid values: ``"12.*"``, ``"13.*"``. If not set, automatically determined based on the installed CUDA Toolkit version. Isolated builds fail if no CUDA Toolkit is found and neither this variable nor pip's ``--config-settings cuda-version=.`` option is set. Not intended for external use. Optional Dependencies ^^^^^^^^^^^^^^^^^^^^^ From 31ee7697ba6c587de93d9f3663da4f94c7a686a1 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 16:44:49 -0600 Subject: [PATCH 16/24] Remove 'not intended for external use' from build isolation docs Signed-off-by: Fred Heinecke --- docs/envvars.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 5b2b5fa25fe..1a748f7e988 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -82,7 +82,7 @@ Build Configuration :Type: ``str`` :Default: Auto-detected based on system CUDA version - :Description: CUDA version to build for. Valid values: ``"12.*"``, ``"13.*"``. If not set, automatically determined based on the installed CUDA Toolkit version. Isolated builds fail if no CUDA Toolkit is found and neither this variable nor pip's ``--config-settings cuda-version=.`` option is set. Not intended for external use. + :Description: CUDA version to build for. Valid values: ``"12.*"``, ``"13.*"``. If not set, automatically determined based on the installed CUDA Toolkit version. Isolated builds fail if no CUDA Toolkit is found and neither this variable nor pip's ``--config-settings cuda-version=.`` option is set. Optional Dependencies ^^^^^^^^^^^^^^^^^^^^^ @@ -93,7 +93,7 @@ In an isolated build, the Python package listed for an option is installed as a requirement unless the option's installation directory variable is set. Without build isolation, the installation directory variable is used if set, and the package is used if installed. ``libtransformer_engine.so`` finds these packages at runtime when they are installed -in the same environment. Not intended for external use. +in the same environment. .. envvar:: NVTE_UB_WITH_MPI From faa1b15fa7c61f89afd07bcb423f3239d7aee791 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 16:49:22 -0600 Subject: [PATCH 17/24] Document installing with build isolation Describe isolated installs from PyPI, GitHub, and source, including pinning the framework version with build constraints, selecting frameworks, CUDA versions, and optional dependencies with config settings, and the cases that still require --no-build-isolation. Signed-off-by: Fred Heinecke --- docs/installation.rst | 64 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 63 insertions(+), 1 deletion(-) diff --git a/docs/installation.rst b/docs/installation.rst index 9d3eada827c..aa61b5152c1 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -36,7 +36,7 @@ Transformer Engine can be directly installed from `our PyPI build-constraints.txt + + # JAX + python3 -c "import jax, jaxlib; print(f'jax=={jax.__version__}'); print(f'jaxlib=={jaxlib.__version__}')" > build-constraints.txt + +If the framework was installed from a separate package index, e.g. the PyTorch index for a specific CUDA version, pass the same `--index-url` or `--extra-index-url` options to pip so that the build environment can find the same version. Framework builds that are not published to a package index, such as the ones in the NGC containers, cannot be used with build isolation. Use `--no-build-isolation` for these. + +Installation (from PyPI) +^^^^^^^^^^^^^^^^^^^^^^^^ + +The core library is shipped as wheels and does not need to be built. Execute the following command to install the PyTorch extension with build isolation: + +.. code-block:: bash + + pip3 install --build-constraint build-constraints.txt transformer_engine[pytorch] + +The JAX extension is installed in the same way: + +.. code-block:: bash + + pip3 install --build-constraint build-constraints.txt transformer_engine[jax] + +With build isolation, the JAX extension is built against JAX with CUDA 12 support and depends on the core library compiled for CUDA 12. For JAX with CUDA 13 support, install the JAX extension with `--no-build-isolation`. + +Installation (from GitHub or source) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +With build isolation, frameworks installed in the current environment are not detected. To build Transformer Engine support for them, set the `framework` config setting to a comma-separated list (e.g. `framework=jax,pytorch`). If no frameworks are specified, only the framework-agnostic C++ API is built. + +.. code-block:: bash + + cd TransformerEngine + pip3 install --build-constraint build-constraints.txt --config-settings framework=pytorch . + +The same options can be used to install from GitHub: + +.. code-block:: bash + + pip3 install --build-constraint build-constraints.txt --config-settings framework=pytorch git+https://github.com/NVIDIA/TransformerEngine.git@stable + +The CUDA compiler and libraries for the build are installed into the build environment, matching the version of the installed CUDA Toolkit. If no CUDA Toolkit is installed, specify the CUDA version with the `cuda-version` config setting (e.g. `cuda-version=13.0`). To build with an installed CUDA Toolkit instead, set `CUDA_HOME` to its location. + +Optional dependencies, such as cuBLASMp and NVSHMEM, are enabled with config settings as well (e.g. `with-cublasmp=1`). In an isolated build, their Python packages are installed into the build environment. Each config setting has an equivalent environment variable, e.g. `NVTE_FRAMEWORK` for `framework` and `NVTE_CUDA_VERSION` for `cuda-version`. See :doc:`envvars` for the complete list. + +Other installers support build constraints as well. For example, with `uv`: + +.. code-block:: bash + + uv pip install --build-constraints build-constraints.txt --config-setting framework=pytorch . + .. include:: troubleshooting.rst From 9613930d2e571688b98da1af6d999d95edef581d Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 17:08:33 -0600 Subject: [PATCH 18/24] Select the CUDA version of JAX for isolated JAX extension builds The JAX extension's build requirements always installed jax[cuda12], so isolated builds required the CUDA 12 core library even for JAX with CUDA 13 support. Add a build backend that installs jax[cuda12] or jax[cuda13] based on the cuda-version config setting or NVTE_CUDA_VERSION. CUDA 12 remains the default so that existing installs are unchanged. setup.py already selects the core library from the CUDA version of the installed JAX. Signed-off-by: Fred Heinecke --- docs/envvars.rst | 2 +- docs/installation.rst | 6 +- transformer_engine/jax/MANIFEST.in | 1 + transformer_engine/jax/_build_backend.py | 152 +++++++++++++++++++++++ transformer_engine/jax/pyproject.toml | 8 +- 5 files changed, 164 insertions(+), 5 deletions(-) create mode 100644 transformer_engine/jax/_build_backend.py diff --git a/docs/envvars.rst b/docs/envvars.rst index 1a748f7e988..e9a3483472d 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -82,7 +82,7 @@ Build Configuration :Type: ``str`` :Default: Auto-detected based on system CUDA version - :Description: CUDA version to build for. Valid values: ``"12.*"``, ``"13.*"``. If not set, automatically determined based on the installed CUDA Toolkit version. Isolated builds fail if no CUDA Toolkit is found and neither this variable nor pip's ``--config-settings cuda-version=.`` option is set. + :Description: CUDA version to build for. Valid values: ``"12.*"``, ``"13.*"``. If not set, automatically determined based on the installed CUDA Toolkit version. Isolated builds fail if no CUDA Toolkit is found and neither this variable nor pip's ``--config-settings cuda-version=.`` option is set. When the JAX extension package is built with build isolation, this variable or the config setting selects the CUDA support of the JAX installed in the build environment, which defaults to CUDA 12. Optional Dependencies ^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/installation.rst b/docs/installation.rst index aa61b5152c1..d2309572a23 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -156,7 +156,11 @@ The JAX extension is installed in the same way: pip3 install --build-constraint build-constraints.txt transformer_engine[jax] -With build isolation, the JAX extension is built against JAX with CUDA 12 support and depends on the core library compiled for CUDA 12. For JAX with CUDA 13 support, install the JAX extension with `--no-build-isolation`. +With build isolation, the JAX extension is built against JAX with CUDA 12 support and depends on the core library compiled for CUDA 12 by default. For JAX with CUDA 13 support, set the `cuda-version` config setting: + +.. code-block:: bash + + pip3 install --build-constraint build-constraints.txt --config-settings cuda-version=13.0 transformer_engine[jax] Installation (from GitHub or source) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/transformer_engine/jax/MANIFEST.in b/transformer_engine/jax/MANIFEST.in index 0c814f95da5..4924c920f51 100644 --- a/transformer_engine/jax/MANIFEST.in +++ b/transformer_engine/jax/MANIFEST.in @@ -1,3 +1,4 @@ +include _build_backend.py recursive-include build_tools *.* recursive-include common_headers *.* recursive-include csrc *.* diff --git a/transformer_engine/jax/_build_backend.py b/transformer_engine/jax/_build_backend.py new file mode 100644 index 00000000000..ce739dcda8a --- /dev/null +++ b/transformer_engine/jax/_build_backend.py @@ -0,0 +1,152 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PEP 517 backend that builds against JAX with support for the requested CUDA version.""" + +import os +import re +from typing import List, Mapping, Optional, Union + +from setuptools import build_meta + +_SETUPTOOLS_BACKEND = build_meta.__legacy__ +_SUPPORTED_CUDA_MAJOR_VERSIONS = ("12", "13") +# Isolated builds have always used JAX with CUDA 12 support. +_DEFAULT_CUDA_MAJOR_VERSION = "12" + +ConfigValue = Union[str, List[str]] +ConfigSettings = Optional[Mapping[str, ConfigValue]] + + +def _setuptools_config_settings(config_settings: ConfigSettings) -> ConfigSettings: + """Remove settings that are private to this backend.""" + if config_settings is None: + return None + + settings = dict(config_settings) + settings.pop("cuda-version", None) + return settings or None + + +def _cuda_major_version(config_settings: ConfigSettings) -> str: + """Get the CUDA major version from config settings or NVTE_CUDA_VERSION.""" + settings: Mapping[str, ConfigValue] = config_settings or {} + value = settings.get("cuda-version") + if isinstance(value, list): + if not value: + raise ValueError("CUDA version cannot be empty") + value = value[-1] + if value is None: + value = os.getenv("NVTE_CUDA_VERSION") + if value is None: + return _DEFAULT_CUDA_MAJOR_VERSION + + value = str(value).lower().removeprefix("cu") + if re.fullmatch(r"\d+\.\d+", value) is None: + raise ValueError( + f"Invalid CUDA version {value!r}; expected ., for example 13.0" + ) + + major = value.split(".", maxsplit=1)[0] + if major not in _SUPPORTED_CUDA_MAJOR_VERSIONS: + supported = ", ".join(_SUPPORTED_CUDA_MAJOR_VERSIONS) + raise ValueError(f"Unsupported CUDA major version {major!r}; expected one of: {supported}") + + return major + + +def _build_requirements(config_settings: ConfigSettings) -> List[str]: + """Get JAX with support for the requested CUDA version.""" + # setup.py detects the CUDA version of the installed JAX to select the core library. + return [f"jax[cuda{_cuda_major_version(config_settings)}]"] + + +################################################################################################### +# PEP 517 and PEP 660 defined functions +################################################################################################### + +# The setuptools hooks for build requirements are not used because they run setup.py, which +# imports JAX before it is installed. + + +# Defined by PEP 517 +def get_requires_for_build_wheel(config_settings: ConfigSettings = None) -> List[str]: + """Get requirements for building a wheel.""" + return _build_requirements(config_settings) + + +# Defined by PEP 517 +def get_requires_for_build_sdist(config_settings: ConfigSettings = None) -> List[str]: + """Get requirements needed while evaluating setup.py for an sdist.""" + return _build_requirements(config_settings) + + +# Defined by PEP 517 +def build_wheel( + wheel_directory: str, + config_settings: ConfigSettings = None, + metadata_directory: Optional[str] = None, +) -> str: + """Build a wheel with setuptools.""" + return _SETUPTOOLS_BACKEND.build_wheel( + wheel_directory, + _setuptools_config_settings(config_settings), + metadata_directory, + ) + + +# Defined by PEP 517 +def prepare_metadata_for_build_wheel( + metadata_directory: str, + config_settings: ConfigSettings = None, +) -> str: + """Prepare wheel metadata with setuptools.""" + return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_wheel( + metadata_directory, + _setuptools_config_settings(config_settings), + ) + + +# Defined by PEP 517 +def build_sdist( + sdist_directory: str, + config_settings: ConfigSettings = None, +) -> str: + """Build an sdist with setuptools.""" + return _SETUPTOOLS_BACKEND.build_sdist( + sdist_directory, + _setuptools_config_settings(config_settings), + ) + + +if hasattr(_SETUPTOOLS_BACKEND, "build_editable"): + + # Defined by PEP 660 + def get_requires_for_build_editable(config_settings: ConfigSettings = None) -> List[str]: + """Get requirements for building an editable wheel.""" + return _build_requirements(config_settings) + + # Defined by PEP 660 + def build_editable( + wheel_directory: str, + config_settings: ConfigSettings = None, + metadata_directory: Optional[str] = None, + ) -> str: + """Build an editable wheel with setuptools.""" + return _SETUPTOOLS_BACKEND.build_editable( + wheel_directory, + _setuptools_config_settings(config_settings), + metadata_directory, + ) + + # Defined by PEP 660 + def prepare_metadata_for_build_editable( + metadata_directory: str, + config_settings: ConfigSettings = None, + ) -> str: + """Prepare editable-wheel metadata with setuptools.""" + return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_editable( + metadata_directory, + _setuptools_config_settings(config_settings), + ) diff --git a/transformer_engine/jax/pyproject.toml b/transformer_engine/jax/pyproject.toml index d5c15b4a8fb..4c9648430cf 100755 --- a/transformer_engine/jax/pyproject.toml +++ b/transformer_engine/jax/pyproject.toml @@ -3,7 +3,9 @@ # See LICENSE for license information. [build-system] -requires = ["setuptools>=61.0", "pybind11[global]", "pip", "jax[cuda12]", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] +requires = ["setuptools>=61.0", "pybind11[global]", "pip", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] -# Use legacy backend to import local packages in setup.py -build-backend = "setuptools.build_meta:__legacy__" +# JAX with CUDA support is added in `_build_backend.py`. Select CUDA with +# --config-settings cuda-version= or NVTE_CUDA_VERSION (default: CUDA 12). +build-backend = "_build_backend" +backend-path = ["."] From 22682e6baaafe56e5035d90bfecb8920d4fbbefc Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 17:15:21 -0600 Subject: [PATCH 19/24] Share build backend helpers between the root and JAX projects Move CUDA version validation, config setting parsing, and the setuptools PEP 517/660 hook wrappers into build_tools/build_backend_utils.py, so the supported CUDA versions and settings handling are defined once. The JAX backend imports the repository's build_tools in a source checkout, like its setup.py, and the copy bundled in its sdist otherwise. Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 195 +++-------------------- build_tools/build_backend_utils.py | 163 +++++++++++++++++++ transformer_engine/jax/_build_backend.py | 159 +++--------------- 3 files changed, 211 insertions(+), 306 deletions(-) create mode 100644 build_tools/build_backend_utils.py diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index 8bf227d0f1f..c639dab692b 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -8,13 +8,17 @@ import os import re import subprocess -from typing import Iterator, List, Mapping, NamedTuple, Optional, Tuple, Union - -from setuptools import build_meta - +from typing import Iterator, List, NamedTuple, Optional, Tuple + +from build_tools.build_backend_utils import ( + ConfigSettings, + config_value, + normalize_cuda_version, + requested_cuda_version, + setuptools_hooks, +) from build_tools.utils import nvcc_path -_SETUPTOOLS_BACKEND = build_meta.__legacy__ _CUDA_BUILD_PACKAGES = { "12": [ # Also contains nvvm and crt @@ -70,57 +74,6 @@ class _Feature(NamedTuple): } _PRIVATE_SETTINGS = ("cuda-version", "framework", *_FEATURES) -ConfigValue = Union[str, List[str]] -ConfigSettings = Optional[Mapping[str, ConfigValue]] - - -def _setuptools_config_settings(config_settings: ConfigSettings) -> ConfigSettings: - """Remove settings that are private to this backend.""" - if config_settings is None: - return None - - settings = dict(config_settings) - for setting in _PRIVATE_SETTINGS: - settings.pop(setting, None) - return settings or None - - -def _normalize_cuda_version(value: str) -> str: - """Validate and normalize a CUDA major/minor version.""" - value = value.lower().removeprefix("cu") - if re.fullmatch(r"\d+\.\d+", value) is None: - raise ValueError( - f"Invalid CUDA version {value!r}; expected ., for example 12.8" - ) - - major = value.split(".", maxsplit=1)[0] - if major not in _CUDA_BUILD_PACKAGES: - supported = ", ".join(sorted(_CUDA_BUILD_PACKAGES)) - raise ValueError(f"Unsupported CUDA major version {major!r}; expected one of: {supported}") - - return value - - -def _config_cuda_version(config_settings: ConfigSettings) -> Optional[str]: - """Get the CUDA version from PEP 517 config settings.""" - settings: Mapping[str, ConfigValue] = config_settings or {} - value = settings.get("cuda-version") - if value is None: - return None - - if isinstance(value, list): - if not value: - raise ValueError("CUDA package version cannot be empty") - value = value[-1] - - return _normalize_cuda_version(str(value)) - - -def _environment_cuda_version() -> Optional[str]: - """Get the CUDA version from the environment.""" - value = os.getenv("NVTE_CUDA_VERSION") - return _normalize_cuda_version(value) if value is not None else None - def _normalize_frameworks(value: str) -> List[str]: """Validate and normalize a framework selection.""" @@ -150,17 +103,8 @@ def _normalize_frameworks(value: str) -> List[str]: def _config_frameworks(config_settings: ConfigSettings) -> Optional[List[str]]: """Get the target frameworks from PEP 517 config settings.""" - settings: Mapping[str, ConfigValue] = config_settings or {} - value = settings.get("framework") - if value is None: - return None - - if isinstance(value, list): - if not value: - raise ValueError("Framework selection cannot be empty") - value = value[-1] - - return _normalize_frameworks(str(value)) + value = config_value(config_settings, "framework") + return _normalize_frameworks(value) if value is not None else None def _environment_frameworks() -> Optional[List[str]]: @@ -177,16 +121,10 @@ def _requested_frameworks(config_settings: ConfigSettings) -> Optional[List[str] def _config_feature(config_settings: ConfigSettings, setting: str) -> Optional[bool]: """Get whether PEP 517 config settings enable or disable an optional feature.""" - settings: Mapping[str, ConfigValue] = config_settings or {} - value = settings.get(setting) + value = config_value(config_settings, setting) if value is None: return None - if isinstance(value, list): - if not value: - raise ValueError(f"{setting} cannot be empty") - value = value[-1] - if value not in ("0", "1"): raise ValueError(f"Invalid value {value!r} for {setting}; expected 0 or 1") return value == "1" @@ -240,7 +178,7 @@ def _nvcc_cuda_version() -> Optional[str]: return None match = re.search(r"release\s+(\d+\.\d+)", result.stdout) - return _normalize_cuda_version(match.group(1)) if match is not None else None + return normalize_cuda_version(match.group(1)) if match is not None else None def _torch_cuda_version() -> Optional[str]: @@ -250,7 +188,7 @@ def _torch_cuda_version() -> Optional[str]: except (ImportError, OSError): return None value = getattr(getattr(torch, "version", None), "cuda", None) - return _normalize_cuda_version(value) if value else None + return normalize_cuda_version(value) if value else None def _framework_cuda_version(config_settings: ConfigSettings) -> Optional[str]: @@ -269,8 +207,7 @@ def _framework_cuda_version(config_settings: ConfigSettings) -> Optional[str]: def _cuda_version(config_settings: ConfigSettings) -> str: """Resolve the CUDA version in descending order of precedence.""" cuda_version = ( - _config_cuda_version(config_settings) - or _environment_cuda_version() + requested_cuda_version(config_settings) or _nvcc_cuda_version() or _framework_cuda_version(config_settings) ) @@ -332,96 +269,12 @@ def _build_requirements(config_settings: ConfigSettings) -> List[str]: ) -################################################################################################### -# PEP 517 and PEP 660 defined functions -################################################################################################### - - -# Defined by PEP 517 -def get_requires_for_build_wheel(config_settings: ConfigSettings = None) -> List[str]: - """Get requirements for building a wheel.""" - return _build_requirements(config_settings) - - -# Defined by PEP 517 -def get_requires_for_build_sdist(config_settings: ConfigSettings = None) -> List[str]: - """Get requirements needed while evaluating setup.py for an sdist.""" - # Build requirements are needed even for sdist because setup.py currently always depends - # on them. It could be refactored to only depend on them when building a wheel, but that - # would require a more invasive refactor. - return _build_requirements(config_settings) - - -# Defined by PEP 517 -def build_wheel( - wheel_directory: str, - config_settings: ConfigSettings = None, - metadata_directory: Optional[str] = None, -) -> str: - """Build a wheel with setuptools.""" - with _build_environment(config_settings): - return _SETUPTOOLS_BACKEND.build_wheel( - wheel_directory, - _setuptools_config_settings(config_settings), - metadata_directory, - ) - - -# Defined by PEP 517 -def prepare_metadata_for_build_wheel( - metadata_directory: str, - config_settings: ConfigSettings = None, -) -> str: - """Prepare wheel metadata with setuptools.""" - with _build_environment(config_settings): - return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_wheel( - metadata_directory, - _setuptools_config_settings(config_settings), - ) - - -# Defined by PEP 517 -def build_sdist( - sdist_directory: str, - config_settings: ConfigSettings = None, -) -> str: - """Build an sdist with setuptools.""" - with _build_environment(config_settings): - return _SETUPTOOLS_BACKEND.build_sdist( - sdist_directory, - _setuptools_config_settings(config_settings), - ) - - -if hasattr(_SETUPTOOLS_BACKEND, "build_editable"): - - # Defined by PEP 660 - def get_requires_for_build_editable(config_settings: ConfigSettings = None) -> List[str]: - """Get requirements for building an editable wheel.""" - return _build_requirements(config_settings) - - # Defined by PEP 660 - def build_editable( - wheel_directory: str, - config_settings: ConfigSettings = None, - metadata_directory: Optional[str] = None, - ) -> str: - """Build an editable wheel with setuptools.""" - with _build_environment(config_settings): - return _SETUPTOOLS_BACKEND.build_editable( - wheel_directory, - _setuptools_config_settings(config_settings), - metadata_directory, - ) - - # Defined by PEP 660 - def prepare_metadata_for_build_editable( - metadata_directory: str, - config_settings: ConfigSettings = None, - ) -> str: - """Prepare editable-wheel metadata with setuptools.""" - with _build_environment(config_settings): - return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_editable( - metadata_directory, - _setuptools_config_settings(config_settings), - ) +# PEP 517 and PEP 660 hooks. setup.py always depends on the build requirements, so they are +# also needed for sdists. +globals().update( + setuptools_hooks( + requirements=_build_requirements, + private_settings=_PRIVATE_SETTINGS, + environment=_build_environment, + ) +) diff --git a/build_tools/build_backend_utils.py b/build_tools/build_backend_utils.py new file mode 100644 index 00000000000..a589cf41bf5 --- /dev/null +++ b/build_tools/build_backend_utils.py @@ -0,0 +1,163 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared helpers for Transformer Engine's PEP 517 build backends. + +Only the standard library and setuptools may be imported here, since the backends import +this module before any other build requirements are installed. +""" + +from contextlib import nullcontext +import os +import re +from typing import Callable, ContextManager, Dict, Iterable, List, Mapping, Optional, Union + +from setuptools import build_meta + +SUPPORTED_CUDA_MAJOR_VERSIONS = ("12", "13") + +ConfigValue = Union[str, List[str]] +ConfigSettings = Optional[Mapping[str, ConfigValue]] + + +def config_value(config_settings: ConfigSettings, setting: str) -> Optional[str]: + """Get a PEP 517 config setting, using the last value if it was given more than once.""" + settings: Mapping[str, ConfigValue] = config_settings or {} + value = settings.get(setting) + if value is None: + return None + + if isinstance(value, list): + if not value: + raise ValueError(f"{setting} cannot be empty") + value = value[-1] + + return str(value) + + +def normalize_cuda_version(value: str) -> str: + """Validate and normalize a CUDA major/minor version.""" + value = value.lower().removeprefix("cu") + if re.fullmatch(r"\d+\.\d+", value) is None: + raise ValueError( + f"Invalid CUDA version {value!r}; expected ., for example 12.8" + ) + + major = value.split(".", maxsplit=1)[0] + if major not in SUPPORTED_CUDA_MAJOR_VERSIONS: + supported = ", ".join(SUPPORTED_CUDA_MAJOR_VERSIONS) + raise ValueError(f"Unsupported CUDA major version {major!r}; expected one of: {supported}") + + return value + + +def requested_cuda_version(config_settings: ConfigSettings) -> Optional[str]: + """Get the CUDA version from the cuda-version config setting or NVTE_CUDA_VERSION.""" + value = config_value(config_settings, "cuda-version") + if value is None: + value = os.getenv("NVTE_CUDA_VERSION") + return normalize_cuda_version(value) if value is not None else None + + +def setuptools_hooks( + requirements: Callable[[ConfigSettings], List[str]], + private_settings: Iterable[str], + environment: Optional[Callable[[ConfigSettings], ContextManager]] = None, +) -> Dict[str, Callable]: + """Create PEP 517 and PEP 660 hooks that build with the legacy setuptools backend. + + Build requirements come only from `requirements`. The setuptools hooks for build + requirements are not used because they run setup.py, which can import build + requirements before they are installed. Private settings are removed before the + config settings are passed to setuptools, and builds run inside `environment`. + """ + backend = build_meta.__legacy__ + private_settings = frozenset(private_settings) + if environment is None: + environment = lambda config_settings: nullcontext() + + def setuptools_config_settings(config_settings: ConfigSettings) -> ConfigSettings: + if config_settings is None: + return None + settings = { + setting: value + for setting, value in config_settings.items() + if setting not in private_settings + } + return settings or None + + def get_requires_for_build_wheel(config_settings: ConfigSettings = None) -> List[str]: + """Get requirements for building a wheel.""" + return requirements(config_settings) + + def get_requires_for_build_sdist(config_settings: ConfigSettings = None) -> List[str]: + """Get requirements needed while evaluating setup.py for an sdist.""" + return requirements(config_settings) + + def build_wheel( + wheel_directory: str, + config_settings: ConfigSettings = None, + metadata_directory: Optional[str] = None, + ) -> str: + """Build a wheel with setuptools.""" + with environment(config_settings): + return backend.build_wheel( + wheel_directory, setuptools_config_settings(config_settings), metadata_directory + ) + + def prepare_metadata_for_build_wheel( + metadata_directory: str, config_settings: ConfigSettings = None + ) -> str: + """Prepare wheel metadata with setuptools.""" + with environment(config_settings): + return backend.prepare_metadata_for_build_wheel( + metadata_directory, setuptools_config_settings(config_settings) + ) + + def build_sdist(sdist_directory: str, config_settings: ConfigSettings = None) -> str: + """Build an sdist with setuptools.""" + with environment(config_settings): + return backend.build_sdist(sdist_directory, setuptools_config_settings(config_settings)) + + hooks = { + "get_requires_for_build_wheel": get_requires_for_build_wheel, + "get_requires_for_build_sdist": get_requires_for_build_sdist, + "build_wheel": build_wheel, + "prepare_metadata_for_build_wheel": prepare_metadata_for_build_wheel, + "build_sdist": build_sdist, + } + + if hasattr(backend, "build_editable"): + + def get_requires_for_build_editable(config_settings: ConfigSettings = None) -> List[str]: + """Get requirements for building an editable wheel.""" + return requirements(config_settings) + + def build_editable( + wheel_directory: str, + config_settings: ConfigSettings = None, + metadata_directory: Optional[str] = None, + ) -> str: + """Build an editable wheel with setuptools.""" + with environment(config_settings): + return backend.build_editable( + wheel_directory, setuptools_config_settings(config_settings), metadata_directory + ) + + def prepare_metadata_for_build_editable( + metadata_directory: str, config_settings: ConfigSettings = None + ) -> str: + """Prepare editable-wheel metadata with setuptools.""" + with environment(config_settings): + return backend.prepare_metadata_for_build_editable( + metadata_directory, setuptools_config_settings(config_settings) + ) + + hooks.update( + get_requires_for_build_editable=get_requires_for_build_editable, + build_editable=build_editable, + prepare_metadata_for_build_editable=prepare_metadata_for_build_editable, + ) + + return hooks diff --git a/transformer_engine/jax/_build_backend.py b/transformer_engine/jax/_build_backend.py index ce739dcda8a..7a981eefb07 100644 --- a/transformer_engine/jax/_build_backend.py +++ b/transformer_engine/jax/_build_backend.py @@ -4,149 +4,38 @@ """PEP 517 backend that builds against JAX with support for the requested CUDA version.""" -import os -import re -from typing import List, Mapping, Optional, Union +from pathlib import Path +import sys +from typing import List + +# Like setup.py, prefer the repository's build_tools in a source checkout. The sdist contains +# a copy next to this file. +_REPOSITORY_ROOT = Path(__file__).resolve().parent.parent.parent +if (_REPOSITORY_ROOT / "build_tools").is_dir(): + sys.path.insert(0, str(_REPOSITORY_ROOT)) + +# pylint: disable=wrong-import-position +from build_tools.build_backend_utils import ( + ConfigSettings, + requested_cuda_version, + setuptools_hooks, +) -from setuptools import build_meta - -_SETUPTOOLS_BACKEND = build_meta.__legacy__ -_SUPPORTED_CUDA_MAJOR_VERSIONS = ("12", "13") # Isolated builds have always used JAX with CUDA 12 support. _DEFAULT_CUDA_MAJOR_VERSION = "12" -ConfigValue = Union[str, List[str]] -ConfigSettings = Optional[Mapping[str, ConfigValue]] - - -def _setuptools_config_settings(config_settings: ConfigSettings) -> ConfigSettings: - """Remove settings that are private to this backend.""" - if config_settings is None: - return None - - settings = dict(config_settings) - settings.pop("cuda-version", None) - return settings or None - - -def _cuda_major_version(config_settings: ConfigSettings) -> str: - """Get the CUDA major version from config settings or NVTE_CUDA_VERSION.""" - settings: Mapping[str, ConfigValue] = config_settings or {} - value = settings.get("cuda-version") - if isinstance(value, list): - if not value: - raise ValueError("CUDA version cannot be empty") - value = value[-1] - if value is None: - value = os.getenv("NVTE_CUDA_VERSION") - if value is None: - return _DEFAULT_CUDA_MAJOR_VERSION - - value = str(value).lower().removeprefix("cu") - if re.fullmatch(r"\d+\.\d+", value) is None: - raise ValueError( - f"Invalid CUDA version {value!r}; expected ., for example 13.0" - ) - - major = value.split(".", maxsplit=1)[0] - if major not in _SUPPORTED_CUDA_MAJOR_VERSIONS: - supported = ", ".join(_SUPPORTED_CUDA_MAJOR_VERSIONS) - raise ValueError(f"Unsupported CUDA major version {major!r}; expected one of: {supported}") - - return major - def _build_requirements(config_settings: ConfigSettings) -> List[str]: """Get JAX with support for the requested CUDA version.""" # setup.py detects the CUDA version of the installed JAX to select the core library. - return [f"jax[cuda{_cuda_major_version(config_settings)}]"] - - -################################################################################################### -# PEP 517 and PEP 660 defined functions -################################################################################################### - -# The setuptools hooks for build requirements are not used because they run setup.py, which -# imports JAX before it is installed. - - -# Defined by PEP 517 -def get_requires_for_build_wheel(config_settings: ConfigSettings = None) -> List[str]: - """Get requirements for building a wheel.""" - return _build_requirements(config_settings) - - -# Defined by PEP 517 -def get_requires_for_build_sdist(config_settings: ConfigSettings = None) -> List[str]: - """Get requirements needed while evaluating setup.py for an sdist.""" - return _build_requirements(config_settings) - - -# Defined by PEP 517 -def build_wheel( - wheel_directory: str, - config_settings: ConfigSettings = None, - metadata_directory: Optional[str] = None, -) -> str: - """Build a wheel with setuptools.""" - return _SETUPTOOLS_BACKEND.build_wheel( - wheel_directory, - _setuptools_config_settings(config_settings), - metadata_directory, - ) - - -# Defined by PEP 517 -def prepare_metadata_for_build_wheel( - metadata_directory: str, - config_settings: ConfigSettings = None, -) -> str: - """Prepare wheel metadata with setuptools.""" - return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_wheel( - metadata_directory, - _setuptools_config_settings(config_settings), - ) - - -# Defined by PEP 517 -def build_sdist( - sdist_directory: str, - config_settings: ConfigSettings = None, -) -> str: - """Build an sdist with setuptools.""" - return _SETUPTOOLS_BACKEND.build_sdist( - sdist_directory, - _setuptools_config_settings(config_settings), + cuda_version = requested_cuda_version(config_settings) + cuda_major = ( + cuda_version.split(".", maxsplit=1)[0] if cuda_version else _DEFAULT_CUDA_MAJOR_VERSION ) + return [f"jax[cuda{cuda_major}]"] -if hasattr(_SETUPTOOLS_BACKEND, "build_editable"): - - # Defined by PEP 660 - def get_requires_for_build_editable(config_settings: ConfigSettings = None) -> List[str]: - """Get requirements for building an editable wheel.""" - return _build_requirements(config_settings) - - # Defined by PEP 660 - def build_editable( - wheel_directory: str, - config_settings: ConfigSettings = None, - metadata_directory: Optional[str] = None, - ) -> str: - """Build an editable wheel with setuptools.""" - return _SETUPTOOLS_BACKEND.build_editable( - wheel_directory, - _setuptools_config_settings(config_settings), - metadata_directory, - ) - - # Defined by PEP 660 - def prepare_metadata_for_build_editable( - metadata_directory: str, - config_settings: ConfigSettings = None, - ) -> str: - """Prepare editable-wheel metadata with setuptools.""" - return _SETUPTOOLS_BACKEND.prepare_metadata_for_build_editable( - metadata_directory, - _setuptools_config_settings(config_settings), - ) +# PEP 517 and PEP 660 hooks +globals().update( + setuptools_hooks(requirements=_build_requirements, private_settings=("cuda-version",)) +) From 5ff3ba5fd72c51aa3b297d1bd2f4508a0816115d Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 17:39:04 -0600 Subject: [PATCH 20/24] Install cuBLAS and cuDNN for isolated builds and declare Python package dependencies Isolated builds installed the CUDA compiler and runtime but not cuBLAS or cuDNN, which the core library links against. Install them through the cuda-toolkit metapackage, which pins all CUDA components to the same release, and nvidia-cudnn-cu*. When CUDA_HOME is set, use the installed CUDA Toolkit and cuDNN instead of installing either. When the core library is built against libraries from Python packages (the CUDA Toolkit, cuDNN, NCCL, cuBLASMp, cuSOLVERMp, NVSHMEM, or openmpi), declare runtime dependencies on those packages, within the major version used for the build so that they stay compatible with frameworks' exact pins. Builds against system libraries, including release wheels, declare no new dependencies. When the CUDA Toolkit comes from Python packages, build NCCL EP against the nvidia-nccl wheel that the core library then depends on and loads at runtime. Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 43 ++++++---------- build_tools/utils.py | 56 +++++++++++++++++---- docs/envvars.rst | 5 +- docs/installation.rst | 4 +- setup.py | 97 ++++++++++++++++++++++++++++++------ 5 files changed, 149 insertions(+), 56 deletions(-) diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index c639dab692b..cca1848f783 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -17,27 +17,11 @@ requested_cuda_version, setuptools_hooks, ) -from build_tools.utils import nvcc_path - -_CUDA_BUILD_PACKAGES = { - "12": [ - # Also contains nvvm and crt - "nvidia-cuda-nvcc-cu12", - "nvidia-cuda-runtime-cu12", - "nvidia-cuda-cccl-cu12", - "nvidia-cuda-profiler-api-cu12", - "nvidia-nvml-dev-cu12", - ], - "13": [ - "nvidia-cuda-nvcc", - "nvidia-cuda-runtime", - "nvidia-cuda-crt", - "nvidia-cuda-cccl", - "nvidia-cuda-profiler-api", - "nvidia-nvml-dev", - "nvidia-nvvm", - ], -} +from build_tools.utils import MIN_CUDNN_VERSION, nvcc_path + +# cuda-toolkit pins all components to the same CUDA release. The nvcc extra includes the +# compiler's crt and nvvm components. +_CUDA_TOOLKIT_EXTRAS = ("nvcc", "cudart", "cccl", "cublas", "nvml", "profiler") _FRAMEWORK_BUILD_PACKAGES = { "pytorch": ["torch>=2.1"], "jax": ["jax>=0.5.0", "flax>=0.7.1"], @@ -221,14 +205,19 @@ def _cuda_version(config_settings: ConfigSettings) -> str: def _cuda_build_requirements(config_settings: ConfigSettings) -> List[str]: - """Get build requirements for the requested CUDA package version.""" + """Get build requirements for the requested CUDA version. + + The CUDA Toolkit and cuDNN are installed from Python packages unless CUDA_HOME selects an + installed CUDA Toolkit, in which case the installed cuDNN is used as well. + """ cuda_version = _cuda_version(config_settings) cuda_major = cuda_version.split(".", maxsplit=1)[0] - component_constraint = f"=={cuda_version}.*" - - packages = _CUDA_BUILD_PACKAGES[cuda_major] - requirements = [f"{package}{component_constraint}" for package in packages] - requirements.append(f"nvidia-nccl-cu{cuda_major}>=2") + requirements = [f"nvidia-nccl-cu{cuda_major}>=2"] + if not os.getenv("CUDA_HOME"): + requirements += [ + f"cuda-toolkit[{','.join(_CUDA_TOOLKIT_EXTRAS)}]=={cuda_version}.*", + f"nvidia-cudnn-cu{cuda_major}>={MIN_CUDNN_VERSION}", + ] return requirements diff --git a/build_tools/utils.py b/build_tools/utils.py index 98c9e64cb45..51bfac9dd28 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -366,6 +366,33 @@ def nccl_lib_path() -> Optional[Path]: return None +# Keep in sync with MIN_CUDNN_VERSION in transformer_engine/common/__init__.py. +MIN_CUDNN_VERSION = "9.12" + + +def python_package_requirement(package_name: str, minimum: Optional[str] = None) -> str: + """Require an installed Python package within its installed major version. + + The lower bound is `minimum` if given, otherwise the installed version. + """ + installed = get_version(package_name) + major = int(installed.split(".", maxsplit=1)[0]) + return f"{package_name}>={minimum or installed},<{major + 1}" + + +def cudnn_from_python_package() -> bool: + """Whether CMake finds cuDNN in the nvidia-cudnn-cu* Python package. + + cmake/cuDNN.cmake prefers CUDNN_INCLUDE_PATH and CUDNN_PATH over the Python package. + """ + if os.getenv("CUDNN_INCLUDE_PATH") or os.getenv("CUDNN_PATH"): + return False + return ( + _python_package_path(f"nvidia-cudnn-cu{cuda_version()[0]}", "nvidia/cudnn/include") + is not None + ) + + def _python_package_path(package_name: str, relative_path: str) -> Optional[Path]: """Return a path inside an installed Python package, or `None` if it does not exist.""" try: @@ -447,26 +474,33 @@ def cuda_toolkit_include_path() -> Optional[Path]: return None -@functools.lru_cache(maxsize=None) -def get_cuda_include_dirs() -> List[Path]: - """Returns the CUDA header directories.""" - - force_wheels = bool(int(os.getenv("NVTE_BUILD_USE_NVIDIA_WHEELS", "0"))) +def _nvidia_package_roots() -> List[Path]: + """Return the directories of the `nvidia` namespace package for NVIDIA Python packages.""" try: import nvidia except ModuleNotFoundError: - nvidia_roots = [] - else: - nvidia_roots = [Path(path).resolve() for path in nvidia.__path__] + return [] + return [Path(path).resolve() for path in nvidia.__path__] + +def cuda_toolkit_from_python_packages() -> bool: + """Whether the CUDA Toolkit used for the build is installed from Python packages.""" cuda_home = cuda_home_path() - cuda_home_from_wheel = cuda_home is not None and any( - cuda_home.is_relative_to(root) for root in nvidia_roots + return cuda_home is not None and any( + cuda_home.is_relative_to(root) for root in _nvidia_package_roots() ) + +@functools.lru_cache(maxsize=None) +def get_cuda_include_dirs() -> List[Path]: + """Returns the CUDA header directories.""" + + force_wheels = bool(int(os.getenv("NVTE_BUILD_USE_NVIDIA_WHEELS", "0"))) + nvidia_roots = _nvidia_package_roots() + # A system toolkit contains all necessary headers under one include # directory. NVIDIA wheels split them across namespace packages. - if not force_wheels and not cuda_home_from_wheel: + if not force_wheels and not cuda_toolkit_from_python_packages(): if (include_path := cuda_toolkit_include_path()) is not None: return [include_path] diff --git a/docs/envvars.rst b/docs/envvars.rst index e9a3483472d..b0d88df51fa 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -92,8 +92,9 @@ Each of the following options can also be set for an isolated PEP 517 build with In an isolated build, the Python package listed for an option is installed as a build requirement unless the option's installation directory variable is set. Without build isolation, the installation directory variable is used if set, and the package is used if -installed. ``libtransformer_engine.so`` finds these packages at runtime when they are installed -in the same environment. +installed. When Transformer Engine is built against a package, the built Transformer Engine +package depends on it, and ``libtransformer_engine.so`` finds it at runtime in the same +environment. .. envvar:: NVTE_UB_WITH_MPI diff --git a/docs/installation.rst b/docs/installation.rst index d2309572a23..e7b36e7e112 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -178,9 +178,9 @@ The same options can be used to install from GitHub: pip3 install --build-constraint build-constraints.txt --config-settings framework=pytorch git+https://github.com/NVIDIA/TransformerEngine.git@stable -The CUDA compiler and libraries for the build are installed into the build environment, matching the version of the installed CUDA Toolkit. If no CUDA Toolkit is installed, specify the CUDA version with the `cuda-version` config setting (e.g. `cuda-version=13.0`). To build with an installed CUDA Toolkit instead, set `CUDA_HOME` to its location. +Unless `CUDA_HOME` is set, the CUDA Toolkit and cuDNN for the build are installed into the build environment from PyPI, matching the version of the installed CUDA Toolkit. If no CUDA Toolkit is installed, specify the CUDA version with the `cuda-version` config setting (e.g. `cuda-version=13.0`). The installed Transformer Engine package then depends on the CUDA Toolkit and cuDNN packages from PyPI, so when building framework support, choose the CUDA version that PyTorch or JAX was built with. To build with an installed CUDA Toolkit and cuDNN instead, set `CUDA_HOME` to the location of the CUDA Toolkit. -Optional dependencies, such as cuBLASMp and NVSHMEM, are enabled with config settings as well (e.g. `with-cublasmp=1`). In an isolated build, their Python packages are installed into the build environment. Each config setting has an equivalent environment variable, e.g. `NVTE_FRAMEWORK` for `framework` and `NVTE_CUDA_VERSION` for `cuda-version`. See :doc:`envvars` for the complete list. +Optional dependencies, such as cuBLASMp and NVSHMEM, are enabled with config settings as well (e.g. `with-cublasmp=1`). In an isolated build, their Python packages are installed into the build environment, and the installed Transformer Engine package depends on them. Each config setting has an equivalent environment variable, e.g. `NVTE_FRAMEWORK` for `framework` and `NVTE_CUDA_VERSION` for `cuda-version`. See :doc:`envvars` for the complete list. Other installers support build constraints as well. For example, with `uv`: diff --git a/setup.py b/setup.py index 8031f8031bb..0f383c0022c 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ import subprocess import time from pathlib import Path -from typing import List, Tuple +from typing import List, Optional, Tuple import setuptools from wheel.bdist_wheel import bdist_wheel @@ -18,10 +18,14 @@ from build_tools.build_ext import CMakeExtension, get_build_ext from build_tools.te_version import te_version from build_tools.utils import ( + MIN_CUDNN_VERSION, cublasmp_home_path, cuda_archs, cuda_home_path, + cuda_toolkit_from_python_packages, + cuda_version, cudnn_frontend_include_path, + cudnn_from_python_package, cusolvermp_home_path, get_bolt_build_flags, get_frameworks, @@ -32,6 +36,7 @@ get_max_jobs_for_parallel_build, nvcc_path, nvshmem_home_path, + python_package_requirement, ) frameworks = get_frameworks() @@ -161,9 +166,62 @@ def setup_requirements() -> Tuple[List[str], List[str]]: install_reqs.extend(install_requirements()) test_reqs.extend(test_requirements()) + install_reqs.extend(python_package_runtime_requirements()) + return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]] +def python_package_runtime_requirements() -> List[str]: + """Runtime requirements on the Python packages that the core library is built against. + + Libraries found through a *_HOME variable or on the system are not declared. Each + requirement stays within the major version of the package used for the build, so that it + is compatible with the exact versions that frameworks pin. + """ + cuda_major = cuda_version()[0] + requirements: List[str] = [] + + def require_installed(package_name: str, minimum: Optional[str] = None) -> None: + try: + requirements.append(python_package_requirement(package_name, minimum)) + except metadata.PackageNotFoundError: + pass + + if cuda_toolkit_from_python_packages(): + cuda_major_minor = ".".join(str(part) for part in cuda_version()[:2]) + requirements.append( + f"cuda-toolkit[cudart,cublas,nvrtc,curand]>={cuda_major_minor},<{cuda_major + 1}" + ) + if cudnn_from_python_package(): + require_installed(f"nvidia-cudnn-cu{cuda_major}", minimum=MIN_CUDNN_VERSION) + + with_cublasmp = bool(int(os.getenv("NVTE_WITH_CUBLASMP", "0"))) + with_cusolvermp = bool(int(os.getenv("NVTE_WITH_CUSOLVERMP", "0"))) + nccl_pip_root = _python_package_nccl_home() + if nccl_pip_root is not None: + if nccl_ep_enabled(archs): + try: + nccl_from_pip = _discover_nccl_home() == nccl_pip_root + except RuntimeError: + nccl_from_pip = False + else: + # CMake looks for NCCL in the Python package before system locations. + nccl_from_pip = with_cublasmp or with_cusolvermp + if nccl_from_pip: + require_installed(f"nvidia-nccl-cu{cuda_major}", "2.30" if with_cublasmp else "2") + + if with_cublasmp and not os.getenv("CUBLASMP_HOME"): + require_installed(f"nvidia-cublasmp-cu{cuda_major}") + if with_cusolvermp and not os.getenv("CUSOLVERMP_HOME"): + require_installed(f"nvidia-cusolvermp-cu{cuda_major}") + if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", "0"))) and not os.getenv("NVSHMEM_HOME"): + require_installed(f"nvidia-nvshmem-cu{cuda_major}") + if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))) and not os.getenv("MPI_HOME"): + require_installed("openmpi") + + return requirements + + def _discover_nccl_home() -> str: """Resolve NCCL_HOME, preferring the NCCL the dynamic loader resolves at runtime. @@ -180,6 +238,11 @@ def _discover_nccl_home() -> str: f"'{env_home}/include/nccl.h' was not found; falling back to system probes." ) + # A build against CUDA Toolkit Python packages also depends on the nvidia-nccl + # wheel at runtime, which libtransformer_engine's RUNPATH prefers. + if cuda_toolkit_from_python_packages() and (pip_root := _python_package_nccl_home()): + return pip_root + 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") @@ -210,24 +273,30 @@ def _discover_nccl_home() -> str: ): 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. + # Fall back to a pip-installed NCCL when no system NCCL is present. + if pip_root := _python_package_nccl_home(): + return pip_root + + raise RuntimeError( + "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." + ) + + +def _python_package_nccl_home() -> Optional[str]: + """Return the NCCL installed from an nvidia-nccl-cu* wheel under nvidia/nccl, if any.""" 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." - ) + return None + if spec is None or not spec.submodule_search_locations: + return None + + pip_root = Path(next(iter(spec.submodule_search_locations))) + if (pip_root / "include" / "nccl.h").exists() and (pip_root / "lib" / "libnccl.so.2").exists(): + return str(pip_root) + return None def build_nccl_ep_submodule() -> str: From aaaf3d549fb90a22537cd922dd542c01c9f4251c Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 17:47:25 -0600 Subject: [PATCH 21/24] Define the minimum cuDNN version once for the build and runtime Move MIN_CUDNN_VERSION into transformer_engine/common/_versions.py, which imports nothing. The runtime imports it as before, and the build loads the file directly so that the transformer_engine package and its shared libraries are not loaded during the build. Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 4 ++-- build_tools/utils.py | 17 +++++++++++++++-- setup.py | 4 ++-- transformer_engine/common/__init__.py | 4 +--- transformer_engine/common/_versions.py | 13 +++++++++++++ transformer_engine/common/cudnn_min_version.h | 2 +- 6 files changed, 34 insertions(+), 10 deletions(-) create mode 100644 transformer_engine/common/_versions.py diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index cca1848f783..18899042938 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -17,7 +17,7 @@ requested_cuda_version, setuptools_hooks, ) -from build_tools.utils import MIN_CUDNN_VERSION, nvcc_path +from build_tools.utils import min_cudnn_version, nvcc_path # cuda-toolkit pins all components to the same CUDA release. The nvcc extra includes the # compiler's crt and nvvm components. @@ -216,7 +216,7 @@ def _cuda_build_requirements(config_settings: ConfigSettings) -> List[str]: if not os.getenv("CUDA_HOME"): requirements += [ f"cuda-toolkit[{','.join(_CUDA_TOOLKIT_EXTRAS)}]=={cuda_version}.*", - f"nvidia-cudnn-cu{cuda_major}>={MIN_CUDNN_VERSION}", + f"nvidia-cudnn-cu{cuda_major}>={min_cudnn_version()}", ] return requirements diff --git a/build_tools/utils.py b/build_tools/utils.py index 51bfac9dd28..8fd8b89eb51 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -7,6 +7,7 @@ import functools import glob import importlib +import importlib.util import os import re import shlex @@ -366,8 +367,20 @@ def nccl_lib_path() -> Optional[Path]: return None -# Keep in sync with MIN_CUDNN_VERSION in transformer_engine/common/__init__.py. -MIN_CUDNN_VERSION = "9.12" +@functools.lru_cache(maxsize=None) +def min_cudnn_version() -> str: + """Return the minimum supported cuDNN version from transformer_engine/common/_versions.py. + + The file is loaded directly because importing the transformer_engine package loads its + shared libraries. The framework sdists do not contain it, so it is only read when needed. + """ + versions_file = ( + Path(__file__).resolve().parent.parent / "transformer_engine" / "common" / "_versions.py" + ) + spec = importlib.util.spec_from_file_location("_transformer_engine_versions", versions_file) + versions = importlib.util.module_from_spec(spec) + spec.loader.exec_module(versions) + return ".".join(str(part) for part in versions.MIN_CUDNN_VERSION) def python_package_requirement(package_name: str, minimum: Optional[str] = None) -> str: diff --git a/setup.py b/setup.py index 0f383c0022c..61d4d5798e2 100644 --- a/setup.py +++ b/setup.py @@ -18,7 +18,6 @@ from build_tools.build_ext import CMakeExtension, get_build_ext from build_tools.te_version import te_version from build_tools.utils import ( - MIN_CUDNN_VERSION, cublasmp_home_path, cuda_archs, cuda_home_path, @@ -31,6 +30,7 @@ get_frameworks, remove_dups, min_python_version_str, + min_cudnn_version, mpi_home_path, nccl_ep_enabled, get_max_jobs_for_parallel_build, @@ -193,7 +193,7 @@ def require_installed(package_name: str, minimum: Optional[str] = None) -> None: f"cuda-toolkit[cudart,cublas,nvrtc,curand]>={cuda_major_minor},<{cuda_major + 1}" ) if cudnn_from_python_package(): - require_installed(f"nvidia-cudnn-cu{cuda_major}", minimum=MIN_CUDNN_VERSION) + require_installed(f"nvidia-cudnn-cu{cuda_major}", minimum=min_cudnn_version()) with_cublasmp = bool(int(os.getenv("NVTE_WITH_CUBLASMP", "0"))) with_cusolvermp = bool(int(os.getenv("NVTE_WITH_CUSOLVERMP", "0"))) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 6237c6732b9..dae276d335d 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -19,9 +19,7 @@ from typing import Optional, Tuple import warnings -# Minimum cuDNN version supported by Transformer Engine, as (major, minor, patch). -# Keep in sync with kMinCudnnVersion in transformer_engine/common/cudnn_min_version.h. -MIN_CUDNN_VERSION = (9, 12, 0) +from ._versions import MIN_CUDNN_VERSION def decode_cudnn_version(encoded_version: int) -> Tuple[int, int, int]: diff --git a/transformer_engine/common/_versions.py b/transformer_engine/common/_versions.py new file mode 100644 index 00000000000..ecab8709707 --- /dev/null +++ b/transformer_engine/common/_versions.py @@ -0,0 +1,13 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Minimum versions of Transformer Engine dependencies. + +The build loads this file directly, without importing the transformer_engine package, so it +must not import anything. +""" + +# Minimum cuDNN version supported by Transformer Engine, as (major, minor, patch). +# Keep in sync with kMinCudnnVersion in transformer_engine/common/cudnn_min_version.h. +MIN_CUDNN_VERSION = (9, 12, 0) diff --git a/transformer_engine/common/cudnn_min_version.h b/transformer_engine/common/cudnn_min_version.h index 1e3ff0a34a3..79fb780ede3 100644 --- a/transformer_engine/common/cudnn_min_version.h +++ b/transformer_engine/common/cudnn_min_version.h @@ -17,7 +17,7 @@ namespace transformer_engine { // Minimum cuDNN version supported by Transformer Engine, encoded as cudnnGetVersion() reports it. -// Keep in sync with MIN_CUDNN_VERSION in transformer_engine/common/__init__.py. +// Keep in sync with MIN_CUDNN_VERSION in transformer_engine/common/_versions.py. constexpr size_t kMinCudnnVersion = 91200; static_assert(static_cast(CUDNN_VERSION) >= kMinCudnnVersion, From 929f70ae8d52335a3686f9c1f82555f8dd30fca0 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 17:51:56 -0600 Subject: [PATCH 22/24] Revert "Define the minimum cuDNN version once for the build and runtime" This reverts commit aaaf3d549fb90a22537cd922dd542c01c9f4251c. Signed-off-by: Fred Heinecke --- build_tools/build_backend.py | 4 ++-- build_tools/utils.py | 17 ++--------------- setup.py | 4 ++-- transformer_engine/common/__init__.py | 4 +++- transformer_engine/common/_versions.py | 13 ------------- transformer_engine/common/cudnn_min_version.h | 2 +- 6 files changed, 10 insertions(+), 34 deletions(-) delete mode 100644 transformer_engine/common/_versions.py diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py index 18899042938..cca1848f783 100644 --- a/build_tools/build_backend.py +++ b/build_tools/build_backend.py @@ -17,7 +17,7 @@ requested_cuda_version, setuptools_hooks, ) -from build_tools.utils import min_cudnn_version, nvcc_path +from build_tools.utils import MIN_CUDNN_VERSION, nvcc_path # cuda-toolkit pins all components to the same CUDA release. The nvcc extra includes the # compiler's crt and nvvm components. @@ -216,7 +216,7 @@ def _cuda_build_requirements(config_settings: ConfigSettings) -> List[str]: if not os.getenv("CUDA_HOME"): requirements += [ f"cuda-toolkit[{','.join(_CUDA_TOOLKIT_EXTRAS)}]=={cuda_version}.*", - f"nvidia-cudnn-cu{cuda_major}>={min_cudnn_version()}", + f"nvidia-cudnn-cu{cuda_major}>={MIN_CUDNN_VERSION}", ] return requirements diff --git a/build_tools/utils.py b/build_tools/utils.py index 8fd8b89eb51..51bfac9dd28 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -7,7 +7,6 @@ import functools import glob import importlib -import importlib.util import os import re import shlex @@ -367,20 +366,8 @@ def nccl_lib_path() -> Optional[Path]: return None -@functools.lru_cache(maxsize=None) -def min_cudnn_version() -> str: - """Return the minimum supported cuDNN version from transformer_engine/common/_versions.py. - - The file is loaded directly because importing the transformer_engine package loads its - shared libraries. The framework sdists do not contain it, so it is only read when needed. - """ - versions_file = ( - Path(__file__).resolve().parent.parent / "transformer_engine" / "common" / "_versions.py" - ) - spec = importlib.util.spec_from_file_location("_transformer_engine_versions", versions_file) - versions = importlib.util.module_from_spec(spec) - spec.loader.exec_module(versions) - return ".".join(str(part) for part in versions.MIN_CUDNN_VERSION) +# Keep in sync with MIN_CUDNN_VERSION in transformer_engine/common/__init__.py. +MIN_CUDNN_VERSION = "9.12" def python_package_requirement(package_name: str, minimum: Optional[str] = None) -> str: diff --git a/setup.py b/setup.py index 61d4d5798e2..0f383c0022c 100644 --- a/setup.py +++ b/setup.py @@ -18,6 +18,7 @@ from build_tools.build_ext import CMakeExtension, get_build_ext from build_tools.te_version import te_version from build_tools.utils import ( + MIN_CUDNN_VERSION, cublasmp_home_path, cuda_archs, cuda_home_path, @@ -30,7 +31,6 @@ get_frameworks, remove_dups, min_python_version_str, - min_cudnn_version, mpi_home_path, nccl_ep_enabled, get_max_jobs_for_parallel_build, @@ -193,7 +193,7 @@ def require_installed(package_name: str, minimum: Optional[str] = None) -> None: f"cuda-toolkit[cudart,cublas,nvrtc,curand]>={cuda_major_minor},<{cuda_major + 1}" ) if cudnn_from_python_package(): - require_installed(f"nvidia-cudnn-cu{cuda_major}", minimum=min_cudnn_version()) + require_installed(f"nvidia-cudnn-cu{cuda_major}", minimum=MIN_CUDNN_VERSION) with_cublasmp = bool(int(os.getenv("NVTE_WITH_CUBLASMP", "0"))) with_cusolvermp = bool(int(os.getenv("NVTE_WITH_CUSOLVERMP", "0"))) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index dae276d335d..6237c6732b9 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -19,7 +19,9 @@ from typing import Optional, Tuple import warnings -from ._versions import MIN_CUDNN_VERSION +# Minimum cuDNN version supported by Transformer Engine, as (major, minor, patch). +# Keep in sync with kMinCudnnVersion in transformer_engine/common/cudnn_min_version.h. +MIN_CUDNN_VERSION = (9, 12, 0) def decode_cudnn_version(encoded_version: int) -> Tuple[int, int, int]: diff --git a/transformer_engine/common/_versions.py b/transformer_engine/common/_versions.py deleted file mode 100644 index ecab8709707..00000000000 --- a/transformer_engine/common/_versions.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Minimum versions of Transformer Engine dependencies. - -The build loads this file directly, without importing the transformer_engine package, so it -must not import anything. -""" - -# Minimum cuDNN version supported by Transformer Engine, as (major, minor, patch). -# Keep in sync with kMinCudnnVersion in transformer_engine/common/cudnn_min_version.h. -MIN_CUDNN_VERSION = (9, 12, 0) diff --git a/transformer_engine/common/cudnn_min_version.h b/transformer_engine/common/cudnn_min_version.h index 79fb780ede3..1e3ff0a34a3 100644 --- a/transformer_engine/common/cudnn_min_version.h +++ b/transformer_engine/common/cudnn_min_version.h @@ -17,7 +17,7 @@ namespace transformer_engine { // Minimum cuDNN version supported by Transformer Engine, encoded as cudnnGetVersion() reports it. -// Keep in sync with MIN_CUDNN_VERSION in transformer_engine/common/_versions.py. +// Keep in sync with MIN_CUDNN_VERSION in transformer_engine/common/__init__.py. constexpr size_t kMinCudnnVersion = 91200; static_assert(static_cast(CUDNN_VERSION) >= kMinCudnnVersion, From 4262d4b57ec08f11ceeea5dabcf720cf17666998 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 17:58:00 -0600 Subject: [PATCH 23/24] Keep Python package build requirements out of runtime dependencies Remove the runtime dependencies that were declared for libraries the core library was built against from Python packages, and restore the original NCCL discovery order for NCCL EP. Build isolation only changes the build environment; runtime requirements are unchanged. The installation guide now notes that the CUDA libraries, cuDNN, and optional dependencies must be installed at runtime. Signed-off-by: Fred Heinecke --- build_tools/utils.py | 52 +++++------------------ docs/envvars.rst | 5 +-- docs/installation.rst | 4 +- setup.py | 97 +++++++------------------------------------ 4 files changed, 29 insertions(+), 129 deletions(-) diff --git a/build_tools/utils.py b/build_tools/utils.py index 51bfac9dd28..c127d1447a2 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -370,29 +370,6 @@ def nccl_lib_path() -> Optional[Path]: MIN_CUDNN_VERSION = "9.12" -def python_package_requirement(package_name: str, minimum: Optional[str] = None) -> str: - """Require an installed Python package within its installed major version. - - The lower bound is `minimum` if given, otherwise the installed version. - """ - installed = get_version(package_name) - major = int(installed.split(".", maxsplit=1)[0]) - return f"{package_name}>={minimum or installed},<{major + 1}" - - -def cudnn_from_python_package() -> bool: - """Whether CMake finds cuDNN in the nvidia-cudnn-cu* Python package. - - cmake/cuDNN.cmake prefers CUDNN_INCLUDE_PATH and CUDNN_PATH over the Python package. - """ - if os.getenv("CUDNN_INCLUDE_PATH") or os.getenv("CUDNN_PATH"): - return False - return ( - _python_package_path(f"nvidia-cudnn-cu{cuda_version()[0]}", "nvidia/cudnn/include") - is not None - ) - - def _python_package_path(package_name: str, relative_path: str) -> Optional[Path]: """Return a path inside an installed Python package, or `None` if it does not exist.""" try: @@ -474,33 +451,26 @@ def cuda_toolkit_include_path() -> Optional[Path]: return None -def _nvidia_package_roots() -> List[Path]: - """Return the directories of the `nvidia` namespace package for NVIDIA Python packages.""" +@functools.lru_cache(maxsize=None) +def get_cuda_include_dirs() -> List[Path]: + """Returns the CUDA header directories.""" + + force_wheels = bool(int(os.getenv("NVTE_BUILD_USE_NVIDIA_WHEELS", "0"))) try: import nvidia except ModuleNotFoundError: - return [] - return [Path(path).resolve() for path in nvidia.__path__] - + nvidia_roots = [] + else: + nvidia_roots = [Path(path).resolve() for path in nvidia.__path__] -def cuda_toolkit_from_python_packages() -> bool: - """Whether the CUDA Toolkit used for the build is installed from Python packages.""" cuda_home = cuda_home_path() - return cuda_home is not None and any( - cuda_home.is_relative_to(root) for root in _nvidia_package_roots() + cuda_home_from_wheel = cuda_home is not None and any( + cuda_home.is_relative_to(root) for root in nvidia_roots ) - -@functools.lru_cache(maxsize=None) -def get_cuda_include_dirs() -> List[Path]: - """Returns the CUDA header directories.""" - - force_wheels = bool(int(os.getenv("NVTE_BUILD_USE_NVIDIA_WHEELS", "0"))) - nvidia_roots = _nvidia_package_roots() - # A system toolkit contains all necessary headers under one include # directory. NVIDIA wheels split them across namespace packages. - if not force_wheels and not cuda_toolkit_from_python_packages(): + if not force_wheels and not cuda_home_from_wheel: if (include_path := cuda_toolkit_include_path()) is not None: return [include_path] diff --git a/docs/envvars.rst b/docs/envvars.rst index b0d88df51fa..e9a3483472d 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -92,9 +92,8 @@ Each of the following options can also be set for an isolated PEP 517 build with In an isolated build, the Python package listed for an option is installed as a build requirement unless the option's installation directory variable is set. Without build isolation, the installation directory variable is used if set, and the package is used if -installed. When Transformer Engine is built against a package, the built Transformer Engine -package depends on it, and ``libtransformer_engine.so`` finds it at runtime in the same -environment. +installed. ``libtransformer_engine.so`` finds these packages at runtime when they are installed +in the same environment. .. envvar:: NVTE_UB_WITH_MPI diff --git a/docs/installation.rst b/docs/installation.rst index e7b36e7e112..791f47929b7 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -178,9 +178,9 @@ The same options can be used to install from GitHub: pip3 install --build-constraint build-constraints.txt --config-settings framework=pytorch git+https://github.com/NVIDIA/TransformerEngine.git@stable -Unless `CUDA_HOME` is set, the CUDA Toolkit and cuDNN for the build are installed into the build environment from PyPI, matching the version of the installed CUDA Toolkit. If no CUDA Toolkit is installed, specify the CUDA version with the `cuda-version` config setting (e.g. `cuda-version=13.0`). The installed Transformer Engine package then depends on the CUDA Toolkit and cuDNN packages from PyPI, so when building framework support, choose the CUDA version that PyTorch or JAX was built with. To build with an installed CUDA Toolkit and cuDNN instead, set `CUDA_HOME` to the location of the CUDA Toolkit. +Unless `CUDA_HOME` is set, the CUDA Toolkit and cuDNN for the build are installed into the build environment from PyPI, matching the version of the installed CUDA Toolkit. If no CUDA Toolkit is installed, specify the CUDA version with the `cuda-version` config setting (e.g. `cuda-version=13.0`). These packages are only used for the build, so the CUDA libraries and cuDNN must also be installed at runtime, e.g. as dependencies of PyTorch or JAX. When building framework support, choose the CUDA version that PyTorch or JAX was built with. To build with an installed CUDA Toolkit and cuDNN instead, set `CUDA_HOME` to the location of the CUDA Toolkit. -Optional dependencies, such as cuBLASMp and NVSHMEM, are enabled with config settings as well (e.g. `with-cublasmp=1`). In an isolated build, their Python packages are installed into the build environment, and the installed Transformer Engine package depends on them. Each config setting has an equivalent environment variable, e.g. `NVTE_FRAMEWORK` for `framework` and `NVTE_CUDA_VERSION` for `cuda-version`. See :doc:`envvars` for the complete list. +Optional dependencies, such as cuBLASMp and NVSHMEM, are enabled with config settings as well (e.g. `with-cublasmp=1`). In an isolated build, their Python packages are installed into the build environment. To use the optional dependencies at runtime, install the same Python packages in the environment where Transformer Engine is installed. Each config setting has an equivalent environment variable, e.g. `NVTE_FRAMEWORK` for `framework` and `NVTE_CUDA_VERSION` for `cuda-version`. See :doc:`envvars` for the complete list. Other installers support build constraints as well. For example, with `uv`: diff --git a/setup.py b/setup.py index 0f383c0022c..8031f8031bb 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ import subprocess import time from pathlib import Path -from typing import List, Optional, Tuple +from typing import List, Tuple import setuptools from wheel.bdist_wheel import bdist_wheel @@ -18,14 +18,10 @@ from build_tools.build_ext import CMakeExtension, get_build_ext from build_tools.te_version import te_version from build_tools.utils import ( - MIN_CUDNN_VERSION, cublasmp_home_path, cuda_archs, cuda_home_path, - cuda_toolkit_from_python_packages, - cuda_version, cudnn_frontend_include_path, - cudnn_from_python_package, cusolvermp_home_path, get_bolt_build_flags, get_frameworks, @@ -36,7 +32,6 @@ get_max_jobs_for_parallel_build, nvcc_path, nvshmem_home_path, - python_package_requirement, ) frameworks = get_frameworks() @@ -166,62 +161,9 @@ def setup_requirements() -> Tuple[List[str], List[str]]: install_reqs.extend(install_requirements()) test_reqs.extend(test_requirements()) - install_reqs.extend(python_package_runtime_requirements()) - return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]] -def python_package_runtime_requirements() -> List[str]: - """Runtime requirements on the Python packages that the core library is built against. - - Libraries found through a *_HOME variable or on the system are not declared. Each - requirement stays within the major version of the package used for the build, so that it - is compatible with the exact versions that frameworks pin. - """ - cuda_major = cuda_version()[0] - requirements: List[str] = [] - - def require_installed(package_name: str, minimum: Optional[str] = None) -> None: - try: - requirements.append(python_package_requirement(package_name, minimum)) - except metadata.PackageNotFoundError: - pass - - if cuda_toolkit_from_python_packages(): - cuda_major_minor = ".".join(str(part) for part in cuda_version()[:2]) - requirements.append( - f"cuda-toolkit[cudart,cublas,nvrtc,curand]>={cuda_major_minor},<{cuda_major + 1}" - ) - if cudnn_from_python_package(): - require_installed(f"nvidia-cudnn-cu{cuda_major}", minimum=MIN_CUDNN_VERSION) - - with_cublasmp = bool(int(os.getenv("NVTE_WITH_CUBLASMP", "0"))) - with_cusolvermp = bool(int(os.getenv("NVTE_WITH_CUSOLVERMP", "0"))) - nccl_pip_root = _python_package_nccl_home() - if nccl_pip_root is not None: - if nccl_ep_enabled(archs): - try: - nccl_from_pip = _discover_nccl_home() == nccl_pip_root - except RuntimeError: - nccl_from_pip = False - else: - # CMake looks for NCCL in the Python package before system locations. - nccl_from_pip = with_cublasmp or with_cusolvermp - if nccl_from_pip: - require_installed(f"nvidia-nccl-cu{cuda_major}", "2.30" if with_cublasmp else "2") - - if with_cublasmp and not os.getenv("CUBLASMP_HOME"): - require_installed(f"nvidia-cublasmp-cu{cuda_major}") - if with_cusolvermp and not os.getenv("CUSOLVERMP_HOME"): - require_installed(f"nvidia-cusolvermp-cu{cuda_major}") - if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", "0"))) and not os.getenv("NVSHMEM_HOME"): - require_installed(f"nvidia-nvshmem-cu{cuda_major}") - if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))) and not os.getenv("MPI_HOME"): - require_installed("openmpi") - - return requirements - - def _discover_nccl_home() -> str: """Resolve NCCL_HOME, preferring the NCCL the dynamic loader resolves at runtime. @@ -238,11 +180,6 @@ def _discover_nccl_home() -> str: f"'{env_home}/include/nccl.h' was not found; falling back to system probes." ) - # A build against CUDA Toolkit Python packages also depends on the nvidia-nccl - # wheel at runtime, which libtransformer_engine's RUNPATH prefers. - if cuda_toolkit_from_python_packages() and (pip_root := _python_package_nccl_home()): - return pip_root - 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") @@ -273,30 +210,24 @@ def _discover_nccl_home() -> str: ): return str(p) - # Fall back to a pip-installed NCCL when no system NCCL is present. - if pip_root := _python_package_nccl_home(): - return pip_root - - raise RuntimeError( - "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." - ) - - -def _python_package_nccl_home() -> Optional[str]: - """Return the NCCL installed from an nvidia-nccl-cu* wheel under nvidia/nccl, if any.""" + # 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): - return None - if spec is None or not spec.submodule_search_locations: - return None - - pip_root = Path(next(iter(spec.submodule_search_locations))) - if (pip_root / "include" / "nccl.h").exists() and (pip_root / "lib" / "libnccl.so.2").exists(): - return str(pip_root) - return None + 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: From 289a162820569a0ec8b27f06e6744b3e323614e6 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 9 Nov 2026 18:00:19 -0600 Subject: [PATCH 24/24] Build NCCL EP against the nvidia-nccl wheel when the CUDA Toolkit comes from Python packages With a CUDA Toolkit from Python packages, the NCCL loaded at runtime is most likely the nvidia-nccl wheel, which the core library's RUNPATH prefers when it is installed. Check for it before system NCCL in that case. Builds with an installed CUDA Toolkit keep the system-first order. Signed-off-by: Fred Heinecke --- build_tools/utils.py | 29 ++++++++++++++++++----------- setup.py | 40 ++++++++++++++++++++++++++-------------- 2 files changed, 44 insertions(+), 25 deletions(-) diff --git a/build_tools/utils.py b/build_tools/utils.py index c127d1447a2..d4af10a55a5 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -451,26 +451,33 @@ def cuda_toolkit_include_path() -> Optional[Path]: return None -@functools.lru_cache(maxsize=None) -def get_cuda_include_dirs() -> List[Path]: - """Returns the CUDA header directories.""" - - force_wheels = bool(int(os.getenv("NVTE_BUILD_USE_NVIDIA_WHEELS", "0"))) +def _nvidia_package_roots() -> List[Path]: + """Return the directories of the `nvidia` namespace package for NVIDIA Python packages.""" try: import nvidia except ModuleNotFoundError: - nvidia_roots = [] - else: - nvidia_roots = [Path(path).resolve() for path in nvidia.__path__] + return [] + return [Path(path).resolve() for path in nvidia.__path__] + +def cuda_toolkit_from_python_packages() -> bool: + """Whether the CUDA Toolkit used for the build is installed from Python packages.""" cuda_home = cuda_home_path() - cuda_home_from_wheel = cuda_home is not None and any( - cuda_home.is_relative_to(root) for root in nvidia_roots + return cuda_home is not None and any( + cuda_home.is_relative_to(root) for root in _nvidia_package_roots() ) + +@functools.lru_cache(maxsize=None) +def get_cuda_include_dirs() -> List[Path]: + """Returns the CUDA header directories.""" + + force_wheels = bool(int(os.getenv("NVTE_BUILD_USE_NVIDIA_WHEELS", "0"))) + nvidia_roots = _nvidia_package_roots() + # A system toolkit contains all necessary headers under one include # directory. NVIDIA wheels split them across namespace packages. - if not force_wheels and not cuda_home_from_wheel: + if not force_wheels and not cuda_toolkit_from_python_packages(): if (include_path := cuda_toolkit_include_path()) is not None: return [include_path] diff --git a/setup.py b/setup.py index 8031f8031bb..020902aa042 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ import subprocess import time from pathlib import Path -from typing import List, Tuple +from typing import List, Optional, Tuple import setuptools from wheel.bdist_wheel import bdist_wheel @@ -21,6 +21,7 @@ cublasmp_home_path, cuda_archs, cuda_home_path, + cuda_toolkit_from_python_packages, cudnn_frontend_include_path, cusolvermp_home_path, get_bolt_build_flags, @@ -180,6 +181,11 @@ def _discover_nccl_home() -> str: f"'{env_home}/include/nccl.h' was not found; falling back to system probes." ) + # With a CUDA Toolkit from Python packages, the NCCL loaded at runtime is most likely the + # nvidia-nccl wheel, which libtransformer_engine's RUNPATH prefers when it is installed. + if cuda_toolkit_from_python_packages() and (pip_root := _python_package_nccl_home()): + return pip_root + 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") @@ -210,24 +216,30 @@ def _discover_nccl_home() -> str: ): 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. + # Fall back to a pip-installed NCCL when no system NCCL is present. + if pip_root := _python_package_nccl_home(): + return pip_root + + raise RuntimeError( + "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." + ) + + +def _python_package_nccl_home() -> Optional[str]: + """Return the NCCL installed from an nvidia-nccl-cu* wheel under nvidia/nccl, if any.""" 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." - ) + return None + if spec is None or not spec.submodule_search_locations: + return None + + pip_root = Path(next(iter(spec.submodule_search_locations))) + if (pip_root / "include" / "nccl.h").exists() and (pip_root / "lib" / "libnccl.so.2").exists(): + return str(pip_root) + return None def build_nccl_ep_submodule() -> str: