GPU 内核与 FlashAttention

在 H100 上，一个形状合适的大矩阵乘能跑到接近峰值，一个 LayerNorm 却只能用到算力的零头；序列一长，标准注意力的大部分时间也花在搬数据上，而不是计算上。多卡并行的账默认单张卡上的每个算子都跑得足够快，这个前提成不成立，要看单张 GPU 内部：硬件长什么样，一个内核（kernel）怎么被调度执行，数据从哪里来、要花多久、怎样少搬几次。FlashAttention 就是把这些知识用到极致的结果。

硬件数字都以 H100 SXM 为准。

1. 为什么要关心内核

先看一组 H100 SXM 的数字：BF16 矩阵乘峰值约 \(9.9\times10^{14}\) FLOPs/s，HBM 带宽约 \(3.35\times10^{12}\) 字节/s。两者相除：

这个比值叫临界算术强度。一个算子每从 HBM 读写 1 字节，如果做不到大约 295 次浮点运算，那它的耗时就由搬数据决定，计算单元会闲着等。这就是 Roofline 模型的拐点。

大多数非矩阵乘的算子离这个数差得很远。两个 FP32 向量相加：每个元素读 8 字节、写 4 字节、做 1 次加法，强度是 \(1/12\)。LayerNorm、softmax、GeLU、dropout、残差加法都差不多是这个量级。矩阵乘 \(C=AB\)（\(A\in\mathbb R^{M\times K}, B\in\mathbb R^{K\times N}\)）做 \(2MNK\) 次运算、最少搬 \(2(MK+KN+MN)\) 字节（BF16），当三个维度都是几千时强度能到上千——但这只是「最少」，内核写得不好，实际搬运量可以比这多几十倍。

举个具体的例子。一个线性层 \(Y=XW\)，\(X\in\mathbb R^{B\times 4096}\)，\(W\in\mathbb R^{4096\times 8192}\)，BF16，在 H100 上：

• \(B=64\)：FLOPs \(=2\cdot64\cdot4096\cdot8192\approx4.3\times10^9\)，按峰值算 4.3 µs；最少搬运 \(2(64\cdot4096+4096\cdot8192+64\cdot8192)\approx6.9\times10^7\) 字节，按带宽算 20.5 µs。搬权重的时间是计算的 5 倍，受带宽限制。
• \(B=512\)：FLOPs \(\approx3.4\times10^{10}\)，34.7 µs；搬运 \(\approx8.0\times10^7\) 字节，23.8 µs。受算力限制。

当 \(B\) 远小于两个权重维度时，强度约等于 \(B\)（每个权重元素读 2 字节、被用 \(2B\) 次），所以 \(B\) 要超过约 295 才能跑满 Tensor Core。换成 B200 算一遍（2.25e15 / 8.0e12），临界值约 281——新一代 GPU 的算力和带宽同步增长，这个门槛并没有降低。

所以写内核的核心问题可以概括成两句话：

• 对于本来就受带宽限制的算子：少碰 HBM。把多个算子合成一个（融合），每个字节只读一次、写一次。
• 对于本来可以受算力限制的算子（矩阵乘、注意力）：让数据在片上被反复用（分块），并且让搬运和计算同时进行（异步、流水线）。

要做到这两点，得先知道硬件里有哪些层次、每层多大多快。

2. GPU 的硬件组织：SM、核心与内存层级

2.1 计算：SM、子分区、CUDA core 与 Tensor Core

一块现代训练 GPU 基本上是「一大堆矩阵乘单元 + 一条很快的内存」。计算部分由许多流式多处理器（Streaming Multiprocessor，SM）组成：H100 有 132 个，B200 有 148 个。每个 SM 基本独立，自己取指令、调度、执行，GPU 因此可以同时跑上百个互不相干的任务。

每个 SM 又分成 4 个相同的子分区（SM sub-partition）。每个子分区里有：

• CUDA core：做标量/向量运算的 ALU。每个子分区有 32 个 FP32 ALU，它们每个周期执行同一条指令。所以 H100 一共有 \(132\times4\times32=16{,}896\) 个 FP32 CUDA core。它们负责逐元素运算、ReLU、求和等。
• Tensor Core：专门做小矩阵乘加的单元，GPU 的大部分 FLOPs 都来自它。H100 上 Tensor Core 的 BF16 峰值约 990 TFLOPs/s，而 CUDA core 只有约 66 TFLOPs/s（按 FMA 算两次运算）。反推一下，每个 H100 Tensor Core 每周期大约做 1024 次 BF16 运算（\(990\text{e}12 / 1.76\text{e}9 / 132 / 4\approx 1066\)），相当于一个 \(8\times8\times8\) 的矩阵乘。
• warp 调度器和寄存器堆：每个子分区 16,384 个 32 位寄存器，整个 SM 64K 个，即 256 KB。

此外每个 SM 有一块 256 KB 的片上 SRAM，可以由程序员当作共享内存（shared memory，SMEM）显式管理，也可以由硬件当作 L1 缓存用。H100 上一个 SM 最多可以把其中 228 KB 配成共享内存，单个 block 最多用 227 KB（CUDA 编程指南的 compute capability 9.0 表）。

两个数字值得记住：Tensor Core 和 CUDA core 的吞吐差了 15 倍左右；FP8 矩阵乘的吞吐又是 BF16 的 2 倍。这意味着在矩阵乘密集的程序里，任何不走 Tensor Core 的运算都很贵。FlashAttention-2 和 -3 很大一部分工作就是在压缩这类运算。

2.2 内存层级：容量、带宽、延迟

从慢到快：

延迟是用 pointer-chasing 微基准在 Hopper（H800）上测得的时钟周期数。NVIDIA 不公布 L2 带宽；在 H100 PCIe 上实测，数据落在本侧 L2 分区时读带宽约 5.5 TB/s，访问跨越整个 50 MB（要走两个分区之间较慢的连接）时降到约 3.8 TB/s。SXM 版 SM 更多、频率更高，数字会再高一些，但量级就是 HBM 的 2 到 3 倍——L2 比 HBM 快，却远没有快一个数量级。SMEM 带宽按每 SM 每周期 128 字节 × 132 个 SM × 1.83 GHz 算出。

几个要点：

• 从 HBM 到 SMEM，带宽差将近一个数量级，延迟也差一个多数量级。HBM 一次访问要等将近 500 个周期，这段时间里 Tensor Core 可以做几十万次运算。
• L2 对程序员不可见、不可显式控制，但它决定了很多内核的实际 HBM 流量：多个 block 读同一块数据时，如果在 L2 里命中，就不用再走 HBM。矩阵乘和 FlashAttention 的实际流量都靠它。
• 寄存器是最快的存储，但每个 SM 只有 64K 个，所有常驻线程分着用。一个线程用的寄存器越多，SM 上能同时驻留的线程越少——这就是「占用率」问题。

各代 GPU 的主要数字（矩阵乘吞吐均为稠密峰值）：

