揭秘内存墙：从存储层次到算子优化的性能瓶颈与拆墙实践

跑GEMM类算子的时候，芯片利用率轻松到70%甚至80%；一旦切到真实模型推理，整体利用率常常掉到20%上下。这个现象干过AI芯片或性能优化的人应该都不陌生。我入行头两年也被这个问题折磨得够呛，总以为是框架调度不行，或者是卡不行，后来把存储层次和内存墙这两件事彻底想明白了，才发现所有看似花哨的优化技巧——从FlashAttention到算子融合再到各种量化——本质都是在跟内存墙“躲猫猫”。这篇就把这个话题摊开讲清楚。

这次要聊的内容，对三类人最有用：一是做AI芯片架构或硬件设计的，二是搞高性能算子开发和模型部署优化的，三是对AI加速底层原理感兴趣、想搞清楚“为什么算力那么高但跑模型还是慢”的读者。理解了存储层次和内存墙，你再看市面上主流加速芯片的设计取舍、论文里各种优化方法的动机，都会通透很多。

1. 先算一笔账：内存墙到底“卡”在哪儿

1.1 算力与带宽的两个量级差

先看一组不算冷门的数字。以NVIDIA H100 SXM为例，FP16稠密矩阵乘法算力大约是990 TFLOPS，HBM3内存带宽是3.35 TB/s。两者一除，得到一个关键比值：约295 FLOP/Byte。意思是，芯片每从外部内存取回1字节数据，理论上只够做295次浮点运算。而一个FP16数据占2字节，换算一下就是每取回一个数，平均只能做大约600次运算。

这600次听着不算少，但真实的AI模型里大量操作根本达不到这个计算密度。比如Softmax、LayerNorm、各种Elementwise操作，它们对每个数据往往只做一两次运算就要读写一遍内存，计算强度可能只有1到5 FLOP/Byte。算力再高，在这种算子面前也只能干等内存把数据喂进来。

业内有个大家常提的观察：算力翻倍的周期已经缩短到一年半到两年，但内存带宽翻倍往往要三到五年甚至更久。这两条曲线的剪刀差越拉越大，就是“内存墙”这个词的由来——算力增长撞在带宽增长的天花板上，跑不快。

1.2 内存墙不只是带宽问题，还有容量、延迟和功耗

很多材料把内存墙简单等同于“带宽不够”，但实际干活的时候，它至少包含四个维度，每个维度在真实场景里的表现完全不一样。

• 带宽墙：单位时间能搬多少字节。这是最常被提到的，典型场景是推理时小batch逐token生成，模型权重全量读一遍但计算量却很小。
• 容量墙：数据能不能一次性放得下。7B模型用FP16存权重是14GB，训练时加上梯度、优化器状态、激活值，经常突破100GB，超出单卡HBM容量就得上多卡或做重计算，都是额外代价。
• 延迟墙：一次访问要等多久。HBM延迟是几百纳秒量级，虽然相比SSD已经是天壤之别，但对于几GHz主频的芯片来说，几百纳秒等于几百上千个时钟周期。如果数据依赖链长且无法预取，计算单元就在空转。
• 功耗墙：搬数据本身非常耗电。一个常被引用的量级对比是：一次浮点运算的能量大概在0.1到1pJ，读一次片上SRAM是几十pJ，而从外部DRAM读一次数据则是几百到上千pJ。换句话说，把数据从内存搬到计算单元所消耗的能量，可能比真正做一百次运算还高。

这四个维度在实际项目里常常同时存在。我做过一个推理引擎优化，单纯提“带宽利用率”还不够，因为模型权重太大放不进SRAM，容量墙先卡死了；换了大显存之后，小batch场景又暴露出延迟问题——比如解码阶段每一步都要串行读KV Cache，延迟直接决定你能跑多快。

1.3 Roofline模型：如何判断一个算子会不会撞墙

判断一个算子到底被算力还是带宽限制，用的工具是Roofline模型。它的核心公式很简单：

实际性能 = min(峰值算力, 带宽 × 计算强度)

其中计算强度 = FLOPs / Bytes，即“每搬运1字节数据，顺便能算多少次运算”。

