把 softmax 装进 SRAM

§2.1FlashAttention v1 · 把 softmax 装进 SRAM

在线 softmax · 一行不变的恒等式

给定已扫的旧最大值 $m$、旧分母 $\ell$ 与新一项 $s$，新值：

数学上 $\ell$ 和 $\sum e^{s_j - m'}$ 相等，所以可以拆任意多个块来算。 把 $V$ 那一侧的"加权和"也照样 rescale 即可。这个恒等式是整篇论文 的核心；剩下的工程都是怎么把它写进一个 GPU kernel。

教学版前向

# 教学版 FlashAttention 前向（PyTorch + Numpy 风格）
def flash_attn(Q, K, V, BQ=64, BK=64):
    N, d = Q.shape
    O = torch.zeros_like(Q)
    m = torch.full((N,), -float('inf'))
    l = torch.zeros((N,))
    for i in range(0, N, BQ):
        Qi = Q[i:i+BQ]
        Oi = torch.zeros((BQ, d))
        mi = torch.full((BQ,), -float('inf'))
        li = torch.zeros((BQ,))
        for j in range(0, N, BK):
            Kj, Vj = K[j:j+BK], V[j:j+BK]
            Sij = (Qi @ Kj.T) / d**0.5
            mij = Sij.max(dim=-1).values
            Pij = torch.exp(Sij - mij[:, None])
            mi_new = torch.maximum(mi, mij)
            scale  = torch.exp(mi - mi_new)
            li     = li * scale + Pij.sum(dim=-1)
            Oi     = Oi * scale[:, None] + Pij @ Vj
            mi     = mi_new
        O[i:i+BQ] = Oi / li[:, None]
    return O

Triton 视角 · 真实 kernel 的骨架

# 简化的 Triton 内核骨架（FA-2 风格，省略 mask/causal）
@triton.jit
def fa_fwd(Q_ptr, K_ptr, V_ptr, O_ptr,
           N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, D: tl.constexpr):
    pid_m = tl.program_id(0)           # 一个 program 处理 BLOCK_M 个 Q-行
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, D)
    # 把整块 Qi 一次 load 到 SRAM（不再回写）
    Qi = tl.load(Q_ptr + offs_m[:, None]*D + offs_d[None, :])
    mi = tl.full([BLOCK_M], -float('inf'), dtype=tl.float32)
    li = tl.zeros([BLOCK_M], dtype=tl.float32)
    Oi = tl.zeros([BLOCK_M, D], dtype=tl.float32)
    for start_n in range(0, N, BLOCK_N):
        offs_n = start_n + tl.arange(0, BLOCK_N)
        Kj = tl.load(K_ptr + offs_n[:, None]*D + offs_d[None, :])
        Vj = tl.load(V_ptr + offs_n[:, None]*D + offs_d[None, :])
        S = tl.dot(Qi, tl.trans(Kj)) * (1.0 / tl.sqrt(D))
        mij = tl.max(S, axis=1)
        mi_new = tl.maximum(mi, mij)
        Pij = tl.exp(S - mi_new[:, None])
        scale = tl.exp(mi - mi_new)
        li = li * scale + tl.sum(Pij, axis=1)
        Oi = Oi * scale[:, None] + tl.dot(Pij.to(Vj.dtype), Vj)
        mi = mi_new
    Oi = Oi / li[:, None]
    tl.store(O_ptr + offs_m[:, None]*D + offs_d[None, :], Oi)

v1 在 A100 上把 GPT-2 1.5B 训练里 attention 的耗时从 50% 砍到 5%， 内存 $O(S^2) \to O(S)$。从此长上下文成为可能。

§2.2FlashAttention v2 · 把工作搬到正确的轴上

§2.3FlashAttention v3 · Hopper / FP8 / 异步

派生分支 · 让 FlashAttention 适应 "怪" mask

§2.4PagedAttention · 把 KV 当虚拟内存

# Page table 长这样（最小化伪代码）
class Sequence:
    block_table: list[int]   # 物理 block 编号
    seq_len:     int

# 当 seq 推进一个 token：
if seq.seq_len % BLOCK == 0:
    seq.block_table.append(allocator.alloc())
write_kv(physical_blocks[seq.block_table[-1]], offset=seq.seq_len % BLOCK)

# Attention kernel 接受 block_table，做 indirect address lookup
attn_output = paged_attention_kernel(
    q, k_cache_phys, v_cache_phys, block_table, seq_len)

§2.5Ring Attention · 跨卡接力，长上下文不再卡死

# Ring Attention 一次 forward 的骨架（教学版）
def ring_attn_forward(Q_local, K_local, V_local, rank, world_size):
    """
    每张卡持有 [S/P, D] 的 Q/K/V。
    一共 world_size 步，每步:
      1. 用本地 K/V 做一次 FlashAttention 累加进 O；
      2. 把 K/V 异步发给右邻居、同时从左邻居收 K/V；
      3. 换上新收到的 K/V，进入下一步。
    通信和计算 overlap, 总通信量 O(S·D), 与 P 无关。
    """
    O = torch.zeros_like(Q_local)
    K, V = K_local, V_local
    m = torch.full((Q_local.shape[0],), -float('inf'))
    l = torch.zeros_like(m)
    for step in range(world_size):
        # 异步 P2P
        K_next, V_next = isend_irecv(K, V, dst=(rank+1) % world_size,
                                      src=(rank-1) % world_size)
        # 与此同时本地 FA 累加
        O, m, l = flash_attn_accumulate(Q_local, K, V, O, m, l)
        wait(K_next, V_next)
        K, V = K_next, V_next
    return O / l[:, None]

§2.6稀疏 · 线性 · Mamba · NSA · attention 之外

上面三节是把 dense $O(S^2)$ 算得更快。 另一条路是从复杂度上把它降下来：