生成式大模型改变了我们在各个行业中应用人工智能的方式。然而，由于模型参数量巨大，为这些模型提供实际服务带来了挑战，即使在昂贵的硬件上，速度可能非常慢，这令人感到沮丧。😓
近日，来自UC Berkeley的团队推出了一个用于加速LLM推理的开源库vLLM，这也是Vicuna在线推理服务的幕后英雄。利用PagedAttention 技术，该库通过有效地管理Attention 模块中的Key 和Value 的Cache ，重新定义了LLM的推理服务。无需更改任何模型架构，它的吞吐量比原生HF Transformers高出24倍。🚀
接下来我们将对这个工作进行一个简单的解读。💡✨
www.bilibili.com/video/BV1zs...
核心要点
-
⚡️为什么切换简单？Easy
- 本质是推理过程中底层内存逻辑，CUDA Op的改变，与模型架构关联小，所以可以无痛切换。
-
🧠为什么节省了内存？Memory-Efficient
- 序列存储的时候物理上分块按需存储，尽量降低了物理空间的浪费，并且原生支持共享内存。
-
🎯为什么更推理的吞吐增加了？Fast
- 通过共享内存，分块内存的方式，降低了单Seq的内存占用，增加BatchSize，提升GPU利用率。
提出问题
随着大模型的参数量不断增加，其智能水平也随之提升。然而，与此同时，这种增长也给推理能力带来了巨大的挑战。在追求更智能的大模型的同时，我们应如何解决这一推理挑战，使得大模型更加易用？
解决方案
KV Cache
KV Cache是大模型推理优化的一个常用技术，该技术以空间换时间的思想，通过使用上次推理的KV缓存，可以在不影响任何计算精度的前提下，提高推理性能，降低端到端的时延。
以GPT为代表的Decoder-Only自回归语言模型在生成每一个新的 token 时，接受所有之前生成的 tokens 作为输入。然而，对于这些先前生成的 tokens，每次生成新的 token 时都需要重新计算他们的表示，这个过程造成了大量的计算浪费。KV Cache 的引入就是为了解决这个问题。
KV Cache实质上是存储了之前计算过的 key-value 对用于下一个Token的生成。在 Transformer 结构中，self-attention 中的k_proj, v_proj会将输入的每个 token 转化为一个 key 和一个 value，然后使用这些 key-value 以及当前的query对来计算下一个 token。引入 KV Cache，我们就可以将之前生成的 tokens 对应的 key-value 对存储起来，当生成新的 token 时，直接从 KV Cache 中取出这些已经计算好的 key-value 对，再把当前token的key-value做一个连结在进行计算，这样就避免了KV的重复计算，大大提高了计算效率。
整体来说，使用KV Cache包含以下两个步骤：
预填充阶段：在计算第一个输出token过程中，此时Cache是空的，计算时需要为每个 transformer layer 计算并保存key cache和value cache，在输出token时Cache完成填充；FLOPs同KV Cache关闭一致，存在大量gemm操作，推理速度慢，这时属于Compute-bound类型计算。
KV Cache阶段：在计算第二个输出token至最后一个token过程中，此时Cache是有值的，每轮推理只需读取Cache，同时将当前轮计算出的新的Key、Value追加写入至Cache；FLOPs降低，gemm变为gemv操作，推理速度相对第一阶段变快，这时属于Memory-bound类型计算。
无KV Cache生成示例
KV Cache生成示例
PagedAttention
通过KV Cache的技术，我们已经可以极大地提升LLM地推理速度，但是现有的Cache仍存在一些问题，
- Large：对于LLaMA-13B中的单个序列，它占用高达1.7GB的内存。
- Dynamic：它的大小取决于序列长度，而序列长度具有高度可变和不可预测的特点。
因此，高效地管理KV Cache是一个重大挑战。现有系统（HuggingFace 默认实现是pytorch的内存分配策略）由于内存碎片化和过度预留而浪费了60%至80%的内存。
为了解决这个问题，我们引入了PagedAttention，这是一种受传统操作系统虚拟内存和分页概念 启发的注意力算法。与传统的注意力算法不同，PagedAttention允许将连续的键和值存储在非连续的内存空间中。具体而言，PagedAttention将每个序列的KV缓存分成多个块，每个块包含固定数量的标记的键和值。在注意力计算过程中，PagedAttention Kernel高效地识别和获取这些块，采用并行的方式加速计算。（和ByteTransformer的思想有点像）
内存布局
由于块在内存中不需要连续存储，我们可以像操作系统的虚拟内存那样以更加灵活的方式管理键和值的缓存：可以将块看作页，标记看作字节，序列看作进程。序列的连续逻辑块通过块表映射到非连续的物理块。随着生成新的标记，序列的边长，物理块按需进行分配。
在PagedAttention中，内存浪费仅发生在序列的最后一个块中。这样就使得我们的方案接近最优的内存使用率，仅有不到4%的浪费。通过内存效率的提升，我们能够显著提升BatchSize，同时进行多个序列的推理，提高GPU利用率，从而显著提高吞吐量。
内存共享
在并行采样中，从相同的提示生成多个输出序列。在这种情况下，可以在输出序列之间共享提示的计算和内存。通过其块表，PagedAttention能够自然地实现内存共享。类似于进程共享物理页，PagedAttention中的不同序列可以通过将它们的逻辑块映射到相同的物理块来共享块。为确保安全共享，PagedAttention跟踪物理块的引用计数并实现 Copy-on-Write 机制。
通过PagedAttention的内存共享机制，极大地降低了复杂采样算法（如ParallelSampling和BeamSearch）的内存开销，使其内存使用量下降了高达55%。这项优化可以直接带来最多2.2倍的吞吐量提升，从而使得LLM服务中使用这些采样方法变得更加实用。
实验验证
吞吐的实验已经经过了官方的验证，我们这里就不在多测。我们对OPT-125m的HuggingFace的实现和vLLM的实现进行Pytorch Timeline测试，通过实验不难发现，相比于原始的HuggingFace的实现，vLLM减少了大量的类似CatArrayBatchedCopy的内存操作，大大提升了每一层前向的速度，降低了端到端的推理延迟，取得了良好的效果。目前，vLLM经过广泛验证并取得显著成果。其应用不仅大幅降低了运营成本，还使LM-Sys在处理相同流量时将所需GPU数量减少了50%。目前，vLLM每天平均能处理30,000个请求，并在峰值时达到60,000个，这充分证明了其鲁棒性。
从FlashAttention到PagedAttention，我们越来越意识到，仅仅掌握算法是不够的。除了深入了解各种模型和技术之外，我们还需要关注整个系统的底层实现。通过深入地了解底层硬件的特性和限制，运用内存管理、数据交换和并行计算等方面的知识，以便优化计算和存储的效率，进而设计出高效、可扩展的解决方案。只有这样，我们才能充分发挥机器学习算法的潜力，并构建出具备良好性能和可维护性的系统。
Reference: