当前位置：首页 > article >正文

[vLLM vs TensorRT-LLM] ：系统调度schedule比较

article 2024/12/24 21:53:26

来源：oldpan

原文：https://medium.com/squeezebits-team-blog/vllm-vs-tensorrt-llm-4-which-scheduler-wins-2dc15283522a

前言

Transformer 和LLMs的时代正在蓬勃发展。除了模型架构的演变之外，工作负载变得愈发动态化，使得系统级优化与模型级优化同等重要（类似于单一的视觉模型加上了前后处理）。特别是请求的调度与批处理方式，已经成为决定服务性能的关键因素。

尽管 vLLM 和 TensorRT-LLM 之间存在多种差异，其中调度器的设计差异更明显。不过优化的请求批处理与管理是提高性能并降低成本的关键，尤其是在计算和内存需求不断变化的情况下。因此，vLLM 和 TensorRT-LLM 都集成了专用的调度组件，以有效管理用户请求。

本文将探讨这些调度器的工作原理，以及它们如何影响性能和资源利用。

调度基础 Scheduling Basics

静态请求级调度

我们首先来看一种简单的调度形式：静态请求级调度。在这种方法中，传入的请求在到达时被分组到批次中并一同处理。批次中的所有请求都会并行处理，新请求需要等到当前批次的所有请求完成后才能被处理。

图 1: 静态请求级调度的 GIF 示例|700x394

这种方法虽然简单，但会导致低效，尤其是当单个请求具有不同的输入和输出长度时。较短序列的请求会因为批次内最长运行请求的完成而受到显著延迟。

迭代级调度

迭代级调度被引入以解决静态请求级调度的局限性。这种方法将任务分解为更小的单位，称为“ iterations”，而非对整个请求进行调度。在自回归模型中，迭代通常被定义为生成一个单独的 token。通过这种更细粒度的调度，迭代级调度最大限度地减少了资源浪费并提高了硬件利用率。

迭代级调度能显著提升计算效率，因为请求的 token 长度通常不同。提前完成的请求可以让新的请求加入批次，而不必等到整个批次完成。这种方式减少了硬件资源的空闲时间，并提高了整体吞吐量，尤其是在请求之间的 token 数量不同的工作负载中。

图 2: 迭代级调度的 GIF 示例|700x394

图 2 展示了调度器如何处理具有不同输出长度的多个请求。例如，六个请求（A 到 F）按批次大小为 2 进行调度。一开始，请求 A 和 B 被分组在一起并同时处理。然后，请求 B 比请求 A 更快完成。通过迭代级调度，请求 C 可以在请求 B 完成后立即开始，而无需等待请求 A 的完成。可以看到请求之间没有空闲等待，从而减少了整体计算时间。

Packed Batching

Packed Batching是高效执行已调度请求的另一个关键组件，尽管它本身并不是一种调度技术。正如图 2 所示，迭代级调度通常需要在同一迭代中处理预填充阶段和解码阶段。这两个阶段在输入 token 大小方面差异显著：预填充阶段一次处理多个 token，而解码阶段每次只处理一个 token，即前一迭代的输出 token （如我们前文所述）。

当预填充阶段的请求和解码阶段的请求被分组到同一批次中时，这种 token 长度的不一致会导致大量填充。即使所有批次请求都处于预填充阶段，由于输入 token 长度不同，也会产生填充。这种填充会降低计算效率，因为填充的额外计算实际上是无效的。

Packed Batching通过沿序列维度而非批次维度连接输入，从而消除不必要的填充并改善硬件利用率。在大多数 LLM 层中，紧凑批处理是简单的，因为这些层对序列长度不敏感。然而，LLM 模型的一个核心组件——注意力层（attention layer），需要稍微改点东西。每个请求需要不同的注意力模式，因此必须单独计算注意力。为此，Packed Batching需要通过切片连接的输入来分别计算每个请求的注意力。

以下是packed attention的简化版本伪代码：

# naive pseudo code of packed attention
# see <https://github.com/vllm-project/vllm/pull/3236> for detail
function packedAttn(Q, K, V, lens):
  out= empty_like(Q)
  s = 0
  for ℓ in lens:
    e = s+ℓ
    q = Q[s : e]
    k = K[s : e]
    v = V[s : e]
    out[s: e] = Attn(q, k, v)
    s = e
  return out