B200 的 Tensor Core 大到输入都放不进 SMEM/寄存器，所以又加了一层 256 KB/SM 的 TMEM（Tensor Memory）专门放累加器。这里能看到一个趋势：Ampere 上一个 warp 就能驱动 Tensor Core，Hopper 需要一个 warpgroup（4 个 warp），Blackwell 则由 2 个 SM 一起驱动。Tensor Core 越来越大，喂饱它越来越难，这也是 FlashAttention-3 必须重写的原因。

3. 编程模型：线程、warp、block、grid

3.1 内核与主机代码

在 GPU 上运行的一段函数叫内核。它可以用 CUDA C++ 或 Triton 写，最终编译成 NVIDIA 的虚拟汇编 PTX，再由驱动翻译成机器码 SASS。CPU 端（host）负责分配显存、拷数据、启动内核。启动是异步的：CPU 把内核放进队列就返回，继续准备下一个。

最简单的例子是向量加法。内核本身只描述「一个线程做什么」：

__global__ void vec_add(const float* A, const float* B, float* C, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;   // 这个线程负责的全局下标
    if (i < n) C[i] = A[i] + B[i];                  // 越界保护：最后一个 block 可能不满
}
// 主机端：每个 block 256 个线程，block 数向上取整
int threads = 256;
int blocks = (n + threads - 1) / threads;
vec_add<<<blocks, threads>>>(dA, dB, dC, n);

所有线程执行同一份代码，靠 blockIdx、threadIdx 区分自己负责哪部分数据。

3.2 线程层级

• 线程（thread）：最小的执行单位，有自己的寄存器和程序计数器。
• warp：32 个连续线程。硬件以 warp 为单位发射指令：一个周期里，一个 warp 的 32 个线程执行同一条指令，作用在各自的数据上。
• 线程块（block，也叫 CTA）：若干个 warp，最多 1024 个线程。一个 block 整体被分配到某一个 SM 上，从开始到结束都不换 SM。同一 block 的线程可以通过共享内存交换数据、用 __syncthreads() 同步。
• 网格（grid）：一次内核启动的所有 block。block 之间不能直接同步，执行顺序也不保证。

Hopper 又加了两层：warpgroup（4 个连续 warp，128 个线程，Hopper 的 wgmma 矩阵乘指令以它为单位发射）和线程块簇（cluster，若干个 block 被保证调度到同一个 GPC 上，可以互相访问对方的共享内存）。

3.3 SIMT 与 SIMD

如果你熟悉 CPU 或 TPU 的向量指令，要注意 GPU 的模型不太一样。TPU 的向量单元是 SIMD（单指令多数据），一条指令处理一个连续的向量寄存器；GPU 的 CUDA core 是 SIMT（单指令多线程）。SIMT 下每个线程有自己的指令指针和寄存器，写代码时可以当成独立的标量程序；但硬件上，同一 warp 的线程每周期还是执行同一条指令。如果 warp 内的线程走了不同的分支，硬件就把两个分支都执行一遍，各自屏蔽掉不该执行的线程——这叫控制分歧（第 8 节）。SIMT 的好处是灵活：线程可以访问任意地址、维护各自的状态；代价是写得不小心时性能会悄悄掉下去，而且掉在哪里不容易看出来。

4. warp 调度、延迟隐藏与占用率

4.1 用更多 warp 掩盖延迟

一个 SM 上可以同时驻留（resident）很多 warp：H100 上最多 64 个（2048 个线程），最多 32 个 block。每个 warp 调度器每个周期只发射一个 warp 的一条指令，但当某个 warp 在等内存时，调度器会在下一个周期换一个就绪的 warp 来执行，切换几乎没有开销（因为每个 warp 的寄存器一直留在寄存器堆里，不需要保存/恢复上下文）。这和 CPU 的超线程类似，只是规模大得多。

要掩盖多少延迟？可以用 Little 定律粗估。为了让 HBM 一直满速运转，任一时刻「在路上」的字节数应当等于带宽乘以延迟：

平摊到 132 个 SM，每个 SM 大约要有 7 KB 的读请求在途。一个 warp 用 16 字节向量化读取一次是 512 字节，所以每个 SM 上大约要有十几个这样的读同时未完成。（时钟按 1.75 GHz、延迟按上表的约 479 周期估算。）获得这么多在途请求的两条路：多驻留一些 warp（提高占用率），或者让每个线程自己发出多个独立的读（指令级并行，比如一次读下一块数据的同时计算当前块）。Hopper 的 TMA 走的是第三条路：用专门的硬件单元做异步拷贝，一个线程就能发起一整块的搬运。

4.2 占用率

占用率（occupancy）定义为一个 SM 上活跃 warp 数与硬件允许的最大 warp 数之比。它受四种资源同时限制：每 SM 的线程数（2048）、block 数（32）、寄存器（64K 个）和共享内存（228 KB）。一个 block 要么整个放得下，要么一个都放不下，所以：

其中 \(T\) 是每个 block 的线程数，\(R\) 是每个线程的寄存器数，\(S\) 是每个 block 的共享内存。举个极端的例子：每个线程用满 255 个寄存器时，一个 SM 只能放 \(65536/(255\times32)\approx 8\) 个 warp，占用率只有 12.5%。

占用率不是越高越好。高性能的矩阵乘和注意力内核往往故意用很多寄存器和共享内存来做大分块，占用率很低，但每个 warp 内部有足够的指令级并行和异步搬运来掩盖延迟。NVIDIA 的 Best Practices Guide 也明确说，占用率高到一定程度后再提高不会带来收益。它的作用是诊断：如果内核是带宽受限的、又没有其他手段掩盖延迟，占用率太低就会是瓶颈。

5. 写内核的四个层次：PyTorch、torch.compile、Triton、CUDA

知道了硬件，下一个问题是用什么写内核。可选的工具大致排成一条线，从易到难、从受限到灵活：

• PyTorch eager：最容易，但每个算子是一个独立内核，中间结果都要写回 HBM。
• torch.compile：同样容易，加一个装饰器；它会捕获计算图，自动融合、生成 Triton 内核。但能做的优化受编译器限制，不够灵活。
• Triton：更难一些，更快、更灵活。编程粒度是「block」：你描述一个 program 实例处理哪一块数据，线程级的事情由编译器决定。
• CUDA：最难，也最快、最灵活——前提是写对了。能控制到每个线程、共享内存布局、异步拷贝和 warp 级指令。

5.1 torch.compile

以 ELU 激活为例：\(\mathrm{ELU}(x)=e^x-1\)（\(x\lt0\)）或 \(x\)（\(x\ge0\)）。直接写 PyTorch 再加装饰器：

@torch.compile
def elu(x, alpha=1.0):
    return torch.where(x < 0, alpha * (torch.exp(x) - 1), x)

eager 模式下这一行会启动好几个内核（比较、exp、减、乘、where），每个都读写一遍整个张量；编译后变成一个内核，读一遍 x、写一遍结果。设置环境变量 TORCH_LOGS="output_code" 可以看到 TorchInductor 生成的 Triton 代码——这是学写 Triton 最快的起点：先让编译器写一个，再在它的基础上改。

5.2 Triton：逐元素内核

把编译器生成的 ELU 内核整理一下，大致是这样：

import triton
import triton.language as tl

