GPU 之上 — CUDA kernel 里训练长什么样
这一章讲三件事: 你写的每一行 PyTorch,在显卡上变成了什么(线程、块、网格); 为什么显卡上的瓶颈往往不是「算」而是「搬数据」;以及四个经典的并行翻车坑。 这是全书的最后一章拆解——到这里,从文字到 checkpoint 到显卡内部,整条链都打开了。
1. 先看现象:同一行代码,两个世界的速度
CPU 像几位教授:每个都极聪明,一次解一道难题。 GPU 像一间有一万六千名小学生的礼堂:每个人只会算一道乘法,但一万六千人同时算1。 神经网络恰好全是「几百万次互不依赖的乘加」,正是小学礼堂的主场—— 原书给的对照:CPU 一般 8-64 核,H100 超过 16,000 个 CUDA 核心2。
2. 主走查:一个 kernel 的诞生
kernel(kernel,在显卡上被成千上万线程同时执行的一个函数)是 GPU 编程的最小单位。 拿原书的第一个例子走一遍——把 1024 个数各翻一倍(线程编号的计算方式照录代码)3:
CPU 视角: for i in range(1024): data[i] *= 2 ← 串行,1024 次
GPU 视角: 写一个 kernel:算出「我是几号」,只翻自己那一个
编号 = 块号 × 每块线程数 + 块内编号
data[我的编号] *= 2
启动时说一句「派 8 个块、每块 128 线程」
→ 1024 个线程同时执行,每个翻自己的格子,一步完成
图说:这就是原书说的 SIMT——单指令、多线程:
所有线程执行同一份代码,只是各自处理数据的不同一角[^4]。
原书给的规模感:1024×1024 的矩阵乘,可以派 1,048,576 个线程,每个线程算输出矩阵的一个元素—— CPU 上毫秒级的活儿,显卡上微秒级完成4。
线程的组织是三层,Reporting 关系固定5:
| 层 | 是什么 | 原书例子 |
|---|---|---|
| 线程 | 最小执行单位,算一个格 | 上例的每个「翻倍工」 |
| 块 | 一组线程,块内可通过共享内存协作 | 每块 256-1024 个线程 |
| 网格 | 一次 kernel 启动的全部块 | 32×32 个块 |
3. 真正的瓶颈:搬数据,不是算
原书这一节推翻了新手的默认想象:现代 GPU 常常是「内存受限」的——算得快,喂不过来6。 显卡的显存层级像食堂:寄存器(自带饭盒,最快)、共享内存(块内公用的取餐台)、全局显存(大仓库,最慢)。 一个朴素的矩阵乘 kernel,每个元素都要从大仓库取一遍——大部分时间在排队等饭,不是在吃。
两个解法,原书都给了代码7:
- 共享内存+分块(tiling):把矩阵切成小方块,整块搬进取餐台,块内线程各取所需—— 每个元素从大仓库只读一次,而不是每个线程都读一遍;
- 算子融合(fusion):把「乘法→加偏置→激活」几个 kernel 合成一个,中间结果不落仓库—— 原书的结论:省下的搬运常常比省下的计算更值钱8。
第 07 章提过一句的 FlashAttention 就是同一思想的代表作:注意力矩阵根本不整体落盘, 分块在高速层内算完就丢,显存从 O(n²) 降到 O(n)9。
4. 训练循环的 kernel 视角
第 10 章的训练循环,在显卡上是一串 kernel 排队10:
前向:矩阵乘 kernel(算层输出)→ 激活 kernel → 归一化 kernel → loss kernel
反向:loss 梯度 kernel → 逐层梯度 kernel → 权重梯度 kernel → 梯度汇总 kernel
更新:参数更新 kernel →(可选)梯度裁剪 kernel
图说:第 02 章「反向传播=链式法则」的那一趟,在这里就是
十几类 kernel 的固定排班;每一步都是全卡并行。
PyTorch 的角色因此清楚了:它是个排班员——你写 x @ W,它挑 cuBLAS(高度调优的矩阵乘库)的 kernel;
你写注意力,它挑 cuDNN 或 FlashAttention 的 fused kernel11。
「库先于自写」是原书的明确建议:那些库是多年的调优结晶;自写 kernel 只在三种情况值得——
全新算子、显存关键路径、可融合的机会12。
5. 四个经典翻车坑
原书列的坑,全可以用「一万六千人同时干活」的类比记住13:
| 坑 | 人话版 | 解法 |
|---|---|---|
| 竞态 | 两个人同时改同一格,结果看谁手快 | 同步点(块内屏障)或原子操作 |
| bank 冲突 | 取餐台分 32 个窗口,全挤一个窗口=退化成排队 | 让相邻线程取相邻地址 |
| 分支发散 | 同组线程走了 if 的两边,两边都得慢慢执行完 | 少分支;同组走同一路 |
| 占用不足 | 每个线程占的饭盒(寄存器)太多,礼堂里挤不下足够多的人 | 限制寄存器/共享内存用量 |
6. 作者的判断与证据
| 说法 | 书里的证据 | 我们的标注 |
|---|---|---|
| 神经网络天然适合 GPU | 矩阵乘/逐元素运算天然并行的机制描述14 | 领域共识 |
| 内存受限、融合最值钱 | 「Modern GPUs are memory-bound」的直接论断+tiling/融合的代码演示6 | 与 HPC 领域共识一致 |
| 库优先于自写 | cuBLAS/cuDNN「代表多年的优化」+自写 kernel 的三个适用条件12 | 工程共识 |
| 四个坑 | 逐一给了成因与解法13 | CUDA 编程的标准清单 |
判断(我们的,不是书里的): 这一章表面上讲 CUDA,实际讲的是**「性能问题要换一个抽象层才能看清」** ——第 04 章的 ε、rsqrt,第 06 章的滑窗,第 10 章的混合精度,单看都是小聪明; 换到「数据搬运是主要成本」这一层看,它们全是同一件事:少搬一次数。 这也是为什么同一份模型代码,换硬件、换实现,速度能差十倍。 如果错,会错在: 如果未来硬件把算与存的边界抹掉(行话叫存算一体:把存储和计算合成一体的硬件思路), 「搬运是主要成本」的框架就会过时——但「换抽象层看性能」这个方法论不会。
7. 边界与局限
- 原书没讲 AMD 的 ROCm 或 TPU 这类非 CUDA 生态——标题里的「Modern Frameworks」实际只覆盖 NVIDIA 系;
- kernel 代码示例是教学级(朴素矩阵乘),原书自己注明「真实实现会用 tiling、共享内存与寄存器分块」15;
- 「如何给一个模型做性能剖析(找出哪个 kernel 慢)」原书没有覆盖——这是把本章知识用起来的第一步;
- Triton 这类「用 Python 写 GPU kernel」的新工具原书只在优化库清单里带过一笔16,没展开。
8. 可带走的
- GPU = 一万六千个只会乘法的小学生;神经网络全是乘法,所以是主场;
- kernel = 被上万线程同时执行的函数;线程→块→网格三层编制;
- 1024×1024 矩阵乘 = 派一百万个线程,每人算一个格;
- 现代 GPU 内存受限:瓶颈在搬数据,不在算;
- 共享内存+分块让每个数从大仓库只读一次;算子融合让中间结果不落仓库;
- FlashAttention = 「注意力不落盘」的融合,显存从 O(n²) 到 O(n);
- 库优先(cuBLAS/cuDNN);自写 kernel 只为全新算子、显存关键路径、可融合三件事;
- 四个坑=一万六千人同时干活的四种乱象:抢格子、挤窗口、走岔路、占座太多;
- 性能问题要下潜一层抽象才能看清——少搬一次数,是所有优化共同的名字。
9. 原文地图
| 主题 | 原书章 | 原文位置 |
|---|---|---|
| CUDA 是什么 | Training Neural Networks with CUDA Kernels and Modern Frameworks | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:9(搜「Compute Unified Device Architecture」) |
| CPU 对照 GPU 核数 | Training Neural Networks with CUDA Kernels and Modern Frameworks | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:13(搜「16,000」) |
| 适合 GPU 的运算 | Training Neural Networks with CUDA Kernels and Modern Frameworks | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:15(搜「Matrix multiplications」) |
| 1000×1000 的并行 | Training Neural Networks with CUDA Kernels and Modern Frameworks | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:25(搜「1000x1000」) |
| kernel 定义与 global | CUDA Kernels Explained | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:31(搜「CUDA kernel is a function」) · text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:47(搜「global」) |
| 线程-块-网格 | Thread Hierarchy | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:35(搜「Thread Hierarchy」) |
| 百万线程例 | Thread Hierarchy | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:43(搜「1,048,576」) |
| 内存层级 | Memory Hierarchy | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:69(搜「Memory Hierarchy」) |
| 训练循环的 kernel 排班 | CUDA Kernels in Neural Network Training | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:87(搜「forward pass」) |
| 矩阵乘 kernel 代码 | Practical CUDA Kernel Examples | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:15(搜「matrix multiplication」) |
| 朴素实现的问题 | Practical CUDA Kernel Examples | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:249(搜「shared memory」) |
| tiling 演示 | Practical CUDA Kernel Examples | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:319(搜「loaded from global memory only once」) |
| 库优先建议 | Integration with Deep Learning Frameworks | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:571(搜「PyTorch Custom CUDA Extensions」) |
| 四个坑 | Common Pitfalls | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:1009(搜「Race Conditions」) · text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:1027(搜「Bank Conflicts」) · text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:1029(搜「Divergent Branches」) · text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:779(搜「Occupancy」) |
| 关键要点清单 | Common Pitfalls 后的 Key Takeaways | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:1035(搜「Key Takeaways」) · text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:1039(搜「Fusion Is Powerful」) |
| 术语表(章末附录) | Appendix: Glossary of Terms | text/95-fm-training-neural-networks-with-cuda-kernels-and-m.txt:1049(搜「Glossary」) |