NVIDIA_TensorRTLLM_0">NVIDIA TensorRT-LLM：高性能大语言模型推理框架详解

引言

随着大语言模型（LLMs）在各行各业的应用日益广泛，如何高效地部署这些模型以实现低延迟、高吞吐量的推理服务成为了一个关键挑战。NVIDIA TensorRT-LLM 是 NVIDIA 专为大语言模型推理优化的高性能框架，它提供了易用的 Python API 来定义大语言模型并构建包含最先进优化的 TensorRT 引擎，使模型能够在 NVIDIA GPU 上高效执行推理。

TensorRT-LLM 不仅提供了构建优化引擎的工具，还包含了用于创建 Python 和 C++ 运行时的组件，以执行这些 TensorRT 引擎。通过利用 NVIDIA 的硬件加速和软件优化，TensorRT-LLM 能够显著提高大语言模型的推理性能，降低延迟，提高吞吐量，从而使企业能够更经济高效地部署大语言模型服务。

本文将详细介绍 TensorRT-LLM 的安装部署步骤、核心功能、代码示例以及性能优化案例，帮助读者全面了解这一强大工具的使用方法和优势。

GTC 2025 中文在线解读｜ CUDA最新特性与未来 [WP72383]
NVIDIA GTC大会火热进行中，一波波重磅科技演讲让人应接不暇，3月24日，NVIDIA 企业开发者社区邀请Ken He、Yipeng Li两位技术专家，面向开发者，以中文深度拆解GTC2025四场重磅开发技术相关会议，直击AI行业应用痛点，破解前沿技术难题!

作为GPU计算领域的基石，CUDA通过其编程语言、编译器、运行时环境及核心库构建了完整的计算生态，驱动着人工智能、科学计算等前沿领域的创新发展。在本次在线解读活动中，将由CUDA架构师深度解析GPU计算生态的核心技术演进。带您了解今年CUDA平台即将推出的众多新功能，洞悉CUDA及GPU计算技术的未来发展方向。

时间：3月24日18:00-19:00
中文解读:Ken He / Developer community
链接：link: https://www.nvidia.cn/gtc-global/session-catalog/?tab.catalogallsessionstab=16566177511100015Kus&search=WP72383%3B%20WP72450%3B%20WP73739b%3B%20WP72784a%20#/session/1739861154177001cMJd=

1. TensorRT-LLM 概述

TensorRT-LLM 是一个为大语言模型推理优化的工具箱，它建立在 NVIDIA 的 TensorRT 基础上，专门针对 Transformer 架构的语言模型进行了深度优化。它的主要特点包括：

• 高性能推理：通过 CUDA 内核优化、混合精度计算和高效内存管理，实现低延迟、高吞吐量的推理
• 易用的 API：提供简洁的 Python API，简化模型定义和推理过程
• 灵活的部署选项：支持单 GPU 和多 GPU 分布式推理
• 广泛的模型支持：支持 Llama、Mistral、Falcon、GPT-J、MPT 等多种流行的大语言模型
• 量化支持：提供 INT8/INT4 量化功能，减少内存占用并提高推理速度
• 自定义优化：允许用户根据特定需求自定义优化策略

2. 安装与部署

TensorRT-LLM 提供了多种安装方式，包括通过 pip 安装预构建的包、使用 Docker 容器以及从源代码构建。以下将详细介绍这些安装方法。

2.1 通过 pip 安装（Ubuntu 24.04）

最简单的安装方式是通过 pip 安装预构建的 wheel 包：

# 安装依赖项
sudo apt-get -y install libopenmpi-dev && pip3 install --upgrade pip setuptools<77.0.1 && pip3 install nvidia-tensorrt-llm

安装完成后，可以通过以下代码验证安装是否成功：

# 导入 tensorrt_llm 模块验证安装
import tensorrt_llm
print(f"TensorRT-LLM version: {tensorrt_llm.__version__}")

2.2 从源代码构建（Linux）