@triton.jit
def elu_kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)                    # 第几个 program（相当于 CUDA 的 block）
    offs = pid * BLOCK + tl.arange(0, BLOCK)  # 这个 program 负责的 BLOCK 个下标
    mask = offs < n                           # 尾部越界保护
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.where(x < 0, tl.exp(x) - 1.0, x)
    tl.store(y_ptr + offs, y, mask=mask)

# 启动：grid 是 program 的个数
grid = (triton.cdiv(n, 1024),)
elu_kernel[grid](x, y, n, BLOCK=1024)

和 CUDA 的区别在于：Triton 里的变量是整块的张量（offs 是长度为 BLOCK 的向量），你不写「线程 i 做什么」，而写「这一块做什么」。tl.arange 产生的连续下标会被编译器映射成合并的访存（第 6 节），向量化宽度、每线程处理几个元素也由编译器决定。

5.3 Triton：一个融合的 RMSNorm

行归约类的算子也很好写：一个 program 处理一行。

@triton.jit
def rmsnorm_kernel(x_ptr, w_ptr, y_ptr, row_stride, D, eps, BLOCK_D: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_D)              # BLOCK_D 取 >= D 的 2 的幂
    mask = cols < D
    x = tl.load(x_ptr + row * row_stride + cols, mask=mask, other=0.0).to(tl.float32)
    ms = tl.sum(x * x, axis=0) / D            # 在寄存器里归约，FP32 累加
    rstd = 1.0 / tl.sqrt(ms + eps)
    w = tl.load(w_ptr + cols, mask=mask, other=0.0)
    y = x * rstd * w
    tl.store(y_ptr + row * row_stride + cols, y.to(y_ptr.dtype.element_ty), mask=mask)

rmsnorm_kernel[(n_rows,)](x, w, y, x.stride(0), D, 1e-6, BLOCK_D=triton.next_power_of_2(D))

整行读进寄存器一次，平方、求和、开方、乘权重都在片上完成，最后写一次。

5.4 Triton：分块矩阵乘

矩阵乘是 Triton 的强项。下面是一个最小的分块实现（结构上和 Triton 官方教程一致，去掉了自动调优和 L2 分组）：

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                  s_am, s_ak, s_bk, s_bn, s_cm, s_cn,
                  BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    rm = pid_m * BM + tl.arange(0, BM)                  # 这块 C 的行
    rn = pid_n * BN + tl.arange(0, BN)                  # 这块 C 的列
    rk = tl.arange(0, BK)
    a_ptrs = a_ptr + rm[:, None] * s_am + rk[None, :] * s_ak   # BM x BK 的指针块
    b_ptrs = b_ptr + rk[:, None] * s_bk + rn[None, :] * s_bn   # BK x BN 的指针块
    acc = tl.zeros((BM, BN), dtype=tl.float32)          # 累加器留在寄存器，FP32
    for k in range(0, tl.cdiv(K, BK)):
        k_left = K - k * BK
        a = tl.load(a_ptrs, mask=(rm[:, None] < M) & (rk[None, :] < k_left), other=0.0)
        b = tl.load(b_ptrs, mask=(rk[:, None] < k_left) & (rn[None, :] < N), other=0.0)
        acc += tl.dot(a, b)                             # 编译成 Tensor Core 指令
        a_ptrs += BK * s_ak
        b_ptrs += BK * s_bk
    c_ptrs = c_ptr + rm[:, None] * s_cm + rn[None, :] * s_cn
    tl.store(c_ptrs, acc.to(tl.bfloat16), mask=(rm[:, None] < M) & (rn[None, :] < N))

grid = (triton.cdiv(M, 128), triton.cdiv(N, 128))
matmul_kernel[grid](a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1),
                    c.stride(0), c.stride(1), BM=128, BN=128, BK=32)

每个 program 负责 C 的一个 \(128\times128\) 块，沿 K 方向每次取 A 的 \(128\times32\) 块和 B 的 \(32\times128\) 块，用 tl.dot 乘加到 FP32 累加器上。这就是「分块」（第 7 节）。Triton 编译器负责把块搬进共享内存、安排流水线、选 Tensor Core 指令。

Triton 也有边界：它的抽象停在 block 这一层，对共享内存的精确布局、SM 内的调度（比如 warp 专门化）控制有限。要榨干最后一截性能，就得下到 CUDA（或者 CUTLASS/CuTe 这类模板库）。访存合并、分块、线程粗化这些技术在 CUDA 层面看得最清楚，但它们背后的道理同样决定 Triton 内核的快慢。

5.5 怎么知道内核慢在哪

写之前先测。常用的两层工具：

• PyTorch Profiler / Nsight Systems：看时间线。哪些内核占了大部分时间，内核之间有没有空隙（CPU 启动跟不上、同步、数据拷贝），通信和计算有没有重叠。
• Nsight Compute（ncu）：深入单个内核。先看它在 roofline 上的位置——实际达到的内存吞吐和计算吞吐分别是峰值的百分之几；再看 warp 停顿原因（warp stall reasons）：大量停在等全局内存的读取返回，说明延迟没掩盖住或访问没合并；停在 mio_throttle，说明共享内存等指令队列太忙（第 8 节的例子）；ncu 也会直接报告未合并的全局访问和共享内存 bank 冲突。

一个内核如果已经接近它该有的 roofline（带宽受限的接近 HBM 带宽，计算受限的接近 Tensor Core 峰值），再优化它本身就没有意义了，该做的是融合或者改算法。

6. 访存合并

6.1 DRAM 按突发传输

全局内存是 DRAM。DRAM 每次被访问时，并不是只读出你要的那个字，而是把一段连续地址并行读出来、以突发（burst）的形式送出。GPU 的内存系统也以固定大小的段为单位服务请求：对 compute capability 6.0 及以上的设备，CUDA Best Practices Guide 的说法是，一个 warp 的并发访问会被合并成若干个 32 字节的事务，数量等于覆盖所有线程地址所需的 32 字节段数。

于是同一个 warp 的 32 个线程如果读的是 32 个相邻的 FP32（共 128 字节、对齐），只需要 4 个 32 字节事务，搬运的每个字节都有用；如果它们读的地址彼此相隔很远，就需要 32 个事务，搬了 1024 字节只用了 128 字节，有效带宽只剩 1/8。这就是访存合并（memory coalescing）：让同一个 warp 里相邻的线程访问相邻的地址。

上面的计算就是 Best Practices Guide 给的规则：线程 \(t\) 访问地址 \(a_t = \text{off} + t\cdot s\cdot b\)（\(s\) 是步长，\(b\) 是元素字节数），事务数等于集合 \(\{\lfloor a_t/32\rfloor\}\) 的大小。几个值得注意的情况：错位 1 个元素时，128 字节跨了 5 个扇区，效率掉到 80%；步长 2 时一半的字节白搬；步长达到 8 个 FP32（32 字节）以后，每个线程独占一个扇区，效率稳定在 1/8；广播时 32 个线程读同一个地址，只要 1 个事务。

6.2 朴素矩阵乘为什么慢

合并规则最经典的反面教材是朴素矩阵乘。最直接的写法是每个线程算 C 的一个元素：

