解锁连续批处理中的异步性

标题：解锁连续批处理中的异步能力

URL 来源：https://huggingface.co/blog/continuous_async

发布时间：2026-05-14T00:00:00.697Z

Markdown 内容：

_TL;DR：我们将解释如何分离 CPU 和 GPU 工作负载，以显著提升推理性能。_

_这是关于高效 LLM 推理系列的第二篇文章。第一篇文章从基本原理介绍了连续批处理。它介绍了一些我们将在此基础上构建的概念：KV 缓存、FlashAttention、注意力掩码等。_

在 Inference Endpoints 上，H200 每小时的成本约为 5 美元。对于一小时来说这很便宜，但如果你使用一整天，费用就达到了 120 美元。在这种情况下，你肯定希望 GPU 能够被充分利用。

我们已经看到，连续批处理通过调度紧密打包的批次来提高 GPU 利用率，从而避免在填充上浪费计算资源。但连续批处理并未解决第二个浪费来源：默认情况下，它是同步的。这意味着 CPU 和 GPU 轮流工作：GPU 计算时，CPU 等待；CPU 准备下一个批次时，GPU 等待。在一个每秒运行数百步的循环中，这些空闲间隙会累积起来，正如我们将要展示的，它们可能占总运行时间的近四分之一。为了确保 GPU 100% 的时间都在忙于计算，我们需要消除这些间隙。

为了实现这一点，我们可以使用异步批处理：我们将 CPU 批次准备与 GPU 批次计算解耦，使两者能够并行运行，从而始终保持 GPU 的高效工作 🔥

同步批处理

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

当 CPU 准备一个新批次时，它会选择要包含的请求，更新 KV 缓存表，驱逐在前几轮运行中完成的请求，并接纳新的请求以填充释放的空间。完成后，它将准备好的输入传输到 GPU。GPU 运行其前向传播并为每个请求采样（即选择）一个新令牌。结果返回给 CPU，这样 CPU 就知道每个请求刚刚生成了什么令牌，然后整个循环再次重复。

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

这就是同步批处理的核心低效之处：CPU 和 GPU 轮流工作。GPU 计算时，CPU 空闲；CPU 更新时，GPU 空闲。在任何情况下，它们都不能同时进行有效工作。对于单次前向传播来说，这可能看起来代价不大，但在一个每秒运行数百步的连续批处理循环中，这些空闲间隙累积起来会导致实际的吞吐量损失。

为了展示这一点，我们分析了在使用 8B 模型、批次大小为 32 生成 8K 令牌时，CPU 和 GPU 所花费的时间：

_如果你想生成类似的图表，可以修改连续批处理代码以转储 CPU 和 GPU 活动跨度，并使用此脚本。_

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

乐观的看法是，如果我们能完全消除 CPU 开销，生成时间将从 300 秒减少到 228 秒（免费获得 24% 的速度提升！）。这不需要任何新的内核或模型更改，只需仔细协调硬件。

从根本上说，这个想法很简单：我们需要找出如何在批次 N 计算的同时，为批次 N+1 进行批次准备。但这个简单的想法隐藏了一些技术难点：

• 我们如何在 GPU 上启动任务并让 CPU 重新获得控制权？
• 我们如何确保在每个任务启动时，数据（无论是 CPU 还是 GPU 任务所需的数据）已经准备就绪？
• 如果批次 N+1 是基于批次 N 的预测结果来准备的，我们该如何处理？

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

创建并发性

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

什么是 CUDA 流？

要理解 CUDA 如何调度其操作，我们需要谈谈 CUDA 流。流是一个有序的 GPU 操作队列（包括内核启动、内存复制、同步屏障等），这些操作会按照提交顺序依次执行。每个 GPU 操作总是被调度到某个流中执行。同一流内的操作是顺序执行的：GPU 必须等待前一个操作完成后才会开始下一个。而*不同*流中的操作则相互独立，可以并发执行。举例来说，如果在 3 个不同的流中各启动一个操作，执行情况如下所示：

这三个操作会同时启动。这里稍作简化：实际上每个 GPU 操作最终都是由 CPU 发起的，而启动过程需要花费少量时间：查找正确的内核、发出调用、将命令从 CPU 传输到 GPU 等。这被称为 CPU 启动开销，更符合实际的示意图如下：

