同一个 GeLU,用 PyTorch 按公式一项一项写,和调用 F.gelu(x, approximate="tanh"),算出来的数完全一样(误差 5e-7),在一台 Apple 芯片的 MPS 后端上前者 206 ms、后者 11 ms,差了将近 20 倍。两者做的浮点运算一样多,差别在于前者把数据在显存和计算单元之间来回搬了七八趟,后者只搬一趟。这件事是理解 GPU 的钥匙:**GPU 的算力多到用不完,真正稀缺的是把数据送到算力面前的带宽。**这一篇从硬件结构讲起,讲到怎么判断一个算子卡在哪里,最后用 Triton 写出几个真正的 kernel。

前置:内存层次与缓存(CPU 上的 cache 与局部性)、资源核算与 Scaling Laws(FLOPs 与算术强度怎么算)

为吞吐而生:SM、warp 与 Tensor Core

CPU 的设计目标是让一个线程尽快跑完:大块芯片面积给了分支预测、乱序执行和大缓存。GPU 反过来,目标是单位时间里处理的总数据量,所以它把面积换成成千上万个简单的计算单元,对分支的支持弱得多。

一块 GPU 由许多个 SM(streaming multiprocessor)组成,每个 SM 可以看成一个独立的小处理器。H100 SXM5 有 132 个 SM,A100 有 108 个(GA100 完整芯片是 128 个,出货时开了 108 个)。写 GPU 程序时,线程分三层组织:

  • thread:最小的执行单位。所有线程跑同一份代码,只是处理的数据不同,这叫 SIMT(single instruction, multiple threads)。
  • thread block(也叫 CTA):一组线程,保证被调度到同一个 SM 上,共享这个 SM 上的一块 shared memory。block 之间不能直接通信。
  • grid:一次 kernel 启动的全部 block。

硬件真正调度的单位是 warp:编号连续的 32 个线程,每个周期执行同一条指令。这带来一个直接后果:如果一个 warp 里有的线程走 if、有的走 else,两段代码只能先后执行,另一半线程被 mask 掉空等,这叫 control divergence。哪怕只有一个线程走了另一个分支,两段代码也都得走一遍。所以 GPU 代码里写 ReLU 这类操作常用乘法或 where,不写分支。

warp 还是 GPU 隐藏延迟的手段。一次读显存要几百个时钟周期,SM 不会傻等:它上面同时驻留着很多 warp(H100 每个 SM 最多 64 个),一个 warp 在等数据,调度器下个周期就发射另一个 warp 的指令。驻留的 warp 越多,延迟越容易被藏住。

最后是算力的来源。从 V100 开始,每个 SM 里有专门做小矩阵乘的 Tensor Core。H100 SXM5 的 BF16 稠密 Tensor Core 峰值是 989 TFLOP/s,而不走 Tensor Core 的 FP32 运算只有约 67 TFLOP/s,算 exp 这类超越函数的 SFU 更少。矩阵乘比其他运算快一个数量级以上,这条硬件事实决定了 Transformer 为什么是一堆矩阵乘:想要从硬件扩展里拿到好处的架构,必须让大部分 FLOPs 落在矩阵乘上。

TPU 走的是同一条路。它的 TensorCore(注意,名字和 GPU 的 Tensor Core 撞了,指的是类似 SM 的处理器)由标量单元、向量单元 VPU 和大型矩阵乘单元 MXU 组成,外面同样挂一块 HBM。区别在粒度:H100 有 132 个 SM、528 个小 Tensor Core,TPU v5p 每芯片只有 2 个 TensorCore、8 个大 MXU,所以 TPU 只擅长大矩阵,维度太小时用不满;而且 TPU 没有 warp 这一层。两家独立设计出几乎一样的结构,说明”轻控制 + 大矩阵乘单元 + 快的片上内存 + 慢的 HBM”差不多是做 matmul 的唯一答案,下面讲的方法在两者上都适用。

内存层级:离 SM 越近,越快也越小