对于需要最佳性能或自定义功能的用户，建议从源代码构建 TensorRT-LLM。

先决条件

1. 安装 Docker（推荐）

   • NVIDIA 平台 Docker 容器安装指南可在此处找到

2. 安装 git-lfs

# TensorRT-LLM 使用 git-lfs，需要提前安装
apt-get update && apt-get -y install git git-lfs
git lfs install

3. 克隆代码库

git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git submodule update --init --recursive
git lfs pull

使用 Docker 构建（推荐方式）

cd docker
# 构建基础开发镜像
./build.sh --build-docker-image# 或者构建包含所有依赖的镜像
./build.sh --build-docker-image --build-wheel# 启动开发容器
./run.sh# 或者指定 GPU
./run.sh --gpus all

手动构建（不使用 Docker）

# 安装基本依赖
apt-get update && apt-get install -y git git-lfs build-essential wget# 安装 Python 依赖
pip install -r requirements.txt# 配置构建
mkdir -p build && cd build
cmake -DCMAKE_BUILD_TYPE=Release ..# 编译
make -j$(nproc)# 安装 Python 包
cd ..
python -m pip install -e .

2.3 Grace Hopper 架构支持

对于使用 NVIDIA Grace Hopper 架构的用户，TensorRT-LLM 提供了专门的优化支持：

# 安装 PyTorch 2.6.0 和相关依赖
pip3 install torch==2.6.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121# 安装基础依赖
sudo apt-get -y install libopenmpi-dev && pip3 install --upgrade pip setuptools<77.0.1 && pip3 install nvidia-tensorrt-llm

注意：如果使用 PyTorch NGC 容器镜像，则不需要安装 CUDA 支持的 PyTorch 包。

2.4 已知限制

使用 pip 预构建的 TensorRT-LLM wheel 包存在一些已知限制：

• C++11 ABI：预构建的 wheel 包链接到了 pypi 上的公共 pytorch，这关闭了 C++11 ABI。而 NVIDIA 优化的 NGC 容器中的 pytorch 启用了 C++11 ABI。因此对于需要 C++11 ABI 支持的用户，建议从源代码构建。

3. 核心功能与 API

TensorRT-LLM 提供了丰富的功能和 API，使用户能够轻松定义、优化和部署大语言模型。以下是一些核心功能和 API 的介绍。

3.1 基本推理 API

TensorRT-LLM 的核心 API 是 LLM 类，它提供了简单直观的接口来加载模型和执行推理：

### 基本推理示例 - 使用 TensorRT-LLM 进行文本生成import tempfilefrom tensorrt_llm import LLM, SamplingParamsdef main():# 模型可以接受 HF 模型名称、本地 HF 模型路径# 或 TensorRT Model Optimizer 量化的检查点，如 nvidia/Llama-3-1-8B-Instruct-FP8 on HFllm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")# 可以将引擎保存到磁盘并稍后加载，LLM 类可以接受 HF 模型或 TRT-LLM 引擎llm.save(tempfile.mkdtemp())# 示例提示prompts = ["Hello, my name is","The president of the United States is","The capital of France is","The future of AI is",]# 创建采样参数sampling_params = SamplingParams(temperature=0.8, top_p=0.95)# 生成文本for output in llm.generate(prompts, sampling_params):print(output)if __name__ == "__main__":main()

3.2 异步推理

对于需要处理大量并发请求的应用，TensorRT-LLM 提供了异步推理 API：

### 异步推理示例 - 使用 TensorRT-LLM 进行异步文本生成import asyncio
import tempfilefrom tensorrt_llm import LLM, SamplingParamsasync def main():# 初始化 LLM 模型llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")# 保存引擎以便后续使用llm.save(tempfile.mkdtemp())# 示例提示prompts = ["Hello, my name is","The president of the United States is","The capital of France is","The future of AI is",]# 创建采样参数sampling_params = SamplingParams(temperature=0.8, top_p=0.95)# 异步生成文本async for output in llm.generate_async(prompts, sampling_params):print(output)if __name__ == "__main__":asyncio.run(main())

