FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

URL

TL;DR

Flash Attention 是标准 Attention 的一种完全数学等价的高效实现
标准 Self-attention $Attention(Q,K,V)=softmax({\frac{QK^T}{\sqrt d}})V$ 是 $\mathcal{O}(n^2)$ 时间复杂度的，随着大模型的发展，序列长度 n 越来越大，带来的时间复杂度和 IO 量越来越恐怖
GPU 或 NPU 等计算设备的 SRAM 空间有限，无法放下所有后续依赖的计算中间结果，而 softmax 运算又需要全局信息（无法直接简单 tiling），所以这里存在巨大的 IO 量（HBM (high bandwidth memory) 和 SRAM 之间）
本文设计了一种 tiling 的局部计算 self-attention 的方法，这种方法在 forward 和 backward 中都可以使用
在 forward 中，FlashAttention 会计算并存储 softmax 操作的归一化因子（下图中的 l 和 m），在 backward 中，利用这些归一化因子，重新计算所需的中间注意力矩阵，而不需要从 HBM 中读取存储的完整矩阵
上述两点改进可以保证和标准 self-attention 计算结果（包括 forward 和 backward）完全一致的情况下，极大降低 IO 量（代价是计算量的小幅上涨）

Algorithm

1. Standard attention implementation

这里之所以要频繁和 HBM 以 block 为单位交互是因为有一个限制是 $b\le M \le Nd$ ，即 SRAM 可存储的数据量在 d 和 Nd 之间
IO 总量：
- 读： $2N^2+3Nd$ $2 N^{2} + 3 N d$
  - 读 Q / K / V
  - 读 S / P
- 写： $2N^2+Nd$ $2 N^{2} + N d$
  - 写 S / P
  - 写 O

2. Flash attention forward implementation

IO 总量：
- 读： $2Nd+2T_cNd$ $2 N d + 2 T_{c} N d$
  - 读 K / V 一次
  - 读 Q / O $T_c$ 次
- 写： $T_cNd+2T_cN$ $T_{c} N d + 2 T_{c} N$
  - 写 O $T_c$ 次
  - 写 l / m $T_c$ 次
由上述可知：
- $4\le T_c\le 4N$
- 因此，读写上下限分别为：
- 读：
  - 上限： $2Nd+8N^2d$
  - 下限： $10Nd$
- 写：
  - 上限： $4N^2d+8N^2$
  - 下限： $4Nd+8N$
- Block size 越大，总 IO 量越小

3. 详解 `Flash attention` 过程

标准 softmax 过程： $softmax(x_{ij})=\frac{e^{x_{ij}-max(x)}}{\sum_{i=0}^{N}\sum_{j=0}^{N}e^{x_{ij}-max(x)}}$ ，为了计算过程的数值稳定，需要对每个值减去序列最大值
Flash attention 中，会将 QKV 都在时序维度切分成若干块，然后在小块之间计算 Attention，直接计算是不等价的，因为需要全局信息，例如：
- 全局最大值 m(x)
- 指数累加和 l(x)

import torch
import math

BLOCK_SIZE = 1024
NEG_INF = -1e10  # -infinity
EPSILON = 1e-10


def flash_attention_forward(Q, K, V):
    batch, seq_len, dim = Q.shape
    O = torch.zeros_like(Q, requires_grad=True, device=Q.device)
    l = torch.zeros((batch, seq_len, 1), device=Q.device)  # 1 for broadcasting
    m = torch.ones((batch, seq_len, 1), device=Q.device) * NEG_INF

    Q_BLOCK_SIZE = min(BLOCK_SIZE, dim)
    KV_BLOCK_SIZE = BLOCK_SIZE

    Q = torch.split(Q, Q_BLOCK_SIZE, dim=1)
    K = torch.split(K, KV_BLOCK_SIZE, dim=1)
    V = torch.split(V, KV_BLOCK_SIZE, dim=1)
    O = list(torch.split(O, Q_BLOCK_SIZE, dim=1))
    l = list(torch.split(l, Q_BLOCK_SIZE, dim=1))
    m = list(torch.split(m, Q_BLOCK_SIZE, dim=1))

    Tr = len(Q)
    Tc = len(K)

    scale = 1 / math.sqrt(dim)

    for j in range(Tc):
        for i in range(Tr):
            Qi_scaled = Q[i] * scale
            S_ij = torch.einsum("... i d, ... j d -> ... i j", Qi_scaled, K[j])
            m_ij, _ = torch.max(S_ij, dim=-1, keepdim=True)
            P_ij = torch.exp(S_ij - m_ij)
            mi_new = torch.maximum(m_ij, m[i])
            l_ij = (
                torch.sum(P_ij * torch.exp(m_ij - mi_new), dim=-1, keepdim=True)
                + EPSILON
            )
            P_ij_Vj = torch.einsum(
                "... i j, ... j d -> ... i d", P_ij * torch.exp(m_ij - mi_new), V[j]
            )
            li_new = torch.exp(m[i] - mi_new) * l[i] + l_ij
            O[i] = (O[i] * l[i] * torch.exp(m[i] - mi_new) + P_ij_Vj) / li_new
            l[i] = li_new
            m[i] = mi_new

    O = torch.cat(O, dim=1)
    l = torch.cat(l, dim=1)
    m = torch.cat(m, dim=1)
    return O, l, m

def flash_attention_backward(Q, K, V, O, l, m, dO):
    # 参考代码中进行了实现，但这里省略，重点关注一下 backward 过程的输入
    pass

if __name__ == "__main__":
    torch.random.manual_seed(0)

    # batch, seq_len, dim
    Q = torch.randn(2, 4096, 128, requires_grad=True).to(device="cuda")
    K = torch.randn(2, 4096, 128, requires_grad=True).to(device="cuda")
    V = torch.randn(2, 4096, 128, requires_grad=True).to(device="cuda")

    out = flash_attention_forward(Q, K, V)
    print(out[0].sum().item())

这里最重要也最难理解的一行代码是：O[i] = (O[i] * l[i] * torch.exp(m[i] - mi_new) + P_ij_Vj) / li_new，这个仔细分析一下
- O[i] * l[i] 是计算得到上次外循环（j-1）时的分子
- O[i] * l[i] * torch.exp(m[i] - mi_new) 是把 j-1 时刻存储的信息更新到最新（调整 max_value）
- O[i] * l[i] * torch.exp(m[i] - mi_new) + P_ij_Vj 因为 O 的初始值是 0，因此这个公式可以拆解为： $O_i = \sum_{j=1}^{T_c} \frac{P_{ij} * V_j}{l_i}$ ，通过不断更新 max_value 和累加求和得到最终结果

Thought

Flash Attention 的效率和 block size 相关，block size 越大，效率提升越明显；在一些 block size 较小的情况下，效率可能会比标准 Attention 更差
Flash Attention 目标是解决 Memory hierarchy 情况下标注 Attention 计算访存比较低导致 IO 开销大的问题，是一种以时间换空间的做法；当 SRAM 相比模型足够大时，标准 Attention 效率比 Flash Attention 更高
backward 过程也是通过几乎相同的方式计算的，forward 阶段计算得到的中间 l / m 都会用于计算 backward

URL

TL;DR

Algorithm

1. Standard attention implementation

2. Flash attention forward implementation

3. 详解 Flash attention 过程

Thought

3. 详解 `Flash attention` 过程