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

NVIDIA TensorRT 是性能最優(yōu)、應用最廣的 GPU 推理框架，但用戶常常因為不支持的算子而面臨手寫插件的痛點。為此，騰訊 TEG 與 NVIDIA 聯(lián)合開發(fā)了一款能自動生成插件的工具：TensorRT Plugin Autogen Tool，簡稱 TPAT。它能支持幾乎所有的 ONNX 算子，在節(jié)省人力成本的同時，性能對比手寫插件也毫不遜色。

騰訊 TEG（技術工程團隊）為騰訊各個團隊提供技術與工程服務，日常任務是為各種產(chǎn)品提供工程設計與實現(xiàn)，以滿足其他團隊的業(yè)務需要。

隨著 AI 應用的普及，騰訊 TEG 在深度學習模型的推理上也開始了技術積累。騰訊 TEG 選擇了TensorRT作為 NVIDIA GPU 上推理應用的基礎軟件，開發(fā)了大量基于 TensorRT 的 AI 模型，并與 NVIDIA 展開了深入合作。

TensorRT 是一個由 NVIDIA 開發(fā)的用于推理的高性能 SDK，包含深度學習推理優(yōu)化器和運行時環(huán)境，可以為深度學習推理應用提供低延遲和高吞吐量，這也是 TensorRT 相比其他框架最大的優(yōu)勢。

不過，TensorRT 也存在一定局限性，當現(xiàn)有模型的算子不能被 TensorRT 直接支持時，工程師不得不手寫插件，比較耗時耗力。

騰訊 TEG 表示，手寫插件往往是 TensorRT 模型開發(fā)中最耗時的一部分，其難點在于：

手寫插件需要熟悉 TensorRT 插件的編程接口，需要學習 GPU 和 CUDA 相關知識。

新入門的工程師通常需要 1~2 周時間來編寫一個算子實現(xiàn)，模型中如果包含多個不支持算子，就需要更多時間來逐個編寫和調(diào)試。

如果有一種自動生成插件的工具，將極大地提高 TensorRT 模型開發(fā)的效率，于是騰訊 TEG 與 NVIDIA 開展了這方面的合作，歷時數(shù)月，共同開發(fā)了TPAT（TensorRT Plugin Autogen Tool）。

TPAT 本質(zhì)上是對 TensorRT 的功能擴展。它要解決的關鍵問題在于如下方面：

為了生成通用的插件，需要開發(fā)出 TensorRT 插件的代碼模板。與此同時，需要提供部分插件的手工參考實現(xiàn)，用于比對計算結(jié)果與性能。這一部分由 NVIDIA 負責。

對于給定的具體算子，其 GPU 上的 CUDA kernel 需要一種自動生成機制。經(jīng)騰訊 TEG 與 NVIDIA 協(xié)商，選擇 ONNX 作為候選算子集合，在編寫中間代碼之后由 TVM 生成 CUDA kernel。這一部分由騰訊 TEG 負責。

在完成上述開發(fā)之后，騰訊 TEG 與 NVIDIA 進行聯(lián)調(diào)，并后續(xù)同步進行版本迭代。

歷經(jīng)騰訊 TEG 與 NVIDIA 工程師的共同努力，TPAT 已正式發(fā)布 1.0 版本。它基于 TVM 和 ONNX，實現(xiàn)了 TensorRT 插件的全自動生成，使得 TensorRT 的部署和上線能基本流程化，不再需要人工參與。

另外，TPAT 全自動生成一個算子插件耗時僅需要 30-60 分鐘的時間（該時間用于搜索算子的高性能 CUDA kernel），TensorRT 會因此成為一個真正端到端的推理框架。

TPAT 亮點:

覆蓋度：幾乎支持 ONNX 所有的算子

全自動：端到端全自動生成用戶指定的 TensorRT Plugin

高性能：大部分算子的性能與手寫插件相比并不遜色

值得一提的是，最新發(fā)布的 TPAT 也支持了動態(tài) batch，從而成為真正具有業(yè)界實用性的工具。

TPAT 已經(jīng)用在了騰訊游戲中心的多個游戲 AI 模型上，體現(xiàn)了優(yōu)秀的好用性與性能。此外，很多其他騰訊團隊也對 TPAT 感興趣，在評估與嘗試部署 TPAT 生成的 TensorRT 模型。

“對內(nèi)部業(yè)務模型里的部分算子進行了測試，TPAT 的性能幾乎全面達到甚至超越 CUDA 工程師手寫的插件，并且端到端的設計能夠大幅減少人力投入。”來自騰訊 TEG 的工程師邱騫表示。

審核編輯：湯梓紅