拿H100来算：当算子的计算强度低于295 FLOP/Byte时，性能上限由带宽决定；高于这个值时，才有机会逼近算力峰值。GEMM之所以是芯片厂商最喜欢展示的场景，正是因为当矩阵足够大时，它的计算强度能达到几千FLOP/Byte，妥妥落在算力限制区。而Softmax这类算子的计算强度基本在个位数，哪怕实现写得再好，也只能贴着带宽天花板走。

Roofline模型最大的价值在于，它让你在动手写代码之前就能判断优化方向：如果是带宽受限，应该优先减少内存搬运量；如果是算力受限，才去抠指令流水和算子切分。很多人一上来就调线程块大小、展开循环，结果瓶颈根本不在计算，纯属白忙活。我在项目里给团队定的规矩是：改代码前先算强度，再对照硬件参数决定策略。

2. 存储层次：一张从寄存器到SSD的“分阶图”

2.1 各级存储的容量-带宽-延迟全景

AI芯片的存储体系，本质上是一层套一层的“缓存漏斗”，从离计算单元最近的寄存器一直延伸到外部存储。用一张表可以把各层的量级看得比较清楚：

这张表的关键信息是：越靠近计算单元，容量越小、带宽越高、延迟越低，但这部分存储的每比特成本也越贵。所以芯片设计的目标不是“把所有数据都放在最快的地方”，而是“让绝大多数访问尽量命中在靠近计算单元的地方”。

2.2 为什么片上SRAM是AI芯片最贵的“厨房台面”

有生活经验的读者可以这样类比：寄存器相当于厨师的双手，SRAM相当于灶台边的操作台，HBM相当于厨房里的冰箱，外部存储则是小区门口的超市。备菜时如果每用一棵葱都要跑一趟小区超市，再好的厨师也快不起来。做AI加速也一样，计算单元能跑多快，很大程度上取决于“操作台”够不够大、东西摆得够不够顺手。

但SRAM为什么不能无限做大？原因有三个：第一是密度低，同样面积下SRAM能存的数据远少于DRAM，做32MB SRAM的面积可以做几百MB的HBM逻辑；第二是漏电功耗，SRAM不用也持续耗电，容量越大静态功耗越难压；第三是良率，片上大块SRAM对工艺缺陷极其敏感，太大的SRAM会让芯片良率直接崩掉。所以你会看到，哪怕是H100这种顶级产品，片上L2缓存也只有50MB，而HBM可以做80GB——两者差了三个数量级。

2.3 Scratchpad与Cache：AI芯片为什么几乎都选前者

CPU里用的是一套复杂的硬件Cache层次，有同位、组相联、一致性协议。但用这套思路做AI芯片，弊端很明显：Cache的命中率靠硬件猜，而AI负载的访问模式是可以静态分析的，没必要让硬件浪费晶体管去猜。因此主流AI芯片几乎都选择了Scratchpad（软件管理片上存储）的方案。

GPU里叫Shared Memory，TPU里叫Unified Buffer，许多NPU里叫Local Memory，名字不同，本质相同：这块SRAM由软件或编译器显式安排数据何时写入、何时读出、何时淘汰。好处是确定性高、面积省、功耗低；代价是“数据搬运”这活儿得人或者编译器来干，谁干得不好，性能就一目了然。这也是为什么AI芯片的性能优化特别依赖编译器、算子库和手写kernel的三层配合。

这里有个实测感受：硬件上Scratchpad做成了，软件安排不好，性能照样拉胯。我见过一个项目里DMA引擎的利用率只有30%，原因就是搬运任务和计算任务没有重叠，芯片大部分时间在干等数据。存储层次不只是硬件设计问题，更是一整套软硬件协同的调度问题。

3. 两类AI工作负载的“访存画像”：GEMM和它的邻居们

3.1 算力密集与访存密集的划分

所有AI算子都可以按计算强度划成两类。算力密集型算子，比如大尺寸GEMM、大卷积核卷积，它们对同一批数据做大量重复运算，数据搬运量相对小，瓶颈在计算单元。访存密集型算子则相反，比如Softmax、LayerNorm、Dropout、残差加、Embedding，它们对每个数据几乎只做少量运算，瓶颈在内存带宽。