__global__ void matmul_naive(int M, int N, int K, const float* A, const float* B, float* C) {
    const uint x = blockIdx.x * blockDim.x + threadIdx.x;   // 行
    const uint y = blockIdx.y * blockDim.y + threadIdx.y;   // 列
    if (x < M && y < N) {
        float tmp = 0.0;
        for (int i = 0; i < K; ++i) tmp += A[x * K + i] * B[i * N + y];
        C[x * N + y] = tmp;
    }
}

问题在 x 用的是 threadIdx.x。一个 warp 由 threadIdx.x 连续的 32 个线程组成（2D block 中 x 变化最快），于是同一 warp 里的线程有相同的 y、连续的 x。在循环第 \(i\) 步：

• 读 A[x*K+i]：相邻线程的地址相差 \(K\) 个元素——这是上面 demo 里「按列」的情况，每个线程独占一个扇区。
• 读 B[i*N+y]：所有线程读同一个地址——广播，1 个事务。

修法只需改下标映射，用一维 block，把 threadIdx.x 的低位分给列：

const int x = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE);
const int y = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);

现在同一 warp 的线程有相同的 x、连续的 y：读 A 变成广播，读 B 和写 C 变成连续访问。只改了两行下标，ncu 里「未合并访问」的警告就消失了，内核执行时间缩短到原来的约 1/10。

6.3 共享内存的 bank 冲突

共享内存有类似但不同的规则。它被分成 32 个 bank，连续的 32 位字依次落在相邻的 bank 上，每个 bank 每周期服务一个 32 位字。一个 warp 的 32 个访问如果落在 32 个不同的 bank 上，一个周期完成；如果有 \(n\) 个地址落在同一 bank 的不同位置，就要串行成 \(n\) 次（同一地址则是广播，不冲突）。

典型的坑是按列读一个 \(32\times32\) 的 FP32 共享内存数组：每列的 32 个元素相隔 32 个字，全都落在同一个 bank 上，32 路冲突。经典修法是把数组声明成 \(32\times33\)：多出的一列 padding 让相邻行错开一个 bank，列访问的步长变成 33，模 32 后等于 1，冲突消失（Best Practices Guide 的矩阵转置例子正是这么做的）。CUTLASS 等库更常用的是对地址做异或 swizzle，不浪费空间。

7. 分块：把数据复用算清楚

7.1 没有复用时的流量

\(C=AB\)，\(A\in\mathbb R^{M\times K}\)，\(B\in\mathbb R^{K\times N}\)。每个输出元素 \(C_{ij}=\sum_k A_{ik}B_{kj}\) 需要读 A 的一行和 B 的一列，共 \(2K\) 个元素。如果每个线程都从全局内存各读各的（即使访问合并了），总读取量是

而总计算是 \(2MNK\) FLOPs。算术强度是每个元素 1 次运算，FP32 下是 \(1/4\) FLOPs/B，比临界值 295 低三个数量级。L1/L2 缓存会挡掉一部分重复读取，但不能指望。

可实际上 A 的每个元素被 \(N\) 个输出用到，B 的每个元素被 \(M\) 个输出用到。理想情况下每个元素只需读一次，总读取 \(MK+KN\)。分块就是在有限的片上存储里尽量逼近这个理想。

7.2 分块后的流量

把 C 切成 \(B_M\times B_N\) 的块，每个 block 负责一块。计算这块需要 A 的 \(B_M\) 行（\(B_M\times K\)）和 B 的 \(B_N\) 列（\(K\times B_N\)），沿 K 方向每次取 \(B_K\) 宽的一片放进共享内存：

每一步从全局内存读 \((B_M+B_N)B_K\) 个元素，做 \(2B_MB_NB_K\) 次运算。一共有 \(\frac{M}{B_M}\cdot\frac{N}{B_N}\cdot\frac{K}{B_K}\) 步，所以：

正方形块 \(B_M=B_N=T\) 时是 \(2MNK/T\)：全局内存读取量缩小到朴素做法的 \(1/T\)。相应的算术强度（对全局内存而言）是

\(b\) 是元素字节数。注意它和 \(B_K\) 无关；\(B_K\) 只决定每一步要多少共享内存：\((B_M+B_N)B_K\cdot b\)。几个具体数字（BF16，\(b=2\)）：

• \(T=32\)：\(I=32\) FLOPs/B。
• \(128\times128\)：\(I=64\) FLOPs/B；\(B_K=64\) 时一片 A + 一片 B 占 32 KB 共享内存。
• \(128\times256\)：\(I\approx85\) FLOPs/B。

这里有个容易忽略的地方：即便 \(128\times256\) 的块，85 也离 H100 的 295 差得远。实际的大矩阵乘能跑到峰值附近，是因为上式算的是「每个 block 从 L2 以下取的数据」，而同时运行的上百个 block 读的是重叠的 A 行条和 B 列条，这些读取大量在 L2 命中。真正落到 HBM 上的流量接近每个元素读一次的理想值。所以高性能 GEMM 还会控制 block 的启动顺序（Triton 教程里叫「grouped ordering」，CUTLASS 叫 swizzle），让同时在跑的 block 共享尽量多的行条和列条，提高 L2 命中率。分块解决 SMEM 层面的复用，启动顺序解决 L2 层面的复用。

7.3 CUDA 里的样子

正方形分块的核心循环长这样：

// 每个线程负责 C 块中的一个元素 (localRow, localCol)
for (int t = 0; t < K; t += TILE) {
    sA[localRow * TILE + localCol] = A[localRow * K + localCol];   // 每个线程搬一个元素
    sB[localRow * TILE + localCol] = B[localRow * N + localCol];
    __syncthreads();                       // 等整块搬完
    A += TILE;  B += TILE * N;             // 指针移到下一片
    for (int i = 0; i < TILE; ++i)
        sum += sA[localRow * TILE + i] * sB[i * TILE + localCol];
    __syncthreads();                       // 等所有人算完再覆盖共享内存
}
C[localRow * N + localCol] = sum;

搬运时 localCol 取自 threadIdx.x，同一 warp 读相邻地址，合并没问题。两次 __syncthreads() 一次保证「读之前已经写完」，一次保证「覆盖之前大家都用完」。在合并访问的基础上加上分块，执行时间又降了约 43%。

7.4 Tile 量化与 wave 量化

分块还带来两个离散化的损失，算 GEMM 效率时经常被忽略。

• Tile 量化：如果 \(M\) 不是 \(B_M\) 的整数倍，最后一行块只有部分有效，但它照样占用一个完整 block 的计算。比如 \(M=4100\)、\(B_M=128\)：需要 33 个行块，最后一块只有 4 行有效，白算了约 3% 的量。这也是为什么词表大小、隐藏维度通常取 64 或 128 的倍数。
• Wave 量化：block 是一批一批（wave）地铺到 SM 上的。\(4096\times4096\) 的输出、\(128\times128\) 的块，一共 1024 个 block；H100 每个 SM 若只能驻留 1 个这样的 block，132 个 SM 需要 \(\lceil1024/132\rceil=8\) 个 wave，最后一个 wave 只有 \(1024-7\times132=100\) 个 block，有 32 个 SM 闲着。总效率上限是 \(1024/(8\times132)\approx97\%\)。矩阵再小一些时损失会更明显：256 个 block 需要 2 个 wave，效率上限 \(256/264\approx97\%\)；而 140 个 block 也需要 2 个 wave，上限只有 53%。