数据从远到近有四层。H100 SXM5 的规格(NVIDIA 白皮书;L1/shared 和 L2 带宽是常被引用的估算值,不同来源略有出入):

层级位置容量带宽(全卡)谁来管
寄存器每个 SM256 KB(64K 个 32 位寄存器)远高于下面各层编译器
L1 / shared memory每个 SM合计 256 KB,shared 最多 228 KB约 31–33 TB/sL1 硬件管,shared 程序管
L2 cache全芯片共享50 MB约 12 TB/s硬件
HBM3(显存)芯片旁80 GB3.35 TB/s程序(cudaMalloc)

两个细节值得记住。第一,L1 和 shared memory 是同一块物理 SRAM,只是管理方式不同:L1 由硬件自动缓存,什么时候被淘汰你控制不了;shared memory 由程序显式写入,在 block 结束前一直在。高性能 kernel 几乎都靠 shared memory,因为一块要被反复用的数据不能冒被 L1 中途换出的风险。第二,132 个 SM 的片上 SRAM 加起来只有 30 多 MB,而一个 4096×4096 的 bf16 矩阵就有 32 MB,所以大张量永远住在 HBM,每次只能搬一小块进 SM。

HBM 本身已经是工程奇迹:DRAM 颗粒竖着堆 8–12 层,用硅通孔打通,和 GPU 一起坐在硅中介层上,每个堆栈 1024 位宽,H100 的 5 个堆栈合起来 3.35 TB/s,是 PCB 上平铺的 GDDR 做不到的。但它相对算力仍然太慢。CS336 引用的 20 年数据:硬件 FLOP/s 涨了约 60000 倍,DRAM 带宽涨了约 100 倍,互连带宽约 30 倍。A100 到 H100 这一代也是一样,BF16 算力从 312 涨到 989 TFLOP/s(约 3.2 倍),HBM 带宽从 2.04 涨到 3.35 TB/s(约 1.6 倍)。每一代新卡都比上一代更”饿”,同一个 kernel 换到新卡上更容易被带宽卡住。

停一下,检查理解单选

H100 每个 SM 的 L1 和 shared memory 合计 256 KB。高性能 kernel 为什么要把反复使用的数据块显式放进 shared memory,而不是靠 L1 自动缓存?

Roofline:先判断算子卡在哪

既然算力和带宽都有上限,一个算子能跑多快就取决于它每读一个字节做多少次运算,也就是算术强度(FLOPs / 字节,怎么数见 资源核算与 Scaling Laws)。Roofline 模型把这件事画成一张图:横轴是算术强度,纵轴是可达到的 FLOP/s,

可达到的 FLOP/s = min( 峰值 FLOP/s ,  算术强度 × 内存带宽 )
拐点 = 峰值 FLOP/s / 内存带宽

拐点左边是斜线,吞吐被带宽限制(memory-bound);右边是平顶,计算单元已经满了(compute-bound)。用 H100 SXM5 的数字算几个典型算子:

peak, bw = 989e12, 3.35e12            # H100 SXM5: BF16 dense, HBM3
print(f"ridge point = {peak / bw:.0f} FLOP/byte")
def ai_matmul(n, b=2): return 2*n**3 / (3*n*n*b)   # read A, B, write C
for name, ai in [("ReLU bf16", 1/4), ("add bf16", 1/6),
                 ("matmul 512", ai_matmul(512)), ("matmul 4096", ai_matmul(4096))]:
    att = min(peak, ai * bw)
    print(f"{name:12s} AI={ai:8.2f}  attainable={att/1e12:7.1f} TFLOP/s  ({att/peak:6.1%} of peak)")
ridge point = 295 FLOP/byte
ReLU bf16    AI=    0.25  attainable=    0.8 TFLOP/s  (  0.1% of peak)
add bf16     AI=    0.17  attainable=    0.6 TFLOP/s  (  0.1% of peak)
matmul 512   AI=  170.67  attainable=  571.7 TFLOP/s  ( 57.8% of peak)
matmul 4096  AI= 1365.33  attainable=  989.0 TFLOP/s  (100.0% of peak)

