3、AI芯片存储层次：片上存储与片外存储
大家好，我是你们的讲师。今天我们来聊聊AI芯片里一个特别关键、也特别容易让人头疼的话题——存储层次。
说实话，我刚开始做AI芯片那会儿，觉得算法才是核心，存储嘛，不就是放数据的地方？后来被现实狠狠教育了一顿。你想想看，计算单元跑得飞快，数据却堵在门口进不来，这芯片性能再强也是白搭。这就是所谓的“存储墙”问题。
今天，我就把我在项目中踩过的坑、积累的经验，一次性讲清楚。
3.1 为什么存储层次这么重要？
AI计算，说白了就是“数据搬运”的游戏。卷积、矩阵乘，这些操作的本质都是：把数据从某个地方搬过来，算完再搬回去。
搬得快，芯片就快。搬得慢，再牛的算力也白费。
我见过一个团队，设计了一款理论算力很高的芯片，结果实测性能只有理论值的30%。查到最后，问题就出在数据搬运上——存储带宽不够，计算单元一直在空等。嗯，这教训挺深刻的。
核心观点：AI芯片的性能瓶颈，往往不在计算单元，而在存储系统。
3.2 片上存储：SRAM与Scratchpad
片上存储，就是芯片内部集成的存储器。它的特点是：离计算单元近，速度快，但容量小。
3.2.1 SRAM（静态随机存取存储器）
SRAM是片上存储的主力。它不需要像DRAM那样定期刷新，所以访问延迟很低，通常能做到1-2个时钟周期。
我个人习惯把SRAM分成两类：
- Cache（缓存）：硬件自动管理，对软件透明。你只管访问地址，硬件帮你判断命中还是缺失。
- Scratchpad（便签存储器）：软件显式管理，程序员自己控制数据搬入搬出。
在AI芯片里，我更倾向于用Scratchpad。为什么？因为AI计算的数据流是高度可预测的。你提前知道下一层卷积要用哪些权重、哪些特征图，完全可以手动调度，避免Cache的替换开销。
我的经验：在项目里，我们给每个计算核配了256KB的Scratchpad。这个大小是反复权衡过的——太小了装不下一个卷积核，太大了芯片面积受不了。
3.2.2 Scratchpad的典型结构
一个典型的Scratchpad包含：
- 存储体（Bank）：多个独立的存储块，支持并行访问
- 地址译码器：将地址映射到具体的Bank
- 控制逻辑：处理读写请求、仲裁冲突
下面这张图是我手绘的Scratchpad内部结构，你一看就明白了：
3.3 片外存储：HBM与GDDR
片上存储容量有限，装不下整个模型。这时候就需要片外存储来救场。
目前主流的选择有两个：HBM（高带宽存储器）和GDDR（图形双倍数据率存储器）。
3.3.1 HBM vs GDDR：怎么选？
| 特性 | HBM | GDDR |
|---|---|---|
| 带宽 | 极高（1TB/s+） | 高（500GB/s+） |
| 容量 | 中等（8-16GB/堆叠） | 大（16-32GB） |
| 功耗 | 较低（每比特） | 较高 |
| 成本 | 高（2.5D/3D封装） | 低（成熟工艺） |
| 典型应用 | 高端AI训练芯片 | 推理芯片、游戏GPU |
我在项目里做过一个对比：同样做ResNet-50推理，用HBM的芯片带宽利用率能达到85%，而GDDR只有60%左右。原因很简单——HBM的堆叠结构让数据路径更短，访问延迟更低。
注意：HBM虽然好，但封装成本高。如果你的芯片目标市场是边缘设备，GDDR可能是更务实的选择。我曾经有个项目，为了追求极致性能选了HBM，结果成本超了预算，最后不得不砍掉一半的HBM堆叠层数。
3.4 存储带宽计算
带宽，就是单位时间内能搬运多少数据。公式很简单：
带宽 = 总线宽度 × 时钟频率 × 传输次数/周期举个例子：
- 总线宽度：512位（64字节）
- 时钟频率：1GHz
- DDR（双倍数据率）：每个时钟传输2次
那么理论带宽 = 64字节 × 1GHz × 2 = 128GB/s
但实际能用到多少？我告诉你，能到理论值的70%就算不错了。为什么？
- 协议开销：读写命令、地址、校验位都要占用带宽
- 刷新开销：DRAM需要定期刷新，这期间不能访问
- 访问冲突：多个请求同时访问同一个Bank，需要排队
我的建议：做带宽规划时，按理论值的60%来估算。这样留出余量，后期不会手忙脚乱。
3.5 数据复用策略
带宽有限，怎么让有限的数据发挥最大价值？答案是：复用。
AI计算里，数据复用主要有三种模式：
3.5.1 权重复用
同一个卷积核，要跟输入特征图的多个位置做卷积。权重只加载一次，反复使用。
我习惯的做法是：把权重存在Scratchpad里，然后让计算单元按行扫描输入特征图。这样权重只搬一次，输入数据流式读取。
3.5.2 输入特征图复用
多个卷积核共享同一个输入特征图。比如3×3卷积，9个权重值要跟同一个输入像素做乘加。
这时候，我会把输入特征图的一小片区域（比如4×4）加载到Scratchpad里，然后让多个卷积核轮流使用。这能显著减少片外访问次数。
3.5.3 输出特征图复用
多个输入通道累加到同一个输出像素。比如输入有64个通道，每个通道的卷积结果都要累加到同一个输出位置。
嗯，这里要注意：输出累加需要额外的存储空间。我一般会在Scratchpad里划出一块区域专门放部分和（partial sum），等所有通道算完再写回片外。
3.6 实战：一个数据流调度示例
说了这么多，我们来个实际的例子。假设我们要做一次3×3卷积，输入是32×32×64，输出是32×32×128，卷积核是3×3×64×128。
数据量有多大？
- 输入特征图：32×32×64 = 65,536个元素
- 权重：3×3×64×128 = 73,728个元素
- 输出特征图：32×32×128 = 131,072个元素
如果全部放片外，每次计算都要搬数据，带宽根本扛不住。
我的调度策略是这样的：
- 分块（Tiling）：把输入特征图分成4×4的小块，每次只处理一个块
- 权重预加载：把当前块需要的卷积核权重提前搬到Scratchpad
- 流水线执行：计算当前块的同时，预取下一个块的权重
- 输出累加：部分和暂存在Scratchpad，等所有通道算完再写回
这样，片外访问量从原来的几十万次，降到了几千次。效果立竿见影。
核心思路：数据复用就是“一次搬入，多次使用”。搬得越少，算得越快。
3.7 避坑指南
最后，分享几个我踩过的坑：
- 别迷信理论带宽：实际带宽受限于访问模式、协议开销、刷新周期。我见过一个团队，理论带宽800GB/s，实测只有300GB/s，就是因为访问模式太随机，导致Bank冲突严重。
- Scratchpad大小要精打细算：太大浪费面积，太小装不下数据。我一般按“能装下最大卷积核+两行输入特征图”来估算。
- 数据复用不是越多越好：过度复用会导致计算单元闲置（等数据）。要找到计算和存储的平衡点。
好了，这一章的内容就到这里。存储层次是AI芯片设计的基石，理解透了，后面的计算单元设计、数据流调度才能游刃有余。
想要更多资料，联系可可