2023/10/16 14:54

别再「浪费」GPU了，FlashAttention重磅升级，实现长文本推理速度8倍提升

处理小说、法律文件等长文本是大模型的一个重要应用方向，但也面临速度上的挑战。FlashAttention 作者 Tri Dao 等人提出的「Flash-Decoding」通过充分利用 GPU，可以将大模型的长上下文推理速度提高至 8 倍。

最近，像 ChatGPT 或 Llama 这样的大型语言模型（LLM）引起了前所未有的关注。然而，它们的运行成本仍然极高。虽然生成单个响应可能仅需 0.01 美元（在 AWS 上的 8xA100 实例上运行几秒钟），但当扩大规模以满足数十亿用户的需求时，成本会迅速累积。而且，这些用户可能每天与 LLM 进行多次互动。某些用例的成本更高，例如代码自动生成，因为它会随着每次输入新字符而运行。随着 LLM 应用的不断增加，即使在生成时间方面实现细微的效率提升，也将产生巨大的影响。

LLM 推理（或「解码」）是一个迭代的过程：token 逐个生成。生成包含 N 个 token 的完整句子需要通过模型进行 N 次前向传递。幸运的是，我们可以缓存先前计算的 token：这意味着单个生成步骤不依赖于上下文长度，除了一个单独的操作 —— 注意力。这个操作导致上下文长度不能很好地扩展。

在 LLM 的重要新兴用例中，有一些需要利用更长的上下文。只有拥有了更长的上下文窗口，LLM 才能对更长的文档进行推理，无论是总结文档还是回答其中的问题。此外，它们还可以保持更长的对话历史，甚至在编写代码之前处理整个代码库。举个例子，在 2022 年，大多数 LLM 的上下文长度最多为 2k（例如 GPT-3），但现在，有些开源 LLM 已经可以扩展到 32k（比如 Llama-2-32k），甚至有些模型已经达到了 100k（比如 CodeLlama）。在这些情境中，注意力操作在推理过程中占据了相当大的时间比例。

在扩展 batch size 维度时，即使上下文相对较短，注意力也可能成为一个瓶颈。这是因为随着 batch 维度的增加，需要读取的内存量也会增加，而对于模型的其余部分，内存需求只取决于模型的大小。

为了解决上述问题，FlashAttention 的作者 Tri Dao 等人提出了一项名为「Flash-Decoding」的技术，它显著加速了推理过程中的注意力计算，使长序列的处理生成速度提高到了原来的 8 倍。其主要思想是以最快的速度并行加载键和值，然后分别重新缩放和合并结果，以维持正确的注意力输出。

解码时的多头注意力

在解码期间，生成的每个新 token 都需要关注所有先前的 token，以计算：softmax (queries @ keys.transpose) @ values

这个操作已经在训练阶段通过 FlashAttention 进行了优化（包括最近的 v1 和 v2 版本），瓶颈是读写中间结果的内存带宽（如 Q @ K^T）。然而，这些优化并不直接适用于推理情况，因为瓶颈不同。在训练中，FlashAttention 并行处理 batch size 和查询长度两个维度。而在推理过程中，查询长度通常为 1：这意味着，如果 batch size 小于 GPU 上的流多处理器（streaming multiprocessor，SM）数量（例如 A100 有 108 个），该操作只会利用 GPU 的一小部分！特别是在处理长上下文时，情况尤为明显，因为它需要较小的 batch size 以适应 GPU 内存。当 batch size 为 1 时，FlashAttention 将使用不到 1% 的 GPU！

^{FlashAttention 只在查询块和 batch size 之间并行，并且在解码期间不会设法占用整个 GPU。}

使用矩阵乘法基元也能执行注意力计算，这样就不需要使用 FlashAttention 了。在这种情况下，该操作会占用整个 GPU，但会启动许多写入和读取中间结果的内核，因此并不是最优的做法。

更快的注意力解码：Flash-Decoding

新方法 Flash-Decoding 基于 FlashAttention，同时引入了一个新的并行维度：键值序列的长度。它综合了上述两种方法的优点。与 FlashAttention 类似，它在全局内存中存储的额外数据很少。然而，只要上下文足够长，即使 batch size 较小，它也能充分利用 GPU。

^{Flash-Decoding 也在键和值之间并行化，代价是一个小的最终归约（reduction 步骤。}

Flash-Decoding 主要有三个工作步骤：

首先，将键 / 值分成更小的块；
使用 FlashAttention 并行计算查询与每个这些分块的注意力，为每行和每个分块额外写入一个标量值：注意力值的 log-sum-exp
最后，通过对所有分块进行归约来计算实际输出，使用 log-sum-exp 来调整每个分块的贡献。

这一切之所以可行，都是因为注意力 /softmax 可以进行迭代计算。在 Flash-Decoding 中，它在两个级别上被使用：在分块内部（类似 FlashAttention），以及跨分块进行最终的归约计算。

