LLM 连续批处理异步优化方案

简而言之：我们将解释如何分离 CPU 和 GPU 工作负载，从而大幅提升推理性能。

这是高效 LLM 推理系列的第二篇文章。第一篇文章从基本原理出发介绍了连续批处理，其中引入了一些本文将继续使用的概念：KV cache、FlashAttention、注意力掩码等。

在 Inference Endpoints 上，一块 H200 每小时的费用约为 5 美元。用一个小时并不贵，但如果用上一整天，你就已经要支付 120 美元。在这种情况下，你自然希望充分利用 GPU。我们已经看到，连续批处理通过调度紧密排列的批次来提高 GPU 利用率，从而避免将计算浪费在填充上。但连续批处理并未解决另一种浪费：默认情况下，它是同步的。这意味着 CPU 和 GPU 会交替工作：GPU 计算时，CPU 等待；CPU 准备下一个批次时，GPU 等待。在每秒运行数百个步骤的循环中，这些空闲间隙会不断累积。正如我们将展示的那样，它们可能占总运行时间的近四分之一。为了确保 GPU 100% 的时间都在忙于计算，我们必须消除这些间隙。

为此，我们可以使用异步批处理：将 CPU 上的批次准备与 GPU 上的批次计算解耦，使二者能够并行运行，确保 GPU 始终在高效工作 🔥

朴素的同步批处理是这样工作的：

CPU 准备新批次时，会选择要纳入其中的请求、更新 KV cache 表、移除在此前运行中已经完成的请求，并接纳新请求以填充释放出来的空间。完成这些操作后，它会将准备好的输入传输到 GPU。GPU 执行前向传播，并为每个请求采样（即选择）一个新 token。结果随后返回 CPU，使其知道每个请求刚刚生成了哪个 token，接着整个周期重新开始。

请注意右侧的红色标注：GPU 完成计算后会进入空闲状态。只有等 CPU 完成更新步骤，下一个批次才能开始：采样输出 token、更新请求状态、重新调度批次。

这就是同步批处理的核心低效之处：CPU 和 GPU 交替工作。GPU 计算时，CPU 处于空闲状态；CPU 更新时，GPU 处于空闲状态。无论在什么情况下，它们都不会同时执行有用的工作。对于单次前向传播来说，这似乎只是很小的代价；但在每秒运行数百个步骤的连续批处理循环中，这些空闲间隙会不断累积，造成实实在在的吞吐量损失。

为了展示这一点，我们对使用 8B 模型、批次大小为 32、生成 8K token 时 CPU 和 GPU 的耗时进行了性能分析：

如果你想生成同类图表，可以在连续批处理代码中添加检测逻辑，以导出 CPU 和 GPU 的活动时间区间，然后使用此脚本。

时间线在绿色（GPU 活跃、CPU 空闲）和红色（CPU 活跃、GPU 空闲）之间交替，二者从不重叠。总生成时间为 300.6 秒，其中有 24.0% 的时间是 GPU 处于空闲状态，等待 CPU 完成工作。从 GPU 的角度来看，全部生成时间中有近四分之一被浪费了。这是一种悲观的看法。

乐观的看法是，如果能够彻底消除 CPU 开销，生成时间就会从 300 秒缩短到 228 秒（免费获得 24% 的加速！）。这不需要任何新的 kernel 或模型改动，只需要仔细协调硬件的工作。

从根本上说，思路很简单：我们需要设法在批次 N 进行计算的同时，为批次 N+1 做准备。但这个简单的想法背后隐藏着几个技术难题：

如何在 GPU 上启动某项工作，同时让 CPU 重新获得控制权？

如何确保每项 CPU 或 GPU 任务启动时，它所需的数据都已准备就绪？

如果批次 N+1 依赖批次 N 的预测结果，我们该如何提前准备批次 N+1？

通过回答这些问题，我们将从零开始构建异步批处理。我们也遵循了相同的步骤，将其作为连续批处理的一部分实现到了 transformers 库中。欢迎查看代码并进行比较！

我们的最终目标是让 CPU 和 GPU 操作并发执行。为此，我们需要一种对操作进行分类的方法，以便告诉机器哪些操作可以并发运行。我们可以使用 CUDA stream 来实现这一点。

什么是 CUDA stream？

要理解 CUDA 如何对操作进行排序，我们需要先介绍 CUDA stream。stream 是一个有序的 GPU 操作队列，其中可以包含 kernel 启动、内存复制和同步屏障等操作，所有操作都会按照提交顺序执行。每个 GPU 操作始终都会被调度到某个 stream 中。同一个 stream 内的操作按顺序执行：前一个操作完成之前，GPU 不会开始执行下一个操作。不同 stream 中的操作彼此独立，可以并发运行。举例来说，如果你在 3 个不同的 stream 中分别启动 3 个操作，其执行过程如下：

