Flash Attention

Why Flash Attention#

標準的 Attention:

\mathbf{S} = \mathbf{QK}^\top \in \mathbb{R}^{N\times N},\quad \mathbf{P} = \text{softmax}(\mathbf{S}),\quad \mathbf{O} = \mathbf{PV}

其他的高效 attention 研究都是在減少 FLOPS, 但其實 Attention 是 memory bound (一直從 HBM 去讀寫矩陣), Flash Attention 稍微犧牲 FLOPS 讓 I/O 更快

How#

只要不把 $N\times N$ 的 S 和 P 寫進 HBM, 就可以加速

Tiling —— 分塊 + softmax scaling, 讓 softmax 可以增量計算
Recomputation —— 只存 $\mathbf{O}, \ell, m$ ，反向時在 SRAM 裡重算 S、P

再加上 kernel fusion: 整條 pipeline (matmul → mask → softmax → dropout → matmul) 塞進一個 CUDA kernel

但分塊的情況下就算不出 softmax (需要看整個row), 所以就需要 Online Softmax

Online Softmax#

Naive softmax (2 遍掃描)

y_i = \frac{e^{x_i}}{\sum_{j=1}^{V} e^{x_j}}

1
d₀ ← 0
2
for j ← 1, V:
3
    d_j ← d_{j-1} + e^{x_j}        # 第一遍：算分母
4
for i ← 1, V:
5
    y_i ← e^{x_i} / d_V            # 第二遍：正規化

但是這種作法是 Numerical-unsafe: $e^{x_j}$ 在 $x_j$ 稍大時直接 overflow (fp32 上 $x > 88$ 就爆了), 稍小時 underflow 成 0

Safe softmax (3 遍掃描)

標準修法是減去最大值:

y_i = \frac{e^{\,x_i - \max_k x_k}}{\sum_{j=1}^{V} e^{\,x_j - \max_k x_k}}

1
m₀ ← -∞
2
for k ← 1, V:                      # 第一遍：找 max
3
    m_k ← max(m_{k-1}, x_k)
4
d₀ ← 0
5
for j ← 1, V:                      # 第二遍：算分母
6
    d_j ← d_{j-1} + e^{x_j - m_V}
7
for i ← 1, V:                      # 第三遍：正規化
8
    y_i ← e^{x_i - m_V} / d_V

當時主流深度學習框架都在用這個安全版本 (TensorFlow v1.7, PyTorch, MxNet, …)

為了把 $m_V$ 找出來所以多了一次掃描, 這個部份沒辦法平行化 ( $m_V$ 要看過整條向量才知道)

Online softmax (2 遍掃描)

1
m₀ ← -∞
2
d₀ ← 0
3
for j ← 1, V:                                          # 第一遍：m 和 d 同時算
4
    m_j ← max(m_{j-1}, x_j)
5
    d_j ← d_{j-1} × e^{m_{j-1} - m_j} + e^{x_j - m_j}   # 合併在這
6
for i ← 1, V:                                          # 第二遍：正規化
7
    y_i ← e^{x_i - m_V} / d_V

維護一個不變式, 使用目前為止的最大值對目前的所有元素求和:

d_j = \sum_{k=1}^{j} e^{\,x_k - m_j}

d_j = \underbrace{\sum_{k=1}^{j-1} e^{\,x_k - m_j}}_{\text{舊元素，但基準換了}} + \underbrace{e^{\,x_j - m_j}}_{\text{新元素}}

但問題是舊的 $d_{j-1}$ 是以 $m_{j-1}$ 為基準算的，不是 $m_j$ , 所以需要 rescale factor: $e^{-m_{j-1}}e^{+m_{j-1}}$

\sum_{k=1}^{j-1} e^{\,x_k - m_j} = \sum_{k=1}^{j-1} e^{\,x_k - m_{j-1}} \cdot e^{\,m_{j-1} - m_j} = d_{j-1}\cdot e^{\,m_{j-1} - m_j}

