解密vLLM推理快的原因

Paged Attention+Continuous Batching+Flash Attention

一、vLLM

vLLM（虚拟大型语言模型）的技术，这是由Kwon等人在2023年9月发表的论文《Efficient Memory Management for Large Language Model Serving with PagedAttention》中提出的。vLLM主要解决了在使用GPU时内存分配的难题，特别是当前大型语言模型（LLM）服务系统中管理键值（KV）缓存内存的低效问题。这些问题导致了GPU资源没有被充分利用，推理速度变慢，并且内存使用量很高。

为了应对这些挑战，论文的作者受到了操作系统中使用的内存和分页技术的启发，提出了一种名为PagedAttention的注意力算法。PagedAttention使用分页技术，这是一种将硬件地址映射到虚拟地址的方法。通过这种方法，PagedAttention可以实现高效的内存管理，因为它允许注意力键和值（KV）在内存中非连续存储。

在推理请求的批处理方面，有两种主要的技术：

1. 客户端侧（静态）批处理：通常，当客户端向服务器发送请求时，服务器会默认按顺序处理每个请求，这样做的效率不高。为了提高效率，客户端可以将多个推理请求打包成一个批处理发送给服务器，服务器再将这个批处理拆分成多个请求并分别处理。这种方法需要客户端修改代码来实现批处理，并且这种解决方案与批处理的大小紧密相关。

2. 服务器端（动态）批处理：另一种方法是让服务器来处理批处理。当独立的推理请求到达服务器时，服务器可以动态地将它们组合成更大的批处理。服务器可以管理这些批处理，以满足指定的延迟目标，在保持所需延迟范围内的同时最大化吞吐量。服务器自动处理这一过程，因此不需要客户端修改代码。服务器端批处理包括多种技术来进一步优化生成语言模型的吞吐量，这些技术包括动态批处理、连续批处理和PagedAttention（vLLM）批处理。vLLM还使用连续批处理，在模型生成输出时动态调整批处理大小。
   
   连续批处理是一种专门用于文本生成的优化技术。它提高了吞吐量，并且不会增加首字节延迟。连续批处理（也称为迭代或滚动批处理）解决了GPU空闲时间的问题，通过不断将新的请求加入批处理中来提高效率。下图展示了连续批处理的工作原理。当请求2和3完成处理时，另一组请求被调度。

二、Paged Attention

vLLM的推理速度之所以较快，主要是依靠Paged Attention 技术。vLLM 是一个高吞吐量、低延迟的大型语言模型（LLM）推理和服务引擎，旨在提高模型在推理阶段的效率。

Paged Attention 是 vLLM 引入的一种注意力机制优化方法。它通过对注意力键值对（key-value caches）进行分页管理，实现了高效的内存利用。当处理长序列或大量并发请求时，Paged Attention 可以将不活跃的键值对临时移出显存，存储在更低成本的内存中，需要时再调入。这种机制避免了显存的过度占用，使得 vLLM 能够在不牺牲性能的情况下处理更大的模型和更长的序列。

上面这张图片描述的是一种名为“PagedAttention”的计算机内存管理技术，它是用于处理自然语言处理任务中的一种技术。在这个例子中，我们有一个句子：“the cat is sleeping in the kitchen and the dog is”。这个句子被分解成了一系列的tokens，每个词语都与一对“键-值”张量（key-value tensors）相关联。这些键-值张量是用于计算注意力（attention）的，注意力机制是一种让模型能够关注到句子中某些部分更多的技术。

在图中，我们可以看到两个主要部分：

1. 连续的虚拟块（Contiguous Virtual Blocks）：这些是逻辑上连续的内存块，用于存储每个词语的键-值张量。在这个例子中，有三个虚拟块（#0, #1, #2），每个块包含了句子中的一部分词语。
   

2. 非连续的GPU内存块（Non-Contiguous Blocks in the GPU Memory）：这些是物理上在GPU内存中的非连续块，它们被用来实际存储数据。由于内存的限制或者优化，这些块可能不会按照顺序存储。
   
   在图的中间部分，我们看到一个索引表（Index），它展示了虚拟块和物理GPU内存块之间的映射关系。例如，虚拟块#0（包含"the cat is sleeping"）映射到物理块#5，虚拟块#1（包含"in the kitchen and"）映射到物理块#2，虚拟块#2（包含"the dog is"）映射到物理块#3。
   
   这种映射允许计算机高效地处理大量数据，即使这些数据在物理内存中不是连续存储的。这对于处理大型模型和复杂任务（如机器翻译、语音识别等）非常重要，因为它们需要大量的内存和计算资源。
   
   简而言之，这张图片展示了如何在计算机的GPU内存中高效地组织和访问用于自然语言处理的数据。