每从 HBM 读一个字节要做约 295 次运算,Tensor Core 才不闲着。逐元素操作只能用到峰值的千分之一,这不是写得不好,是它本质上就被带宽锁死了。A100 的拐点是 312 / 2.04 ≈ 153,所以一个强度 200 的 kernel 在 A100 上是 compute-bound,到了 H100 上反而变成 memory-bound。

这张图直接告诉你该往哪个方向优化。memory-bound 的算子想变快,只有两条路:减少要搬的字节(低精度、融合、重算),或者让每个字节被多用几次(tiling)。compute-bound 的算子再省内存也没用,只能换更快的计算路径(Tensor Core、更低精度的矩阵乘)。看一张卡的参数表也是同样的道理:H20 的算力只有 H100 的约六分之一,但保留了 96 GB HBM3 和完整带宽,而推理 decode 阶段的强度只有个位数(见 推理系统),所以它跑推理远比算力数字看上去好。

停一下,检查理解单选

H100 上一个 bf16 算子的算术强度是 30 FLOP/字节。把它的数据改成 fp8(字节数减半、FLOPs 不变),按 roofline 大约能快多少?

少搬几趟:kernel fusion

回到开头的 GeLU。PyTorch 的 eager 模式里,每一个基本运算都是一个独立的 kernel:读输入出 HBM,算完写回 HBM,下一个 kernel 再读回来。按 tanh 近似公式写的 GeLU 有乘方、乘、加、tanh、再乘再加,每一步都往返一次显存:

import torch, time, math
x = torch.randn(64 * 1024 * 1024, device="mps")      # 64M floats = 256 MB
def naive_gelu(x):
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * x ** 3)))
def bench(f, n=20):
    for _ in range(3): f()
    torch.mps.synchronize(); t = time.perf_counter()
    for _ in range(n): f()
    torch.mps.synchronize()      # GPU 是异步的,不同步只量到了发射 kernel 的时间
    return (time.perf_counter() - t) / n * 1e3
print(f"naive  {bench(lambda: naive_gelu(x)):6.2f} ms")
print(f"fused  {bench(lambda: torch.nn.functional.gelu(x, approximate='tanh')):6.2f} ms")
naive  206.28 ms
fused   10.63 ms

(Apple M 系列 MPS 后端的真实结果,数字随机器变化,比例是重点。)融合版本只读一次 x、写一次结果,所有中间量留在寄存器里。这就是 kernel fusion。它不需要什么数学技巧,F.gelu 就是有人手写了一个融合好的 CUDA kernel;torch.compile 则会分析计算图,把相邻的逐元素操作自动生成一个融合的 Triton kernel。

数一下 softmax 能看清收益的上限。对 M 行 N 列的矩阵,按”求 max、减 max、exp、求和、相除”五个 kernel 写,读约 5MN、写约 3MN,合计约 8MN 次元素读写;融合成一个 kernel 只读 MN、写 MN。理论上快 4 倍。

benchmark 本身有几个必须遵守的规矩,上面的代码都做了:先 warmup(第一次调用可能包含编译和分配);计时前后都同步(GPU 异步执行,CPU 发出 kernel 就往下走);多跑几次取平均或中位数;先用 torch.allclose 确认结果正确再谈速度;在 CUDA 上用 torch.cuda.Event 计时以避开 CPU 开销。找瓶颈时先用 PyTorch profiler 或 Nsight 看时间花在哪些 kernel 上,再决定要不要自己写。

降低精度是另一种”少搬”。同样的 ReLU,fp32 每个元素读写 8 字节,bf16 只要 4 字节,算术强度翻倍。对 compute-bound 的矩阵乘,低精度的收益来自 Tensor Core 的专用路径:H100 上 FP8 峰值是 BF16 的 2 倍。哪些运算能降精度、怎么保持训练稳定,是 并行训练 里混合精度训练的话题。第三种”少搬”是重算:与其把中间结果写进显存给反向用,不如反向时从输入重新算一遍,这个思路在 FlashAttention 里发挥到了极致。