这三个操作会同时开始。这里稍微做了一些简化：每个 GPU 操作最终都由 CPU 发起，而发起操作需要少量时间，例如查找正确的 kernel、发出调用、将命令从 CPU 传输到 GPU 等。这称为 CPU 启动开销。更符合实际情况的示意图如下：

这些操作仍然是并发执行的，但由于每次 CPU 启动都存在成本，它们的开始时间会依次错开。后文中我们会继续展示这些 CPU 启动事件，因为它们确实会消耗时间；随着我们转向异步工作流，它们还可以帮助我们追踪“什么操作在什么时候启动”。例如，我们会经常检查某个 stream 是否已清空：这表示该 stream 中的所有操作均已执行完毕。

默认 stream 和非默认 stream

如果你从未在 PyTorch 中显式使用过 CUDA stream，可能会惊讶于它们的存在。典型的 PyTorch 脚本从不提及 stream，而且使用时也感觉不到 GPU 操作是异步的：CPU 似乎总是在等待 GPU 完成后才继续执行。这种感觉是准确的，其原因就在于默认 stream。

当你调用 PyTorch 操作而未指定 stream 时，该操作会进入默认 stream。默认 stream 有一个特殊属性：它具有同步性。如果一个操作被调度到默认 stream，它会等待其他所有 stream 清空，也就是说，在默认 stream 中的任何一个操作开始之前，GPU 上的全部工作都必须结束。反过来也一样：任何操作，无论位于哪个 stream，在启动前都必须等待默认 stream 清空。

因此，如果你将默认 stream 中某个操作的结果传输到 CPU，即便使用的传输方式按理说不会阻塞 CPU，CPU 仍然会阻塞，直至所有 GPU 操作完成，因为这些操作被调度到了默认 stream。这实际上会破坏构建并发执行所做的任何努力。

这就是我们需要使用非默认 stream 的原因。将一次 kernel 启动或非阻塞内存复制加入队列后，控制权会立即返回 CPU。GPU 会在后台运行该操作，但 CPU 不会等待。这回答了我们的第一个问题：为了在启动 GPU 工作后重新获得 CPU 控制权，我们需要使用非默认 stream。

在本文后续内容中，我们将假设设备之间的所有内存传输都是非阻塞的。因此，我们必须自行对它们进行同步。

回到连续批处理

我们已经确定，不应让任何 GPU 操作进入默认 stream。但问题仍然存在：如果不使用默认 stream，我们应该使用哪些 stream？让我们回到同步批处理的示意图：

我们可以识别出三种不同的 GPU 操作：

将输入从 CPU 传输到 GPU

在 GPU 上执行计算

将输出从 GPU 传输到 CPU

这意味着我们需要三个 stream：一个用于计算，一个用于 CPU 到 GPU 的传输，另一个用于 GPU 到 CPU 的传输。这两种传输彼此独立，没有理由让它们串行执行，因此每种传输都使用自己的 stream。

关于术语的说明：讨论 CPU 和 GPU 时，CUDA 文档采用的惯例是将 CPU 称为主机（host），将 GPU 称为设备（device）。从现在开始，我们也将遵循这一惯例。CPU 到 GPU 的传输称为主机到设备（host-to-device，H2D）传输，GPU 到 CPU 的传输称为设备到主机（device-to-host，D2H）传输。因此，这三个 stream 分别是 H2D stream、计算 stream 和 D2H stream。

现在，让我们尝试使用 stream 在 GPU 上异步启动一个批次，并让 CPU 重新获得控制权。在 CPU 端，我们执行以下操作：

在 CPU 上准备批次输入数据（不使用 stream，仅执行 CPU 操作）

将其传输到 GPU（使用 H2D stream）

在 GPU 上执行计算（使用计算 stream）

取回批次输出（使用 D2H stream）

查看结果（不使用 stream）

如果只使用 CUDA stream 来完成这些操作，结果几乎会立刻可用，但结果是错误的。为了理解原因，让我们看看实际发生了什么：

由于各个流彼此独立，三个 GPU 操作几乎同时启动。计算流没有等待 H2D 传输完成，因此前向传播使用的是当时 GPU 内存中已有的数据。D2H 流也没有等待计算完成，因此它传回的是尚未完成计算的结果。步骤 5 会立即返回，因为没有任何操作阻塞 CPU：不存在需要与之同步的默认流。

