4. 显存管理基础
各位同学，今天我们聊聊显存管理。说实话，很多做算子开发的同学，一开始都只关注计算逻辑，觉得显存嘛，不就是个存数据的地方？我刚开始也这么想，直到有一次在NPU上调一个超大模型，显存直接爆了，排查了整整两天才发现是分配策略出了问题。从那以后，我养成了一个习惯：写算子之前，先把显存结构摸清楚。
NPU的显存，不是一块简单的存储空间。它是有层级的，就像我们电脑的缓存一样。你想想看，如果所有数据都堆在一个地方，那访问速度肯定跟不上计算速度。所以，NPU设计了一套层级结构，来平衡容量和速度。
4.1 显存层级结构：HBM、SRAM、寄存器
NPU的显存层级，从上到下大致可以分为三层：HBM（高带宽内存）、SRAM（静态随机存取存储器）和寄存器。每一层都有自己的特点，说白了就是：越靠近计算单元，速度越快，但容量越小；越远离，容量越大，但速度越慢。
| 层级 | 典型容量 | 访问延迟 | 带宽 | 主要用途 |
|---|---|---|---|---|
| 寄存器 | 几百字节 | ~1 个时钟周期 | 极高 | 存储中间计算结果、循环变量 |
| SRAM（L1/L2） | 几MB到几十MB | ~10 个时钟周期 | 高 | 存放Tile数据、权重、激活值 |
| HBM（全局内存） | 几十GB | ~几百个时钟周期 | 高（但延迟大） | 存储模型参数、大数据块 |
嗯，这里要注意，不同厂商的NPU，SRAM的大小和层级划分可能不一样。比如华为的昇腾芯片，它的SRAM就分成了L1和L2两级。我个人习惯，在写算子时，会优先把最常用的数据放到寄存器里，其次是SRAM，最后才是HBM。为什么？因为寄存器访问几乎不花时间，而HBM的延迟，你想想看，几百个时钟周期，足够计算单元干很多事了。
核心原则：数据离计算单元越近，计算效率越高。算子优化的本质，就是尽可能把数据往上层搬。
我在项目中遇到过这样一个场景：一个卷积算子，原始实现是把所有输入特征图都放在HBM里，每次计算都去HBM读。结果性能瓶颈不在计算，而在访存。后来我把输入特征图分块（Tiling），每块先搬到SRAM里，计算完再写回HBM。就这么一个改动，性能提升了3倍多。说白了，这就是利用SRAM的带宽优势，来掩盖HBM的高延迟。
4.2 显存带宽与延迟：为什么它们如此重要？
带宽和延迟，是衡量显存性能的两个核心指标。很多同学容易混淆，我简单解释一下：
- 带宽：单位时间内能传输的数据量，好比高速公路的车道数。车道越多，同时能跑的车就越多。
- 延迟：发起一次访问请求到数据到达的时间，好比从入口到出口的通行时间。即使车道再多，如果路很长，第一辆车到达的时间还是固定的。
在NPU上，HBM的带宽通常很高（比如1TB/s以上），但延迟也很大（几百纳秒）。SRAM的带宽虽然不如HBM，但延迟极低（几纳秒）。所以，我们做算子优化时，要尽量做到：
- 减少对HBM的随机访问：随机访问会浪费带宽，因为每次访问都要经历完整的延迟。
- 利用SRAM做数据复用：把数据从HBM搬到SRAM后，尽量多次使用，避免反复搬运。
- 使用DMA（直接内存访问）：让DMA负责数据搬运，计算单元专心做计算，实现“计算-访存”重叠。
避坑指南：我曾经在写一个矩阵乘算子时，为了省事，直接把所有数据都放在HBM里，然后让计算单元一条一条指令去读。结果性能惨不忍睹。后来我改用DMA把数据分块搬到SRAM，再让计算单元从SRAM读，性能直接翻倍。记住：不要让计算单元直接访问HBM，除非万不得已。
为什么会这样？因为计算单元访问HBM时，每次都要等待几百个时钟周期。而DMA搬运数据时，计算单元可以同时做其他计算。这就是所谓的“流水线”思想。你想想看，如果计算单元每次都要等数据到了才能干活，那效率能高吗？
4.3 显存分配策略：静态分配 vs 动态分配
显存分配，说白了就是决定什么时候、在哪里、分配多大空间。常见的策略有两种：静态分配和动态分配。
4.3.1 静态分配
静态分配，就是在编译时或算子初始化时，一次性把需要的显存都申请好。优点是简单、高效，没有运行时开销。缺点是灵活性差，如果模型输入尺寸变化，可能需要重新分配。
// 伪代码示例：静态分配
void init_operator() {
// 在初始化时，一次性分配所有需要的显存
hbm_buffer = alloc_hbm(1024 * 1024 * 4); // 4MB
sram_buffer = alloc_sram(256 * 1024); // 256KB
// 后续计算中，直接使用这些buffer，不再申请
}我个人习惯，对于固定尺寸的模型，比如输入是224x224的图像分类网络，我会优先使用静态分配。因为这样没有额外的分配开销，性能更稳定。
4.3.2 动态分配
动态分配，就是在运行时根据需要临时申请显存。优点是灵活，可以适应不同的输入尺寸。缺点是每次分配都有开销，而且容易产生内存碎片。
// 伪代码示例：动态分配
void compute(int input_size) {
// 根据输入尺寸，动态计算需要的显存大小
int buffer_size = input_size * sizeof(float);
void* buffer = alloc_hbm(buffer_size);
// ... 使用buffer ...
free_hbm(buffer); // 使用完后立即释放
}警告：动态分配一定要小心内存泄漏。我曾经在一个循环里忘记释放显存，结果跑了几个小时后，显存被吃光，程序直接崩溃。排查了好久才发现是这个问题。所以，我建议：谁分配，谁释放；用完即释放。
4.3.3 混合策略
在实际项目中，我通常采用混合策略：对于生命周期长、尺寸固定的数据（如模型权重），使用静态分配；对于生命周期短、尺寸变化的数据（如中间激活值），使用动态分配。这样既保证了性能，又保留了灵活性。
嗯，这里还要提一个技巧：显存池（Memory Pool）。你可以预先申请一大块显存，然后自己管理内部的分配和释放。这样可以避免频繁调用系统API，减少开销。很多NPU框架（比如华为的CANN）内部就是这么做的。
总结一下：显存管理不是简单的“申请-使用-释放”，而是一门平衡艺术。你要在速度、容量、灵活性之间找到最佳点。多想想数据在哪里、怎么搬、什么时候搬，你的算子性能就能上一个台阶。
这张图展示了NPU显存的层级结构和数据流向。你可以看到，数据从HBM出发，经过SRAM，最终到达寄存器，然后被计算单元使用。DMA引擎负责异步搬运，让计算和访存可以重叠。嗯，这就是我们做显存优化的核心思路。