SmoothQuant模型在AMD Instinct MI300X上使用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 计算中的计算能力。实例的实现包含两个阶段：

• 模板参数定义；

• 模板化内核实例化和运行。

模板参数定义

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

图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++ 结构函数绑定来实现。

上述 PassThrough 表示不对目标矩阵进行任何操作。`CDEElementOp` 确定应用于 CShuffle 输出和矩阵 D 的操作。以下绑定结构 AddRelu 显示了对 CShuffle 输出和矩阵 D 进行相加、向相加结果应用 ReLU（Rectified Linear Unit）操作的示例。 然后，将结果传递给矩阵 E。

struct AddRelu
{
    __host__ __device__ void operator()(ck::half_t& e, const ck::half_t& c, const ck::half_t& d) const
    {
        const ck::half_t x = c + d;
        e = x > 0 ? x : 0;
    }
};

可调节参数

CK 实例包括一系列可调节的模板参数，以控制工作负载的并行粒度，从而在不同的硬件平台上实现负载均衡。

这些参数包括 Block Size、M/N/K 每块、M/N 每个 XDL、AK1、BK1 等。

• Block Size 决定线程块中的线程数量。

• M/N/K  每块 决定每个线程块负责计算的矩形区域的大小。

• M/N 每个 XDL 指的是每个 wavefront 基础上 Instinct 加速器矩阵融合乘加（MFMA）指令操作的 M/N 大小。

• A/B K1 与数据类型相关。 它可以是从 1 到 KPerBlock 的任何值。为了实现最佳的加载/存储性能，建议每次加载 128 位。此外，A/B 加载参数必须相应更改以匹配 A/B K1 值，否则会导致编译错误。

在不同硬件平台上实现计算负载均衡的条件可能会有所不同。

实例化并运行模板化内核

在确定模板参数后，我们使用实际参数来实例化内核。可以执行以下操作之一：

• 使用 CK 的自定义结构 DeviceMem 中的 GetDeviceBuffer 来传递需要计算的矩阵的元素值。

• 通过 hipMalloc 分配设备缓冲区。确保设备缓冲区大小能够容纳矩阵的大小。

• 如果需要计算的矩阵是 Tensor 类型，通过 Tensor 对象中的 data_ptr 方法传递矩阵元素。

输入矩阵的行、列和步幅信息也会传递给实例。对于批量 GEMM，必须传入额外的批次计数和批次步幅值。用于前后处理的额外操作也会使用实际参数传递；例如，GEMM 缩放操作的 α 和 β。然后，通过调用器来启动实例化的内核，如图 2 所示。

图 2：模板化内核启动包括内核实例化，通过传递实际应用参数来创建参数，创建调用器，通过调用器运行实例。

开发用于SmoothQuant模型的融合INT8内核

 (SQ) 是一种量化算法，能够实现权重和激活量的INT8量化，用于LLM中的所有矩阵乘法。下表显示了在Instinct加速器上加速SQ模型推理所需的GPU内核功能。

操作流程分析

以下部分讨论了`Linear_ReLU_ABDE_I8`的操作流程分析。表1中的其余包装器的分析方法类似。

该流程的第一个操作是进行输入矩阵A和B的乘法运算。结果矩阵C与α相乘，以获得T1。同时，对D元素执行缩放操作以获得T2。随后，执行T1和T2的矩阵加法、使用ReLU进行元素激活计算以及逐元素舍入。这些生成E1、E2和E的操作被封装并由CK中的用户定义模板函数在单个GPU内核中完成。此模板函数在编译阶段直接集成到基本实例中，因此这些步骤可以融合在一个内核中完成。

图3：操作流程。

CK库包含许多实现不同功能的基本实例。熟悉各种CK实例的名称并确定它们是否满足目标功能需求。

其次，考虑输入数据的格式是否满足实际计算需求。对于SQ模型，矩阵计算应用8位整数数据格式（INT8）。

第三，考虑实现CK实例的平台。以`xdl`结尾的实例仅在AMD Instinct加速器上编译后运行，无法在Radeon系列GPU上运行。这是因为这些基本实例的底层设备特定指令集是针对特定设备实现的。

在这里，我们使用作为基本实例来实现之前表中的功能。

