Hugging Face 详解连续批处理中的异步化：用 CUDA 流与事件消除 CPU-GPU 串行等待

简而言之：我们解释如何分离 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 运行其前向传播，并为每个请求采样（即选择）一个新 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 streams 来实现这一点。

什么是 CUDA stream？

要理解 CUDA 如何对其操作进行排序，我们需要谈谈 CUDA streams。一个 stream 是 GPU 操作（kernel 启动、内存拷贝、同步屏障）的有序队列，按照提交顺序执行。每个 GPU 操作总是在某个 stream 中调度。同一 stream 内的操作是顺序执行的：GPU 不会在前一个操作完成之前启动下一个。而不同 stream 中的操作彼此独立，可以并发运行。举例来说，如果你在 3 个不同的 stream 中启动 3 个操作，执行情况如下：

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

这些操作仍然是并发的，但它们的启动时间因每次 CPU 启动的开销而错开。我们会在后续内容中持续展示这些 CPU 启动事件，因为它们确实消耗时间，并且在我们转向异步工作流时，它们能帮助我们追踪“什么在何时被启动”。例如，我们经常会检查一个 stream 是否已刷新（flushed）：这意味着该 stream 中的所有操作都已执行完毕。

默认 stream 和非默认 stream

如果你在 PyTorch 中从未显式使用过 CUDA streams，你可能会惊讶于它们的存在。典型的 PyTorch 脚本从不提及它们，而且 GPU 操作感觉上也不是异步的：CPU 似乎在继续之前会等待 GPU 完成。这种感觉是准确的，它来自于 默认 stream。

当你在不指定流的情况下调用 PyTorch 操作时，它会被放到默认流上。默认流有一个特殊属性：它是同步的。如果某个操作被安排在默认流上，它会等待所有其他流被刷新完毕，也就是说，默认流上的单个操作开始之前，GPU 上的所有工作都必须已经结束。反过来也同样成立：任何操作，无论属于哪个流，都会等待默认流被刷新后才启动。

因此，如果你把默认流操作的结果传输回 CPU，即使该传输本应对 CPU 是非阻塞的，你的 CPU 仍会阻塞，直到所有 GPU 操作都完成为止，因为这些操作被安排在默认流上。这实际上摧毁了任何构建并发的努力。

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

在本文余下的部分，我们将假设所有从一个设备到另一个设备的内存传输都是非阻塞的。因此，我们必须自己进行同步。

回到连续批处理

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

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

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

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

关于命名的一点说明：在谈论 CPU 和 GPU 时，整个 CUDA 文档中使用的惯例是将 CPU 称为主机，将 GPU 称为设备。从现在起我们将沿用这一惯例。CPU 到 GPU 的传输称为主机到设备（H2D）传输，GPU 到 CPU 的传输称为设备到主机（D2H）传输。因此，这三个流分别是 H2D 流、计算流和 D2H 流。

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

在 CPU 上准备批次输入数据（无流，仅 CPU 操作）
将其传输到 GPU（使用 H2D 流）
在 GPU 上运行计算（使用计算流）
取回批次输出（使用 D2H 流）
查看结果（无流）

如果我们只使用 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 侧强制执行的。

在 Continuous Batching 中使用事件

应用到我们的场景中，修复很直接。在将 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 可用的窗口位于将批次 N 分派到 GPU 和将批次 N+1 分派到 GPU 之间。它自然的用途是准备批次 N+1 的输入，这样我们就可以将它们分派到 GPU，并在批次 N 的计算结束后让它们准备就绪。让我们看看如何做到这一点。

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

数据损坏：批次 N+1 的设备侧输入缓冲区不能与批次 N 的相同：否则我们会破坏 GPU 仍在读取的数据
数据传输：如果一个请求同时存在于批次 N 和 N+1 中，并且它在批次 N 的输出中产生了一个新 token，那么该 token 在批次 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 批的输出 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 的目标位置，对于不需要传递的则为 -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。循环此时已经运转起来，此后每一步都遵循完全相同的模式。
我们在下面展示了完整的工作负载。每个 slot 对 CPU 和 GPU 操作以及事件（也是 slot 专属的）都有专门的配色。为便于阅读，我们没有展示 CPU 对 GPU 操作（如计算或数据搬运）的启动，但它们依然会发生。这是合理的，因为与图中展示的操作相比，启动一次 GPU 操作的延迟可以忽略不计。

只要批次 N+1 的输入在批次 N 完成时已在 GPU 上就绪，GPU 就不会在批次之间空闲。唯一的问题是 CPU 能否在 GPU 完成计算之前完成其工作。通常情况正是如此：模型不断增大，而批次调度相对仍然廉价，因此 GPU 计算才是瓶颈，而不是 CPU。

它真的有效吗？

为了验证，我们运行与之前相同的实验：8K tokens，batch size 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 的帮助和富有洞见的审阅。