腾讯 TEG 与 NVIDIA 团队共同开发 TPAT
2022-07-15 20:10
分享到微信
分享到微博
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 的工程师邱骞表示。
