第 13 课：GPU 内核与 FlashAttention

对应官方章节：Diving into the GPUs（A primer on GPUs / Improving performance with kernels / Fused kernels / FlashAttention）+ 附录 A1（kernel 剖析）
前置：第 2 课（激活、softmax 的注意力）、第 8 课（online softmax 概念）
状态：未开始

本课目标

完成后你需要能够：

画出 GPU 的层次结构（SM/warp/block、寄存器/共享内存/L1/L2/HBM）并说明带宽差异。
解释 memory coalescing、tiling、thread coarsening、控制分歧各自解决的问题。
说明 torch.compile → Triton → CUDA 的工具谱系与取舍。
用代码实现 FlashAttention 的 online softmax 分块算法并验证与标准实现一致。

核心概念

1. GPU 硬件微架构层次与存储金字塔（Memory Hierarchy）

要写出达到硬件峰值利用率的深度学习内核，必须深入理解现代 GPU 的物理组织。以 NVIDIA Hopper（H100）架构为例：

计算拓扑（SM / Warp / Thread）：
- 一张 GPU 芯片由若干个流式多处理器（Streaming Multiprocessor, SM）构成（H100 SXM 配备 132 个 SM）。
- 每个 SM 包含若干个调度器与执行单元（如 Tensor Core、FP32/FP64 CUDA Core、SFU 特殊函数单元）。
- 硬件执行的最基本不可分割调度单位是 Warp（线程束，固定为 32 个线程），Warp 内部 32 个线程严格以单指令多线程（SIMT）步调锁步执行。若干个 Warp 组织为逻辑上的 Thread Block（线程块，如 128~1024 线程），调度驻留在一个 SM 上。
物理存储金字塔（带宽与延迟的极端鸿沟）：
1. 寄存器堆（Registers）：单 SM 容量约 256 KB，全芯片总带宽高达 $\approx 30 \sim 40\text{ TB/s}$ ，0~1 个时钟周期延迟，由线程独占；
2. 共享内存 / L1 缓存（Shared Memory / L1 Cache, SRAM）：单 SM 容量约 228 KB，物理带宽高达 $\approx 15 \sim 20\text{ TB/s}$ ，约 20 周期延迟，供同一个 Block 内的所有线程高速协同共享；
3. L2 缓存（L2 Cache）：全芯片共享（H100 为 50 MB），带宽 $\approx 5 \sim 8\text{ TB/s}$ ，约 100~200 周期延迟；
4. 高带宽全局显存（High Bandwidth Memory, HBM3）：容量 80GB，虽然标称带宽高达 $3.35\text{ TB/s}$ ，但其访问延迟高达 $400 \sim 800$ 个时钟周期

内存墙第一性原理： H100 的 Tensor Core 标称计算峰值高达近 $1000\text{ TFLOPS}$ ，但外部 HBM 带宽仅为 $3.35\text{ TB/s}$ 。若一个算子每加载 1 字节数据只做 1 次运算（算术强度极低），硬件必然陷入访存受限（Memory-Bound），绝大部分 SM 将长期处于等待内存数据回包的饥饿停顿状态。

2. 经典 GPU 内核优化四大铁律

高效算子开发的核心目标，是竭尽一切可能将数据截留在高速片上缓存（Registers 与 SRAM）中反复复用，消灭对低速 HBM 的往返读写：

访存合并（Memory Coalescing）：
- GPU 全局内存控制器以固定突发（Burst，通常为 32 或 128 字节对齐的缓存行 Cache Line）为单位进行事务传输。
- 当一个 Warp 内的 32 个线程并发访问连续且对齐的内存地址时，硬件能够将这 32 个内存请求合并为单次 128 字节的物理传输事务，总线利用率达到 100%。
- 若线程访问地址产生跨度（Strided）或随机离散访问，硬件将被迫发射多达 32 次独立的孤立内存请求，导致总线利用率跌破 10%，引发近十倍的吞吐断崖。
分块与共享内存复用（Tiling & Shared Memory Staging）：
- 在矩阵乘法（GEMM）中，利用分块算法将大矩阵拆解为适应 SRAM 容量的子块（Tile）。
- 将输入 Tile 从 HBM 协同搬运到 Shared Memory 后，块内各个线程反复读取该 Tile 进行乘累加，使 HBM 访存次数降低 $\text{BlockSize}$ 倍。
线程粗化与寄存器分块（Thread Coarsening & Register Tiling）：
- 让每个 CUDA 线程不再仅负责计算 1 个输出元素，而是负责计算一个微型子块（例如 $4 \times 4$ 或 $8 \times 8$ 个输出）。
- 累加过程中的中间变量全程保存在线程私有寄存器中，消除对 Shared Memory 的读写冲突（Bank Conflicts）与指令发射瓶颈。
消除分支分歧（Warp Divergence Elimination）：
- 避免在 Warp 内部编写基于线程号的条件分支（if (threadIdx.x % 2 == 0)）。
- 一旦同一 Warp 内线程执行路径分叉，硬件只能串行执行所有分支路径并用掩码屏蔽无效线程，导致算力有效利用率腰斩。

3. 算子开发工具谱系的权衡与演进

开发体系	易用性与开发效率	性能天花板	底层掌控力与灵活性	典型适用场景
Eager PyTorch	极高（Python 语法）	较低（缺乏算子融合，中间激活频繁写回 HBM）	极低（黒盒封装）	算法快速验证、原型设计
`torch.compile` (Inductor)	高（单行装饰器）	较高（自动生成融合 Triton Kernel）	中等（支持导出代码二次修改）	标准模型的一键端到端加速
OpenAI Triton	中（Block 级 Python 编程）	极高（逼近手写 CUDA 90%~95% 性能）	高（显式控制 Block、分块加载、Masking）	自定义算子开发、前沿注意力实现
CUDA C++ / CUTLASS	低（极其繁琐复杂）	最高（100% 榨干硬件极致潜力）	绝对掌控（寄存器分配、异步管道 TMA、Warp 级 MMA）	基础算子库底层核心、极致性能压榨

