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

这是高效大语言模型推理系列文章的第二篇。第一篇从基本原理出发介绍了连续批处理，其中引入了一些我们将在此基础上深入探讨的概念：KV 缓存、FlashAttention、注意力掩码等。

在 Inference Endpoints 上，一块 H200 每小时的成本约为 5 美元。按小时算确实不贵，但如果用上一整天，就要支付 120 美元。既然如此，你当然希望 GPU 能被充分利用。

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

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

同步批处理

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

当 CPU 准备一个新批次时，它会选择要包含哪些请求，更新 KV 缓存表，驱逐上一轮运行中已完成的请求，并接纳新请求来填补释放的空间。完成这些操作后，它将准备好的输入传输到 GPU。GPU 执行前向传播并对每个请求进行采样（即选择）一个新 token。结果返回给 CPU，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% 的加速！）。这不需要任何新的内核或模型改动，只需对硬件进行精心协调。

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

• 我们如何在 GPU 上启动任务，同时将控制权交还给 CPU？
• 我们如何确保在 CPU 或 GPU 任务启动时，数据已经准备就绪？
• 如果批次 N+1 是基于批次 N 的预测结果，我们该如何准备它？

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

创建并发

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

什么是 CUDA 流？

要理解 CUDA 如何对其操作进行排序，我们需要讨论 CUDA 流。流是一个有序的 GPU 操作队列（包括内核启动、内存拷贝、同步屏障），这些操作按照提交的顺序执行。每个 GPU 操作总是在某个流中被调度。同一流中的操作是顺序执行的：GPU 在前一个操作完成之前不会启动下一个操作。不同流中的操作彼此独立，可以并发运行。举例来说，如果你在 3 个不同的流中启动 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 称为主机（host），将 GPU 称为设备（device）。从现在起我们将沿用这一惯例。CPU 到 GPU 的数据传输称为主机到设备（H2D）传输，GPU 到 CPU 的数据传输称为设备到主机（D2H）传输。因此，这三个流分别是 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 传输，并在每个阶段之间插入记录和等待调用。之后，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，那么这个 token 需要出现在第 N+1 个批次的输入中。

我们将在接下来的两个部分中解决这些问题，即数据损坏和数据传输。

竞态条件

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

假设批次 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 图来降低延迟。简而言之，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 开始前向传播之前，我们会替换掉这个占位符。我们将这一步称为“结转”（carry-over），因为我们正在将批次 N 的新 token 结转到批次 N+1。结转背后的思路如下图所示：

要执行结转，我们只需要三样东西：批次 N 的输出 token ID、批次 N+1 的输入 token ID，以及一个包含如何执行结转指令的张量。我们将这个张量称为结转掩码（carry-over mask）。它包含需要结转的 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 处于空闲状态。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 token，批次大小 32，8B 模型。

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

Conclusion

我们从一个同步工作负载开始，CPU 和 GPU 依次运行，导致两者都未得到充分利用。通过将基于调度的依赖关系转变为基于数据的依赖关系，并优化同步点，我们成功解耦了 CPU 和 GPU 的工作负载，使两者能够并行执行。因此，我们能够填满 GPU 的工作队列，确保它始终处于运行状态。这最终大幅提升了生成速度，同时保持了模型的准确性。这几乎是一个毫无悬念的胜利。

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

异步批处理让我们向解锁长序列生成（如强化学习中 16K+ 的生成长度）的 SOTA 吞吐量又迈进了一步。但要实现这一目标，还需要一些其他更小的改进。在下一篇文章中，我们将介绍这些内容：请求卸载、解码专用内核或细粒度编译等。敬请期待！

致谢：衷心感谢 Pedro Cuenca 和 Aritra Roy Gosthipaty 提供的帮助和富有洞见的审阅。