• Vadim Markovtsev's avatar
    Support building with headers from nvidia wheels (#2623) · 94ba75d7
    Vadim Markovtsev authored
    
    
    * Support building with headers from nvidia wheels
    
    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 },
    ]
    ```
    Signed-off-by: default avatarVadim Markovtsev <vadim@poolside.ai>
    
    * Apply suggestion from @ksivaman
    Signed-off-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>
    Signed-off-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>
    
    ---------
    Signed-off-by: default avatarVadim Markovtsev <vadim@poolside.ai>
    Signed-off-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>
    Co-authored-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>
    94ba75d7
utils.py 10.5 KB