真正麻烦的是，一个真实模型是这两类算子串在一起跑的。哪怕GEMM本身算得很爽，周围的访存密集算子一旦拖后腿，整体性能就被拉下来了。有公开分析统计过，在BERT这类Transformer训练任务里，非GEMM算子能占到三四成甚至更高的时间比例。这也是为什么仅仅优化GEMM核函数远远不够。

3.2 训练场景：容量墙和梯度通信是双重的“墙”

训练比推理更容易撞墙，而且是不同类型的墙一起撞。以7B模型、FP16混合精度训练为例，粗略算一笔账：权重FP16要14GB，梯度FP16要14GB，Adam优化器为每个参数维护一份FP32的一阶动量和二阶动量，共56GB，master weight FP32还要28GB。光是这几项就已经超过110GB，再加激活值显存，单卡80GB的H100根本放不下。

即便你上了多卡，数据并行训练每次迭代都要做全量梯度同步，AllReduce通信的数据量和模型参数量是同一个量级。也就是说，每训练一步，整个模型的体量都要在卡间搬一遍。这里的“墙”已经不只是芯片到内存之间的墙，而是延伸到了服务器互联网络上。

针对训练场景的存储层次优化，常见的思路有：激活重计算（不存中间结果，反向时重新算）、混合精度通信（梯度压缩）、ZeRO系列优化器状态切分等。它们的本质仍然是同一个逻辑——减少跨层搬移的字节数，把数据尽量留在离计算更近的地方。

3.3 推理场景：小batch和长上下文都是带宽杀手

推理阶段的访存画像和训练又不一样，它最典型的特征是小batch下的“权重饥饿”。

举个例子：一个7B模型，FP16权重14GB。如果推理时batch size等于1，每生成一个token，理论上只需要对权重做一轮极稀疏的计算，但权重这14GB几乎必须从HBM里过一遍。假如单卡HBM带宽是3TB/s，光把权重读一遍就得4.7毫秒；要是QPS要求10，光搬权重就已经吃掉一半多带宽。这也解释了为什么LLM推理都在拼命做量化：4bit权重把同样的模型从14GB压到3.5GB，同样带宽下能服务的用户直接翻几倍。

长上下文推理还有一个容易被忽略的点：KV Cache。上下文越长，KV Cache越大，每生成一个新token就要把所有历史的KV读一遍。上下文从2K涨到32K，KV Cache读取量增长十几倍，连解码阶段都从“算力密集”变成了“访存密集”。我之前调过一版长文本推理，用了KV Cache的量化加分页管理，才勉强把预填充和解码的性能拉回正常水平。

3.4 快速定位访存瓶颈的实操方法

想确认某个算子到底卡在哪面墙，我的习惯是先做两个动作。第一，用手算一遍计算强度，和硬件参数对照，心里先有谱；第二，用profiler看实际运行时的行为数据。GPU上NVIDIA Nsight Compute能直接给出Memory Throughput和Compute Throughput；如果是自研NPU，厂商通常也有类似的内存吞吐计数器。

判断逻辑很简单：如果Memory Throughput已经接近90%以上，而Compute Throughput只有百分之二三十，就是带宽墙，别再优化计算了，去减搬运量。反过来则去优化计算效率。这一步看着基础，但真能坚持做的人不多，很多团队优化半天发现方向错了，根源就是漏了最开始的这步判断。

4. 硬件的“拆墙法”：HBM、片上SRAM与近存思路

4.1 HBM的高带宽秘密：堆叠、TSV和超宽位宽

HBM和普通DDR、GDDR最大的区别，在于它把容量和带宽同步拉高了。普通DDR或GDDR颗粒的位宽一般是32bit或64bit，频率再高，单颗粒带宽也就几十GB/s。HBM通过硅通孔把多层DRAM裸片叠在一起，每层之间用TSV垂直互连，底部再通过硅中介层和计算芯片连在一起，单个HBM堆叠的位宽能做到1024bit，再加上比较高的I/O速率，单位带宽一下子就上去了。

以HBM3为例，单个堆叠的带宽大致能到几百GB/s甚至更高，H100用了多个HBM堆叠，才凑出3.35TB/s的总带宽。这种方案之所以没有被无限扩展，一是钱的问题，2.5D封装和TSV工艺成本远高于普通DDR；二是良率和散热问题，堆叠层数越高，制造和散热压力越大。所以HBM从诞生起就定位在高性能计算和数据中心市场，民用设备基本用不起。

