圖優(yōu)化在降低 AI 模型的訓(xùn)練和推理使用的時(shí)間和資源方面起著重要作用。圖優(yōu)化的一個(gè)重要功能是模型中將可以融合的算子進(jìn)行融合，通過(guò)降低內(nèi)存占用和減少數(shù)據(jù)在低速內(nèi)存中的搬運(yùn)來(lái)提高計(jì)算效率。然而，實(shí)現(xiàn)一套能夠提供各種算子融合的后端方案難度很大，導(dǎo)致在實(shí)際硬件上 AI 模型能夠使用的算子融合非常有限。

Composable Kernel (CK)庫(kù)旨在提供一套在 AMD GPU 上的算子融合的后端方案。CK 使用通用編程語(yǔ)言 HIP C++，完全開(kāi)源。其設(shè)計(jì)理念包括：

• 高性能 & 高生產(chǎn)力：CK 的核心是一組精心設(shè)計(jì)，高度優(yōu)化，可復(fù)用的基礎(chǔ)模塊。CK 庫(kù)內(nèi)所有的算子都是通過(guò)組合這些基礎(chǔ)模塊實(shí)現(xiàn)的。復(fù)用這些基礎(chǔ)模塊大大縮短開(kāi)發(fā)后端算法的周期，同時(shí)還能保證高性能。
• 精通當(dāng)前的 AI 問(wèn)題，快速適應(yīng)未來(lái)的 AI 問(wèn)題：CK 旨在提供一套完整的 AI 算子后端方案，這讓復(fù)雜的算子融合成為可能，因?yàn)檫@樣讓整個(gè)后端都可以用 CK 實(shí)現(xiàn)，而不需依賴(lài)外部算子庫(kù)。CK 的可復(fù)用基礎(chǔ)模塊足以實(shí)現(xiàn)常見(jiàn) AI 模型（機(jī)器視覺(jué)，自然語(yǔ)言處理，等等）所需的各種算子及其融合。當(dāng)新出現(xiàn)的 AI 模型需要新的算子時(shí)，CK 也將會(huì)提供所需的基礎(chǔ)模塊。
• AI 系統(tǒng)專(zhuān)家的簡(jiǎn)單但強(qiáng)大的工具：CK 所有的算子都是用 HIP C++ 模版實(shí)現(xiàn)的。AI 系統(tǒng)專(zhuān)家可以通過(guò)實(shí)例化模版來(lái)定制這些算子的屬性，比如數(shù)據(jù)類(lèi)型，元操作類(lèi)型，張量存儲(chǔ)格式，等等。這通常只需要幾行代碼。
• 友好的 HIP C++ 界面：HPC 算法開(kāi)發(fā)者一直在推動(dòng)著 AI 計(jì)算加速的前沿。CK 的一個(gè)重要設(shè)計(jì)理念就是要讓 HPC 算法開(kāi)發(fā)者更容易對(duì) AI 加速作出貢獻(xiàn)。因此 CK 所有核心模塊都是用 HIP C++ 實(shí)現(xiàn)，而不是 Intermediate Representation （IR）。HPC 算法開(kāi)發(fā)者直接以他們熟悉的編寫(xiě) C++ 代碼的形式編寫(xiě)算法，而無(wú)需像基于 IR 的算子庫(kù)那樣，以通過(guò)編寫(xiě)針對(duì)某種特定算法的 Compiler Pass 來(lái)實(shí)現(xiàn)。這樣做可以大大提高算法的迭代速度。
• 可移植性：今天使用 CK 作為后端的圖優(yōu)化將能夠移植到未來(lái) AMD 的所有的 GPU 上，并且最終也可以被移植到 AMD CPU 上【2】。
• CK 源代碼：https://github.com/ROCmSoftwarePlatform/composable_kernel

核心概念

CK 引入了兩個(gè)概念以提高后端開(kāi)發(fā)者的生產(chǎn)力：

