这一周的目标不是背 PyTorch 源码路径,而是建立一条主线: Python 表达式如何变成 GPU 上的一串 kernel,这些 kernel 又为什么能被编译器合并。 今天先把这条链路补齐,后面看分布式训练时才不会只盯着 API。

GPU 的峰值 FLOPS 很吓人,但很多模型片段并不是算力瓶颈,而是内存带宽瓶颈。 对 pointwise / reduction / softmax 这类算子,每多拆一个 kernel, 就可能多一次 HBM 读写、多一次 kernel launch、多一次中间张量分配。 融合的本质是:把中间值留在寄存器或 shared memory 里,不要写回 HBM。

三类常见融合

Roadmap 要求今天读一篇 FlashAttention 论文。不要一上来抠每个 CUDA 细节, 先抓住论文最重要的工程判断:标准 attention 慢,不是因为公式复杂,而是因为把 N×N 的 attention score 和 probability 写进 HBM。 FlashAttention 用 tiling 把 Q/K/V 分块搬进 SRAM,用 online softmax 保持数值稳定,最后只写输出 O。

online softmax 的最小直觉

softmax 需要每一行的最大值和指数和。分块后,你看不到完整一行,所以要维护两个状态: m 表示到目前为止的最大值,l 表示按当前最大值缩放后的 exp 和。 新 block 来了以后,如果最大值变大,旧的 l 和输出累积量都要乘一个缩放因子, 这样结果和一次性对整行做 softmax 等价。

# 伪代码:一行 attention 的 online softmax 直觉,不是可直接运行实现
m = -inf
l = 0
o = 0

for k_block, v_block in blocks(K, V):
    scores = q @ k_block.T
    m_new = max(m, scores.max())
    alpha = exp(m - m_new)
    p = exp(scores - m_new)
    l_new = alpha * l + p.sum()
    o = (alpha * l * o + p @ v_block) / l_new
    m, l = m_new, l_new

今天不用手写 FlashAttention kernel。目标是建立工程直觉: 用 benchmark 看出融合减少了 kernel 数和中间张量压力,再用 PyTorch 的 SDPA 路径体验现代 attention backend 的选择。

实验 1:eager vs torch.compile 的 pointwise fusion

import torch
from torch.utils.benchmark import Timer

torch.manual_seed(0)
device = "cuda"
x = torch.randn(8192, 8192, device=device, dtype=torch.float16)
bias = torch.randn(8192, device=device, dtype=torch.float16)

def eager_fn(x, bias):
    y = x + bias
    y = torch.nn.functional.gelu(y)
    y = y * 0.5 + 0.1
    return y

compiled_fn = torch.compile(eager_fn, mode="max-autotune")

# warmup
for _ in range(10):
    eager_fn(x, bias)
    compiled_fn(x, bias)
torch.cuda.synchronize()

print(Timer(stmt="eager_fn(x, bias)", globals=globals()).blocked_autorange())
print(Timer(stmt="compiled_fn(x, bias)", globals=globals()).blocked_autorange())

实验 2:检查 attention backend

import torch
import torch.nn.functional as F

B, H, N, D = 2, 16, 2048, 64
q = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
k = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
v = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)

torch.cuda.reset_peak_memory_stats()
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
torch.cuda.synchronize()
print(out.shape)
print(torch.cuda.max_memory_allocated() / 1024**2, "MiB peak")

# 可选:打开 profiler,看是否走到 flash / efficient attention backend
# with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
#     F.scaled_dot_product_attention(q, k, v, is_causal=True)
# print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

实验 3:估算 attention 的中间矩阵成本

def mib(numel, bytes_per_elem=2):
    return numel * bytes_per_elem / 1024**2

for n in [1024, 2048, 4096, 8192, 16384]:
    score = 2 * 16 * n * n
    print(f"N={n:5d} attention score/prob matrix: {mib(score):8.1f} MiB each")

1. 画出从 y = gelu(x + bias) 到 CUDA kernel 的路径:Tensor / Dispatcher / ATen / Autograd / allocator 分别在哪一步出现?
2. 解释为什么连续 pointwise op 通常是 memory-bound,并用"读几次、写几次"估算融合前后的 HBM 流量。
3. 用自己的话讲清 FlashAttention 为什么不需要把完整 attention score 矩阵写到 HBM。
4. 写出 online softmax 维护的两个状态,并解释当新 block 最大值变大时旧贡献为什么要缩放。
5. 用 profiler 对比 eager 和 compile 后的 kernel 列表:kernel 数量、总 CUDA 时间、显存峰值分别有什么变化?
6. 列出 3 个不适合盲目融合的场景,说明风险是寄存器压力、shape 动态、数值稳定还是语义 side effect。