其中 $e^{m_{j-1} - m_j}$ 可以提到求和符號外面, 因為它跟 $k$ 無關, 所以可以得:

d_j = d_{j-1}\cdot e^{\,m_{j-1} - m_j} + e^{\,x_j - m_j}

平行化#

Online softmax 還是循序的, 但 GPU 要平行, 所以論文把它推廣成一個二元運算子, 把狀態寫成一對 $(m, d)$ ，定義:

\begin{pmatrix} m_i \\ d_i \end{pmatrix} \oplus \begin{pmatrix} m_j \\ d_j \end{pmatrix} = \begin{pmatrix} \max(m_i, m_j) \\[4pt] d_i\, e^{\,m_i - \max(m_i,m_j)} + d_j\, e^{\,m_j - \max(m_i,m_j)}\end{pmatrix}

所以整條向量的統計量就是:

\begin{pmatrix} m_V \\ d_V\end{pmatrix} = \begin{pmatrix} x_1 \\ 1\end{pmatrix} \oplus \begin{pmatrix} x_2 \\ 1\end{pmatrix} \oplus \cdots \oplus \begin{pmatrix} x_V \\ 1\end{pmatrix}

運算 $\oplus$ 是可結合的, 也是可交換的, 可結合 → 可以隨便加括號; 可交換 → 順序無所謂

但 Online softmax 本身還是兩遍, 第一遍算出 $(m_V, d_V)$ , 第二遍才能算 $y_i = e^{x_i - m_V}/d_V$ , 因為要輸出整條 $y$ 向量, 但每個 $y_i$ 都要除以只有掃完才知道的 $d_V$

論文自己也承認這點, 但有指出一個特例：當 softmax 後面接的是 TopK 時, 可以把兩者融合成一遍, 因為 TopK 不需要算出所有 $y_i$

但在 attention 裡, softmax 的輸出 $\mathbf{P}$ 從來就不是我們要的東西, 要的是 $\mathbf{O} = \mathbf{PV}$ ; $\mathbf{P}$ 只是個中間產物，它馬上就會被 $\mathbf{V}$ 加權求和掉, 所以這個跟 TopK 的性質就有點像, 可以塞進online的遞推裡面, 定義一個未正規化的累積輸出:

\tilde{o}_j := \sum_{k=1}^{j} e^{\,x_k - m_j}\,v_k

用一樣的方法可以得到:

\tilde{o}_j = \sum_{k=1}^{j-1} e^{\,x_k - m_{j-1}}v_k \cdot e^{\,m_{j-1}-m_j} + e^{\,x_j - m_j}v_j = \tilde{o}_{j-1}\cdot e^{\,m_{j-1}-m_j} + e^{\,x_j - m_j}\,v_j

o = \frac{\tilde{o}_V}{d_V}

狀態從 $(m, d)$ 擴充成三元組 $(m, d, \tilde{o})$ ， $\oplus$ 也擴充一下:

\begin{pmatrix} m_a \\ d_a \\ \tilde{o}_a\end{pmatrix} \oplus \begin{pmatrix} m_b \\ d_b \\ \tilde{o}_b\end{pmatrix} = \begin{pmatrix} m := \max(m_a, m_b) \\[3pt] d_a e^{m_a - m} + d_b e^{m_b - m} \\[3pt] \tilde{o}_a\, e^{m_a - m} + \tilde{o}_b\, e^{m_b - m}\end{pmatrix}

$\tilde{o}$ 那一行跟 $d$ 那一行結構完全相同，所以結合律與交換律自動繼承

Algorithm 1: Forward pass#

分塊大小 ( $M$ = SRAM 容量):

B_c = \left\lceil \frac{M}{4d}\right\rceil,\qquad B_r = \min\left(\left\lceil \frac{M}{4d}\right\rceil,\; d\right)

