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

来源：https://huggingface.co/blog/continuous_async 返回文章列表 (https://huggingface.co/blog)

• 同步批处理 (https://huggingface.co/blog/continuous_async#synchronous-batching)
• 创建并发 (https://huggingface.co/blog/continuous_async#creating-concurrency)
  • 什么是 CUDA 流？ (https://huggingface.co/blog/continuous_async#what-is-a-cuda-stream)
  • 默认流与非默认流 (https://huggingface.co/blog/continuous_async#default-and-non-default-streams)
  • 回到连续批处理 (https://huggingface.co/blog/continuous_async#back-to-continuous-batching)
• 强制同步 (https://huggingface.co/blog/continuous_async#enforcing-synchronization)
  • 什么是 CUDA 事件？ (https://huggingface.co/blog/continuous_async#what-is-a-cuda-event)
  • 在连续批处理中使用事件 (https://huggingface.co/blog/continuous_async#using-events-in-continuous-batching)
• 填补真空 (https://huggingface.co/blog/continuous_async#filling-the-vacuum)
  • 竞态条件 (https://huggingface.co/blog/continuous_async#race-conditions)
  • 延续 (https://huggingface.co/blog/continuous_async#carry-over)
• 完整的异步循环 (https://huggingface.co/blog/continuous_async#the-full-async-loop)
• 它真的有效吗？ (https://huggingface.co/blog/continuous_async#does-it-actually-work)
• 结论 (https://huggingface.co/blog/continuous_async#conclusion)

标题图 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/continuous_async/banner.png)

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

这是关于高效 LLM 推理系列的第二篇文章。第一篇 (https://huggingface.co/blog/continuous_batching) 从基本原理出发介绍了连续批处理。它涵盖了我们在此构建所需的一些概念：KV 缓存、FlashAttention、注意力掩码等。

一台 H200 在 Inference Endpoints (https://endpoints.huggingface.co/) 上每小时成本约 5 美元。一小时算便宜，但用一天就要支付 120 美元。既然如此，你肯定希望 GPU 被充分利用。我们已经看到连续批处理通过调度紧密打包的批次来提高 GPU 利用率，这样就不会因填充而浪费算力。但还有第二个浪费源是连续批处理无法解决的：默认情况下，它是同步的。这意味着 CPU 和 GPU 轮流工作：当 GPU 计算时，CPU 等待；当 CPU 准备下一个批次时，GPU 等待。在一个每秒运行数百步的循环中，这些空闲间隙会累积起来，正如我们将展示的，它们可能占总运行时间的近四分之一。为了确保 GPU 100% 的时间都在忙于计算，我们需要消除这些间隙。

为此，我们可以使用异步批处理：我们将把 CPU 的批次准备与 GPU 的批次计算分离开来，这样两者可以并行运行，并且我们始终有一个在工作的 GPU 🔥

https://huggingface.co/blog/continuous_async#synchronous-batching 同步批处理

这就是简单的同步批处理的工作原理：

同步批处理 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/continuous_async/synchronous_batching.png)

当 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 活动时间线 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/continuous_async/cpu_gpu_phases_sync.png)

如果你想生成类似的图表，可以在连续批处理代码中添加埋点，导出 CPU 和 GPU 活动的时间段，并使用此脚本 (https://gist.github.com/remi-or/8de44738629c4d3c72451aa01df1a2ab)。

时间线在绿色（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 (https://github.com/huggingface/transformers) 库中将其实现为连续批处理的一部分。欢迎查看代码并进行比较！

https://huggingface.co/blog/continuous_async#creating-concurrency 创建并发

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

https://huggingface.co/blog/continuous_async#what-is-a-cuda-stream 什么是 CUDA 流？

要理解 CUDA 如何对其操作进行排序，我们需要讨论 CUDA 流。流是 GPU 操作（内核启动、内存复制、同步屏障）的有序队列，这些操作按照提交的顺序执行。每个 GPU 操作总是在某个流中调度。同一流中的操作是顺序执行的：GPU 会等待前一个操作完成后才启动下一个。不同流中的操作相互独立，可以并发运行。举例说明，如果你在 3 个不同的流中启动 3 个操作，执行情况如下：

CUDA 流并发 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/continuous_async/stream_concurrency.png)

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

真实的 CUDA 流并发 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/continuous_async/realistic_concurrency.png)

操作仍然是并发的，但它们的启动时间因每次 CPU 启动的开销而错开。我们将在整个过程中展示这些 CPU 启动事件，因为它们占用真实时间，并且能帮助我们跟踪“何时启动了什么”，以便我们过渡到异步工作流。例如，我们经常检查流是否刷新：这意味着流中的所有操作都已执行完毕。

https://huggingface.co/blog/continuous_async#default-and-non-default-streams 默认流与非默认流

如果你从未在 PyTorch 中显式使用过 CUDA 流，你可能会惊讶它们竟然存在。典型的 PyTorch 脚本从不提及它们，而且你也不会感觉 GPU 操作是异步的：CPU 似乎在 GPU 完成之前一直等待。这种感受是准确的，这要归功于默认流。

当你调用 PyTorch 操作而不指定流时，它会进入默认流。默认流有一个特殊性质：它是同步的。如果某个操作在默认流上调度，它会等待所有其他流被刷新，即 GPU 上的所有工作都必须在默认流上的任何操作开始之前结束。反过来也是如此：任何操作，无论它在哪个流上，都需要等待默认流刷新后才能启动。

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

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

阻塞传输 vs 非阻塞传输 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/continuous_async/block_or_not.png)

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

https://huggingface.co/blog/continuous_async#back-to-continuous-batching 回到连续批处理

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

同步批处理 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/continuous_async/synchronous_batching.png)

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

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

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

关于术语的说明：在谈论 CPU 和 GPU 时，CUDA 文档中通常将 CPU 称为主机 (host)，将 GPU 称为设备 (device)。从现在起我们将沿用这一约定。CPU 到 GPU 的传输称为主机到设备 (host-to-device, H2D) 传输，GPU 到 CPU 的传输称为设备到主机 (device-to-host, D2H) 传输。因此，三个流分别是 H2D 流、计算流和 D2H 流。

现在让我们尝试使用流来异步地在 GPU 上启动一个批次，并立即恢复 CPU 控制。在 CPU 上，我们执行以下操作：

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

如果我们只使用 CUDA 流来执行这些操作，结果几乎会立刻返回，但却是错误的。为了理解原因，让我们看看发生了什么：

失败的异步批处理 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/continuous_async/failed_async.png)

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

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

https://huggingface.co/blog/continuous_async#enforcing-synchronization 强制同步

为了在流之间强制同步，我们将使用 CUDA 事件。

https://huggingface.co/blog/continuous_async#what-is-a-cuda-event 什么是 CUDA 事件？

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

CUDA 事件 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/continuous_async/events.png)

上图显示了一个事件同步两个流的过程。CPU 快速连续发出三个操作（三个小方块）：在流 1 上启动输入准备，在流 1 上记录事件，然后告诉流 2 等待该事件。然后 CPU 立即继续。流 1 运行其操作，完成后事件被设置。流 2 在此期间一直停留在等待标记处，只有在事件被标记为完成后才开始计算。CPU 并未参与其中：整个顺序完全在 GPU 端强制执行。

https://huggingface.co/blog/continuous_async#using-events-in-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)。结果就是一个具有显式顺序的流水线：

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

CPU 按顺序将所有这些放入队列，然后继续执行。CPU 在任何时候都不会阻塞。GPU 通过事件强制执行顺序，一旦它们的依赖条件得到满足，所有三个流都会活跃起来。