PagedAttention。目标是在 GPU VRAM 的非连续空间中更有效地存储键值张量。简而言之，PagedAttention 背后的想法是创建映射到 GPU 内存中物理块的连续虚拟块。

每个区块都旨在存储预定义数量的令牌的键值对张量。所有块实际上都是连续的，并映射到物理非连续块，在推理期间按需分配，在碎片化的 GPU 内存中。还会在内存中创建一个简单的索引表，以将虚拟块与物理块相关联。

PagedAttention 的内核会根据需要获取这些块。这很有效，因为由于块的大小有限，系统获取的键值张量数量较少。

让我们以以下提示为例进行说明：

the cat is sleeping in the kitchen and the dog is

我们为每个令牌设置了键值张量。使用 PageAttention，我们可以（任意）将块大小设置为 4。每个块包含 4 个键值张量，但最后一个仅包含 3 个键值张量。这些块实际上是连续的，但在 GPU 内存中不一定是连续的。

为了计算注意力，对于每个查询令牌，系统会逐个获取块，如下图所示。

通过按块而不是整个张量序列获取键值张量，注意力的计算速度要快得多。

PagedAttention 的另一个优点是，在推理过程中采样时可以共享虚拟块。通过采样或光束搜索并行生成的所有序列都可以使用相同的虚拟模块，从而避免重复。

享虚拟块是 PagedAttention 的一项技术。PagedAttention 通过将显存划分为多个小块（block），并使用虚拟内存和分页技术来管理这些小块。在推理过程中，所有并行生成的序列（例如通过采样或光束搜索）可以共享这些虚拟块，从而避免重复存储相同的数据。

这种方法不仅节省了显存，还提高了内存管理的效率。通过共享虚拟块，PagedAttention 可以在不增加显存使用的情况下处理更多的并行请求，从而提高整体推理性能。

伯克利在使用 vLLM 中实现的 PagedAttention 与 Hugging Face 开发的文本生成推理库相比时报告的性能。

根据这些结果，vLLM 看起来要快得多，尤其是在多个输出完成的情况下。TGI 和 vLLM 之间的差异随着模型的增加而增加。这是意料之中的，因为较大的模型需要更多的内存，因此受内存碎片的影响更大。

三、vLLM默认预分配90%的显存

vLLM 预分配 90% 显存的主要原因是为了优化内存管理和提高推理效率。具体来说，vLLM 使用了一种称为 PagedAttention 的机制来管理注意力键值（KV）缓存。这种机制受到操作系统中虚拟内存和分页的启发，通过将显存划分为多个小块（block），按需分配内存，从而减少内存碎片和浪费。

默认情况下，vLLM 的 gpu_memory_utilization 参数设置为 0.9，这意味着它会预分配 90% 的显存来存储这些 KV 缓存。这种预分配方式可以确保在处理长序列或大批量数据时，有足够的显存来存储所有必要的中间结果，从而提高推理的速度和效率。此外，预分配 90% 的显存可以减少内存管理的开销，避免频繁的内存分配和释放操作。

PagedAttention 确实是为了优化内存使用。它通过将显存划分为多个小块（block），按需分配内存，从而减少内存碎片和浪费。这种方法不仅节约了内存，还提高了内存管理的效率。

然而，vLLM 预分配 90% 显存的主要原因是为了确保在处理长序列或大批量数据时，有足够的显存来存储所有必要的中间结果，从而提高推理的速度和效率。这种预分配方式可以减少内存管理的开销，避免频繁的内存分配和释放操作。

四、Flash attention

FlashAttention：具有 IO 意识的快速且内存效率高的精确注意力

我们查看论文中的下图：

左边的图表按照速度和容量从上到下排列了三种GPU的内存：

1. GPU的SRAM（静态随机存取存储器）：这是最快的内存类型，带宽高达19TB/s，但容量最小，只有20MB。

