[CI] Improve build time dependency resolution

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,33 @@ 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.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() and (cuda_full_version := cuda_version()):
+                cuda_major_version = cuda_full_version[0]
+                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:
+            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 +214,11 @@ 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"]

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,
     )

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.
@@ -176,52 +176,137 @@ def found_pybind11() -> bool:
 
 
 @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
+def nvcc_path() -> Optional[Path]:
+    """Get the NVCC binary path.
+
+    Returns `None` if NVCC is not found.
+    """
+
+    def lookup_via_cuda_home() -> Optional[Path]:
+        if cuda_home := os.getenv("CUDA_HOME"):
+            return Path(cuda_home) / "bin" / "nvcc"
+        return None
+
+    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
+
+            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)
+
+            for cuda_root in cuda_roots:
+                nvcc_bin = cuda_root / "bin" / "nvcc"
+                if nvcc_bin.is_file():
+                    return nvcc_bin
+
+        return None
+
+    def lookup_via_distribution() -> Optional[Path]:
+        try:
+            cuda_nvcc_distribution = distribution("nvidia-cuda-nvcc")
+        except PackageNotFoundError:
+            return None
+
+        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 nvcc_path() -> Tuple[str, str]:
-    """Returns 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}")
-
-    return nvcc_bin
+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.
+
+    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() -> 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
 
 
 @functools.lru_cache(maxsize=None)
@@ -324,13 +409,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 +420,20 @@ 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]:

setup.py

@@ -19,13 +19,15 @@
 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,
     remove_dups,
     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.