这些操作单独来看都运行正确。问题在于，我们从未要求这些流相互等待。我们知道，计算必须在 H2D 完成后开始，而 D2H 必须在计算完成后开始，但我们并没有强制执行这一顺序。我们需要一种能够跨越流边界表达“在那个操作完成之前，不要启动这个操作”的机制。

强制同步

为了强制各个流之间进行同步，我们将使用 CUDA 事件。

什么是 CUDA 事件？

CUDA 事件是一种可以记录到流中的标记。当 GPU 在执行过程中到达该标记时，它会将事件设置为已完成。随后，可以要求其他任意流在启动下一个操作之前等待该事件。具体来说，有两个操作：stream.record(event)，它会在流的当前位置插入标记；以及 stream.wait(event)，它会阻止流继续执行，直到该事件被标记为完成。需要特别注意的是，wait 阻塞的是流，而不是 CPU 或其他正在并行运行的流：CPU 调用会立即返回，只有等待中的流会被暂停。

上图展示了一个事件如何同步两个流。CPU 快速连续地发出三个操作（图中的三个小方块）：在流 1 上启动输入准备，在流 1 上记录事件，然后让流 2 等待该事件。之后，CPU 会立即继续执行。流 1 运行其操作，并在操作完成时设置该事件。在此期间，流 2 始终停在等待标记处，只有当事件被标记为完成后才开始计算。CPU 完全没有参与这一过程：执行顺序完全由 GPU 端强制保证。

在连续批处理中使用事件

应用到我们的场景中，修复方法很直接。将 H2D 传输加入队列后，我们调用 h2d_stream.record(h2d_done)：只有传输完成后，该事件才会被标记为已完成。在将前向传播加入队列之前，我们调用 compute_stream.wait(h2d_done)，这样计算流就不会在 h2d_done 被设置之前启动。计算与 D2H 之间也采用相同的方式：通过 model.forward 启动前向传播后，我们调用 compute_stream.record(compute_done)，然后在将输出传输加入队列之前调用 d2h_stream.wait(compute_done)。由此得到一个具有明确执行顺序的流水线：

• H2D 传输在 h2d_stream 上运行

• compute_stream 等待 h2d_done，然后运行前向传播

• d2h_stream 等待 compute_done，然后将输出传回

CPU 按顺序将所有这些操作加入队列，然后继续执行。整个过程中它都不会阻塞。GPU 通过事件强制保证执行顺序，每个流都会在其依赖条件满足后立即开始工作。

上图展示了这一过程如何展开。CPU 准备批次，然后快速将所有 GPU 工作加入队列：H2D 传输、前向传播和 D2H 传输，并在各阶段之间插入 record 和 wait 调用。此后，CPU 便恢复空闲。GPU 接管执行，在相应的依赖事件被设置后，按顺序执行每个流。请注意右侧的绿色标注：D2H 传输完成后，CPU 会回来读取结果。这个最终同步点是整个步骤中 CPU 唯一会阻塞的地方。为了实现它，我们在输出传输之后，在 D2H 流上记录第三个事件，然后在 CPU 端调用 d2h_done_event.synchronize()。synchronize 会阻塞 CPU，直到 D2H 流到达该标记。

这就是它与同步批处理之间的关键区别：以前，CPU 会在每个操作之后阻塞。现在，GPU 工作时，CPU 可以自由地做“某些事情”。我们需要确定这些“某些事情”究竟是什么，因为从 GPU 利用率的角度看，目前还没有发生任何变化。

CPU 的可用时间窗口位于向 GPU 分派批次 N 与分派批次 N+1 之间。最自然的用途是准备批次 N+1 的输入，这样我们就可以将它们分派到 GPU，并在批次 N 的计算完成后立即使用。下面来看看如何实现这一点。

为了准备批次 N+1，我们可以复用准备批次 N 时使用的 CPU 端对象：当前请求列表、缓存状态、主机端张量缓冲区等。不过，我们需要注意两个问题：

• 数据损坏：批次 N+1 的设备端输入缓冲区不能与批次 N 使用相同的缓冲区，否则会破坏 GPU 仍在读取的数据

• 数据传递：如果某个请求同时存在于批次 N 和 N+1 中，并且它在批次 N 的输出中生成了一个新 token，那么批次 N+1 的输入就需要使用该 token

接下来的两节将分别解决数据损坏和数据传递问题。

