FlashAttention 与 Triton：为什么 Attention 还能更快

本章定位：第二批第二模块（Kernel Bridge）。第 15 章建立了 Roofline 心智模型，本章是它的第一个真实案例：为什么 Attention 的数学没变、FLOPs 没少，却可以显著更快？

前置：第 7 章（Attention 数学）、第 15 章（GPU 存储层次 / Roofline）。本章不重讲 Attention 数学，只回答「数据在哪、怎么搬」。

19.1 换一个视角：从「数学」到「数据在哪」

回顾 Attention 三步（数学完全不变）：

新增的问题：$S$ 是 $S_{len} \times S_{len}$ 的矩阵。在 naive / eager 分解实现（显式执行 scores = QKᵀ → softmax → PV 三个步骤）里，它被物化（materialize）到 HBM，然后 softmax 再读回来、写回去，再和 V 相乘……数据在 HBM 与计算单元之间来回搬运。

19.2 Naive（Eager 分解）Attention 的真实问题：HBM I/O

scores = QKᵀ     → scores 写入 HBM
P = softmax(scores)  → 读回 → 再写 HBM
O = PV           → 再读回

F.scaled_dot_product_attention(q, k, v)   # PyTorch SDPA

① S = QKᵀ          → 写出 S（S² 个元素）到 HBM
② softmax(S)       → 从 HBM 读 S → 计算 → 把 P 写回 HBM
③ O = P·V          → 从 HBM 读 P（和 V）→ 计算 → 写出 O

19.3 HBM vs SRAM：第 15 章的 Roofline 终于用上了

Roofline 联系：朴素 Attention 的算术强度低（第 15 章表中「数 ~ 数十」），落在带宽墙上；分块融合把算术强度抬高，性能向上移动。

19.4 Tiling：分块计算

for each Q_tile (在 SRAM 中):
    for each K_tile, V_tile:
        载入 SRAM → 算局部 scores → 更新 running 统计 → 累积到输出
    输出 O_tile 写出 HBM（只写一次）

19.5 Online Softmax：本章最难但必须讲懂的地方

普通 softmax（一次看全）

Online 版本（分块看）

维护三个量：running max $m$、running denominator $\ell$、partial output $o$。新块到来时：

关键：旧结果乘上修正因子 $e^{m - m_{new}}$，等价于「用新 max 重新归一化」。

超小数值例子：scores = [1, 2] 与 [3, 4]

一次性计算的 softmax：

Online 分两块算：

19.6 FlashAttention 的核心（正确表述）

19.7 FA1 → FA2：高层演进（不深入 kernel 细节）

学习顺序建议：先把「tiling + online softmax + 不物化」讲清楚，再去看论文细节。kernel 级的调度优化属于工程深入，不影响对核心思想的理解。

19.8 什么时候收益最大

在很短的序列上（$S$ 小），$S^2$ 本来就不大，收益有限——这不是「永远快 3 倍」的魔法。

19.9 Triton：定位

为什么需要 Triton：像本章这种「多个算子融合成一个 kernel」的优化（softmax 的 max/sub/exp/sum/div 融合、attention 的 QKᵀ+softmax+PV 融合），PyTorch 组合写法会产生多次 HBM 往返；Triton 让你在一个 kernel 里完成。

19.10 第一个 Triton Kernel：Vector Add

# 类型：【Runnable on CUDA】需要 GPU + pip install triton
import torch, triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)                     # 我是第几个 program（块）
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 本块负责的元素下标
    mask = offsets < n_elements                # 边界掩码（最后一块可能不满）
    x = tl.load(x_ptr + offsets, mask=mask)    # 从 HBM 载入本块
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)   # 写回本块

