大模型推理加速之Flash Decoding：更小子任务提升并行度

Flash Decoding（FD)是FlashAttention（FA）针对推理场景的改进版本，它的设计思想在2023.10.13发布在如下PyTorch官方blog。如果大家了解FA原理的话会觉得FD改进非常自然。

Flash-Decoding for long-context inference

关于FlashAttention V1和V2如何加速LLM训练的技术，可以看我之前的文章：

方佳瑞：大模型训练加速之FlashAttention系列：爆款工作背后的产品观

大语言模型LLM的推理和训练任务的计算特性不同，大家可以看图1和图2的绿色框的部分，也就是FA和FD加速的范围，其中黑线上标注了Activation Tensor的维度。

2. 推理分为Prefilling+Decoding阶段，Prefilling完成prompt sequence length长度token一次forward计算，Decoding阶段一个token一个token地生成。图2和图1区别，在于图2红框中Q矩阵大小为d/h*1，是训练的 1/l ，其中 l 是序列长度。

FA面向训练设计的Multihead Attention的优化方案不适用于推理Decoding阶段，因为query sequence length=1，导致Decoding可并行的维度减少。

先回顾FA V2中对MHA的并行方案，下图中展示一个head的Attention三个矩阵乘法如何做Tiling，图中N是sequence length维度，d是hidden dimension。FA沿着Q、K、V的N切分成若干块，将Outer Loop切分分配给CUDA thread block，每个thread block完成Q的一块和所有K、V块的运算，结果存在一块O中。

下图是Flash Decoding博客中的图，它展示了一个thread block中的运算流程，用一个固定Q矩阵块顺序地和不同K，V块进行运算。注意，沿着outer loop切分任务，可以构成很多个这样的thread block的子任务，它们会被GPU硬件调度到不同的SM上执行。比如，A100有108个SM，如果outer loop维度不足108，则会导致GPU利用率不足。不过训练时候sequence_length * batch_size * head很大，切分绰绰有余。

在推理Decoding阶段，由于N=1，图1中Outer Loop消失了，只有一个inner loop循环。如果还按照FA V2老办法并行，则thread block数目远小于SM数目，导致GPU利用率偏低。

为了保证并行度，有足够的thread block，只能硬着头皮沿着inner loop维度切分。但是inner loop需要做规约操作的，不同thread block任务之间需要通信。因此，就有了下图的执行逻辑。沿着inner loop切成5分，Split X/5表示第X个thread block，在thread block内部计算局部结果，最后5个thread block还需要Reduction Step得到最终结果。

最后作者实验验证了，Flash Decoding比FA、PyTorch和FasterTransformer的attention kernel都快。其实我比较好奇FasterTransformer怎么优化MHA，希望有了解读者不吝赐教。

Flash Decoding是FlashAttention针对推理场景的适配。由于Q矩阵sequence length维度变成1，导致原来可以切分的维度消失，必须沿着之前的内层循环沿着K/V的sequence length维度切分，以获得更多的可并行子任务，这带来的后果是不同子任务之间需要一次规约操作来同步数据。