一行代碼的威力到底有多大？今天我們要介紹的這個 Kernl 庫，用戶只需一行代碼，在 GPU 上就能以快幾倍的速度運行 Pytorch transformer 模型，從而極大的加快了模型的推理速度。?

具體而言，有了 Kernl 的加持，Bert 的推理速度比 Hugging Face 基線快了 12 倍。這一成果主要得益于 Kernl 用新的 OpenAI 編程語言 Triton 和 TorchDynamo 編寫了定制的 GPU 內核。項目作者來自 Lefebvre Sarrut。

GitHub 地址：https://github.com/ELS-RD/kernl/

以下是 Kernl 與其他推理引擎的比較，橫坐標中括號里的數字分別表示 batch size、序列長度，縱坐標為推理加速情況。

基準測試在 3090 RTX GPU 運行，以及 12 核 Intel CPU。?

由上述結果可得，在長序列輸入這一塊，Kernl 可以說是最快的推理引擎（上圖中的右半部分），在短輸入序列上接近英偉達的 TensorRT（上圖中的左半部分）。除此以外，Kernl 內核代碼非常簡短，易于理解和修改。該項目甚至添加了 Triton 調試器和工具 (基于 Fx) 來簡化內核替換，因此不需要修改 PyTorch 模型源代碼。?

項目作者 Micha?l Benesty 對這一研究進行了總結，他們發布的 Kernl 是一個用于加速 transformer 推理的庫，速度非常快，有時會到達 SOTA 性能，可破解以匹配大多數 transformer 架構。

他們還在 T5 上做了測試，速度提高 6 倍，Benesty 表示這僅僅是個開始。

為什么創建 Kernl？?

在 Lefebvre Sarrut，項目作者在生產中運行幾個 transformers 模型，其中一些對延遲敏感，主要是搜索和 recsys。他們還在使用 OnnxRuntime 和 TensorRT，甚至創建了 transformer-deploy OSS 庫來與社區分享知識。?

最近，作者在測試生成語言，并努力加速它們。然而事實證明，使用傳統工具做到這些非常困難。在他們看來，Onnx 是另一種有趣的格式，它是一種針對機器學習所設計的開放式文件格式，用于存儲訓練好的模型，具有廣泛的硬件支持。

但是，當他們處理新的 LLM 架構時，Onnx 生態系統（主要是推理引擎）存在以下幾種限制：?

• 沒有控制流的模型導出到 Onnx 很簡單，這是因為可以依賴跟蹤。但是動態行為更難獲得；
• 與 PyTorch 不同，ONNX Runtime/TensorRT 還沒有原生支持實現張量并行的多 GPU 任務；
• TensorRT 無法為具有相同配置文件的 transformer 模型管理 2 個動態軸。但由于通常希望能夠提供不同長度的輸入，因此需要每個批大小構建 1 個模型；
• 非常大的模型很常見，但 Onnx（作為 protobuff 文件）在文件大小方面有一些限制，需要將權重存儲在模型之外來解決問題。?

一個非常煩人的事實是新模型永遠不會被加速，你需要等著其他人來為此編寫自定義 CUDA 內核。現有解決方案并不是不好，OnnxRuntime 的一大優點是它的多硬件支持，TensorRT 則以非常快速著稱。

所以，項目作者想要在 Python/PyTorch 上有像 TensorRT 一樣快的優化器，這也是他們創建 Kernl 的原因。

如何做到？?

內存帶寬通常是深度學習的瓶頸，為了加速推理，減少內存訪問往往是一個很好的策略。在短輸入序列上，瓶頸通常與 CPU 開銷有關，它必須被消除。項目作者主要利用了以下 3 項技術：?

首先是 OpenAI Triton，它是一種編寫 CUDA 等 GPU 內核的語言，不要將它與 Nvidia Triton 推理服務器混淆，它的效率更高。幾個操作的融合實現了改進，使得他們不在 GPU 內存中保留中間結果的情況下鏈接計算。作者使用它重寫注意力（由 Flash Attention 替換）、線性層和激活以及 Layernorm/Rmsnorm。?

其次是 CUDA 圖。在預熱（warmup）步驟中，它將保存每個啟動的內核及它們的參數。然后，項目作者重建了整個推理過程。?

最后是 TorchDynamo，這個由 Meta 提出的原型機幫助項目作者應對動態行為。在預熱步驟中，它會跟蹤模型并提供一個 Fx 圖（靜態計算圖）。他們使用自己的內核替換了 Fx 圖的一些操作，并在 Python 中重新編譯。

未來，項目路線圖將涵蓋更快的預熱、ragged 推理（padding 中沒有損失計算）、訓練支持（長序列支持）、多 GPU 支持（多并行化模式）、量化（PTQ）、新 batch 的 Cutlass 內核測試以及提升硬件支持等。

更多詳細內容請參閱原項目。