TensorRT-LLM-00-总览

一、项目摘要

1.1 项目目标

TensorRT-LLM 是由 NVIDIA 开发的开源大语言模型（LLM）推理优化工具库，旨在为 NVIDIA GPU 上的大语言模型提供高性能推理能力。项目核心目标包括：

  • 提供最先进的 LLM 推理优化能力，包括自定义注意力内核、inflight batching、分页 KV 缓存等
  • 支持多种量化技术（FP8、FP4、INT4 AWQ、INT8 SmoothQuant）降低推理成本
  • 支持推测解码（Speculative Decoding）提升吞吐量
  • 提供基于 PyTorch 的模块化架构，便于开发者扩展和定制
  • 支持单 GPU 到多 GPU、多节点的灵活部署方案
  • 集成主流推理生态系统（Triton Inference Server、NVIDIA Dynamo）

1.2 项目边界

核心功能：

  • LLM 模型的 TensorRT 引擎构建（Builder）
  • 高性能推理运行时（Runtime/Executor）
  • 多种模型架构支持（Llama、GPT、Qwen、DeepSeek、Mamba 等）
  • 自动并行化策略（Tensor Parallel、Pipeline Parallel、Expert Parallel）
  • 量化和模型优化（Quantization、LoRA、PEFT）
  • 多模态支持（视觉、音频输入）

非核心功能（不包含）：

  • 模型训练
  • 模型微调（仅支持 LoRA/PEFT 的推理）
  • 端到端的模型服务系统（由 Triton 等外部工具提供）
  • 数据预处理和后处理业务逻辑

1.3 运行环境

  • 硬件： NVIDIA GPU（Ampere、Ada Lovelace、Hopper、Blackwell 架构）
  • 操作系统： Linux（主要）、Windows（部分支持）、Grace Hopper
  • CUDA 版本： CUDA 13.0.0+
  • TensorRT 版本： TensorRT 10.13.2+
  • Python 版本： Python 3.10、3.12
  • 依赖库： PyTorch、transformers、NCCL、MPI

1.4 版本信息

  • 当前版本： v1.2.0rc1
  • 许可证： Apache 2.0

二、整体架构

TensorRT-LLM 采用九层架构设计：

  1. 用户层： 应用程序入口
  2. Python 高层 API 层： LLM 类提供易用接口
  3. 构建层： 模型编译和 TensorRT 引擎构建
  4. 模型层： 各种 LLM 架构实现
  5. PyTorch 后端层： PyTorch 原生实现
  6. 执行层： 请求管理和多进程协调
  7. C++ 核心层： 高性能 C++ 实现
  8. TensorRT 层： NVIDIA TensorRT 引擎
  9. 硬件层： NVIDIA GPU

2.1 关键数据流

用户输入 → 分词 → Token IDs → Executor → BatchManager → TensorRT → GPU
                                                                    ↓
用户输出 ← 解码 ← Token IDs ← Executor ← BatchManager ← Sampling ← Logits

2.2 核心组件

2.2.1 LLM API（Python 高层接口）

  • 职责：提供用户友好的 Python API
  • 关键类：LLM、MultimodalEncoder
  • 主要方法：generate()、generate_async()

2.2.2 Builder（引擎构建器）

  • 职责：将模型编译为优化的 TensorRT 引擎
  • 关键类：Builder、BuildConfig
  • 优化：层融合、精度校准、内核自动调优

2.2.3 Executor（执行器）

  • 职责：管理推理请求，协调多进程执行
  • 关键类：GenerationExecutor、Worker、Proxy
  • 特性：Inflight Batching、请求队列管理

2.2.4 Runtime（运行时）

  • 职责：管理 TensorRT 推理会话
  • 关键类：ModelRunner、Session、Generation
  • 特性：KV Cache 管理、采样策略

2.2.5 BatchManager（批处理管理器，C++）

  • 职责：动态批处理和 KV Cache 管理
  • 特性：Capacity Scheduling、Paged KV Cache
  • 优化：减少内存碎片，提高利用率

三、模块清单

