AI推理加速原理解析与工程实践分享

GEEK TALK

如果我们从端到端的视角再来分析下整个 AI 推理过程，会发现这两类用户的痛点目前没有得到很好的解决。

用户对 GPU 的使用初始于业务系统，用户根据业务需求搭建模型，并为最终模型的效果负责。

业务系统构建完成后，会从资源管理系统中申请资源，而资源管理器则会将 GPU 卡分配给业务系统，这个管理器只会为资源分配率负责，而不会关心资源分配后的业务使用效率。

用户在申请到资源后，会通过 AI 框架执行模型的计算过程。AI 框架更专注为用户提供易用的模型构建接口，而不会为业务的推理效率和资源利用率负责。

最后 AI 框架在使用异构硬件算力时，只能使用基础的加速包或工具，而不会专门结合业务特点进行优化。总的来看，整个过程中没有专门的工具为 GPU 算力的利用效率负责。

为此，我们需要 AI 推理加速，针对用户训练好的模型进行针对性的加速，缩短业务推理时间，同时提升资源利用率。

为了系统性的分析和进行推理加速方案，我们首先需要能够定义推理加速的优化目标。为此我们先简单回顾下 GPU 的硬件架构和执行模式。

从硬件上看，GPU 的强大算力来源于多 SM 处理器，每个 SM 中包含多个 ALU 计算单元和专有的 Tensor Core 处理单元。对 GPU 来说，当所有 SM 上的所有计算单元都在进行计算时，我们认为其算力得到了充分的发挥。

GPU 无法自己独立工作，其工作任务还是由 CPU 进行触发的。

整体的工作流程可以看做是 CPU 将需要执行的计算任务异步的交给 GPU，GPU 拿到任务后，会将 Kernel 调度到相应的 SM 上，而 SM 内部的线程则会按照任务的描述进行执行。当 CPU 不发起新的任务时，则 GPU 的处理单元就处在空闲状态。

通过这两个层面的分析，我们知道要想将 GPU 的算力充分发挥，其核心就是保持GPU 上有任务，同时对单个 GPU 计算任务，使其可以尽量充分的用好 SM 处理器。

针对这两个层面的使用情况，NVIDIA 提供了相应的牵引指标 GPU 利用率和 SM 利用率：GPU 利用率被定义为在采样间隔内，GPU 上有任务在执行的时间。而 SM 利用率则被定义为在采样间隔内，每个 SM 有 warp 活跃时间的平均值。

我们可以通过 2 个 case 来比较下这两个指标的差异。在下图的 case 1 中，由于 CPU 异步发射任务到 GPU 上，GPU 很快就处理完了，于是就出现了等待下一个任务的空隙。在这种情况下，按照定义，GPU 利用率比较低，SM 利用率也相对较低。两个指标都能反应这种情况没有充分利用 GPU 资源。

针对下图的第二个 case，对于某一个 Kernel 来说，由于计算实现、参数配置等一系列问题，它只使用了 1 个 SM 处理器，而剩下的 3 个 SM 处理器（假设只有 4 个 SM 处理器）空闲。

对于这种情况，由于 GPU 上有 Kernel 在执行，在这个时间段内 GPU 利用率仍然是 100%，但 SM 利用率只有 25%。可以看到这种场景下，SM 利用率可以反应计算任务效率不高的问题，而 GPU 利用率则无法反应此类问题。

我们认为 SM 利用率可以更精细的反应 GPU 算力发挥情况，所以把 SM 利用率当做 AI 推理加速的牵引指标。

有了牵引指标，结合模型执行的流程，我们就可以在逻辑上将优化方案分为三类：

第一类优化是模型精简类，即在模型真正执行之前就对模型的计算量进行精简，从而提升推理速度。这部分业界常见的优化方向包括量化、减枝、蒸馏和 NAS 等；

第二类和第三类则是当模型已经交由推理引擎在 GPU 上执行时，如何更好的提升 GPU 的利用效率。

我们由 SM 利用率公式做一个近似可以导出这两类优化方案，分别是尽可能让 GPU 上有计算任务和单个计算任务在 GPU 上执行效率更高。这两类优化方案常见的手段分别是算子融合和单算子优化。

• 离线量化是指在模型训练完成后，离线的对计算算子进行量化，这种方案通常易用性较好，对算法开发人员几乎透明，但对模型精度会有一定损失；
• 量化训练则是在模型训练过程中就显示插入量化相关的操作，这样通常会有更好的精度，但需要算法开发同学准备相关数据。

而减枝则是通过将模型中对结果影响较小的一些计算进行移除，从而降低计算量。

蒸馏则通常是将一个复杂的大模型通过降维的知识传递层，将大模型中的复杂计算，减少为效果相当的更小规模的计算，从而实现降低计算量，提升推理效率的效果。下图中是百度文心 3.0 大模型知识蒸馏的过程。

这些模型精简的方案，由于涉及到对精度的影响，通常需要算法工程师介入，协同优化。

第二类优化则是算子融合，算子融合顾名思义是指将多个计算算子合并成一个大算子的过程。

例如对于 BERT Base 这个模型，经过 PyTorch 原生 jit 编译生成的 TorchScript 图中有 800 多个小算子，这些小算子会带来 2 类问题：一是这些算子通常执行过程较短，因此会造成大量的 GPU 空闲时间；二是由于不同的任务之间还有数据的依赖，因此也会带来额外的访存开销。

目前针对算子融合，业界通常会采用手写或自动化生成的方式发现可融合的模式，并提供融合后的算子。例如针对 Transformer 结构，NVIDIA 开发了 FasterTransformer 这个库，其中包括针对多种 Transformer 类结构模型的具体融合算子。

例如下图中，针对 batch GEMM Q x K、Softmax、batch GEMM for QK x V 和长尾的 transpose 等操作，FasterTransformer 提供 Fused Multi-head Attention 等融合后算子。BERT Base 在经过 FasterTransformer 的算子融合优化后，数量可以降到 100 个左右。

第三类优化则是单算子优化。单算子优化是根据单个算子，结合计算模式和硬件架构特点，调整 GPU 核函数的实现方法，从而提升具体算子的执行效率。单算子优化中，最经典的例子就是对通用矩阵乘（GEMM）的优化。

下图是 NVIDIA Cutlass 库对 GEMM 操作的抽象：

Cutlass 结合 GPU Global Memory、Shared Memory、Register File 这几层存储架构和 block、warp、thread 和 tensor core 这几层计算抽象，设计了一系列计算模板，并提供相关可优化参数（切分大小等），方便开发者开发高性能的 GEMM 实现。

以上就是业界常见的几类 AI 推理加速的方法和业界的一些解决方案。接下来我们重点介绍下 AIAK-Inference 是如何站在巨人的肩膀上，提供性能更好的推理加速方案。

GEEK TALK

GEEK TALK