跑GEMM类算子的时候芯片利用率轻松到70%甚至80%一旦切到真实模型推理整体利用率常常掉到20%上下。这个现象干过AI芯片或性能优化的人应该都不陌生。我入行头两年也被这个问题折磨得够呛总以为是框架调度不行或者是卡不行后来把存储层次和内存墙这两件事彻底想明白了才发现所有看似花哨的优化技巧——从FlashAttention到算子融合再到各种量化——本质都是在跟内存墙“躲猫猫”。这篇就把这个话题摊开讲清楚。这次要聊的内容对三类人最有用一是做AI芯片架构或硬件设计的二是搞高性能算子开发和模型部署优化的三是对AI加速底层原理感兴趣、想搞清楚“为什么算力那么高但跑模型还是慢”的读者。理解了存储层次和内存墙你再看市面上主流加速芯片的设计取舍、论文里各种优化方法的动机都会通透很多。1. 先算一笔账内存墙到底“卡”在哪儿1.1 算力与带宽的两个量级差先看一组不算冷门的数字。以NVIDIA H100 SXM为例FP16稠密矩阵乘法算力大约是990 TFLOPSHBM3内存带宽是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芯片的存储体系本质上是一层套一层的“缓存漏斗”从离计算单元最近的寄存器一直延伸到外部存储。用一张表可以把各层的量级看得比较清楚层级容量量级带宽量级延迟量级相对功耗量级寄存器堆KB级极高多读写口1ns基准片上SRAM几十KB到几十MB数十TB/s级几纳秒到几十纳秒数十倍于寄存器HBM高带宽内存几十GB到上百GB3~8TB/s级数百纳秒级远高于SRAM外部存储NVMe/SSDTB级GB/s级毫秒级更高这张表的关键信息是越靠近计算单元容量越小、带宽越高、延迟越低但这部分存储的每比特成本也越贵。所以芯片设计的目标不是“把所有数据都放在最快的地方”而是“让绝大多数访问尽量命中在靠近计算单元的地方”。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与CacheAI芯片为什么几乎都选前者CPU里用的是一套复杂的硬件Cache层次有同位、组相联、一致性协议。但用这套思路做AI芯片弊端很明显Cache的命中率靠硬件猜而AI负载的访问模式是可以静态分析的没必要让硬件浪费晶体管去猜。因此主流AI芯片几乎都选择了Scratchpad软件管理片上存储的方案。GPU里叫Shared MemoryTPU里叫Unified Buffer许多NPU里叫Local Memory名字不同本质相同这块SRAM由软件或编译器显式安排数据何时写入、何时读出、何时淘汰。好处是确定性高、面积省、功耗低代价是“数据搬运”这活儿得人或者编译器来干谁干得不好性能就一目了然。这也是为什么AI芯片的性能优化特别依赖编译器、算子库和手写kernel的三层配合。这里有个实测感受硬件上Scratchpad做成了软件安排不好性能照样拉胯。我见过一个项目里DMA引擎的利用率只有30%原因就是搬运任务和计算任务没有重叠芯片大部分时间在干等数据。存储层次不只是硬件设计问题更是一整套软硬件协同的调度问题。3. 两类AI工作负载的“访存画像”GEMM和它的邻居们3.1 算力密集与访存密集的划分所有AI算子都可以按计算强度划成两类。算力密集型算子比如大尺寸GEMM、大卷积核卷积它们对同一批数据做大量重复运算数据搬运量相对小瓶颈在计算单元。访存密集型算子则相反比如Softmax、LayerNorm、Dropout、残差加、Embedding它们对每个数据几乎只做少量运算瓶颈在内存带宽。算子类型计算强度瓶颈典型代表算力密集高算力大GEMM、卷积访存密集低带宽Softmax、LayerNorm、激活函数混合中等视分块策略Attention融合前、GELU相关操作真正麻烦的是一个真实模型是这两类算子串在一起跑的。哪怕GEMM本身算得很爽周围的访存密集算子一旦拖后腿整体性能就被拉下来了。有公开分析统计过在BERT这类Transformer训练任务里非GEMM算子能占到三四成甚至更高的时间比例。这也是为什么仅仅优化GEMM核函数远远不够。3.2 训练场景容量墙和梯度通信是双重的“墙”训练比推理更容易撞墙而且是不同类型的墙一起撞。以7B模型、FP16混合精度训练为例粗略算一笔账权重FP16要14GB梯度FP16要14GBAdam优化器为每个参数维护一份FP32的一阶动量和二阶动量共56GBmaster 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涨到32KKV 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 片上SRAMAI芯片真正的胜负手别看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 L2H100做到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举例。假设计算CA×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是40MBH100是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性能打交道这两张图就是你的底层地图。先看数据怎么流再看计算怎么算性能优化才能少走弯路。 揭秘内存墙：从存储层次到算子优化的性能瓶颈与拆墙实践
跑GEMM类算子的时候芯片利用率轻松到70%甚至80%一旦切到真实模型推理整体利用率常常掉到20%上下。这个现象干过AI芯片或性能优化的人应该都不陌生。我入行头两年也被这个问题折磨得够呛总以为是框架调度不行或者是卡不行后来把存储层次和内存墙这两件事彻底想明白了才发现所有看似花哨的优化技巧——从FlashAttention到算子融合再到各种量化——本质都是在跟内存墙“躲猫猫”。这篇就把这个话题摊开讲清楚。这次要聊的内容对三类人最有用一是做AI芯片架构或硬件设计的二是搞高性能算子开发和模型部署优化的三是对AI加速底层原理感兴趣、想搞清楚“为什么算力那么高但跑模型还是慢”的读者。理解了存储层次和内存墙你再看市面上主流加速芯片的设计取舍、论文里各种优化方法的动机都会通透很多。1. 先算一笔账内存墙到底“卡”在哪儿1.1 算力与带宽的两个量级差先看一组不算冷门的数字。以NVIDIA H100 SXM为例FP16稠密矩阵乘法算力大约是990 TFLOPSHBM3内存带宽是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芯片的存储体系本质上是一层套一层的“缓存漏斗”从离计算单元最近的寄存器一直延伸到外部存储。用一张表可以把各层的量级看得比较清楚层级容量量级带宽量级延迟量级相对功耗量级寄存器堆KB级极高多读写口1ns基准片上SRAM几十KB到几十MB数十TB/s级几纳秒到几十纳秒数十倍于寄存器HBM高带宽内存几十GB到上百GB3~8TB/s级数百纳秒级远高于SRAM外部存储NVMe/SSDTB级GB/s级毫秒级更高这张表的关键信息是越靠近计算单元容量越小、带宽越高、延迟越低但这部分存储的每比特成本也越贵。所以芯片设计的目标不是“把所有数据都放在最快的地方”而是“让绝大多数访问尽量命中在靠近计算单元的地方”。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与CacheAI芯片为什么几乎都选前者CPU里用的是一套复杂的硬件Cache层次有同位、组相联、一致性协议。但用这套思路做AI芯片弊端很明显Cache的命中率靠硬件猜而AI负载的访问模式是可以静态分析的没必要让硬件浪费晶体管去猜。因此主流AI芯片几乎都选择了Scratchpad软件管理片上存储的方案。GPU里叫Shared MemoryTPU里叫Unified Buffer许多NPU里叫Local Memory名字不同本质相同这块SRAM由软件或编译器显式安排数据何时写入、何时读出、何时淘汰。好处是确定性高、面积省、功耗低代价是“数据搬运”这活儿得人或者编译器来干谁干得不好性能就一目了然。这也是为什么AI芯片的性能优化特别依赖编译器、算子库和手写kernel的三层配合。这里有个实测感受硬件上Scratchpad做成了软件安排不好性能照样拉胯。我见过一个项目里DMA引擎的利用率只有30%原因就是搬运任务和计算任务没有重叠芯片大部分时间在干等数据。存储层次不只是硬件设计问题更是一整套软硬件协同的调度问题。3. 两类AI工作负载的“访存画像”GEMM和它的邻居们3.1 算力密集与访存密集的划分所有AI算子都可以按计算强度划成两类。算力密集型算子比如大尺寸GEMM、大卷积核卷积它们对同一批数据做大量重复运算数据搬运量相对小瓶颈在计算单元。访存密集型算子则相反比如Softmax、LayerNorm、Dropout、残差加、Embedding它们对每个数据几乎只做少量运算瓶颈在内存带宽。算子类型计算强度瓶颈典型代表算力密集高算力大GEMM、卷积访存密集低带宽Softmax、LayerNorm、激活函数混合中等视分块策略Attention融合前、GELU相关操作真正麻烦的是一个真实模型是这两类算子串在一起跑的。哪怕GEMM本身算得很爽周围的访存密集算子一旦拖后腿整体性能就被拉下来了。有公开分析统计过在BERT这类Transformer训练任务里非GEMM算子能占到三四成甚至更高的时间比例。这也是为什么仅仅优化GEMM核函数远远不够。3.2 训练场景容量墙和梯度通信是双重的“墙”训练比推理更容易撞墙而且是不同类型的墙一起撞。以7B模型、FP16混合精度训练为例粗略算一笔账权重FP16要14GB梯度FP16要14GBAdam优化器为每个参数维护一份FP32的一阶动量和二阶动量共56GBmaster 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涨到32KKV 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 片上SRAMAI芯片真正的胜负手别看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 L2H100做到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举例。假设计算CA×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是40MBH100是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性能打交道这两张图就是你的底层地图。先看数据怎么流再看计算怎么算性能优化才能少走弯路。