圖最佳化在降低 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 的可複用的基礎模組。
程式碼結構
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 的最佳化管線。
基於 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
原文連結：https://community.amd.com/t5/instinct-accelerators/amd-composable-kernel-library-efficient-fused-kernels-for-ai/ba-p/553224