图 3: 简化packed attention实现中查询-键乘法的 GIF 示例|700x394

如图 3 所示，Packed 请求会在注意力层切片计算后沿序列维度再次连接。尽管为注意力层切片引入了一些开销，但消除不必要填充的好处通常超过了这些开销。因此，这种方法通过改善硬件利用率显著增强了服务性能。

连续批处理 (Continuous Batching 或 In-flight Batching)

通过集成迭代级批处理和Packed Batching，我们得到了 vLLM 和 TensorRT-LLM 调度器的核心：Continuous Batching（也称为“In-flight Batching”）。这种方法旨在最大限度地减少队列等待时间并减少填充开销，从而提高硬件利用率和服务性能。

vLLM 和 TensorRT-LLM 的调度策略在本质上是相同的，但在具体实现，特别是内存管理方面有所不同。这些差异是导致两个框架性能变化的关键因素。一个重要的影响因素是 KV 缓存（KV Cache）的管理，它在决定请求调度效率方面发挥了重要作用。下一节中，我们将深入探讨 KV 缓存管理如何影响调度及整体性能。

内存感知调度 (Memory-aware Scheduling)

在前文中，我们讨论了两个关键参数，这些参数决定了请求如何被分组到批次中：

最大批次大小（max batch size）
最大 token 数量（max number of tokens）

这两个参数限制了可以分组到一个批次中的请求数量。在某些情况下，批次受最大批次大小限制，而在其他情况下，受最大 token 数量限制。

除此之外，还有一个未提到的重要限制：KV 缓存（KV Cache）的大小。如果没有足够的剩余 KV 缓存存储请求的上下文，该请求将无法被调度。这是一个显著的挑战，因为在大多数 LLM 服务场景中，KV 缓存所需的内存通常超过加载模型本身所需的内存。尽管内存限制非常苛刻，但重新计算每个输出 token 的整个 KV 缓存必须被避免，因为这会带来巨大的计算开销。因此，即使批次大小或 token 数量不是限制因素，剩余 KV 缓存的大小仍然会限制能够一起分组的请求数量。

图 4: 在内存限制下的调度示例 GIF|700x394

然而，与其他限制因素不同，管理 KV 缓存大小并非确定性的——它会随着每个生成的 token 增长，最终可能扩展到最大输出 token 长度。因此，管理剩余 KV 缓存涉及一定程度的估算。与其他估算挑战类似，我们可以采用悲观或乐观的方式分配 KV 缓存。这两种策略分别称为预分配（preallocation）和按需分配（on-demand allocation）。

在预分配中，一旦请求被调度，其 KV 缓存的内存空间会基于输入 token 数量和最大生成 token 数量之和进行保留。这种方法确保了在解码阶段不会出现内存不足的情况，因为所需的最大 KV 缓存内存已经提前分配。TensorRT-LLM 使用预分配作为其默认策略，被称为 GUARANTEED_NO_EVICT。

预分配内存以支持最大 KV 缓存大小可能会导致内存使用效率低下，因为并非所有请求都会生成其最大 token 限制的内容。相比之下，按需分配会动态分配 KV 缓存内存，而不是预先保留最大值。

按需分配随着 token 的生成动态分配 KV 缓存内存，而不是预先为最大值保留内存。这种方法是 vLLM 的唯一策略，也是 TensorRT-LLM 的另一种策略，被称为 MAX_UTILIZATION。这种策略帮助最小化内存浪费，并允许更大的批次大小，但它引入了 KV 缓存耗尽（preemption）的风险。

Preemption

当活动批次中的请求上下文长度随着更多文本生成而增长时，可能会需要额外的按需 KV 缓存分配。如果可用 KV 缓存内存不足，就会发生预警。在这种情况下，批次中的某些请求必须被中断，其 KV 缓存需要被清除以释放内存，从而避免死锁。清除可以通过两种方式实现：将 KV 缓存交换到主存储器（host memory）或完全丢弃缓存。

• 交换（Swapping）：当请求被中断时，将 KV 缓存转移到主存储器，随后在请求恢复时再将其加载回内存。

• 丢弃（Dropping）：直接丢弃 KV 缓存，同时存储上下文和已生成的 token。在请求恢复时，KV 缓存在预填充阶段重新计算。

