原文：https://zhuanlan.zhihu.com/p/699027700

一、内存墙与硬件的发展

内存墙（memory wall）问题是指当处理器的计算速度超过数据传入和传出内存系统的速度时，计算机架构中出现的一种现象。在这种情况下，处理器必须等待从内存中获取数据，这会降低其性能并限制其速度。

随着处理器变得更快、更强大，而内存速度却跟不上这些进步，内存墙问题变得越来越严重。这意味着即使处理器可以快速执行指令，它也会花费大量时间等待数据传入和传出内存。

如下图所示，在过去 20 年中，

• 服务器硬件 FLOPS 峰值以每 2 年 3.0 倍的速度增长

• DRAM 带宽以每 2 年 1.6 倍的速度增长

• 片间互连带宽以每 2 年 1.4 倍的速度增长

这种发展是不均衡的，在 20 年的时间里，算力提升了超过 1000000 倍，而 DRAM 带宽只提高了约100倍，互联带宽仅提高了 30 余倍。这严重制约了硬件的效率。

缓解这一问题的典型方法就是使用多级缓存，如下图所示。数据缓存有几个级别（主要是L1到L3），其次是DRAM。每个存储器的存储器带宽差异很大。一级缓存 (L1) 的内存带宽最高，而 DRAM 的内存带宽最低。然而，L1 的容量相当低（ KB 量级）。相比之下，DRAM 可能要大得多（ GB 量级）。下表中列举了当前几款主流的 GPU 产品的内存及带宽数据。

H100 SXM 80GB 使用 HBM3，A100 PCIe 80GB 使用 HBM2e，两者都比使用 GDDR6X 的 RTX Ada 6000 和 RTX 4090 具有更大的内存带宽。RTX Ada 6000 和 RTX 4090 的非 Tensor Core 峰值性能（分别为 90 和 80 TFLops/s）比 A100（20 TFLops/s）大得多。后三款 GPU 在 f16 累积方面具有相似的峰值 Tensor Core（fp16 输入，无稀疏性）性能，RTX 4090 的 Tensor Core 在 fp16 累积方面的吞吐量是 fp32 积累的 2 倍，而其他 GPU 在 fp16 和 fp32 积累方面的吞吐量相同。以下是 4 个 GPU 的 Tensor Core 和 CUDA Core 的 roofline 曲线，其中横坐标为计算强度(Arithmetic intensity or Operational intensity) 即算法对于内存带宽的需求，单位是 FLOPs/Byte。意思是在算法中平均每读入单位数据，能支持多少次运算操作。

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

最近几年，计算机视觉(CV)，自然语言处理(NLP)和语音识别领域最新模型的训练运算量，以大约每两年翻15倍数的速度在增长。而 Transformer 类的模型运算量的增长则更为夸张，约为每两年翻 750 倍。这种接近指数增长的趋势驱动了 AI 硬件的研发，这些 AI 硬件更专注于提高硬件的峰值算力，但是通常以简化或者删除其他部分（例如内存的分层架构）为代价。

然而，在应付最新 AI 模型的训练时，这些设计上的趋势已经显得捉襟见肘，特别是对于 NLP 和 推荐系统相关的模型：有通信带宽瓶颈。事实上，芯片内部、芯片间还有 AI 硬件之间的通信，都已成为不少 AI 应用的瓶颈。特别是基于 Transformer 的大语言模型，模型参数量平均每两年增长 410 倍（如下图所示）。类似的，大规模的推荐系统模型，模型大小已经达到了 O(10) TB 的级别了。与之相比，AI 硬件上的内存大小仅仅是以每两年翻 2 倍的速率在增长。

值得注意的是，训练 AI 模型时候所需要的内存一般比模型参数量还要多几倍。这是因为训练时候需要保存中间层的输出激活值，通常需要增加3到4倍的内存占用。在 2018-2022 年的时间框架内，训练这些最新模型所需的计算/浮点运算（FLOPs）增加了 750x /2 年，但瓶颈并非计算，而是带宽，理由如下：

• 对于单个芯片，芯片内寄存器、二级缓存、全局存储器等之间的内存传输仍日益成为瓶颈。得益于 tensor core等专用计算单元的最新进展，大量计算的算术运算可以在几个周期内完成。因此，要让这些运算单元始终处于可用状态，就需要快速向它们输入大量数据，而这正是芯片内存带宽成为瓶颈的原因

• 当模型超过了单芯片可用的内存，需要利用分布式内存并行技术，将训练/服务扩展到多个加速器，以避免单个硬件的内存容量和带宽有限。由于片间带宽小于片内带宽，此时会面临更严重的内存墙问题

下图展示了最新的 AI 模型训练时候，内存占用大小逐年的增长变化趋势。从中能清楚地看到，神经网络模型的设计是如何受 AI 硬件内存大小影响的。

三、以 Transformer 为例进行分析

Transformer 是当前最流行的模型结构，而其又包含 2 种变体：Encoder 和 Decoder，在推理过程中，前者可并发处理所有 token ，后者则在每次迭代时生成一个 token。现在主要以 Decoder 生成过程为例进行分析，首先因为生成的特点，Decoder 生成过程的计算强度远远小于 Encoder。

更加详细地，典型 Decoder 生成过程还可以分为 2 个阶段：

• 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 阶段使用 FlashAttention V2 与 FlashInfer 算法在不同芯片上的带宽使用情况，可以得到以下结论：

• 带宽利用率随着长度的增加快速增加，并快速达到稳定值，说明此时处于 IO-bound 区间；

• 带宽利用率的最大值在 80% 上下

PageAttention 的主要特性是对 KV cache 高效存储与访问，以下是 FlashInfer PageAttention 内核和 vLLM PageAttention 内核的性能对比，需要说明的是，能够在更小的文本长度达到更高的带宽利用率，即代表对 GPU 的整体利用率的同等提高。