FlashAttention 怎么做？online softmax 怎么更新？
不减 FLOPs，为什么还快？

英文问法Explain how FlashAttention works, including the tiling and the shapes. · How does online softmax work? · Why is it faster if it doesn't reduce FLOPs?

考过字节（讲讲 Flash Attention 的原理，矩阵 tiling 的过程说细一点，最好把 shape 也讲清楚。）· Kimi（Ring Attention 里 online softmax 的 LSE 合并，营销号转发，可信度低）

先弄清楚：每个概念是什么

HBM vs SRAM
HBM 是片外的 GPU memory，大但慢。SRAM 是每个 SM 上的片上存储，快但小（数字见 GPU memory hierarchy）。
标准 attention 慢在哪
\(S = QK^\top\)、\(P = \mathrm{softmax}(S)\) 都是 \(s\times s\)，算完写回 HBM、下一步再读出来。softmax 计算很少，时间花在搬运上。例：\(s = 4096\)，每个 head 的 \(S\) 就有约 1678 万个数。
Online softmax
一行分块读进来，只记 running max \(m\) 和 running sum \(\ell\)。新块让 max 变大，就把旧的 \(\ell\) 和输出乘 \(e^{m_{\text{old}}-m_{\text{new}}}\) 校正。例：一行 (1, 3 | 2, 5)。第一块后 \(m=3\)，\(\ell = e^{-2}+1 = 1.135\)；第二块后 \(m=5\)，\(\ell = 1.135\,e^{-2} + e^{-3} + 1 = 1.203\)，和一次算完相同。
Tiling
Q 切成 \(B_r\) 行一块，K、V 切成 \(B_c\) 行一块，每次在 SRAM 里算一个 \(B_r\times B_c\) 的 score block，不写回 HBM。
Recomputation
backward 要用 \(P\)，但只存输出 \(O\) 和每行的 logsumexp，用时按块重算 \(S\)、\(P\)。

看图：Q、K、V 怎么分块，什么留在 SRAM

内层循环 j：Kj、Vj 一块块读进来外层循环 i：Qi 一块块来KᵀKjᵀ每块 [Bc, d]QQiSij[Br, Bc]算完就丢s × s 的 S、P从来不完整存在，不写回 HBM第 i 块一直带着mi、ℓi：[Br]Oi：[Br, d]每来一个 j 更新一次Qi：[Br, d]　Kj、Vj：[Bc, d]　Sij：[Br, Bc]
图 1：按 FlashAttention-2 的循环顺序画（外层 Q 块、内层 K/V 块）。蓝色跟着外层 i，橙红色跟着内层 j，每次在 SRAM 里算出一个 score 块 \(S_{ij}\)。虚线格是从来没有同时存在过的完整 \(s\times s\) 矩阵。
HBM（大、慢）Q：[s, d]K、V：[s, d]O：[s, d]S、P：[s, s]从不写到这里SRAM（一个 SM 上，小、快）Qi：[Br, d]Kj、Vj：[Bc, d]Oi：[Br, d]mi、ℓi：[Br]Sij、P̃ij：[Br, Bc]每个 i 读 1 次每个 j 读 1 次j 走完写回 1 次每个 j：① Sij = QiKjᵀ② 新 max m′；旧 ℓi、Oi × exp(m − m′)③ ℓi += rowsum P̃ij，Oi += P̃ijVjj 走完：Oi ← Oi / ℓi，写回 HBM
图 2：同一个过程里数据在哪。左边是 HBM，右边是一个 SM 上的 SRAM。\(Q_i\) 读一次、\(O_i\) 最后写一次，\(K_j\)、\(V_j\) 一块块流过 SRAM，\(S_{ij}\)、\(\tilde P_{ij}\) 和 \(m_i\)、\(\ell_i\) 只留在 SRAM。右下 4 行是 online softmax 的公式。

30 秒回答

Standard attention writes the s×s score and probability matrices to HBM and reads them back, so it needs O(s²) memory and is bound by memory traffic, not FLOPs. FlashAttention splits Q into blocks of B_r rows and K, V into blocks of B_c rows, computes each B_r×B_c score block in SRAM, and never writes it out. Softmax needs the whole row, so it uses online softmax: keep a running row max m and sum ℓ, and when a new block raises the max, rescale the old sum and output by exp(m_old − m_new). It's exact, FLOPs stay O(s²d), extra memory is O(s), and backward recomputes the blocks.

