2、环境搭建:CUDA、cuDNN、TensorRT、TensorRT-LLM的安装与配置,验证开发环境

2.1 核心原理:版本匹配是环境搭建的第一要务

TensorRT-LLM 并非一个完全独立的软件包,它深度依赖底层 CUDA、cuDNN 和 TensorRT 的运行时库。其编译过程需要调用 TensorRT 的 API 来构建引擎,推理时则通过 CUDA 驱动管理 GPU 显存与内核执行。因此,版本矩阵的严格匹配是环境搭建的核心矛盾。若 CUDA 版本过新而 TensorRT 未适配,或 cuDNN 的 major 版本与 TensorRT 预期不符,都会在编译阶段抛出“找不到 cudnn_version.h”或运行时出现“symbol lookup error”等难以排查的问题。

下图展示了 TensorRT-LLM 开发环境的分层依赖关系,以及各组件在编译与运行阶段的作用边界:

TensorRT-LLM 开发环境分层依赖关系 自定义插件 / 模型推理应用 编译期:调用 TensorRT API 构建引擎 TensorRT-LLM(源码编译) 依赖 TensorRT 头文件、cuDNN 头文件、CUDA 工具链 通过 CMake 检测各组件版本,生成编译配置 TensorRT 8.6 / 9.x 提供 nvinfer.h 头文件 libnvinfer.so 运行时库 cuDNN 8.9 / 9.x 提供 cudnn.h 头文件 libcudnn.so 深度学习原语库 CUDA 11.8 / 12.x nvcc 编译器、cuda_runtime.h libcudart.so 运行时 NVIDIA GPU 驱动(Driver) 驱动版本必须 ≥ CUDA 最低要求,否则 CUDA 运行时无法初始化 验证命令 nvcc --version dpkg -l | grep cudnn dpkg -l | grep TensorRT python -c "import tensorrt"

从图中可以看出,TensorRT-LLM 位于中间层,它同时向上承接应用、向下调用三个基础库。任何一层的版本缺失或错位,都会导致整个链路编译失败或运行崩溃。因此,在动手安装前,务必先查阅 TensorRT-LLM 官方 GitHub 的版本对应表,确认你选择的 CUDA 版本与 TensorRT 版本是否在支持列表中。

2.2 方法流程:从零搭建可编译的插件开发环境

本节以 Ubuntu 22.04 + CUDA 12.1 + cuDNN 8.9 + TensorRT 9.2 + TensorRT-LLM 0.9.0 为例,展示一套经过验证的安装流程。该组合在社区中较为成熟,适合作为插件开发的基础环境。

2.2.1 安装 NVIDIA 驱动与 CUDA

首先确认 GPU 型号并安装匹配的驱动。推荐使用 nvidia-smi 查看当前驱动版本,若驱动过旧,先通过 sudo apt install nvidia-driver-535 安装较新驱动并重启。随后安装 CUDA 12.1:

# 下载 CUDA 12.1 runfile(约 3GB)
wget https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run
sudo sh cuda_12.1.0_530.30.02_linux.run --toolkit --silent --override

# 配置环境变量
echo 'export PATH=/usr/local/cuda-12.1/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

# 验证
nvcc --version

注意:--toolkit 参数只安装 CUDA 工具包而不覆盖驱动,避免与已有驱动冲突。安装完成后,nvcc --version 应显示 Cuda compilation tools, release 12.1。

2.2.2 安装 cuDNN 与 TensorRT

cuDNN 和 TensorRT 均通过 NVIDIA 的 apt 仓库安装最为便捷。首先配置 NVIDIA 官方仓库:

# 添加 NVIDIA 仓库密钥与源
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update

# 安装 cuDNN 8.9(与 CUDA 12.1 匹配)
sudo apt install libcudnn8=8.9.7.29-1+cuda12.1
sudo apt install libcudnn8-dev=8.9.7.29-1+cuda12.1

# 安装 TensorRT 9.2
sudo apt install tensorrt=9.2.0.12-1+cuda12.1
sudo apt install python3-libnvinfer-dev

安装完成后,验证 TensorRT 的 Python 接口是否可用:

python3 -c "import tensorrt as trt; print(trt.__version__)"
# 期望输出:9.2.0.12

2.2.3 源码编译 TensorRT-LLM

TensorRT-LLM 需要从源码编译,因为官方暂未提供预编译的 pip 包。克隆仓库并执行编译:

git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git checkout v0.9.0
git submodule update --init --recursive

# 创建虚拟环境
python3 -m venv trt-llm-env
source trt-llm-env/bin/activate
pip install --upgrade pip

# 编译安装(耗时约 20-40 分钟)
python3 scripts/build_wheel.py --cuda_architectures "80;86;89" --trt_root /usr/local/tensorrt
pip install ./build/tensorrt_llm-0.9.0-*.whl

其中 --cuda_architectures 参数需根据你的 GPU 计算能力填写:A100 为 80,RTX 3090 为 86,RTX 4090 为 89。编译完成后,通过以下命令验证安装:

python3 -c "import tensorrt_llm; print(tensorrt_llm.__version__)"
# 期望输出:0.9.0

2.3 案例:验证自定义插件开发的最小环境

环境搭建是否成功,不能只看 import 是否报错,而应通过一个最小插件编译测试来验证。下面创建一个简单的 ReLU 插件,确认 TensorRT-LLM 的插件注册机制能正常工作。

首先创建插件源文件 relu_plugin.cpp:

#include <tensorrt_llm/plugins/common/plugin.h>
#include <NvInfer.h>

using namespace nvinfer1;

class ReluPlugin : public IPluginV2DynamicExt {
public:
    ReluPlugin() = default;
    // ... 省略接口实现 ...
};

// 注册插件
REGISTER_TENSORRT_PLUGIN(ReluPluginCreator);

然后编写 CMakeLists.txt 并编译:

cmake_minimum_required(VERSION 3.18)
project(relu_plugin)

find_package(TensorRT-LLM REQUIRED)
add_library(relu_plugin SHARED relu_plugin.cpp)
target_link_libraries(relu_plugin PRIVATE tensorrt_llm::plugin)

编译成功后,在 Python 中加载并测试:

import tensorrt as trt
import tensorrt_llm

# 加载插件动态库
trt.init_libnvinfer_plugins(trt.Logger(), "")
plugin_registry = trt.get_plugin_registry()
creator = plugin_registry.get_plugin_creator("Relu", "1", "")
assert creator is not None, "插件注册失败!"
print("自定义插件注册成功,环境验证通过。")

若上述代码能正常输出“插件注册成功”,则说明 CUDA、cuDNN、TensorRT、TensorRT-LLM 四层环境均已正确配置,且头文件路径、库链接、插件注册机制全部打通。

2.4 常见误区与复盘

误区 典型表现 正确做法
盲目安装最新 CUDA TensorRT-LLM 编译时报错找不到 cuda_runtime.h 或版本不匹配 严格对照官方版本表,选择 TensorRT-LLM 明确支持的 CUDA 版本
忽略 cuDNN 的 dev 包 编译时提示 cudnn.h: No such file 必须同时安装 libcudnn8-dev,仅装 runtime 包不够
使用系统 Python 直接编译 pip 安装时权限错误或污染系统环境 始终使用 venv 或 conda 创建独立虚拟环境
忘记设置 LD_LIBRARY_PATH import tensorrt_llm 时报 libnvinfer.so 找不到 将 TensorRT 的 lib 目录加入动态库搜索路径
GPU 架构参数写错 编译出的引擎无法在目标 GPU 上运行 用 nvidia-smi --query-gpu=compute_cap --format=csv 查询真实计算能力
特别提醒:不要使用 pip 直接安装 tensorrt 包,它与 apt 安装的 TensorRT 可能产生冲突。统一通过 NVIDIA apt 仓库管理 TensorRT 的版本,能避免大量头文件路径混乱的问题。

复盘整个安装过程,最关键的验证节点有三个:nvcc 版本输出、TensorRT Python 接口版本、TensorRT-LLM import 成功。只有这三个节点全部通过,才能进入后续的插件开发环节。建议将上述验证命令整理为一个 shell 脚本,每次切换环境或升级组件后重新执行,确保开发环境始终处于健康状态。