图4：使用 ‘DeviceBatchedGemmMultiD_Xdl’ 实例作为根。

DeviceBatchedGemmMultiD_Xdl 实例直接通过使用适当的输入和输出数据精度类型实现了批量GEMM BMM_ABE_I8 和 BMM_AB_I8_E_F32 内核。

基于这两个批量GEMM内核，GEMM内核 Linear_ABDE_I8 和 Linear_AB_I8_DE_F32 可以通过对其输入2D张量进行非压缩为3D张量来实现。然后，在返回之前，将根实例生成的3D输出张量压缩回2D输出张量。

例如，在将A$(M, K)$分配到根实例之前，将A $(M, K)$ 非压缩为A$(1, M, K)$，在根实例返回计算结果后，将E$(1, M, N)$ 压缩到$(M, N)$。`Linear_ReLU_ABDE_I8` 通过在 Linear_ABDE_I8的结果输出上添加一个ReLU操作来实现。

开发完整的函数

SQ量化模型的推理依赖于使用PyTorch和Transformer库。torch中的tensor类型用来表示矩阵和向量，因此CK(C++库)中的数据类型需要替换为torch::tensor类型。输入和输出矩阵的dtype应该是tensor类型。

在GEMM（广义矩阵乘法）中，A和B输入是二维矩阵，而选择的基础CK实例需要三维矩阵输入。因此，在将这些矩阵传递给实例之前，我们必须使用tensor的unsqueeze()方法将输入的二维tensor转换为三维tensor。对于上述表中的批量GEMM，忽略此步骤。