$T_r = \lceil N/B_r\rceil$ (Q 的塊數), $T_c = \lceil N/B_c\rceil$ (K/V 的塊數)

HBM 初始化: $\mathbf{O} = (0)_{N\times d}$ ， $\ell = (0)_N$ ， $m = (-\infty)_N$

1
for j = 1 … T_c:                              ← 外層：K/V
2
    載入 K_j, V_j → SRAM
3

4
    for i = 1 … T_r:                          ← 內層：Q
5
        載入 Q_i, O_i, ℓ_i, m_i → SRAM        ← O、ℓ、m 存在 HBM，每步都要搬
6

7
        S_ij = Q_i K_jᵀ                        ∈ R^{B_r × B_c}   (on chip)
8

9
        m̃_ij = rowmax(S_ij)
10
        P̃_ij = exp(S_ij − m̃_ij)
11
        ℓ̃_ij = rowsum(P̃_ij)
12

13
        m_i^new = max(m_i, m̃_ij)
14
        ℓ_i^new = e^{m_i − m_i^new} ℓ_i + e^{m̃_ij − m_i^new} ℓ̃_ij
15

16
        寫回 O_i ← diag(ℓ_i^new)⁻¹ ( diag(ℓ_i) e^{m_i − m_i^new} O_i
17
                                     + e^{m̃_ij − m_i^new} P̃_ij V_j )
18
        寫回 ℓ_i ← ℓ_i^new，m_i ← m_i^new

Algorithm 1 回傳 $\mathbf{O} = \text{softmax}(\mathbf{QK}^\top)\mathbf{V}$ , FLOPs 為 $O(N^2d)$ , 額外記憶體 $O(N)$

FLOPs 的算法: 內層每次做兩個 matmul, 各 $O(B_rB_cd)$ ; 內層總共執行 $T_cT_r = \lceil N/B_c\rceil\lceil N/B_r\rceil$ 次, 所以

O\!\left(\frac{N^2}{B_cB_r}\cdot B_rB_cd\right) = O(N^2d)

Algorithm 2：加上 mask 與 dropout 的完整版#

實際的 kernel 還要處理:

\mathbf{S} = \tau\mathbf{QK}^\top,\quad \mathbf{S}^{\text{masked}} = \text{mask}(\mathbf{S}),\quad \mathbf{P} = \text{softmax}(\mathbf{S}^{\text{masked}}),\quad \mathbf{P}^{\text{dropped}} = \text{dropout}(\mathbf{P}, p),\quad \mathbf{O} = \mathbf{P}^{\text{dropped}}\mathbf{V}

其中 $\tau$ 通常是 $1/\sqrt{d}$ , Algorithm 1 為了簡潔省略了它

Forward pass 回傳： $\mathbf{O}, \ell, m, R$ 。

Algorithm 4: backward pass#

1
for j = 1 … T_c:
2
    載入 K_j, V_j → SRAM
3
    在 SRAM 初始化 dK̃_j = 0, dṼ_j = 0        ← 留在 SRAM 累加
4

5
    for i = 1 … T_r:
6
        載入 Q_i, O_i, dO_i, dQ_i, ℓ_i, m_i
7

8
        S_ij = τ Q_i K_jᵀ
9
        S_ij^masked = mask(S_ij)
10
        P_ij = diag(ℓ_i)⁻¹ exp(S_ij^masked − m_i)     ← 重算 P（只用 ℓ 和 m）
11

12
        Z_ij = 重新生成的 dropout mask                 ← 用存下的 R
13
        P_ij^dropped = P_ij ∘ Z_ij
14

15
        dṼ_j  += (P_ij^dropped)ᵀ dO_i
16
        dP_ij^dropped = dO_i V_jᵀ
17
        dP_ij = dP_ij^dropped ∘ Z_ij
18
        D_i   = rowsum(dO_i ∘ O_i)
19
        dS_ij = P_ij ∘ (dP_ij − D_i)