cuBLAS 会针对形状选不同的块大小和 split-K（把 K 维也切给多个 block，最后再归约）来缓解这些问题；自己写内核或选模型维度时，最好心里有这笔账。

8. 线程粗化与控制分歧

8.1 线程粗化（寄存器分块）

分块之后再用 Nsight Compute 看 warp 状态，会发现大量周期停在 mio_throttle：warp 在等共享内存指令队列有空位。原因是上面的内核里每做一次 FMA 都要从共享内存读两个数，共享内存成了新的瓶颈。

办法是线程粗化（thread coarsening）：让一个线程算 C 的 \(T_M\times T_N\) 个元素而不是一个。在内层循环的每个 \(k\)，线程从共享内存读 A 的 \(T_M\) 个数和 B 的 \(T_N\) 个数放进寄存器，然后做它们的外积——\(T_MT_N\) 次 FMA：

\(T_M=T_N=1\) 时是 \(1/2\)，\(8\times8\) 时是 4，提高了 8 倍。这其实是同一个分块思想在下一层的重复：全局内存 → 共享内存是 block 级分块，共享内存 → 寄存器是线程级分块。Tensor Core 指令本身（比如一个 warp 做 \(16\times8\times16\) 的 mma）又是再下一层。代价是每个线程要更多寄存器存 \(T_MT_N\) 个累加器，占用率会下降——又回到第 4 节的权衡。

8.2 控制分歧

同一 warp 的线程如果在 if 上走了不同的分支，硬件就依次执行两个分支，每次屏蔽掉另一半线程。最坏情况下 32 条路径串行，吞吐掉到 1/32。常见的避免方法：

• 让分支条件在 warp 内一致。例如按 warp 而不是按线程分配任务；FlashAttention 里因果掩码只在对角线附近的块上做，其他块要么整块计算、要么整块跳过（第 14 节）。
• 用无分支的写法：where/select，或者用谓词执行（predication），两边都算、最后按条件选。对于很短的分支，编译器通常会自动这么做。
• 把数据重排，让走同一路径的元素聚在一起。

9. 内核融合

GPU 和 CPU 是异步的，CPU 可以不停地往队列里塞内核。但每个内核都遵循同一个模式：从 HBM 读输入到 SM，计算，把输出写回 HBM。连续几个逐元素算子各自成一个内核时，中间结果就要在 HBM 和 SM 之间来回走好几趟，而计算本身几乎不花时间。

算一笔账：RMSNorm，输入 \(x\in\mathbb R^{B\times D}\)。eager 下大致是 x.pow(2)（读 \(x\)、写临时量）、mean（读临时量）、x * rstd（读 \(x\)、写临时量）、* w（读临时量、写 \(y\)），大约 6 次全尺寸的 HBM 读写；如果中间还要转 FP32，还会更多。融合后只有读 \(x\)、写 \(y\) 两次。由于这类算子完全受带宽限制，运行时间基本按 3 倍缩短。

Transformer 里能融合的地方很多：bias + GeLU、残差加 + 归一化、RoPE、softmax + 掩码 + dropout、交叉熵（不把整个 logits 的 softmax 写回）。torch.compile 会自动做逐元素融合。

但融合在训练中有一个限制：反向传播需要的中间值，前向时仍然得写回 HBM 存起来，融合省不掉这部分写入。比如 softmax 的输出 \(P\) 在反向里要用，朴素地融合了「掩码 + softmax」也还是要把 \(N\times N\) 的 \(P\) 写回去。解决方法是重计算：不存中间值，反向时从更小的输入重新算出来。如果重算的代价（多出来的 FLOPs）比存取的代价（HBM 流量）小，那它反而更快。FlashAttention 就是把分块、融合、重计算三件事同时用在注意力上。

10. 标准注意力的访存账

单个注意力头，\(Q,K,V\in\mathbb R^{N\times d}\)（\(N\) 是序列长度，\(d\) 是头维度）：

（通常 \(S\) 还要乘 \(1/\sqrt d\)，这里省略，不影响讨论。）标准实现分三步，每步一个内核：

1. 读 \(Q,K\)（\(2Nd\)），算 \(S\)，写 \(S\)（\(N^2\)）。
2. 读 \(S\)（\(N^2\)），算 softmax，写 \(P\)（\(N^2\)）。
3. 读 \(P\)（\(N^2\)）和 \(V\)（\(Nd\)），算 \(O\)，写 \(O\)（\(Nd\)）。

总 HBM 访问 \(4N^2+4Nd\) 个元素，即 \(\Theta(N^2+Nd)\)。而 FLOPs 是 \(4N^2d\)（两个矩阵乘，各 \(2N^2d\)）。当 \(N\gg d\) 时，平方项主导，每个 \(S\) 元素带来 \(4d\) 次运算、却要搬 4 次（BF16 共 8 字节）：

\(d=128\) 时只有 64，远低于 295，所以标准注意力是带宽受限的。实际情况更糟：掩码、dropout 往往各是一个额外的 \(N^2\) 读写。

以 \(N=8192\)、\(d=128\)、32 个头、BF16 为例：每个头的 \(S\) 有 \(N^2=6.7\times10^7\) 个元素，134 MB；32 个头是 4.3 GB——这只是一层、一个样本，而且 \(P\) 还要存下来给反向用。在 A100 上测 GPT-2 medium（\(N=1024\)，\(d=64\)，16 头，batch 64）前向 + 反向：标准注意力 HBM 读写 40.3 GB、41.7 ms；FlashAttention 4.4 GB、7.3 ms——尽管后者 FLOPs 更多（75.2 vs 66.6 GFLOPs，多出来的是反向重计算）。

11. 在线 softmax 的推导

要不物化 \(S\)，就得把 softmax 拆成按块计算。难点在于 softmax 的分母要对整行求和，看起来必须先看到整行。

11.1 安全 softmax

先回忆数值稳定的写法。对向量 \(x\in\mathbb R^B\)：

减去最大值不改变结果（分子分母同乘 \(e^{-m}\)），但保证指数不会上溢：所有指数都 \(\le 0\)，\(e^{\cdot}\le1\)。朴素实现要扫三遍：一遍求最大值，一遍求和，一遍归一化。

11.2 两块的合并

把 \(x\) 拆成两段 \(x=[x^{(1)}, x^{(2)}]\)，分别算出 \((m_1,\ell_1)\) 和 \((m_2,\ell_2)\)。合并：

验证：\(\ell_1=\sum_{i\in 1}e^{x_i-m_1}\)，乘 \(e^{m_1-m}\) 得 \(\sum_{i\in1}e^{x_i-m}\)，两段加起来正是 \(\sum_i e^{x_i-m}=\ell(x)\)。也就是说，每段只需带着自己的 \((m,\ell)\) 两个数，就能在之后正确合并，而且合并满足结合律，可以一块接一块地做。这叫在线（online）softmax，Milakov 和 Gimelshein 2018 年把它用在了单遍 softmax 上。

11.3 把输出也一起在线化

