From 2ea5b3afa41ff280609ce49457ae07fa35eac2b4 Mon Sep 17 00:00:00 2001 From: Yuan Tong <13075180+tongyuantongyu@users.noreply.github.com> Date: Thu, 7 Aug 2025 17:30:25 +0800 Subject: [PATCH 1/2] [None][fix] Migrate to new cuda binding package name Signed-off-by: Yuan Tong <13075180+tongyuantongyu@users.noreply.github.com> --- tensorrt_llm/_ipc_utils.py | 13 ++++++++----- tensorrt_llm/_mnnvl_utils.py | 6 +++++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 6 +++++- tensorrt_llm/auto_parallel/cluster_info.py | 6 +++++- tensorrt_llm/runtime/generation.py | 5 ++++- tensorrt_llm/runtime/multimodal_model_runner.py | 7 ++++++- tests/microbenchmarks/all_reduce.py | 5 ++++- 7 files changed, 37 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_ipc_utils.py b/tensorrt_llm/_ipc_utils.py index c19fa516babd..d0e30a61d6c5 100644 --- a/tensorrt_llm/_ipc_utils.py +++ b/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: diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index 50610c304089..39f6deac4c54 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index d87dbef4e7d6..03bb6c9a7fb4 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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 diff --git a/tensorrt_llm/auto_parallel/cluster_info.py b/tensorrt_llm/auto_parallel/cluster_info.py index 514504dbe81d..8bdf5eba29e3 100644 --- a/tensorrt_llm/auto_parallel/cluster_info.py +++ b/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 diff --git a/tensorrt_llm/runtime/generation.py b/tensorrt_llm/runtime/generation.py index ec223f31c1a3..bf6e228f7699 100755 --- a/tensorrt_llm/runtime/generation.py +++ b/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 diff --git a/tensorrt_llm/runtime/multimodal_model_runner.py b/tensorrt_llm/runtime/multimodal_model_runner.py index bb3a5480fcb6..a3105bb94b97 100644 --- a/tensorrt_llm/runtime/multimodal_model_runner.py +++ b/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 diff --git a/tests/microbenchmarks/all_reduce.py b/tests/microbenchmarks/all_reduce.py index 8fdcc0a827c1..ca74165a3db8 100644 --- a/tests/microbenchmarks/all_reduce.py +++ b/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 From f328e64b0e9c96dd1af055e24441a9ba81d07cf6 Mon Sep 17 00:00:00 2001 From: Yuan Tong <13075180+tongyuantongyu@users.noreply.github.com> Date: Thu, 7 Aug 2025 17:30:25 +0800 Subject: [PATCH 2/2] more occurrences Signed-off-by: Yuan Tong <13075180+tongyuantongyu@users.noreply.github.com> --- cpp/kernels/fmha_v2/fmha_test.py | 7 ++++++- tests/integration/defs/sysinfo/get_sysinfo.py | 6 +++++- tests/microbenchmarks/build_time_benchmark.py | 6 +++++- .../_torch/multi_gpu/test_lowprecision_allreduce.py | 5 ++++- tests/unittest/trt/functional/test_allreduce_norm.py | 5 ++++- .../functional/test_allreduce_prepost_residual_norm.py | 5 ++++- tests/unittest/trt/functional/test_nccl.py | 5 ++++- tests/unittest/trt/functional/test_pp_reduce_scatter.py | 5 ++++- tests/unittest/utils/util.py | 8 +++++++- 9 files changed, 43 insertions(+), 9 deletions(-) diff --git a/cpp/kernels/fmha_v2/fmha_test.py b/cpp/kernels/fmha_v2/fmha_test.py index bd255bdd8563..d02e3cc31c0f 100644 --- a/cpp/kernels/fmha_v2/fmha_test.py +++ b/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): diff --git a/tests/integration/defs/sysinfo/get_sysinfo.py b/tests/integration/defs/sysinfo/get_sysinfo.py index a7de1871e517..4efd791d7420 100644 --- a/tests/integration/defs/sysinfo/get_sysinfo.py +++ b/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__) diff --git a/tests/microbenchmarks/build_time_benchmark.py b/tests/microbenchmarks/build_time_benchmark.py index 133be635335f..8cea08205c6b 100644 --- a/tests/microbenchmarks/build_time_benchmark.py +++ b/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, diff --git a/tests/unittest/_torch/multi_gpu/test_lowprecision_allreduce.py b/tests/unittest/_torch/multi_gpu/test_lowprecision_allreduce.py index be9c10ebdfe1..331e250b349c 100644 --- a/tests/unittest/_torch/multi_gpu/test_lowprecision_allreduce.py +++ b/tests/unittest/_torch/multi_gpu/test_lowprecision_allreduce.py @@ -36,7 +36,10 @@ def run_single_rank(dtype, strategy, message_size): import numpy as np import torch - from cuda import cuda + try: + from cuda.bindings import driver as cuda + except ImportError: + from cuda import cuda import tensorrt_llm from tensorrt_llm._torch.distributed import AllReduce, AllReduceStrategy diff --git a/tests/unittest/trt/functional/test_allreduce_norm.py b/tests/unittest/trt/functional/test_allreduce_norm.py index a6a8ee477e11..dbd4448c45d1 100644 --- a/tests/unittest/trt/functional/test_allreduce_norm.py +++ b/tests/unittest/trt/functional/test_allreduce_norm.py @@ -21,7 +21,10 @@ import torch # isort: on -from cuda import cudart +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart from parameterized import parameterized from utils.util import create_session, run_session, unittest_name_func diff --git a/tests/unittest/trt/functional/test_allreduce_prepost_residual_norm.py b/tests/unittest/trt/functional/test_allreduce_prepost_residual_norm.py index 295e92c0aec9..bf0592193d59 100644 --- a/tests/unittest/trt/functional/test_allreduce_prepost_residual_norm.py +++ b/tests/unittest/trt/functional/test_allreduce_prepost_residual_norm.py @@ -21,7 +21,10 @@ import torch # isort: on -from cuda import cudart +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart from parameterized import parameterized from utils.util import create_session, run_session, unittest_name_func diff --git a/tests/unittest/trt/functional/test_nccl.py b/tests/unittest/trt/functional/test_nccl.py index 760e6538123f..61b9eacc825d 100644 --- a/tests/unittest/trt/functional/test_nccl.py +++ b/tests/unittest/trt/functional/test_nccl.py @@ -21,7 +21,10 @@ import torch # isort: on -from cuda import cudart +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart from parameterized import parameterized from utils.util import create_session, run_session, unittest_name_func diff --git a/tests/unittest/trt/functional/test_pp_reduce_scatter.py b/tests/unittest/trt/functional/test_pp_reduce_scatter.py index fba41cbf76fc..660a499c9226 100644 --- a/tests/unittest/trt/functional/test_pp_reduce_scatter.py +++ b/tests/unittest/trt/functional/test_pp_reduce_scatter.py @@ -21,7 +21,10 @@ import torch # isort: on -from cuda import cudart +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart from parameterized import parameterized from utils.util import create_session, run_session, unittest_name_func diff --git a/tests/unittest/utils/util.py b/tests/unittest/utils/util.py index ba75e57e7cf0..893af5d93b8d 100644 --- a/tests/unittest/utils/util.py +++ b/tests/unittest/utils/util.py @@ -9,7 +9,13 @@ import pytest import tensorrt as trt import torch -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 + from parameterized import parameterized import tensorrt_llm