GPU计算基础：CUDA编程与GPU架构
要玩转分子生成训练，GPU你不是用着，而是得“懂”它。我刚开始接触CUDA时，以为把for循环扔到GPU上就完事了，结果跑了三天三夜还没收敛……后来才明白，不懂架构的调优就是瞎蒙。
1. GPU架构概览——为什么它那么快
先看一张我经常画给新同事的手绘图。它解释了一个核心问题：为什么GPU能同时处理上千个线程？
这张图藏着几个关键点：
- SM（Streaming Multiprocessor）：GPU 的最小计算单元，每 SM 有多个 CUDA Core。我用的 A100 每 SM 有 64 个 FP32 Core。
- 共享内存（Shared Memory）：每 SM 独享，延迟远低于全局内存。但只有几十 KB，得精打细算。
- 线程块（Block）：分配到某个 SM，块内线程能共享共享内存并进行同步。
一个小经验：我调分子生成网络时，刚开始把所有参数都放全局内存里，结果带宽瓶颈严重。后来把一些频繁访问的权重塞到共享内存里，迭代速度提升了 30%。
2. CUDA 线程模型与内存层次
写 CUDA 你得理解三个层次：Grid（网格）→ Block（线程块）→ Thread（线程）。每个线程都有自己独一无二的 ID，靠它来区分处理哪部分数据。
看一个最简单的向量加法 kernel，我猜你没写过也见过：
__global__ void vecAdd(float *a, float *b, float *c, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
c[idx] = a[idx] + b[idx];
}
}这里 blockIdx.x 是块索引，threadIdx.x 是块内线程索引，blockDim.x 是每块线程数。组合起来就得到全局索引。
为什么这么设计？说白了就是让成千上万个线程能并行跑，每个线程只管自己那一小块数据。
| 内存类型 | 作用域 | 速度 | 典型大小 |
|---|---|---|---|
| 全局内存 | 所有线程 | ~200 GB/s 带宽 | 几 GB ~ 80 GB |
| 共享内存 | 单个 Block | ~1 TB/s 带宽 | 48~164 KB / SM |
| 寄存器 | 单个 Thread | 零延迟 | 256 / SM（每个 SM） |
| 局部内存 | 单个 Thread | 同全局内存 | 由编译器决定 |
你想想看，如果你把线程内频繁用到的变量不加限制地声明，可能就被编译器推到局部内存里了。局部内存其实在全局内存空间，速度慢很多。我曾经遇到过核函数里声明了一个大数组，结果性能直接腰斩。解决方案？改成动态分配，或者尽量用寄存器变量（__shared__ 不行，那就用普通局部变量但长度小一点）。
注意：共享内存虽然快，但太小。而且同一 warp 内的线程访问共享内存时如果发生 bank conflict（比如多个线程访问同一个 bank 的不同地址），性能会下降。我调一个分子力场计算时，就因为共享内存 bank conflict 导致 occupancy 高但实际性能低，后来把数据排成 float2 结构对齐才解决。
3. 性能考量与调优方向
调优 GPU 程序，我习惯从三个维度找瓶颈：
- Occupancy – 每个 SM 活跃的 Warp 数量。占有率越高，隐藏延迟能力越强。
- 内存带宽利用率 – 全局内存访问是否合并（coalesced）？比如
float4一次读 16 字节比四个单独读快得多。 - 计算与访存比 – 计算密集还是访存密集？分子生成里能量损失计算通常是访存瓶颈，而反向传播可能是计算瓶颈。
举个例子，如果你想计算 occupancy，可以用 NVIDIA 的 cuda-occupancy-calculator 工具，或者直接在代码里调用 cudaOccupancyMaxActiveBlocksPerMultiprocessor：
int minGridSize, blockSize;
cudaOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize,
vecAdd, 0, n);
// blockSize 就是建议的线程块大小
// 然后在 kernel 启动时用这个 blockSize我个人习惯在训练循环里动态调整 blockSize，因为不同层次的网络层，计算密度不一样。拿分子生成模型来说，注意力层和 MLP 层的线程需求就不同。
一条调优经验：把 blockSize 设成 32 的倍数（比如 128、256），而且最好不超过 1024。记不清了？就设 256。大部分情况它都不是最差的。
好了，GPU 的基础架构和 CUDA 编程要点就聊这么多。你大概心里有数了：CPU 负责控制流，GPU 负责数据并行。SM 是计算元，共享内存是加速利器，线程块是调度单位。后面我们会把这一套用在实际分子生成训练里，比如怎么把分子图数据塞进 GPU，怎么用 CUDA 加速消息传递层。嗯，先消化这些吧。