腾讯 TEG 与 NVIDIA 团队共同开发 TPAT：TensorRT 插件的全自动生成工具

案例简介

• NVIDIA TensorRT 是性能最优、应用最广的 GPU 推理框架，但用户常常因为不支持的算子而面临手写插件的痛点。为此，腾讯 TEG 与 NVIDIA 联合开发了一款能自动生成插件的工具：TensorRT Plugin Autogen Tool，简称 TPAT。它能支持几乎所有的 ONNX 算子，在节省人力成本的同时，性能对比手写插件也毫不逊色。

• 本案例主要应用到 NVIDIA TensorRT。

Case Introduction

• NVIDIA TensorRT is the most efficient and widely used GPU inference framework. However, users often suffer the pain of handwritten plugins because of unsupported operators. To solve the problem, Tencent TEG and NVIDIA jointly developed a tool that can automatically generate plugins: TensorRT Plugin Autogen Tool, or TPAT for short. It can support almost all ONNX operators and save labor costs significantly, without loss of the performance.

• The major products utilized in the case is NVIDIA TensorRT.

客户简介及应用背景

腾讯 TEG（技术工程团队）为腾讯各个团队提供技术与工程服务，日常任务是为各种产品提供工程设计与实现，以满足其他团队的业务需要。

随着 AI 应用的普及，腾讯 TEG 在深度学习模型的推理上也开始了技术积累。腾讯 TEG 选择了 TensorRT 作为 NVIDIA GPU 上推理应用的基础软件，开发了大量基于 TensorRT 的 AI 模型，并与 NVIDIA 展开了深入合作。

客户挑战

TensorRT 是一个由 NVIDIA 开发的用于推理的高性能 SDK，包含深度学习推理优化器和运行时环境，可以为深度学习推理应用提供低延迟和高吞吐量，这也是 TensorRT 相比其他框架最大的优势。

TensorRT 也存在一定局限性，当现有模型的算子不能被 TensorRT 直接支持时，工程师不得不手写插件，比较耗时耗力。

腾讯 TEG 表示，手写插件往往是 TensorRT 模型开发中最耗时的一部分，其难点在于：

手写插件需要熟悉 TensorRT 插件的编程接口，需要学习 GPU 和 CUDA 相关知识。
新入门的工程师通常需要 1~2 周时间来编写一个算子实现，模型中如果包含多个不支持算子，就需要更多时间来逐个编写和调试。

如果有一种自动生成插件的工具，将极大地提高 TensorRT 模型开发的效率，于是腾讯 TEG 与 NVIDIA 开展了这方面的合作，历时数月，共同开发 TPAT（TensorRT Plugin Autogen Tool）。

应用方案

TPAT 本质上是对 TensorRT 的功能扩展。它要解决的关键问题在于如下方面：

为了生成通用的插件，需要开发出 TensorRT 插件的代码模板。与此同时，需要提供部分插件的手工参考实现，用于比对计算结果与性能。这一部分由 NVIDIA 负责。
对于给定的具体算子，其 GPU 上的 CUDA kernel 需要一种自动生成机制。经腾讯 TEG 与 NVIDIA 协商，选择 ONNX 作为候选算子集合，在编写中间代码之后由 TVM 生成 CUDA kernel。这一部分由腾讯 TEG 负责。
在完成上述开发之后，腾讯 TEG 与 NVIDIA 进行联调，并后续同步进行版本迭代。

历经腾讯 TEG 与 NVIDIA 工程师的共同努力，TPAT 已正式发布 1.0 版本。它基于 TVM 和 ONNX，实现了 TensorRT 插件的全自动生成，使得 TensorRT 的部署和上线能基本流程化，不再需要人工参与。

另外，TPAT 全自动生成一个算子插件耗时仅需要 30-60 分钟的时间（该时间用于搜索算子的高性能 CUDA kernel），TensorRT 会因此成为一个真正端到端的推理框架。