AI芯片：GPU、TPU、NPU，以及真正的瓶颈——内存带宽

为什么芯片的话题变成了内存的话题

提到AI芯片，人们通常想到的是“运算有多快”。这也是芯片规格表最上方写着FLOPS（每秒浮点运算次数）的原因。然而，真正把大语言模型（LLM）投入服务时，运算单元闲置的时间比想象中长得多。原因是计算所需的数据还没有从内存送达。

原因在于LLM生成文本的方式。LLM是自回归模型，每次只生成一个token，再把这个token接到输入后面，生成下一个token。每生成一个token，都需要从内存中重新读取模型的大部分权重。但读取一次权重后所做的计算，大约只是一次矩阵与向量的乘法。关于DRAM内部并行计算的研究（paper）指出，在边缘设备上，自回归解码的大部分时间都花在受内存限制的矩阵-向量乘法（GeMV）上。

面向NPU的推测解码研究（paper）也把在NPU硬件上进行LLM推理时遇到的问题称为“内存墙（Memory Wall）”。无论运算单元变得多快，只要内存带宽（每秒可搬运的数据量）跟不上，整体速度就会被内存一侧拖住。因此，如今AI芯片的竞争，很多时候是围绕内存和数据搬运展开的，而不是运算单元。

核心思路：食材运送比厨师还慢的厨房

想象一个餐厅厨房。厨师（运算单元）手非常快，但食材每次只能从远处的仓库（内存）运来固定的量。如果为一位客人一勺一勺地做菜，而每做一勺都要把全部食材从仓库重新运来一次，那么无论增加多少厨师，出菜速度都不会超过运送速度。

用这个比喻来看，芯片设计者能做的事大致有三类。第一，拓宽通道。在芯片旁边堆叠内存以提高带宽的HBM（高带宽内存）就是这个方向。第二，用一次运来的食材做好几份菜。把多个请求合并处理的批处理，以及一次验证多个token的推测解码都属于这一类。第三，在仓库里做菜。在内存内部进行计算的PIM（Processing-In-Memory）就是这个思路。

同样的比喻也能解释GPU、TPU、NPU的区别。三者都是配备了大量“快手厨师”的厨房，但不同之处在于建造时主要打算做什么菜，以及厨具（软件）配备得是否齐全。

GPU、TPU、NPU有何不同

GPU（图形处理器）最初是为同时计算屏幕上大量像素点而设计的芯片。它大规模并行处理相同计算的结构与神经网络的矩阵乘法非常契合，因此被广泛用于AI训练和推理。GPU最大的优势与其说在硬件本身，不如说在名为CUDA的编程生态。开发者、库和推理引擎大多首先在这一生态上构建。

TPU（张量处理器）是Google针对神经网络矩阵运算设计的专用芯片。它降低了通用性，转而专注于高效地处理矩阵乘法。NPU（神经网络处理器）是对专门用于神经网络运算的芯片的统称，范围从智能手机、边缘设备中的小芯片，到数据中心用的推理芯片。像Groq的TSP那样完全去掉缓存和硬件调度器、由编译器预先确定所有运算时序的静态调度芯片，也是这一类芯片中的一个极端。

芯片越专用，在特定运算上的功耗和成本效率就可能越好。但如果软件跟不上，这些优势就无法发挥。NPU推测解码研究（paper）的出发点，正是国内（中国）基础设施并不原生支持主流推测解码算法这一问题。在GPU上理所当然可用的优化，到了新芯片上就成了需要亲手重新实现的功课。这正是通用芯片与专用芯片之间取舍的核心。

工作原理：一次推理如何使用内存

LLM推理分为两个阶段。第一阶段是预填充（prefill），一次性处理用户发送的整个提示词。由于同时计算多个token，读取一次的权重可以多次复用，运算单元也容易被填满。在这一阶段，每个token的中间计算结果会被存入KV缓存，目的是避免在下一阶段重复相同的计算。

第二阶段是解码，逐个生成token。要生成一个token，需要从内存中读取模型权重和迄今积累的KV缓存，进行计算，再把生成的token重新作为输入。由于读取一次数据所做的计算量很少，这一阶段受内存带宽限制。对话越长，KV缓存也越大。因此LoopServe（paper）指出，对话历史越长，LLM的计算和内存负担会一同增加。

总结来说，“内存容量”决定了模型权重和KV缓存能否装进芯片，“内存带宽”决定了解码速度。换成用户的直观感受，预填充主要影响首个token的等待时间（TTFT），解码则主要影响此后文字输出的速度。

发展脉络：突破内存墙的三种尝试

第一种是减少内存读取的算法。推测解码让小模型预先猜测接下来的几个token，再由大模型一次性验证这些候选。读取一次大模型的权重就能确定多个token，从而减少内存往返。NPU推测解码研究（paper）为了让这一技术在NPU上运行，在OpenPangu-7B模型上实现了一套端到端（end-to-end）的推测推理加速体系。同一类研究中，分层验证树研究（paper）报告称，通过优先验证可能性高的候选、尽早剪除没有希望的候选，减少了推理时间和能耗。

