SmoothQuant模型在AMD Instinct MI300X上使用Composable Kernel进行推理

SmoothQuant model inference on AMD Instinct MI300X using Composable Kernel

AMD ROCm™ Composable Kernel (CK)库为编写机器学习工作负载的性能关键内核提供了一种编程模型。它通过C++模板在编译阶段生成通用内核，允许开发人员在不同数据精度上实现操作融合。
本文概述了基于CK设计示例`03_gemm_bias_relu`的CK通用矩阵乘法 (GEMM) 内核的高级概述。它还概述了构建和运行内核的步骤。此外，本文还详细介绍了使用CK在AMD Instinct MI300X加速器上运行SmoothQuant量化的INT8模型的实现。

高层概述：一个 CK GEMM 实例

GEMM 是线性代数、机器学习和深度神经网络中的一个基本模块。它被定义为操作：E=α×(A×B)+β×(D)，其中 A 和 B 是矩阵输入，α 和 β 是标量输入，D 是一个预先存在的矩阵。以全连接层中常用的线性变换为例，这些术语分别对应输入激活 (A)、权重 (B)、偏置 (D) 和输出 (E)。(DeviceGemmMultipleD_Xdl_CShuffle)` 结构体作为基本实例，探索 AMD Instinct 加速器在 GEMM 计算中的计算能力。实例的实现包含两个阶段：

模板参数定义；
模板化内核实例化和运行。

模板参数定义

实例的模板参数分为四种参数类型：

../../_images/template_parameters.jpg

图1: 所选GEMM核的模板参数被分类为四组。在运行实例之前，应正确定义这些模板参数组。

矩阵数据精度

A、B、D 和 E 被定义为半精度浮点数据类型。矩阵 A 和 B 的乘加结果与已有的矩阵 D（半精度）相加，最终的 GEMM 结果也是半精度浮点。

using ADataType        = F16;
using BDataType        = F16;
using AccDataType      = F32;
using CShuffleDataType = F16;
using DDataType        = F16;
using EDataType        = F16;

ADataType 和 BDataType 表示输入矩阵 A 和 B 的数据精度。`AccDataType` 确定用于表示 A 和 B 元素乘加结果的数据精度。这些结果存储在局部数据共享（LDS）中的 CShuffle 模块中，LDS 是一种低延迟、高带宽的显式寻址内存，用于在工作组内同步 LDS 以供后续使用。

CShuffleDataType 表示 LDS 中 CShuffle 的数据精度。

DDataType 表示存储在 GPU 全局内存中的已有 D 矩阵的数据精度，而 EDataType 表示最终输出的数据精度。CK 内核支持一种融合策略，因此 CShuffle 可以与相同 GPU 内核中的单个已有矩阵相加，以获得更好的性能。

矩阵数据布局

using ALayout = Row;
using BLayout = Col;
using DLayout = Row;
using ELayout = Row;

根据各种线性代数库的惯例，CK 假设输入矩阵 A 是一个 M x K 矩阵，这意味着矩阵有 M 行和 K 列。同样地，矩阵 B 被假设为 K x N，这意味着它有 K 行和 N 列。在计算中，行优先顺序和列优先顺序是常用的用于将矩阵存储在线性存储中的方式。在理解了矩阵存储模式之后，可以根据这些矩阵的存储顺序应用底层优化的内存访问方式，以实现更好的性能。

矩阵元素操作

using AElementOp   = PassThrough;
using BElementOp   = PassThrough;
using CDEElementOp = AddRelu;

CK 支持在计算矩阵乘法（GEMM）之前对矩阵进行预处理，即 C = AElementOp(A) * BElementOp(B)。同样，它也支持以相同方式处理 GEMM 结果后的操作，即 E = CDEElementOp(C, D)。

AElementOp 和 BElementOp 确定在进行 GEMM 之前分别应用于矩阵 A 和 B 的操作，这通过将操作与 C++ 结构函数绑定来实现。

上述