FlashDecoding++

论文地址：https://arxiv.org/pdf/2311.01282.pdf

背景

LLM推理中一个常用的优化技巧是 KV Cache，通过在计算时缓存 KV 的结果，来节约大部分的运算时间；所以 LLM 推理可以分为两个部分，分别是预测第一个 token 的 prefill 阶段和使用 KV Cache 的 decode 阶段

Flash Decoding 为了解决 Flash attention 在小 batch 推理时 GPU 利用率低的问题，引入了 key/value 的并行化，提出分块处理 attention output 后合并的思路。虽然这样增加了总体的计算量，但是在 GPU 利用率不高时提速明显。

本文参考 Flash Decoding 的思路，在此基础上提出了 Flash Decoding++ ，主要解决了下面三个问题：

1. 分块处理过程中同步更新 softmax 耗时过高
2. Decode阶段输入只有一个 token ，存在大量的 flat GEMM 操作，GPU利用率低
3. 静态的 kernel 无法应对动态的输入与硬件配置，需要找到一个方法平衡 内存 memory-bounded 和 计算 compute-bounded

论文效果

在 llama 7B 上测试了不同 GPU 的推理速度（主要关注 A100 ），相比于 TensorRT-LLM 和其他的推理框架有较大提升。从测试结果上来看比我目前使用的 vllm 要快 20%

具体方法

Partial softmax with unified max value

Flash Decoding 合并 softmax 结果的操作如下：
$$\begin{array}{c}m(x)=max(m(x_1),m(x_2))\\ f(x_1)=e^{m(x_1)-m(x)}f(x_1)\\ f(x_2)=e^{m(x_2)-m(x)}f(x_2)\\ l(x)=f(x_1)+f(x_2)\\ softmax(x_1,x_2)=(\frac{f(x_1)}{l(x)},\frac{f(x_2)}{l(x)})\end{array}$$

Flash Decoding++为了节约同步 m(x) 的时间，选择使用统一的最大值 ϕ 来代替。softmax 公式替换为
$$softmax(x_i)=\frac{e^{x_i-\varphi }}{{\sum }_j{e}^{x_j-\varphi }}$$
设置 ϕ 的目的是为了对 softmax 后的结果进行放缩，结果太大可能会溢出 float32，太小又可能出现精度问题影响效果，所以 ϕ 应满足 $a<x_i-\varphi <b。$在实际推理过程中我们需要根据不同的模型和想要的精度设置对应的$a、\varphi 、b$。以 llama-7B 为例，作者发现 99.99%的$x_i$位于[-16.8,6.5] ，ϕ 就可以在这个区间内进行选择。

如果块间差距过大可能会导致不存在 ϕ 使得 $a<x_i-\varphi <b$ 都被满足，这时该策略会选择退化成 Flash Decoding 的动态合并 softmax，确保不会出现误差

Flat GEMM Optimization with Double Buffering

假设 GEMM 中两个相乘的矩阵大小分别为 M * K 和 K * N，同时每个 GEMM Tile 会对 K * N 的矩阵进行分块，每块大小为 Bn * Bk （不足则进行填充），那么每个 GEMM Tile的计算量为 2MBn*Bk，内存访问量为 M * Bk + Bn * Bk，共有 N * K / Bn * Bk 块。算上把乘法结果写入的内存访问，整个 GEMM 过程中计算与内存的比值为
$$\begin{array}{rl}\frac{2\ast M\ast B_n\ast B_k\ast \frac{N\ast K}{B_n\ast B_k}}{(M\ast B_k+B_n\ast B_k)\ast \frac{N\ast K}{B_n\ast B_k}+M\ast N}& =\frac{2\ast M\ast K}{K+\frac{M\ast K}{B_n}+M}\end{array}$$
由于 GEMM 运算的并行度为 N/Bn ，因此计算与内存的比值与 Bn 成正比，并行度与 Bn 成反比，这就让计算和内存成为同时制约 GEMM 速度的两个要素。

Tensor Core在进行 GEMM 时，会将 M padding 到 64 去减少内存访问的延时，但这会浪费大量的计算资源。本文选择将这个数值调整到 8，同时为了解决内存访问的问题，增加了 Double Buffering ，将共享的内存分成两个 buffer ，一个用于当前的 GEMM 计算，另一个加载下一次 GEMM 需要的数据。

Heuristic Dataflow with Hardware Resource Adaption

第三点其实是对第二点的补充，作者说影响 LLM 推理的因素有很多，像第二点优化 Flat GEMM的操作实测下来也不一定是最优的，还需要具体情况具体分析，但幸运的是 LLM 中 GEMM 的结构都比较类似，以 llama-7b 为例，一共只有4种情况：

那就好办很多了，直接每个都测一遍。于是作者测了 FaseGEMV、flat GEMM（本文）、CUTLASS 三种方法在不同 M N K 下的速度，并且在实际推理中动态的切换：

结论

通篇介绍的三种方法并没有对整个推理过程做大的改动，更像是针对 llama-7B 做了特殊的代码优化，并且都存在使用条件，参数的设置都需要大量实验来确定，引入了不少选择和判断。直观感觉实用性不大，这点加速效果并不难达到，针对任何一个主流的推理框架，为某个模型单独的设计 kernel 和并行化策略应该都有这样的效果。但其中工程性的手段值得借鉴，在不改变 decoding 结构的前提下，能做的优化手段确实非常有限了。