操作仍然是并发的，但它们的启动时间因每次 CPU 启动的开销而错开。我们将在后续内容中持续展示这些 CPU 启动事件，因为它们占用实际时间，并且在我们转向异步工作流时有助于追踪“何时启动了何操作”。例如，我们经常需要检查一个流是否已被清空：这意味着该流中的所有操作都已执行完毕。

默认流与非默认流

如果你从未在 PyTorch 中显式使用过 CUDA 流，可能会惊讶于它们的存在。典型的 PyTorch 脚本从不提及流，而且 GPU 操作看起来并不*像是*异步的：CPU 似乎会等待 GPU 完成操作后才继续执行。这种感觉是准确的，它源于默认流的特性。

当你在不指定流的情况下调用 PyTorch 操作时，该操作会进入默认流。默认流有一个特殊属性：它是同步的。如果某个操作被调度到默认流上，它会等待所有其他流被清空，即 GPU 上的所有工作必须全部完成，默认流上的单个操作才能开始。反之亦然：任何操作，无论它在哪个流中，在启动前都会等待默认流被清空。

因此，如果你将默认流操作的结果传输到 CPU，即使该传输本应对 CPU 是非阻塞的，你的 CPU 仍然会阻塞，直到所有 GPU 操作完成，因为这些操作被调度在默认流上。这实际上破坏了任何构建并发性的努力。

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

在本文的后续部分，我们将假设所有设备间的内存传输都是非阻塞的。因此，我们需要自行对它们进行同步。

回到连续批处理

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

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

1. 将输入从 CPU 传输到 GPU
2. 在 GPU 上进行计算
3. 将输出从 GPU 传输回 CPU

这意味着我们需要三个流：一个用于计算，一个用于 CPU 到 GPU 的传输，一个用于 GPU 到 CPU 的传输。传输操作是独立的，没有理由将它们串行化，因此每个传输操作都有自己的流。

关于命名法的说明：在讨论 CPU 和 GPU 时，整个 CUDA 文档中使用的约定是将 CPU 称为主机，将 GPU 称为设备。从现在开始，我们将采用这一约定。CPU 到 GPU 的传输称为主机到设备传输，GPU 到 CPU 的传输称为设备到主机传输。因此，这三个流分别是 H2D 流、计算流和 D2H 流。

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

1. 在 CPU 上准备批次输入数据（无流，仅 CPU 操作）
2. 将其传输到 GPU（使用 H2D 流）
3. 在 GPU 上运行计算（使用计算流）
4. 检索批次输出（使用 D2H 流）
5. 查看结果（无流）

如果我们仅使用 CUDA 流来实现这一点，结果几乎会立即返回，但却是错误的。要理解原因，让我们看看发生了什么：

由于流彼此独立，三个 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)。最终形成一个具有明确顺序的流水线：

1. H2D 传输在 h2d_stream 上运行
2. compute_stream 等待 h2d_done，然后运行前向传播
3. 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 工作时自由执行“某些操作”。

我们需要弄清楚这个“某些操作”是什么，因为目前从 GPU 利用率的角度来看，没有任何改变。

填补空白

CPU 可用的时间窗口位于调度批次 N 和调度批次 N+1 到 GPU 之间。其自然用途是准备批次 N+1 的输入，这样我们可以在批次 N 计算结束后立即将它们调度到 GPU 并使其就绪。接下来让我们看看如何实现这一点。

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

• 数据损坏：批次 N+1 的设备端输入缓冲区不能与批次 N 相同，否则会破坏 GPU 仍在读取的数据
• 数据传输：如果某个请求同时存在于批次 N 和 N+1 中，并且在批次 N 的输出中生成新令牌，则该令牌需要出现在批次 N+1 的输入中

我们将在接下来两节中解决数据损坏和数据传输这两个问题。

竞态条件

首先，我们要解决潜在的数据损坏问题。

假设批次 N 和批次 N+1 共享相同的设备端输入缓冲区，且批次 N+1 的 H2D 传输在批次 N 仍在计算时启动。CPU 可能在 GPU 仍从同一内存读取批次 N 数据时写入批次 N+1 的输入。这将导致 GPU 可能读取到部分被覆盖的数据，造成数据损坏。这就是竞态条件。主机端同样存在此风险：在批次 N 的 H2D 复制尚未完成时复用同一复制源会破坏传输。

解决方案是使用两组张量并交替使用。当 GPU 从槽 A 处理批次 N 时，CPU 使用批次 N-1 的结果更新请求状态。接着 CPU 在输入槽 B 中准备批次 N+1。下一步，两者交换位置。下图展示了这一过程：