图 5: 由于上下文增长导致请求预警的 GIF 示例|700x394

与交换相比，丢弃更常被优先选择，因为主存储器的读写操作会引入显著的开销。而丢弃仅需要一次预填充迭代，将先前生成的 token 与原始输入 token 连接即可，因此在大多数情况下是一种更高效的选项。

TensorRT-LLM 如何调度请求

由于 TensorRT-LLM 部分闭源（proprietary），其确切的调度策略无法直接从源码中确定。然而，根据仔细的观察，它似乎采用了连续批处理方法，并且几乎没有修改。尽管源码不可公开访问，我们可以通过分析每次迭代中的请求模式进行推断。

TensorRT-LLM 使用 GUARANTEED_NO_EVICT 策略调度多个请求的示例|700x394

图中，每个请求以不同颜色表示。可以看到，每个请求需要 512 次迭代（输出长度设置为 512）。新的请求会在前一个请求完成后立即被加入。这种调度行为展示了连续批处理的核心原则，特别是迭代级调度，因为新请求会在完成的请求后立刻被引入。

当策略更改为 MAX_UTILIZATION 后，行为有所变化。

TensorRT-LLM 使用 MAX_UTILIZATION 策略调度多个请求的示例|700x393

在图中，可以观察到预警的存在。被中断的请求会从调度中移除，并在后续迭代中恢复。尽管存在预警，连续批处理的模式在图 7 中仍然清晰可见。

vLLM 如何调度请求

与 TensorRT-LLM 不同，vLLM 的调度器完全透明，因为其代码库是开源的。vLLM 同样采用了迭代级调度（iteration-level scheduling），这是实现连续批处理（continuous batching）的核心组成部分。但它在此基础上引入了两项独特的改进：不使用混合批处理（no mixed batching）以及优先处理 prefill 请求（prefill prioritization）。

不使用混合批处理

目前，vLLM 默认不支持混合批处理。这意味着 prefill 请求只会与其他 prefill 请求一起进行批处理，而 decode 请求只会与其他 decode 请求一起处理。这种设计简化了计算路径，因为每个批次仅处理相同阶段的请求。由于没有混合批处理，调度器必须采用另一种策略：优先处理 prefill 请求。

Prefill 请求的优先级

以下通过一个场景来说明为什么需要优先处理 prefill 请求：假设当前批次中的某个请求完成了，而请求池中还有新的请求等待加入批次。由于不支持混合批处理，新的请求无法直接加入当前批次，因为它需要先完成 prefill 阶段，才能进入 decode 阶段。因此，新的请求无法与当前正在处理的 decode 请求一起被批处理。

这种限制破坏了连续批处理的概念。为了解决这一问题，当前批次的 decode 请求需要暂时延后处理，先处理 prefill 请求，以确保连续批处理的流程不被中断。因此，为了在后续的 decode 迭代中确保有足够的 decode 请求可以处理，必须优先调度 prefill 请求。

有限的混合批处理支持

值得注意的是，当启用了分块 prefill（chunked prefill）时，vLLM 对混合批处理有一定程度的支持。这种方法可能实现类似于 TensorRT-LLM 的 MAX_UTILIZATION 策略的调度行为。然而，如果未启用分块 prefill，混合批处理将不可用。

实验设置

我们设计了实验来比较 vLLM 和 TensorRT-LLM 的不同调度策略。主要比较了以下几种策略：TensorRT-LLM 的 GUARANTEED_NO_EVICTION 策略、TensorRT-LLM 的 MAX_UTILIZATION 策略以及 vLLM。

框架版本、模型与硬件

• vLLM: v0.6.2

• TensorRT-LLM: 0.14.0.dev24092401（C++ API）

• 模型: Llama-3–8B（BF16）

• 硬件: NVIDIA A100-SXM 80G GPU，Intel Xeon(R) 2.20GHz（12 核心），128 GB 内存

基准测试数据集

在基准测试数据集中，我们使用了以 prefill 为主和以 decode 为主的数据集，并通过变化序列长度来评估在不同条件下的性能。

• NK prefill-heavy: NK 输入 tokens 和 1K 输出 tokens

• NK decode-heavy: 1K 输入 tokens 和 NK 输出 tokens

