NVIDIA TensorRT插件的全自动生成工具

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

腾讯 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 会因此成为一个真正端到端的推理框架。

TPAT 亮点:

覆盖度：几乎支持 ONNX 所有的算子

全自动：端到端全自动生成用户指定的 TensorRT Plugin

高性能：大部分算子的性能与手写插件相比并不逊色

值得一提的是，最新发布的 TPAT 也支持了动态 batch，从而成为真正具有业界实用性的工具。

TPAT 已经用在了腾讯游戏中心的多个游戏 AI 模型上，体现了优秀的好用性与性能。此外，很多其他腾讯团队也对 TPAT 感兴趣，在评估与尝试部署 TPAT 生成的 TensorRT 模型。

“对内部业务模型里的部分算子进行了测试，TPAT 的性能几乎全面达到甚至超越 CUDA 工程师手写的插件，并且端到端的设计能够大幅减少人力投入。”来自腾讯 TEG 的工程师邱骞表示。

　　审核编辑：汤梓红