1. 開(kāi)創(chuàng)性的引入“張量坐標(biāo)變換” （Tensor Coordinate Transformation）降低 AI 算子的編寫(xiě)復(fù)雜度。該研究開(kāi)創(chuàng)性地定義了一組可復(fù)用的 Tensor Coordinate Transformation 基礎(chǔ)模塊，并且用它們把復(fù)雜的 AI 算子（比如卷積，group normalization reduction，Depth2Space，等等）以數(shù)學(xué)嚴(yán)謹(jǐn)?shù)姆绞街匦卤磉_(dá)成了最基礎(chǔ)的 AI 算子（GEMM，2D reduction，tensor transfer，等等）。這項(xiàng)技術(shù)可以讓為基礎(chǔ) AI 算子編寫(xiě)的算法直接被用到所有與之對(duì)應(yīng)的復(fù)雜的 AI 算子上，而無(wú)需重寫(xiě)算法。

2. 基于 Tile 的編程范式：開(kāi)發(fā)算子融合的后端算法可以被看成先將每一個(gè)融合前的算子（獨(dú)立算子）拆解成許多 “小塊” 的數(shù)據(jù)操作，然后再把這些 “小塊” 操作組合成融合的算子。每一個(gè)這樣的 “小塊” 操作都對(duì)應(yīng)一個(gè)原始的獨(dú)立算子，但是被操作的數(shù)據(jù)只是原始張量的一部分（tile），因此這樣的 “小塊” 操作被稱(chēng)為 Tile Tensor Operator。CK 庫(kù)包含一組針對(duì) Tile Tensor Operator 的高度優(yōu)化的實(shí)現(xiàn)，CK 里所有的 AI 獨(dú)立算子和融合算子都是用它們實(shí)現(xiàn)的。目前，這些 Tile Tensor Operators 包括 Tile GEMM，Tile Reduction 和 Tile Tensor Transfer。每一個(gè) Tile Tensor Operator 都有針對(duì) GPU thread block，warp 和 thread 的實(shí)現(xiàn)。

Tensor Coordinate Transformation 和 Tile Tensor Operator 共同組成了 CK 的可復(fù)用的基礎(chǔ)模塊。

圖 1，使用 CK 的 Tensor Coordinate Transformation 基礎(chǔ)模塊將 convolution 算子表達(dá)成 GEMM 算子

圖 2，CK 的組成（下：可復(fù)用的基礎(chǔ)模塊；上：獨(dú)立算子與融合算子）

代碼結(jié)構(gòu)

CK 庫(kù)結(jié)構(gòu)分為四層，從下到上分別是：Templated Tile Operator，Templated Kernel and Invoker，Instantiated Kernel and Invoker 和 Client API【3】。每一層對(duì)應(yīng)不同的開(kāi)發(fā)者。

• AI 系統(tǒng)專(zhuān)家：“我需要一個(gè)后端方案提供高性能的獨(dú)立和融合算子讓我可以直接使用”。這個(gè)例子【4】里用的 Client API 和 Instantiated Kernel and Invoker 提供了預(yù)先實(shí)例化并編譯好的對(duì)象，以滿(mǎn)足這類(lèi)開(kāi)發(fā)者的需求。
• AI 系統(tǒng)專(zhuān)家：“我為一個(gè)開(kāi)源的 AI 框架做最先進(jìn)的圖優(yōu)化工作。我需要一個(gè)能夠?yàn)閳D優(yōu)化所需的所有融合算子提供高性能 kernel 的后端方案。同時(shí)我也需要定制這些 kernel，所以像 “要么接受，要么棄用” 的黑盒解決方案不能滿(mǎn)足我的需求”。Templated Kernel and Invoker 層能滿(mǎn)足這類(lèi)開(kāi)發(fā)者。比如這個(gè)例子【5】中開(kāi)發(fā)者可以自己使用 Templated Kernel and Invoker 層實(shí)例化出所需的 FP16 的 GEMM + Add + Add + FastGeLU 的 kernel。
• HPC 算法專(zhuān)家：“我的團(tuán)隊(duì)為公司內(nèi)部不斷迭代的 AI 模型開(kāi)發(fā)高性能后端算法。我們團(tuán)隊(duì)中有 HPC 算法專(zhuān)家，但我們?nèi)匀幌Ｍ梢酝ㄟ^(guò)復(fù)用和改進(jìn)硬件供應(yīng)商提供的高度優(yōu)化的源代碼來(lái)提高我們的生產(chǎn)力，并且讓我們的代碼可以被移植到未來(lái)的硬件構(gòu)架上。我們希望可以不用通過(guò)與硬件提供商分享我們的代碼來(lái)做到這點(diǎn)”。Templated Tile Operator 層可以幫助到這一類(lèi)開(kāi)發(fā)者。比如這個(gè)代碼【6】中開(kāi)發(fā)者使用 Templated Tile Operator 來(lái)實(shí)現(xiàn) GEMM 的優(yōu)化管線(xiàn)。

