[None][fix] Migrate to new cuda binding package name (#6700)

Signed-off-by: Yuan Tong <[email protected]>

Author: tongyuantongyu

cpp/kernels/fmha_v2/fmha_test.py

@@ -1,7 +1,12 @@
 import subprocess
 
 import pytest
-from cuda import cuda, nvrtc
+
+try:
+    from cuda.bindings import driver as cuda
+    from cuda.bindings import nvrtc
+except ImportError:
+    from cuda import cuda, nvrtc
 
 
 def ASSERT_DRV(err):

tensorrt_llm/_ipc_utils.py

@@ -17,17 +17,20 @@
 import sys
 from typing import List, Tuple
 
-from cuda import cuda, cudart
-from cuda.cudart import cudaError_t
+try:
+    from cuda.bindings import driver as cuda
+    from cuda.bindings import runtime as cudart
+except ImportError:
+    from cuda import cuda, cudart
 
 from ._utils import mpi_comm
 from .logger import logger
 from .mapping import Mapping
 
 
-def _raise_if_error(error: cudaError_t | cuda.CUresult):
-    if isinstance(error, cudaError_t):
-        if error != cudaError_t.cudaSuccess:
+def _raise_if_error(error: cudart.cudaError_t | cuda.CUresult):
+    if isinstance(error, cudart.cudaError_t):
+        if error != cudart.cudaError_t.cudaSuccess:
             raise RuntimeError(f"CUDA Runtime API error: {repr(error)}")
     if isinstance(error, cuda.CUresult):
         if error != cuda.CUresult.CUDA_SUCCESS:

tensorrt_llm/_mnnvl_utils.py

@@ -21,7 +21,11 @@
 
 import pynvml
 import torch
-from cuda import cuda
+
+try:
+    from cuda.bindings import driver as cuda
+except ImportError:
+    from cuda import cuda
 
 from ._dlpack_utils import pack_strided_memory
 from ._utils import mpi_comm

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -11,7 +11,11 @@
 from typing import Dict, Iterable, List, Optional, Tuple, Union
 
 import torch
-from cuda import cudart
+
+try:
+    from cuda.bindings import runtime as cudart
+except ImportError:
+    from cuda import cudart
 
 from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType
 from tensorrt_llm._torch.pyexecutor.seq_slot_manager import SeqSlotManager

tensorrt_llm/auto_parallel/cluster_info.py

@@ -5,7 +5,11 @@
 
 import pynvml
 import torch
-from cuda import cudart
+
+try:
+    from cuda.bindings import runtime as cudart
+except ImportError:
+    from cuda import cudart
 
 from tensorrt_llm._utils import DictConversion
 from tensorrt_llm.logger import logger

tensorrt_llm/runtime/generation.py

@@ -29,7 +29,10 @@
 import torch
 import tensorrt as trt
 # isort: on
-from cuda import cudart
+try:
+    from cuda.bindings import runtime as cudart
+except ImportError:
+    from cuda import cudart
 
 from tensorrt_llm.runtime.memory_pools.memory_pools_allocator import \
     MemoryPoolsAllocator

tensorrt_llm/runtime/multimodal_model_runner.py

@@ -13,7 +13,12 @@
 from typing import Optional, Tuple
 
 import torch.nn.functional as F
-from cuda import cudart
+
+try:
+    from cuda.bindings import runtime as cudart
+except ImportError:
+    from cuda import cudart
+
 from huggingface_hub import hf_hub_download
 from PIL import Image, UnidentifiedImageError
 from safetensors import safe_open

tests/integration/defs/sysinfo/get_sysinfo.py

@@ -24,7 +24,11 @@
 
 import psutil
 import pynvml
-from cuda import cuda
+
+try:
+    from cuda.bindings import driver as cuda
+except ImportError:
+    from cuda import cuda
 
 # Logger
 logger = logging.getLogger(__name__)

tests/microbenchmarks/all_reduce.py

@@ -18,7 +18,10 @@
 # isort: off
 import torch
 # isort: on
-from cuda import cudart
+try:
+    from cuda.bindings import runtime as cudart
+except ImportError:
+    from cuda import cudart
 
 import tensorrt_llm as tllm
 from tensorrt_llm import Mapping

tests/microbenchmarks/build_time_benchmark.py

@@ -7,7 +7,11 @@
 import traceback
 
 import tensorrt as trt
-from cuda import cudart
+
+try:
+    from cuda.bindings import runtime as cudart
+except ImportError:
+    from cuda import cudart
 
 import tensorrt_llm
 from tensorrt_llm import (AutoConfig, AutoModelForCausalLM, BuildConfig,