From f4df883a34cab9d7437002f6152f180f7864db88 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 24 Jul 2026 10:10:29 -0500 Subject: [PATCH 1/5] [CI] Improve build time dependency resolution Signed-off-by: Fred Heinecke --- build_tools/build_ext.py | 28 +++++++ build_tools/jax.py | 12 ++- build_tools/utils.py | 175 +++++++++++++++++++++++++++++++-------- 3 files changed, 179 insertions(+), 36 deletions(-) diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index cbb8838b00..704354ecfd 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -19,6 +19,8 @@ from .utils import ( cmake_bin, + cuda_home_path, + cuda_version, debug_build_enabled, found_ninja, get_frameworks, @@ -61,6 +63,29 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: f"-DCMAKE_BUILD_TYPE={build_type}", f"-DCMAKE_INSTALL_PREFIX={install_dir}", ] + + discovered_cuda_home = cuda_home_path() + if discovered_cuda_home is not None: + configure_command += [ + f"-DCUDAToolkit_ROOT={discovered_cuda_home}", + # CUDA wheels use `lib`, while nvcc's host linker expects `lib64`. + f"-DCMAKE_CUDA_FLAGS=-L{discovered_cuda_home}/lib", + ] + + if cuda_full_version := cuda_version(): + cuda_major_version = cuda_full_version[0] + cuda_lib_dir = discovered_cuda_home / "lib" + configure_command += [ + f"-DCUDA_CUDART={cuda_lib_dir / f'libcudart.so.{cuda_major_version}'}", + f"-DCUDA_cudart_LIBRARY={cuda_lib_dir / f'libcudart.so.{cuda_major_version}'}", + f"-DCUDA_cublas_LIBRARY={cuda_lib_dir / f'libcublas.so.{cuda_major_version}'}", + f"-DCUDA_cublasLt_LIBRARY={cuda_lib_dir / f'libcublasLt.so.{cuda_major_version}'}", + ] + + discovered_nvcc_path = nvcc_path() + if discovered_nvcc_path is not None: + configure_command.append(f"-DCMAKE_CUDA_COMPILER={discovered_nvcc_path}") + if bool(int(os.getenv("NVTE_USE_CCACHE", "0"))): ccache_bin = os.getenv("NVTE_CCACHE_BIN", "ccache") configure_command += [ @@ -185,6 +210,9 @@ def _compile_fn(obj, src, ext, cc_args, extra_postargs, pp_opts) -> None: and not framework_extension_only ): nvcc_bin = nvcc_path() + if nvcc_bin is None: + raise RuntimeError(f"NVCC not found and is required for building CUDA source {src}") + self.compiler.set_executable("compiler_so", str(nvcc_bin)) if isinstance(cflags, dict): cflags = cflags["nvcc"] diff --git a/build_tools/jax.py b/build_tools/jax.py index d61c26c128..031432e6f9 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -16,6 +16,8 @@ cudnn_frontend_include_path, debug_build_enabled, setup_mpi_flags, + nccl_include_path, + nccl_lib_path, nccl_ep_enabled, ) from typing import List @@ -90,6 +92,8 @@ def setup_jax_extension( # Header files include_dirs = get_cuda_include_dirs() + if (discovered_nccl_include_path := nccl_include_path()) is not None: + include_dirs.append(discovered_nccl_include_path) include_dirs.append(cudnn_frontend_include_path()) include_dirs.extend( [ @@ -117,6 +121,12 @@ def setup_jax_extension( if nccl_ep_enabled(): cxx_flags.append("-DNVTE_WITH_NCCL_EP") + kwargs = {} + if (discovered_nccl_lib_path := nccl_lib_path()) is not None: + kwargs["extra_objects"] = [str(discovered_nccl_lib_path)] + else: + kwargs["libraries"] = ["nccl"] + # Define TE/JAX as a Pybind11Extension from pybind11.setup_helpers import Pybind11Extension @@ -125,5 +135,5 @@ def setup_jax_extension( sources=[str(path) for path in sources], include_dirs=[str(path) for path in include_dirs], extra_compile_args=cxx_flags, - libraries=["nccl"], + **kwargs, ) diff --git a/build_tools/utils.py b/build_tools/utils.py index 6a8168e7d0..f741359957 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -16,7 +16,7 @@ from pathlib import Path from importlib.metadata import PackageNotFoundError, distribution, version as get_version from subprocess import CalledProcessError -from typing import List, Optional, Tuple, Union +from typing import Callable, List, Optional, Tuple, Union # Needs to stay consistent with .pre-commit-config.yaml config. @@ -175,6 +175,90 @@ def found_pybind11() -> bool: return False +@functools.lru_cache(maxsize=None) +def cuda_home_path() -> Optional[Path]: + """Returns the CUDA home path. This path should contain binaries (e.g. nvcc), headers, and libraries. + + Returns `None` if CUDA is not found.""" + if cuda_home := os.getenv("CUDA_HOME"): + return Path(cuda_home) + + # Check if site-packages contains an `nvidia` directory + for site_package in sys.path: + if not Path(site_package).is_dir(): + continue + + nvidia_dir = Path(site_package) / "nvidia" + if not nvidia_dir.is_dir(): + continue + + if Path(nvidia_dir / "bin").is_dir(): + return nvidia_dir + + # In this case there must be a "CUDA version directory" like `cu12` or `cu13` + # that contains the binaries, headers, and libraries + for cuda_version_dir in nvidia_dir.iterdir(): + if not cuda_version_dir.is_dir(): + continue + + # Verify that the directory name matches the expected `cu##` format + if not re.match(r"cu\d+", cuda_version_dir.name): + continue + + if not (cuda_version_dir / "bin").is_dir(): + continue + + return cuda_version_dir + + return None + + +@functools.lru_cache(maxsize=None) +def nccl_root_path() -> Optional[Path]: + """Return the NCCL installation root. + + Returns `None` if NCCL is not found.""" + if (cuda_home := cuda_home_path()) is not None: + nccl_root = cuda_home.parent / "nccl" + if nccl_root.is_dir(): + return nccl_root + + # Check NCCL Python packages + for package_name in ["nvidia-nccl-cu13", "nvidia-nccl-cu12"]: + try: + nccl_distribution = distribution(package_name) + except PackageNotFoundError: + continue + + nccl_root = Path(nccl_distribution.locate_file("nvidia/nccl")) + if nccl_root.is_dir(): + return nccl_root + + return None + + +@functools.lru_cache(maxsize=None) +def nccl_include_path() -> Optional[Path]: + """Return the NCCL include directory.""" + if (nccl_root := nccl_root_path()) is not None: + include_path = nccl_root / "include" + if include_path.is_dir(): + return include_path + + return None + + +@functools.lru_cache(maxsize=None) +def nccl_lib_path() -> Optional[Path]: + """Return the NCCL shared library path.""" + if (nccl_root := nccl_root_path()) is not None: + lib_path = nccl_root / "lib" / "libnccl.so.2" + if lib_path.is_file(): + return lib_path + + return None + + @functools.lru_cache(maxsize=None) def cuda_toolkit_include_path() -> Tuple[str, str]: """Returns root path for cuda toolkit includes. @@ -198,30 +282,49 @@ def cuda_toolkit_include_path() -> Tuple[str, str]: @functools.lru_cache(maxsize=None) -def nvcc_path() -> Tuple[str, str]: - """Returns the NVCC binary path. +def nvcc_path() -> Optional[Path]: + """Get the NVCC binary path. - Throws FileNotFoundError if NVCC is not found.""" - # Try finding NVCC - nvcc_bin: Optional[Path] = None - if nvcc_bin is None and os.getenv("CUDA_HOME"): - # Check in CUDA_HOME - cuda_home = Path(os.getenv("CUDA_HOME")) - nvcc_bin = cuda_home / "bin" / "nvcc" - if nvcc_bin is None: - # Check if nvcc is in path - nvcc_bin = shutil.which("nvcc") - if nvcc_bin is not None: - cuda_home = Path(nvcc_bin.rstrip("/bin/nvcc")) - nvcc_bin = Path(nvcc_bin) - if nvcc_bin is None: - # Last-ditch guess in /usr/local/cuda - cuda_home = Path("/usr/local/cuda") - nvcc_bin = cuda_home / "bin" / "nvcc" - if not nvcc_bin.is_file(): - raise FileNotFoundError(f"Could not find NVCC at {nvcc_bin}") + Returns `None` if NVCC is not found. + """ + + def lookup_via_cuda_home() -> Optional[str]: + if (cuda_home := cuda_home_path()) is not None: + return cuda_home / "bin" / "nvcc" + return None + + def lookup_via_path() -> Optional[str]: + if (nvcc_bin := shutil.which("nvcc")) is not None: + return nvcc_bin + return None + + def lookup_via_distribution() -> Optional[str]: + try: + return distribution("nvidia-cuda-nvcc").locate_file("bin/nvcc") + except PackageNotFoundError: + return None + + def lookup_via_local_cuda() -> Optional[str]: + return "/usr/local/cuda/bin/nvcc" + + nvcc_lookup_funcs: List[Callable[[], Optional[str]]] = [ + lookup_via_cuda_home, + lookup_via_path, + lookup_via_distribution, + lookup_via_local_cuda, + ] + + for nvcc_lookup_func in nvcc_lookup_funcs: + nvcc_bin = nvcc_lookup_func() + + if nvcc_bin is None: + continue + + nvcc_bin_path = Path(nvcc_bin) + if nvcc_bin_path.is_file(): + return nvcc_bin_path - return nvcc_bin + return None @functools.lru_cache(maxsize=None) @@ -324,13 +427,9 @@ def cuda_version() -> Tuple[int, ...]: and check pip version. """ - try: - nvcc_bin = nvcc_path() - except FileNotFoundError as e: - pass - else: + if (nvcc_bin := nvcc_path()) is not None: output = subprocess.run( - [nvcc_bin, "-V"], + [str(nvcc_bin), "-V"], capture_output=True, check=True, universal_newlines=True, @@ -339,12 +438,18 @@ def cuda_version() -> Tuple[int, ...]: version = match.group(1).split(".") return tuple(int(v) for v in version) - try: - version_str = get_version("nvidia-cuda-runtime-cu12") - version_tuple = tuple(int(part) for part in version_str.split(".") if part.isdigit()) - return version_tuple - except importlib.metadata.PackageNotFoundError: - raise RuntimeError("Could neither find NVCC executable nor CUDA runtime Python package.") + version_str: Optional[str] = None + package_names = ["nvidia-cuda-runtime", "nvidia-cuda-runtime-cu13", "nvidia-cuda-runtime-cu12"] + + for package_name in package_names: + try: + version_str = get_version(package_name) + except PackageNotFoundError: + pass + else: + return tuple(int(part) for part in version_str.split(".") if part.isdigit()) + + raise RuntimeError(f"Could neither find NVCC executable nor CUDA runtime Python package for {package_names}.") def get_frameworks() -> List[str]: From 4e0b2d0b49f2d746718fe9f471da6e0c876579d2 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 15:29:34 +0000 Subject: [PATCH 2/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- build_tools/build_ext.py | 8 ++++++-- build_tools/utils.py | 8 +++++--- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index 704354ecfd..d6c8ddb315 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -79,7 +79,9 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: f"-DCUDA_CUDART={cuda_lib_dir / f'libcudart.so.{cuda_major_version}'}", f"-DCUDA_cudart_LIBRARY={cuda_lib_dir / f'libcudart.so.{cuda_major_version}'}", f"-DCUDA_cublas_LIBRARY={cuda_lib_dir / f'libcublas.so.{cuda_major_version}'}", - f"-DCUDA_cublasLt_LIBRARY={cuda_lib_dir / f'libcublasLt.so.{cuda_major_version}'}", + ( + f"-DCUDA_cublasLt_LIBRARY={cuda_lib_dir / f'libcublasLt.so.{cuda_major_version}'}" + ), ] discovered_nvcc_path = nvcc_path() @@ -211,7 +213,9 @@ def _compile_fn(obj, src, ext, cc_args, extra_postargs, pp_opts) -> None: ): nvcc_bin = nvcc_path() if nvcc_bin is None: - raise RuntimeError(f"NVCC not found and is required for building CUDA source {src}") + raise RuntimeError( + f"NVCC not found and is required for building CUDA source {src}" + ) self.compiler.set_executable("compiler_so", str(nvcc_bin)) if isinstance(cflags, dict): diff --git a/build_tools/utils.py b/build_tools/utils.py index f741359957..da4ecf957b 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -16,7 +16,7 @@ from pathlib import Path from importlib.metadata import PackageNotFoundError, distribution, version as get_version from subprocess import CalledProcessError -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable, List, Optional, Tuple, Union # Needs to stay consistent with .pre-commit-config.yaml config. @@ -255,7 +255,7 @@ def nccl_lib_path() -> Optional[Path]: lib_path = nccl_root / "lib" / "libnccl.so.2" if lib_path.is_file(): return lib_path - + return None @@ -449,7 +449,9 @@ def cuda_version() -> Tuple[int, ...]: else: return tuple(int(part) for part in version_str.split(".") if part.isdigit()) - raise RuntimeError(f"Could neither find NVCC executable nor CUDA runtime Python package for {package_names}.") + raise RuntimeError( + f"Could neither find NVCC executable nor CUDA runtime Python package for {package_names}." + ) def get_frameworks() -> List[str]: From 7e915b9a6b0f15e2cc92703166f3a38670a41042 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 3 Aug 2026 18:11:10 -0500 Subject: [PATCH 3/5] Only target wheel-installed deps when they are available (fix lib/lib64 bug) Signed-off-by: Fred Heinecke --- build_tools/build_ext.py | 33 ++++++++++++++++++--------------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index d6c8ddb315..e2ee4fe581 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -66,23 +66,26 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: discovered_cuda_home = cuda_home_path() if discovered_cuda_home is not None: - configure_command += [ - f"-DCUDAToolkit_ROOT={discovered_cuda_home}", - # CUDA wheels use `lib`, while nvcc's host linker expects `lib64`. - f"-DCMAKE_CUDA_FLAGS=-L{discovered_cuda_home}/lib", - ] + configure_command.append(f"-DCUDAToolkit_ROOT={discovered_cuda_home}") - if cuda_full_version := cuda_version(): + # 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() and (cuda_full_version := cuda_version()): cuda_major_version = cuda_full_version[0] - cuda_lib_dir = discovered_cuda_home / "lib" - configure_command += [ - f"-DCUDA_CUDART={cuda_lib_dir / f'libcudart.so.{cuda_major_version}'}", - f"-DCUDA_cudart_LIBRARY={cuda_lib_dir / f'libcudart.so.{cuda_major_version}'}", - f"-DCUDA_cublas_LIBRARY={cuda_lib_dir / f'libcublas.so.{cuda_major_version}'}", - ( - f"-DCUDA_cublasLt_LIBRARY={cuda_lib_dir / f'libcublasLt.so.{cuda_major_version}'}" - ), - ] + cuda_libraries = { + "CUDA_CUDART": cuda_lib_dir / f"libcudart.so.{cuda_major_version}", + "CUDA_cudart_LIBRARY": cuda_lib_dir / f"libcudart.so.{cuda_major_version}", + "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() + ) discovered_nvcc_path = nvcc_path() if discovered_nvcc_path is not None: From eb5358faa25f3b6d31e40147da790322021453e8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 3 Aug 2026 23:12:12 +0000 Subject: [PATCH 4/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- build_tools/build_ext.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index e2ee4fe581..e8a60f91e7 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -83,8 +83,7 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: 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() + f"-D{variable}={library}" for variable, library in cuda_libraries.items() ) discovered_nvcc_path = nvcc_path() From d2f9f1d6ce21d3d330b6ac871dad54d64055dd91 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Mon, 31 Aug 2026 13:43:21 -0500 Subject: [PATCH 5/5] Standardize CUDA home path and NVCC lookup Signed-off-by: Fred Heinecke --- build_tools/utils.py | 160 +++++++++++++++++++------------------------ setup.py | 6 ++ 2 files changed, 77 insertions(+), 89 deletions(-) diff --git a/build_tools/utils.py b/build_tools/utils.py index da4ecf957b..3b02fcea2b 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -176,43 +176,83 @@ def found_pybind11() -> bool: @functools.lru_cache(maxsize=None) -def cuda_home_path() -> Optional[Path]: - """Returns the CUDA home path. This path should contain binaries (e.g. nvcc), headers, and libraries. +def nvcc_path() -> Optional[Path]: + """Get the NVCC binary path. - Returns `None` if CUDA is not found.""" - if cuda_home := os.getenv("CUDA_HOME"): - return Path(cuda_home) + Returns `None` if NVCC is not found. + """ - # Check if site-packages contains an `nvidia` directory - for site_package in sys.path: - if not Path(site_package).is_dir(): - continue + def lookup_via_cuda_home() -> Optional[Path]: + if cuda_home := os.getenv("CUDA_HOME"): + return Path(cuda_home) / "bin" / "nvcc" + return None - nvidia_dir = Path(site_package) / "nvidia" - if not nvidia_dir.is_dir(): - continue + def lookup_via_python_path() -> Optional[Path]: + for python_path in sys.path: + nvidia_dir = Path(python_path) / "nvidia" + if not nvidia_dir.is_dir(): + continue - if Path(nvidia_dir / "bin").is_dir(): - return nvidia_dir + cuda_roots = [nvidia_dir] + cuda_version_dirs = [ + path + for path in nvidia_dir.iterdir() + if path.is_dir() and re.fullmatch(r"cu\d+", path.name) + ] + cuda_version_dirs.sort(key=lambda path: int(path.name[2:]), reverse=True) + cuda_roots.extend(cuda_version_dirs) - # In this case there must be a "CUDA version directory" like `cu12` or `cu13` - # that contains the binaries, headers, and libraries - for cuda_version_dir in nvidia_dir.iterdir(): - if not cuda_version_dir.is_dir(): - continue + for cuda_root in cuda_roots: + nvcc_bin = cuda_root / "bin" / "nvcc" + if nvcc_bin.is_file(): + return nvcc_bin - # Verify that the directory name matches the expected `cu##` format - if not re.match(r"cu\d+", cuda_version_dir.name): - continue + return None - if not (cuda_version_dir / "bin").is_dir(): - continue + def lookup_via_distribution() -> Optional[Path]: + try: + cuda_nvcc_distribution = distribution("nvidia-cuda-nvcc") + except PackageNotFoundError: + return None - return cuda_version_dir + for package_path in cuda_nvcc_distribution.files or []: + package_path = Path(package_path) + if package_path.name == "nvcc" and package_path.parent.name == "bin": + return Path(cuda_nvcc_distribution.locate_file(package_path)) + + return None + + def lookup_via_path() -> Optional[Path]: + if (nvcc_bin := shutil.which("nvcc")) is not None: + return Path(nvcc_bin) + return None + + def lookup_via_local_cuda() -> Path: + return Path("/usr/local/cuda/bin/nvcc") + + nvcc_lookup_funcs: List[Callable[[], Optional[Path]]] = [ + lookup_via_cuda_home, + lookup_via_python_path, + lookup_via_distribution, + lookup_via_path, + lookup_via_local_cuda, + ] + + for nvcc_lookup_func in nvcc_lookup_funcs: + if (nvcc_bin := nvcc_lookup_func()) is not None and nvcc_bin.is_file(): + return nvcc_bin.resolve() return None +@functools.lru_cache(maxsize=None) +def cuda_home_path() -> Optional[Path]: + """Return the CUDA Toolkit root containing NVCC.""" + if (nvcc_bin := nvcc_path()) is not None: + return nvcc_bin.parent.parent + return None + + @functools.lru_cache(maxsize=None) def nccl_root_path() -> Optional[Path]: """Return the NCCL installation root. @@ -260,70 +300,12 @@ def nccl_lib_path() -> Optional[Path]: @functools.lru_cache(maxsize=None) -def cuda_toolkit_include_path() -> Tuple[str, str]: - """Returns root path for cuda toolkit includes. - - return `None` if CUDA is not found.""" - # Try finding CUDA - cuda_home: Optional[Path] = None - if cuda_home is None and os.getenv("CUDA_HOME"): - # Check in CUDA_HOME - cuda_home = Path(os.getenv("CUDA_HOME")) / "include" - if cuda_home is None: - # Check in NVCC - nvcc_bin = shutil.which("nvcc") - if nvcc_bin is not None: - cuda_home = Path(nvcc_bin.rstrip("/bin/nvcc")) / "include" - if cuda_home is None: - # Last-ditch guess in /usr/local/cuda - if Path("/usr/local/cuda").is_dir(): - cuda_home = Path("/usr/local/cuda") / "include" - return cuda_home - - -@functools.lru_cache(maxsize=None) -def nvcc_path() -> Optional[Path]: - """Get the NVCC binary path. - - Returns `None` if NVCC is not found. - """ - - def lookup_via_cuda_home() -> Optional[str]: - if (cuda_home := cuda_home_path()) is not None: - return cuda_home / "bin" / "nvcc" - return None - - def lookup_via_path() -> Optional[str]: - if (nvcc_bin := shutil.which("nvcc")) is not None: - return nvcc_bin - return None - - def lookup_via_distribution() -> Optional[str]: - try: - return distribution("nvidia-cuda-nvcc").locate_file("bin/nvcc") - except PackageNotFoundError: - return None - - def lookup_via_local_cuda() -> Optional[str]: - return "/usr/local/cuda/bin/nvcc" - - nvcc_lookup_funcs: List[Callable[[], Optional[str]]] = [ - lookup_via_cuda_home, - lookup_via_path, - lookup_via_distribution, - lookup_via_local_cuda, - ] - - for nvcc_lookup_func in nvcc_lookup_funcs: - nvcc_bin = nvcc_lookup_func() - - if nvcc_bin is None: - continue - - nvcc_bin_path = Path(nvcc_bin) - if nvcc_bin_path.is_file(): - return nvcc_bin_path - +def cuda_toolkit_include_path() -> Optional[Path]: + """Return the CUDA Toolkit include directory.""" + if (cuda_home := cuda_home_path()) is not None: + include_path = cuda_home / "include" + if include_path.is_dir(): + return include_path return None diff --git a/setup.py b/setup.py index 2a8a9d7688..5187179498 100644 --- a/setup.py +++ b/setup.py @@ -19,6 +19,7 @@ from build_tools.te_version import te_version from build_tools.utils import ( cuda_archs, + cuda_home_path, cuda_version, cudnn_frontend_include_path, get_frameworks, @@ -26,6 +27,7 @@ min_python_version_str, nccl_ep_enabled, get_max_jobs_for_parallel_build, + nvcc_path, ) frameworks = get_frameworks() @@ -253,6 +255,10 @@ def build_nccl_ep_submodule() -> str: nproc = get_max_jobs_for_parallel_build() env = os.environ.copy() + if (cuda_home := cuda_home_path()) is not None: + env.setdefault("CUDA_HOME", str(cuda_home)) + if (nvcc_bin := nvcc_path()) is not None: + env.setdefault("NVCC", str(nvcc_bin)) env["NVCC_GENCODE"] = gencode # NCCL EP needs the core NCCL headers + libnccl.so; write NCCL EP build # outputs to the submodule's local build/ tree.