浅谈FlashAttention优化原理

FlashAttention原理简介

背景：

在Transformer 结构中，自注意力机制的时间和存储复杂度与序列的长度呈平方的关系，因此占用了大量的计算设备内存和并消耗大量计算资源。如何优化自注意力机制的时空复杂度、增强计算效率是大语言模型需要面临的重要问题

1.FlashAttention

NVIDIA GPU 中的内存（显存）按照它们物理上是在GPU 芯片内部还是板卡RAM 存储芯片上，决定了它们的速度、大小以及访问限制。

GPU 显存分为：全局内存（Global memory）、本地内存（Local memory）、共享内存（Shared memory，SRAM）、寄存器内存（Register memory）、常量内存（Constant memory）、纹理内存（Texture memory）等六大类。

下图给出了NVIDIA GPU 内存的整体结构。其中全局内存、本地内存、共享内存和寄存器内存具有读写能力。全局内存和本地内存使用的高带宽显存（High Bandwidth Memory，HBM）位于板卡RAM 存储芯片上，该部分内存容量很大。

全局内存是所有线程都可以访问，而本地内存则只能当前线程访问。NVIDIA H100 中全局内存有80GB 空间，其访问速度虽然可以达到3.35TB/s，但是如果全部线程同时访问全局内存时，其平均带宽仍然很低。

共享内存和寄存器位于GPU 芯片上，因此容量很小，并且共享内存只有在同一个GPU 线程块（Thread Block）内的线程才可以共享访问，而寄存器仅限于同一个线程内部才能访问。NVIDIA H100 中每个GPU 线程块在流式多处理器（Stream Multi-processor，SM）可以使用的共享存储容量仅有228KB，但是其速度非常快，远高于全局内存的访问速度。

通过自注意力机制的原理可知，在GPU 中进行计算时，传统的方法还需要引入两个中间矩阵S 和P 并存储到全局内存中。具体计算过程如下：

S = Q ×K, P = Softmax(S), O = P × V (2.27)

按照上述计算过程:

• 需要首先从全局内存中读取矩阵Q 和K，并将计算好的矩阵S 再写入全局内存

• 再从全局内存中获取矩阵S，计算Softmax 得到矩阵P，再写入全局内容

• 读取矩阵P 和矩阵V ，计算得到矩阵O。

这样的过程会极大占用显存的带宽。在自注意力机制中，计算速度比内存速度快得多，因此计算效率越来越多地受到全局内存访问的瓶颈。

FlashAttention就是通过利用GPU 硬件中的特殊设计，针对全局内存和共享存储的I/O 速度的不同，尽可能的避免HBM 中读取或写入注意力矩阵。

FlashAttention 目标:是尽可能高效地使用SRAM 来加快计算速度，避免从全局内存中读取和写入注意力矩阵。达成该目标需要能做到在不访问整个输入的情况下计算Softmax 函数，并且后向传播中不能存储中间注意力矩阵。

标准Attention 算法中:

• Softmax 计算按行进行，即在与V 做矩阵乘法之前，需要将Q、K 的各个分块完成一整行的计算。

• 在得到Softmax 的结果后，再与矩阵V 分块做矩阵乘。

而在FlashAttention 中:

• 将输入分割成块，并在输入块上进行多次传递，从而以增量方式执行Softmax 计算。

自注意力算法的标准实现将计算过程中的矩阵S、P 写入全局内存中，而这些中间矩阵的大小与输入的序列长度有关且为二次型。

因此，FlashAttention 就提出了不使用中间注意力矩阵，通过存储归一化因子来减少全局内存的消耗。FlashAttention 算法并没有将S、P 整体写入全局内存，而是通过分块写入，存储前向传递的Softmax 归一化因子，在后向传播中快速重新计算片上注意力，这比从全局内容中读取中间注意力矩阵的标准方法更快。由于大幅度减少了全局内存的访问量，即使重新计算导致FLOPs 增加，但其运行速度更快并且使用更少的内存。

PyTorch 2.0 中已经可以支持FlashAttention，使用“torch.backends.cuda.enable_flash_sdp()”启
用或者关闭FlashAttention 的使用。