4. 算子融合（Operator Fusion）与 FlashAttention 的划时代变革

传统的 Attention 机制在数学上定义为： $\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{Q K^T}{\sqrt{d}}\right) V$

A. 传统自注意力的访存灾难（如图 13-1 所示）

传统注意力 HBM 访存灾难

Q, K, V (显存)

得分矩阵 S = QK^T
(写回显存: O(N^2))

Softmax 概率 P
(写回显存: O(N^2))

输出 O = PV (写回显存)

FlashAttention 片上融合

GPU SRAM 高速缓存
(瓦片分块)

在线流式 Softmax
(零中间显存写回)

直接输出最终 O

图 13-1：传统注意力大量读写高带宽内存 vs FlashAttention 片上瓦片融合对比

在原生 PyTorch 实现中，该公式被拆解为三个分离的物理内核：

GEMM 1：从 HBM 加载 $Q$ 与 $K$ ，计算注意力点积分数 $S = Q K^T$ ，将庞大的 $N \times N$ 矩阵完整写入 HBM（ $O(N^2)$ 显存）；
Softmax：从 HBM 重新完整读出 $S$ ，计算归一化权重矩阵 $P = \text{Softmax}(S)$ ，再次将 $N \times N$ 矩阵完整写回 HBM；
GEMM 2：从 HBM 重新读出 $P$ 与 $V$ ，计算上下文输出 $O = P V$ ，最终写回 HBM。

痛点曝光：长序列下 $N \times N$ 矩阵极其庞大（例如 $N=16\text{k}$ 时，单头浮点数高达 2.56 亿个）。整个过程对极其缓慢的 HBM 进行了多达 3 次往返读写。计算时间完全被 HBM 访存延迟统治。

B. FlashAttention 的破局机制（SRAM 瓦片融合 + Online Softmax）

Dao et al. 提出的 FlashAttention 颠覆了上述流程（如图 13-2 所示）：

高带宽显存 (HBM)

片上共享内存 (SRAM)

Tensor Core 计算核心

1分块加载 Q_i 瓦片 (驻留外层循环)

loop内层循环遍历 K, V 瓦片 (j = 1 to T_c)

2分块加载 K_j, V_j

3计算局部得分 S_ij = Q_i K_j^T

4在线更新局部最大值 m_i 与归一化分母 l_i

5增量累加更新注意力输出 O_i

6将最终归一化结果 O_i 写回全局显存

图 13-2：FlashAttention 前向分块双层循环流水与增量更新流

不物化（No Materialization）：不在 HBM 中生成完整的 $N \times N$ 矩阵。
分块载入 SRAM：将 $Q$ 划分为行块 $Q_i \in \mathbb{R}^{B_r \times d}$ ，将 $K, V$ 划分为列块 $K_j, V_j \in \mathbb{R}^{B_c \times d}$ 。将这些紧凑的小块载入片上共享内存（SRAM）。
增量式 Online Softmax：在 SRAM 内部计算局部点积分数 $S_{ij} = Q_i K_j^T$ ，并实时维护行最大值 $m$ 与归一化分母 $\ell$ 。当遍历下一个 $K, V$ 块时，利用代数恒等式，直接将旧输出按指数比值 $\exp(m_{\text{old}} - m_{\text{new}})$ 重新缩放，就地累加新块贡献。
计算效益：
- HBM 访存量从标准实现的 $\Theta(Nd + N^2)$ 降为 $\Theta\!\left(\frac{N^2 d^2}{M}\right)$ （ $M$ 为 SRAM 容量，常见 $d^2 \ll M$ ，因此远小于前者；仍随 $N$ 二次增长，并不是 $O(N)$ ）；
- 除输入输出外的额外显存从 $O(N^2)$ 降为 $O(N)$ （每行只保留 $m, \ell$ 统计量）；
- 浮点计算量（FLOPs）基本不变，但 HBM 往返大幅减少，原论文报告墙钟速度提升约 $2\times \sim 4\times$ 。

具体演示与数值推演

以单头注意力计算为例，设输入序列长度 $N = 4096$ ，特征维度 $d = 128$ ，采用 BF16 精度（每元素 2 字节）：

传统原生注意力的显存与访存量推演：
- 中间分数矩阵 $S$ 尺寸为 $4096 \times 4096$ ，需物化存储： $\text{Memory}_{S} = 4096 \times 4096 \times 2\text{ 字节} = 33,554,432\text{ 字节} = 32\text{ MiB} \approx 33.6\text{ MB}$
- 加上 Softmax 权重矩阵 $P$ （同为 $32\text{ MiB}$ ），单头单层即产生 $64\text{ MiB}$ 的临时中间显存。若网络拥有 32 个注意力头与 32 层深度，全网单样本的注意力中间激活为： $\text{Total Memory} = 32 \times 32 \times 64\text{ MiB} = 65,536\text{ MiB} = 64\text{ GiB} \approx 68.7\text{ GB}$
- 显存直接被中间激活耗尽。同时，读写这批矩阵需要往返吞吐数百 GB 的 HBM 流量。
FlashAttention 分块融合推演：
- 设置 SRAM 分块尺寸 $B_r = B_c = 64$ 。
- 单个分块在片上 SRAM 中仅占用： $\text{Tile Memory} = 64 \times 64 \times 2\text{ 字节} = 8192\text{ 字节} = 8\text{ KB}$
- 8 KB 的小张量可以完全装进每个 SM 的 Shared Memory（容量高达 228 KB）中。
- 全过程没有 $N \times N$ 中间矩阵写回 HBM，原本 $64\text{ GiB}$ 的二次方激活不再需要保存（只保留线性规模的输出 $O$ 与每行统计量，反向时在片上重算）。这为训练省出几十 GiB 显存，也使注意力层更接近算力受限（Compute Bound）。

代码填空题

用 numpy 实现分块 online softmax 注意力，验证与标准实现数值一致。

Code

import numpy as np