当然这需要代价：存储输入输出张量所需的 RAM 和 VRAM 将翻倍。这是可接受的权衡，特别是在使用 FlashAttention 时，因为它不需要注意力掩码——这是目前最大的输入张量。

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

解决方案是内存池：两个图从中分配内存的共享内存缓冲区。唯一约束是同一池中的两个图绝不能并发执行。由于批次 N 必须在批次 N+1 开始前完成，这一条件始终满足。实践中，两个图共同使用的 VRAM 量与单个图几乎相同。我们仅在初始化时承担两次捕获的开销。

我们可以在同一内存池中创建任意数量的 CUDA 图，总内存使用量仍受限于各图中的最大值。下图展示了这一机制：

现在我们已经知道如何防止数据损坏，接下来可以解决第二个问题：将批次 N 的输出令牌传递到批次 N+1 的输入中。

令牌传递

考虑同时出现在批次 N 和 N+1 中的请求。在批次 N 中，它产生新令牌。该令牌将成为批次 N+1 的输入。问题在于当我们准备批次 N+1 的输入缓冲区时，尚未获得该令牌：批次 N 仍在运行。为此，我们在构建批次 N+1 时使用占位符令牌。我们将使用 0 作为占位符，具体原因后文会说明。在批次 N 计算完成之后、批次 N+1 开始前向传播之前，我们会替换该占位符。我们将此步骤称为令牌传递，因为这是将新令牌从批次 N 传递到批次 N+1。令牌传递的原理如下图所示：

执行令牌传递只需三要素：批次 N 的输出令牌 ID、批次 N+1 的输入令牌 ID，以及指示如何执行传递的张量。我们将此张量称为传递掩码。它包含需要传递令牌的目标位置，不需要的位置标记为 -1。下图展示了其结构：

令牌传递包含四个操作步骤：

• 从批次 N 的输出中筛选需要传递的令牌至新张量 T
• 将 T 中不需要传递的令牌清零
• 截断 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+1 的输入在批次 N 完成时已经在 GPU 上准备就绪，GPU 在批次之间就不会空闲。唯一的问题是 CPU 是否能在 GPU 完成计算之前完成其工作。通常情况确实如此：模型规模持续增长，而批次调度的开销相对较小，因此 GPU 计算是瓶颈，而不是 CPU。

实际效果如何？

为了验证这一点，我们运行了与之前相同的实验：8K tokens，批次大小为 32，使用 8B 模型。

时间线几乎完全被深绿色覆盖：CPU 和 GPU 同时运行。偶尔出现的浅绿色细条表示 GPU 处于活动状态，但 CPU 已经完成准备工作并处于等待状态。几乎不可见的红色标记是批次之间的同步点，CPU 在此处阻塞以采样批次 N 的输出。GPU 在总运行时间中的活动占比从 76.0% 提升至 99.4%。总生成时间从 300.6 秒降至 234.5 秒，提升了 22% 的速度。我们之前预测如果完全消除 CPU 开销，速度将提升 24%。剩余的小差距是不可避免的同步点。没有新增内核，没有修改模型：仅仅是让 CPU 和 GPU 同时工作。

结论

我们从同步工作负载开始，其中 CPU 和 GPU 依次工作，导致两者均未充分利用。通过从基于调度的依赖转向基于数据的依赖，并优化同步点，我们成功解耦了 CPU 和 GPU 的工作负载，使两者的并行执行成为可能。因此，我们能够饱和 GPU 的工作队列，确保其始终运行。最终，在保持模型准确性的同时，大幅提升了生成速度。这几乎是一次完美的成功。

完整实现位于 transformers 代码库中。若想了解具体代码实现，持续批处理的通用入口点位于 continuous_batching.py 文件。而更偏向异步处理的代码则位于 ContinuousBatchingAsyncIOs 类中。

异步批处理让我们在实现长文本生成（如强化学习中16K+生成长度）的SOTA吞吐量目标上更近一步。但要完全达成这一目标，仍需其他一些技术要素。在下一篇文章中，我们将深入探讨这些内容：请求卸载、解码专用内核或细粒度编译等技术。敬请关注！

*致谢：衷心感谢 Pedro Cuenca 和 Aritra Roy Gosthipaty 的宝贵帮助与深刻见解。*