GPU内存层次与融合：全局内存、共享内存、寄存器、内存带宽瓶颈
聊算子融合，绕不开GPU的内存层次。说白了，你写的每个算子都在跟内存打交道。我见过太多团队，算子融合做得花里胡哨，结果性能没提多少——为什么？因为根本没搞明白数据在GPU里是怎么流动的。
今天咱们就掰开揉碎，把GPU的内存层次讲透。你想想看，一个典型的深度学习训练任务，数据从全局内存读到寄存器，中间要经过多少道关卡？每一道关卡都有瓶颈。融合的本质，就是减少这些关卡之间的来回折腾。
GPU内存层次概览
先画个图，把整体结构看清楚。我习惯用这张图来跟团队新人解释GPU内存的“阶级差异”。
这张图我每次培训必讲。你看，从全局内存到寄存器，每往上走一层，带宽就上一个台阶，延迟就降一个数量级。但代价是什么？容量越来越小。
核心矛盾：全局内存容量大但慢，寄存器快但容量极小。算子融合要做的，就是让数据尽量在“快”的层次里多待一会儿，少去“慢”的层次里折腾。
全局内存：最大的瓶颈
全局内存是GPU里最大的存储区域，所有数据最初都在这。但它的带宽，说实话，跟计算单元的需求比起来，差得远。
我举个例子。一个A100 GPU，全局内存带宽大概是2 TB/s。听起来很吓人对吧？但它的计算能力呢？FP16算力接近312 TFLOPS。算一笔账：
- 每个FP16操作需要读取2字节数据
- 312 TFLOPS × 2字节 = 624 TB/s 的数据需求
- 实际带宽只有2 TB/s，差了300多倍
这就是所谓的“内存墙”。说白了，计算单元大部分时间都在等数据，而不是真正在算。
我的经验：在优化一个BERT训练任务时，我发现GPU利用率只有30%出头。用NVIDIA Nsight一查，大部分时间花在了全局内存访问上。后来做了算子融合，把多个小算子的中间结果留在寄存器里，利用率直接飙到75%。
共享内存：程序员可控的缓存
共享内存是GPU里最灵活的一层。它不像L1/L2缓存那样完全由硬件控制，你可以手动管理数据在共享内存里的生命周期。
共享内存的带宽大约是全局内存的10倍，延迟只有后者的1/20。但容量有限——每个SM只有48KB到164KB（取决于架构）。
怎么用？我习惯把共享内存当作“数据中转站”。比如做矩阵乘法时，把子矩阵先搬到共享内存里，让同一个block里的线程反复使用，而不是每次计算都去全局内存里取。
// 一个典型的共享内存使用模式
__global__ void matmul_shared(float* A, float* B, float* C, int N) {
__shared__ float As[TILE_SIZE][TILE_SIZE];
__shared__ float Bs[TILE_SIZE][TILE_SIZE];
int bx = blockIdx.x, by = blockIdx.y;
int tx = threadIdx.x, ty = threadIdx.y;
float sum = 0.0f;
for (int t = 0; t < N / TILE_SIZE; t++) {
// 从全局内存加载到共享内存
As[ty][tx] = A[by * TILE_SIZE + ty][t * TILE_SIZE + tx];
Bs[ty][tx] = B[t * TILE_SIZE + ty][bx * TILE_SIZE + tx];
__syncthreads();
// 从共享内存计算
for (int k = 0; k < TILE_SIZE; k++)
sum += As[ty][k] * Bs[k][tx];
__syncthreads();
}
C[by * TILE_SIZE + ty][bx * TILE_SIZE + tx] = sum;
}注意：共享内存用不好反而会拖慢性能。我曾经犯过一个错误——把太大的数据塞进共享内存，导致每个SM能同时运行的block数量减少，最终吞吐反而下降了。记住，共享内存是“以空间换时间”，但空间本身也是稀缺资源。
寄存器：最快的存储，最稀缺的资源
寄存器是GPU里最快的存储，延迟只有1个时钟周期。但每个线程能用的寄存器数量有限——通常是32到255个（取决于架构和编译器设置）。
算子融合的一个关键技巧，就是尽量把中间结果留在寄存器里，而不是写回全局内存。比如把两个element-wise操作融合成一个kernel：
// 不融合：两个kernel，两次全局内存读写
// kernel1: y = relu(x)
// kernel2: z = y * 2 + 1
// 融合：一个kernel，中间结果在寄存器
__global__ void fused_relu_scale(float* x, float* z, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
float y = x[idx]; // 从全局内存读一次
y = y > 0 ? y : 0; // relu，结果在寄存器
z[idx] = y * 2.0f + 1.0f; // 写回全局内存一次
}
}你看，原来要两次全局内存读写，现在变成一次读一次写。对于大张量来说，这个优化能省下几百微秒。
内存带宽瓶颈：融合为什么能破局
现在我们把所有东西串起来。内存带宽瓶颈的本质是什么？是数据搬运的速度跟不上计算的速度。
算子融合解决这个问题的思路很简单：减少数据搬运的次数。
| 场景 | 不融合 | 融合后 | 带宽节省 |
|---|---|---|---|
| 两个element-wise算子 | 2次读 + 2次写 | 1次读 + 1次写 | 50% |
| Conv + ReLU + BN | 3次读 + 3次写 | 1次读 + 1次写 | 67% |
| MatMul + Add + Gelu | 3次读 + 3次写 | 1次读 + 1次写 | 67% |
数据不会骗人。融合后，全局内存的访问次数大幅减少，带宽压力自然就小了。
核心原则：每次从全局内存读数据，都要尽可能多地利用它。一个数据读进来，做完所有能做的计算，再写回去。这就是融合的哲学。
实战：如何判断一个融合是否有效
我总结了一个简单的判断标准：计算密度。计算密度 = 计算量 / 访存量。单位是FLOPs/Byte。
- 计算密度 < 1：访存密集型，融合效果显著
- 计算密度 1-10：混合型，融合有一定效果
- 计算密度 > 10：计算密集型，融合效果有限
举个例子。两个element-wise算子，每个算子计算量是O(n)，访存量也是O(n)。计算密度约等于1（每个元素做一次乘加）。这种场景下，融合能省掉一半的访存，效果立竿见影。
但如果是两个矩阵乘法融合呢？每个矩阵乘的计算量是O(n³)，访存量是O(n²)。计算密度是O(n)。对于大矩阵，计算密度很高，融合带来的访存节省相对有限。这时候更应该关注的是如何利用共享内存和寄存器来优化单个矩阵乘。
我的建议：拿到一个模型，先用profiler跑一遍，看看哪些算子的计算密度低。这些就是融合的“潜力股”。我一般会优先处理计算密度低于5的算子序列。
嗯，说到这，你应该对GPU内存层次和融合的关系有了清晰的认识。记住一句话：融合的本质，是让数据在快层次里多待一会儿，在慢层次里少跑几趟。下一节我们会聊具体的融合策略和代码实现，到时候我会拿几个真实案例来拆解。
本章要点回顾：
- GPU内存层次：全局内存（慢、大）→ 共享内存（中、中）→ 寄存器（快、小）
- 内存带宽瓶颈：计算需求远超内存供给，数据搬运成为瓶颈
- 融合破局：减少全局内存访问次数，让数据在寄存器/共享内存中完成更多计算
- 判断标准：计算密度越低，融合收益越大
想要更多资料，联系可可