def standard_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray):
    """参照实现：物化 S 和 P。Q,K,V: (n, d)。"""
    n = Q.shape[0]
    S = Q @ K.T                    # (n, n) —— 长序列下这是内存炸弹
    S = S - S.max(axis=-1, keepdims=True)
    P = np.exp(S)
    P /= P.sum(axis=-1, keepdims=True)
    return P @ V, S

def flash_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray,
                    block_size: int = 64) -> np.ndarray:
    """
    分块实现：不物化完整 S。
    对每个 query 块，逐步处理 kv 块，维护每个 query 行的：
      m: 见过的分数的 running max
      l: 见过的 exp(score - m) 之和（softmax 分母）
      O: 已累加的加权输出
    """
    n, d = Q.shape
    O = np.zeros((n, d))
    l = np.zeros(n)                       # 分母（running sum）
    m = np.full(n, -np.inf)               # running max

    for i in range(0, n, block_size):
        Qi = Q[i:i + block_size]
        Oi = O[i:i + block_size]
        li = l[i:i + block_size]
        mi = m[i:i + block_size]

        for j in range(0, n, block_size):
            Kj = K[j:j + block_size]
            Vj = V[j:j + block_size]

            S_ij = Qi @ Kj.T              # 当前块分数 (bs, bs)

            # online softmax 更新：
            # 1) 新 running max（把当前块的最大分数并进来）
            m_new = ______
            # 2) 旧分母/旧输出都要按 exp(mi - m_new) 缩放，新块贡献 p_ij
            scale = ______
            p_ij = np.exp(S_ij - m_new[:, None])   # 新块的"原始"权重
            l_new = ______                         # 分母更新：旧项缩放 + 新块和
            # 3) 输出更新：旧输出按比例缩放，加上新块贡献
            Oi = ______
            mi, li = m_new, l_new

        # 把块的统计量写回全局数组
        O[i:i + block_size] = Oi
        l[i:i + block_size] = li
        m[i:i + block_size] = mi

    # 最终归一化：O / l（把 running max 移回）
    return O / l[:, None]

if __name__ == "__main__":
    rng = np.random.default_rng(0)
    n, d = 256, 32
    Q = rng.standard_normal((n, d)) * 0.5
    K = rng.standard_normal((n, d)) * 0.5
    V = rng.standard_normal((n, d))

    O_ref, S_ref = standard_attention(Q, K, V)
    O_flash = flash_attention(Q, K, V, block_size=32)

    print("最大绝对误差:", np.abs(O_ref - O_flash).max())
    print("S 矩阵元素数:", S_ref.size, "（FlashAttention 不物化它）")
    assert np.allclose(O_ref, O_flash, atol=1e-6)
    print("与标准注意力一致 ✓（内存从 O(n²) 降到 O(n)）")

三个问答题

Q1

为什么”先算完整 S = QKᵀ 再 softmax”在长序列下既慢又占内存？请从 HBM 带宽与显存两个角度回答，并把它与第 2 课激活公式中的二次项联系起来。FlashAttention 最终在内存与访存上各省掉了什么？

Q2

教材的 naive matmul 中，为什么”每个线程算一个输出元素”会导致访存不合并？改成一维 block 映射后为什么吞吐提升约 10 倍？Tiling 为什么能进一步减少 HBM 访问？Thread coarsening 解决的是哪一类停顿（提示：MIO throttle）？

Q3

请对比 torch.compile、Triton、CUDA 三条写 kernel 的路径（难度/性能/灵活性），并说明什么时候必须离开 Triton 写 CUDA。FA-1 → FA-2 → FA-3 各自改进的重点是什么（各一句话）？

检查与通过标准

总分 10 分：代码正确 4 分（m_new/l_new/O 更新、归一化、与标准实现一致）、三题各 2 分、通过线 8 分。

一票否决项：

认为 FlashAttention 降低了计算量（它降低的是访存与显存）。
说不出 online softmax 为什么必须维护 running max（数值稳定性）。
认为”分别 softmax 再拼接”正确。
分不清 HBM 与共享内存在带宽上的量级差异。

💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先完成题目再核对；核对后必须解释关键公式，并改一个规模重新计算。

Code

import numpy as np

def standard_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray):
    """参照实现：物化 S 和 P。Q,K,V: (n, d)。"""
    n = Q.shape[0]
    S = Q @ K.T                    # (n, n) —— 长序列下这是内存炸弹
    S = S - S.max(axis=-1, keepdims=True)
    P = np.exp(S)
    P /= P.sum(axis=-1, keepdims=True)
    return P @ V, S

def flash_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray,
                    block_size: int = 64) -> np.ndarray:
    """
    分块实现：不物化完整 S。
    对每个 query 块，逐步处理 kv 块，维护每个 query 行的：
      m: 见过的分数的 running max
      l: 见过的 exp(score - m) 之和（softmax 分母）
      O: 已累加的加权输出
    """
    n, d = Q.shape
    O = np.zeros((n, d))
    l = np.zeros(n)                       # 分母（running sum）
    m = np.full(n, -np.inf)               # running max

    for i in range(0, n, block_size):
        Qi = Q[i:i + block_size]
        Oi = O[i:i + block_size]
        li = l[i:i + block_size]
        mi = m[i:i + block_size]

        for j in range(0, n, block_size):
            Kj = K[j:j + block_size]
            Vj = V[j:j + block_size]

            S_ij = Qi @ Kj.T              # 当前块分数 (bs, bs)

            # online softmax 更新：
            # 1) 新 running max（把当前块的最大分数并进来）
            m_new = np.maximum(mi, S_ij.max(axis=-1))
            # 2) 旧分母/旧输出都要按 exp(mi - m_new) 缩放，新块贡献 p_ij
            scale = np.exp(mi - m_new)
            p_ij = np.exp(S_ij - m_new[:, None])   # 新块的"原始"权重
            l_new = li * scale + p_ij.sum(axis=-1)                         # 分母更新：旧项缩放 + 新块和
            # 3) 输出更新：旧输出按比例缩放，加上新块贡献
            Oi = Oi * scale[:, None] + p_ij @ Vj
            mi, li = m_new, l_new

        # 把块的统计量写回全局数组
        O[i:i + block_size] = Oi
        l[i:i + block_size] = li
        m[i:i + block_size] = mi

    # 最终归一化：O / l（把 running max 移回）
    return O / l[:, None]

