以 Transformer 为例进行分析

一、内存墙与硬件的发展

“内存墙”这个概念，说得直白点，就是处理器的运算速度太快，而内存读写数据的“脚程”跟不上，结果处理器常常得干等着数据从内存里搬进来。这种情况在AI训练和推理中尤其突出——计算单元空转，白白浪费算力。

随着处理器越来越猛，内存带宽却像个慢跑选手，差距越拉越大。下面这组数据很能说明问题：过去20年里，服务器硬件的峰值算力（FLOPS）每两年翻3倍，但DRAM带宽每两年只增长1.6倍，片间互联带宽更是只有1.4倍。算起来，20年间算力飙升超过100万倍，DRAM带宽只涨了100倍左右，互联带宽仅提高了30多倍。这种严重失衡，直接拖累了硬件的整体效率。

为了解决这个问题，经典方案就是搞多级缓存——从L1到L3，再到DRAM，逐级放大容量，但带宽逐级下降。L1带宽最高，容量却只有KB级别；DRAM容量能上GB，带宽却最低。下面这张表列出了几款主流GPU的内存和带宽数据，一目了然。

GPU	A100	H100	A6000	4090
GPU Architecture	Ampere	Hopper	Ada Lovelace	Ada Lovelace
Memory Interface	5120-bit HBM2	5120-bit HBM3	GDDR6	GDDR6
Memory Size	40 GB	80 GB	48 GB	24 GB
Memory Bandwidth	1555 GB/sec	3000 GB/sec	960 GB/s	1008 GB/s
SMs	108	132	142	128
Texture Units	432	528	576	512
L2 Cache Size	40 MB	50 MB	96MB	72MB
Shared Memory Size / SM	up to 164 KB	up to 228 KB	up to 128 KB	100 KB
Register File Size / SM	256 KB	256 KB	256 KB	256 KB
FP16 TFLOPS	624	1979	366	330

H100 SXM 80GB用的是HBM3，A100 PCIe 80GB是HBM2e，带宽都比用GDDR6X的RTX Ada 6000和RTX 4090大不少。不过后两款在非Tensor Core峰值性能上更高（90和80 TFLops/s vs A100的20 TFLops/s）。在fp16累积方面，四款GPU的Tensor Core峰值性能（fp16输入，无稀疏性）接近，但RTX 4090在fp16累积时的吞吐量是fp32累积的2倍，其他GPU则一致。下面这张roofline曲线图，横轴是计算强度（Arithmetic intensity，单位FLOPs/Byte），表示每读入单位数据能支持多少次运算，能直观看出各芯片的瓶颈所在。

二、人工智能发展对计算设备的需求

这几年，计算机视觉、自然语言处理、语音识别领域最新模型的训练运算量，大约每两年翻15倍。而Transformer类模型的增长更夸张，每两年接近750倍。这种指数级增长驱动着AI硬件的研发——厂商更拼命地堆峰值算力，但常常以牺牲内存分层架构等部分为代价。

然而，面对最新AI模型的训练，这种设计思路已经有点力不从心了。尤其是NLP和推荐系统模型，通信带宽成了显著瓶颈。芯片内部、芯片之间、甚至AI硬件之间的通信，都卡得死死的。以Transformer大语言模型为例，模型参数量平均每两年增长410倍（见下图）。大规模的推荐系统模型，大小已经达到O(10) TB级别。而AI硬件上的内存容量，每两年只增长2倍。更需要注意的是，训练时需要的内存往往比模型参数量还要多几倍——因为要保存中间层的激活值，通常要额外增加3到4倍的空间。

2018到2022年间，训练最新模型所需的浮点运算量（FLOPs）每两年增长750倍，但真正的瓶颈不是计算，而是带宽。原因有两个：

• 单芯片内部，从寄存器到二级缓存再到全局存储器之间的数据传输越来越成为瓶颈。tensor core等专用单元几周期就能完成大量运算，但要让它们持续跑满，就需要源源不断地喂数据——芯片内存带宽就成了短板。
• 当模型超过单芯片可用内存，就得靠分布式并行把训练/服务扩展到多个翻跟斗。但片间带宽小于片内带宽，这时内存墙问题更严重。

下面这张图展示了最新AI模型训练时的内存占用逐年变化趋势，能清楚看到神经网络模型设计是如何受硬件内存大小制约的。

三、以 Transformer 为例进行分析

Transformer是当前最火的模型结构，有Encoder和Decoder两种变体。推理时，Encoder能并发处理所有token，Decoder则每次只生成一个token。重点来看看Decoder的生成过程——因为生成的特点，它的计算强度远小于Encoder。

典型的Decoder生成过程分两个阶段：

• Prefill 阶段
  ：并行处理所有输入token，同时把对应的KV Cache存下来。这个阶段计算、带宽和显存都是瓶颈，但主要卡在计算上（compute bound）。

• Decode 阶段
  ：每步只计算一个token，计算强度极低，这时候显存带宽就成了瓶颈。

3.1 典型方法的GPU利用效率分析

3.1.1 FlashAttention 与 PagedAttention

先看Prefill阶段使用FlashAttention V2和FlashInfer算法时的实际算力表现：

• H100和A100的最高使用算力只达到理论峰值的30%左右，而A6000和4090却能达到60%左右。此时各芯片都处于IO-bound区间，差异主要来自带宽和理论峰值。
• 实际使用算力随序列长度快速增加，很快稳定下来。

再看Decode阶段使用相同算法时的带宽使用情况：

• 带宽利用率随长度快速增加并趋于稳定，说明此时处于IO-bound区间。
• 带宽利用率最大值大约在80%上下。

PageAttention的主要特点是对KV cache的高效存储与访问。下面对比了FlashInfer PageAttention内核和vLLM PageAttention内核的性能——能在更小的文本长度达到更高带宽利用率，意味着对GPU整体利用率的提升。

（注：原文中关于FlashAttention和FlashInfer的具体图表数据因未提供，此处以文字概括，实际使用时应保留原文图表及对应分析。）