内存墙之痛：为什么AI加速不能只靠堆计算单元？

本文来自微信公众号： 歪睿老哥 ，作者：歪睿老哥

很多人一说起AI加速，第一反应就是堆算力、堆更多浮点计算单元.

错了

现在的瓶颈根本不是计算本身，是数据搬运、通信和不规则算子，峰值算力再高也没用。

举个最直观的例子，大语言模型推理的时候，每个新生成的token都要读写一遍已经存起来的KV缓存，这玩意儿根本不怎么缺计算，缺的是内存带宽——你带宽不够，就算计算单元堆得再多，也得等着数据慢悠悠从内存运过来，大部分时间都在摸鱼。

还有现在流行的混合专家模型（MoE），每次token只走部分专家，这种动态路由的不规则计算，原来的通用硬件也根本跑不利索。

所以现在要跑最新的SOTA模型，不管是训练还是推理，没有专门的硬件加速优化，要么成本高到离谱，要么延迟慢到没法用，根本就没法落地。

现在整个硬件加速赛道已经铺开了，从带张量核心的通用GPU，到谷歌TPU、手机NPU这类专用张量处理器，还有FPGA可重构设计、ASIC专用推理芯片，最近还出来了专门给大语言模型推理做的LPU，甚至存算一体、神经形态计算这些新技术也在跑。

说白了，都是被刚需逼出来的——没有这些加速，现在的AI根本没法从实验室走到你我手机里。

不管是训练还是推理，现在的大模型早就把压力给到内存和通信了，纯堆计算单元根本解决不了问题。咱们拆开来给你讲清楚每个维度的瓶颈到底在哪。

首先就是功耗和能量瓶颈。

你知道吗？

移动一个字节数据消耗的能量，比做一次浮点计算高好几个数量级。

现在绝大多数场景下，整个芯片的功耗大部分都花在数据搬移上，不是花在计算上。

数据移动能耗显著高于算术运算能耗，

咱们分不同硬件看：

GPU功耗基本都绑在内存流量上，不做算子融合的话，大部分能量都浪费在寄存器、缓存和显存之间来回搬数据；

TPU/Edge NPU靠专用数据流省能量，但一旦工作负载不符合设计预期，需要频繁访问片外存储，能量效率马上掉下来；

ASIC和FPGA倒是能靠流水线和片上缓存减少搬移，但只要模型出了新算子需要回退到CPU，或者片上存不下要频繁换数据，能量优势立刻没了；

哪怕是专门针对大模型推理的LPU，还有存算一体、神经形态这些新架构，能量瓶颈也都绕不开KV缓存的搬移成本——本质还是数据移动的问题。

讲完功耗，再讲大家最关心的延迟和吞吐量。

很多厂商吹峰值TOPS，实际跑起来根本达不到，为什么？

因为吞吐量和延迟从来不是单靠算力堆出来的，瓶颈永远在数据那端。

LLM推理的预填和解码阶段瓶颈对比，解码阶段明显是内存带宽瓶颈

就拿大语言模型推理举例：整个过程分预填和解码两个阶段，预填是计算密集型，解码是妥妥的带宽密集型。

每个新生成的token都要读写整个KV缓存，数据搬不完，计算单元再强也只能闲着等数据。

不同硬件的痛点也不一样：

GPU靠批处理提吞吐量，但批越大排队越久，尾延迟根本控不住；

TPU/NPU对固定形状的dense算子跑的快，形状一变、序列一长，编译调度跟不上，延迟直接飙升；

ASIC低延迟靠固定算子，新算子不支持，延迟直接崩；

FPGA能做确定性低延迟，但路由拥堵和片上内存不够的问题，稍微长点的序列就掉吞吐量；

专门做LLM推理的LPU，哪怕把调度做到硬件里，也绕不开一个死规律：上下文越长，需要搬的KV缓存越多，延迟下限就是由内存带宽决定的，再怎么优化架构也破不了这个物理限制。

接下来是面积和成本瓶颈。

加速器设计中的面积，成本，性能权衡