if __name__ == "__main__":
    rng = np.random.default_rng(0)
    n, d = 256, 32
    Q = rng.standard_normal((n, d)) * 0.5
    K = rng.standard_normal((n, d)) * 0.5
    V = rng.standard_normal((n, d))

    O_ref, S_ref = standard_attention(Q, K, V)
    O_flash = flash_attention(Q, K, V, block_size=32)

    print("最大绝对误差:", np.abs(O_ref - O_flash).max())
    print("S 矩阵元素数:", S_ref.size, "（FlashAttention 不物化它）")
    assert np.allclose(O_ref, O_flash, atol=1e-6)
    print("与标准注意力一致 ✓（内存从 O(n²) 降到 O(n)）")

补全后的代码（关键填空处）

# online softmax 更新
            m_new = np.maximum(mi, S_ij.max(axis=-1))        # 1) 新 running max
            scale = np.exp(mi - m_new)                       # 旧输出的缩放因子
            p_ij = np.exp(S_ij - m_new[:, None])             # 新块的"原始"权重
            l_new = li * scale + p_ij.sum(axis=-1)           # 2) 分母更新
            Oi = Oi * scale[:, None] + p_ij @ Vj             # 3) 输出更新

验证结果：

最大绝对误差: ~1e-15
S 矩阵元素数: 65536（标准实现物化；FlashAttention 不物化）
与标准注意力一致 ✓（内存从 O(n²) 降到 O(n)）

三个问答题要点

Q1

慢：S 与 P 都是 n×n，计算它们要写回 HBM 再读回（每步都要走 HBM 带宽）；n 大时访存时间远超计算（memory-bound）；
占内存：S 是 O(n²)，正是第 2 课激活公式中 5·n_heads·seq²·bs 项——长序列时是激活最大头，全量重计算也救不了（层边界激活仍随 seq 线性增长，但注意力分数矩阵是二次项）；
FA 省掉：不物化 S/P（O(n²) → O(n) 内存）；HBM 访存量从 Θ(nd + n²) 降到 Θ(n²d²/M)（M 为 SRAM 容量），额外显存从 O(n²) 降到 O(n)——计算量不变，但 memory-bound 瓶颈被大幅缓解。

Q2

naive matmul：线程 (x,y) 都算一个输出元素时，同 warp 的线程 x 相邻但读的是 A 的不同行（行距 = K）→ 不连续 → 无法合并成 DRAM burst；B 的列访问同理；
一维 block 映射（x = blockIdx·BLOCK + threadIdx/BLOCKSIZE，y = threadIdx%BLOCKSIZE）：warp 内相邻线程的 y 相邻 → 读 B 的连续列（行主序下连续）→ 合并；吞吐提升 ~10 倍；
Tiling：把 A/B 的 tile 先搬进共享内存，块内所有线程复用 → HBM 访问从”每个线程各读一遍”变成”每 tile 读一次”；
Thread coarsening：每个线程算多个输出 → 共享内存指令数减少 → 缓解 MIO 指令队列阻塞（ncu 的 mio_throttle 停顿）。

Q3

谱系：torch.compile（易用/快/不灵活，可导出 Triton 代码作起点）→ Triton（block 级、较快、较灵活）→ CUDA（最难、最快、完全可控）；
必须写 CUDA 的场景：需要精细控制共享内存布局、SM 内调度/warp 划分、特殊指令（如 FA-2/FA-3 的 warp 分工与 FP8 Tensor Core 路径），Triton 只暴露 block 级抽象；
FA-1→2→3：FA-1 确立分块 + online softmax + 融合（内存 O(n)）；FA-2 减少非 matmul 操作、按 warp 精细划分 Q/K/V 工作量（提高利用率）；FA-3 面向 Hopper 优化 FP8/Tensor Core 与异步流水线。

通过要点

FlashAttention 不降低 FLOPs，降低的是 HBM 访存（降为 Θ(n²d²/M)）与额外显存（O(n²)→O(n)）；
online softmax 的 running max 是为了数值稳定（防止 exp 溢出），running sum 是为了正确归一化；
与第 8 课 Ring Attention 共享同一套 online softmax 思想。

官方主参考

← 上一课
第 12 课：配置搜索与性能诊断 下一课 →
第 14 课：混合精度与综合面试

