Flash Attention

一、GPU的HBM和SRAM

速度：SRAM的读写速度通常比HBM快一个数量级，因为SRAM利用晶体管存储数据，而HBM则是通过堆叠DRAM芯片实现高带宽和容量。然而，在实际应用中，由于HBM的堆叠设计和高速串行接口，其整体性能仍然非常出色。

容量：HBM通过堆叠多个DRAM芯片实现了更高的容量和带宽，而SRAM的容量则相对较小。因此，在需要处理大规模数据集和复杂计算任务的应用中，HBM更具优势

二、attention理解

attention的本质：带权特征筛选。注意力机制描述了（序列）元素的加权平均值，其权重是根据输入的query和元素的键值进行动态计算的。

seq2seq模型：

三、FlashAttention原理解读

• fast：加快模型训练速度
• memory-efficient：减少显存占用
• exact attention：和标准attention计算得到的结果完全一致，精度不变
• IO-awareness：算法是改进IO的效率

以一个attention的标准计算过程为例进行存储方面的说明：从下图中可看到随着序列长度N的增长，中间矩阵的缓存是以N的平方进行增长的。而且这些中间矩阵是必须要存的，因为反向传播时需要它们来计算梯度。 flash attention可以将N的复杂度降低到一次。

正常attention 的读取过程:

Flash Attention原理：避免attention matrix从HBM的读写:

• 传统减少HBM的访问，将QKV切分为小块后放入SRAM中
• 反向传播时，重新计算中间结果，以此来解决不缓存后梯度无法计算的问题

 tiling(平铺): 分块计算：

四、补充

RAM（Random Access Memory）随机存取存储器，是计算机中的一种主要的存储设备。RAM用于暂时存储和访问数据，以便中央处理单元（CPU）快速处理。以下是RAM的一些关键特点：

1. 临时存储：RAM是一种易失性存储介质，这意味着当计算机断电时，存储在RAM中的数据会丢失。

2. 快速访问：与硬盘驱动器或固态驱动器等长期存储设备相比，RAM的读写速度非常快，这使得它非常适合作为操作系统、应用程序和当前正在处理的数据的临时存储位置。

3. 随机访问：与顺序访问存储设备（如磁带）不同，RAM允许以任何顺序读取和写入数据，这使得数据访问非常高效。

4. 容量：现代计算机通常配备了几GB到数百GB的RAM，具体取决于计算机的用途和用户的需求。

5. 类型：常见的RAM类型包括：

   • SRAM（静态RAM）：速度快，但成本较高，通常用于缓存。

   • DRAM（动态RAM）：速度较SRAM慢，但更便宜，通常用于主内存。

RAM是计算机性能的关键组成部分，因为它直接影响到计算机运行多个程序和处理大量数据的能力。足够的RAM可以显著提高计算机的响应速度和整体性能。

参考说明

Flash Attention原理详解(含代码讲解) - 知乎 (zhihu.com)