实际操作中，步骤（1）不涉及任何 GPU 操作，因为键 / 值块是完整键 / 值张量的视图。然后，有两个独立的核函数，分别用于执行步骤（2）和（3）。

在 CodeLlama 34B 上进行的基准测试

为了验证上述新方法，研究者对 CodeLLaMa-34b 的解码吞吐量进行了基准测试。该模型与 Llama 2 具有相同的架构，一般来说，结果应该适用于许多大型语言模型。研究者在不同序列长度下（从 512 到 64k），以 tok/s 为单位来测量解码速度，并比较了多种计算注意力的方式：

Pytorch：使用纯粹的 PyTorch 基元来运行注意力计算（不使用 FlashAttention）；
FlashAttention v2；
FasterTransformer：使用 FasterTransformer 的注意力内核；
Flash-Decoding；
以及一个上限值，该值计算了从内存中读取整个模型和 KV-cache 所需的时间

对于非常大的序列，Flash-Decoding 可以将解码速度提高至 8 倍，并且比其他方法的扩展性要好得多。

^{在 prompt 比较小时，所有方法表现接近。但是当序列长度从 512 增加到 64k 时，除了 Flash-Decoding，其他方法的可扩展性都很差。在 Flash-Decoding 的这种模式下（batch size 为 1），扩展序列长度对生成速度的影响很小。}

组件级微基准测试

研究者还在 A100 上对多头注意力进行了微基准测试，输入为 f16，考虑了不同的序列长度和 batch size。他们将 batch size 设置为 1，并且使用 16 个 128 维的查询头，以及 2 个键 / 值头（分组查询注意力），这与在 4 个 GPU 上运行的 CodeLLaMa-34b 使用的维度相匹配。

^{上述微基准测试展示了多头注意力的运行时间，单位为微秒。Flash-Decoding 在序列长度扩展到高达 64k 时，几乎实现了恒定的运行时间。}

之前测量的高达 8 倍的端到端加速是可能的，因为注意力本身的速度比 FlashAttention 快高达 50 倍。在序列长度达到 32k 之前，注意力的时间大致是恒定的，因为 Flash-Decoding 能够完全利用 GPU。

使用 Flash-Decoding

Flash-decoding 可以在以下链接中找到：

FlashAttention 包，从 v2.2 开始：https://github.com/Dao-AILab/flash-attention/tree/main
xFormers 包（搜索 xformers.ops.memory_efficient_attention），从 0.0.22 开始：调度程序将根据问题的大小自动使用 Flash-Decoding 或 FlashAttention 方法。当这些方法不受支持时，它可以调度到一个高效的 triton 内核，该内核实现了 Flash-Decoding 算法。

一个完整的使用 LLaMa v2 / CodeLLaMa 的解码示例可以在 FlashAttention repo 和 xFormers repo 中找到。此外，作者还提供了一个简单的 LLaMa v1/v2 模型的高效解码代码示例，旨在快速、易读、有教育意义和易于修改。

^{参考链接：https://princeton-nlp.github.io/flash-decoding/}

产业Flash-Decoding

相关数据

调度技术

调度在计算机中是分配工作所需资源的方法。资源可以指虚拟的计算资源，如线程、进程或数据流；也可以指硬件资源，如处理器、网络连接或扩展卡。进行调度工作的程序叫做调度器。调度器通常的实现使得所有计算资源都处于忙碌状态，允许多位用户有效地同时共享系统资源，或达到指定的服务质量。 see planning for more details

来源：维基百科

核函数技术

核函数包括线性核函数、多项式核函数、高斯核函数等，其中高斯核函数最常用，可以将数据映射到无穷维，也叫做径向基函数（Radial Basis Function 简称 RBF），是某种沿径向对称的标量函数。最常应用于SVM支持向量机中

来源：百度百科

基准技术

一种简单的模型或启发法，用作比较模型效果时的参考点。基准有助于模型开发者针对特定问题量化最低预期效果。

来源：Google ML Glossary

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数，这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在维空间内，有个分量的一种量，其中每个分量都是坐标的函数，而在坐标变换时，这些分量也依照某些规则作线性变换。称为该张量的秩或阶（与矩阵的秩和阶均无关系）。在数学里，张量是一种几何实体，或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达，记作标量的数组，但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中，表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了，它们都是二阶张量，对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

来源：维基百科

查询技术

一般来说，查询是询问的一种形式。它在不同的学科里涵义有所不同。在信息检索领域，查询指的是数据库和信息系统对信息检索的精确要求

来源：Wikipedia

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用，如语音识别，机器翻译，词性标注，句法分析和资讯检索。由于字词与句子都是任意组合的长度，因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题)，也使得在语料库中估算字串的机率变得很困难，这也是要使用近似的平滑n元语法(N-gram)模型之原因。

来源：维基百科

分块技术

将标注好词性的句子按句法结构把某些词聚合在一起形成比如主语、谓语、宾语等等。

来源：掘金