首先，我们来处理潜在的数据损坏问题。假设批次 N 和批次 N+1 共享同一组设备端输入缓冲区，并且批次 N+1 输入的 H2D 传输在批次 N 仍在计算时就已经开始。此时，GPU 可能仍在从同一块内存中读取批次 N 的输入，而 CPU 已经开始写入批次 N+1 的输入。因此，GPU 可能会读取到部分被覆盖的数据，导致结果损坏。这就是一种竞态条件。主机端也存在同样的风险：如果批次 N 的 H2D 复制仍在进行，却复用了同一个复制源，就会破坏传输的数据。

解决方法是使用两组张量，并在它们之间交替切换。当 GPU 使用槽位 A 处理批次 N 时，CPU 使用批次 N-1 的结果更新请求状态。接下来，CPU 在输入槽位 B 中准备批次 N+1。到了下一个步骤，两者再进行交换。下图展示了这一过程：

当然，这会带来一定成本：用于存储输入和输出张量的 RAM 与 VRAM 用量都会翻倍。不过，这是一项可以接受的权衡，尤其是在使用 FlashAttention 时，因为它不需要注意力掩码，而注意力掩码通常是迄今为止最大的输入张量。

但是，拥有两个槽位又会带来另一个问题。在推理中，我们通常使用 CUDA Graph 来降低延迟。简而言之，CUDA Graph 是一段预先录制的 CUDA 操作序列。它针对特定的内存地址进行录制：为槽位 A 捕获的图无法使用槽位 B 的缓冲区进行重放。因此，我们需要两个图。如果每个图都有自己的内存缓冲区，VRAM 用量就会再次翻倍。

解决方案是使用内存池：两个图都从同一个共享内存缓冲区中分配内存。唯一的约束是，同一个内存池中的两个图绝不能并发执行。由于批次 N 必须在批次 N+1 启动之前完成，因此始终能够满足这一条件。实际中，这两个图加起来使用的 VRAM 几乎与单个图相同。我们只需要在初始化时承担两次捕获的成本。我们可以在同一个内存池中创建任意数量的 CUDA Graph，而总内存用量仍然不会超过所有图中的最大值。下图展示了这一点。

现在，我们已经知道如何防止数据损坏，接下来可以处理第二个问题：如何将批次 N 的输出 token 放入批次 N+1 的输入中。

考虑一个同时出现在批次 N 和批次 N+1 中的请求。在批次 N 中，它会生成一个新 token。这个 token 就是它在批次 N+1 中的输入。问题在于，当我们准备批次 N+1 的输入缓冲区时，还没有拿到这个 token，因为批次 N 仍在运行。为了解决这个问题，我们在构建批次 N+1 时使用一个占位 token。我们会使用 0 作为占位符，原因稍后会变得显而易见。在批次 N 完成计算之后、批次 N+1 开始前向传播之前，我们再替换该占位符。我们将这一步称为“结转”，因为我们正在将新 token 从批次 N 结转到批次 N+1。下图展示了结转背后的思路：

为了执行结转，我们只需要三样东西：批次 N 的输出 token ID、批次 N+1 的输入 token ID，以及一个包含结转执行指令的张量。我们将这个张量称为结转掩码。它包含需要结转的 token 的目标位置，而不需要结转的 token 对应位置则为 -1。下面展示了一个这样的掩码：

结转本身包含四个操作：

• 从批次 N 的输出中选择需要结转的 token，并放入一个新的张量 T

• 将 T 中不需要结转的 token 清零

• 截断 T，使其长度与批次 N+1 的输入长度一致

• 将 T 加到批次 N+1 的输入 ID 上（这就是占位输入 ID 的值为零的原因）

由于这四个操作的开销很小，我们在每个新批次的开始执行它们，并在 CUDA 图中捕获进位。如果进位掩码只包含 -1（-1 值表示：不在此位置进行进位），则最后一步是与零张量的加法。这种情况不常发生，因为跨越多个批次的解码请求通常被调度在连续的批次中。

让我们把所有内容整合在一起，并跟踪前两个步骤。

第 0 步是冷启动：没有前面的批次运行，所以 CPU 在槽位 A 中准备批次 0，并像同步批处理那样进行调度。还没有重叠。

第 1 步是异步循环开始的地方。GPU 现在在槽位 A 上运行批次 0，CPU 是空闲的。它立即开始在槽位 B 中准备批次 1：驱逐已完成的请求、接纳新请求、更新 KV 缓存路由表、构建进位掩码。所有这些都与 GPU 完全重叠运行。一旦批次 1 的输入准备好，CPU 按顺序入队工作：它为槽位 B 启动 H2D 传输，记录并等待计算和 D2H 流的事件，然后继续。

