📚 课程目录第2章 TensorRT-LLM基础

2. TensorRT-LLM基础

2.1 TensorRT-LLM是什么

TensorRT-LLM是NVIDIA推出的开源大语言模型推理框架，它在TensorRT高性能推理引擎之上，为LLM的预填充（Prefill）与解码（Decode）两阶段做了深度定制优化。与直接使用TensorRT或vLLM不同，TensorRT-LLM的核心设计理念是：将LLM的复杂结构（Embedding、Attention、MLP、LayerNorm等）编译为高度融合的CUDA内核，并通过显存池、KV Cache管理、inflight batching等机制，最大化GPU在推理时的利用率。

从架构视角看，TensorRT-LLM并非一个简单的Python库，而是一个包含模型定义层、编译优化层、运行时执行层的完整技术栈。用户通过Python API定义模型结构并导出为TensorRT引擎文件，C++/Python运行时负责加载引擎并执行推理。这种“定义-编译-执行”的分离模式，使得引擎文件可以脱离原始模型权重独立部署，也为多GPU并行推理提供了基础。

  核心认知：TensorRT-LLM的“编译期”与“运行期”是解耦的。编译期将模型结构、权重、精度、GPU拓扑等信息固化到引擎文件中；运行期只负责加载引擎、管理KV Cache和调度请求。因此，多GPU推理的并行策略（如TP、PP）在编译期就已确定，运行期无法动态更改。

2.2 核心组件详解

TensorRT-LLM的功能由多个核心组件协同完成，理解这些组件是掌握多GPU推理架构的前提。下表总结了各组件的职责与关键接口：

组件	职责	关键类/接口	多GPU相关
Model Definition	定义模型结构（Transformer Block、Attention等）	`ModelConfig`, `TransformerLayer`	指定TP/PP大小
Weight Converter	将HuggingFace等权重映射为TensorRT-LLM格式	`HFCheckpointLoader`	按TP切分权重
Engine Builder	编译优化生成.engine文件	`build()`, `EngineConfig`	生成多GPU引擎
Runtime Session	加载引擎、执行推理、管理KV Cache	`LLM`, `GenerationSession`	NCCL通信初始化
KV Cache Manager	分页管理KV Cache显存	`KVCacheManager`	按GPU分配Cache
Executor	调度请求、管理inflight batching	`Executor`, `Request`	跨GPU请求分发

实操提示：在多GPU场景中，Weight Converter和Engine Builder是决定并行策略能否生效的关键。例如，当设置 tensor_parallel_size=2 时，Weight Converter会将每个权重矩阵按列/行切分为2份，Engine Builder则生成2个独立的引擎文件（每个GPU一个），运行时通过NCCL实现梯度/激活同步。

2.3 安装与配置流程

TensorRT-LLM的安装需要严格匹配NVIDIA驱动、CUDA、cuDNN和TensorRT版本。推荐使用NVIDIA官方Docker镜像，避免环境冲突。以下为标准的安装与配置流程：

环境准备：确认GPU型号（推荐A100/H100/A800等数据中心GPU），安装NVIDIA驱动（≥535.104.05），并拉取官方镜像：

# 拉取TensorRT-LLM官方镜像（以24.05版本为例）
docker pull nvcr.io/nvidia/tritonserver:24.05-trtllm-python-py3

# 启动容器并挂载模型目录
docker run --gpus all -it --shm-size=2g \
  -v /data/models:/models \
  nvcr.io/nvidia/tritonserver:24.05-trtllm-python-py3 /bin/bash

验证安装：在容器内检查关键库版本：

python -c "import tensorrt_llm; print(tensorrt_llm.__version__)"
# 预期输出：0.10.0 或类似版本号
python -c "import tensorrt; print(tensorrt.__version__)"
# 预期输出：10.0.1 或更高
nvidia-smi  # 确认GPU可见

配置多GPU环境：设置NCCL环境变量以优化多卡通信：

export NCCL_DEBUG=INFO          # 调试NCCL通信
export NCCL_P2P_LEVEL=LOCAL     # 启用P2P直连（同节点）
export NCCL_SHM_DISABLE=0       # 启用共享内存通信
export CUDA_VISIBLE_DEVICES=0,1 # 指定可见GPU

常见误区：很多初学者在安装时忽略CUDA版本匹配。TensorRT-LLM对CUDA版本有严格限制（通常要求CUDA 12.x），如果使用系统自带的CUDA 11.x，会导致编译失败或运行时崩溃。务必使用官方镜像或严格遵循版本兼容矩阵。

2.4 多GPU推理的配置案例

以一个7B模型（如LLaMA-2-7B）在两块A100（40GB）上进行张量并行推理为例，展示从权重转换到引擎构建的完整流程：

# 1. 权重转换（将HF权重转为TensorRT-LLM格式）
python convert_checkpoint.py \
  --model_dir /models/llama2-7b-hf \
  --output_dir /models/llama2-7b-trtllm \
  --dtype float16 \
  --tp_size 2 \
  --pp_size 1

# 2. 构建TensorRT引擎
trtllm-build \
  --checkpoint_dir /models/llama2-7b-trtllm \
  --output_dir /models/llama2-7b-engine \
  --gemm_plugin float16 \
  --max_batch_size 8 \
  --max_input_len 2048 \
  --max_seq_len 4096

# 3. 运行推理（指定TP=2）
python run.py \
  --engine_dir /models/llama2-7b-engine \
  --tokenizer_dir /models/llama2-7b-hf \
  --max_output_len 256 \
  --input_text "Explain the concept of tensor parallelism"

上述命令中，--tp_size 2 是关键参数，它告诉构建器将模型权重按张量并行方式切分到2个GPU上。生成的引擎目录中会包含 rank0.engine 和 rank1.engine 两个文件，分别对应GPU 0和GPU 1。

2.5 常见误区与复盘

在TensorRT-LLM的安装与配置阶段，以下问题出现频率最高，需要特别注意：

误区一：忽略GPU架构兼容性。TensorRT-LLM的某些特性（如FP8、Hopper架构专属优化）仅在特定GPU上可用。在V100上使用FP8会直接报错，需回退到FP16。
误区二：TP与PP混淆。TP（张量并行）适合单节点内多卡（NVLink互联），PP（流水线并行）适合跨节点或多卡场景。7B模型在2卡上通常用TP即可，盲目使用PP反而增加通信开销。
误区三：KV Cache配置不当。默认配置下KV Cache可能占用过多显存，导致OOM。建议通过 --kv_cache_free_gpu_memory_fraction 参数控制，例如设为0.3表示预留30%显存给KV Cache。
误区四：NCCL通信未验证。多GPU推理前应执行 nccl-tests 验证通信带宽，否则可能因网络配置问题导致推理速度极慢。

  复盘建议：在正式部署前，建议先使用单GPU构建引擎并验证推理正确性，再切换到多GPU模式。这样可以将“模型转换问题”与“并行通信问题”分离排查。如果多GPU推理结果与单GPU不一致，优先检查权重切分是否正确、NCCL是否报错，以及是否启用了确定性模式（--deterministic）。