显存管理入门：从硬件到实践的平衡艺术
做算子优化，说白了就是在跟显存打交道。我刚开始接触这个领域时，总觉得显存就是个「大仓库」，往里塞数据就行。结果第一次跑大模型训练，直接OOM（显存溢出），那叫一个惨。后来我才明白，不了解显存的脾气，你写的算子再快也没用。
这一章，咱们就聊聊显存管理的底层逻辑。你想想看，GPU里的显存和CPU的内存，虽然都是存数据，但脾气秉性完全不同。搞懂了这些，你才能写出既快又省显存的算子。
显存层级结构：为什么不是所有数据都放HBM？
GPU的显存不是铁板一块。它分好几层，每一层的速度和容量都不一样。我习惯把它比作一个「三级仓库」：
- 全局显存（HBM/DRAM）：容量最大，但速度最慢。比如A100有80GB HBM2e，带宽约2TB/s。所有数据默认都放这儿。
- L2缓存：容量小一些（几十MB），速度快一些。所有SM共享。
- L1缓存/共享内存：每个SM独享，容量极小（几十KB到几百KB），但速度极快。这是算子优化的主战场。
- 寄存器：每个线程独享，最快，但数量有限（每个线程最多255个32位寄存器）。
核心观点：算子优化的本质，就是尽量让数据待在「离计算单元近」的地方。别动不动就去HBM里搬数据，那太慢了。
我在项目中遇到过这样一个场景：一个简单的element-wise加法，如果每次从HBM读两个数，加完再写回去，带宽利用率可能只有10%。但如果用共享内存做数据复用，带宽利用率能冲到80%以上。差距就是这么大。
显存带宽与延迟：为什么说「带宽是命，延迟是病」？
显存带宽和延迟，是两个完全不同的概念。我经常跟团队新人说：
- 带宽：单位时间内能搬多少数据。好比高速公路的车道数。
- 延迟：从发出请求到拿到数据的时间。好比从入口到出口的通行时间。
GPU的显存延迟很高（几百个时钟周期），但带宽也很大。所以优化的核心思路是：用高带宽掩盖高延迟。说白了，就是一次发起大量内存请求，让数据在总线上「排队」传输，而不是等一个数据到了再发下一个。
| 层级 | 典型容量 | 典型带宽 | 典型延迟 |
|---|---|---|---|
| HBM2e (A100) | 80 GB | ~2 TB/s | ~400 cycles |
| L2缓存 | 40 MB | ~4 TB/s | ~200 cycles |
| L1/共享内存 | 192 KB per SM | ~20 TB/s | ~30 cycles |
| 寄存器 | 256 KB per SM | ~100 TB/s | ~1 cycle |
我的经验：写算子时，先算一下「计算强度」（计算量/访存量）。如果计算强度低（比如<1），那就是访存瓶颈，优化重点在数据复用和合并访问。如果计算强度高（比如>10），那就是计算瓶颈，优化重点在指令级并行和流水线。
显存分配与释放：别让「内存泄漏」找上门
显存分配和释放，听起来简单，但坑特别多。我见过太多人用cudaMalloc分配了显存，忘了cudaFree，结果程序跑着跑着就OOM了。
嗯，这里要注意几个关键点：
- cudaMalloc vs cudaMallocManaged：前者分配设备显存，后者分配统一内存（可被CPU/GPU共同访问）。我建议在性能敏感场景用
cudaMalloc，因为cudaMallocManaged有隐式迁移开销。 - cudaFree：释放显存。一定要成对出现。我曾经在循环里忘了释放，结果显存越占越多，最后卡死。
- cudaMemcpy：主机和设备之间的数据拷贝。注意方向：
cudaMemcpyHostToDevice和cudaMemcpyDeviceToHost。
// 一个典型的显存分配与释放流程
float *d_data;
size_t size = 1024 * sizeof(float);
// 分配设备显存
cudaMalloc(&d_data, size);
// 从主机拷贝数据到设备
cudaMemcpy(d_data, h_data, size, cudaMemcpyHostToDevice);
// ... 执行核函数 ...
// 释放显存
cudaFree(d_data);避坑指南：我曾经在调试一个算子时，发现每次调用都会多分配几百MB显存。查了半天，原来是cudaMalloc放在循环里了，每次迭代都分配新显存，但只释放最后一次的。正确的做法是在循环外分配，循环内复用。
显存分配策略：池化与复用
频繁调用cudaMalloc和cudaFree是有代价的。每次分配都要和驱动交互，开销不小。所以工业级的框架（比如PyTorch、TensorFlow）都用了显存池技术。
说白了，就是提前分配一大块显存，然后自己管理里面的「小块」。需要时从池子里取，用完还回去。这样避免了频繁的系统调用。
- 预分配：启动时一次性分配大块显存（比如总显存的90%）。
- 缓存复用：释放的显存块不还给系统，而是放入空闲列表，下次分配时优先复用。
- 碎片整理：长时间运行后，显存可能碎片化。需要定期合并相邻的空闲块。
我建议你在写高性能算子时，也考虑用显存池。别每次都cudaMalloc，那样太慢了。尤其是推理场景，延迟敏感，显存分配的开销不能忽视。
知识体系总览
下面这张图，是我自己总结的显存管理知识体系。你可以把它当作一个「思维导图」，快速定位每个知识点的位置。
这张图把显存管理的三个核心维度串起来了。你写算子时，可以对照着看：当前操作主要跟哪一层打交道？带宽够不够？有没有频繁分配释放？
一个小技巧：用nvidia-smi监控显存使用情况。如果发现显存占用持续增长，大概率是忘了释放。配合cuda-memcheck可以定位具体哪一行代码泄漏了。
好了，显存管理的基础就聊到这儿。记住一句话：显存是有限的，带宽是宝贵的，延迟是要掩盖的。下一章咱们会深入算子优化的具体技巧，到时候这些概念都会用上。
想要更多资料，联系可可