闪光注意
Flash Attention 是一种计算 Transformer 内部注意力步骤的巧妙方法，无需编写巨大的注意力矩阵来减慢内存速度。
概述
它使长上下文模型速度更快、内存效率更高，而无需改变其数学原理。
深入探讨
标准注意力将每个标记与其他每个标记进行比较，生成一个 N×N 分数矩阵，该矩阵随序列长度呈二次方增长。简单地说，该矩阵被写入 GPU 高带宽内存 (HBM) 并从其中读回，而真正的瓶颈是穿梭（而不是乘法）。 Tri Dao 及其同事于 2022 年推出的 Flash Attention 重新组织了计算，因此矩阵永远不会完全存储。它处理适合快速片上 SRAM 的小块中的查询、键和值，计算部分结果，并使用在线运行的 softmax 技巧将它们拼接在一起。输出在数学上与普通注意力相同，但使用线性存储器并且运行速度快几倍，尤其是在长序列上。
技术洞察
关键技巧是平铺加上在线 softmax。 Softmax 通常需要整行分数来计算其分母，但 Flash Attention 在流式传输每个图块时会保留运行最大值和运行总和，重新缩放早期的部分输出，以便最终结果准确。由于中间分数保留在 SRAM 中（比 HBM 快几个数量级），因此该算法具有 IO 感知能力：它最大限度地减少了内存读取和写入，而不是原始算术运算。
战略影响
成本与预算
多年来，架构决策决定着性能和运营成本。
更清晰的判决
技术教育帮助团队选择正确的堆栈，而不仅仅是最新的堆栈。
质量控制
更好的工程选择可以减少生产中的可靠性事故。
Flash 注意力的未来
Flash Attention 已成为默认构建块，FlashAttention-2 和 FlashAttention-3 通过改进工作分区和利用低精度 FP8 路径，从 H100 等新型 GPU 中获得更多吞吐量。预计将继续与硬件进行协同设计，更紧密地集成到训练和推理框架中，以及针对稀疏、滑动窗口和超长上下文注意力进行调整的变体。随着上下文窗口扩展到数百万个令牌，像这样的 IO 感知内核对于保持内存和速度实用仍然至关重要。
现实世界的实施
以更低的内存成本训练大型语言模型，例如 Llama 和 GPT 类系统，具有更长的上下文窗口。
通过加快首次阅读长提示的预填充阶段，更快地为聊天助理提供服务。
通过在单个 GPU 上实现长序列注意力，启用可摄取整本书或代码库的文档分析工具。
为视觉和音频变压器提供动力，其中高分辨率输入创建非常长的令牌序列。
风险与防护栏
优化一项基准测试可以隐藏更广泛的系统弱点。
基础设施和维护成本常常被低估。
随着系统变得更加复杂，安全性和可观察性差距可能会扩大。
实施路线图
在实施之前定义延迟、质量和成本目标。
在实际负载和数据条件下进行基准测试。
仪器监控错误、漂移和用户影响。
在扩展之前准备回滚和事件响应路径。
不断探索
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the Flash Attention quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
常见问题
什么是Flash注意力？
Flash Attention 是一种计算 Transformer 内部注意力步骤的巧妙方法，无需编写巨大的注意力矩阵来减慢内存速度。它使长上下文模型速度更快、内存效率更高，而无需改变其数学原理。
Flash Attention的主要瓶颈是什么？
Flash Attention 具有 IO 感知能力：它减少了快速片上 SRAM 和慢速高带宽内存之间传输的数据，这是真正的瓶颈，而不是算法本身。
Flash Attention 如何避免存储完整的 N×N 注意力矩阵？
它平铺计算，因此每个块都适合快速 SRAM，计算和累积部分输出，而无需在 HBM 中具体化整个矩阵。
什么技术可以让 Flash Attention 正确计算 softmax，而无需一次看到整行？
在线 softmax 在流式传输瓦片时保持运行最大值和运行分母，重新缩放早期的部分输出，以便最终标准化是准确的。
为什么 Flash Attention 对于长序列最有帮助？
朴素注意力矩阵在内存中按 N 平方缩放，因此当序列很长时，避免其完全存储会产生最大的节省。
Flash Attention 的“IO 感知”设计优先考虑最小化什么？
IO 感知意味着算法是围绕内存层次结构中移动数据的成本设计的，最大限度地减少 HBM 流量而不是算术运算。