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

](https://huggingface.co/ror)

• 同步批处理

• 创建并发

  • 什么是 CUDA stream？

  • 默认 stream 与非默认 stream

  • 回到 Continuous Batching

• 强制同步

  • 什么是 CUDA event？

  • 在 Continuous Batching 中使用 events

• 填补真空

  • 竞态条件

  • 结转

• 完整的异步循环

• 它真的有效吗？

• 结论

TL;DR：我们解释了如何分离 CPU 和 GPU 的工作负载，从而大幅提升推理性能。

这是关于高效 LLM 推理系列的第二篇文章。第一篇文章从基本原理出发介绍了 continuous batching。它引入了一些我们在此基础之上构建的概念：KV cache、FlashAttention、attention mask 等。

一台 H200 在 Inference Endpoints 上每小时大约花费 5 美元。一小时听起来不贵，但用一天就要支付 120 美元。既然如此，你肯定希望 GPU 被充分利用。

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

为此，我们可以使用异步批处理：我们将 CPU 的 batch 准备工作与 GPU 的 batch 计算解耦，这样两者可以并行运行，GPU 始终处于高效工作状态 🔥

同步批处理

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

当 CPU 准备一个新的 batch 时，它会选择包含哪些请求，更新 KV cache 表，驱逐在前几次运行中已完成的请求，并接纳新请求以填充释放的空间。完成后，它将准备好的输入传输到 GPU。GPU 运行其 forward pass 并为每个请求采样（即选择）一个新 token。结果返回给 CPU，以便 CPU 知道每个请求刚刚产生了哪个 token，然后整个循环再次重复。

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

这就是同步批处理的核心低效之处：CPU 和 GPU 轮流工作。GPU 计算时，CPU 空闲；CPU 更新时，GPU 空闲。在任何情况下，它们都不会同时做有用功。对于单次 forward pass，这似乎代价不大，但在一个每秒运行数百步的 continuous batching 循环中，这些空闲间隙会累积成实际的吞吐量损失。

为了说明这一点，我们分析了在使用 8B 模型、batch size 为 32、生成 8K token 时，CPU 和 GPU 各自花费的时间：

如果你想生成类似的图表，可以对 continuous batching 代码进行插桩，以转储 CPU 和 GPU 活动区间，并使用这个脚本。

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

乐观的观点是，如果我们能完全消除 CPU 开销，生成时间将从 300 秒降至 228 秒（免费获得 24% 的加速！）。这不需要任何新的 kernel 或模型更改，只需要仔细协调硬件。

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

• 我们如何在 GPU 上启动某些操作，同时将控制权交还给 CPU？
• 我们如何确保在启动每个任务时，数据对于 CPU 或 GPU 任务来说已经准备就绪？
• 如果 batch N+1 基于 batch N 的预测结果，我们如何准备它？

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

创建并发

我们的最终目标是实现 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 launch overhead，更现实的图景如下：

操作仍然是并发的，但它们的开始时间因每次 CPU 启动的开销而错开。我们将在整个过程中继续展示这些 CPU 启动事件，因为它们会占用实际时间，并且在我们转向异步工作流时，它们将帮助我们追踪“何时启动了何物”。例如，我们经常会检查一个 stream 是否被 flushed：这意味着该 stream 中的所有操作都已被执行。

默认 stream 与非默认 stream

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

当你调用 PyTorch 操作而未指定 stream 时，它会落在默认 stream 上。默认 stream 有一个特殊属性：它是同步的。如果一个操作被调度到默认 stream 上，它会等待所有其他 stream 被 flushed，即 GPU 上的所有工作必须在默认 stream 上的单个操作开始之前结束。反之亦然：任何操作，无论其属于哪个 stream，都会在启动前等待默认 stream 被 flushed。

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

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

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

回到 Continuous Batching

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

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

1. 将输入从 CPU 传输到 GPU
2. 在 GPU 上计算
3. 将输出从 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、compute stream 和 D2H stream。

现在，让我们尝试使用 streams 在 GPU 上异步启动一个 batch 并取回 CPU 控制权。从 CPU 的角度，我们执行以下操作：

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

如果我们只使用 CUDA streams 来执行此操作，结果几乎会立即返回，并且是不正确的。要理解原因，让我们看看发生了什么：

由于 streams 彼此独立，所有三个 GPU 操作几乎同时启动。compute stream 没有等待 H2D 传输完成，因此 forward pass 在 GPU 内存中已有的任何数据上运行。D2H stream 没有等待计算完成，因此它传输了尚未计算出的结果。第 5 步立即返回，因为没有任何东西阻塞 CPU：没有默认 stream 需要同步。

这些操作在隔离状态下都正确运行。问题在于我们从未告诉 streams 要相互等待。我们知道计算必须在 H2D 完成后开始，D2H 必须在计算完成后开始，但我们没有强制执行该顺序。我们需要一种机制来跨 stream 边界说“在这个操作完成之前，不要开始那个操作”。

强制同步

为了在 streams 之间强制同步，我们将使用 CUDA events。

什么是 CUDA event？

CUDA event 是一个可以记录到 stream 中的标记。当 GPU 在执行过程中到达该标记时，它会将该 event 设置为已完成。然后可以告诉任何其他 stream 在开始其下一个操作之前等待该 event。具体来说，有两个操作：stream.record(event)，它将标记插入到 stream 的当前位置；以及 stream.wait(event)，它阻止一个 stream 继续执行，直到该 event 被标记为完成。重要的是，wait 阻塞的是 stream，而不是 CPU 或其他并行运行的 stream：CPU 调用会立即返回，只有等待的 stream 被延迟。

上图显示了一个 event 同步两个 stream。CPU 快速连续发出三个操作（三个小方块）：在 stream 1 上启动输入准备，在 stream 1 上记录 event，然后告诉 stream 2 等待它。然后 CPU 立即继续。Stream 1 运行其操作，完成后，event 被设置。Stream 2 在此期间一直被阻塞在 wait 标记处，只有在 event 被标记为完成后才开始计算。CPU 没有参与其中任何一个：排序完全在 GPU 端强制执行。

在 Continuous Batching 中使用 events

应用于我们的情况，修复方法很简单。在将 H2D 传输入队后，我们调用 h2d_stream.record(h2d_done)：只有当传输完成时，该 event 才会被标记为完成。在将 forward pass 入队之前，我们调用 compute_stream.wait(h2d_done)，这样 compute stream 在 h2d_done 被设置之前不会开始。我们在计算和 D2H 之间也做同样的事情：在通过 model.forward 启动 forward pass 后，我们调用 compute_stream.record(compute_done)，然后在将输出传输入队之前调用 d2h_stream.wait(compute_done)。结果是一个具有显式排序的流水线：

1. H2D 传输在 h2d_stream 上运行
2. compute_stream 等待 h2d_done，然后运行 forward pass
3. d2h_stream 等待 compute_done，然后将输出传输回来

CPU 按顺序将所有内容入队，然后继续。它从未阻塞。GPU 通过 events 强制执行排序，并且所有三个 stream 在其依赖项满足后立即变为活跃状态。

上图展示了这是如何展开的。CPU 准备 batch，然后快速将所有 GPU 工作入队：H2D 传输、forward pass、D2H 传输，并在每个阶段之间插入 record 和 wait 调用。之后，CPU 就空闲了。GPU 接管，按顺序执行每个 stream，因为其依赖的 event 已被设置。注意右侧的绿色标注：一旦 D2H 传输完成，CPU 回来读取结果。这个最终的同步是整个步骤中 CPU 唯一阻塞的点。为了实现它，我们在输出传输后在 D2H stream 上记录第三个 event，然后在 CPU 端调用 d2h_done_event.synchronize()。synchronize 会阻塞 CPU，直到 D2H stream 到达该标记。

这是与同步批处理的关键区别：以前，CPU 在每个操作后都会阻塞。现在，它可以在 GPU 工作时自由地做“某事”。

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

填补真空

CPU 可用的时间窗口位于将 batch N 分派给 GPU 和分派 batch N+1 之间。它的自然用途是准备 batch N+1 的输入，这样我们就可以将它们分派给 GPU，并在 batch N 计算结束时让它们准备就绪。让我们看看如何做到这一点。

为了准备 batch N+1，我们可以重用准备 batch N 时使用的相同 CPU 端对象：当前请求列表、cache 状态、host 端张量缓冲区等。但是，我们需要注意两件事：

• 数据损坏：batch N+1 的设备端输入缓冲区不能与 batch N 的相同：我们会损坏 GPU 仍在读取的数据
• 数据传输：如果一个请求同时出现在 batch N 和 N+1 中，并且它在 batch N 的输出中产生了一个新 token，那么该 token 在 batch N+1 的输入中是必需的

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

竞态条件

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

假设 batch N 和 batch N+1 共享相同的设备端输入缓冲区，并且 batch N+1 输入的 H2D 传输在 batch N 仍在计算时开始。CPU 可能会在 GPU 仍在从同一内存读取 batch N 的输入时写入 batch N+1 的输入。因此，GPU 可能会读取到部分被覆盖的数据，结果就是数据损坏。这是一个竞态条件。在 host 端也存在同样的风险：在 batch N 的 H2D 拷贝仍在进行时重用相同的拷贝源会损坏传输。

解决方法是使用两组张量并在它们之间交替。当 GPU 从插槽 A 处理 batch N 时，CPU 使用 batch N-1 的结果更新请求状态。CPU 接下来在输入插槽 B 中准备 batch N+1。下一步，它们交换。下图说明了这一点：

当然，这是有代价的：它使存储输入和输出张量所需的 RAM 和 VRAM 翻倍。这是一个可以接受的权衡，尤其是在使用 FlashAttention 时，因为它不需要 attention mask，而 attention mask 是迄今为止最大的输入张量。

但是拥有两个插槽会带来另一个问题。在推理中，我们通常使用 CUDA graphs 来减少延迟。简而言之，CUDA graph 是一个预先记录的 CUDA 操作序列。它是针对特定内存地址记录的：为插槽 A 捕获的 graph 不能针对插槽 B 的缓冲区重放。因此我们需要两个 graphs。如果每个 graph 都有自己的内存缓冲区，那又是双倍的 VRAM。

解决方案是内存池：一个共享的内存缓冲区，两个 graphs 都从中分配。唯一的约束是同一池中的两个 graphs 绝不能同时执行。由于 batch N 必须在 batch N+1 开始之前完成，这总是成立的。在实践中，两个 graphs 一起使用的 VRAM 几乎与一个 graph 相同。我们只在初始化时付出两次捕获的代价。

我们可以在同一个池中创建任意数量的 CUDA graphs，总内存使用量仍然以 graphs 中的最大值为上限。如下所示。

现在我们知道如何防止数据损坏了，我们可以解决第二个问题：将 batch N 的输出 token 放入 batch N+1 的输入中。

结转

考虑一个同时出现在 batch N 和 batch N+1 中的请求。在 batch N 中，它产生一个新 token。该 token 是它在 batch N+1 中的输入。问题在于，当我们准备 batch N+1 的输入缓冲区时，我们还没有那个 token：batch N 仍在运行。为了解决这个问题，我们在构建 batch N+1 时使用一个占位符 token。我们将使用 0 作为占位符，原因稍后会变得明显。我们在 batch N 计算完成之后、batch N+1 开始 forward pass 之前替换该占位符。我们将这一步称为结转，因为我们正在将新 token 从 batch N 结转到 batch N+1。结转背后的想法如下所示：

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

结转本身由四个操作组成：

• 我们从 batch N 的输出中选择要结转的 token，放入一个新张量 T
• 我们将 T 中不想结转的 token 置零
• 我们截断 T 以匹配 batch N+1 的输入长度
• 我们将 T 加到 batch N+1 的输入 id 上（这就是为什么占位符输入 id 的值为零）

由于这四个操作非常廉价，我们在每个新 batch 开始时执行它们，并将结转捕获到 CUDA graph 中。如果结转掩码只包含 -1（值为 -1 意味着：不要结转此位置），那么最后一步就是与零张量相加。这并不经常发生，因为跨越多个 batch 的解码请求通常会被调度到连续的 batch 中。

完整的异步循环

让我们把所有内容整合起来，并追踪前两步。

第 0 步是冷启动：没有前一个 batch 在运行，因此 CPU 在插槽 A 中准备 batch 0，并像同步批处理一样分派它。还没有重叠。

第 1 步是异步循环开始的地方。GPU 现在在插槽 A 上运行 batch 0，CPU 空闲。它立即开始在插槽 B 中准备 batch 1：驱逐已完成的请求、接纳新请求、更新 KV cache 路由表、构建结转掩码。所有这些都与 GPU 完全重叠运行。一旦 batch 1 的输入准备就绪，CPU 按顺序将工作入队：它为插槽 B 启动 H2D 传输，记录并等待 compute 和 D2H streams 的 events，然后继续。

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

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

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

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

它真的有效吗？

为了找出答案，我们运行了与之前相同的实验：8K token，batch size 32，8B 模型。

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

结论

我们从 CPU 和 GPU 依次工作的同步工作负载开始，导致两者都未被充分利用。通过从基于调度的依赖关系转向基于数据的依赖关系并细化同步点，我们成功地将 CPU 和 GPU 的工作负载解耦，使得两个硬件可以并行执行。因此，我们能够饱和 GPU 工作队列，确保它始终在运行。这最终在保持模型准确性的同时，大幅提升了生成速度。这几乎是一个稳赢的局面。

完整的实现在 transformers 库中。如果你想了解这如何转化为实际代码，continuous batching 的通用入口点是 continuous_batching.py。更侧重于异步的代码位于 ContinuousBatchingAsyncIOs 类中。

异步批处理让我们离解锁长序列生成（例如强化学习中 16K+ 的生成长度）的 SOTA 吞吐量更近了一步。但还有一些其他较小的事情也需要实现才能达到这个目标。在下一篇文章中，我们将介绍这些内容：卸载请求、解码专用 kernel 或细粒度编译等。敬请期待！

致谢：非常感谢 Pedro Cuenca 和 Aritra Roy Gosthipaty 的帮助和富有洞察力的审阅。