[论文分享]LLM推理加速——FLASHDECODING++

FLASHDECODING++: FASTER LARGE LANGUAGE MODEL INFERENCE ON GPUS:
论文地址：https://arxiv.org/pdf/2311.01282.pdf

1 背景

大型语言模型（LLM）在各个领域的应用日益增多，其推理性能对于大规模应用至关重要。LLM为多个layer叠加的结构，每个layer可分为线性GEMM(通用矩阵乘法)操作(例如:
K, Q, V, O权重投影和前馈)和注意力/softmax计算。

加速LLM推理的挑战主要包括以下几个方面：

1. 同步部分softmax更新：由于attention-softmax需要统计整个向量的max，计算量不大，但较耗时，现有优化方案采用并行化处理。为了提高并行性，以前的设计将注意力矩阵划分为更小的块，并且还拆分行以计算部分softmax结果。当计算新的部分softmax结果时，采用同步softmax操作更新先前的部分softmax结果。据统计在1024输入长度的NVIDIA Tesla A100 GPU上的分析，这种同步的部分softmax更新占Llama2-7B推理的注意力计算的18.8%。

2. 未充分利用flat GEMM计算：在LLM推理中，矩阵执行GEMM操作的形状是扁平的，这导致未充分利用计算和超过50%的性能损失（在之前的设计中，如cuBLAS、CUTLASS等）。

3. 性能损失由于静态数据流：较小的批处理大小使LLM推理的解码阶段具有内存限制，而较大的批处理大小使其具有计算限制。在LLM推理中，对于不同形状的gem，单个静态数据流可能导致50.25%的性能损失。。

为解决这些问题，研究团队提出了一种名为FlashDecoding++的更快的GPU实现方法，通过解决不同输入的最大值问题、优化flat GEMM操作和根据输入动态性和硬件配置进行启发式数据流优化，实现了对主流LMs的加速。

2 方法

2.1 异步并行Softmax统一最大值

当前针对softmax的优化策略是通过分块并行加速计算，但该过程需要同步各分块的max值，导致注意力操作的开销约为20%。
由softmax的计算公式可知，指数中减去ϕ值，并不会影响计算结果。
因此作者提出去设置一个全局max，达到不需要其他部分softmax计算结果的情况下计算每个部分softmax结果。
最终问题转化为：如何设置一个全局max，取代分块max的同步过程，同时确保计算过程不溢出。
作者通过统计分析发现，Softmax的输入的数值分布会集中在一定范围内，所以设置一个全局max，在出现溢出时退化回常规的的softmax计算。
假设每个xi的a < xi−φ < b，以确保精度并避免溢出。然后，单独处理部分softmax操作。然而，当xi−φ≤a或xi−φ≥b时，对于xi所在的向量x，终止异步部分softmax计算。然后使用同步部分softmax方案(在FlashAttention和FlashDecoding中使用)重新计算softmax。根据图5所示的统计数据，这种重新计算方案避免了溢出，同时引入了可以忽略不计的开销。

2.2 双缓冲的“矮胖”矩阵乘优化

在不失一般性的前提下，GEMM运算可以用M、N、K来表示，其中两个相乘矩阵的大小分别为M × K和K × N。
对于GEMM，以前的LLM推理引擎利用Tensor Core使用cuBLAS和CUTLASS等库来加速这些操作。尽管现代Tensor Core架构处理的GEMM M=8，但这些库通常将M维度进行til-size设置为64，即padding到64以隐藏内存延迟。
在llm的decoding阶段，由于输入token为1，该过程的矩阵乘法主要由GEMV (batch size=1)或flat GEMM (batch size>1)操作组成。此时仍然使用tensorcore进行gemm，会将M填充到64，会导致实际的计算利用率不高。
作者提出把矩阵成的tile size设置为8，以降低padding的资源浪费。同时隐藏内存访问延迟，作者引入双缓冲技术。
用于“矮胖”GEMM操作。作者在共享内存中分配两个单独的缓冲区。一个缓冲区中的块执行GEMM操作，而另一个缓冲区为下一个GEMM操作加载一个新块。因此，计算和内存访问是重叠的。在我们的实践中，我们在N较大时应用这种技术。

2.3 基于硬件资源自适应的启发式数据流

根据硬件算力和GEMM的矩阵维度自适应选择使用CUDA Core或者Tensor Core的矩阵乘实现方式。

对一个固定的llm，四个全连接（proj）操作对应的矩阵乘维度是确定的。
但GEMM/GEMV的不同[k、n]在计算性能上分别对应不同的实现方式：FASTGEMV，flatGEMM（2中的实现），CUTLASS
对于运行时LLM推理，FlashDecoding++在M < M1时采用使用CUDA Core的ImplA，在M1≤M < M2/M2≤M时采用使用Tensor Core的ImplB/ImplC。注意决策流是离线执行的，不影响运行时LLM推理的性能。

3 实验

作者选择的硬件平台和模型如下表：

并使用业内主流的优化框架vllm、tensorrt-llm、deepspeed、ppl等作为对照组，实际结果如下：

4 结论

FlashDecoding++提出了具有统一最大值的异步softmax、具有双缓冲的平面GEMM优化和具有硬件资源自适应的启发式数据流三种新颖设计。FlashDecoding++在主流llm和硬件上均取得了有效的加速效果：
与hug-Face实现的相比，在NVIDIA和AMD上实现了高达4.86倍和2.18倍的gpu加速，
与最先进的LLM推理引擎FlashDecoding相比，FlashDecoding在各种LLM上也实现了平均1.37倍的加速。

5 思考

论文中作者使用了三个技巧来加速llm推理，但是在实际应用中的问题也很明显：
异步并行Softmax统一最大值，需要根据设置一个max取代max计算，防止softmax计算过程中的溢出，如果根据经验拍一个max，有效性很难保障所以这个方法很难应用到实际生产中；
shuangbuffer技术在gemm中也是广泛使用的策略；
同时作者为了选择softmax值而去做的统计分析，对于量化工作具有一定借鉴意义。