圖 3，CK 庫(kù)四層結(jié)構(gòu)

基于 AITemplate + CK 的端到端模型推理

Meta 的 AITemplate 【7】（AIT）是一個(gè)統(tǒng)一 AMD 和 Nvidia GPU 的 AI 推理系統(tǒng)。AITemplate 使用 CK 作為其 AMD GPU 上的后端，它使用的是 CK 的 Templated Kernel and Invoker 層。

AITemplate + CK 在 AMD Instinct? MI250 上取得了多個(gè)重要 AI 模型最先進(jìn)的推理性能。CK 里大多數(shù)先進(jìn)的融合算子的定義，都是在 AITemplate 團(tuán)隊(duì)的遠(yuǎn)見(jiàn)下推動(dòng)的。許多融合算子的算法也是由 CK 和 AITemplate 團(tuán)隊(duì)共同設(shè)計(jì)。

本文比較了幾個(gè)端到端模型在 AMD Instinct MI250 和同級(jí)別產(chǎn)品【8】的性能表現(xiàn)。本文中所有 AMD Instinct MI250 的 AI 模型的性能數(shù)據(jù)都是用 AITemplate【9】 + CK【10】取得的。

實(shí)驗(yàn)

ResNet-50

下圖顯示了 AMD Instinct MI250 上的 AIT + CK 與 A100-PCIe-40GB 和 A100-DGX-80GB 上的 TensorRT v8.5.0.12 【11】（TRT）的性能比較。結(jié)果顯示 AMD Instinct MI250 上的 AIT + CK 取得了相比于 A100-PCIe-40GB 上的 TRT 1.08 倍的加速。

BERT

一個(gè)基于 CK 實(shí)現(xiàn)的 Batched GEMM + Softmax + GEMM 融合算子模版，可以完全消除掉中間結(jié)果在 GPU 計(jì)算單元（Compute Unit）與 HBM 之間的搬運(yùn)。通過(guò)使用這個(gè)融合算子模版，attention layer 許多原本是帶寬瓶頸（bandwidth bound）的問(wèn)題變成了計(jì)算瓶頸（compute bound）的問(wèn)題，這樣可以更好發(fā)揮 GPU 的計(jì)算能力。這個(gè) CK 的實(shí)現(xiàn)深受 FlashAttention 【12】的啟發(fā)，并比原始的 FlashAttention 的實(shí)現(xiàn)減少了更多的數(shù)據(jù)搬運(yùn)。

下圖顯示了 AMD Instinct MI250 上的 AIT + CK 與 A100-PCIe-40GB 和 A100-DGX-80GB 上的 FasterTransformer v5.1.1 bug fix 【13】（FT）的 Bert Base 模型（uncased)的性能比較。當(dāng) Sequence 是 4096 時(shí)，F(xiàn)T 在 A100-PCIe-40GB 和 A100-DGX-80GB 上會(huì)在 Batch 32 時(shí) GPU 內(nèi)存溢出。因此，在 Sequence 是 4096 時(shí)，本文只顯示 Batch 16 的結(jié)果。結(jié)果顯示 AMD Instinct MI250 上的 AIT + CK 取得了相比于 A100-PCIe-40GB 上的 FT 3.28 倍，以及相比于 A100-DGX-80GB 上的 FT 2.91 倍的加速。

Vision Transformer （VIT）

下圖顯示了 AMD Instinct MI250 上的 AIT + CK 與 A100-PCIe-40GB 和 A100-DGX-80GB 上的 TensorRT v8.5.0.12（TRT）的 Vision Transformer Base （224x224 圖片）的性能比較。結(jié)果顯示 AMD Instinct MI250 上的 AIT + CK 取得了相比于 A100-PCIe-40GB 上的 TRT 1.8 倍，以及相比于 A100-DGX-80GB 上的 TRT 1.4 倍的加速。

Stable Diffusion

端到端的 Stable Diffusion

下表顯示 AIT + CK 在 AMD Instinct MI250 上 Stable Diffusion 端到端（Batch 1，2，4， 6）的性能數(shù)據(jù)。當(dāng) Batch 是 1 時(shí)，在 MI250 上只有一個(gè) GCD 被使用，而在 Batch 2，4，6 時(shí)，兩個(gè) GCD 都被使用了。