现在 GPU 上发生两件事情并行进行。在槽位 A 上，GPU 完成计算并设置 compute_done，这释放了批次 0 输出的 D2H 传输。在槽位 B 上，批次 1 输入的 H2D 传输正在运行。完成后，h2d_done 事件被设置，批次 1 的计算开始。从批次 0 到批次 1 的进位是该计算的一部分：它发生在常规前向传播之前。由于槽位 A 和槽位 B 是独立的，所有这些都可以自由重叠。

与此同时，CPU 在 d2h_done_event.synchronize() 上阻塞，直到批次 0 的输出到达。然后它处理输出，更新批次 0 中所有请求的状态，并开始调度批次 2。循环现在正在运行，每个后续步骤都遵循完全相同的模式。

我们在下面说明完整的工作负载。每个槽位对 CPU 和 GPU 操作以及事件（也是特定于槽位的）都有专用的颜色。为了可读性，我们不显示 CPU 启动 GPU 操作（如计算或数据移动），但它们仍然发生。这是合理的，因为启动 GPU 操作相对于所示操作的延迟可以忽略不计。

只要当批次 N 完成时批次 N+1 的输入已在 GPU 上准备好，GPU 在批次之间就不会闲置。唯一的问题是 CPU 是否在 GPU 完成计算之前完成其工作。这通常是这样的情况：模型继续增长，而批处理调度保持相对便宜，因此 GPU 计算是瓶颈，而不是 CPU。

它真的有效吗？

为了找出答案，我们运行与之前相同的实验：8K 令牌，批大小 32，8B 模型。

时间线几乎完全是深绿色：CPU 和 GPU 同时运行。偶尔出现的浅绿色细条是 GPU 活跃但 CPU 已完成准备并等待的时刻。几乎看不见的红色标记是批次之间的同步点，CPU 在该点阻塞以采样批次 N 的输出。GPU 的活跃时间占总运行时间的 99.4%，从 76.0% 上升。总生成时间从 300.6 秒下降到 234.5 秒，加速了 22%。如果 CPU 开销完全消除，我们预测为 24% 的加速。剩余的小差距就是那个不可避免的同步点。没有新的内核，没有模型更改：只是让 CPU 和 GPU 同时工作。

我们从一个同步工作负载开始，其中 CPU 和 GPU 一个接一个地工作，导致两者都未被充分利用。通过从基于调度的依赖关系转移到基于数据的依赖关系，并优化同步点，我们设法解缠 CPU 和 GPU 工作负载，使两个硬件的并行执行成为可能。因此，我们能够饱和 GPU 工作队列并确保它始终运行。这最终导致生成速度大幅提高，同时保持模型的准确性。完美的解决方案。

完整的实现在 transformers 库中。如果你想看看这如何转化为实际代码，连续批处理的通用入口点是 continuous_batching.py。更多以异步为中心的代码位于 ContinuousBatchingAsyncIOs 类中。

异步批处理使我们离解锁长生成的 SOTA 吞吐量更近一步，对于强化学习中的 16K+ 生成长度。但仍然需要一些其他较小的东西来达到该目标。在下一篇文章中，我们将讨论这些：请求卸载、解码特定内核或细粒度编译等。敬请期待！

致谢：感谢 Pedro Cuenca 和 Aritra Roy Gosthipaty 的帮助和深入的审阅。

博客中的更多文章

Continuous batching from first principles

Native-speed vLLM transformers modeling backend

---

太棒了！有一个澄清问题，关于 cuda 图共享内存池的好处。根据我的理解，共享内存池将是 Batch(N) + Batch(N+1) 的输入/输出所需的内存。所以每个 cuda 图可能会捕获这个内存并重新使用它（尽管不同的槽位）。共享内存池的好处是它减少了任何额外的 cudagraph 开销吗？

还有另一个问题，N 和 N+1 的批大小不同不是很常见吗？如果是这样，如何准确地为批次 N+1 预分配所需的内存，是否只是最大批大小？

输入和输出实际上是在创建任何 CUDA 图之前预分配的静态张量。

至于为什么内存池很有用：假设你捕获一个 CUDA 图 A：它需要内存来执行。这个内存是为在图 A 中启动的内核存储激活或工作空间张量所需的。这个分配的内存由图 A 拥有。然后，如果你创建另一个图 B，它也会为其执行分配一些内存。因为 CUDA 不能确定你不会同时运行图 A 和 B，为图 A 和 B 分配的内存不能相同：你会冒数据损坏的风险。但如果你能保证图 A 和 B 不会同时运行，那么就没有理由不为图 A 和 B 分配相同的内存。在图 A 和 B 之间共享的那个内存就是内存池。

附注，如果你知道哪个图需要最多的内存，你最好