来源： 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 for detail function packedAttn(Q, K, V, lens):for ℓ in lens: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) 在前文中，我们讨论了两个关键参数，这些参数决定了请求如何被分组到批次中：
最大 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 为主的数据集，并通过变化序列长度来评估在不同条件下的性能。
• N K prefill-heavy: N K 输入 tokens 和 1K 输出 tokens
• N K decode-heavy: 1K 输入 tokens 和 N K 输出 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 缓存，限制了批处理大小的扩展。然而，该策略确保了没有抢占，从而可能在最终的吞吐量上表现更优。
