1. AI训练与GPU计算基础：为什么需要GPU？
大家好，我是你们这趟旅程的向导。咱们直接开门见山——搞AI训练，为什么非得用GPU？CPU不行吗？
说实话，我早年刚入行时也这么想过。那时候我还在用CPU跑一个简单的两层神经网络，训练一次要等一宿。后来换了块入门级GPU，同样的模型，十几分钟就跑完了。你想想看，这差距从哪来的？
核心就一句话：CPU是“通才”，GPU是“专才”。
CPU擅长处理复杂的逻辑分支、串行任务，它把大量晶体管用在了缓存和控制单元上。而GPU呢？它把绝大部分晶体管都堆成了计算单元。说白了，GPU就是为“简单、重复、大规模并行”的计算而生的。
核心观点： AI训练的本质，就是海量的矩阵运算和卷积操作。这些操作天然具备“数据并行”的特性——同一个操作，作用在成千上万个不同的数据上。这正是GPU的强项。
1.1 GPU架构概览：SM、Warp、Memory Hierarchy
要写好GPU算子，你得先摸清它的“脾气”。我习惯把GPU架构拆成三个层次来看：计算单元、调度单元、存储层次。
1.1.1 SM（Streaming Multiprocessor）—— 真正的“干活的人”
一个GPU由多个SM组成。每个SM内部包含若干CUDA Core（或者叫Shader Core），以及共享内存、寄存器文件、调度器等。你可以把SM想象成一个“小工厂”，里面有好几条生产线（CUDA Core）。
我在项目中遇到过一个问题：某个算子性能死活上不去，后来发现是SM利用率太低。说白了，就是没让SM“吃饱”。每个SM能同时驻留的线程块（Thread Block）数量是有限的，你得尽量让每个SM塞满线程块，才能压榨出它的全部性能。
1.1.2 Warp —— GPU调度的基本单位
这个点特别重要，我建议你记牢。GPU并不是一条一条地执行线程，而是32个线程绑在一起，组成一个Warp。这32个线程执行同一条指令，但处理不同的数据。这就是所谓的SIMT（单指令多线程）模型。
避坑指南： 我曾经因为没注意Warp发散（Warp Divergence），导致性能直接腰斩。如果Warp内的32个线程走了不同的分支（比如if-else），那GPU只能串行执行这些分支，其他线程只能干等着。所以，写算子时尽量让同一个Warp内的线程走相同的代码路径。
1.1.3 Memory Hierarchy —— 存储的“金字塔”
GPU的存储结构，说白了就是一个金字塔。越往上，速度越快，容量越小；越往下，速度越慢，容量越大。
| 存储层级 | 位置 | 延迟（大约） | 容量 | 作用域 |
|---|---|---|---|---|
| 寄存器 | SM内部 | 1 cycle | 几十KB/SM | 单个线程 |
| 共享内存 | SM内部 | ~20 cycles | 几十~上百KB/SM | 线程块内 |
| L1缓存 | SM内部 | ~30 cycles | 与共享内存共用 | SM内 |
| L2缓存 | 芯片全局 | ~200 cycles | 几MB | 全局 |
| 全局内存 | 显存 | ~400-800 cycles | 几GB~几十GB | 全局 |
嗯，这里要注意：全局内存的延迟是寄存器的几百倍。所以，优化算子的核心思路之一，就是尽量减少对全局内存的访问，把数据尽量往“金字塔”的上层搬。
我个人习惯，在写任何算子之前，先画一张数据流图，标清楚哪些数据该放共享内存，哪些该放寄存器。这一步做好了，后面能省很多事。
1.2 AI训练中的计算模式
AI训练里，来来去去就那么几种计算模式。我把它们总结成三类：矩阵乘法、卷积、归一化。你把这三种吃透了，大部分算子的优化思路就通了。
1.2.1 矩阵乘法 —— 算力的“试金石”
矩阵乘法（GEMM）是深度学习最核心的操作。全连接层、注意力机制里的QKV计算，本质上都是矩阵乘法。
为什么矩阵乘法这么重要？因为它计算密度高，容易把GPU的算力喂饱。一个优化的好的矩阵乘法，可以达到理论峰值算力的80%以上。
我举个例子，一个简单的矩阵乘法 C = A * B，其中A是MxK，B是KxN：
// 伪代码：最朴素的实现
for (int i = 0; i < M; i++) {
for (int j = 0; j < N; j++) {
float sum = 0.0f;
for (int k = 0; k < K; k++) {
sum += A[i * K + k] * B[k * N + j];
}
C[i * N + j] = sum;
}
}
这个三重循环，直接跑在GPU上，性能会非常差。为什么？因为对B的访问是“列优先”的，而B在内存中是“行优先”存储的，导致缓存命中率极低。
优化技巧： 我常用的方法是“分块”（Tiling）。把大矩阵切成小块，让每个小块能完全塞进共享内存里。这样，计算时从共享内存读数据，速度比从全局内存读快几十倍。具体做法是：把A和B都切成大小为BLOCK_SIZE的小块，每个线程块负责计算一个小块的结果。
1.2.2 卷积 —— 计算机视觉的“心脏”
卷积操作，说白了就是用一个小的“滤波器”（Kernel）在输入图像上滑动，做点积运算。CNN里的卷积层、池化层，都离不开它。
卷积的计算模式，其实可以转化成矩阵乘法来做。这就是著名的im2col（image to column）方法。把输入图像中每个卷积窗口拉成一个列向量，所有列向量拼成一个矩阵，然后跟滤波器矩阵做乘法。
我在项目中遇到过一个问题：im2col虽然能利用高效的矩阵乘法库，但它会引入大量的内存冗余（因为同一个像素可能出现在多个窗口中）。对于大模型，这种内存开销是难以接受的。所以，现在更流行的是直接实现高效的卷积算子，比如使用Winograd算法或FFT。
1.2.3 归一化 —— 训练的“稳定器”
归一化操作，比如Batch Normalization、Layer Normalization，是让训练更稳定的关键。它们的作用是把中间层的输出拉回到一个均值为0、方差为1的分布上。
归一化的计算模式，核心是统计量的计算：求均值、求方差。这些操作需要跨多个维度做规约（Reduction）。
举个例子，Batch Normalization在训练时，需要对每个通道，计算所有样本、所有空间位置的均值和方差：
// 伪代码：Batch Normalization 前向
// 输入 x: [N, C, H, W]
// 输出 y: [N, C, H, W]
// 第一步：对每个通道c，计算均值
for (int c = 0; c < C; c++) {
float mean = 0.0f;
for (int n = 0; n < N; n++)
for (int h = 0; h < H; h++)
for (int w = 0; w < W; w++)
mean += x[n][c][h][w];
mean /= (N * H * W);
// 第二步：计算方差，类似
// 第三步：归一化
}
这个规约操作，在GPU上实现时，最常用的方法是“树形规约”（Tree Reduction）。每个线程块内，先把数据两两相加，再四四相加，直到得到一个最终结果。我习惯用共享内存来做这个事，因为共享内存的带宽比全局内存高得多。
总结一下： 矩阵乘法、卷积、归一化，这三种计算模式覆盖了AI训练中90%以上的计算量。理解它们的本质——矩阵乘法是“密集计算”，卷积是“滑动窗口+点积”，归一化是“统计规约”——是写出高性能算子的第一步。
好了，这一章的内容就到这里。记住，GPU不是万能的，但用对了地方，它就是你的“性能加速器”。下一章，我们会深入CUDA编程模型，手把手带你写出第一个高性能算子。