--- title: "第 13 课：GPU 内核与 FlashAttention" jupyter: python3 --- - 对应官方章节：Diving into the GPUs（A primer on GPUs / Improving performance with kernels / Fused kernels / FlashAttention）+ 附录 A1（kernel 剖析） - 前置：第 2 课（激活、softmax 的注意力）、第 8 课（online softmax 概念） - 状态：未开始 ## 本课目标完成后你需要能够： - 画出 GPU 的层次结构（SM/warp/block、寄存器/共享内存/L1/L2/HBM）并说明带宽差异。 - 解释 memory coalescing、tiling、thread coarsening、控制分歧各自解决的问题。 - 说明 torch.compile → Triton → CUDA 的工具谱系与取舍。 - 用代码实现 FlashAttention 的 online softmax 分块算法并验证与标准实现一致。 ## 核心概念 ### 1. GPU 硬件微架构层次与存储金字塔（Memory Hierarchy）要写出达到硬件峰值利用率的深度学习内核，必须深入理解现代 GPU 的物理组织。以 NVIDIA Hopper（H100）架构为例： - **计算拓扑（SM / Warp / Thread）**： - 一张 GPU 芯片由若干个**流式多处理器（Streaming Multiprocessor, SM）**构成（H100 SXM 配备 132 个 SM）。 - 每个 SM 包含若干个调度器与执行单元（如 Tensor Core、FP32/FP64 CUDA Core、SFU 特殊函数单元）。 - 硬件执行的最基本不可分割调度单位是 **Warp（线程束，固定为 32 个线程）**，Warp 内部 32 个线程严格以单指令多线程（SIMT）步调锁步执行。若干个 Warp 组织为逻辑上的 **Thread Block（线程块，如 128~1024 线程）**，调度驻留在一个 SM 上。 - **物理存储金字塔（带宽与延迟的极端鸿沟）**： 1. **寄存器堆（Registers）**：单 SM 容量约 256 KB，全芯片总带宽高达 $\approx 30 \sim 40\text{ TB/s}$，0~1 个时钟周期延迟，由线程独占； 2. **共享内存 / L1 缓存（Shared Memory / L1 Cache, SRAM）**：单 SM 容量约 228 KB，物理带宽高达 $\approx 15 \sim 20\text{ TB/s}$，约 20 周期延迟，供同一个 Block 内的所有线程高速协同共享； 3. **L2 缓存（L2 Cache）**：全芯片共享（H100 为 50 MB），带宽 $\approx 5 \sim 8\text{ TB/s}$，约 100~200 周期延迟； 4. **高带宽全局显存（High Bandwidth Memory, HBM3）**：容量 80GB，虽然标称带宽高达 $3.35\text{ TB/s}$，但其访问延迟高达 $400 \sim 800$ 个时钟周期 **内存墙第一性原理**： H100 的 Tensor Core 标称计算峰值高达近 $1000\text{ TFLOPS}$，但外部 HBM 带宽仅为 $3.35\text{ TB/s}$。若一个算子每加载 1 字节数据只做 1 次运算（算术强度极低），硬件必然陷入**访存受限（Memory-Bound）**，绝大部分 SM 将长期处于等待内存数据回包的饥饿停顿状态。 ### 2. 经典 GPU 内核优化四大铁律高效算子开发的核心目标，是竭尽一切可能**将数据截留在高速片上缓存（Registers 与 SRAM）中反复复用，消灭对低速 HBM 的往返读写**： 1. **访存合并（Memory Coalescing）**： - GPU 全局内存控制器以固定突发（Burst，通常为 32 或 128 字节对齐的缓存行 Cache Line）为单位进行事务传输。 - 当一个 Warp 内的 32 个线程并发访问连续且对齐的内存地址时，硬件能够将这 32 个内存请求**合并为单次 128 字节的物理传输事务**，总线利用率达到 100%。 - 若线程访问地址产生跨度（Strided）或随机离散访问，硬件将被迫发射多达 32 次独立的孤立内存请求，导致总线利用率跌破 10%，引发近十倍的吞吐断崖。 2. **分块与共享内存复用（Tiling & Shared Memory Staging）**： - 在矩阵乘法（GEMM）中，利用分块算法将大矩阵拆解为适应 SRAM 容量的子块（Tile）。 - 将输入 Tile 从 HBM 协同搬运到 Shared Memory 后，块内各个线程反复读取该 Tile 进行乘累加，使 HBM 访存次数降低 $\text{BlockSize}$ 倍。 3. **线程粗化与寄存器分块（Thread Coarsening & Register Tiling）**： - 让每个 CUDA 线程不再仅负责计算 1 个输出元素，而是负责计算一个微型子块（例如 $4 \times 4$ 或 $8 \times 8$ 个输出）。 - 累加过程中的中间变量全程保存在线程私有寄存器中，消除对 Shared Memory 的读写冲突（Bank Conflicts）与指令发射瓶颈。 4. **消除分支分歧（Warp Divergence Elimination）**： - 避免在 Warp 内部编写基于线程号的条件分支（`if (threadIdx.x % 2 == 0)`）。 - 一旦同一 Warp 内线程执行路径分叉，硬件只能串行执行所有分支路径并用掩码屏蔽无效线程，导致算力有效利用率腰斩。 ### 3. 算子开发工具谱系的权衡与演进 | 开发体系 | 易用性与开发效率 | 性能天花板 | 底层掌控力与灵活性 | 典型适用场景 | |---|---|---|---|---| | **Eager PyTorch** | 极高（Python 语法） | 较低（缺乏算子融合，中间激活频繁写回 HBM） | 极低（黒盒封装） | 算法快速验证、原型设计 | | **`torch.compile` (Inductor)** | 高（单行装饰器） | 较高（自动生成融合 Triton Kernel） | 中等（支持导出代码二次修改） | 标准模型的一键端到端加速 | | **OpenAI Triton** | 中（Block 级 Python 编程） | **极高（逼近手写 CUDA 90%~95% 性能）** | 高（显式控制 Block、分块加载、Masking） | 自定义算子开发、前沿注意力实现 | | **CUDA C++ / CUTLASS** | 低（极其繁琐复杂） | **最高（100% 榨干硬件极致潜力）** | 绝对掌控（寄存器分配、异步管道 TMA、Warp 级 MMA） | 基础算子库底层核心、极致性能压榨 | ### 4. 算子融合（Operator Fusion）与 FlashAttention 的划时代变革传统的 Attention 机制在数学上定义为： $$\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{Q K^T}{\sqrt{d}}\right) V$$ #### A. 传统自注意力的访存灾难（如**图 13-1** 所示） ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-lr"><div class="fd-sg" data-nid="Standard" data-depth="0"><div class="fd-sg-title">传统注意力 HBM 访存灾难</div><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-rank"><div class="fd-node fd-box" data-nid="QKV">Q, K, V (显存)</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="S">得分矩阵 S = QK^T<br>(写回显存: O(N^2))</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="P">Softmax 概率 P<br>(写回显存: O(N^2))</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="O">输出 O = PV (写回显存)</div></div></div></div></div></div><div class="fd-comp fd-lr"><div class="fd-sg" data-nid="FA" data-depth="0"><div class="fd-sg-title">FlashAttention 片上融合</div><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-rank"><div class="fd-node fd-box" data-nid="SRAM">GPU SRAM 高速缓存<br>(瓦片分块)</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Online">在线流式 Softmax<br>(零中间显存写回)</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Out">直接输出最终 O</div></div></div></div></div></div></div><i class="fd-e" data-s="QKV" data-t="S" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="S" data-t="P" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="P" data-t="O" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="SRAM" data-t="Online" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Online" data-t="Out" data-k="solid" data-h="1" data-b="0"></i></div> ``` <p class="caption" align="center"><em>图 13-1：传统注意力大量读写高带宽内存 vs FlashAttention 片上瓦片融合对比</em></p> 在原生 PyTorch 实现中，该公式被拆解为三个分离的物理内核： 1. **GEMM 1**：从 HBM 加载 $Q$ 与 $K$，计算注意力点积分数 $S = Q K^T$，将庞大的 $N \times N$ 矩阵完整**写入 HBM**（$O(N^2)$ 显存）； 2. **Softmax**：从 HBM 重新完整读出 $S$，计算归一化权重矩阵 $P = \text{Softmax}(S)$，再次将 $N \times N$ 矩阵完整**写回 HBM**； 3. **GEMM 2**：从 HBM 重新读出 $P$ 与 $V$，计算上下文输出 $O = P V$，最终写回 HBM。 - **痛点曝光**：长序列下 $N \times N$ 矩阵极其庞大（例如 $N=16\text{k}$ 时，单头浮点数高达 2.56 亿个）。整个过程对极其缓慢的 HBM 进行了多达 **3 次往返读写**。计算时间完全被 HBM 访存延迟统治。 #### B. FlashAttention 的破局机制（SRAM 瓦片融合 + Online Softmax） Dao et al. 提出的 FlashAttention 颠覆了上述流程（如**图 13-2** 所示）： ```{=html} <div class="fd fd-seq" role="img"><div class="sd-grid" style="--n:3"><div class="sd-life" style="grid-column:1;grid-row:1/10"></div><div class="sd-life" style="grid-column:2;grid-row:1/10"></div><div class="sd-life" style="grid-column:3;grid-row:1/10"></div><div class="sd-actor" style="grid-column:1;grid-row:1">高带宽显存 (HBM)</div><div class="sd-actor" style="grid-column:2;grid-row:1">片上共享内存 (SRAM)</div><div class="sd-actor" style="grid-column:3;grid-row:1">Tensor Core 计算核心</div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:1/3;grid-row:2;--k:2"><div class="sd-text"><span class="sd-num">1</span>分块加载 Q_i 瓦片 (驻留外层循环)</div><div class="sd-line"></div></div><div class="sd-frame" style="grid-column:1/4;grid-row:3"><span class="sd-kw">loop</span>内层循环遍历 K, V 瓦片 (j = 1 to T_c)</div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:1/3;grid-row:4;--k:2"><div class="sd-text"><span class="sd-num">2</span>分块加载 K_j, V_j</div><div class="sd-line"></div></div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:2/4;grid-row:5;--k:2"><div class="sd-text"><span class="sd-num">3</span>计算局部得分 S_ij = Q_i K_j^T</div><div class="sd-line"></div></div><div class="sd-msg sd-l sd-h-arrow" style="grid-column:2/4;grid-row:6;--k:2"><div class="sd-text"><span class="sd-num">4</span>在线更新局部最大值 m_i 与归一化分母 l_i</div><div class="sd-line"></div></div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:2/4;grid-row:7;--k:2"><div class="sd-text"><span class="sd-num">5</span>增量累加更新注意力输出 O_i</div><div class="sd-line"></div></div><div class="sd-frame-end" style="grid-column:1/4;grid-row:8"></div><div class="sd-msg sd-l sd-h-arrow" style="grid-column:1/3;grid-row:9;--k:2"><div class="sd-text"><span class="sd-num">6</span>将最终归一化结果 O_i 写回全局显存</div><div class="sd-line"></div></div></div></div> ``` <p class="caption" align="center"><em>图 13-2：FlashAttention 前向分块双层循环流水与增量更新流</em></p> - **不物化（No Materialization）**：**不在 HBM 中生成完整的 $N \times N$ 矩阵**。 - **分块载入 SRAM**：将 $Q$ 划分为行块 $Q_i \in \mathbb{R}^{B_r \times d}$，将 $K, V$ 划分为列块 $K_j, V_j \in \mathbb{R}^{B_c \times d}$。将这些紧凑的小块载入片上共享内存（SRAM）。 - **增量式 Online Softmax**：在 SRAM 内部计算局部点积分数 $S_{ij} = Q_i K_j^T$，并实时维护行最大值 $m$ 与归一化分母 $\ell$。当遍历下一个 $K, V$ 块时，利用代数恒等式，直接将旧输出按指数比值 $\exp(m_{\text{old}} - m_{\text{new}})$ 重新缩放，就地累加新块贡献。 - **计算效益**： - HBM 访存量从标准实现的 $\Theta(Nd + N^2)$ 降为 $\Theta\!\left(\frac{N^2 d^2}{M}\right)$（$M$ 为 SRAM 容量，常见 $d^2 \ll M$，因此远小于前者；仍随 $N$ 二次增长，并不是 $O(N)$）； - 除输入输出外的额外显存从 $O(N^2)$ 降为 $O(N)$（每行只保留 $m, \ell$ 统计量）； - 浮点计算量（FLOPs）基本不变，但 HBM 往返大幅减少，原论文报告墙钟速度提升约 **$2\times \sim 4\times$**。 ## 具体演示与数值推演以单头注意力计算为例，设输入序列长度 $N = 4096$，特征维度 $d = 128$，采用 BF16 精度（每元素 2 字节）： - **传统原生注意力的显存与访存量推演**： - 中间分数矩阵 $S$ 尺寸为 $4096 \times 4096$，需物化存储： $$\text{Memory}_{S} = 4096 \times 4096 \times 2\text{ 字节} = 33,554,432\text{ 字节} = 32\text{ MiB} \approx 33.6\text{ MB}$$ - 加上 Softmax 权重矩阵 $P$（同为 $32\text{ MiB}$），单头单层即产生 $64\text{ MiB}$ 的临时中间显存。若网络拥有 32 个注意力头与 32 层深度，全网单样本的注意力中间激活为： $$\text{Total Memory} = 32 \times 32 \times 64\text{ MiB} = 65,536\text{ MiB} = 64\text{ GiB} \approx 68.7\text{ GB}$$ - 显存直接被中间激活耗尽。同时，读写这批矩阵需要往返吞吐数百 GB 的 HBM 流量。 - **FlashAttention 分块融合推演**： - 设置 SRAM 分块尺寸 $B_r = B_c = 64$。 - 单个分块在片上 SRAM 中仅占用： $$\text{Tile Memory} = 64 \times 64 \times 2\text{ 字节} = 8192\text{ 字节} = 8\text{ KB}$$ - 8 KB 的小张量可以完全装进每个 SM 的 Shared Memory（容量高达 228 KB）中。 - 全过程没有 $N \times N$ 中间矩阵写回 HBM，原本 $64\text{ GiB}$ 的二次方激活不再需要保存（只保留线性规模的输出 $O$ 与每行统计量，反向时在片上重算）。这为训练省出几十 GiB 显存，也使注意力层更接近算力受限（Compute Bound）。 ## 代码填空题用 numpy 实现分块 online softmax 注意力，验证与标准实现数值一致。 ```{python} #| course_role: exercise import numpy as np def standard_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray): """参照实现：物化 S 和 P。Q,K,V: (n, d)。""" n = Q.shape[0] S = Q @ K.T # (n, n) —— 长序列下这是内存炸弹 S = S - S.max(axis=-1, keepdims=True) P = np.exp(S) P /= P.sum(axis=-1, keepdims=True) return P @ V, S def flash_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray, block_size: int = 64) -> np.ndarray: """ 分块实现：不物化完整 S。对每个 query 块，逐步处理 kv 块，维护每个 query 行的： m: 见过的分数的 running max l: 见过的 exp(score - m) 之和（softmax 分母） O: 已累加的加权输出 """ n, d = Q.shape O = np.zeros((n, d)) l = np.zeros(n) # 分母（running sum） m = np.full(n, -np.inf) # running max for i in range(0, n, block_size): Qi = Q[i:i + block_size] Oi = O[i:i + block_size] li = l[i:i + block_size] mi = m[i:i + block_size] for j in range(0, n, block_size): Kj = K[j:j + block_size] Vj = V[j:j + block_size] S_ij = Qi @ Kj.T # 当前块分数 (bs, bs) # online softmax 更新： # 1) 新 running max（把当前块的最大分数并进来） m_new = ______ # 2) 旧分母/旧输出都要按 exp(mi - m_new) 缩放，新块贡献 p_ij scale = ______ p_ij = np.exp(S_ij - m_new[:, None]) # 新块的"原始"权重 l_new = ______ # 分母更新：旧项缩放 + 新块和 # 3) 输出更新：旧输出按比例缩放，加上新块贡献 Oi = ______ mi, li = m_new, l_new # 把块的统计量写回全局数组 O[i:i + block_size] = Oi l[i:i + block_size] = li m[i:i + block_size] = mi # 最终归一化：O / l（把 running max 移回） return O / l[:, None] if __name__ == "__main__": rng = np.random.default_rng(0) n, d = 256, 32 Q = rng.standard_normal((n, d)) * 0.5 K = rng.standard_normal((n, d)) * 0.5 V = rng.standard_normal((n, d)) O_ref, S_ref = standard_attention(Q, K, V) O_flash = flash_attention(Q, K, V, block_size=32) print("最大绝对误差:", np.abs(O_ref - O_flash).max()) print("S 矩阵元素数:", S_ref.size, "（FlashAttention 不物化它）") assert np.allclose(O_ref, O_flash, atol=1e-6) print("与标准注意力一致 ✓（内存从 O(n²) 降到 O(n)）") ``` ## 三个问答题 ### Q1 为什么"先算完整 S = QKᵀ 再 softmax"在长序列下既慢又占内存？请从 HBM 带宽与显存两个角度回答，并把它与第 2 课激活公式中的二次项联系起来。FlashAttention 最终在内存与访存上各省掉了什么？ ### Q2 教材的 naive matmul 中，为什么"每个线程算一个输出元素"会导致访存不合并？改成一维 block 映射后为什么吞吐提升约 10 倍？Tiling 为什么能进一步减少 HBM 访问？Thread coarsening 解决的是哪一类停顿（提示：MIO throttle）？ ### Q3 请对比 torch.compile、Triton、CUDA 三条写 kernel 的路径（难度/性能/灵活性），并说明什么时候必须离开 Triton 写 CUDA。FA-1 → FA-2 → FA-3 各自改进的重点是什么（各一句话）？ ## 检查与通过标准总分 10 分：代码正确 4 分（m_new/l_new/O 更新、归一化、与标准实现一致）、三题各 2 分、通过线 8 分。一票否决项： - 认为 FlashAttention 降低了计算量（它降低的是访存与显存）。 - 说不出 online softmax 为什么必须维护 running max（数值稳定性）。 - 认为"分别 softmax 再拼接"正确。 - 分不清 HBM 与共享内存在带宽上的量级差异。 ::: {.callout-tip collapse="true" title="💡 点击展开参考答案与详细代码实现" icon="false"} ## 参考答案（仅 answer 分支）先完成题目再核对；核对后必须解释关键公式，并改一个规模重新计算。 ```{python} #| course_role: solution import numpy as np def standard_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray): """参照实现：物化 S 和 P。Q,K,V: (n, d)。""" n = Q.shape[0] S = Q @ K.T # (n, n) —— 长序列下这是内存炸弹 S = S - S.max(axis=-1, keepdims=True) P = np.exp(S) P /= P.sum(axis=-1, keepdims=True) return P @ V, S def flash_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray, block_size: int = 64) -> np.ndarray: """ 分块实现：不物化完整 S。对每个 query 块，逐步处理 kv 块，维护每个 query 行的： m: 见过的分数的 running max l: 见过的 exp(score - m) 之和（softmax 分母） O: 已累加的加权输出 """ n, d = Q.shape O = np.zeros((n, d)) l = np.zeros(n) # 分母（running sum） m = np.full(n, -np.inf) # running max for i in range(0, n, block_size): Qi = Q[i:i + block_size] Oi = O[i:i + block_size] li = l[i:i + block_size] mi = m[i:i + block_size] for j in range(0, n, block_size): Kj = K[j:j + block_size] Vj = V[j:j + block_size] S_ij = Qi @ Kj.T # 当前块分数 (bs, bs) # online softmax 更新： # 1) 新 running max（把当前块的最大分数并进来） m_new = np.maximum(mi, S_ij.max(axis=-1)) # 2) 旧分母/旧输出都要按 exp(mi - m_new) 缩放，新块贡献 p_ij scale = np.exp(mi - m_new) p_ij = np.exp(S_ij - m_new[:, None]) # 新块的"原始"权重 l_new = li * scale + p_ij.sum(axis=-1) # 分母更新：旧项缩放 + 新块和 # 3) 输出更新：旧输出按比例缩放，加上新块贡献 Oi = Oi * scale[:, None] + p_ij @ Vj mi, li = m_new, l_new # 把块的统计量写回全局数组 O[i:i + block_size] = Oi l[i:i + block_size] = li m[i:i + block_size] = mi # 最终归一化：O / l（把 running max 移回） return O / l[:, None] if __name__ == "__main__": rng = np.random.default_rng(0) n, d = 256, 32 Q = rng.standard_normal((n, d)) * 0.5 K = rng.standard_normal((n, d)) * 0.5 V = rng.standard_normal((n, d)) O_ref, S_ref = standard_attention(Q, K, V) O_flash = flash_attention(Q, K, V, block_size=32) print("最大绝对误差:", np.abs(O_ref - O_flash).max()) print("S 矩阵元素数:", S_ref.size, "（FlashAttention 不物化它）") assert np.allclose(O_ref, O_flash, atol=1e-6) print("与标准注意力一致 ✓（内存从 O(n²) 降到 O(n)）") ``` ## 补全后的代码（关键填空处） ```python # online softmax 更新 m_new = np.maximum(mi, S_ij.max(axis=-1)) # 1) 新 running max scale = np.exp(mi - m_new) # 旧输出的缩放因子 p_ij = np.exp(S_ij - m_new[:, None]) # 新块的"原始"权重 l_new = li * scale + p_ij.sum(axis=-1) # 2) 分母更新 Oi = Oi * scale[:, None] + p_ij @ Vj # 3) 输出更新 ``` 验证结果： ```python 最大绝对误差: ~1e-15 S 矩阵元素数: 65536（标准实现物化；FlashAttention 不物化）与标准注意力一致 ✓（内存从 O(n²) 降到 O(n)） ``` ## 三个问答题要点 ### Q1 - 慢：S 与 P 都是 n×n，计算它们要写回 HBM 再读回（每步都要走 HBM 带宽）；n 大时访存时间远超计算（memory-bound）； - 占内存：S 是 O(n²)，正是第 2 课激活公式中 `5·n_heads·seq²·bs` 项——长序列时是激活最大头，全量重计算也救不了（层边界激活仍随 seq 线性增长，但注意力分数矩阵是二次项）； - FA 省掉：不物化 S/P（O(n²) → O(n) 内存）；HBM 访存量从 Θ(nd + n²) 降到 Θ(n²d²/M)（M 为 SRAM 容量），额外显存从 O(n²) 降到 O(n)——计算量不变，但 memory-bound 瓶颈被大幅缓解。 ### Q2 - naive matmul：线程 (x,y) 都算一个输出元素时，同 warp 的线程 x 相邻但读的是 A 的不同行（行距 = K）→ 不连续 → 无法合并成 DRAM burst；B 的列访问同理； - 一维 block 映射（x = blockIdx·BLOCK + threadIdx/BLOCKSIZE，y = threadIdx%BLOCKSIZE）：warp 内相邻线程的 y 相邻 → 读 B 的连续列（行主序下连续）→ 合并；吞吐提升 ~10 倍； - Tiling：把 A/B 的 tile 先搬进共享内存，块内所有线程复用 → HBM 访问从"每个线程各读一遍"变成"每 tile 读一次"； - Thread coarsening：每个线程算多个输出 → 共享内存指令数减少 → 缓解 MIO 指令队列阻塞（ncu 的 mio_throttle 停顿）。 ### Q3 - 谱系：torch.compile（易用/快/不灵活，可导出 Triton 代码作起点）→ Triton（block 级、较快、较灵活）→ CUDA（最难、最快、完全可控）； - 必须写 CUDA 的场景：需要精细控制共享内存布局、SM 内调度/warp 划分、特殊指令（如 FA-2/FA-3 的 warp 分工与 FP8 Tensor Core 路径），Triton 只暴露 block 级抽象； - FA-1→2→3：FA-1 确立分块 + online softmax + 融合（内存 O(n)）；FA-2 减少非 matmul 操作、按 warp 精细划分 Q/K/V 工作量（提高利用率）；FA-3 面向 Hopper 优化 FP8/Tensor Core 与异步流水线。 ## 通过要点 - FlashAttention 不降低 FLOPs，降低的是 HBM 访存（降为 Θ(n²d²/M)）与额外显存（O(n²)→O(n)）； - online softmax 的 running max 是为了数值稳定（防止 exp 溢出），running sum 是为了正确归一化； - 与第 8 课 Ring Attention 共享同一套 online softmax 思想。 ## 官方主参考 - [Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook) - [PyTorch distributed documentation](https://pytorch.org/docs/stable/distributed.html) ::: ::: {.course-pagination} [**← 上一课**<br>第 12 课：配置搜索与性能诊断](lesson12.qmd){.pagination-link .pagination-prev} [**下一课 →**<br>第 14 课：混合精度与综合面试](lesson14.qmd){.pagination-link .pagination-next} :::