3.3 分布式推理

对于大型模型，TensorRT-LLM 支持分布式推理，可以跨多个 GPU 执行模型：

### 分布式推理示例 - 使用 TensorRT-LLM 进行多 GPU 分布式推理import os
import tempfileimport torch.distributedfrom tensorrt_llm import LLM, SamplingParams
from tensorrt_llm.mapping import Mappingdef main():# 初始化分布式环境torch.distributed.init_process_group(backend="nccl")# 获取当前设备信息local_rank = int(os.environ.get("LOCAL_RANK", "0"))world_size = int(os.environ.get("WORLD_SIZE", "1"))# 创建映射配置mapping = Mapping(world_size=world_size, rank=local_rank, tp_size=world_size)# 初始化分布式 LLM 模型llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", mapping=mapping)# 保存引擎以便后续使用llm.save(tempfile.mkdtemp())# 示例提示prompts = ["Hello, my name is","The president of the United States is","The capital of France is","The future of AI is",]# 创建采样参数sampling_params = SamplingParams(temperature=0.8, top_p=0.95)# 分布式生成文本for output in llm.generate(prompts, sampling_params):if local_rank == 0:  # 只在主进程打印结果print(output)if __name__ == "__main__":main()

3.4 自定义推理

TensorRT-LLM 允许用户自定义模型配置和推理流程：

### 自定义推理示例 - 使用 TensorRT-LLM 进行自定义推理流程import tempfilefrom tensorrt_llm import LLM, SamplingParams
from tensorrt_llm.models import LLaMAForCausalLMdef main():# 使用自定义模型配置model_config = {"architecture": "LLaMAForCausalLM","dtype": "float16","num_hidden_layers": 32,"num_attention_heads": 32,"hidden_size": 4096,"vocab_size": 32000,"max_position_embeddings": 4096,"tensor_parallel_size": 1,}# 初始化自定义 LLM 模型llm = LLM(model_config=model_config)# 保存引擎以便后续使用llm.save(tempfile.mkdtemp())# 示例提示prompts = ["Hello, my name is","The president of the United States is","The capital of France is","The future of AI is",]# 创建自定义采样参数sampling_params = SamplingParams(temperature=0.7,top_p=0.9,top_k=50,repetition_penalty=1.1,max_new_tokens=128)# 生成文本for output in llm.generate(prompts, sampling_params):print(output)if __name__ == "__main__":main()

3.5 流式输出

对于需要实时显示生成结果的应用，TensorRT-LLM 提供了流式输出功能：

### 流式输出示例 - 使用 TensorRT-LLM 进行流式文本生成import asyncio
import tempfilefrom tensorrt_llm import LLM, SamplingParamsasync def main():# 初始化 LLM 模型llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")# 保存引擎以便后续使用llm.save(tempfile.mkdtemp())# 示例提示prompt = "Write a short story about artificial intelligence:"# 创建采样参数，设置流式输出sampling_params = SamplingParams(temperature=0.8,top_p=0.95,max_new_tokens=200,streaming=True  # 启用流式输出)# 流式生成文本print(f"Prompt: {prompt}")print("Generated text: ", end="")async for output in llm.generate_async([prompt], sampling_params):# 只打印新生成的 tokenif output.token_ids and len(output.token_ids) > 0:print(output.text, end="", flush=True)print("\nGeneration completed.")if __name__ == "__main__":asyncio.run(main())

4. 性能优化与案例分析

TensorRT-LLM 通过多种优化技术显著提高了大语言模型的推理性能。以下是一些性能分析和案例研究。

4.1 性能概述

TensorRT-LLM 在多种 GPU 上进行了性能测量，以下是主要性能数据和案例分析。需要注意的是，从 TensorRT-LLM v0.10 开始，性能基准测试方法已更改为使用实时批处理，不再使用静态基准测试。这些数据是初步测量结果，预计在未来版本中会有所改进。

