圖優化在降低 AI 模型的訓練和推理使用的時間和資源方面起著重要作用。圖優化的一個重要功能是模型中將可以融合的算子進行融合，通過降低內存占用和減少數據在低速內存中的搬運來提高計算效率。然而，實現一套能夠提供各種算子融合的后端方案難度很大，導致在實際硬件上 AI 模型能夠使用的算子融合非常有限。

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

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

核心概念

CK 引入了兩個概念以提高后端開發者的生產力：

1. 開創性的引入“張量坐標變換” （Tensor Coordinate Transformation）降低 AI 算子的編寫復雜度。該研究開創性地定義了一組可復用的 Tensor Coordinate Transformation 基礎模塊，并且用它們把復雜的 AI 算子（比如卷積，group normalization reduction，Depth2Space，等等）以數學嚴謹的方式重新表達成了最基礎的 AI 算子（GEMM，2D reduction，tensor transfer，等等）。這項技術可以讓為基礎 AI 算子編寫的算法直接被用到所有與之對應的復雜的 AI 算子上，而無需重寫算法。

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

Tensor Coordinate Transformation 和 Tile Tensor Operator 共同組成了 CK 的可復用的基礎模塊。

圖 1，使用 CK 的 Tensor Coordinate Transformation 基礎模塊將 convolution 算子表達成 GEMM 算子

圖 2，CK 的組成（下：可復用的基礎模塊；上：獨立算子與融合算子）

代碼結構

CK 庫結構分為四層，從下到上分別是：Templated Tile Operator，Templated Kernel and Invoker，Instantiated Kernel and Invoker 和 Client API【3】。每一層對應不同的開發者。

• AI 系統專家：“我需要一個后端方案提供高性能的獨立和融合算子讓我可以直接使用”。這個例子【4】里用的 Client API 和 Instantiated Kernel and Invoker 提供了預先實例化并編譯好的對象，以滿足這類開發者的需求。
• AI 系統專家：“我為一個開源的 AI 框架做最先進的圖優化工作。我需要一個能夠為圖優化所需的所有融合算子提供高性能 kernel 的后端方案。同時我也需要定制這些 kernel，所以像 “要么接受，要么棄用” 的黑盒解決方案不能滿足我的需求”。Templated Kernel and Invoker 層能滿足這類開發者。比如這個例子【5】中開發者可以自己使用 Templated Kernel and Invoker 層實例化出所需的 FP16 的 GEMM + Add + Add + FastGeLU 的 kernel。
• HPC 算法專家：“我的團隊為公司內部不斷迭代的 AI 模型開發高性能后端算法。我們團隊中有 HPC 算法專家，但我們仍然希望可以通過復用和改進硬件供應商提供的高度優化的源代碼來提高我們的生產力，并且讓我們的代碼可以被移植到未來的硬件構架上。我們希望可以不用通過與硬件提供商分享我們的代碼來做到這點”。Templated Tile Operator 層可以幫助到這一類開發者。比如這個代碼【6】中開發者使用 Templated Tile Operator 來實現 GEMM 的優化管線。

圖 3，CK 庫四層結構

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

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

AITemplate + CK 在 AMD Instinct? MI250 上取得了多個重要 AI 模型最先進的推理性能。CK 里大多數先進的融合算子的定義，都是在 AITemplate 團隊的遠見下推動的。許多融合算子的算法也是由 CK 和 AITemplate 團隊共同設計。

本文比較了幾個端到端模型在 AMD Instinct MI250 和同級別產品【8】的性能表現。本文中所有 AMD Instinct MI250 的 AI 模型的性能數據都是用 AITemplate【9】 + CK【10】取得的。

實驗

ResNet-50

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

BERT

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

下圖顯示了 AMD Instinct MI250 上的 AIT + CK 與 A100-PCIe-40GB 和 A100-DGX-80GB 上的 FasterTransformer v5.1.1 bug fix 【13】（FT）的 Bert Base 模型（uncased)的性能比較。當 Sequence 是 4096 時，FT 在 A100-PCIe-40GB 和 A100-DGX-80GB 上會在 Batch 32 時 GPU 內存溢出。因此，在 Sequence 是 4096 時，本文只顯示 Batch 16 的結果。結果顯示 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 圖片）的性能比較。結果顯示 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）的性能數據。當 Batch 是 1 時，在 MI250 上只有一個 GCD 被使用，而在 Batch 2，4，6 時，兩個 GCD 都被使用了。

Stable Diffusion 中的 UNet

不過本文還沒有關于使用 TensorRT 運行 Stable Diffusion 端到端模型的公開的信息。但這篇文章“Make stable diffusion 25% faster using TensorRT” 【14】說明了怎么使用 TensorRT 加速 Stable Diffusion 中的 UNet 模型。UNet 是 Stable Diffusion 中最重要最花時間的部分，因此 UNet 的性能大致反應了 Stable Diffusion 的性能。

下圖顯示了 AMD Instinct MI250 上的 AIT + CK 與 A100-PCIe-40GB 和 A100-DGX-80GB 上的 TensorRT v8.5.0.12（TRT）的 UNet 的性能比較。結果顯示 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