2. HBM（高带宽存储器）：这种内存的带宽为1.5TB/s，容量为40GB，是一种在GPU中用于高性能计算的内存。

3. DRAM（动态随机存取存储器）：这是主内存中的一种，带宽为12.8GB/s，容量超过1TB。

FlashAttention是一种优化的注意力机制计算方法，它通过将数据分块处理，避免了在GPU的HBM上生成大型的注意力矩阵。流程图显示了如何将K和V矩阵的块复制到快速的SRAM中，并在每个块中通过Q矩阵的块进行计算，然后将结果写回HBM。

右边的柱状图比较了使用FlashAttention技术和传统PyTorch实现在GPT-2模型上的性能差异。图中展示了矩阵乘法、Dropout、Softmax、Mask和融合内核等操作的耗时对比。结果表明，FlashAttention在所有这些操作上都显著减少了耗时，总体上实现了7.6倍的加速，这意味着FlashAttention大大提高了模型的计算效率。

在传统的注意力机制实现中，模型需要计算一个很大的注意力矩阵，这个矩阵的大小通常是输入序列长度的平方（N x N）。这种计算非常消耗内存和计算资源，尤其是对于长序列。

FlashAttention通过一种称为“分块”（tiling）的方法来优化这个过程。它不是一次性计算整个大型注意力矩阵，而是将输入序列分成小块，并且分别计算每个小块的注意力。这样做可以减少在GPU的高带宽存储器（HBM）上一次性需要存储的数据量，从而减少内存使用和提高计算效率。

在描述的流程图中，外循环（用红色箭头表示）指的是对K（键）和V（值）矩阵的块进行迭代，将这些块复制到快速的片上SRAM（一种高速缓存内存）。然后，对于每个K和V的块，内循环（用蓝色箭头表示）通过Q（查询）矩阵的块进行迭代，执行计算，并将结果存储回SRAM。最后，计算得到的注意力输出被写回到HBM。

总的来说，FlashAttention通过这种分块和优化的内存管理，能够显著减少执行注意力机制所需的时间，从而提高模型的整体性能。右边的柱状图展示了这种方法在实际应用中相比传统方法的性能提升，即在执行矩阵乘法、Dropout、Softmax等操作时的加速效果。

想象一下，你有一本非常厚的书，你的任务是找到所有提到“苹果”的句子。这本书太厚了，你不能一次性记住所有的内容，所以你决定采取一种策略：

1. 外循环（红色箭头）：你把书分成几个部分（我们称之为“块”）。每次你只拿出一部分（一个块）来查找“苹果”。这就像是你在大脑中有一个小的记忆空间（SRAM），你只在这个小空间里处理一部分书的内容。

2. 内循环（蓝色箭头）：在你处理每个部分的时候，你会逐页（或逐段）去查找“苹果”。每找到一个“苹果”，你就在一个小本子上做个记号，这个小本子代表你的工作空间（SRAM），它可以快速地记录和更新信息。

3. 写回HBM：当你完成了这部分书的查找后，你会把你在小本子上的记号整理好，然后把它们转移到一个大本子上（HBM），这样你就可以释放小本子的空间，准备处理下一个块。

在这个过程中，你的小本子（SRAM）是用来快速处理和记录信息的地方，而大本子（HBM）是用来存储所有你已经完成的工作的地方。通过这种方式，你可以有效地管理你的记忆和记录空间，使得查找“苹果”的任务变得更加高效。

在FlashAttention中，K（键）和V（值）矩阵的块就像是书的不同部分，Q（查询）矩阵的块就像是你在每个部分中查找的页或段。通过这种分块和循环的方法，FlashAttention能够高效地处理大量数据，而不会超出内存限制，从而加快了注意力机制的计算速度。

背后原理

FlashAttention 背后的主要理念是在 GPU 的 SRAM（上图中金字塔的顶端）上尽可能多地执行注意力计算。SRAM 是一种片上存储器，速度比 GPU HBM 存储器（我们通常称为 "VRAM"）快得多，但价格也贵得多，因此仅有极少量（通常少于 100 MB）可用。

FlashAttention 将注意力计算分解成可加载到 SRAM 上的小块。换句话说，它避免了在 HBM 上写入大型注意力矩阵。如上图中间部分所示。

减少非数学运算的数量

FlashAttention-2 通过减少非 Matmul 运算的数量，进一步加快了注意力计算的速度。