4.2 片上SRAM：AI芯片真正的胜负手

别看HBM带宽已经到TB/s级别，跟片上SRAM一比还是小巫见大巫。H100的L2缓存有50MB，这个容量远不能把大模型权重装下，但这50MB的L2带宽业内评估在10TB/s量级以上，是HBM的好几倍。再加上每个SM内部的Shared Memory，整颗芯片的片上聚合带宽非常可观。

AI芯片的优化本质上就是一件事：把数据尽量在片上有层次地“玩”起来，减少对HBM的依赖。一个典型的GEMM kernel，要是完全不做分块，每个输出元素都要去全局内存取操作数，带宽立刻被打爆；做了分块后，数据块被加载到Shared Memory或寄存器里反复复用，对全局内存的访问次数能减少几个数量级。

这也是为什么各家AI芯片在SRAM容量上卷得很凶：A100是40MB L2，H100做到50MB，许多专用NPU更是在每个计算单元旁边堆了大块Local Memory。因为芯片的实际算力能发挥出多少，很大程度上不是看峰值FLOPs，而是看片上的“数据周转率”够不够。

4.3 近存计算和存内计算：物理拆墙的尝试

既然搬数据的代价太大，那干脆把计算搬到存储旁边，甚至搬进存储里面。近存计算（Processing-Near-Memory）把计算逻辑放在存储颗粒附近，典型代表是三星的HBM-PIM和SK海力士的AiM，它们把计算单元以“存储体”为单位嵌入到HBM堆叠中，让一部分计算直接在内存侧完成，省掉数据搬出搬进的开销。存内计算（Computing-in-Memory）更进一步，直接用存储单元的物理特性完成乘加运算，比如用电阻变化器件的电导代表权重，输入电压后直接得到模拟域的部分和。

这两条路理想很美好，但商用进展都比较缓慢。近存计算受限于芯片面积和散热，能放进去的计算功能有限；存内计算则要面对模拟计算精度低、ADC/DAC转换开销大、器件一致性难保证等一堆问题。我的看法是：在2020年代的中期，想靠存算一体在主流场景里替代传统架构还不现实，但它给“功耗墙”指了一个远期方向——毕竟把搬家耗电直接降为零的诱惑实在太大了。

5. 软件才是真正的主力“拆墙者”：局部性、算子融合与流水

5.1 分块（Tiling）：让数据在SRAM里多活一会儿

硬件把存储层次搭好后，真正决定性能的是软件怎么利用。分块（Tiling）是其中最基本也最有效的技术，思路一句话：不要一次处理整个矩阵，切成小块让数据在片上反复用。

拿一个朴素GEMM举例。假设计算C=A×B，矩阵都是4096×4096，每个数是FP16。光是B矩阵就有32MB，远超片上SRAM，所以必须分块。每取出一块128×128的子块，数据量是32KB，加载到Shared Memory后，这块数据会被反复读取用于计算多个输出元素。分块越小，片上装载压力越小，但复用率也越低；分块越大，复用率越高，但SRAM放不下，还会出现bank冲突。实际工程里块大小通常要做一轮参数扫描，结合具体矩阵形状和硬件SRAM容量来定。

CUDA里一个简化版的分块GEMM长这样：

__global__ void gemm_tiled(const float* A, const float* B, float* C, int N) { __shared__ float As[TILE][TILE]; __shared__ float Bs[TILE][TILE]; int row = blockIdx.y * TILE + threadIdx.y; int col = blockIdx.x * TILE + threadIdx.x; float acc = 0.0f; for (int k = 0; k < N; k += TILE) { As[threadIdx.y][threadIdx.x] = A[row * N + k + threadIdx.x]; Bs[threadIdx.y][threadIdx.x] = B[(k + threadIdx.y) * N + col]; __syncthreads(); for (int i = 0; i < TILE; i++) for (int j = 0; j < TILE; j++) acc += As[threadIdx.y][i] * Bs[i][threadIdx.x]; __syncthreads(); } C[row * N + col] = acc; }

