TensorRT-LLM：大模型推理加速引擎的架构与实践

前言：技术背景与发展历程：

随着GPT-4、LLaMA等千亿级参数模型的出现，传统推理框架面临三大瓶颈：显存占用高（单卡可达80GB）、计算延迟大（生成式推理需迭代处理）、硬件利用率低（Transformer结构存在计算冗余）。根据MLPerf基准测试，原始PyTorch推理的token生成速度仅为12.3 tokens/s（A100显卡）。

---

一、TensorRT-LLM 介绍：

TensorRT-LLM 是一款由 NVIDIA 推出的大语言模型（LLMs）推理加速框架，为用户提供了一个易于使用的 Python API，并使用最新的优化技术将大型语言模型构建为 TensorRT 引擎文件，以便在 NVIDIA GPU 上高效地进行推理。

TensorRT-LLM是一个由Nvidia设计的开源框架，用于在生产环境中提高大型语言模型的性能。该框架是基于 TensorRT 深度学习编译框架来构建、编译并执行计算图，并借鉴了许多 FastTransformer 中高效的 Kernels 实现，并且可以利用 NCCL 完成设备之间的通讯。

TensorRT-LLM 也提供了支持被 NVIDIA Triton Inference Server 集成的后端，用于将模型部署成在线推理服务，并且支持 In-Flight Batching 技术，可以显著提升推理服务的吞吐率并降低时延。

TensorRT-LLM可加速并优化 NVIDIA GPU 上最新大型语言模型 (LLM) 的推理性能。LLM 彻底改变了人工智能领域，并创造了与数字世界互动的全新方式。然而，随着世界各地的组织和应用程序开发者纷纷将 LLM 融入到他们的工作中，运行这些模型的一些挑战也逐渐显现。简而言之，LLM 规模庞大。如果没有合适的技术，这一事实可能会导致其运行成本高昂且速度缓慢。

---

二、TensorRT-LLM 的优势：

TensorRT-LLM 是 NVIDIA 用于做 LLM（Large Language Model）的可扩展推理方案。该方案是基于 TensorRT 深度学习编译框架来构建、编译并执行计算图，并借鉴了许多 FastTransformer 中高效的 Kernels 实现，然后利用 NCCL 完成设备之间的通讯。

考虑到技术的发展和需求的差异，开发者还可以定制算子来满足定制需求，比如基于 cutlass 开发定制 GEMM，TensorRT-LLM 是一款致力于提供高性能并不断完善其实用性的 NVIDIA 官方推理方案。

TensorRT-LLM 还提供了类似于 Pytorch 的 API 来降低开发者的学习成本，并提供了许多预定义好的模型供用户使用。

考虑到大语言模型比较大，有可能单卡放不下，需要多卡甚至多机推理，因此 TensorRT-LLM 还提供了 Tensor Parallelism 和 Pipeline Parallelism 两种并行机制来支持多卡或多机推理。

---

三、Linux环境安装TensorRT-LLM：

3.1 环境准备‌：

‌硬件要求‌：

NVIDIA GPU（需支持目标架构，如 Ampere/Ada Lovelace）
磁盘空间 ≥ 63GB（TensorRT-LLM 全量编译场景）

软件依赖‌：

CUDA 12.2+‌（推荐，确保与 TensorRT 兼容）
‌cuDNN 8.9+‌
‌NVIDIA 驱动‌（与 CUDA 版本匹配）
‌Docker 运行时‌（推荐
‌Git LFS‌（管理大文件
‌Python 3.10+‌ 及 pip310
编译工具链（gcc/cmake/make

---

3.2 Docker全自动构建：

Git LFS（Large File Storage）是一个开源的Git大文件版本控制的解决方案和工具集，旨在帮助管理大型文件，避免Git仓库体积过大，提高克隆和拉取的速度

拉取源码：

sudo apt-get update && sudo apt-get -y install git git-lfs
git lfs install
git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git submodule update --init --recursive
git lfs pull

---

3.3 执行构建‌：

‌完整构建‌（支持所有 GPU 架构）：

make -C docker release_build

‌限定 GPU 架构‌（加速编译，如 RTX 4090）：

make -C docker release_build CUDA_ARCHS="89-real;90-real"

编译时长约 1-2 小时，请确保 /tmp 目录有足够空间。

验证安装‌，进入编译容器：

docker run --gpus all -it --rm -v $(pwd):/workspace tensorrt-llm:latest

在容器内执行：

python3 -c "import tensorrt_llm; print(tensorrt_llm.__version__)"

若输出版本号（如 0.10.0）则安装成功。

---

四、Ollama 安装与配置‌：

4.1 安装 Ollama‌：

# 官方脚本安装（网络不佳时可替换下载源为 GitHub 加速代理）
curl -fsSL https://ollama.com/install.sh | sh
# 验证安装
ollama --version  # 应输出版本号（如 0.6.6）:ml-citation{ref="4,9" data="citationList"}

4.2 下载并运行模型‌：

# 下载模型（以 llama3 为例）
ollama pull llama3

# 启动模型服务
ollama run llama3

模型默认存储在 ~/.ollama/models/，可通过 OLLAMA_MODELS 环境变量修改路径。

---

五、整合 TensorRT-LLM 与 Ollama‌：

以下是 TensorRT-LLM 在 Linux 上的安装及整合 Ollama 运行 LLM 的完整流程，将 TensorRT-LLM 引擎接入 Ollama‌：

5.1 修改 Ollama 配置文件 ~/.ollama/config.json，指定推理后端：

使用 TensorRT-LLM 编译模型引擎（参考 examples/llama 示例）

{
  "model": {
    "backend": "tensorrt_llm",
    "engine_path": "/path/to/your/trt_engine"
  }
}

需重启 Ollama 服务生效。

---

5.2 通过 API 调用验证‌：

curl http://localhost:11434/api/generate -d '{
  "model": "llama3",
  "prompt": "为什么天空是蓝色的？"
}'

若返回 JSON 格式的生成结果，则整合成功。

---

5.3 服务自启动配置（可选）‌：

# 创建 Ollama 系统服务
sudo tee /etc/systemd/system/ollama.service > /dev/null <<EOF
[Unit]
Description=Ollama Service
After=network.target

[Service]
ExecStart=/usr/local/bin/ollama serve
User=$USER
Restart=always

[Install]
WantedBy=default.target
EOF

启用服务：

sudo systemctl enable ollama
sudo systemctl start ollama

---

5.4 常见问题解决‌：

(1). TensorRT-LLM 编译失败‌：

检查 CUDA/cuDNN 版本兼容性（需严格匹配 TensorRT 版本），增加 /tmp 分区空间或设置 TMPDIR 环境变量

(2). Ollama 模型下载缓慢‌：

使用 GitHub 代理下载离线模型包，手工放入 ~/.ollama/models/

(3). GPU 显存不足‌：

为 Ollama 添加 --num-gpu-layers 20 参数限制 GPU 层数（如 ollama run llama3 --num-gpu-layers 20）

以上流程通过容器化隔离环境依赖，确保 TensorRT-LLM 的高性能推理与 Ollama 的易用性结合。