From 494da4ae899d177062a804f7eed29149d99ef013 Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Mon, 15 Jun 2026 16:15:22 -0400 Subject: [PATCH] [CUDA] Narrow the `cuda` extra from cuda-python to cuda-bindings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TVM's shipped code only uses cuda.bindings — cuda.bindings.nvrtc for the NVRTC JIT path and cuda.bindings.driver for the NVSHMEM link path, both in python/tvm/support/nvcc.py; it never uses cuda.core. cuda-python is now a metapackage that pulls in cuda-bindings + cuda-core (and cuda-pathfinder), so depending on it drags in cuda-core that TVM does not need. Depend directly on cuda-bindings, which provides exactly the nvrtc and driver submodules TVM imports, and update the user-facing 'pip install cuda-python' hints to match. A plain cuda-bindings install pulls no nvidia-* toolkit wheels (those live behind the [all] extra); libnvrtc is loaded from the system / TVM's CUDA install as before. --- pyproject.toml | 2 +- python/tvm/support/nvcc.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3c61fa389fc3..221b3c3383b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,7 @@ dependencies = [ [project.optional-dependencies] torch = ["torch"] -cuda = ["cuda-python"] +cuda = ["cuda-bindings"] meta-schedule = ["xgboost"] popen-pool = ["psutil", "cloudpickle"] rpc = ["tornado", "psutil", "cloudpickle"] diff --git a/python/tvm/support/nvcc.py b/python/tvm/support/nvcc.py index 859dbd3077c7..ea5939fceffc 100644 --- a/python/tvm/support/nvcc.py +++ b/python/tvm/support/nvcc.py @@ -65,7 +65,7 @@ def compile_cuda( Notes ----- - NVRTC is a "runtime" compilation library and can be faster for JIT compilation. - - NVRTC requires cuda-python: pip install cuda-python + - NVRTC requires cuda-bindings: pip install cuda-bindings """ use_nvshmem = "#include " in code or "#include " in code @@ -289,9 +289,9 @@ def _compile_cuda_nvrtc( from cuda.bindings import nvrtc # pylint: disable=import-outside-toplevel except ImportError as e: raise RuntimeError( - "Failed to compile CUDA with NVRTC because the `cuda-python` package " + "Failed to compile CUDA with NVRTC because the `cuda-bindings` package " "is not available.\n" - "Please install it with: pip install cuda-python\n" + "Please install it with: pip install cuda-bindings\n" "See: https://nvidia.github.io/cuda-python/" ) from e @@ -301,9 +301,9 @@ def _compile_cuda_nvrtc( if importlib.util.find_spec("cuda.bindings.driver") is None: raise RuntimeError( - "Failed to compile CUDA with NVRTC+NVSHMEM because the `cuda-python` package " + "Failed to compile CUDA with NVRTC+NVSHMEM because the `cuda-bindings` package " "is not available.\n" - "Please install it with: pip install cuda-python\n" + "Please install it with: pip install cuda-bindings\n" "See: https://nvidia.github.io/cuda-python/" ) @@ -812,7 +812,7 @@ def tvm_callback_cuda_compile(code): TVM_CUDA_COMPILE_MODE : str Compiler backend: "nvcc" (default) or "nvrtc" - "nvcc": Use nvcc subprocess, generates fatbin - - "nvrtc": Use NVRTC via cuda-python for faster JIT, generates cubin + - "nvrtc": Use NVRTC via cuda-bindings for faster JIT, generates cubin TVM_KERNEL_DUMP : str If set, dump generated CUDA/intermediate files and append "-lineinfo" so profilers can correlate SASS back to the dumped source.