模型的计算速度除了与每秒浮点运算次数（FLOPS，Floating Point Operations Per Second）有很大关系，同时与存储访问开销（MAC，Memory Access Cost）也有关。尤其是当计算本身已经很高效的情况下，存储访问开销更加不可忽略。
存储访问开销主要来自两个方面：
- 从存储中读取数据
- 向存储中写数据
与CPU的情况类似，在GPU中当需要计算时，要从显存中读取数据并由计算单元进行计算操作；计算单元完成计算之后，再将结果写入显存中。因此可以将计算分为两类：1）计算密集型。整个计算的耗时主要在于计算本身，对于显存的读写几乎可以忽略，典型的计算包括大矩阵乘法、通道维度较大的卷积操作等。2）存储访问密集型。整个计算的耗时主要集中在存储的访问上，计算本身耗时很短，典型的计算包括逐元素操作（ReLU，Dropout等）以及Reduce操作（求和、softmax和BatchNorm等）。
在绝大多数神经网络中，因为含有大量的存储访问密集型操作，所以存储访问开销都不能忽略不计。但是绝大多数Efficient Transformer都忽略了存储访问开销，所以即使它们整体的FLOPS降低了，但是计算耗时并没有降低。
Flash Attention
虽然目前Transformer已经成为深度学习领域最广泛的架构，但是由于其固有的O(N^2)复杂度和内存限制的键值缓存，在推理过程中表现出次优效率。
特别是对于长序列来说，这种高复杂度使得模型的实际部署变得复杂，这就是大模型在发展初期其输入输出往往只支持2K或者4K长度的原因。基于Transformer的大模型不能处理长序列的本质原因就是，Transformer的计算复杂度和空间复杂度随着序列长度N呈二次方增长。具体计算可以参考苏剑林的博客。
GPU中存储单元主要有HBM和SRAM，HBM容量大但是访问速度慢，SRAM容量小却有着较高的访问速度。例如，A100 GPU有40-80GB的HBM，带宽为1.5~2.0TB/s；每108个流式多核处理器各有192KB的片上SRAM，带宽估计约为19TB/s。可以看出，片上的SRAM比HBM快一个数量级，但是尺寸要小许多数量级。
FlashAttention是一种基于GPU硬件对attention的优化技术，它的加速原理是，通过利用更高速的上层存储计算单元，减少对低速、更下层存储器的访问次数，从而提升模型的训练性能。也就是说Flash Attention的目的不是节约FLOPS，而是减少对HBM的访问，即降低存储访问开销。重点是Flash Attention在训练和预测过程中的结果和标准Attention是一样的，对用户是无感的，而其它加速方法做不到这一点。
标准Attention计算
首先，从HBM中读取完整的$Q$和$K$矩阵（大小均为$n \times d$），计算点积得到相似度得分$S$（大小为$n \times n$），需要进行$O(n * d + n^2)$次HBM访问。
其次，计算注意力权重$P$（大小为$n \times n$）时，需要对相似度得分$S$进行Softmax操作，这需要进行$O(n^2)$次HBM访问。
最后，将注意力权重$P$和值向量$V$（每个大小为$n \times d$）加权求和得到输出向量$O$（大小为$n \times d$）时，需要进行$O(n * d)$次HBM访问。
因此，标准Attention算法的HBM访问总次数为$O(n * d + n ^ 2)$。当$n$比较大时，总的HBM访问次数会比较多。
标准的Attention算法在GPU内存分级存储的架构下，存在以下缺陷：
- 对HBM过多的访问，如注意力得分$S$和注意力权重$P$需要在存入HMB之后又立即被访问，由于HBM带宽较低，这会导致算法性能受限
- 注意力得分$S$和注意力权重$P$需要占用$O(n^2)$的存储空间，显存占用较高
我们可以从以下思路出发：
- 之所以需要大量访问HBM，一个原因就是Attention的计算中存在三个Kernel，每个Kernel的计算过程都存在从HBM读取数据，计算完成之后又要写会HBM。如果将三个Kernel融合为一个，则可以减少部分访问HBM的次数。
- 在计算过程中，要尽量利用SRAM进行计算，避免访问HBM操作。
但是SRAM的带宽虽然很大，它的计算可存储的数据量较小，如果采取分治策略将数据进行Tilling处理放到SRAM中计算，会导致长序列被截断，标准的Softmax无法正常工作。
FlashAttention计算
Flash Attention的实现主要在两点：
- 平铺（Tilling）
将原始的注意力矩阵分解为更小的子矩阵，然后分别对这些子矩阵进行计算，只要这个子矩阵的大小可以存放在SRAM中，那么就可以在计算过程中只访问SRAM了。
前向传播和反向传播时，都使用 - 重算（Recomputation）
重算是深度学习优化中的常见概念，以算力换内存，不要存储那么多梯度和每一层正向传播的中间状态，而是在计算到反向传播的某一层时再临时从头开始重新计算正向传播的中间状态。
仅在反向传播时使用
对于矩阵乘法而言，可以直接通过分块来达到分块计算的目的，但是self-attention中有softmax计算，softmax的分母包含与所有元素相关的求和项，所以对self-attention进行分块计算的真正难点在于对softmax的分块计算。Flash Attention提出了分块Softmax算法，确保了整个算法的正确性，这也是Flash Attention的核心。
分块Softmax算法：
假设有维度为$D$的向量$x \in \mathbb{R}^D$，对于向量$x$的softmax的计算公式如下：
其中$x_i$为向量$x$的第$i$个分量。
因为softmax的计算公式中含有指数项，当$x_i$较大时指数项$e^{x_i}$也会很大，从而在计算中出现溢出。为了避免溢出，大多数深度学习框架中都使用了softmax的稳定版本，稳定版的softmax的计算如下：
其中分子$f(x)$是向量，分母$\sum _j f(x)_j$是标量，所以这里的除法是逐元素相除。
直觉上，softmax难以分块计算的主要原因是它的分母依赖于输入向量$x$中的每一个值。
如果我们将一个向量进行分块，那么在分块计算中我们可以得到局部softmax，这个局部的softmax显然不是最终结果，但是我们可以保存它对应的分子（局部最大值）和分母（局部EXP求和项）。同时我们需要维护一个全局最大值和全局EXP求和项。在进行后续的分块计算时，我们需要根据此时的局部最大值和局部EXP求和项，更新全局最大值和全局EXP求和项，然后再处理下一个分块，然后再更新。当处理完所有分块之后，此时的所有分块的softmax值都是全局的。详细推导过程可参考这篇博客。
Paged Attention
vLLM主要用于LLM推理服务，其核心是Paged Attention，这是一种新颖的注意力算法，它将操作系统的虚拟内存中分页的经典思想引入到LLM服务中，可以有效地管理大语言模型的注意力内存，从而提升大模型的吞吐量和内存使用效率。vLLM的吞吐量可以达到HuggingFace实现的24倍，文本生成推理（TGI）性能高出3.5倍，且不需要对模型结构进行任何修改。
Paged Attention允许在不连续的内存空间中存储连续的keys和values，通过将每个序列的KV cache划分为块，每个块包含固定数量的tokens的keys和values。在注意力计算期间，Paged Attention内核可以有效地识别和获取这些块。因为块在内存中不需要连续，因而可以用一种更加灵活的方式管理key和value，就像在操作系统的虚拟内存中一样：将块视为页面，将token视为字节，将序列视为进程。序列的连续逻辑块通过块表映射到非连续物理块中，物理块在生成新的token时按需分配。这种方式可以更灵活地对KV Cache进行管理，内存浪费只会发生在序列的最后一个块中，这使得在实践中可以实现接近最佳的内存使用，仅浪费不到4%。通过这样的策略，系统将更多序列进行批处理，提高GPU使用率，显著提升吞吐量。可参考官方博客提供的动图进一步了解细节。
此外，Paged Attention还有另一个关键优势——高效的内存共享。例如，在并行采样中多个输出序列是由同一个prompt生成的，在这种情况下，prompt的计算和内存可以在输出序列中共享。
Paged Attention的内存共享大大减少了复杂采样算法的内存开销，例如并行采样和集束搜索的内存使用量降低了55%。这可以转化为高达2.2倍的吞吐量提升。这种采样方法也在LLM服务中变得实用起来。Paged Attention是vLLM背后的核心技术。
【参考文献】