What is TensorRT?

Definition

TensorRT 是 NVIDIA 的专有 SDK，用于优化和部署深度学习模型在 NVIDIA GPU 上的推理。它接收来自 PyTorch 或 TensorFlow 等框架的训练模型，应用激进的硬件特定优化——包括层融合、精度校准、内核自动调优和内存优化——并生成高度优化的推理引擎，从 NVIDIA 硬件中提取最大性能。

对于 LLM 推理，TensorRT-LLM 是 NVIDIA 的专用扩展，增加了 Transformer 特定的优化。这些包括 flash attention 实现、在途批处理（在其他请求仍在生成时处理新请求）、KV cache 管理、跨多个 GPU 的张量并行以及用于注意力和前馈层的自定义 CUDA 内核。TensorRT-LLM 可以在相同硬件上为相同模型提供比标准推理框架高 2-5 倍的吞吐量。

TensorRT 在不同于框架级推理的层面上运行。PyTorch 使用其即时执行模型逐个执行操作，而 TensorRT 分析整个计算图，识别优化机会，并将计算图编译为针对特定 GPU 架构量身定制的整体执行计划。这种全图优化方法是 TensorRT 能够实现如此显著性能提升的原因——它消除了更灵活执行模型中固有的开销。

Why It Matters

对于生产 LLM 服务，推理成本通常是最大的支出。一个使用标准推理每秒生成 50 个 token 但使用 TensorRT 优化后每秒生成 200 个 token 的模型，意味着每 token 服务成本降低 4 倍。在大规模部署中，这意味着每年节省数十万美元的 GPU 成本。

TensorRT 对于延迟敏感的应用尤为重要，如实时聊天、代码补全和语音助手，用户期望亚秒级响应时间。优化内核、高效内存管理和硬件特定调优的组合使 TensorRT 能够实现通用推理框架无法达到的延迟。

How It Works

TensorRT 优化遵循多阶段管道。首先，从原始格式（ONNX、PyTorch 或 TensorFlow）将模型解析到 TensorRT 的内部图表示中。接下来，图优化遍历将兼容的相邻操作融合在一起——例如，将卷积、偏置加法和激活合并为单个内核启动，消除中间内存分配和内核启动开销。

精度校准然后确定每层的最佳精度。TensorRT 可以在单个模型中混合 FP32、FP16 和 INT8 精度，在精度至关重要的地方使用更高精度，在不重要的地方使用更低精度。最后，内核自动调优通过基准测试多个实现并选择最优者，为目标 GPU 架构上的每个操作选择最快的 CUDA 内核实现。结果是一个针对特定 GPU 型号优化的序列化引擎文件。

Example Use Case

一家 SaaS 公司在 8 块 A100 GPU 上为数千个并发用户提供 7B 参数模型的服务，部署 TensorRT-LLM。使用标准 vLLM 服务时，他们实现了每秒 800 token 的总吞吐量。在进行 TensorRT-LLM 优化（FP8 精度、在途批处理和张量并行）后，吞吐量增加到每秒 2400 token——在不增加额外硬件的情况下将服务能力提高了三倍，每 token 成本降低了 67%。

Key Takeaways

• TensorRT 是 NVIDIA 的推理优化器，在 NVIDIA GPU 上最大化吞吐量。
• TensorRT-LLM 增加了 Transformer 特定的优化，如 flash attention 和在途批处理。
• 通过全图优化，它实现了比标准推理框架高 2-5 倍的吞吐量。
• 混合精度支持和内核自动调优从特定 GPU 架构中提取最大性能。
• 性能提升直接转化为大规模部署中更低的每 token 推理成本。