NVIDIA TensorRT插件的全自動(dòng)生成工具

NVIDIA TensorRT 是性能最優(yōu)、應(yīng)用最廣的 GPU 推理框架，但用戶常常因?yàn)椴恢С值乃阕佣媾R手寫(xiě)插件的痛點(diǎn)。為此，騰訊 TEG 與 NVIDIA 聯(lián)合開(kāi)發(fā)了一款能自動(dòng)生成插件的工具：TensorRT Plugin Autogen Tool，簡(jiǎn)稱 TPAT。它能支持幾乎所有的 ONNX 算子，在節(jié)省人力成本的同時(shí)，性能對(duì)比手寫(xiě)插件也毫不遜色。

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

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

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

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

騰訊 TEG 表示，手寫(xiě)插件往往是 TensorRT 模型開(kāi)發(fā)中最耗時(shí)的一部分，其難點(diǎn)在于：

手寫(xiě)插件需要熟悉 TensorRT 插件的編程接口，需要學(xué)習(xí) GPU 和 CUDA 相關(guān)知識(shí)。

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

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

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

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

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

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

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

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

TPAT 亮點(diǎn):

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

全自動(dòng)：端到端全自動(dòng)生成用戶指定的 TensorRT Plugin

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

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

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

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

審核編輯：湯梓紅