4.2 案例分析：Llama 3.1 8B 模型性能

测试环境

• GPU: NVIDIA H100 (80GB)
• 批处理大小: 1-32
• 输入长度: 1024 tokens
• 输出长度: 128 tokens

性能结果

• 单 GPU 吞吐量: 约 2,500 tokens/sec (批处理大小=32)
• 延迟: 约 50ms (批处理大小=1)
• 与 PyTorch 原生实现相比，TensorRT-LLM 提供了约 2-3 倍的性能提升

优化技术

• 使用 FP8 量化将内存占用减少约 40%，同时保持模型质量
• 使用张量并行(Tensor Parallelism)在多 GPU 上扩展模型
• 使用 KV 缓存优化减少内存占用和提高推理速度

4.3 案例分析：多模型并行推理

测试环境

• GPU: 8x NVIDIA A100 (40GB)
• 模型: 多个 7B 参数模型并行运行

性能结果

• 使用 TensorRT-LLM 可以在单个 A100 上同时运行多达 4 个 7B 参数模型
• 与单模型运行相比，总吞吐量提高了约 3.5 倍
• 延迟增加约 20%，但总体资源利用率显著提高

优化技术

• 使用模型并行和流水线并行技术
• 动态批处理和请求调度
• 内存优化以最大化 GPU 利用率

4.4 案例分析：企业级部署优化

客户场景

某大型科技公司需要在其生产环境中部署 LLM 服务，要求高吞吐量和低延迟。

解决方案

• 使用 TensorRT-LLM 优化模型并部署在 NVIDIA H100 集群上
• 实施自定义批处理策略和请求调度
• 使用量化技术减少内存占用

性能结果

• 与原始 PyTorch 部署相比，吞吐量提高了 4 倍
• 95% 的请求延迟降低了 60%
• 服务器利用率从 40% 提高到 85%
• 每个 GPU 可以同时处理的请求数量增加了 3 倍

关键优化

• 使用 TensorRT-LLM 的自定义 CUDA 内核
• 实施连续批处理技术
• 优化内存管理和 KV 缓存
• 使用混合精度计算

4.5 性能优化最佳实践

1. 选择合适的批处理策略

   • 对于延迟敏感的应用，使用小批量或无批处理
   • 对于吞吐量敏感的应用，使用大批量和动态批处理

2. 量化技术

   • 对于大多数模型，FP16 提供良好的性能和精度平衡
   • 对于内存受限场景，考虑使用 INT8 或 FP8 量化
   • 对于极端性能要求，可以探索 INT4 量化，但需要注意精度损失

3. 并行策略

   • 对于小型模型(7B 以下)，单 GPU 通常足够
   • 对于中型模型(7B-70B)，使用张量并行
   • 对于大型模型(70B 以上)，结合张量并行和流水线并行

4. 内存优化

   • 使用注意力机制优化
   • 实施 KV 缓存管理
   • 考虑使用激活检查点(Activation Checkpointing)

5. 高级功能与应用场景

TensorRT-LLM 不仅提供了基本的推理功能，还支持多种高级功能和应用场景。

5.1 多 LoRA 适配器支持

TensorRT-LLM 支持 LoRA（Low-Rank Adaptation）技术，允许用户在推理时动态切换不同的 LoRA 适配器：

# 使用 LoRA 适配器的示例
from tensorrt_llm import LLM, SamplingParams# 加载基础模型
llm = LLM(model="meta-llama/Llama-2-7b-hf")# 加载 LoRA 适配器
lora_weights = ["path/to/lora_adapter_1","path/to/lora_adapter_2"
]# 使用第一个 LoRA 适配器生成文本
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
prompt = "Translate the following English text to French: 'Hello, how are you?'"
for output in llm.generate([prompt], sampling_params, lora_weights=lora_weights[0]):print(output)# 切换到第二个 LoRA 适配器
prompt = "Write a poem about nature:"
for output in llm.generate([prompt], sampling_params, lora_weights=lora_weights[1]):print(output)