Stable Diffusion 中的 UNet

不過(guò)本文還沒(méi)有關(guān)于使用 TensorRT 運(yùn)行 Stable Diffusion 端到端模型的公開(kāi)的信息。但這篇文章“Make stable diffusion 25% faster using TensorRT” 【14】說(shuō)明了怎么使用 TensorRT 加速 Stable Diffusion 中的 UNet 模型。UNet 是 Stable Diffusion 中最重要最花時(shí)間的部分，因此 UNet 的性能大致反應(yīng)了 Stable Diffusion 的性能。

下圖顯示了 AMD Instinct MI250 上的 AIT + CK 與 A100-PCIe-40GB 和 A100-DGX-80GB 上的 TensorRT v8.5.0.12（TRT）的 UNet 的性能比較。結(jié)果顯示 AMD Instinct MI250 上的 AIT + CK 取得了相比于 A100-PCIe-40GB 上的 TRT 2.45 倍，以及相比于 A100-DGX-80GB 上的 TRT 2.03 倍的加速。

更多信息

ROCm webpage: AMD ROCm? Open Software Platform | AMD

ROCm Information Portal: AMD Documentation - Portal

AMD Instinct Accelerators: AMD Instinct? Accelerators | AMD

AMD Infinity Hub: AMD Infinity Hub | AMD

Endnotes:

1.Chao Liu is PMTS Software Development Engineer at AMD. Jing Zhang is SMTS Software Development Engineer at AMD. Their postings are their own opinions and may not represent AMD’s positions, strategies, or opinions. Links to third party sites are provided for convenience and unless explicitly stated, AMD is not responsible for the contents of such linked sites and no endorsement is implied. GD-5

2.CK for CPU is in early development phase.

3.C++ APIs for now, Python APIs are under planning.

4.Example of CK “Client API” for GEMM + Add + Add + FastGeLU fused operator. https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/685860c2a9483c9e909d2f8bfb95056672491...

5.Example of CK “Templated Kernel and Invoker” of GEMM + Add + Add + FastGeLU fuse operator. https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/685860c2a9483c9e909d2f8bfb95056672491...

6.Example of using CK “Templated Tile Operator” primitives to write a GEMM pipeline. https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/685860c2a9483c9e909d2f8bfb95056672491...

7.Meta’s AITemplate GitHub repository. https://github.com/facebookincubator/AITemplate

8.MI200-71: Testing Conducted by AMD MLSE 10.23.22 using AITemplate https://github.com/ROCmSoftwarePlatform/AITemplate, commit f940d9b) + Composable Kernel  https://github.com/ROCmSoftwarePlatform/composable_kernel, commit 40942b9) with ROCm?5.3 running on 2x AMD EPYC 7713 64-Core Processor server with 4x AMD Instinct MI250 OAM (128 GB HBM2e) 560W GPU with AMD Infinity Fabric? technology vs. TensorRT v8.5.0.12 and FasterTransformer (v5.1.1 bug fix) with CUDA? 11.8 running on 2x AMD EPYC 7742 64-Core Processor server with 4x Nvidia A100-PCIe-40GB (250W) GPU and TensorRT v8.5.0.12 and FasterTransformer (v5.1.1 bug fix) with CUDA? 11.8 running on 2xAMD EPYC 7742 64-Core Processor server with 8x NVIDIA A100 SXM 80GB (400W) GPU. Server manufacturers may vary configurations, yielding different results. Performance may vary based on factors including use of latest drivers and optimizations. 

9.https://github.com/ROCmSoftwarePlatform/AITemplate/tree/f940d9b7ac8b976fba127e2c269dc5b368f30e4e

10.https://github.com/ROCmSoftwarePlatform/composable_kernel/tree/40942b909801dd721769834fc61ad201b5795...

11.TensorRT GitHub repository. https://github.com/NVIDIA/TensorRT

12.FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. https://arxiv.org/abs/2205.14135

13.FasterTransformer GitHub repository. https://github.com/NVIDIA/FasterTransformer

14.Making stable diffusion 25% faster using TensorRT. https://www.photoroom.com/tech/stable-diffusion-25-percent-faster-and-save-seconds/

15.During their time in AMD