模块ID 模块名称 目录 职责
01 LLMAPI tensorrt_llm/llmapi/ 高层 Python API
02 Builder tensorrt_llm/builder.py 引擎构建
03 Runtime tensorrt_llm/runtime/ Python 运行时
04 Executor tensorrt_llm/executor/ 请求管理
05 Models tensorrt_llm/models/ 模型实现
06 Layers tensorrt_llm/layers/ 神经网络层
07 Quantization tensorrt_llm/quantization/ 量化支持
08 AutoParallel tensorrt_llm/auto_parallel/ 自动并行
09 TorchBackend tensorrt_llm/_torch/ PyTorch 后端
10 C++ Core cpp/ C++ 核心实现

四、关键技术与优化

4.1 性能优化

Attention 优化：

  • Flash Attention V2（FMHA）：Context Phase 优化
  • XQA Kernel：Generation Phase 优化
  • Multi-Block Attention：长序列优化

内存优化：

  • Paged KV Cache：减少内存碎片
  • Prefix Caching：共享公共前缀
  • Block Reuse：跨请求共享

量化技术：

  • FP8：Hopper 架构原生支持
  • FP4：Blackwell 架构优化
  • INT4/INT8：AWQ、SmoothQuant

批处理优化：

  • Inflight Batching：动态批次调整
  • Continuous Batching：降低延迟
  • Chunked Context：长上下文处理

4.2 并行策略

  • Tensor Parallel： 跨 GPU 切分权重
  • Pipeline Parallel： 跨 GPU 切分层
  • Expert Parallel： MoE 模型专家并行
  • Auto Parallel： 自动策略搜索

4.3 性能指标

  • Llama 3.1 405B： 400 tok/s（H100 单节点）
  • Llama 4 Maverick： 1000+ TPS/User（Blackwell）
  • DeepSeek R1： 世界纪录性能（Blackwell）

五、推理流程

5.1 冷启动（首次运行）

  1. 模型加载： 从 HuggingFace 或本地加载
  2. 引擎构建： TensorRT 优化和编译（5-30 分钟）
  3. 引擎缓存： 保存到磁盘，避免重复构建
  4. 执行器初始化： Worker 进程启动，KV Cache 分配

5.2 热路径（推理执行）

  1. 请求提交： 用户提交 prompt
  2. 分词： 文本转 Token IDs
  3. 批次调度： BatchManager 分配资源
  4. Context Phase： 处理输入序列（首 Token）
  5. Generation Phase： 自回归生成（逐 Token）
  6. 采样： Top-K/Top-P/Temperature
  7. 解码： Token IDs 转文本
  8. 流式返回： 实时返回生成结果

六、模块交互

6.1 主要调用关系

LLM API
  ├─→ Builder（构建时，同步）
  └─→ Executor（推理时，异步）
        ├─→ Worker（多进程）
        │     └─→ C++ Executor（跨语言）
        │           └─→ BatchManager（批处理）
        │                 └─→ TensorRT（推理）
        │                       └─→ GPU（执行）
        └─→ PostprocWorker（后处理，可选）

6.2 数据流

Text → Tokens → GenerationRequest → Batch → TensorRT → Logits → Tokens → Text

6.3 同步与异步

  • 同步操作： 引擎构建、模型加载、单步推理
  • 异步操作： 请求提交、响应获取、流式生成

七、配置管理

7.1 构建配置（BuildConfig）

  • max_batch_size：最大批次大小
  • max_input_len：最大输入长度
  • max_seq_len：最大序列长度
  • max_num_tokens：最大 Token 数
  • precision：精度（FP16/BF16/FP8/FP4）

7.2 执行配置（ExecutorConfig）

  • max_batch_size：运行时最大批次
  • kv_cache_config：KV Cache 配置
  • gpu_weights_percent：权重 GPU 占比
  • enable_chunked_context：启用分块上下文

7.3 采样配置（SamplingParams）

  • temperature：温度系数
  • top_k：Top-K 采样
  • top_p：Top-P（Nucleus）采样
  • repetition_penalty：重复惩罚
  • max_new_tokens：最大生成长度

八、总结

TensorRT-LLM 是一个高度模块化、高性能的 LLM 推理框架：

核心优势：

  • 世界级推理性能
  • 灵活并行化策略
  • 丰富模型支持
  • PyTorch 原生架构

适用场景：

  • 大规模生产推理
  • 极致性能需求
  • 多模型服务
  • 研究开发