停一下,检查理解单选

sin(x)**2 + cos(x)**2 不融合时是 5 个 kernel(sin、cos、两次平方、相加)。按"每个 kernel 读输入、写输出都走 HBM"数,不融合和融合后各读写多少个长度为 n 的张量?

访问模式:合并访存、bank conflict 与 tiling

字节数相同,读的顺序不同,速度也能差很多。下面在 CPU 上做个对照:拷贝一个 4096×4096 的 fp32 矩阵,和把它转置成连续存储。两者读写的字节数完全一样,计算量都是零:

import torch, time
x = torch.randn(4096, 4096)
def bench(f, n=20):
    f(); t = time.perf_counter()
    for _ in range(n): f()
    return (time.perf_counter() - t) / n * 1e3
print(f"clone            {bench(lambda: x.clone()):6.2f} ms")
print(f"t().contiguous() {bench(lambda: x.t().contiguous()):6.2f} ms")
clone              2.35 ms
t().contiguous()  31.92 ms

CPU 上的原因是 cache line(见 内存层次与缓存),GPU 上的机制类似但更严格。DRAM 按一整段(burst)读取,一个 warp 的 32 个线程同时发出访存请求时,硬件会看这 32 个地址落在几个段里。如果相邻线程读相邻地址(线程 0 读 M[0][0],线程 1 读 M[0][1]……),32 × 4 字节正好是一条 128 字节的事务,一次取完,这叫合并访存(coalescing)。转置时读端连续,写端的 32 个线程却写到 32 个不同的行,每个地址相隔 4096 个元素,各占一次至少 32 字节的事务,却只用了 4 字节,有效带宽只剩 12.5%。

关键在于同一时刻、相邻线程的地址是否相邻,而不是单个线程自己往哪走。在行主序矩阵上,让”每个线程负责一行、沿行往右走”反而不合并:第一步四个线程读的是四行的开头,落在四个段里。

转置的解法是用 shared memory 转一手:一个 block 先按行连续地读一块进 shared memory,在片上完成转置(片上随机访问几乎免费),再按行连续地写回 HBM,两端都恢复合并。片上也有自己的规矩:shared memory 分成 32 个 bank,每个 4 字节宽,多个线程同时访问同一个 bank 的不同地址就要排队,这叫 bank conflict。按列读 float tile[32][32] 时,一列的 32 个元素恰好都在同一个 bank 上;声明成 tile[32][33],每行错开一个 bank,冲突消失。CUTLASS 和 FlashAttention-3 用的是更一般的 swizzle(用异或打乱列地址)。

把”先搬进片上、再反复用”推到极致,就是 tiling,这是写 kernel 最重要的一招。朴素的 N×N 矩阵乘里,每个输入元素要被 N 个输出元素用到,就从 HBM 读 N 次。把输出切成 T×T 的块,每个 block 负责一块,沿 k 方向一对一对地把 A、B 的 tile 读进 shared memory、在片上算完部分和再读下一对,每个输入元素只从 HBM 读 N/T 次,剩下的 T 次复用都发生在 shared memory 里。N = 4096、T = 128 时,读取次数从 4096 降到 32。算术强度从 O(1) 提到 O(T),T 受 shared memory 和寄存器容量限制。

tiling 也带来两种量化损失。一是维度整除不了 tile:256×257 的矩阵用 128×128 的 tile 要开 2×3 = 6 个 block,多出的两个几乎是空的。二是 wave quantization:block 按”波”分配到 SM 上。A100 上用 256×128 的 tile 算 1792 方阵,tile 数是 7 × 14 = 98,不超过 108 个 SM,一波算完;边长加 1 变成 1793,tile 数变成 8 × 15 = 120,第一波 108 个,第二波只有 12 个 SM 在干活,96 个空等,时间接近翻倍。实践建议很朴素:自己定的维度取 64 或 128 的倍数。Karpathy 把 nanoGPT 的词表从 50257 补到 50304(64 × 786),据他报告提速约 25%。

