Add CUDA wheels as build-time requirements for build isolation by fheinecke · Pull Request #3253 · NVIDIA/TransformerEngine

fheinecke · 2026-07-24T17:57:53Z

Description

One of the issues blocking wheel build isolation is that CUDA build-time dependencies (e.g. nvcc) are not available today. These cannot easily be statically defined in pyproject.toml, because the exact dependencies required depend on whether the wheel is being built for CUDA 12 or CUDA 13.

This PR adds a PEP 517/660 compliant build_system.py that defines these requirements based on specific build-time flag, env var, or system dependencies. As with today, this logic is completely ignored when --no-build-isolation is set, so it should have no impact on existing release processes. This does not completely enable build isolation, but fixes most of the remaining blockers.

I have intentionally listed this as "not for external use" in the documentation. I'll update this and the install docs once build isolation is completely supported.

Depends on #3251

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Enables installing CUDA version-specific, build-time dependencies via wheels

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

fheinecke · 2026-07-24T17:58:06Z

/te-ci

ksivaman · 2026-08-27T01:21:15Z

/te-ci

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

Match the nvidia-cudnn-frontend version that main requires. Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

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 <fheinecke@nvidia.com>

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 <fheinecke@nvidia.com>

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<name> 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 <fheinecke@nvidia.com>

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 <fheinecke@nvidia.com>

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 <fheinecke@nvidia.com>

… 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 <fheinecke@nvidia.com>

…d 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 <fheinecke@nvidia.com>

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

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 <fheinecke@nvidia.com>

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 <fheinecke@nvidia.com>

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 <fheinecke@nvidia.com>

…ge 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 <fheinecke@nvidia.com>

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 <fheinecke@nvidia.com>

This reverts commit aaaf3d5. Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

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 <fheinecke@nvidia.com>

…es 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 <fheinecke@nvidia.com>

fheinecke force-pushed the fred/add-build-time-requirements-1 branch from e01577c to d8270b8 Compare September 4, 2026 16:42

fheinecke mentioned this pull request Sep 14, 2026

[Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI #3137

Merged

10 of 13 tasks

fheinecke and others added 12 commits November 6, 2026 14:15

Add CUDA wheels as build-time requirements for build isolation

1a6efa1

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

573e5eb

for more information, see https://pre-commit.ci

Support framework-specific isolated builds

e8e4d2a

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

Exclude cmake 3.25.0 due to but with finding cuda toolkit

ddb3178

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

fixed cuda library resolution when using wheels as shared lib source

ecfe945

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

Standardize CUDA lookup order based on feedback from NVIDIA#3251

1501892

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

Fix CUDA include dir resolution

aca8e20

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

Require cuDNN frontend 1.29 for isolated builds

da9f4af

Match the nvidia-cudnn-frontend version that main requires. Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

fheinecke force-pushed the fred/add-build-time-requirements-1 branch from c60111f to 4a5b76a Compare November 6, 2026 21:43

fheinecke added 2 commits November 9, 2026 16:07

fheinecke force-pushed the fred/add-build-time-requirements-1 branch from 190f8c2 to 152b4fb Compare November 9, 2026 22:19

fheinecke added 10 commits November 9, 2026 16:44

Remove 'not intended for external use' from build isolation docs

31ee769

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

Revert "Define the minimum cuDNN version once for the build and runtime"

929f70a

This reverts commit aaaf3d5. Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

Navigation Menu

Repository navigation

Add CUDA wheels as build-time requirements for build isolation - #3253

Add CUDA wheels as build-time requirements for build isolation#3253
fheinecke wants to merge 24 commits into
NVIDIA:mainfrom
fheinecke:fred/add-build-time-requirements-1

fheinecke commented Jul 24, 2026

Uh oh!

fheinecke commented Jul 24, 2026

Uh oh!

ksivaman commented Aug 27, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

fheinecke commented Jul 24, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

fheinecke commented Jul 24, 2026

Uh oh!

ksivaman commented Aug 27, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants