Support building with headers from nvidia wheels

[vmarkovtsev]

There are two changes:

1. import nvidia returns a namespace package with __file__ equal to None
2. Add the way to force headers from nvidia wheels. Without that envvar, it's practically impossible with CUDA installed system-wide.

I successfully built the package with torch using the following uv configuration:

[tool.uv.extra-build-dependencies]
"transformer-engine-torch" = [
    "ninja",
    "nvidia-cuda-crt==13.0.88",
    "nvidia-cuda-cccl==13.0.85",
    { requirement = "torch", match-runtime = true },
    { requirement = "pytorch-triton", match-runtime = true },
    { requirement = "nvidia-cusolver", match-runtime = true },
    { requirement = "nvidia-curand", match-runtime = true },
    { requirement = "nvidia-cublas", match-runtime = true },
    { requirement = "nvidia-cusparse", match-runtime = true },
    { requirement = "nvidia-cudnn-cu13", match-runtime = true },
    { requirement = "nvidia-nvtx", match-runtime = true },
    { requirement = "nvidia-cuda-nvrtc", match-runtime = true },
    { requirement = "nvidia-cuda-runtime", match-runtime = true },
]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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

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

[greptile-apps]

Greptile Overview

Greptile Summary

Enhanced CUDA header discovery to support building with NVIDIA wheels in environments where CUDA toolkit is installed system-wide. The PR introduces two key improvements:

• Namespace package handling: Fixed handling of nvidia package when it's a namespace package (where __file__ is None), falling back to __path__[0] to locate headers
• Force wheels option: Added NVTE_BUILD_USE_NVIDIA_WHEELS environment variable to explicitly prefer headers from pip wheels over system-wide CUDA installation

The changes enable building with specific CUDA component versions from wheels (like nvidia-cuda-crt, nvidia-cuda-cccl, etc.) rather than relying on system CUDA, which is particularly useful for reproducible builds and version pinning.

Confidence Score: 4/5

• Safe to merge with minor documentation consideration
• The implementation correctly handles namespace packages and adds proper default handling for the environment variable. The fix follows existing patterns in the codebase (similar bool(int(os.getenv())) pattern used elsewhere). The only minor consideration is lack of documentation for the new environment variable.
• No files require special attention

Important Files Changed

Filename	Overview
build_tools/utils.py	Enhanced CUDA header discovery to handle namespace packages and added env var to force nvidia wheels usage

Sequence Diagram

sequenceDiagram
    participant Build as Build System
    participant Utils as get_cuda_include_dirs()
    participant Env as Environment
    participant FS as File System
    participant Nvidia as nvidia package

    Build->>Utils: Request CUDA include dirs
    Utils->>Env: Check NVTE_BUILD_USE_NVIDIA_WHEELS
    
    alt force_wheels=True OR toolkit not found
        Utils->>Nvidia: import nvidia
        
        alt nvidia.__file__ exists
            Nvidia-->>Utils: Return __file__ path
            Utils->>FS: cuda_root = Path(__file__).parent
        else nvidia.__file__ is None (namespace pkg)
            Nvidia-->>Utils: Return None
            Utils->>FS: cuda_root = Path(__path__[0])
        end
        
        Utils->>FS: Find subdirs with include/
        FS-->>Utils: Return list of include dirs
        Utils-->>Build: Return wheel-based includes
    else force_wheels=False AND toolkit found
        Utils->>FS: cuda_toolkit_include_path()
        FS-->>Utils: Return toolkit include path
        Utils-->>Build: Return toolkit includes
    end

Loading

[greptile-apps]

_{No files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[cyanguwa]

/te-ci L0

[vmarkovtsev]

It looks like the red jobs are unrelated timeouts.

[ksivaman]

/te-ci L0

[ksivaman]

LGTM

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Consider documenting the new NVTE_BUILD_USE_NVIDIA_WHEELS environment variable in docs/envvars.rst under the "CUDA Configuration" section for consistency with other build-time environment variables.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}