FlashAttention：Fast and Memory-Efficient Exact Attention with IO-Awareness

FlashAttention

序：

基于 FlashAttention 技术，清华将基座模型的上下文长度（Context Length）由 ChatGLM-6B 的 2K 扩展到了 32K，并在对话阶段使用 8K 的上下文长度训练，允许更多轮次的对话。

概述：

• 给 Transformer 配备更长的上下文仍然很困难，因为 Transformer 核心自注意力模块的时间复杂度以及内存复杂度在序列长度上是二次方的。

• FlashAttention是一种新的注意力机制，旨在解决Transformer在处理长序列时速度慢且内存需求大的问题。

• FlashAttention的创新之处在于引入了IO（输入/输出）感知的设计原则，专注于减少GPU内存（高带宽内存）和GPU片上内存（SRAM）之间的读写次数。它使用平铺（tiling）的方法来实现这一目标，从而降低了数据在不同级别存储器之间的传输次数。

• 此外，研究者还将FlashAttention扩展到块稀疏注意力（block-sparse attention），产生比现有的近似注意力方法更快的近似注意力算法。

简介：

1. SRAM：SRAM（Static Random Access Memory）是一种高速、低延迟的随机访问存储器。在GPU（Graphics Processing Unit，图形处理器）中，SRAM通常用于存储临时数据、缓存以及计算中的中间结果。GPU SRAM是指嵌入在GPU芯片上的SRAM存储器。
2. HBM（High-Bandwidth Memory）：HBM是一种高带宽、低功耗的内存技术。GPU HBM是指嵌入在GPU芯片上的HBM存储器。相对于传统的GDDR（Graphics Double Data Rate）内存，HBM具有更高的带宽和更低的能耗，使得GPU能够更快地读取和写入数据。
3. DARM（Dynamic Random Access Memory）：DARM是一种动态随机访问存储器，用于存储CPU（Central Processing Unit，中央处理器）中的数据和指令。DARM是一种较为常见的内存类型，具有较高的存储密度，但相对于SRAM和HBM而言，它的访问速度较慢，延迟较高。CPUDARM指的是嵌入在CPU芯片上的DARM存储器。

FlashAttention

这是一种新的注意力算法，主要目标是避免从 HBM 读取和写入注意力矩阵。

• 在不访问整个输入的情况下计算 softmax 缩减。

我们重组注意力计算以将输入分成块并在输入块上进行多次传递，从而逐步执行 softmax 缩减（也称为平铺）

• 不为后向传递存储大的中间注意力矩阵。

存储前向传播的 softmax 归一化因子，以便在后向传播中快速重新计算片上注意力，这比从 HBM 读取中间注意力矩阵的标准方法更快。我们在 CUDA 中实现了FlashAttention以实现对内存访问的细粒度控制，并将所有注意力操作融合到一个 GPU 内核中。

即使由于重新计算而增加了 FLOP，我们的算法也比标准注意力更快地运行（在 GPT-2 [67] 上高达 7.6 倍，图1 右）并且使用更少的内存（序列长度呈线性），这要归功于大量减少 HBM 访问量。

块稀疏 FlashAttention