// 函数输入和输出 
torch::Tensor linear_relu_abde_i8(
			torch::Tensor A_,
			torch::Tensor B_,
			torch::Tensor D_,
			float alpha,
			float beta)
{
  // 将torch::Tensor A_(M, K) 转换为 torch::Tensor A (1, M, K) 
  auto A = A_.unsqueeze(0);

  // 将torch::Tensor B_(K, N) 转换为 torch::Tensor B (1, K, N) 
  auto B = B_.unsqueeze(0);
...

如以下代码块所示，我们使用输入tensor的size值来获取M, N和K值。此步长大小信息用于重塑输入向量D并分配tensor E的存储空间。步长反映了内存中连续元素的精确大小，它们作为GPU内核使用的重要参数传递给基础实例。

  // 返回第0维的批次计数
  int batch_count = A.size(0);

  // 返回第1和第2维的M, N, K值
  int M = A.size(1);
  int N = B.size(1);
  int K = A.size(2);

  // 初始化A, B, D和E的步长大小
  int stride_A = K;
  int stride_B = K;
  int stride_D0 = N;
  int stride_E = N;

  // 初始化批量A, B, D和E的步长大小
  long long int batch_stride_A = M * K;
  long long int batch_stride_B = K * N;
  long long int batch_stride_D0 = M * N;
  long long int batch_stride_E = M * N;

  // 将2-D tensor转换为3-D tensor	
  auto D = D_.view({1,-1}).repeat({M, 1});

  // 分配E的内存
  auto E = torch::empty(	{batch_count, M, N}, 
			torch::dtype(torch::kInt8).device(A.device()));

在以下代码块中，ADataType, BDataType和D0DataType用于表示输入tensors A, B和D的数据精度。EDataType用于表示输出tensor E的数据精度。这些参数应指定为I8数据格式（8-bit integer）以满足内核设计要求。

AccDataType决定用于表示A和B元素的乘加结果的数据精度。通常，应用更大的范围数据类型来存储A和B的乘加结果以避免结果溢出；此处应用I32。CShuffleDataType I32数据类型表示乘加结果继续以I32数据格式存储在LDS中。所有这些通过以下代码块实现。

  // 数据精度 
  using ADataType        = I8;
  using BDataType        = I8;
  using AccDataType      = I32;
  using CShuffleDataType = I32;
  using D0DataType 	      = I8;
  using DsDataType       = ck::Tuple<D0DataType>;
  using EDataType        = I8;

按各种线性代数库的惯例，行主要（row-major）和列主要（column-major）顺序用于表示以线性存储方式存储矩阵的方式。将矩阵B指定为列主要的好处是，当B中的一列与A中的一行相乘时，所有相关矩阵元素连续存储在GPU全局内存中，这可以帮助GPU实现数据一致性访问以提高访问性能。

  // 指定tensor顺序
  using ALayout  = RowMajor;
  using BLayout  = ColumnMajor;
  using D0Layout = RowMajor;
  using DsLayout = ck::Tuple<D0Layout>;
  using ELayout  = RowMajor;

在CK中，PassThrough是一个结构体表示是否对绑定的tensor应用操作。为了融合 中介绍的E1, E2和E之间的操作，我们定义一个自定义C++结构体ScaleScaleAddRelu，并将其绑定到CDEElementOp。它决定了哪些操作会应用到CShuffle（A×B结果）、张量D、α和β上。

  // No operations bound to the elements of A and B 
  using AElementOp   = PassThrough;
  using BElementOp   = PassThrough;

  // Operations bound to the elements of C, D and E
  using CDEElementOp = ScaleScaleAddRelu;

在绑定结构体中，`operator()`执行一个加法操作，使`CShuffle`与矩阵`D`相加，对相加结果进行ReLU操作，然后对输出元素进行四舍五入操作。然后将结果返回给`E`。

struct ScaleScaleAddRelu {
// 此模板函数执行缩放、加法、ReLU和四舍五入操作
  template <>
  __host__ __device__ constexpr void
  operator()<I8, I32, I8>(I8& e, const I32& c, const I8& d) const
  {
      // 用alpha缩放AxB结果
      const F32 c_scale = ck::type_convert<F32>(c) * alpha;

      // 用beta缩放D
      const F32 d_scale = ck::type_convert<F32>(d) * beta;

      // 执行加法操作
      F32 temp = c_scale + d_scale;
      
      // 执行ReLU操作
      temp = temp > 0 ? temp : 0;

      // 执行四舍五入操作
      temp = temp > 127 ? 127 : temp;
      
      // 返回给E
      e = ck::type_convert<I8>(temp);
  }
    
  F32 alpha;
  F32 beta;
};

原始的输入tensors需要扩展以满足GPU的基于块(tile-based)并行处理。

static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::MNKPadding;

目标基础实例的模板参数使用上述参数进行初始化，并包括默认的可调参数。关于特定的调优方法，请参阅。

using DeviceOpInstance = ck::tensor_operation::device::DeviceBatchedGemmMultiD_Xdl< 
    // Tensor layout
    ALayout, BLayout, DsLayout, ELayout, 
    // Tensor data type
    ADataType, BDataType, AccDataType, CShuffleDataType, DsDataType, EDataType,  
    // Tensor operation
    AElementOp,  BElementOp, CDEElementOp,  
    // Padding strategy  
    GemmDefault,
    // Tunable parameters        
    tunable parameters>;

 auto A_ref = A.data_ptr<ADataType>();
 auto B_ref = B.data_ptr<BDataType>();
 auto D0_ref = D.data_ptr<D0DataType>();
 auto E_ref = E.data_ptr<EDataType>();

然后使用实际参数初始化和运行基础实例:

 auto device_op    = DeviceOpInstance{};
 auto invoker = device_op.MakeInvoker();
 auto argument = device_op.MakeArgument(
		A_ref, B_ref, {D0_ref}, E_ref,
		M, N, K,
		batch_count,
		stride_A,	stride_B,	{stride_D0}, stride_E,
		batch_stride_A, batch_stride_B, {batch_stride_D0}, batch_stride_E,
		AElementOp{}, BElementOp{}, CDEElementOp{alpha, beta});

invoker.Run(argument, StreamConfig{nullptr, 0});

基础实例的输出是一个计算出的批量矩阵`E` (batch, M, N)。在返回之前，如果需要常规的GEMM结果，需要将其转换为二维矩阵。

// 将 (1, M, N) 转换为 (M, N)
return E.squeeze(0);

将其绑定到Python

由于这些函数是用 C++ 和 torch::Tensor 编写的，您可以使用 pybind11 绑定这些函数并将其作为 Python 模块导入。以下示例中，用于公开表中的函数的绑定代码仅需几行。

#include <torch/extension.h>

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
  m.def("linear_ab_i8_de_f32", &linear_ab_i8_de_f32);
  m.def("linear_relu_abde_i8", &linear_relu_abde_i8);
  m.def("linear_abde_i8", &linear_abde_i8);
  m.def("bmm_abe_i8", &bmm_abe_i8);
  m.def("bmm_ab_i8_e_f32", &bmm_ab_i8_e_f32);
}

