From 5f2889ebd6ad096c9d3dfa0f82447ddd7bef94f4 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 24 Jul 2026 12:26:41 -0500 Subject: [PATCH 1/2] 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 0000000000..0b89294979 --- /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 b3765a06bd..014720206b 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 2c9f224c14..4f7278a7fc 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.25.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 06aa62aa19367799e843074bc65f9e464c3a1ba4 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 2/2] [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 0b89294979..f340b47ba9 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,