所有实验的最大序列长度设置为 (N+1)K。最大批处理大小设置为 256，最大 tokens 数量设置为 16384，请求速率设置为无限。

结果

为了分析调度器的行为，我们首先评估平均运行批处理大小，然后再分析端到端性能。调度器通过决定每次迭代处理的请求数量，展示了 vLLM 和 TensorRT-LLM 之间的关键差异。

尽管在混合批处理或 prefill 优先级上存在一些差异，但 vLLM 和采用 MAX_UTILIZATION 策略的 TensorRT-LLM 在序列长度增加时，平均批处理大小的下降趋势类似（如图 8 所示）。这是由于 KV 缓存的内存限制，随着序列长度的增长，调度器会进行抢占处理。

图 8 TensorRT-LLM（GUARANTEED_NO_EVICT）、TensorRT-LLM（MAX_UTILIZATION）和 vLLM 在各种场景下的平均批处理大小比较|700x423

在与其他两种策略的比较中，GUARANTEED_NO_EVICT 的平均批处理大小始终较小。这是因为 GUARANTEED_NO_EVICT 策略会预先分配 KV 缓存，限制了批处理大小的扩展。然而，该策略确保了没有抢占，从而可能在最终的吞吐量上表现更优。

通常情况下，较大的批处理大小会带来更高的吞吐量，但可能会降低 TPOT。然而，从图 9 和图 10 中可以看出，情况并非总是如此。

图 9: TensorRT-LLM（GUARANTEED_NO_EVICT）、TensorRT-LLM（MAX_UTILIZATION）和 vLLM 在各种场景下的吞吐量比较|700x423

在以 prefill 为主的场景中（输出长度限制为 1K），采用 GUARANTEED_NO_EVICT 策略的 TensorRT-LLM 实现了最高的吞吐量。这是因为抢占对性能的影响超过了增加批处理大小所带来的收益。当输出被限制为 1K 时，总解码步数较少，因此批处理大小的增加不会带来显著差异。

然而，在以 decode 为主的场景中，随着解码步数的增加，较大的批处理大小带来的收益更加明显。在这种情况下，采用 MAX_UTILIZATION 策略的 TensorRT-LLM 的吞吐量高于 GUARANTEED_NO_EVICT。

图 10: TensorRT-LLM（GUARANTEED_NO_EVICT）、TensorRT-LLM（MAX_UTILIZATION）和 vLLM 在各种场景下的 TPOT 比较|700x423

TPOT 数据还表明，关于“更长的序列长度或更大的批处理大小会降低 TPOT”的常见观点并不总是成立。在以 prefill 为主的数据集中，对比 MAX_UTILIZATION 和 GUARANTEED_NO_EVICT 策略时，这种假设似乎是成立的：MAX_UTILIZATION 的 TPOT 高于 GUARANTEED_NO_EVICT。然而，在分析 MAX_UTILIZATION 策略下不同输入长度时，这一行为表现得不那么一致。这是因为更长序列长度的影响可能会超过某些输入序列长度下批处理大小减少所带来的效果。

在以 decode 为主的数据集中，尤其是 GUARANTEED_NO_EVICT 策略的 TensorRT-LLM，TPOT 随输出长度增加而下降，尽管上下文长度更长。这是因为 GUARANTEED_NO_EVICT 即使在解码早期阶段也会减少批处理大小。而相比之下，采用 MAX_UTILIZATION 策略的 TensorRT-LLM 和 vLLM 则表现出更直接的趋势。这两种策略由于其激进的内存管理政策，调度了较大的批处理大小，从而错失了在序列长度较短时优化 TPOT 的机会。

最终思考

在本文中，我们探索了 vLLM 和 TensorRT-LLM 的调度策略，并通过固定长度的基准测试分析了它们对性能的影响。结果表明，调度策略的效果会因场景的不同而变化。

然而，由于我们使用的是固定输出长度的测试集，因此未能充分展现内存分配上的差异。当请求的序列长度变化显著时，一个更智能的调度器才能真正显示其优势。在这种情况下，动态调整批处理和资源分配的能力至关重要。如果我们使用动态长度的数据集，基准测试结果可能会有所不同。例如，将固定长度测试中的长度分布从均匀分布调整为正态分布，会显示出不同的趋势。