通过编写一个 setup.py 脚本来编译 C++ 代码以构建 C++ 扩展，该脚本使用 setuptools 进行编译。以下是 setup.py 脚本的参考实现。

import os
from setuptools import setup, find_packages
from torch.utils import cpp_extension
from torch.utils.cpp_extension import BuildExtension

os.environ["CC"] = "hipcc"
os.environ["CXX"] = "hipcc"

sources = [
    'torch_int/kernels/linear.cpp',
    'torch_int/kernels/bmm.cpp',
    'torch_int/kernels/pybind.cpp', 
]

include_dirs = ['torch_int/kernels/include']
extra_link_args = ['libutility.a']
extra_compile_args = ['-O3','-DNDEBUG', '-std=c++17', '--offload-arch=gfx942', '-DCK_ENABLE_INT8', '-D__HIP_PLATFORM_AMD__=1']

setup(
    name='torch_int',
    ext_modules=[
        cpp_extension.CUDAExtension(
            name='torch_int.rocm',
            sources=sources,
            include_dirs=include_dirs,
            extra_link_args=extra_link_args,
            extra_compile_args=extra_compile_args
            ),
    ],
    cmdclass={
        'build_ext': BuildExtension.with_options(use_ninja=False)
    },
    packages=find_packages(
        exclude=['notebook', 'scripts', 'tests']),
)

运行 python setup.py install 来构建和安装扩展。它看起来应该像图 5 所示：

图 5：编译和安装 INT8 内核。

INT8模型推理和性能

图6展示了在MI300X GPU上运行SmoothQuant模型的实现架构，其中(a)展示了目标模型的解码器层组成部分，(b)展示了解码器层组件的主要实现类，(c)表示由CK实例实现的底层GPU内核。

图6：在AMD MI300X加速器上运行SmoothQuant模型的实现架构。

对于目标 ，每个解码器层包含三个主要组件：注意力计算、层归一化和全连接层中的线性变换。这些组件的相应实现类是：

• Int8OPTAttention

• W8A8B8O8LinearReLU

• W8A8BF32OF32Linear

这些类的底层实现逻辑将利用前表中的函数。需要注意的是，对于该示例，`LayerNormQ`模块是通过torch自带模块实现的。

测试环境：用于测试的硬件平台配备了256个AMD EPYC 9534 64核处理器、8个AMD Instinct MI300X加速器和1.5TB内存。测试在一个从Docker Hub公开提供的Docker镜像中完成: 

测试的模型是OPT-1.3B、2.7B、6.7B和13B FP16模型，相应的SmoothQuant INT8 OPT模型来自Hugging Face。

需要注意的是，由于基础实例的可调参数使用了默认值，因此INT8内核的性能是不优化的。

图7显示了在单个MI300X加速器上原始FP16模型和SmoothQuant量化的INT8模型之间的性能比较。SmoothQuant量化模型的GPU内存占用显著减少。同时还表明，对于所有SmoothQuant量化的OPT模型，单样本推理延迟显著降低（如(b)所示）。值得注意的是，基于CK实例的INT8内核性能随着模型尺寸的增加而稳步提升。

图7：在单个MI300X加速器上原始FP16模型和SmoothQuant量化INT8模型之间的性能比较。

为了比较原始FP16模型和INT8模型的准确性，使用了LAMBADA数据集验证集中前1,000个样本进行评估。我们采用了与中引入的相同的Last Token Prediction Accuracy方法作为评估指标。比较结果如表2所示。

结论

CK提供了丰富的模板参数集，可以为不同的应用场景生成灵活的加速计算内核。

CK支持AMD Instinct GPU的多种指令集、操作融合和不同数据精度。它的可组合性帮助用户快速构建操作性能验证。

借助CK，您可以在不同的AMD加速平台上构建更高效、更灵活和性能更好的AI应用。