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 的工程師邱騫表示。

審核編輯：湯梓紅