def add(x, y):
    out = torch.empty_like(x)
    n = x.numel()
    grid = lambda meta: ((n + meta["BLOCK_SIZE"] - 1) // meta["BLOCK_SIZE"],)
    add_kernel[grid](x, y, out, n, BLOCK_SIZE=1024)
    return out

x = torch.randn(1_000_000, device="cuda")
y = torch.randn(1_000_000, device="cuda")
assert torch.allclose(add(x, y), x + y)

视觉对应：

N = 10,000,000，BLOCK_SIZE = 1024
→ grid = 9766 个 program，每个 program 处理 1024 个连续元素
program 0: offsets [0..1023]
program 1: offsets [1024..2047]
...
最后一个 program: mask 掉越界元素

19.11 Triton Softmax：融合的意义

# 类型：【Runnable on CUDA】row-wise softmax 的 fused kernel（教学简化版）
@triton.jit
def softmax_kernel(out_ptr, in_ptr, row_stride, n_cols, BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)                     # 每个 program 处理一行
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols
    x = tl.load(in_ptr + row * row_stride + cols, mask=mask, other=-float("inf"))
    x = x - tl.max(x, axis=0)                  # 稳定 softmax：减最大值
    num = tl.exp(x)
    denom = tl.sum(num, axis=0)
    tl.store(out_ptr + row * row_stride + cols, num / denom, mask=mask)

对比 PyTorch 组合写法：

# PyTorch：5 个算子 → 可能 5 次以上的 kernel 启动与 HBM 往返
x = x - x.max(dim=-1, keepdim=True).values
e = torch.exp(x)
out = e / e.sum(dim=-1, keepdim=True)

在算子是小算子、带宽受限时（第 15 章：softmax 是 memory-bound），融合成一个 kernel 意味着中间结果留在片上、读写 HBM 的次数从多次降到一次左右。

19.12 Mini Lab：Attention Tiling Simulator

# 类型：【Runnable】依赖：numpy
# tiled_attention.py —— 分块注意力模拟器（CPU）
import numpy as np

def standard_attention(Q, K, V):
    S = Q @ K.T / np.sqrt(Q.shape[-1])
    P = np.exp(S - S.max(-1, keepdims=True))
    P = P / P.sum(-1, keepdims=True)
    return P @ V

def tiled_attention(Q, K, V, q_tile=4, k_tile=4):
    """分块 + online softmax：不物化完整 S×S 矩阵"""
    n, d = Q.shape
    out = np.zeros((n, d))
    for qi in range(0, n, q_tile):
        q = Q[qi:qi + q_tile]                 # [tq, d]
        m = np.full((q.shape[0], 1), -np.inf) # running max
        l = np.zeros((q.shape[0], 1))         # running denominator
        acc = np.zeros((q.shape[0], d))       # partial output
        for ki in range(0, n, k_tile):
            k = K[ki:ki + k_tile]             # [tk, d]
            v = V[ki:ki + k_tile]
            s = q @ k.T / np.sqrt(d)          # 局部 scores [tq, tk]
            m_new = np.maximum(m, s.max(-1, keepdims=True))
            alpha = np.exp(m - m_new)         # 旧结果修正因子
            # 更新统计与输出（全部在"片上"完成，s 不写回"HBM"）
            p = np.exp(s - m_new)
            l = l * alpha + p.sum(-1, keepdims=True)
            acc = acc * alpha + p @ v
            m = m_new
        out[qi:qi + q_tile] = acc / l         # 最后才归一化并写出
    return out

# 验证
rng = np.random.default_rng(0)
Q, K, V = rng.normal(size=(16, 8)), rng.normal(size=(16, 8)), rng.normal(size=(16, 8))
ref = standard_attention(Q, K, V)
sim = tiled_attention(Q, K, V, q_tile=4, k_tile=4)
print("max diff =", np.abs(ref - sim).max())    # < 1e-12
assert np.allclose(ref, sim, atol=1e-10)
print("Tiled attention 与 naive 分解 attention 数值一致 ✅")

19.13 Benchmark 规范（写进任何 kernel 对比）

# 类型：【Runnable on CUDA】
import torch, time

def bench(fn, warmup=10, rep=50):
    for _ in range(warmup):                      # ① warmup
        fn()
    torch.cuda.synchronize()                     # ② 同步（CUDA 异步！）
    t0 = time.perf_counter()
    for _ in range(rep):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / rep      # ③ 平均（更稳可用 median / min）