推导 / 关键公式

online softmax：每来一块 \(K_j, V_j\)

\[ m' = \max\big(m,\ \operatorname{rowmax} S_{ij}\big),\qquad \tilde P_{ij} = e^{S_{ij} - m'} \] \[ \ell' = e^{m - m'}\,\ell + \operatorname{rowsum} \tilde P_{ij} \] \[ O' = e^{m - m'}\,O + \tilde P_{ij} V_j \]

全部块走完再 \(O \leftarrow O/\ell\)。\(e^{m-m'}\le 1\)，不会 overflow。

为什么快：数 HBM 读写，不数 FLOPs

\(d\) 是 head dim，\(M\) 是 SRAM 大小：

\[ \text{standard: } \Theta(sd + s^2),\qquad \text{Flash: } \Theta(s^2d^2/M) \]

\(d\) = 64–128、\(M\) ≈ 100 KB 时 \(d^2 \ll M\)（FLOPs 见 Transformer 复杂度）。

shape

块shape放在哪
\(Q_i,\ O_i\)\([B_r, d]\)SRAM；\(O_i\) 最后写回 HBM
\(K_j,\ V_j\)\([B_c, d]\)SRAM
\(S_{ij},\ \tilde P_{ij}\)\([B_r, B_c]\)只在 SRAM，用完就丢
\(m_i,\ \ell_i\)\([B_r]\)SRAM；最后每行存一个 logsumexp，共 \([s]\)，给 backward

手撕代码（NumPy，按 K/V 分块）

def flash_attn(Q, K, V, Bc=16):
    s, d = Q.shape
    m = np.full(s, -np.inf)           # 每行 running max
    l = np.zeros(s)                   # 每行 running exp-sum
    O = np.zeros((s, d))              # 还没除以 l 的输出
    for j in range(0, s, Bc):         # 每次只读一块 K_j, V_j
        S = Q @ K[j:j+Bc].T / np.sqrt(d)       # [s, Bc]
        m_new = np.maximum(m, S.max(1))
        P = np.exp(S - m_new[:, None])
        a = np.exp(m - m_new)         # 旧结果的 rescale 系数
        l = a * l + P.sum(1)
        O = a[:, None] * O + P @ V[j:j+Bc]
        m = m_new
    return O / l[:, None]

Q, K, V = np.random.default_rng(0).normal(size=(3, 100, 32))
S = Q @ K.T / np.sqrt(32)             # 对照：直接 softmax
P = np.exp(S - S.max(1, keepdims=True))
ref = P / P.sum(1, keepdims=True) @ V
assert np.allclose(flash_attn(Q, K, V), ref)

必答点

追问

FLOPs 更多了，为什么反而更快？

论文 Fig. 2（A100、GPT-2 medium）：recompute 让 FLOPs 从 66.6 增到 75.2 GFLOPs，但 HBM 读写从 40.3 降到 4.4 GB，forward + backward 从 41.7 ms 降到 7.3 ms。attention 是 memory-bound（见 roofline）。

\(B_r\)、\(B_c\) 怎么定？

由 SRAM 大小 \(M\) 定：\(B_c = \lceil M/4d\rceil\)，\(B_r = \min(\lceil M/4d\rceil, d)\)，保证 \(Q_i, K_j, V_j, O_i\) 同时放得下。

FlashAttention-2 改了什么？

中途不除 \(\ell\)、最后除一次，只存 logsumexp \(L = m + \log\ell\)，减少用不上 Tensor Core 的 non-matmul FLOPs。外层循环改成 Q 块，沿 sequence 维分给不同 thread block。causal mask 下整块被 mask 的直接跳过。

多卡（ring attention）怎么合并两段结果？

每段算出归一化好的 \(O_k\) 和 logsumexp \(L_k\)：\(L = \log(e^{L_1} + e^{L_2})\)，\(O = e^{L_1-L}O_1 + e^{L_2-L}O_2\)，和 online softmax 是同一个式子（见 Context Parallel）。

来源Dao et al. 2022, FlashAttention §2.1（p.3）、§3.1 Algorithm 1 + Theorem 1（p.5）、Theorem 2 + Fig. 2（p.6）、App. B.3–B.4（pp.19–20）· Dao 2023, FlashAttention-2 §3.1–3.2（pp.5–8）· Milakov & Gimelshein 2018 §3 Algorithm 3（pp.2–3）