这段代码只是示意，真实库里的GEMM比这复杂得多，但核心逻辑就是：先从全局内存加载一块到Shared Memory，再在片上做密集计算，算完再加载下一块。目的只有一个——减少对全局内存的访问次数。

5.2 算子融合：中间结果不落地，直接消灭一轮“大搬家”

分块解决的是单个算子内部的复用，而算子融合解决的是算子之间的搬运。一个模型里经常出现“LayerNorm -> 残差加 -> Dropout -> 下一个Block”这种组合，如果不做融合，每个算子的输出都要先写回HBM，下一个算子再把它读出来。中间张量动辄几十MB甚至更大，一来一回就是巨额的带宽消耗。

融合的核心思路是：在同一个kernel里把多个算子的计算流程串起来，中间结果留在寄存器或Shared Memory里，不写回全局内存。现在框架里的vertical fusion、pointwise fusion都是这个套路。比如LayerNorm加残差加Dropout，完全可以合成一个kernel，读一次输入，依次完成归一化、残差累加、随机丢弃，一次写回。

最著名的例子是FlashAttention。传统Attention实现里，注意力分数矩阵是N×N大小，计算过程需要把QK^T结果写回HBM，再做Softmax，再读出来和V乘，中间张量极大。FlashAttention通过分块计算和在线Softmax，把整个Attention过程压缩成一次从HBM读Q/K/V、一次写回输出的流程，中间那N×N矩阵永远不会在HBM里落地。对长序列来说，内存访问量从O(N²)被压到O(N)，这个收益在多轮迭代里能直接反映成好几倍的端到端提速。

5.3 双缓冲与流水线：把搬运时间藏进计算时间

减少搬运量是一方面，把必须在搬的那部分时间“藏”起来则是另一方面。如果CPU或专用处理器只是“搬一块、算一块、等一块”，那么搬运期间计算单元就在睡觉。双缓冲的解法很简单：准备两块缓冲区，当前tile在计算时，下一块tile的数据已经在后台搬运了。

在GPU上，Ampere架构之后引入的cp.async指令就是干这事的。它允许数据从全局内存异步拷贝到Shared Memory，不需要等计算完当前块再发起搬运。一个双缓冲循环的伪代码如下：