注意力要的不是 softmax 本身，而是 \(o=\sum_j p_j v_j\)。设已经处理了前 \(j-1\) 块，维护三个量：当前最大值 \(m\)、当前分母 \(\ell\)、和一个未归一化的输出

来了新的一块分数 \(s^{(j)}\)（以及对应的 \(V_j\)），令 \(m' = \max(m, \max s^{(j)})\)，则：

第二个式子的道理和 \(\ell\) 一样：旧的每一项 \(e^{s_k-m}v_k\) 乘上 \(e^{m-m'}\) 就变成了以新最大值为基准的 \(e^{s_k-m'}v_k\)。处理完所有块后，\(o = \tilde o/\ell\)。用归纳法很容易验证：每一步之后，\(\tilde o\) 和 \(\ell\) 都恰好等于「对已处理的所有项、以当前 \(m\) 为基准」的和，所以最后一步得到精确结果，不是近似。

每一块要做的事：一个小矩阵乘得到 \(s^{(j)}\)；按行求最大值；指数；按行求和；把旧的 \(\tilde o\) 和 \(\ell\) 乘一个修正因子 \(e^{m-m'}\)；再一个小矩阵乘 \(\tilde P_jV_j\) 加到 \(\tilde o\) 上。上面的步进动画里切到 FlashAttention 模式，就能逐块看到 \(m\)、\(\ell\) 和输出的变化。

12. FlashAttention：分块 + 重计算

12.1 算法

有了在线 softmax，注意力可以写成一个内核。设片上 SRAM 能放 \(M\) 个元素。FlashAttention（第一版）的算法 1：

1. 取块大小 \(B_c=\lceil M/4d\rceil\)，\(B_r=\min(\lceil M/4d\rceil, d)\)。把 \(Q\) 切成 \(T_r=\lceil N/B_r\rceil\) 块，\(K,V\) 切成 \(T_c=\lceil N/B_c\rceil\) 块。在 HBM 中初始化 \(O=0\)、\(\ell=0\)、\(m=-\infty\)。
2. 外层循环 \(j=1..T_c\)：把 \(K_j,V_j\) 载入 SRAM。
   • 内层循环 \(i=1..T_r\)：载入 \(Q_i,O_i,\ell_i,m_i\)；片上算 \(S_{ij}=Q_iK_j^\top\)（\(B_r\times B_c\)）、行最大值、\(\tilde P_{ij}=\exp(S_{ij}-\tilde m_{ij})\)、行和；按在线 softmax 的公式更新 \(m_i,\ell_i\)，并更新 \(O_i\)；把 \(O_i,\ell_i,m_i\) 写回 HBM。

块大小的选法来自 SRAM 预算：\(K_j,V_j\) 各占 \(B_cd\)，\(Q_i,O_i\) 各占 \(B_rd\)，加起来不超过 \(4\cdot (M/4d)\cdot d = M\)。

代入 H100 的数字：228 KB 共享内存，BF16 下 \(M\approx1.17\times10^5\) 个元素；\(d=128\) 时 \(B_c=\lceil M/4d\rceil=228\)，\(B_r=\min(228,128)=128\)。实际实现里块大小还受寄存器限制，一般取 64 或 128。

12.2 IO 复杂度

先看结论：标准注意力需要 \(\Theta(Nd+N^2)\) 次 HBM 访问，FlashAttention 需要 \(\Theta(N^2d^2M^{-1})\) 次。推导：外层循环每一轮载入一块 K/V（大小 \(\Theta(M)\)），所以总共 \(T_c=\Theta(Nd/M)\) 轮；每一轮内层都要把整个 \(Q\) 和 \(O\) 过一遍，\(\Theta(Nd)\)。乘起来 \(\Theta(N^2d^2/M)\)。

典型参数下 \(d=64\sim128\)，\(M\) 约 100 KB 量级，\(d^2\) 比 \(M\) 小很多倍，所以 \(N^2d^2/M\ll N^2\)。不过常数和 \(d\) 的大小都要紧。代入具体数字：\(N=8192\)、\(d=128\)、\(B_c=228\) 时，FA1 的外层循环有 \(T_c=36\) 轮，每轮读 \(Q\)、读写 \(O\) 共约 \(3Nd\)，再加 \(K,V\) 各读一次，一共约 \(36\times3Nd+2Nd\approx1.2\times10^8\) 个元素；标准实现是 \(4N^2+4Nd\approx2.7\times10^8\)。如果不算 L2 的帮助，只差 2 倍多。\(d=64\) 时 \(B_c\) 翻倍、\(T_c\) 减半、每轮的量也减半，差距就到了 9 倍左右。所以在 \(d=128\) 的现代模型上，FlashAttention 的收益很大一部分来自：重复读取的 \(Q,O\)（FA1）或 \(K,V\)（FA2）在 L2 里命中；不再有掩码、dropout、softmax 各自的额外 \(N^2\) 读写；以及 \(O(N)\) 的显存。反过来还可以证明一个下界：不存在对所有 \(M\in[d,Nd]\) 都只需要 \(o(N^2d^2/M)\) 次 HBM 访问的精确注意力算法——在这个意义上 FlashAttention 是渐近最优的。

换个角度看片上的算术强度会更直观。以 FA2 的顺序为例：每载入一块 \(K_j,V_j\)（\(2B_cd\) 个元素），对一个 \(Q_i\) 块要做 \(4B_rB_cd\) 次运算，所以对 L2/HBM 的强度是

BF16、\(B_r=128\) 时是 128 FLOPs/B。它与 \(N\) 无关，只和 Q 块大小有关——这就是「分块让数据复用」在注意力上的体现。同一个头的 K/V 被所有 Q 块重复读，只要它们在 L2 里（\(N=8192,d=128\) 时一个头的 K+V 是 4 MB，远小于 50 MB 的 L2），真实 HBM 流量会接近「每个元素读写一次」。上面的计算器同时给出了不命中 L2 的上界和全部命中的下界。

12.3 显存

除了输入输出，FlashAttention 只需要 \(O(N)\) 的额外显存：每行的 \(m\) 和 \(\ell\)（FA2 合并成一个 logsumexp）。标准实现要存 \(O(N^2)\) 的 \(P\)。这使得长序列训练的激活显存从平方变成线性。这也是线性注意力等近似方法后来大多被搁置的原因之一：精确的注意力已经足够快、足够省。

13. 反向传播与重计算

13.1 推导

设损失为 \(\phi\)，已知 \(dO=\partial\phi/\partial O\)。逐个推（以下 \(q_i,k_j,v_j,o_i\) 等表示对应矩阵的第 \(i\)/\(j\) 行）：

• \(O=PV\)，所以 \(dV=P^\top dO\)，\(dP = dO\,V^\top\)，即 \(dP_{ij}=do_i^\top v_j\)。
• softmax 的雅可比：\(y=\mathrm{softmax}(x)\) 时 \(\partial y/\partial x=\mathrm{diag}(y)-yy^\top\)。所以第 \(i\) 行 $$dS_{i:} = \big(\mathrm{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 P_{ij}\,do_i^\top v_j = do_i^\top\Big(\sum_j P_{ij}v_j\Big) = do_i^\top o_i.$$ 所以 \(D=\mathrm{rowsum}(dO\circ O)\)，只要 \(O\) 和 \(dO\) 就能算，不需要整行的 \(P\)。于是 \(dS_{ij}=P_{ij}(dP_{ij}-D_i)\)。
• \(S=QK^\top\)，所以 \(dQ=dS\,K\)，\(dK=dS^\top Q\)。

13.2 分块反向

前向只存 \(O\) 和每行的 logsumexp \(L_i=m_i+\log\ell_i\)。反向时任何一块 \(P_{ij}\) 都可以从 \(Q_i,K_j\) 重新算出来：\(P_{ij}=\exp(Q_iK_j^\top - L_i)\)，不需要再做在线修正，因为最终的归一化常数已经知道了。FA2 的反向（算法 2）：

1. 预处理：\(D=\mathrm{rowsum}(dO\circ O)\)，写回 HBM。
2. 外层循环 K/V 块 \(j\)：载入 \(K_j,V_j\)，片上初始化 \(dK_j=dV_j=0\)。
   • 内层循环 Q 块 \(i\)：载入 \(Q_i,O_i,dO_i,L_i,D_i\)；重算 \(S_{ij}\) 和 \(P_{ij}\)；\(dV_j \mathrel{+}= P_{ij}^\top dO_i\)；\(dP_{ij}=dO_iV_j^\top\)；\(dS_{ij}=P_{ij}\circ(dP_{ij}-D_i)\)；\(dQ_i \mathrel{+}= dS_{ij}K_j\)（读写 HBM）；\(dK_j \mathrel{+}= dS_{ij}^\top Q_i\)。
3. 写回 \(dK_j,dV_j\)。

反向里有 5 个矩阵乘（重算 \(S\)、\(dV\)、\(dP\)、\(dQ\)、\(dK\)），前向 2 个，所以反向 FLOPs 约为前向的 2.5 倍（报告反向 TFLOPs/s 时通常就用这个系数）。多出来的那次 \(QK^\top\) 就是重计算的代价。由于标准实现的反向同样是带宽受限的（要读 \(P\)、写 \(dS\) 等若干个 \(N^2\) 矩阵），用多出的 FLOPs 换掉这些 HBM 流量是划算的：FlashAttention 的反向不但省显存，而且更快。这和一般的激活重计算（用时间换显存）不同。

14. FlashAttention-2：少做非矩阵运算，多做并行

FlashAttention 在 A100 上只能达到峰值的 25–40%，而优化好的 GEMM 能到 80–90%。FA2 的改进不在算法的数学上，而在怎么把工作分给 GPU。

14.1 减少非矩阵乘 FLOPs

A100 的 FP16/BF16 矩阵乘峰值是 312 TFLOPs/s，非矩阵乘的 FP32 只有 19.5 TFLOPs/s：一次非矩阵乘运算的代价相当于 16 次矩阵乘运算。所以哪怕只占总 FLOPs 很小比例的 exp、缩放、比较，也可能拖慢整个内核。FA2 做了两处调整：

• 循环中不再每一步都除以 \(\ell\)：维护未归一化的 \(\tilde O\)，只在最后除一次 \(\mathrm{diag}(\ell)^{-1}\)。FA1 每一步都对 \(O_i\) 做 \(\mathrm{diag}(\ell^{\text{new}})^{-1}(\mathrm{diag}(\ell)\,e^{\cdot}O_i+\dots)\) 这种双重缩放。
• 前向只为反向存 logsumexp \(L=m+\log\ell\) 一个数，而不是 \(m\) 和 \(\ell\) 两个。

14.2 交换循环顺序，在序列维度上并行

FA1 的并行只在 batch 和头数两个维度上：每个 (样本, 头) 一个 thread block。A100 有 108 个 SM，block 数 ≥ 80 左右时才能填满。长序列训练时 batch 往往很小，block 数不够，很多 SM 闲着。

FA2 把外层循环换成 Q 块、内层换成 K/V 块（这一顺序最早由 Phil Tillet 在 Triton 的实现中提出）。这样不同的 Q 块之间完全独立——每个 Q 块各自维护自己的 \(m,\ell,\tilde O\)，可以分给不同的 thread block，并行度变成 batch × 头数 × \(\lceil N/B_r\rceil\)。另外，\(O_i\) 累加器可以一直留在寄存器里，不必每一步读写 HBM。

反向时，按 K/V 列块并行：每个 block 负责一个 \(j\)，独立积累 \(dK_j,dV_j\)；唯一的共享写入是 \(dQ_i\)，不同 block 用原子加（atomic add）累加到 HBM 上。

MQA/GQA（多个 query 头共享一组 K/V 头）也在这一层处理：不在 HBM 里复制 K/V，而是在计算下标时让多个 query 头指向同一个 K/V 头；反向时，被共享的 \(dK,dV\) 要把各个 query 头的贡献加起来。

14.3 warp 之间的分工

一个 block 里通常有 4 或 8 个 warp。FA1 的做法是把 \(K,V\) 按列切给各个 warp，\(Q\) 大家共享（这种分法叫 split-K）：每个 warp 算出 \(QK^\top\) 的一部分列，乘以 \(V\) 的对应行后，得到的是输出的部分和，必须写进共享内存、同步、再加起来。FA2 改为把 \(Q\) 按行切给各个 warp，\(K,V\) 共享：每个 warp 算出自己那几行的完整输出，warp 之间不需要通信。这减少了共享内存读写和同步。

14.4 因果掩码

因果注意力里 \(j\gt i\) 的分数是 \(-\infty\)。按块处理时，列下标全部大于行下标的块（序列长时约一半）直接跳过；行下标严格大于列下标的块完全不需要掩码；只有对角线上的块要逐元素掩码（方块时每行只有 1 块）。FA2 报告相比无掩码约有 1.7–1.8 倍加速。这也是第 8 节「让分支在块级一致」的例子。

14.5 块大小与结果

块越大，共享内存读写越少，但寄存器和共享内存需求越大，过了某个点会寄存器溢出或者放不进共享内存。FA2 通常在 \(\{64,128\}\times\{64,128\}\) 中按头维度和设备手动选。结果：比 FA1 快约 2 倍，在 A100 上前向 + 反向最高 230 TFLOPs/s，约为峰值的 73%；端到端训练 GPT 风格模型时达到 225 TFLOPs/s（72% MFU）。

15. FlashAttention-3：Hopper 上的异步与 FP8

FA2 在 H100 上只能到峰值的约 35%。原因在于 FA2 的设计基于一个简化的同步模型，没有利用 Hopper 的两项新能力——异步和低精度。

15.1 Hopper 的新硬件

• TMA（Tensor Memory Accelerator）：专门在 HBM 和共享内存之间做异步拷贝的硬件单元。一个线程发出指令，就能搬运一整个多维块（地址计算、越界处理都由 TMA 完成），其他线程不用参与。
• WGMMA：warpgroup（128 个线程）级别的矩阵乘指令，异步执行，可以直接从共享内存取操作数。发出之后 warp 可以去做别的事，稍后再等结果。
• setmaxnreg：在 warpgroup 之间动态重分配寄存器。
• FP8 Tensor Core：吞吐是 FP16/BF16 的 2 倍。

15.2 warp 专门化：生产者与消费者

FA3 把一个 block 的 warp 分成两种角色：

• 生产者 warpgroup 只负责用 TMA 把 \(Q_i\) 和一块块的 \(K_j,V_j\) 搬进共享内存。它几乎不需要寄存器，所以用 setmaxnreg 把寄存器让出去。
• 消费者 warpgroup 只做计算：WGMMA 算 \(S\)，softmax，WGMMA 算 \(\tilde PV\)。它们拿到更多寄存器，可以用更大的块。

两者之间是一个 \(s\) 级的环形共享内存缓冲区：生产者等第 \(j\bmod s\) 级被消费完后填入 \(K_j,V_j\)，填好后通知消费者；消费者用完后释放。只要缓冲区不空，消费者就永远不必等 HBM——搬运和计算完全重叠。

15.3 乒乓调度：让 softmax 藏在 GEMM 后面

非矩阵乘的代价在 H100 上更尖锐。H100 SXM5 的 FP16 矩阵乘 989 TFLOPs/s，而 exp 这类特殊函数由多功能单元（MUFU）执行，每个 SM 每周期 16 次，\(16\times132\times1.83\text{ GHz}\approx3.9\) TFLOPs/s。头维度 128 时，每个分数元素对应 \(4d=512\) 次矩阵乘运算和 1 次 exp：矩阵乘运算多 512 倍，但吞吐高约 256 倍，于是 exp 要花掉相当于矩阵乘 50% 的时间。FP8 下矩阵乘吞吐翻倍而 exp 不变，情况更糟。

因为 exp 和 Tensor Core 是不同的硬件单元，理想情况是一边做 softmax 一边做矩阵乘。FA3 用两个消费者 warpgroup 做乒乓调度：用 bar.sync 强制 warpgroup 1 的两个 GEMM（本轮的 \(\tilde PV\) 和下一轮的 \(QK^\top\)）排在 warpgroup 2 的 GEMM 之前，于是 warpgroup 1 做 softmax 的时候 warpgroup 2 正在用 Tensor Core，然后角色互换。

实际的调度没有图上这么整齐，但在 FP16、头维度 128、序列长 8192 的前向上，把吞吐从约 570 提高到 620–640 TFLOPs/s。

15.4 warpgroup 内部的两级流水

单个 warpgroup 内部也可以重叠。一次迭代里 softmax 依赖第一个 GEMM 的结果，第二个 GEMM 又依赖 softmax，看起来是串行的。FA3 跨迭代打破这个依赖：在第 \(j\) 轮，先异步发出 \(S_{\text{next}}=Q_iK_j^\top\)（不等），再异步发出 \(O\mathrel{+}=\tilde P_{\text{cur}}V_{j-1}\)（不等），然后等第一个完成、对 \(S_{\text{next}}\) 做 softmax——这时第二个 GEMM 还在 Tensor Core 上跑。代价是要多存一份 \(S_{\text{next}}\)，每个 block 多占 \(B_r\times B_c\times 4\) 字节寄存器，和「用更大的块」争寄存器。还可以做成三级流水，寄存器压力更大。

15.5 FP8

FP8 带来两个问题。效率上，FP8 WGMMA 要求共享内存中的操作数在 K 维上连续（k-major），而 \(V\) 通常在头维度上连续，做 \(\tilde PV\) 时需要转置；FA3 在内核里用 ldmatrix/stmatrix 指令转置 \(V\) 的块，并安排在前面两个 WGMMA 的阴影里执行。另外 FP8 WGMMA 的 FP32 累加器在寄存器中的布局和下一次 WGMMA 期望的操作数布局不一致，要用字节置换指令重排。

精度上，E4M3 只有 3 位尾数，大模型的激活又有离群值（outlier）。FA3 用了两个手段：

• 块量化：\(Q,K,V\) 每个 \(B_r\times d\) 或 \(B_c\times d\) 的块各有一个缩放因子，而不是整个张量一个。因为算法本来就按块处理，在 \(S\) 的块上乘回缩放因子几乎没有额外开销。
• 非相干处理（incoherent processing）：量化前把 \(Q\) 和 \(K\) 都乘以同一个随机正交矩阵 \(M\)。因为 \(MM^\top=I\)，\((QM)(KM)^\top=QK^\top\)，结果不变；但 \(QM\) 的每个元素是 \(Q\) 一行的随机线性组合，离群值被「摊开」了，量化误差变小。实际用随机 \(\pm1\) 对角阵乘 Hadamard 矩阵，\(O(d\log d)\) 就能算，还可以和 RoPE 融合。

这两点和 FP8 训练里的细粒度缩放、离群值问题是同一回事。

15.6 结果

在 H100 SXM5 上，FA3 的 FP16 前向比 FA2 快 1.5–2.0 倍，最高 740 TFLOPs/s（约 75% 利用率），反向快 1.5–1.75 倍；FP8 接近 1.2 PFLOPs/s。在有离群值的情形下，带块量化和非相干处理的 FP8 FA3 比按张量量化的 FP8 基线数值误差低 2.6 倍。

FlashAttention 的局限是只支持固定的几种注意力模式；需要自定义掩码或分数修改时，可以看 PyTorch 的 FlexAttention，它用 torch.compile 为用户给的 score_mod/mask_mod 生成融合内核。

16. 小结

实际工作中的顺序大概是：先用现成的（FlashAttention、cuBLAS、融合好的 Norm/RoPE/交叉熵内核），再用 torch.compile 收拾零散的逐元素算子；profiler 里还看到哪个算子明显低于它应有的 roofline，再考虑用 Triton 写；只有在 Triton 也表达不了的调度上（warp 专门化、精确的共享内存布局），才值得下到 CUDA/CUTLASS。

参考与延伸阅读

想看更多实验和 profiler 截图，推荐读 Ultra-Scale Playbook 的 GPU 一章；想要更系统的硬件模型，推荐 Scaling Book 的 GPU 章节。

• Nouamane Tazi et al., The Ultra-Scale Playbook，「Diving in the GPUs – fusing, threading, mixing」一章（A primer on GPUs、Improving performance with kernels、Memory coalescing、Tiling、Thread coarsening、Minimizing control divergence、Fused kernels、FlashAttention）。
• Jacob Austin et al., How to Think About GPUs，JAX Scaling Book：What Is a GPU?、Memory、Summary of GPU specs。
• Tri Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（2022），第 2–3 节与附录 B。
• Tri Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning（2023），第 3 节。
• Jay Shah et al., FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision（2024），第 2–3 节。
• Maxim Milakov, Natalia Gimelshein, Online normalizer calculation for softmax（2018）。
• Weile Luo et al., Benchmarking and Dissecting the Nvidia Hopper GPU Architecture（2024），表 IV（各级内存延迟）。
• NVIDIA, CUDA C++ Best Practices Guide：Coalesced Access to Global Memory、Shared Memory and Memory Banks、Occupancy；CUDA Programming Guide：Compute Capabilities。
• Simon Boehm, How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance。
• Triton 教程：Matrix Multiplication（分块与 L2 分组启动顺序）。