第二种是更聪明地缩减KV缓存。LoopServe（paper）在预填充阶段针对每个输入只挑选注意力矩阵中重要的部分进行计算，在解码阶段则以最近生成的token为基准逐步压缩KV缓存。其出发点是：仅根据位置、按固定规则截断的方式无法跟上真实对话的走向。该研究还一并发布了一个包含11个多轮数据集的基准，反映了真实的提问位置和对话依赖关系。

第三种是改变硬件结构本身，即在内存内部计算的PIM。PRIMAL（paper）将LoRA（低秩适应）与基于PIM的LLM推理加速器结合。它用二维网格网络连接多种PIM处理单元，并通过SRAM重编程和关闭空闲部分电源的方式，让LoRA切换与计算重叠进行。在Llama-13B上应用LoRA秩8的设置下，报告的吞吐量为NVIDIA H100的1.5倍，能效为其25倍。这是一个说明不搬运数据、在仓库里做菜可以大幅提升能效的例子。

最新趋势

近期研究正在更深入地探索内存内计算。DRAM内部并行计算研究（paper）用直接在DRAM内部计算的PUD（Processing-using-DRAM）来处理解码中的GeMV。DRAM的并行度很高，可以在65,536列上同时计算。但正如Llama3 8B的QKV矩阵只有6,144行一样，LLM单层的权重太小，无法填满这种并行度。研究团队同时设计了把下一层计算与当前层打包在一起执行的算法，以及适配DRAM结构的数据布局。结果报告称，解码速度比PUD基线最高快1.41倍，比基于处理器的执行快2.45倍。

第二个趋势是不再只看“单块芯片”，而是衡量“芯片组合”的效率。集群效率研究（paper）认为，廉价加速器单芯片内存小，必须把多块组合起来使用，而这时芯片间通信决定成败。在由Blackhole芯片组成的Tenstorrent Galaxy上进行通信与计算重叠等优化后，按3年购置费和电费计算，其每美元token数是配备8块RTX PRO 6000 Blackwell GPU服务器的1.07~2.24倍。相反，在同一芯片上把通信模拟成TCP级别的低速后，吞吐量下降了23.5~67.4%，并且在所有评估设置中成本优势都消失了。这说明连接芯片的互连与内存带宽同样重要。

第三个趋势是尝试用AI弥补新芯片的软肋——软件。多智能体内核优化研究（paper）发现，最新的LLM大多能正确编写单算子CUDA内核，但面对几乎没见过的芯片和语言时，准确率会大幅崩溃。研究团队以使用文档匮乏的基于Rust的语言编程的FuriosaAI RNGD NPU为对象，构建了一个智能体循环。智能体自行撰写并验证文档摘要、测量芯片性能，并记录编译器规则。在12天内，它优化了3个Gemma-4-12B解码内核，相对官方基准取得了几何平均9.293倍的加速。每个内核都逼近到仅搬运权重所需周期数的1.25倍以内。这再次证实，优化的天花板最终也是内存搬运量。

局限与争议，以及实务判断标准

最常见的错误是只用FLOPS数字来比较芯片。像聊天机器人、编程智能体这类生成长回答的服务，解码占比大，内存带宽和容量决定了体感速度。如果处理长上下文，首先要考虑是否有足够的内存容量容纳KV缓存。如果打算把多块芯片组合使用，正如集群效率研究（paper）所显示的，互连较慢时成本优势可能会全部消失。

第二个错误是把论文数据泛化。PRIMAL（paper）的25倍能效是在Llama-13B和LoRA秩8这一特定设置下测得的。Galaxy的结果同样以特定模型和3年成本模型为前提。还要注意，相当多的PIM研究还处于设计和模拟实验阶段，而不是实际量产的芯片。也一定要确认比较基准是什么（是H100，还是同类基线）。DRAM内部计算研究（paper）中1.41倍和2.45倍的差异，就是基准不同造成的。

第三个是低估软件成熟度。专用NPU即使硬件效率高，也可能不原生支持推测解码之类的优化（paper）。编写内核也需要投入大量精力。在多智能体内核研究（paper）中，由于设备时间不足、测量值噪声大，也花了12天。在判断是否引入时，比较稳妥的做法是先确认当前使用的模型和推理引擎能否直接在该芯片上运行。

在实务指标上，最好养成不只看“每秒token数”，同时也看“每美元token数”和“每焦耳token数”的习惯。即便单块芯片很快，如果需要组合多块或耗电量大，总成本也会不同。

延伸学习

若想先对内存墙这个问题本身建立直观认识，建议先读DRAM内部并行计算研究（paper）中的问题描述部分。其中用具体数字说明了解码为什么受内存限制，以及硬件并行度与模型规模为什么不匹配。

在算法层面的应对上，依次阅读NPU推测解码研究（paper）和LoopServe（paper），可以比较“少读”策略的两条路线（打包验证token、缩减KV缓存）。在硬件层面，PRIMAL（paper）是了解PIM加速器设计的良好入口。

对于正在考虑实际引入的读者，推荐集群效率研究（paper）。它展示了芯片价格、电费和互连如何汇总成一个成本指标。如果你关注新NPU的软件问题，多智能体内核优化研究（paper）坦率地记录了一线的困难。