概述
它使用平铺和在线softmax来减少内存流量；速度和内存增益取决于序列长度、硬件和实现。
深入探讨
标准注意力将每个标记与其他每个标记进行比较，生成一个 N×N 分数矩阵，该矩阵随序列长度呈二次方增长。简单来说，该矩阵被写入 GPU 高带宽内存 (HBM) 并从其中读回，而真正的瓶颈是穿梭（而不是乘法）。 Tri Dao 及其同事于 2022 年推出的 Flash Attention 重新组织了计算，因此矩阵永远不会完全存储。它处理适合快速片上 SRAM 的小块中的查询、键和值，计算部分结果，并使用在线运行的 softmax 技巧将它们拼接在一起。该算法计算相同的注意力操作，同时减少高带宽内存流量和中间存储。它可以更快，尤其是在长序列上，但测量的增益取决于 GPU、序列长度、数据类型和实现。
现实世界的实施
以更低的内存成本训练大型语言模型，例如 Llama 和 GPT 类系统，具有更长的上下文窗口。
通过加快首次阅读长提示的预填充阶段，更快地为聊天助理提供服务。
通过在单个 GPU 上实现长序列注意力，启用可摄取整本书或代码库的文档分析工具。
为视觉和音频变压器提供动力，其中高分辨率输入创建非常长的令牌序列。
Flash 注意力的未来
FlashAttention-2 和 FlashAttention-3 研究探索了改进的工作分区、异步执行和特定于硬件的低精度路径。结果取决于 GPU、序列长度、数据类型和实现。这些算法可以减少密集注意力的内存流量，但它们不会消除底层的查询密钥对计算。在假设加速之前对确切的训练或服务工作量进行基准测试。
付诸实践
为什么这很重要
- 成本与预算. 多年来，架构决策决定着性能和运营成本。
- 更清晰的判决. 技术教育帮助团队选择正确的堆栈，而不仅仅是最新的堆栈。
- 质量控制. 更好的工程选择可以减少生产中的可靠性事故。
留意
- 优化一项基准测试可以隐藏更广泛的系统弱点。
- 基础设施和维护成本常常被低估。
- 随着系统变得更加复杂，安全性和可观察性差距可能会扩大。
后续步骤
- 在实施之前定义延迟、质量和成本目标。
- 在实际负载和数据条件下进行基准测试。
- 仪器监控错误、漂移和用户影响。
- 在扩展之前准备回滚和事件响应路径。
不断探索
常见问题
什么是Flash注意力？
FlashAttention 是一种 IO 感知方式，可以计算精确的注意力，而无需在高带宽内存中实现完整的注意力矩阵。它使用平铺和在线softmax来减少内存流量；速度和内存增益取决于序列长度、硬件和实现。
Flash Attention的主要瓶颈是什么？
Flash Attention 具有 IO 感知能力：它减少了快速片上 SRAM 和慢速高带宽内存之间传输的数据，这是真正的瓶颈，而不是算法本身。
Flash Attention 如何避免存储完整的 N×N 注意力矩阵？
它平铺计算，因此每个块都适合快速 SRAM，计算和累积部分输出，而无需在 HBM 中具体化整个矩阵。
什么技术可以让 Flash Attention 正确计算 softmax，而无需一次看到整行？
在线 softmax 在流式传输瓦片时保持运行最大值和运行分母，重新缩放早期的部分输出，以便最终标准化是准确的。
为什么 Flash Attention 对于长序列最有帮助？
朴素注意力矩阵在内存中按 N 平方缩放，因此当序列很长时，避免其完全存储会产生最大的节省。
Flash Attention 的“IO 感知”设计优先考虑最小化什么？
IO 感知意味着算法是围绕内存层次结构中移动数据的成本设计的，最大限度地减少 HBM 流量而不是算术运算。