for (int tile = 0; tile < num_tiles; tile++) { if (tile + 1 < num_tiles) cp_async(next_buffer, global_ptr + (tile + 1) * tile_size); compute(current_buffer); // 当前tile计算 cp_async_wait_group(1); // 确保下一块数据就绪 swap(current_buffer, next_buffer); }

算法层面还有更细的流水线切分，把加载、计算、写回拆成不同阶段，让整条链路像工厂流水线一样运转起来。我做算子优化时，碰到带宽利用率提不上去的情况，第一个检查的就是有没有用双缓冲，以及拷贝和计算的切分粒度够不够细。很多时候问题不在计算代码，而在流水没铺满，计算单元一直在等。

5.4 代码级优化中容易踩的几个存储坑

实操里有一些坑特别隐蔽，写出来给后来者提个醒。

• 读写混合会把有效带宽打下来。HBM的读带宽和写带宽不是完全对称的，读写切换频繁时会损失不少有效带宽。如果算法需要频繁的小块写回，尽量攒一批再写，别一个个写。
• bank冲突会让Shared Memory带宽暴跌。GPU的Shared Memory按bank组织，一次能同时服务多个不同bank的访问；但如果同一warp里的线程恰好访问同一个bank的不同地址，就会被串行化。调整数据布局和访问方式（比如加padding）是常见解法。
• cache line对齐问题。很多硬件按128字节的粒度搬运数据，如果你的数据不对齐，一个搬运请求只装了很少的有效数据，带宽利用率就上不去。数组填充、内存对齐这些细节在性能悬殊时能差出30%以上。
• 不要忽略地址空间和profiler的读值误差。有些profiler给出的带宽数据是平均值，可能掩盖峰值抖动；建议同时看延迟分布和执行单元的空闲率，多角度交叉验证。

这些坑在文档里基本找不到，都是实际操作中踩出来的。我在好几个项目里都是先怀疑算法，最后发现是bank冲突或者对齐问题，修完之后性能直接翻倍。所以调优别急着推倒重来，先查存储访问的微观行为。

6. GPU、TPU、NPU和存算一体：不同流派怎么“拆墙”

6.1 NVIDIA：大L2、Tensor Core和软件生态三线并举

NVIDIA的GPU是通用AI加速的标杆，它的拆墙策略是三管齐下。第一，硬件上依靠大L2缓存（A100是40MB，H100是50MB）和每SM的Shared Memory承接片上复用；第二，用Tensor Core这种专用矩阵单元把单位功耗的算力顶上去，同时H100的Transformer Engine支持FP8，精度经过缩放处理，把同带宽下能搬运的数据量直接翻倍；第三，也是最厉害的，用cuBLAS、cuDNN、CUTLASS和FlashAttention这些库把优化经验固化下来，普通用户不需要自己处理存储层次，调用库函数就基本能拿到不错的性能。

对我个人而言，研究GPU上那些库的源码本身就是很好的学习素材。CUTLASS把分块、流水、寄存器级复用全部做到了极致，看一遍等于把存储层次的优化方法系统过了一遍。

6.2 Google TPU：脉动阵列的“数据用到底”哲学

TPU选择了和GPU很不一样的路线。它的核心计算单元是脉动阵列，数据像水一样在相邻的PE之间流动，每个PE算完一个数后把结果直接传给下一个PE，中间几乎不用反复去SRAM里取数。这种设计让数据在片上被尽可能多地复用，铝合权重和输入都尽量“滑”过整个阵列，而不是每算一步都回存储取数。

TPU的SRAM是统一Buffer，由编译器（XLA）显式安排数据流和调度方案。好处是省掉了硬件Cache的复杂逻辑，面积和功耗都有优势；代价是编译器能否把各类负载都安排得明明白白成了性能天花板。对形状规整的矩阵运算，TPU效率惊人；但遇到动态形状、稀疏访问或者不规则的算子，脉动阵列就有劲使不上。这也是为什么TPU比较“挑活”。

6.3 专用NPU与边缘部署：直接把模型塞进SRAM

专用NPU在拆墙策略上更激进，因为它可以为了特定场景牺牲通用性。很多边缘端NPU的典型做法是：把片上SRAM做大一点，把整个小模型——比如几MB到几十MB——全部提前装进SRAM，推理过程中几乎不碰外部DRAM。这样功耗可以压得非常低，几瓦甚至几百毫瓦就能跑完一个图像分类或语音识别模型。

这是存储层次设计里最直接的“容量换带宽”的思路：既然外部带宽吃紧，就把所有数据都搬到离计算最近的地方。数据中心里的NPU（比如昇腾系列）则采取另一种做法，把片上SRAM分成多级，比如L0/L1/L2，配合专门的DMA引擎，由编译器按数据流模式把数据在各层之间搬移。这种软硬协同的做法，针对性极强，芯片面积和功耗也更可控，代价是编译器复杂度非常高。

6.4 存内计算：在物理层面把墙“拆掉”还有多远

最后说下存内计算（CIM）。它的思路和前面所有方案都不同，前面是“让数据离计算更近”，CIM是“让计算发生在数据所在的地方”。用存储单元阵列做乘加，权重以电导或电荷的形式物理存储在单元中，输入信号一进来，物理定律直接帮你完成乘累加。

这个方向最诱人的地方是它能从根本上消除存储和计算之间的数据搬运——理论上能把能耗降低一到两个数量级。但它面临的工程化挑战非常大：模拟计算的精度受工艺误差影响、ADC/DAC的转换开销和延迟、阵列规模和一致性的权衡、写寿命问题。目前更多落地在低精度、低功耗的推理场景，比如一些语音唤醒和简单传感器应用。距离成为AI训练的主流方案，还有很长的路要走，但它是“内存墙”终极难题下少数值得长期跟踪的方向之一。

我这几年调内核下来最大的体会就是：动手优化之前，一定要先把算子的计算强度算清楚，再结合硬件存储层次决定方向。很多人一上来就抠指令级并行，结果发现卡的根本是数据没搬到位。存储层次和内存墙不是芯片架构师专属的话题，只要你在和AI性能打交道，这两张图就是你的底层地图。先看数据怎么流，再看计算怎么算，性能优化才能少走弯路。