5.2 Medusa 解码

TensorRT-LLM 支持 Medusa 解码技术，这是一种可以显著提高生成速度的并行解码方法：

# 使用 Medusa 解码的示例
from tensorrt_llm import LLM, SamplingParams# 加载支持 Medusa 的模型
llm = LLM(model="path/to/medusa/model")# 创建采样参数，启用 Medusa 解码
sampling_params = SamplingParams(temperature=0.8,top_p=0.95,use_medusa=True,  # 启用 Medusa 解码medusa_choices=5   # Medusa 头数量
)# 生成文本
prompt = "Explain the concept of artificial intelligence:"
for output in llm.generate([prompt], sampling_params):print(output)

5.3 Lookahead 解码

TensorRT-LLM 还支持 Lookahead 解码，这是另一种提高生成速度的技术：

# 使用 Lookahead 解码的示例
from tensorrt_llm import LLM, SamplingParams# 加载模型
llm = LLM(model="meta-llama/Llama-2-7b-hf")# 创建采样参数，启用 Lookahead 解码
sampling_params = SamplingParams(temperature=0.8,top_p=0.95,use_lookahead=True,  # 启用 Lookahead 解码lookahead_factor=2    # Lookahead 因子
)# 生成文本
prompt = "Write a summary of quantum computing:"
for output in llm.generate([prompt], sampling_params):print(output)

5.4 应用场景

TensorRT-LLM 适用于多种应用场景，包括但不限于：

1. 聊天机器人和对话系统：利用低延迟推理提供流畅的对话体验
2. 内容生成：用于生成文章、摘要、翻译等内容
3. 企业知识库问答：结合检索增强生成（RAG）技术，提供基于企业知识库的精准回答
4. 代码辅助和生成：用于代码补全、生成和解释
5. 多模态应用：结合图像和文本进行多模态理解和生成
6. 大规模服务部署：在云端或边缘设备上部署高性能 LLM 服务

6. 总结与展望

NVIDIA TensorRT-LLM 为大语言模型的高效推理提供了强大的解决方案。通过其优化的引擎和灵活的 API，用户可以显著提高模型推理性能，降低部署成本，并支持各种复杂的应用场景。

主要优势包括：

• 显著的性能提升：与原生 PyTorch 实现相比，可实现 2-4 倍的性能提升
• 内存优化：通过量化和优化的内存管理，减少 GPU 内存占用
• 灵活的部署选项：支持从单 GPU 到多 GPU 分布式部署
• 丰富的功能：支持异步推理、流式输出、多 LoRA 适配器等高级功能
• 易于使用的 API：提供简洁直观的 Python API，简化开发流程

随着大语言模型在各行各业的应用不断扩展，TensorRT-LLM 将继续发挥重要作用，帮助开发者和企业高效部署和优化大语言模型服务。未来，我们可以期待 TensorRT-LLM 在以下方面的发展：

• 支持更多新兴的模型架构和技术
• 进一步优化性能和内存使用
• 提供更丰富的工具和 API，简化开发和部署流程
• 增强与其他 NVIDIA 生态系统工具的集成

对于希望在生产环境中部署大语言模型的开发者和企业来说，TensorRT-LLM 是一个不可或缺的工具，它能够帮助充分发挥 NVIDIA GPU 的性能潜力，实现高效、经济的大语言模型推理服务。

参考资料

1. NVIDIA TensorRT-LLM GitHub 仓库：https://github.com/NVIDIA/TensorRT-LLM
2. TensorRT-LLM 官方文档：https://github.com/NVIDIA/TensorRT-LLM/tree/main/docs/source
3. NVIDIA 开发者博客：https://developer.nvidia.com/blog
4. TensorRT 官方文档：https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html