图优化在降低 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