什么是 Matmul 运算呢？

当我们训练和运行 LLM 时，GPU 会执行大量的矩阵乘法运算，也就是 Matmul 运算。

最近的 GPU 都有专门用于加速计算的内核。例如，NV GPU tensor core专门用于计算矩阵乘法运算。张量内核是在 RTX 20xx 代开始使用的，但最近的 GPU（如 RTX 40xx）拥有更多的张量内核，速度也更快。不过，请注意只有安培（Ampere）或更新的 GPU（RTX 30xx）才支持 FlashAttention。

A100 GPU 的 FP16/BF16 matmul 最大理论吞吐量为 312 TFLOPs/s，但非 matmul FP32 吞吐量仅为 19.5 TFLOPs/s。另一种理解方式是，每个非 matmul FLOP 的成本是一个 matmul FLOP 的 16 倍。为了保持高吞吐量，我们希望在 matmul FLOP 上花费尽可能多的时间。

改进长序列的并行化

FlashAttention 的第一个版本针对批量大小和注意力头的并行计算进行了优化。它在处理大批量数据时效果尤佳。但是，在很多情况下，我们无法进行大批量处理，例如，当我们处理长序列的标记时。

现在，FlashAttention-2 也能对序列长度进行并行处理。因此，即使在使用小批量长序列时，我们仍能从 FlashAttention 中受益。

改进分区并支持更多 LLM

FlashAttention-2 还能更好地在计算线程之间进行分区。它减少了线程之间（更准确地说，是 32 个线程组成的 "warps "之间）的通信和同步量。

现在，多达 256 个heads的模型也支持 FlashAttention-2。使用多查询注意（MQA）和分组查询注意（GQA）的模型也可以从 FalshAttention-2 中受益。

虽然 Flash Attention 也是一种用于优化 Transformer 模型注意力计算的方法，旨在减少计算量和内存占用，提高训练和推理的效率，但 vLLM 的核心加速技术主要是基于 Paged Attention。

五、连续批处理

vLLM除了page attention之外，Continuous Batching也会大幅加快推理速度。

这张图展示了连续批处理（Continuous Batching）在处理多个推理请求时的工作原理。图中展示了三个请求（Request 1、Request 2 和 Request 3）如何在时间轴上被处理和生成响应。

图解说明：

1. 请求输入：

• Request 1: "Capital of"

• Request 2: "The diamondback turtle is"

• Request 3: "Largest Mammal is"
  

• Request 1: 从T3到T6，系统生成响应令牌，直到T7结束。

• Request 2: 从T3到T5，系统生成响应令牌，T6结束。

• Request 3: 从T3到T5，系统生成响应令牌，T5结束。
  

• Request 1: "of"

• Request 2: "diamondback"

• Request 3: "mammal"

• Request 1: "Capital"

• Request 2: "The"

• Request 3: "Largest"

• T1: 在第一个时间步（T1），系统开始处理三个请求的第一个词：

• T2: 在第二个时间步（T2），系统处理三个请求的第二个词：

• T3到T7: 从第三个时间步（T3）开始，系统开始生成响应令牌（Response token）：

详细解释：

• 连续批处理：图中展示了连续批处理的概念，即在处理请求时，系统不断地将新的请求加入到批处理中，而不是等待所有请求都完成后再开始处理新的请求。这种方法可以最大化GPU的利用率，减少空闲时间，提高整体吞吐量。
  

• 响应生成：在T3到T7的时间步中，系统开始为每个请求生成响应令牌。每个请求的响应生成时间不同，具体取决于请求的复杂性和长度。例如，Request 1的响应生成时间最长，直到T7才结束，而Request 3在T5就结束了。
  

• 并行处理：图中展示了多个请求在同一时间步内被并行处理的情况。这种并行处理方式可以显著提高系统的处理效率，减少每个请求的等待时间。
  

总结：

 
这张图通过展示连续批处理的工作流程，说明了如何在处理多个推理请求时提高系统的效率。通过不断地将新的请求加入批处理中，系统可以最大化GPU的利用率，减少空闲时间，从而提高整体吞吐量和响应速度。

参考：https://mrmaheshrajput.medium.com/deploy-llm-with-vllm-on-sagemaker-in-only-13-lines-of-code-1601f780c0cf