20

21
        寫回 dQ_i ← dQ_i + τ dS_ij K_j  → HBM        ← 唯一要反覆讀寫 HBM 的梯度
22
        dK̃_j += τ dS_ijᵀ Q_i                        ← 留在 SRAM
23

24
    寫回 dK_j, dV_j → HBM                            ← 各只寫一次

d\mathbf{V} = \mathbf{P}^\top d\mathbf{O} \;\Longrightarrow\; dv_j = \sum_i P_{ij}\,do_i

d\mathbf{P} = d\mathbf{O}\mathbf{V}^\top \;\Longrightarrow\; dP_{ij} = do_i^\top v_j

softmax 的 Jacobian 是 $\text{diag}(y) - yy^\top$ , 所以:

dS_{i:} = \big(\text{diag}(P_{i:}) - P_{i:}P_{i:}^\top\big)dP_{i:} = P_{i:} \circ dP_{i:} - (P_{i:}^\top dP_{i:})P_{i:}

定義一個 $D_i := P_{i:}^\top dP_{i:}$ , 可以得:

$D_i = \sum_j \frac{e^{q_i^\top k_j}}{L_i}\,do_i^\top v_j = do_i^\top \sum_j \frac{e^{q_i^\top k_j}}{L_i}v_j = do_i^\top o_i$

$D_i$ 需要對整列的 $\mathbf{P}$ 和 $d\mathbf{P}$ 做 reduce, 但這樣根本塞不進 SRAM, 化簡後只需要兩個長度 $d$ 的向量做內積

可以得 $dS_{ij} = P_{ij}(dP_{ij} - D_i)$ , 接著可得:

dq_i = \sum_j dS_{ij}k_j,\qquad dk_j = \sum_i dS_{ij}q_i

I/O Complexity#

設 $d \le M \le Nd$ , 標準 attention 需要 $\Theta(Nd + N^2)$ 次 HBM 存取, FlashAttention 需要 $\Theta(N^2d^2M^{-1})$ 次

Proof:

K 和 V 的每個元素只被載入一次
Q 和 O 則要被掃過 $T_c$ 遍, 每遍 $\Theta(Nd)$ 個元素
總計 $\Theta(Nd + Nd\,T_c) = \Theta(Nd\,T_c)$

分塊大小的三個 SRAM 約束:

\underbrace{B_cd = O(M)}_{\mathbf{K}_j, \mathbf{V}_j},\qquad \underbrace{B_rd = O(M)}_{\mathbf{Q}_i, \mathbf{O}_i},\qquad \underbrace{B_rB_c = O(M)}_{\mathbf{S}_{ij}}

可以得 $B_c = \Theta(M/d)$ ， $B_r = \Theta\big(\min(M/d,\ M/B_c)\big) = \Theta\big(\min(M/d,\ d)\big)$ , 接著可得:

T_c = \frac{N}{B_c} = \Theta\!\left(\frac{Nd}{M}\right) \;\Longrightarrow\; \Theta(Nd\,T_c) = \Theta\!\left(\frac{N^2d^2}{M}\right)

典型情況 $d = 64$ – $128$ , $M \approx 100$ KB, $d^2 \ll M$ , 所以 HBM 存取少了很多倍 (最多 9×)

下界: 在 $M \in [d, Nd]$ 的範圍內，不存在任何演算法能用 $o(N^2d^2M^{-1})$ 次 HBM 存取算出精確 attention

假設存在這樣的演算法, 那取 $M = \Theta(Nd)$ , 他的存取次數會是 $o(N^2d^2/(Nd)) = o(Nd)$ , 但 Q、K、V、O 本身就佔 $Nd$ 的空間且一開始就在 HBM, 任何演算法至少得碰它們一次 → $\Omega(Nd)$ (矛盾)

反向的 IO 複雜度同樣是 $\Theta(N^2d^2M^{-1})$