停一下,检查理解单选

8192×8192 的矩阵乘法,用 64×64 的 tile。每个输入元素从 HBM 读取的次数大约从多少降到多少?

写一个 kernel:Triton

CUDA 要你写一个线程做什么:自己用 threadIdx、blockIdx 算下标,自己声明 __shared__ 数组,自己插 __syncthreads()。Triton(OpenAI)把编程单位提升到一个 block:你写的是对一整块数据的向量运算,线程怎么分工、数据放寄存器还是 shared memory、何时同步,都交给编译器。代价是够不着最新硬件的全部特性。

每个 Triton kernel 都是同一个结构:我是第几个 block → 我负责哪些下标 → load → 计算 → store。下面是 GeLU(示意:需要 NVIDIA GPU 和 triton 包,本机无法运行):

import torch, triton
import triton.language as tl
 
@triton.jit
def gelu_kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr):   # 块形状参与代码生成,必须编译期已知
    pid = tl.program_id(axis=0)                           # 相当于 CUDA 的 blockIdx.x
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n                                       # 最后一个 block 可能越界
    x = tl.load(x_ptr + offs, mask=mask)                  # HBM -> 片上
    a = 0.79788456 * (x + 0.044715 * x * x * x)
    e = tl.exp(2 * a)
    y = 0.5 * x * (1 + (e - 1) / (e + 1))                 # tanh 用 exp 拼出来
    tl.store(y_ptr + offs, y, mask=mask)                  # 片上 -> HBM,只写一次
 
def gelu(x):
    y = torch.empty_like(x)                               # kernel 没有返回值,输出先分配好
    grid = (triton.cdiv(x.numel(), 1024),)
    gelu_kernel[grid](x, y, x.numel(), BLOCK=1024)
    return y

mask 的作用用一个具体数字看:n = 5000、BLOCK = 1024 时启动 cdiv(5000, 1024) = 5 个 block,最后一个 pid = 4 负责下标 4096..5119,其中只有 904 个小于 5000,其余位置不读不写。

softmax 是下一个台阶:它要对整行求 max 和 sum,不再是逐元素。一行放得进一个 block 时,让每个 block 负责一行,BLOCK = next_power_of_2(N),越界的列用 other=-float("inf") 填充,exp 之后正好是 0,不影响求和。如果误填成 0,[5, 5, 5] 这一行的分母会多出一项 e^(0−5),每个结果变成约 0.3326 而不是 1/3。一行放不进一个 block 时(比如 4096 列、BLOCK 1024),block 仍负责一整行,在内部循环 4 次,每次读一个 tile 加进同一个累加器,最后 tl.sum 归约。不能把 4 个 tile 分给 4 个 block,否则一行的部分和散在不同 block 里,只能经 HBM 再合并。

matmul 是两者的组合。每个 block 负责 C 的一个 BLOCK_M × BLOCK_N tile,沿 k 方向循环:

acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)     # 累加器用 fp32,输入可以是 bf16
for k in range(0, K, BLOCK_K):
    a = tl.load(a_ptrs, mask=..., other=0.0)             # [BLOCK_M, BLOCK_K]
    b = tl.load(b_ptrs, mask=..., other=0.0)             # [BLOCK_K, BLOCK_N]
    acc += tl.dot(a, b)                                  # 编译器把它映射到 Tensor Core
    a_ptrs += BLOCK_K * stride_ak
    b_ptrs += BLOCK_K * stride_bk
acc = tl.maximum(acc, 0.0)                               # ReLU 在写回前做,又省一次 HBM 往返
tl.store(c_ptrs, acc, mask=...)

(示意。)每个 tile 读 (BLOCK_M + BLOCK_N) × K 个元素,总读取量是 MNK × (1/BLOCK_M + 1/BLOCK_N),64×64 的 tile 时为 MNK/32,比朴素写法的 2MNK 少 64 倍。最后一行的 ReLU 展示了 tiling 附带的好处:激活函数直接融合进矩阵乘的尾巴(epilogue)。