硅片面积就那么大，你把面积给计算单元还是给内存，这是永恒的Trade-off。

现在越来越多的设计发现，把面积分给内存和互联，比分给更多计算单元回报高的多。

比如现在做大模型推理，哪怕你有满片的计算单元，要放下70B模型的权重和KV缓存，你就得买更多卡，哪怕每张卡的计算利用率只有30%，这个“容量税”你不得不交。

专门做低延迟的SRAM架构LPU，延迟确实牛，但要存大模型权重就得堆很多芯片，成本直接炸上天，根本不是一般人用得起的。

新架构也一样，存算一体要做很多ADC/DAC周边电路，这些东西占的面积比计算交叉阵还大；

神经形态芯片把大量面积分给存储和路由，做密集模型的时候面积效率远不如传统的脉动阵列，成本根本收不回来。

然后就是最核心的内存和通信瓶颈。

现在端到端性能，基本都是内存容量、内存带宽和通信限制死的，峰值算力就是个摆看的数字。

大模型训练和推理中的内存与通信瓶颈，有限HBM带宽和互联延迟是主要限制

训练的时候要存激活、梯度、优化器状态，大模型这些东西分分钟把显存占满；

推理的时候LLM的KV缓存，随着上下文长度和并发数线性涨，哪怕计算管够，带宽不够照样卡成PPT。

分布式训练更不用说，多卡之间做all-reduce，通信时间经常比计算时间还长，互联带宽不够，堆再多卡也扩不动。

不同硬件的具体情况也不同：

GPU现在算力涨的比内存带宽快得多，“内存墙”越来越明显，

TPU编译器只要工作集稍微超过片上SRAM容量，性能直接跳崖，比GPU降级快多了；

ASIC靠固定数据流省带宽，碰到Attention这种不规则访问，直接卡成狗；

LPU做多芯片互联，要求纳秒级同步，普通PCIe根本满足不了，必须做专用互联，成本又上去了；

存算一体倒是解决了片外搬移的问题，但片上网络又成了新瓶颈，交叉阵算完传不出来，计算单元照样闲着。

最后还有资源利用率的问题，很多人忽略这点，但实际体验里这个影响极大。

不规则的工作负载——比如非结构化稀疏、动态形状、MoE路由，都会导致负载不均衡，哪怕峰值算力很高，实际能用起来的可能不到一半。

不规则workload导致的负载不均衡和计算单元利用率下降示意图

举个例子，小批量推理的时候，很多ASIC为了峰值TOPS做了很宽的向量，batch=1的时候利用率连十分之一都不到；

TPU的脉动阵列跑matmul快的飞起，跑LayerNorm、Softmax这些小算子的向量单元不够用，整个大脉动阵列只能闲着；

大模型的预填和解构阶段，预填占满计算，解码阶段计算单元90%时间在等数据，利用率根本上不去，就算是专门做LPU的，也要靠特殊的调度硬件才能把利用率拉回来。

最后说下基准测试的坑，现在很多厂商吹的成绩都是挑最优场景测的，换个实际负载直接不对味。

性能结果对软件栈、精度、模型大小、batch大小、序列长度都特别敏感，同一个硬件不同测法能出差好几倍的结果。

比如很多存算一体的论文只测计算单元的峰值，不算ADC转换和非matmul算子的开销；很多NPU只测峰值TOPS，不说有多少算子不支持要回退CPU；

测大模型只说tokens per second，不告诉你并发多少、尾延迟多少，上下文多长，这些都是耍流氓。

神经网络加速器基准测试的挑战，结果对软件版本、编译选项、测量方法高度敏感

总结下来就是：

不管什么架构，什么场景，现在神经网络加速的核心矛盾从来都不是“不够算力”，而是“数据搬不过来”。

所有有效的优化，本质都是减少不必要的数据移动，让数据尽量待在离计算单元近的地方，这个逻辑从上世纪的“内存墙”提出到今天AI大爆发，从来没变过。

参考：Hardware Acceleration for Neural Networks:A Comprehensive Survey