Flash Attention：少搬几趟，比算得更快重要

Flash Attention：少搬几趟，比算得更快重要

GPU 的算力每代翻倍，但显存带宽十年没怎么涨。

所以 GPU 大部分时间不是在算，是在等数据搬过来。

Flash Attention 做的事情只有一件：少搬。

标准做法：算一步，搬一趟

左边的 Standard，每算一个中间结果都要写回显存再读回来。8 个 cycle 里 5 个在等。

这不是算得慢，是搬得慢。

Flash 做法：搬一次，算到底

右边。切成 Tile，一次搬进 SRAM，在 SRAM 里把 S、Softmax、O 全算完，再搬下一块。中间结果不落地。

搬运量从 O(N²) 降到 O(N)：

N=8K 时，45 GB vs 2 GB。22.5 倍，仅仅是少搬了几趟。

但分块之后，Softmax 怎么算？

传统 Softmax 需要看到整行才能算分母。分块之后只能看到局部。

Online Softmax 的解法：维护两个局部变量，增量更新。

• row_m：到目前为止的最大值
• row_l：到目前为止的累加和

每来一个 Tile，用旧的 m 和 l 算出新的 m 和 l。不需要全局视野，局部状态正确演进就能到达全局正确。

图里 Tile 1 那行是关键：m 从 2.0 变成 3.0，旧的 l 和 O 必须乘 exp(m_prev - m_new) 修正。m 变了，l 就得跟着变，O 也得跟着缩放。 这条因果链是 Online Softmax 的全部秘密。

对应代码就是 row_m_prev → row_m_new，row_l_prev → row_l_new。

反向传播：重新算，比搬回来便宜

Flash Attention 反向传播时不存中间矩阵 S，而是从 Q、K、V 重新算。

看似浪费，实则最优：A100 的 312 TFLOPS 算力是过剩的，2 TB/s 的 HBM 带宽是稀缺的。用富余换稀缺。

一句话总结

在受限系统里，管理"移动"比管理"计算"重要。 少搬一趟的收益，远大于算得快一点。

这个规律不只在 GPU：

• CPU 优化 = 减少 Cache Miss
• 数据库优化 = 减少磁盘 IO
• 分布式优化 = 计算向数据靠拢

都是同一件事：尊重物理约束，让数据少跑路。