块大小怎么选没有通用答案。块太大,每个线程占的寄存器多,SM 上能驻留的 warp 少(occupancy 低),延迟藏不住;块太小,复用不够。例如每个 block 128 线程、每线程用 160 个寄存器,一个 SM 的 65536 个寄存器只够放 3 个 block,即 12 个 warp,occupancy 为 12/64 = 18.75%。occupancy 低不一定慢,让每个线程多干点活(thread coarsening)有时更快。实际做法是 @triton.autotune 列出若干组 BLOCK_*、num_warps、num_stages,第一次调用时都跑一遍,按 shape 缓存最快的那组。

停一下,检查理解多选

关于 Triton 的 block 编程模型,哪些说法正确?

从 PyTorch 到机器码:CUDA 软件栈

写完 kernel,还要知道它怎么跑到卡上。从下往上四层:

层是什么怎么看
driver内核模块 + 用户态库 libcuda,直接驱动硬件nvidia-smi 的 Driver Version
CUDA runtime / toolkitlibcudart(cudaMalloc 等)、编译器 nvccnvcc --version
加速库cuBLAS(矩阵乘)、cuDNN(神经网络算子)、NCCL(多卡通信)各自的 .so
PyTorch把 a @ b 分派给上面的库或自己的 kerneltorch.version.cuda

nvidia-smi 右上角的 CUDA Version 是 driver 最高支持的 CUDA 版本,不代表装了 toolkit。Linux 上 pip 装的 CUDA 版 PyTorch 把 runtime、cuBLAS、cuDNN、NCCL 都作为依赖装进了 venv,运行时只依赖宿主机的 driver;只有编译自己的 CUDA 扩展才需要 nvcc(编译和链接的机制见 编译链接与 Python 扩展)。由此得到第一条兼容规则:程序自带的 runtime 不能比 driver 新,driver 则向后兼容旧程序。容器隔离整个用户态,但不隔离 driver,容器里看到的是宿主机的 driver。

第二条规则关于架构。nvcc 先把 .cu 编译成 PTX(一种面向虚拟架构 compute_XY 的中间汇编,Triton 也编译到这一层),再由 ptxas 编译成某个具体架构 sm_XY 的机器码 SASS。XY 是 compute capability:A100 是 8.0,H100 是 9.0,B200 是 10.0。SASS 只在同一主版本内向上兼容(sm_80 能跑在 8.6 上,不能跑在 9.0 上);PTX 可以被 driver 当场 JIT 成任何更新架构的 SASS。两者都缺,就是熟悉的 no kernel image is available for execution on the device。所以编译扩展时用 TORCH_CUDA_ARCH_LIST="8.0;9.0+PTX" 这样的写法,既给已知的卡放 SASS,也给未来的卡留一份 PTX。

kernel 名字也透露了这些信息。profiler 里看到 cutlass3x_sm100_..._64x64x16_... 这样的名字:CUTLASS 是 NVIDIA 的模板库,sm100 是 Blackwell,64x64x16 是 tile 形状。PyTorch 会按张量形状挑不同的 kernel,同一个 a @ b 在 dim = 128 和 dim = 2048 时调用的就不是同一个。

停一下,检查理解单选

一个 CUDA 扩展只用 -gencode arch=compute_80,code=sm_80 编译(只有 SASS,没有 PTX)。它能在哪些卡上直接运行?

记住这三件事

  1. GPU 的算力远多于带宽能喂饱的量(H100 SXM5 每读 1 字节 HBM 要做约 295 次 BF16 运算),所以先用 roofline 判断算子是 memory-bound 还是 compute-bound,再决定优化方向。
  2. 几乎所有 kernel 优化都是在少碰 HBM:融合减少往返,低精度减少字节,tiling 把重复读挪到 shared memory,合并访存让每次事务不浪费。
  3. Triton 让你按 block 写 kernel,结构永远是”定位 block → 算下标和 mask → load → 片上计算 → store”,块大小靠 autotune 而不是猜。

延伸阅读