🟣 Flash Attention：一个数学技巧把内存省 5 倍
🟣 Flash Attention：一个数学技巧把内存省 5 倍
Tri Dao 2022 年的一篇论文， 没改 Attention 的数学， 也没让模型变更聪明， 却让全行业训练 / 推理一夜之间快 2-4 倍。 这就是 Flash Attention。
🟣 AI 深度派专栏 · 第 S2-06 期 📅 2026年9月26日 ✍️ 小敏说 AI
💡 本文要点：本文用人话讲清 Flash Attention 的核心思想——重排计算顺序、用 SRAM 而不是 HBM、Online Softmax，以及它如何成为”事实标准”。
🎯 本文导读
- 🔹 朴素 Attention 的内存瓶颈
- 🔹 GPU 内存层级速览
- 🔹 Flash Attention 的核心三招
- 🔹 三代演进
- 🔹 它为什么是”基础设施级”突破
一、🧱 朴素 Attention 的内存瓶颈
标准 Attention 公式：
S = Q × K^T # 形状 [N, N]
P = softmax(S) # 还是 [N, N]
O = P × V # 形状 [N, d]
问题：中间矩阵 S 和 P 大小都是 N²。 当 N = 8192：S 占 256 MB（FP16），N=64K 时就 ≈ 16 GB。
💬 关键观点：Attention 慢，不是因为算得多，是因为”搬来搬去”。
二、📐 GPU 内存层级（关键背景）
GPU 不是只有”显存”，是分层的：
| 层级 | 容量 | 速度 |
|---|---|---|
| HBM（显存） | 80GB | 较慢（3 TB/s） |
| SRAM（片上） | 几 MB / SM | 超快（几十 TB/s） |
| Register | 几十 KB | 最快 |
朴素 Attention 反复把 N² 大矩阵在 HBM ↔ 计算单元 间搬运——搬运耗时 > 计算耗时。
三、💡 Flash Attention 的三招
招式 1：分块（Tiling）
把 Q、K、V 切成小块（如 64×64），一次只处理一小块——块小到能塞进 SRAM。
招式 2：在 SRAM 里算完
每块的 Attention 全在 SRAM 里完成，不落地 HBM。
招式 3：Online Softmax
朴素 Softmax 需要看到一整行才能算。Flash Attention 用一个数学等价变换，让 Softmax 可以增量计算：
每来一块，更新 (running_max, running_sum, output)
最后一块算完时，结果完全等价于"一整行算 Softmax"
🎯 一句话总结：Flash Attention = 用分块 + Online Softmax，把”N² 中间矩阵”完全去掉。
四、📊 效果
| 指标 | 朴素 | Flash Attention |
|---|---|---|
| HBM 读写 | O(N²) | O(N) |
| 显存占用 | O(N²) | O(N) |
| 速度 | 1× | 2-4× |
| 数学结果 | — | 完全等价（无精度损失） |
⚠️ 重要点：Flash Attention 不是近似算法——它给出完全相同的数值结果，只是计算更聪明。
五、🪜 三代演进
| 版本 | 关键突破 | 加速比 |
|---|---|---|
| FlashAttention v1 (2022) | 分块 + Online Softmax | 2-4× |
| FlashAttention v2 (2023) | 更好并行 + 减少非矩阵乘 | 再快 ~2× |
| FlashAttention v3 (2024) | 用 H100 异步特性（TMA、WGMMA） | 在 H100 上再快 1.5-2× |
每一代都建立在新一代 GPU 架构上。
六、🌐 它的”基础设施级”意义
Flash Attention 出来后：
| 影响 | 描述 |
|---|---|
| 训练成本 | 一夜降 2-4× |
| 上下文窗口 | 16K → 128K → 1M 才有了可能 |
| 集成进框架 | PyTorch 2.0 自带、HuggingFace 默认开 |
| 所有 LLM 都用 | GPT-4 / Claude / Llama / DeepSeek 都基于它 |
💬 关键观点：Flash Attention 是”那种一旦出现，再也回不去”的工程突破。
七、🤝 配套技术
| 技术 | 作用 |
|---|---|
| Flash Attention | 加速 Attention 本身 |
| KV Cache | 复用过去的 K/V |
| PagedAttention | 像虚拟内存管理 KV Cache |
| Flash Decoding | 长上下文推理优化 |
八、❌ 常见误区
| 误区 | 真相 |
|---|---|
| ❌ Flash Attention 损精度 | ✅ 完全等价 |
| ❌ 它是新算法 | ✅ 是 IO 感知的实现 |
| ❌ 只有 Nvidia 卡能用 | ✅ AMD、华为都已支持 |
| ❌ 自己写代码会更快 | ✅ 已是行业基础设施，没必要重造 |
九、🎁 给非工程读者的”FA 心法”
🔑 大模型瓶颈不是"算"，是"搬"
🔑 FA 把中间矩阵彻底消除
🔑 它让长上下文从"理论可行"变成"商业可跑"
🔑 同样的 GPU，开 FA = 1.5-2 张卡的算力
📮 今日话题：你听过的另一个”看似小、却撑起整个行业”的工程突破是什么？
🔮 下一篇预告：s2ep07《Speculative Decoding：小模型猜、大模型审》
🟣 本系列是「小敏说 AI · AI 深度派专栏」第 S2-06 期。 关注公众号 小敏说 AI，回复「深度派」领取第二季合集。